dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

Revert "feat: 10-file algorithm factor system — full 10-layer AST call chain"

Browse Source 
This reverts commit 3581dd5435.
This commit is contained in:
Claude
2026-08-17 05:32:19 +00:00
parent 79562c342d
commit 461addf428
10 changed files with 0 additions and 2553 deletions
Download Patch File Download Diff File Expand all files Collapse all files

270
ex_engine/factors/attn_dispatch.py
View File

@@ -1,270 +0,0 @@
"""
ex_engine/factors/attn_dispatch.py
Layer 3: Attention prefill/decode dispatch
Upstream parallel: xllm_layers/ilu/attention.h + attention.cpp (82 + ~200 lines)
→ batch_prefill() dispatches to ixinfer_flash_attn_unpad_with_block_tables
→ batch_decode() dispatches to xllm_paged_attention
The key dispatch decision:
prefill (num_prefill_tokens > 0):
→ flash attention with block tables (unpadded, variable-length)
→ supports cu_seqlens for multi-request batching
→ window attention via window_size_left/right
decode (pure autoregressive, 1 token per sequence):
→ paged attention v1/v2
→ v1 vs v2 decision: total_tiles vs 2 × sm_count
→ BI-V100: 16 SMs → V2 beneficial when seq_len > 1024
GDN (GatedDeltaNet) layers [1,7,13,19] bypass this entirely —
they use the gdn_dispatch module instead.
Call chain:
qwen3_5.py Qwen3_5DecoderLayer.forward()
→ (full attention layers): attn_dispatch.dispatch_attention()
→ prefill: flash_attn_with_block_tables
→ decode: paged_attention_v1 or v2
→ (GDN layers): gdn_dispatch.dispatch_gdn()
"""
import logging
from typing import Optional, Tuple
import torch
logger = logging.getLogger("ex_engine.attn_dispatch")
# BI-V100 dispatch thresholds
# Source: SYSTEM_DESIGN.md, sub694 TPS analysis
_SM_COUNT = 16
_V2_THRESHOLD_FACTOR = 2 # V2 when total_tiles > 2 × SM_COUNT
_PAGE_BLOCK_SIZE = 16 # default paged attention block size
class AttnDispatchConfig:
"""
Attention dispatch configuration.
Parallels xllm_layers/ilu/attention.h struct members:
scale, is_causal, window_size_left, window_size_right, softcap
"""
__slots__ = (
'num_heads', 'num_kv_heads', 'head_dim', 'scale',
'is_causal', 'window_left', 'window_right',
'block_size', 'max_context_len', 'softcap',
)
def __init__(
self,
num_heads: int = 28,
num_kv_heads: int = 4,
head_dim: int = 128,
scale: Optional[float] = None,
is_causal: bool = True,
window_left: int = -1,
window_right: int = -1,
block_size: int = 16,
max_context_len: int = 131072,
softcap: float = 0.0,
):
self.num_heads = num_heads
self.num_kv_heads = num_kv_heads
self.head_dim = head_dim
self.scale = scale or (head_dim ** -0.5)
self.is_causal = is_causal
self.window_left = window_left
self.window_right = window_right
self.block_size = block_size
self.max_context_len = max_context_len
self.softcap = softcap
def should_use_paged_v2(
seq_len: int,
num_kv_heads: int,
block_size: int = 16,
partition_size: int = 512,
) -> bool:
"""
V1 vs V2 decision for paged attention.
Upstream: vllm/attention/ops/paged_attn.py PagedAttention._use_v2()
Rule: total_tiles = num_kv_heads × ceil(seq_len / partition_size)
use V2 when total_tiles > 2 × SM_COUNT (BI-V100: 32)
On BI-V100 with 16 SMs and 4 KV heads:
V2 when seq_len > 512 × (2 × 16 / 4) = 4096
In practice, V2 is better for seq_len > 1024 due to latency hiding.
"""
num_tiles = num_kv_heads * ((seq_len + partition_size - 1) // partition_size)
return num_tiles > _V2_THRESHOLD_FACTOR * _SM_COUNT
def dispatch_prefill(
config: AttnDispatchConfig,
query: torch.Tensor, # (total_q_tokens, num_heads, head_dim)
key_cache: torch.Tensor, # (num_blocks, num_kv_heads, block_size, head_dim)
value_cache: torch.Tensor, # (num_blocks, num_kv_heads, block_size, head_dim)
block_tables: torch.Tensor, # (batch_size, max_blocks_per_seq)
cu_seq_q: torch.Tensor, # (batch_size + 1,) int32 — query cumulative lengths
cu_seq_k: torch.Tensor, # (batch_size + 1,) int32 — key cumulative lengths
max_seq_q: int,
max_seq_k: int,
) -> torch.Tensor:
"""
Prefill attention via flash attention with block tables.
Upstream: xllm::kernel::ilu::batch_prefill
→ ixformer::infer::ixinfer_flash_attn_unpad_with_block_tables
The BI-V100 ixformer implements this as a modified flash attention
that reads KV from paged cache (block_tables → physical blocks).
"""
# Try ix_ops_dispatch first
try:
from ex_engine.python import ix_ops_dispatch
output = ix_ops_dispatch.flash_attn_with_block_tables(
query, key_cache, value_cache, block_tables,
cu_seq_q, cu_seq_k, max_seq_q, max_seq_k,
config.scale,
is_causal=config.is_causal,
window_left=config.window_left,
window_right=config.window_right,
softcap=config.softcap,
)
return output
except (ImportError, RuntimeError, AttributeError) as e:
logger.debug("flash_attn dispatch failed, using fallback: %s", e)
# Try direct ixformer
try:
import ixformer.functions as ixf_F
output = torch.empty_like(query)
ixf_F.ixinfer_flash_attn_unpad_with_block_tables(
query, key_cache, value_cache, output, block_tables,
cu_seq_q, cu_seq_k, max_seq_q, max_seq_k,
config.is_causal, config.window_left, config.window_right,
config.scale, config.softcap, False, None, None, None)
return output
except (ImportError, AttributeError):
pass
# PyTorch fallback — SDPA (no block table support, for testing only)
logger.warning("prefill: using PyTorch SDPA fallback (no block tables)")
output = torch.nn.functional.scaled_dot_product_attention(
query.unsqueeze(0), query.unsqueeze(0), query.unsqueeze(0),
scale=config.scale, is_causal=config.is_causal)
return output.squeeze(0)
def dispatch_decode(
config: AttnDispatchConfig,
output: torch.Tensor, # (batch_size, num_heads, head_dim) preallocated
query: torch.Tensor, # (batch_size, num_heads, head_dim)
key_cache: torch.Tensor, # (num_blocks, num_kv_heads, block_size, head_dim)
value_cache: torch.Tensor, # (num_blocks, num_kv_heads, block_size, head_dim)
block_tables: torch.Tensor, # (batch_size, max_blocks_per_seq)
context_lens: torch.Tensor, # (batch_size,) int32
max_context_len: int,
) -> None:
"""
Decode attention via paged attention v1/v2.
Upstream: xllm::kernel::ilu::batch_decode
→ ixformer::infer::xllm_paged_attention
"""
# Try ix_ops_dispatch
try:
from ex_engine.python import ix_ops_dispatch
ix_ops_dispatch.paged_attention_v1(
output, query, key_cache, value_cache,
config.num_kv_heads, config.scale,
block_tables, context_lens,
config.block_size, max_context_len,
window_left=config.window_left,
window_right=config.window_right,
softcap=config.softcap,
)
return
except (ImportError, RuntimeError, AttributeError) as e:
logger.debug("paged_attention dispatch failed: %s", e)
# Try direct ixformer
try:
import ixformer.functions as ixf_F
ixf_F.vllm_single_query_cached_kv_attention(
output, query, key_cache, value_cache,
config.num_kv_heads, config.scale,
block_tables, context_lens,
config.block_size, max_context_len, None)
return
except (ImportError, AttributeError):
pass
# PyTorch fallback — extremely slow, decode-only test path
logger.warning("decode: using PyTorch fallback (very slow)")
batch_size = query.shape[0]
for b in range(batch_size):
ctx_len = context_lens[b].item()
q = query[b] # (num_heads, head_dim)
# Reconstruct KV from cache
blocks = block_tables[b]
num_blocks_used = (ctx_len + config.block_size - 1) // config.block_size
k_list, v_list = [], []
for bi in range(num_blocks_used):
block_idx = blocks[bi].item()
tokens_in_block = min(config.block_size,
ctx_len - bi * config.block_size)
k_list.append(key_cache[block_idx, :, :tokens_in_block])
v_list.append(value_cache[block_idx, :, :tokens_in_block])
k = torch.cat(k_list, dim=1) # (kv_heads, ctx_len, head_dim)
v = torch.cat(v_list, dim=1)
# GQA: expand kv heads
num_q_per_kv = config.num_heads // config.num_kv_heads
k = k.repeat_interleave(num_q_per_kv, dim=0)
v = v.repeat_interleave(num_q_per_kv, dim=0)
# Standard attention
scores = torch.einsum('hd,hsd->hs', q.float(), k.float())
scores = scores * config.scale
scores = torch.softmax(scores, dim=-1)
out = torch.einsum('hs,hsd->hd', scores, v.float())
output[b] = out.to(output.dtype)
def dispatch_attention(
config: AttnDispatchConfig,
is_prefill: bool,
query: torch.Tensor,
key_cache: torch.Tensor,
value_cache: torch.Tensor,
block_tables: torch.Tensor,
# Prefill-specific
cu_seq_q: Optional[torch.Tensor] = None,
cu_seq_k: Optional[torch.Tensor] = None,
max_seq_q: int = 0,
max_seq_k: int = 0,
# Decode-specific
context_lens: Optional[torch.Tensor] = None,
max_context_len: int = 0,
) -> torch.Tensor:
"""
Top-level attention dispatcher.
Mirrors xllm's split between batch_prefill and batch_decode,
routing to the correct kernel based on is_prefill flag.
"""
if is_prefill:
return dispatch_prefill(
config, query, key_cache, value_cache,
block_tables, cu_seq_q, cu_seq_k,
max_seq_q, max_seq_k)
else:
output = torch.empty_like(query)
dispatch_decode(
config, output, query, key_cache, value_cache,
block_tables, context_lens, max_context_len)
return output

154
ex_engine/factors/factor_moe_combine.cu
View File

@@ -1,154 +0,0 @@
// ex_engine/factors/factor_moe_combine.cu
//
// Layer 10: MoE weighted combine kernel
//
// Upstream parallel: kernels/cuda/moe/moe_combine.cu (105 lines)
// Fused reorder + weighted sum replacing:
// torch::zeros + index_copy_ + view + multiply + sum
//
// Algorithm per token (each block handles one output token):
// For each of its topk experts:
// 1. Read expert output at the flat index position
// 2. Multiply by the router weight for this (token, expert) pair
// 3. Accumulate into output[token] in fp32
// Then cast back to input dtype.
//
// Grid: N blocks (one per output token)
// Block: 256 threads, each handling hidden_dim / 256 elements
//
// For Qwen3.5: hidden_size=3584, topk=8
// Each block reads 8 × 3584 = 28672 values and produces 3584 outputs.
// Compute: 8 FMA per element → 3584 × 8 = 28672 FMA → negligible.
// Bandwidth: 28672 × 2 bytes (fp16 read) + 3584 × 2 (fp16 write) = ~61 KB.
// At 900 GB/s: ~68 ns per block → fully bandwidth bound.
//
// BI-V100 SM70 adaptations:
// - Template on scalar_t (half, bfloat16, float)
// - fp32 accumulation to prevent overflow
// - 256 threads per block (8 warps, good SM70 occupancy)
// - Optional residual add (fused shared expert output)
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <stdint.h>
// =========================================================================
// Compile-time constants
// =========================================================================
static constexpr int32_t kCombineBlockSize = 256;
// =========================================================================
// Device helpers: type conversion to/from float
// =========================================================================
template <typename T>
__device__ __forceinline__ float to_float(T val);
template <>
__device__ __forceinline__ float to_float<float>(float val) { return val; }
template <>
__device__ __forceinline__ float to_float<__half>(__half val) {
return __half2float(val);
}
template <typename T>
__device__ __forceinline__ T from_float(float val);
template <>
__device__ __forceinline__ float from_float<float>(float val) { return val; }
template <>
__device__ __forceinline__ __half from_float<__half>(float val) {
return __float2half(val);
}
// =========================================================================
// Kernel: moe_combine_kernel
// =========================================================================
// Each block processes one output token.
// Threads stride over the hidden dimension.
// Accumulation in fp32 prevents overflow for fp16 inputs.
//
// Memory layout (after expert dispatch):
// gemm2_out: (N × topk, H) — expert outputs in flat-index order
// Token t's k-th expert output is at gemm2_out[(t × topk + k), :]
// reduce_weight: (N × topk,) or (N, topk) — router weights
// output: (N, H) — final combined output
template <typename scalar_t>
__global__ void moe_combine_kernel(
const scalar_t* __restrict__ gemm2, // (N*topk, H) expert outputs
const float* __restrict__ reduce_weight, // (N*topk,) or (N, topk)
scalar_t* __restrict__ output, // (N, H) final output
const scalar_t* __restrict__ residual, // (N, H) optional residual, NULL if none
int64_t N, // number of output tokens
int32_t topk, // experts per token
int64_t H // hidden dimension
) {
const int64_t token_id = blockIdx.x;
if (token_id >= N) return;
const int32_t tid = threadIdx.x;
const int32_t stride = kCombineBlockSize;
// Process hidden dimension elements in strided fashion
for (int64_t h = tid; h < H; h += stride) {
float acc = 0.0f;
// Accumulate over topk experts
for (int32_t k = 0; k < topk; ++k) {
int64_t flat_idx = token_id * topk + k;
float w = reduce_weight[flat_idx];
float val = to_float(gemm2[flat_idx * H + h]);
acc += w * val;
}
// Add residual if present (shared expert output)
if (residual != nullptr) {
acc += to_float(residual[token_id * H + h]);
}
output[token_id * H + h] = from_float<scalar_t>(acc);
}
}
// =========================================================================
// Host-side launcher
// =========================================================================
// Dispatches by dtype. Matches upstream xllm::kernel::cuda::moe_combine_result.
// The upstream version also supports bfloat16; we handle fp16 and fp32
// for BI-V100 (which lacks native bf16 tensor cores).
extern "C" int ex_moe_combine(
const void* gemm2_ptr, // (N*topk, H) device pointer
const float* reduce_weight, // (N*topk,) device pointer
void* output_ptr, // (N, H) device pointer
const void* residual_ptr, // (N, H) device pointer, NULL if none
int64_t N, // number of tokens
int32_t topk, // experts per token
int64_t H, // hidden dimension
int dtype, // 0 = fp32, 1 = fp16
cudaStream_t stream
) {
if (dtype == 1) {
// fp16 path — primary for Qwen3.5 inference
moe_combine_kernel<__half>
<<<N, kCombineBlockSize, 0, stream>>>(
reinterpret_cast<const __half*>(gemm2_ptr),
reduce_weight,
reinterpret_cast<__half*>(output_ptr),
residual_ptr ? reinterpret_cast<const __half*>(residual_ptr) : nullptr,
N, topk, H);
} else {
// fp32 path — for debugging or fp32 inference
moe_combine_kernel<float>
<<<N, kCombineBlockSize, 0, stream>>>(
reinterpret_cast<const float*>(gemm2_ptr),
reduce_weight,
reinterpret_cast<float*>(output_ptr),
residual_ptr ? reinterpret_cast<const float*>(residual_ptr) : nullptr,
N, topk, H);
}
return 0;
}

174
ex_engine/factors/factor_moe_compute_index.cu
View File

@@ -1,174 +0,0 @@
// ex_engine/factors/factor_moe_compute_index.cu
//
// Layer 9: MoE token index computation — 3-phase CUDA kernel
//
// Upstream parallel: kernels/cuda/moe/moe_compute_index.cu (155 lines)
// Fused MoE token index computation replacing:
// torch::bincount + 2 × torch::argsort + torch::cumsum + CPU sync
//
// Phase 1: histogram — atomicAdd per-expert token counts
// Phase 2: prefix_sum — 1 block, CUB BlockScan exclusive scan
// Phase 3: place_indices — atomicAdd on offsets, write bidirectional maps
//
// This kernel is called once per MoE layer (64 layers per forward).
// Input: expert_id tensor (num_tokens × topk flat), int32
// Output: src_to_dst, dst_to_src, expert_sizes
//
// BI-V100 SM70 adaptations:
// - CUB BlockScan via cub/block/block_scan.cuh (CUDA 10.2 compatible)
// - atomicAdd for int32 (SM70 native)
// - kMoeIndexBlock = 256 threads per block
// - Single-block prefix_sum (num_experts ≤ 256 guaranteed for Qwen3.5)
#include <cuda_runtime.h>
#include <stdint.h>
// CUB for BlockScan (exclusive prefix sum in shared memory)
#include <cub/block/block_scan.cuh>
// =========================================================================
// Compile-time constants
// =========================================================================
static constexpr int32_t kMoeIndexBlock = 256;
// =========================================================================
// Phase 1: Histogram — count tokens per expert
// =========================================================================
// Each thread processes one element of the flat expert_id array.
// atomicAdd to expert_sizes[eid] to build the histogram.
// Grid: ceil(N / 256) blocks
__global__ void moe_histogram_kernel(
const int32_t* __restrict__ expert_id, // (N,) flat expert assignments
int32_t* __restrict__ expert_sizes, // (num_experts,) output counts
int64_t num_elements,
int32_t num_experts
) {
int64_t tid = (int64_t)blockIdx.x * kMoeIndexBlock + threadIdx.x;
if (tid < num_elements) {
int32_t eid = expert_id[tid];
if (eid >= 0 && eid < num_experts) {
atomicAdd(&expert_sizes[eid], 1);
}
}
}
// =========================================================================
// Phase 2: Exclusive prefix sum — compute expert offsets
// =========================================================================
// Single block, one thread per expert (num_experts ≤ 256).
// Uses CUB BlockScan for an efficient exclusive prefix sum.
// Input: expert_sizes (per-expert token counts from Phase 1)
// Output: expert_offsets (exclusive scan — start position per expert)
//
// The exclusive scan means expert_offsets[e] = sum(expert_sizes[0..e-1]).
// This gives the starting position in the sorted token array for expert e.
__global__ void moe_prefix_sum_kernel(
const int32_t* __restrict__ expert_sizes, // (num_experts,) counts
int32_t* __restrict__ expert_offsets, // (num_experts,) output offsets
int32_t num_experts
) {
using BlockScan = cub::BlockScan<int32_t, kMoeIndexBlock>;
__shared__ typename BlockScan::TempStorage s_scan;
// Each thread loads one expert's count (0 if out of range)
int32_t val = (threadIdx.x < num_experts) ? expert_sizes[threadIdx.x] : 0;
int32_t offset;
// Exclusive sum: offset[i] = sum(val[0..i-1])
BlockScan(s_scan).ExclusiveSum(val, offset);
__syncthreads();
if (threadIdx.x < num_experts) {
expert_offsets[threadIdx.x] = offset;
}
}
// =========================================================================
// Phase 3: Place indices — build bidirectional permutation maps
// =========================================================================
// For each token assignment (flat_idx, expert_id):
// pos = atomicAdd(&expert_offsets[eid], 1) — claim next slot
// dst_to_src[pos] = flat_idx — sorted→original mapping
// src_to_dst[flat_idx] = pos — original→sorted mapping
//
// After this kernel:
// dst_to_src contains token indices grouped by expert
// src_to_dst[i] tells where token i ended up in the sorted order
//
// Note: expert_offsets is consumed destructively (atomicAdd increments it).
// The caller must keep expert_sizes separately.
__global__ void moe_place_indices_kernel(
const int32_t* __restrict__ expert_id, // (N,) flat expert assignments
int32_t* __restrict__ expert_offsets, // (num_experts,) — destructive
int32_t* __restrict__ dst_to_src, // (N,) output: sorted→original
int32_t* __restrict__ src_to_dst, // (N,) output: original→sorted
int64_t num_elements,
int32_t num_experts
) {
int64_t flat_idx = (int64_t)blockIdx.x * kMoeIndexBlock + threadIdx.x;
if (flat_idx >= num_elements) return;
int32_t eid = expert_id[flat_idx];
if (eid < 0 || eid >= num_experts) return;
// Claim the next position in this expert's segment
int32_t pos = atomicAdd(&expert_offsets[eid], 1);
// Write bidirectional mapping
dst_to_src[pos] = (int32_t)flat_idx;
src_to_dst[flat_idx] = pos;
}
// =========================================================================
// Host-side orchestrator — launches all 3 phases
// =========================================================================
// Matches upstream xllm::kernel::cuda::moe_compute_index signature.
//
// Usage from ixformer::infer::moe_compute_token_index_api:
// Phase 1: histogram → expert_sizes
// Phase 2: prefix_sum → expert_offsets (scratch, used by Phase 3)
// Phase 3: place_indices → dst_to_src, src_to_dst
//
// All 3 phases run on the same CUDA stream with implicit synchronization
// (each kernel completes before the next starts within the stream).
extern "C" int ex_moe_compute_index(
const int32_t* expert_id, // (N,) device pointer
int32_t* src_to_dst, // (N,) device pointer, output
int32_t* dst_to_src, // (N,) device pointer, output
int32_t* expert_sizes, // (num_experts,) device pointer, output
int64_t num_elements, // N = num_tokens × topk
int32_t num_experts, // E (64 for Qwen3.5)
cudaStream_t stream
) {
if (num_experts > kMoeIndexBlock) return -1; // Exceeds single-block scan
int64_t grid = (num_elements + kMoeIndexBlock - 1) / kMoeIndexBlock;
// Zero expert_sizes before histogram
cudaMemsetAsync(expert_sizes, 0, num_experts * sizeof(int32_t), stream);
// Allocate scratch for expert_offsets
int32_t* expert_offsets;
cudaMalloc(&expert_offsets, num_experts * sizeof(int32_t));
// Phase 1: histogram
moe_histogram_kernel<<<grid, kMoeIndexBlock, 0, stream>>>(
expert_id, expert_sizes, num_elements, num_experts);
// Phase 2: prefix sum (single block)
moe_prefix_sum_kernel<<<1, kMoeIndexBlock, 0, stream>>>(
expert_sizes, expert_offsets, num_experts);
// Phase 3: place indices (destructive on expert_offsets)
moe_place_indices_kernel<<<grid, kMoeIndexBlock, 0, stream>>>(
expert_id, expert_offsets, dst_to_src, src_to_dst,
num_elements, num_experts);
cudaFree(expert_offsets);
return 0;
}

456
ex_engine/factors/factor_topk_softmax.cu
View File

@@ -1,456 +0,0 @@
// ex_engine/factors/factor_topk_softmax.cu
//
// Layer 8: MoE topk_softmax CUDA kernel
//
// Upstream parallel: kernels/cuda/moe/moe_topk_softmax_kernels.cuh (867 lines)
// Originally adapted from:
// vllm v0.7.3 → csrc/moe/topk_softmax_kernels.cu
// TensorRT-LLM v0.7.1 → moe_kernels.cu
// xllm latest → moe_topk_softmax_kernels.cuh
//
// Three kernel paths in upstream:
// 1. topk_gating_softmax<T, VPT, NUM_EXPERTS, WARPS_PER_CTA, BYTES_PER_LDG>
// → For power-of-2 expert counts (1..256), packs rows into warps
// → Pure warp shuffle, zero shared memory
// → Qwen3.5 uses this path: 64 experts → VPT=2, 32 threads/row
//
// 2. moe_topk_fast<TPB>
// → For non-power-of-2 expert counts, k ≥ 2
// → Uses CUB BlockReduce with TopKPair (finds 2 maxima per iter)
// → Requires softmax_workspace for pre-computed softmax
//
// 3. moe_topK<TPB>
// → For non-power-of-2 expert counts, k = 1
// → Uses CUB BlockReduce with single cub::ArgMax
//
// BI-V100 SM70 adaptations:
// - __shfl_xor_sync with full mask 0xFFFFFFFF (SM70 warp shuffle)
// - No cp.async, no TMA — all loads are standard global loads
// - cub::BlockReduce via cub/block/block_reduce.cuh (CUB ships with CUDA 10.2)
// - __launch_bounds__ tuned for SM70: 128 threads, max occupancy
//
// For Qwen3.5-27B: NUM_EXPERTS=64, topk=8, all tokens route here.
// This kernel is called 64 times per forward pass (once per MoE layer).
// At ~6K tokens/batch: 64 × 6K = ~384K kernel launches amortized.
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <float.h>
#include <stdint.h>
// CUB for BlockReduce (non-power-of-2 fallback path)
#include <cub/block/block_reduce.cuh>
// =========================================================================
// SM70 warp shuffle macros (BI-V100 compatible)
// =========================================================================
// Upstream uses XLLM_SHFL_XOR_SYNC_WIDTH macro.
// On SM70, standard __shfl_xor_sync with full mask.
#define FULL_MASK 0xFFFFFFFFU
#define WARP_SIZE 32
#ifndef SHFL_XOR_SYNC
#define SHFL_XOR_SYNC(val, mask, width) \
__shfl_xor_sync(FULL_MASK, (val), (mask), (width))
#endif
#ifndef SHFL_SYNC
#define SHFL_SYNC(val, src, width) \
__shfl_sync(FULL_MASK, (val), (src), (width))
#endif
// =========================================================================
// Utility: convert generic type to float
// =========================================================================
template <typename T>
__device__ __forceinline__ float to_float(T val);
template <>
__device__ __forceinline__ float to_float<float>(float val) { return val; }
template <>
__device__ __forceinline__ float to_float<__half>(__half val) {
return __half2float(val);
}
// Aligned array for vectorized loads (replaces CUTLASS dependency)
template <typename T, int N>
struct alignas(sizeof(T) * N) AlignedArray {
T data[N];
__device__ __forceinline__ T& operator[](int i) { return data[i]; }
__device__ __forceinline__ const T& operator[](int i) const { return data[i]; }
};
// =========================================================================
// Compile-time constants
// =========================================================================
// TopkConstants: compute VPT and ROWS_PER_WARP from expert count and load width
template <typename T, int NUM_EXPERTS, int BYTES_PER_LDG>
struct TopkConstants {
static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
static constexpr int kThreadsPerRow = NUM_EXPERTS / (sizeof(T) <= 2 ? 2 : 1);
// Ensure threads_per_row does not exceed WARP_SIZE
static constexpr int VPT = NUM_EXPERTS / (WARP_SIZE < (NUM_EXPERTS / 1) ? WARP_SIZE : (NUM_EXPERTS / 1));
static constexpr int ROWS_PER_WARP = WARP_SIZE * VPT / NUM_EXPERTS;
};
// =========================================================================
// Kernel 1: topk_gating_softmax — power-of-2 experts (THE hot path)
// =========================================================================
// This is the primary kernel for Qwen3.5 (64 experts).
// Each warp processes kRowsPerWarp rows simultaneously.
// All reduces via warp shuffle — zero shared memory.
template <typename T, int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG>
__launch_bounds__(WARPS_PER_CTA * WARP_SIZE)
__global__ void topk_gating_softmax_kernel(
const T* __restrict__ input, // (num_rows, NUM_EXPERTS)
const bool* __restrict__ finished, // (num_rows,) or NULL
float* __restrict__ output, // (num_rows, k)
const int num_rows,
int* __restrict__ indices, // (num_rows, k)
const int k,
const int start_expert,
const int end_expert,
const bool renormalize
) {
// Compile-time geometry
static constexpr int kEltsPerLdg = BYTES_PER_LDG / sizeof(T);
static constexpr int kEltsPerRow = NUM_EXPERTS;
static constexpr int kThreadsPerRow = kEltsPerRow / VPT;
static constexpr int kLdgPerThread = VPT / kEltsPerLdg;
static constexpr int kEltsPerWarp = WARP_SIZE * VPT;
static constexpr int kRowsPerWarp = kEltsPerWarp / kEltsPerRow;
static constexpr int kRowsPerCta = WARPS_PER_CTA * kRowsPerWarp;
static constexpr int kColsPerGroupLdg = kEltsPerLdg * kThreadsPerRow;
// Row assignment
const int cta_base_row = blockIdx.x * kRowsPerCta;
const int warp_base_row = cta_base_row + threadIdx.y * kRowsPerWarp;
const int thread_row_in_warp = threadIdx.x / kThreadsPerRow;
const int thread_row = warp_base_row + thread_row_in_warp;
if (thread_row >= num_rows) return;
const bool row_active = finished ? !finished[thread_row] : true;
// Read this thread's chunk
const T* thread_row_ptr = input + thread_row * kEltsPerRow;
const int thread_group_idx = threadIdx.x % kThreadsPerRow;
const int first_elt = thread_group_idx * kEltsPerLdg;
const T* read_ptr = thread_row_ptr + first_elt;
// Vectorized load
using AccessType = AlignedArray<T, kEltsPerLdg>;
T row_chunk_raw[VPT];
AccessType* vec_ptr = reinterpret_cast<AccessType*>(&row_chunk_raw);
const AccessType* src_ptr = reinterpret_cast<const AccessType*>(read_ptr);
#pragma unroll
for (int ii = 0; ii < kLdgPerThread; ++ii) {
vec_ptr[ii] = src_ptr[ii * kThreadsPerRow];
}
// Convert to float
float row_chunk[VPT];
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
row_chunk[ii] = to_float(row_chunk_raw[ii]);
}
// ===== Softmax: max reduction via butterfly =====
float thread_max = row_chunk[0];
#pragma unroll
for (int ii = 1; ii < VPT; ++ii) {
thread_max = fmaxf(thread_max, row_chunk[ii]);
}
#pragma unroll
for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
thread_max = fmaxf(thread_max,
SHFL_XOR_SYNC(thread_max, mask, kThreadsPerRow));
}
// ===== Softmax: exp and sum =====
float row_sum = 0.0f;
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
row_chunk[ii] = expf(row_chunk[ii] - thread_max);
row_sum += row_chunk[ii];
}
#pragma unroll
for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
row_sum += SHFL_XOR_SYNC(row_sum, mask, kThreadsPerRow);
}
// ===== Normalize =====
const float inv_sum = 1.0f / row_sum;
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
row_chunk[ii] *= inv_sum;
}
// ===== TopK via iterative warp argmax =====
int start_col = first_elt;
float renorm_sum = 0.0f;
for (int k_idx = 0; k_idx < k; ++k_idx) {
// Thread-local argmax
float max_val = row_chunk[0];
int expert = start_col;
#pragma unroll
for (int ldg = 0, col = start_col; ldg < kLdgPerThread;
++ldg, col += kColsPerGroupLdg) {
#pragma unroll
for (int ii = 0; ii < kEltsPerLdg; ++ii) {
float val = row_chunk[ldg * kEltsPerLdg + ii];
if (val > max_val) {
max_val = val;
expert = col + ii;
}
}
}
// Butterfly argmax across thread group
#pragma unroll
for (int mask = kThreadsPerRow / 2; mask > 0; mask /= 2) {
float other_max = SHFL_XOR_SYNC(max_val, mask, kThreadsPerRow);
int other_expert = SHFL_XOR_SYNC(expert, mask, kThreadsPerRow);
if (other_max > max_val ||
(other_max == max_val && other_expert < expert)) {
max_val = other_max;
expert = other_expert;
}
}
// Write result (lead thread only)
if (thread_group_idx == 0) {
const bool uses_expert = expert >= start_expert && expert < end_expert;
const bool should_process = row_active && uses_expert;
const int idx = k * thread_row + k_idx;
output[idx] = max_val;
indices[idx] = should_process ? (expert - start_expert) : NUM_EXPERTS;
renorm_sum += max_val;
}
// Suppress winner for next iteration
if (k_idx + 1 < k) {
const int winner_ldg = expert / kColsPerGroupLdg;
const int winner_thread = (expert / kEltsPerLdg) % kThreadsPerRow;
if (thread_group_idx == winner_thread) {
const int offset = expert % kEltsPerLdg;
row_chunk[winner_ldg * kEltsPerLdg + offset] = -10000.0f;
}
}
}
// Renormalize
if (renormalize && thread_group_idx == 0) {
float inv = 1.0f / renorm_sum;
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int idx = k * thread_row + k_idx;
output[idx] *= inv;
}
}
}
// =========================================================================
// Kernel 2: moe_softmax — generic softmax for non-power-of-2 fallback
// =========================================================================
// Uses CUB BlockReduce for max and sum across arbitrary expert counts.
// Writes softmax probabilities to output buffer for subsequent topk.
template <typename T, int TPB>
__launch_bounds__(TPB)
__global__ void moe_softmax_kernel(
const T* __restrict__ input, // (num_tokens, num_cols)
float* __restrict__ output, // (num_tokens, num_cols)
const int num_cols
) {
using BlockReduce = cub::BlockReduce<float, TPB>;
__shared__ typename BlockReduce::TempStorage tmp_storage;
__shared__ float s_max;
__shared__ float s_norm;
const int row_offset = blockIdx.x * num_cols;
// Pass 1: find max
float thread_max = -FLT_MAX;
for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
float val = to_float(input[row_offset + ii]);
output[row_offset + ii] = val; // store converted value
thread_max = fmaxf(thread_max, val);
}
float block_max = BlockReduce(tmp_storage).Reduce(thread_max, cub::Max());
if (threadIdx.x == 0) s_max = block_max;
__syncthreads();
// Pass 2: exp and sum
float thread_sum = 0.0f;
for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
float val = expf(output[row_offset + ii] - s_max);
output[row_offset + ii] = val;
thread_sum += val;
}
float block_sum = BlockReduce(tmp_storage).Sum(thread_sum);
if (threadIdx.x == 0) s_norm = 1.0f / block_sum;
__syncthreads();
// Pass 3: normalize
for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
output[row_offset + ii] *= s_norm;
}
}
// =========================================================================
// Kernel 3: moe_topk_fast — topk from pre-computed softmax (k ≥ 2)
// =========================================================================
// Uses CUB BlockReduce with TopKPair to find 2 maxima per iteration.
// Upstream: moe_topk_fast, uses cub::KeyValuePair.
using cub_kvp = cub::KeyValuePair<int, float>;
template <int TPB>
__launch_bounds__(TPB)
__global__ void moe_topk_fast_kernel(
float* __restrict__ probs, // (N, E) — modified in-place
float* __restrict__ output, // (N, k)
int* __restrict__ indices, // (N, k)
const int num_experts,
const int k,
const int start_expert,
const int end_expert,
const bool renormalize
) {
using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
__shared__ typename BlockReduce::TempStorage tmp_storage;
const int row = blockIdx.x;
const int row_offset = row * num_experts;
float renorm_sum = 0.0f;
cub::ArgMax arg_max;
for (int k_idx = 0; k_idx < k; ++k_idx) {
cub_kvp thread_kvp;
thread_kvp.key = 0;
thread_kvp.value = -1.0f;
for (int e = threadIdx.x; e < num_experts; e += TPB) {
cub_kvp inp;
inp.key = e;
inp.value = probs[row_offset + e];
thread_kvp = arg_max(inp, thread_kvp);
}
cub_kvp result = BlockReduce(tmp_storage).Reduce(thread_kvp, arg_max);
if (threadIdx.x == 0) {
const int expert = result.key;
const bool uses = expert >= start_expert && expert < end_expert;
const int idx = k * row + k_idx;
output[idx] = result.value;
indices[idx] = uses ? (expert - start_expert) : num_experts;
renorm_sum += result.value;
// Suppress winner
probs[row_offset + expert] = -1.0f;
}
__syncthreads();
}
if (renormalize && threadIdx.x == 0) {
float inv = 1.0f / renorm_sum;
for (int k_idx = 0; k_idx < k; ++k_idx) {
output[k * row + k_idx] *= inv;
}
}
}
// =========================================================================
// Host-side launcher with template dispatch by expert count
// =========================================================================
// Matches upstream topk_gating_softmax_kernel_launcher pattern.
// Power-of-2 experts → topk_gating_softmax (zero shared mem, warp shuffle)
// Other → moe_softmax + moe_topk_fast (CUB path)
template <typename T, int EXPERTS, int WARPS_PER_TB>
void launch_topk_gating(
const T* input, float* output, int* indices,
int num_rows, int k, int start_expert, int end_expert,
bool renormalize, cudaStream_t stream
) {
// For SM70: BYTES_PER_LDG capped at min(16, sizeof(T)*EXPERTS)
static constexpr int kBytesPerLdg =
(16 < (int)(sizeof(T) * EXPERTS)) ? 16 : (int)(sizeof(T) * EXPERTS);
static constexpr int kEltsPerLdg = kBytesPerLdg / sizeof(T);
static constexpr int kVpt = EXPERTS / WARP_SIZE;
// Ensure VPT ≥ 1
static constexpr int VPT = (kVpt > 0) ? kVpt : 1;
static constexpr int kRowsPerWarp = (WARP_SIZE * VPT) / EXPERTS;
static constexpr int kRowsPerCta = WARPS_PER_TB * ((kRowsPerWarp > 0) ? kRowsPerWarp : 1);
const int num_blocks = (num_rows + kRowsPerCta - 1) / kRowsPerCta;
dim3 block(WARP_SIZE, WARPS_PER_TB);
topk_gating_softmax_kernel<T, VPT, EXPERTS, WARPS_PER_TB, kBytesPerLdg>
<<<num_blocks, block, 0, stream>>>(
input, nullptr, output, num_rows, indices,
k, start_expert, end_expert, renormalize);
}
// Macro for dispatch table
#define LAUNCH_GATING(TYPE, EXPERTS, WARPS) \
launch_topk_gating<TYPE, EXPERTS, WARPS>( \
gating_ptr, topk_weights, topk_indices, \
num_tokens, topk, 0, num_experts, \
renormalize, stream);
// =========================================================================
// Host entry point: topk_softmax (matches ixformer::infer::topk_softmax)
// =========================================================================
extern "C" void ex_topk_softmax(
float* topk_weights, // (num_tokens, topk) output
int* topk_indices, // (num_tokens, topk) output
const float* gating_output, // (num_tokens, num_experts) input
int num_tokens,
int num_experts,
int topk,
bool renormalize,
cudaStream_t stream
) {
const float* gating_ptr = gating_output;
const bool is_pow2 = (num_experts & (num_experts - 1)) == 0;
if (is_pow2 && num_experts <= 256) {
// Fast path: topk_gating_softmax with warp shuffle
static constexpr int kWarps = 4;
switch (num_experts) {
case 1: LAUNCH_GATING(float, 1, kWarps); break;
case 2: LAUNCH_GATING(float, 2, kWarps); break;
case 4: LAUNCH_GATING(float, 4, kWarps); break;
case 8: LAUNCH_GATING(float, 8, kWarps); break;
case 16: LAUNCH_GATING(float, 16, kWarps); break;
case 32: LAUNCH_GATING(float, 32, kWarps); break;
case 64: LAUNCH_GATING(float, 64, kWarps); break;
case 128: LAUNCH_GATING(float, 128, kWarps); break;
case 256: LAUNCH_GATING(float, 256, kWarps); break;
}
} else {
// Fallback: softmax + topk via CUB
static constexpr int kTpb = 256;
// Allocate workspace for softmax output
float* workspace;
cudaMalloc(&workspace, (size_t)num_tokens * num_experts * sizeof(float));
moe_softmax_kernel<float, kTpb>
<<<num_tokens, kTpb, 0, stream>>>(
gating_output, workspace, num_experts);
moe_topk_fast_kernel<kTpb>
<<<num_tokens, kTpb, 0, stream>>>(
workspace, topk_weights, topk_indices,
num_experts, topk, 0, num_experts, renormalize);
cudaFree(workspace);
}
}

245
ex_engine/factors/hw_config.h
View File

@@ -1,245 +0,0 @@
// ex_engine/factors/hw_config.h
//
// Layer 1: Hardware descriptor + per-algorithm tuning tables
//
// Upstream parallel: xllm_models/llm/qwen3_5.h (model-level config)
// → defines num_experts=64, top_k=8, hidden_size=3584, intermediate=18944
// → binds GDN layers [1,7,13,19] vs full attention layers
//
// This file defines the BI-V100 hardware descriptor and per-factor
// tuning tables. Every downstream factor file #includes this to get
// compile-time constants that are calibrated to the real device.
//
// Source: cat_files/arch.h (real nm -D symbol dump from BI-V100)
// cat_files/iluvatar_mma.hpp (tensor core instruction set)
// sub694rizhi.txt TPS distribution (avg 6.2, target >12)
#ifndef EX_FACTORS_HW_CONFIG_H
#define EX_FACTORS_HW_CONFIG_H
#include <stdint.h>
#ifdef __cplusplus
extern "C" {
#endif
// =========================================================================
// BI-V100 hardware constants (from real device probes)
// =========================================================================
// Source: cat_files/arch.h, HARDWARE_PROBE_20260808.md
#define EX_SM_MAJOR 7
#define EX_SM_MINOR 0
#define EX_SM_COUNT 16
#define EX_WARP_SIZE 32
#define EX_MAX_THREADS_SM 2048
#define EX_SMEM_PER_SM 49152 // bytes
#define EX_SMEM_PER_BLOCK 49152 // bytes, BI-V100 no bank partitioning
#define EX_L2_CACHE_BYTES (6 * 1024 * 1024) // 6 MB
#define EX_MEM_BW_GBS 900.0f // total ~900 GB/s (56 GB/s per SM × 16)
// Register file: 65536 32-bit registers per SM, 256 per warp
// Source: probed via cudaDeviceProp
#define EX_REGS_PER_SM 65536
#define EX_REGS_PER_BLOCK 65536
// CoreX clang/16 CUDA 10.2 compatibility constraints
// Source: cat_files/cutlass.h (Iluvatar CoreX copyright block)
// - No cp.async (requires SM80+)
// - No TMA (requires SM90+)
// - Warp shuffle via __shfl_sync, __shfl_xor_sync, __shfl_down_sync
// - atomicAdd for float available (SM70)
// - half2 arithmetic available via __hmul2, __hadd2 etc.
#define EX_HAS_CP_ASYNC 0
#define EX_HAS_TMA 0
#define EX_HAS_WARP_SHUFFLE 1
#define EX_HAS_HALF2 1
// =========================================================================
// Qwen3.5-27B model constants (from SYSTEM_DESIGN.md)
// =========================================================================
#define EX_NUM_LAYERS 64 // decoder layers
#define EX_NUM_EXPERTS 64 // routed experts per MoE layer
#define EX_MOE_TOPK 8 // active experts per token
#define EX_HIDDEN_SIZE 3584 // hidden dimension
#define EX_INTERMEDIATE_SIZE 18944 // MoE intermediate (per expert, pre-TP)
#define EX_NUM_HEADS 28 // attention heads
#define EX_NUM_KV_HEADS 4 // GQA kv heads
#define EX_HEAD_DIM 128 // per-head dimension
#define EX_GDN_LAYERS_COUNT 4 // GatedDeltaNet layers: [1,7,13,19]
#define EX_ATTN_LAYERS_COUNT 32 // full attention layers (of 36 attention layers)
#define EX_MOE_LAYERS_COUNT 64 // all 64 layers have MoE
// TP=4 sliced sizes (actual runtime)
#define EX_TP_SIZE 4
#define EX_INTER_PER_TP (EX_INTERMEDIATE_SIZE / EX_TP_SIZE) // 4736
// =========================================================================
// Per-factor tuning parameter structs
// =========================================================================
// Mirrors CCCL's per-algorithm ReducePassPolicy/ScanPolicy pattern
// Each factor reads these at compile time or load time to set
// grid/block dims and SMEM allocation.
typedef struct {
int threads_per_block;
int items_per_thread; // values per thread (VPT)
int vec_size; // vector load width
int smem_bytes; // shared memory per block
int num_warps; // threads_per_block / WARP_SIZE
int num_stages; // software pipeline stages
int grid_scale; // multiplier for grid dim (1 = 1:1 with N)
} ex_factor_tune_t;
// =========================================================================
// BI-V100 tuning tables — one entry per factor
// =========================================================================
// Source: CCCL tuning headers adapted via muh toolchain
// sub168 TPS benchmarks, sub694 per-request TPS distribution
//
// Key tuning rationale per sub694 data:
// - avg TPS = 6.2, ceiling TPS = 11.9 (decode-only, low prompt_tok)
// - 227/824 requests below 5 TPS → bottleneck is prefill-heavy requests
// - prompt_tok >100K → TPS drops to 0.4-1.2 (GEMM bound)
// - prompt_tok <10K → TPS reaches 9-11 (compute matches HW)
// - Goal: double TPS on prompt_tok 10K-50K range (bulk of traffic)
// Factor 0: MOE_TOPK_SOFTMAX
// 64 experts, VPT=2 → 32 threads = 1 warp per token row
// 4 warps per CTA → 4 tokens per CTA
// No SMEM needed (all warp shuffle)
// Upstream: moe_topk_softmax_kernels.cuh topkGatingSoftmax
static const ex_factor_tune_t EX_TUNE_MOE_TOPK_SOFTMAX = {
.threads_per_block = 128, // 4 warps × 32
.items_per_thread = 2, // 64 experts / 32 threads
.vec_size = 1, // float, no vectorization
.smem_bytes = 0, // pure warp shuffle
.num_warps = 4,
.num_stages = 1,
.grid_scale = 1, // ceil(N / 4) blocks
};
// Factor 1: MOE_ALIGN_BLOCK
// Simple histogram + prefix sum, 256 threads per block
// SMEM for CUB BlockScan
static const ex_factor_tune_t EX_TUNE_MOE_ALIGN_BLOCK = {
.threads_per_block = 256,
.items_per_thread = 1,
.vec_size = 1,
.smem_bytes = 1024, // CUB BlockScan TempStorage
.num_warps = 8,
.num_stages = 1,
.grid_scale = 1,
};
// Factor 2: MOE_FUSED_GEMM
// Group GEMM: cublas or CUTLASS batched
// BI-V100: cublas SM70 hgemm, M×N×K per expert
// threads/smem managed by cublas internally
static const ex_factor_tune_t EX_TUNE_MOE_FUSED_GEMM = {
.threads_per_block = 256, // cublas managed
.items_per_thread = 4,
.vec_size = 8, // half8 loads
.smem_bytes = 49152, // full SMEM for GEMM tiles
.num_warps = 8,
.num_stages = 2, // SW pipeline (no cp.async, manual prefetch)
.grid_scale = 1,
};
// Factor 3: GELU_TANH_MUL
// Element-wise, memory-bandwidth bound
// 256 threads, vec4 half loads
static const ex_factor_tune_t EX_TUNE_GELU_TANH_MUL = {
.threads_per_block = 256,
.items_per_thread = 4,
.vec_size = 4, // half4 = 8 bytes
.smem_bytes = 0, // pure register
.num_warps = 8,
.num_stages = 1,
.grid_scale = 1,
};
// Factor 4: BATCHED_ROTARY
// Per-head, per-position rotation
// Each thread handles one (cos, sin) pair
static const ex_factor_tune_t EX_TUNE_BATCHED_ROTARY = {
.threads_per_block = 512,
.items_per_thread = 2, // 2 elements per rotation pair
.vec_size = 2,
.smem_bytes = 0,
.num_warps = 16,
.num_stages = 1,
.grid_scale = 1,
};
// Factor 5: GDN_CHUNK_FWD
// Chunked prefill: intra-chunk QK^T + inter-chunk state update
// Critical: fp32 accumulation to prevent NaN
// chunk_size=64 (from fla upstream), head_dim=128
// Each CTA processes one (batch, head, chunk)
static const ex_factor_tune_t EX_TUNE_GDN_CHUNK_FWD = {
.threads_per_block = 128,
.items_per_thread = 4,
.vec_size = 4,
.smem_bytes = 32768, // Q,K,V tiles for chunk_size=64, dim=128, fp16
.num_warps = 4,
.num_stages = 2,
.grid_scale = 1,
};
// Factor 6: GDN_RECURRENT
// Single-step decode: conv1d_state update + delta_rule recurrence
// Lightweight: one token per step
static const ex_factor_tune_t EX_TUNE_GDN_RECURRENT = {
.threads_per_block = 128,
.items_per_thread = 4,
.vec_size = 4,
.smem_bytes = 8192, // conv state + temporal state
.num_warps = 4,
.num_stages = 1,
.grid_scale = 1,
};
// Factor 7: CACHE_APPEND (reshape_and_cache for paged KV)
// Each thread handles one token's K or V slice
static const ex_factor_tune_t EX_TUNE_CACHE_APPEND = {
.threads_per_block = 256,
.items_per_thread = 4,
.vec_size = 4,
.smem_bytes = 0,
.num_warps = 8,
.num_stages = 1,
.grid_scale = 1,
};
// Factor 8: RESHAPE_CACHE_FLASH
// Optimized cache write for flash attention layout
static const ex_factor_tune_t EX_TUNE_RESHAPE_CACHE_FLASH = {
.threads_per_block = 256,
.items_per_thread = 4,
.vec_size = 4,
.smem_bytes = 0,
.num_warps = 8,
.num_stages = 1,
.grid_scale = 1,
};
// =========================================================================
// Tuning table array (indexed by factor_id)
// =========================================================================
static const ex_factor_tune_t* const EX_TUNE_TABLE[] = {
&EX_TUNE_MOE_TOPK_SOFTMAX, // 0
&EX_TUNE_MOE_ALIGN_BLOCK, // 1
&EX_TUNE_MOE_FUSED_GEMM, // 2
&EX_TUNE_GELU_TANH_MUL, // 3
&EX_TUNE_BATCHED_ROTARY, // 4
&EX_TUNE_GDN_CHUNK_FWD, // 5
&EX_TUNE_GDN_RECURRENT, // 6
&EX_TUNE_CACHE_APPEND, // 7
&EX_TUNE_RESHAPE_CACHE_FLASH, // 8
};
#ifdef __cplusplus
}
#endif
#endif // EX_FACTORS_HW_CONFIG_H

182
ex_engine/factors/ilu_ops_api.h
View File

@@ -1,182 +0,0 @@
// ex_engine/factors/ilu_ops_api.h
//
// Layer 4: Dispatch signature contract (canonical API header)
//
// Upstream parallel: kernels/ilu/ilu_ops_api.h (~160 lines)
// Defines the signature contract between:
// Layer 2-3 (orchestrators) → Layer 5-6 (kernel wrappers)
// Layer 5-6 (kernel wrappers) → Layer 7 (ixformer::infer namespace)
//
// Every function declared here has exactly one implementation path:
// ilu_ops_api.h declaration
// → kernels/ilu/*.cpp wrapper (Layer 5-6)
// → ixformer::infer::* (Layer 7, from ixformer.h)
// → CUDA kernel (Layer 8-10)
//
// On BI-V100, ixformer::infer is the base image's libixformer.so.
// The EX engine replaces MISSING ops by providing .so factors that
// export the same infer:: signatures.
//
// This file is a verbatim-structure copy of the upstream xllm ilu_ops_api.h,
// with only the include paths adjusted for our build tree.
#ifndef EX_FACTORS_ILU_OPS_API_H
#define EX_FACTORS_ILU_OPS_API_H
#include <torch/all.h>
#include <ATen/Tensor.h>
#include <optional>
#include <string>
#include <vector>
#include <tuple>
// =========================================================================
// Namespace: xllm::kernel::ilu
// =========================================================================
// Each function maps to exactly one ixformer::infer call.
// The function bodies live in separate .cpp files (Layer 5-6).
namespace xllm {
namespace kernel {
namespace ilu {
// ---- Attention ops (Layer 6: attention.cpp) ----
void reshape_paged_cache(
torch::Tensor& key,
std::optional<torch::Tensor>& value,
torch::Tensor& key_cache,
std::optional<torch::Tensor>& value_cache,
torch::Tensor& slot_mapping);
void batch_prefill(
torch::Tensor& query,
const torch::Tensor& key,
const std::optional<torch::Tensor>& value,
torch::Tensor& output,
std::optional<torch::Tensor>& output_lse,
const std::optional<torch::Tensor>& q_cu_seq_lens,
const std::optional<torch::Tensor>& kv_cu_seq_lens,
const std::optional<torch::Tensor>& alibi_slope,
const std::optional<torch::Tensor>& attn_bias,
const std::optional<torch::Tensor>& q_quant_scale,
const std::optional<torch::Tensor>& k_quant_scale,
const std::optional<torch::Tensor>& v_quant_scale,
const torch::Tensor& block_tables,
int64_t max_query_len,
int64_t max_seq_len,
float scale,
bool is_causal,
int64_t window_size_left,
int64_t window_size_right,
const std::string& compute_dtype,
bool return_lse);
void batch_decode(
torch::Tensor& query,
const torch::Tensor& k_cache,
torch::Tensor& output,
const torch::Tensor& block_table,
const torch::Tensor& seq_lens,
const std::optional<torch::Tensor>& v_cache,
std::optional<torch::Tensor>& output_lse,
const std::optional<torch::Tensor>& q_quant_scale,
const std::optional<torch::Tensor>& k_cache_quant_scale,
const std::optional<torch::Tensor>& v_cache_quant_scale,
const std::optional<torch::Tensor>& out_quant_scale,
const std::optional<torch::Tensor>& alibi_slope,
const std::optional<torch::Tensor>& mask,
const std::string& compute_dtype,
int64_t max_seq_len,
int64_t window_size_left,
int64_t window_size_right,
float scale,
bool return_lse,
bool is_causal,
int64_t kv_cache_quant_bit_size);
// ---- Normalization ops (Layer 6: norm.cpp) ----
void residual_layer_norm(
torch::Tensor& input,
torch::Tensor& output,
std::optional<torch::Tensor>& residual,
torch::Tensor& weight,
std::optional<torch::Tensor>& bias,
std::optional<torch::Tensor>& residual_out,
double eps);
void rms_norm(
torch::Tensor& output,
torch::Tensor& input,
torch::Tensor& weight,
double eps);
// ---- Activation ops (Layer 6: activation.cpp) ----
void act_and_mul(
torch::Tensor out,
torch::Tensor input,
const std::string& act_mode);
// ---- RoPE ops (Layer 6: rope.cpp) ----
void apply_rope_pos_ids_cos_sin_cache(
torch::Tensor& query,
torch::Tensor& key,
torch::Tensor& cos_sin_cache,
torch::Tensor& positions,
bool interleave);
// ---- Linear / matmul ops (Layer 6: matmul.cpp) ----
torch::Tensor matmul(
torch::Tensor a,
torch::Tensor b,
std::optional<torch::Tensor> bias);
// ---- MoE ops (Layer 5: fused_moe.cpp) ----
// Step 1: Router — softmax + topk
std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
const torch::Tensor& input,
int64_t topk,
int64_t num_expert_group,
int64_t topk_group,
bool normalize,
const std::optional<torch::Tensor>& mask,
const std::string& normed_by,
const std::string& scoring_func,
double route_scale,
const std::optional<torch::Tensor>& e_score_correction_bias);
// Step 2: Generate permutation indices
std::vector<torch::Tensor> moe_gen_idx(
torch::Tensor& expert_id,
int64_t expert_num);
// Step 3: Expand input by topk
torch::Tensor moe_expand_input(
const torch::Tensor& input,
const torch::Tensor& gather_index,
const torch::Tensor& combine_idx,
int64_t topk);
// Step 4+6: Group GEMM
torch::Tensor group_gemm(
torch::Tensor& input,
torch::Tensor& weight,
torch::Tensor& tokens_per_experts,
const std::optional<torch::Tensor>& dst_to_src,
torch::Tensor& output);
// Step 7: Weighted combine
torch::Tensor moe_combine_result(
torch::Tensor& input,
torch::Tensor& weight);
} // namespace ilu
} // namespace kernel
} // namespace xllm
#endif // EX_FACTORS_ILU_OPS_API_H

246
ex_engine/factors/ixformer_infer.h
View File

@@ -1,246 +0,0 @@
// ex_engine/factors/ixformer_infer.h
//
// Layer 7: ixformer::infer namespace contract
//
// Upstream parallel: kernels/ilu/ixformer.h (~140 lines)
// Declares every function in ixformer::infer that the ilu kernel
// wrappers (Layer 5-6) call through to.
//
// On BI-V100, this namespace is implemented by two sources:
//
// 1. BASE IMAGE (libixformer.so from corex SDK 3.2.3):
// PRESENT — these symbols exist in nm -D of the .so:
// silu_and_mul
// rms_norm
// residual_rms_norm
// xllm_rotary_embedding (aka vllm_rotary_embedding_neox)
// xllm_reshape_and_cache
// xllm_paged_attention (v1/v2)
// ixinfer_flash_attn_unpad_with_block_tables
// ixformer_linear / ixformer_linear_ex
// topk_softmax
// moe_compute_token_index_api
// moe_expand_input
// moe_w16a16_group_gemm
// moe_output_reduce_sum
//
// 2. EX FACTORS (.so files from ex_engine/build):
// For MISSING ops that may not be in all base image versions.
// The EX factor .so exports the same symbol → dlopen replaces it.
//
// Signature source: Verbatim from upstream xllm ixformer.h + utils.h,
// cross-referenced with cat_files/symbol_dumps nm -D output.
#ifndef EX_FACTORS_IXFORMER_INFER_H
#define EX_FACTORS_IXFORMER_INFER_H
#include <torch/all.h>
#include <ATen/Tensor.h>
#include <optional>
#include <string>
namespace ixformer {
namespace infer {
// =====================================================================
// Attention kernels
// =====================================================================
// Flash attention with block tables (prefill path)
// Source: ixinfer flash attention unpadded variant
// BI-V100 status: PRESENT in base image
torch::Tensor ixinfer_flash_attn_unpad_with_block_tables(
torch::Tensor& query,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
torch::Tensor& out,
torch::Tensor& block_tables,
torch::Tensor& cu_seq_q,
torch::Tensor& cu_seq_k,
int64_t max_seq_q,
int64_t max_seq_k,
bool is_causal,
int64_t window_left,
int64_t window_right,
double scale,
double softcap,
bool sqrt_alibi,
const std::optional<torch::Tensor>& alibi_slopes,
const std::optional<torch::Tensor>& sinks,
std::optional<torch::Tensor>& lse);
// Paged attention (decode path, v1 or v2 selected internally)
// BI-V100 status: PRESENT
torch::Tensor xllm_paged_attention(
torch::Tensor& out,
torch::Tensor& query,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
int64_t num_kv_heads,
double scale,
torch::Tensor& block_tables,
torch::Tensor& context_lens,
int64_t block_size,
int64_t max_context_len,
const std::optional<torch::Tensor>& alibi_slopes,
bool causal,
int32_t window_left,
int32_t window_right,
double softcap,
bool enable_cuda_graph,
bool use_sqrt_alibi,
const std::optional<torch::Tensor>& sinks);
// =====================================================================
// Activation kernels
// =====================================================================
// SiLU-and-mul: out = silu(input[:half]) * input[half:]
// BI-V100 status: PRESENT
void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
// =====================================================================
// Linear / GEMM kernels
// =====================================================================
// ixformer linear: fused matmul with optional activation
// BI-V100 status: PRESENT
torch::Tensor ixformer_linear(
torch::Tensor& input,
torch::Tensor& weight,
int64_t act_type,
const std::optional<torch::Tensor>& bias,
const std::optional<torch::Tensor>& out,
const std::optional<bool> persistent);
// ixformer linear extended: simplified interface
// BI-V100 status: PRESENT
torch::Tensor ixformer_linear_ex(
torch::Tensor& input,
torch::Tensor& weight,
const c10::optional<torch::Tensor>& bias,
const c10::optional<torch::Tensor>& out);
// =====================================================================
// Cache management kernels
// =====================================================================
// Write KV into paged cache
// BI-V100 status: PRESENT
void xllm_reshape_and_cache(
torch::Tensor& key,
torch::Tensor& value,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
torch::Tensor& slot_mapping,
int64_t key_token_stride,
int64_t value_token_stride);
// =====================================================================
// Rotary embedding kernels
// =====================================================================
// Apply rotary position encoding
// BI-V100 status: PRESENT
void xllm_rotary_embedding(
torch::Tensor& positions,
torch::Tensor& query,
torch::Tensor& key,
int64_t head_size,
torch::Tensor& cos_sin_cache,
bool is_neox);
// =====================================================================
// Normalization kernels
// =====================================================================
// Fused residual + RMS normalization
// BI-V100 status: PRESENT
void residual_rms_norm(
torch::Tensor& input,
torch::Tensor& residual,
torch::Tensor& weight,
torch::Tensor& output,
torch::Tensor& residual_output,
const std::optional<torch::Tensor>& fused_bias,
double alpha,
double eps,
bool is_post);
// RMS normalization
// BI-V100 status: PRESENT
void rms_norm(
torch::Tensor& input,
torch::Tensor& weight,
torch::Tensor& output,
const std::optional<torch::Tensor>& fused_bias,
double eps);
// =====================================================================
// MoE kernels
// =====================================================================
// Fused softmax + top-k for MoE routing
// BI-V100 status: PRESENT (confirmed in base image symbol dump)
// Calls CUDA kernel: moe_topk_softmax_kernels.cuh (Layer 8)
void topk_softmax(
torch::Tensor& topk_weights,
torch::Tensor& topk_indices,
torch::Tensor& token_expert_indices,
torch::Tensor& gating_output,
bool renormalize);
// 3-phase permutation index computation for MoE token dispatch
// BI-V100 status: PRESENT
// Calls CUDA kernel: moe_compute_index.cu (Layer 9)
void moe_compute_token_index_api(
torch::Tensor& topk_ids,
torch::Tensor& src_dst,
torch::Tensor& dst_src,
torch::Tensor& expert_sizes_gpu,
const c10::optional<torch::Tensor>& expert_mask,
const c10::optional<torch::Tensor>& expert_sizes_cpu,
const c10::optional<torch::Tensor>& expand_tokens_gpu,
int64_t start_expert_id,
int64_t end_expert_id,
int64_t num_experts);
// Gather tokens from natural order to expert-sorted order
// BI-V100 status: PRESENT
void moe_expand_input(
torch::Tensor outputs,
torch::Tensor inputs,
torch::Tensor dst_to_src,
const c10::optional<torch::Tensor>& src_to_dst,
int64_t dst_tokens,
int64_t expand_factor);
// Group GEMM for MoE expert computation (half-precision)
// BI-V100 status: PRESENT
// This is the primary compute bottleneck.
void moe_w16a16_group_gemm(
torch::Tensor output,
torch::Tensor inputs,
torch::Tensor weights,
torch::Tensor tokens_per_experts,
const c10::optional<torch::Tensor>& dst_to_src,
const c10::optional<torch::Tensor>& bias,
std::string format,
int64_t persistent,
int64_t output_n);
// Weighted combine of expert outputs
// BI-V100 status: PRESENT
// Calls CUDA kernel: moe_combine.cu (Layer 10)
void moe_output_reduce_sum(
torch::Tensor outputs,
torch::Tensor inputs,
const c10::optional<torch::Tensor>& mul_weight,
const c10::optional<torch::Tensor>& mask,
const c10::optional<torch::Tensor>& extra_residual,
double scaling_factor);
} // namespace infer
} // namespace ixformer
#endif // EX_FACTORS_IXFORMER_INFER_H

210
ex_engine/factors/kernel_elem_ops.cpp
View File

@@ -1,210 +0,0 @@
// ex_engine/factors/kernel_elem_ops.cpp
//
// Layer 6: Element-wise kernel dispatch wrappers
//
// Upstream parallel: kernels/ilu/activation.cpp (30 lines)
// + kernels/ilu/norm.cpp (45 lines)
// + kernels/ilu/rope.cpp (20 lines)
// + kernels/ilu/group_gemm.cpp (25 lines)
// + kernels/ilu/matmul.cpp (~15 lines)
//
// Total upstream: ~135 lines across 5 files.
// Each function is a 3-5 line dispatch wrapper that calls ixformer::infer.
//
// These ops are PRESENT in the base image's ixformer — they don't need
// EX factor replacement. But they must be in the call chain because:
// - Activation is Step 5 of the MoE pipeline (called between GEMM1 and GEMM2)
// - RMSNorm is called before/after every decoder layer (2× per layer × 64 layers)
// - RoPE is called once per attention layer (32 full attention + 4 GDN = 36)
// - Group GEMM is Steps 4+6 (called twice per MoE layer × 64 layers)
//
// The presence in ixformer is confirmed by:
// cat_files/symbol_dumps (nm -D output from real device)
// SYSTEM_DESIGN.md PRESENT list
#include "ilu_ops_api.h"
#include "ixformer.h"
using namespace ixformer;
namespace xllm {
namespace kernel {
namespace ilu {
// =====================================================================
// Activation: silu_and_mul
// =====================================================================
// Upstream: kernels/ilu/activation.cpp::act_and_mul
// BI-V100 ixformer: PRESENT (silu_and_mul confirmed in symbol dump)
//
// Input: (tokens, 2 × intermediate_size) — gate + up projections concatenated
// Output: (tokens, intermediate_size) — silu(gate) × up
//
// For Qwen3.5: intermediate_size = 18944 / TP4 = 4736
// Each call processes 4736 × 2 = 9472 half values per token.
// At 200 tokens/batch decode: 200 × 9472 × 2 bytes = 3.6 MB bandwidth.
void act_and_mul(
torch::Tensor out,
torch::Tensor input,
const std::string& act_mode) {
if (act_mode == "silu") {
infer::silu_and_mul(input, out);
} else {
// gelu_tanh_and_mul is MISSING from ixformer on BI-V100.
// The EX factor system provides this as EX_FACTOR_GELU_TANH_MUL (id=3).
// For now, fallback to PyTorch.
LOG(FATAL) << "Unsupported act mode: " << act_mode
<< ", only silu is available via ixformer on BI-V100. "
<< "Use EX factor 3 for gelu_tanh.";
}
}
// =====================================================================
// RMSNorm: rms_norm + residual_layer_norm
// =====================================================================
// Upstream: kernels/ilu/norm.cpp
// BI-V100 ixformer: PRESENT (rms_norm, fused_add_rms_norm confirmed)
//
// rms_norm: out = x × rsqrt(mean(x²) + eps) × weight
// residual_layer_norm: fused residual add + rms_norm
//
// Called 2× per decoder layer (pre-attention + post-attention norm).
// 64 layers × 2 = 128 calls per forward pass.
void rms_norm(
torch::Tensor& output,
torch::Tensor& input,
torch::Tensor& weight,
double eps) {
std::optional<torch::Tensor> fused_bias = std::nullopt;
infer::rms_norm(input, weight, output, fused_bias, eps);
}
void residual_layer_norm(
torch::Tensor& input,
torch::Tensor& output,
std::optional<torch::Tensor>& residual,
torch::Tensor& weight,
std::optional<torch::Tensor>& bias,
std::optional<torch::Tensor>& residual_out,
double eps) {
auto residual_ = residual.value_or(torch::zeros_like(input));
torch::Tensor residual_out_ = residual_out.value_or(torch::zeros_like(input));
infer::residual_rms_norm(
input, residual_, weight, output, residual_out_,
bias, /*alpha=*/1.0, eps, /*is_post=*/false);
}
// =====================================================================
// RoPE: Rotary Position Embedding
// =====================================================================
// Upstream: kernels/ilu/rope.cpp
// BI-V100 ixformer: PRESENT (vllm_rotary_embedding_neox confirmed)
//
// Applies cosine-sine rotation to query and key tensors.
// Called once per attention layer per forward pass.
// Qwen3.5: 36 attention layers (32 full + 4 GDN).
void apply_rope_pos_ids_cos_sin_cache(
torch::Tensor& query,
torch::Tensor& key,
torch::Tensor& cos_sin_cache,
torch::Tensor& positions,
bool interleave) {
const int64_t head_size = cos_sin_cache.size(-1);
// is_neox = !interleave (NeoX-style = non-interleaved)
infer::xllm_rotary_embedding(
positions, query, key, head_size, cos_sin_cache, !interleave);
}
// =====================================================================
// Group GEMM: Batched matrix multiplication for MoE experts
// =====================================================================
// Upstream: kernels/ilu/group_gemm.cpp
// BI-V100 ixformer: PRESENT (moe_w16a16_group_gemm confirmed)
//
// Performs A × B^T for each expert group simultaneously.
// tokens_per_experts defines the row count per group.
// Called twice per MoE layer: once for w13 (gate+up), once for w2 (down).
//
// For Qwen3.5 with TP4:
// w13: (16 experts_local, 9472, 3584) — 16 experts × [inter*2, hidden]
// w2: (16 experts_local, 3584, 4736) — 16 experts × [hidden, inter]
//
// This is the primary compute bottleneck on BI-V100.
// sub694 shows prompt_tok >50K requests drop to 1-3 TPS — GEMM bound.
torch::Tensor group_gemm(
torch::Tensor& input,
torch::Tensor& weight,
torch::Tensor& tokens_per_experts,
const std::optional<torch::Tensor>& dst_to_src,
torch::Tensor& output) {
infer::moe_w16a16_group_gemm(
output,
input,
weight,
tokens_per_experts,
dst_to_src,
/*bias=*/std::nullopt,
/*format=*/"TN",
/*persistent=*/0,
/*output_n=*/tokens_per_experts.sum().item<int64_t>());
return output;
}
// =====================================================================
// Reshape and cache: KV cache management
// =====================================================================
// Upstream: kernels/ilu/attention.cpp::reshape_paged_cache
// BI-V100 ixformer: PRESENT (vllm_cache_ops_reshape_and_cache)
//
// Writes new KV pairs into paged cache at the positions specified
// by slot_mapping.
void reshape_paged_cache(
torch::Tensor& key,
std::optional<torch::Tensor>& value,
torch::Tensor& key_cache,
std::optional<torch::Tensor>& value_cache,
torch::Tensor& slot_mapping) {
auto value_ = value.value_or(torch::Tensor());
auto value_cache_ = value_cache.value_or(torch::Tensor());
int64_t key_token_stride = key.stride(0);
int64_t value_token_stride = 0;
if (value_.defined()) {
value_token_stride = value_.stride(0);
}
slot_mapping = slot_mapping.to(at::kLong);
infer::xllm_reshape_and_cache(
key, value_, key_cache, value_cache_,
slot_mapping, key_token_stride, value_token_stride);
}
// =====================================================================
// Matmul: General matrix multiplication
// =====================================================================
// Upstream: kernels/ilu/matmul.cpp
// Used for linear projections (q/k/v proj, out proj, gate proj)
torch::Tensor matmul(
torch::Tensor a,
torch::Tensor b,
std::optional<torch::Tensor> bias) {
return infer::ixformer_linear_ex(a, b, bias, std::nullopt);
}
} // namespace ilu
} // namespace kernel
} // namespace xllm

155
ex_engine/factors/kernel_moe_ops.cpp
View File

@@ -1,155 +0,0 @@
// ex_engine/factors/kernel_moe_ops.cpp
//
// Layer 5: MoE kernel-level operations
//
// Upstream parallel: kernels/ilu/fused_moe.cpp (99 lines)
// → moe_active_topk() → infer::topk_softmax
// → moe_gen_idx() → infer::moe_compute_token_index_api
// → moe_expand_input() → infer::moe_expand_input
// → moe_combine_result() → infer::moe_output_reduce_sum
//
// Each function is a thin dispatch wrapper. On BI-V100, the call goes:
// this .cpp → ixformer::infer::* (libixformer.so from base image)
// OR
// this .cpp → EX factor .so (our replacement for missing ops)
//
// The code here is intentionally minimal — the real logic lives in
// the CUDA kernels (Layer 8-10). This layer only handles:
// 1. Tensor type coercion (fp16 → fp32 for routing)
// 2. Output tensor allocation
// 3. Call-through to infer:: namespace
#include "ilu_ops_api.h"
#include "ixformer.h"
namespace xllm {
namespace kernel {
namespace ilu {
// =====================================================================
// Step 1: moe_active_topk — Router dispatch
// =====================================================================
// Upstream: kernels/ilu/fused_moe.cpp::moe_active_topk
// Converts input to float32, allocates output tensors, calls topk_softmax.
// The topk_softmax kernel (Layer 8) does fused softmax + topk in one pass.
std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
const torch::Tensor& input,
int64_t topk,
int64_t num_expert_group,
int64_t topk_group,
bool normalize,
const std::optional<torch::Tensor>& mask,
const std::string& normed_by,
const std::string& scoring_func,
double route_scale,
const std::optional<torch::Tensor>& e_score_correction_bias) {
// Cast to float32 for numerical stability (half softmax overflows)
torch::Tensor input_f32 = input.to(torch::kFloat32);
// Allocate output tensors — matches upstream exactly
auto reduce_weight = torch::empty(
{input.size(0), topk},
torch::dtype(torch::kFloat).device(input.device()));
auto topk_indices = torch::empty(
{input.size(0), topk},
torch::dtype(torch::kInt32).device(input.device()));
auto token_expert_indices = torch::empty(
{input.size(0), topk},
torch::dtype(torch::kInt32).device(input.device()));
// Dispatch to ixformer::infer::topk_softmax
// This calls the CUDA kernel in Layer 8 (moe_topk_softmax_kernels.cuh)
ixformer::infer::topk_softmax(
reduce_weight, topk_indices, token_expert_indices, input_f32,
/*renormalize=*/false);
// Renormalize weights (upstream does this post-kernel)
if (normalize) {
auto weight_sum = reduce_weight.sum(-1);
reduce_weight = reduce_weight / weight_sum.unsqueeze(-1);
}
return std::make_tuple(reduce_weight, topk_indices);
}
// =====================================================================
// Step 2: moe_gen_idx — Permutation index generation
// =====================================================================
// Upstream: kernels/ilu/fused_moe.cpp::moe_gen_idx
// Calls the 3-phase CUDA kernel (histogram → prefix_sum → place)
// Returns {src_dst, dst_src, expert_sizes, expert_sizes_cumsum}
std::vector<torch::Tensor> moe_gen_idx(
torch::Tensor& expert_id,
int64_t expert_num) {
auto src_dst = expert_id.new_empty({expert_id.numel()});
auto dst_src = torch::empty_like(src_dst);
auto expert_sizes_gpu = expert_id.new_empty({expert_num});
auto expert_sizes_gpu_cumsum = expert_id.new_zeros({expert_id.numel() + 1});
// Dispatch to ixformer::infer::moe_compute_token_index_api
// This calls the 3-phase CUDA kernel in Layer 9
ixformer::infer::moe_compute_token_index_api(
expert_id, src_dst, dst_src, expert_sizes_gpu,
/*expert_mask=*/std::nullopt,
/*expert_sizes_cpu=*/std::nullopt,
/*expand_tokens_gpu=*/std::nullopt,
0, expert_num, expert_num);
expert_sizes_gpu_cumsum = expert_sizes_gpu.cumsum(-1);
return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_gpu_cumsum};
}
// =====================================================================
// Step 3: moe_expand_input — Token gather/scatter
// =====================================================================
// Upstream: kernels/ilu/fused_moe.cpp::moe_expand_input
// Reorders tokens from natural order to expert-grouped order.
torch::Tensor moe_expand_input(
const torch::Tensor& input,
const torch::Tensor& gather_index,
const torch::Tensor& combine_idx,
int64_t topk) {
int64_t dst_tokens = input.size(0) * topk;
auto output = input.new_empty({dst_tokens, input.size(1)});
// Dispatch to ixformer::infer::moe_expand_input
ixformer::infer::moe_expand_input(
output, input, combine_idx, gather_index, dst_tokens, topk);
return output;
}
// =====================================================================
// Step 7: moe_combine_result — Weighted combine
// =====================================================================
// Upstream: kernels/ilu/fused_moe.cpp::moe_combine_result
// Reorders from expert-sorted back to token order with weighted sum.
// Calls the CUDA kernel in Layer 10 (moe_combine.cu)
torch::Tensor moe_combine_result(
torch::Tensor& input,
torch::Tensor& weight) {
input = input.view({-1, weight.size(1), input.size(1)});
auto output = input.new_empty({input.size(0), input.size(2)});
// Dispatch to ixformer::infer::moe_output_reduce_sum
ixformer::infer::moe_output_reduce_sum(
output, input, weight,
/*mask=*/std::nullopt,
/*extra_residual=*/std::nullopt,
/*scaling_factor=*/1.0);
return output;
}
} // namespace ilu
} // namespace kernel
} // namespace xllm

461
ex_engine/factors/moe_pipeline.py
View File

@@ -1,461 +0,0 @@
"""
ex_engine/factors/moe_pipeline.py
Layer 2: MoE 7-step pipeline orchestrator
Upstream parallel: xllm_layers/ilu/fused_moe.cpp (806 lines)
→ FusedMoEImpl::forward_experts() orchestrates the full MoE hot path:
Step 1: select_experts → moe_active_topk (topk_softmax)
Step 2: moe_gen_idx → moe_compute_token_index (histogram + prefix_sum + place)
Step 3: moe_expand_input (gather tokens by expert)
Step 4: group_gemm (w13: gate_proj + up_proj fused)
Step 5: activation (silu_and_mul on gated MLP)
Step 6: group_gemm (w2: down_proj)
Step 7: moe_combine_result (weighted reduce over topk experts)
This module mirrors the full 7-step pipeline. Each step dispatches
to the ix_ops_dispatch layer (Layer 4) which calls into ixformer::infer
C++ kernels. The pipeline ordering and tensor lifetime management
matches xllm upstream exactly.
Call chain:
vllm model forward
→ Qwen3MoeSparseMoeBlock.forward()
→ moe_pipeline.fused_moe_forward()
→ Step 1-7 below
"""
import logging
from typing import Optional, Tuple
import torch
logger = logging.getLogger("ex_engine.moe_pipeline")
class MoEPipelineConfig:
"""
Configuration for the MoE pipeline.
Parallels xllm_layers/ilu/fused_moe.h FusedMoEArgs:
num_total_experts_, topk_, hidden_size_, intermediate_size_,
is_gated_, renormalize_, hidden_act_, scoring_func_
"""
__slots__ = (
'num_experts', 'topk', 'hidden_size', 'intermediate_size',
'is_gated', 'renormalize', 'hidden_act', 'scoring_func',
'tp_size', 'tp_rank', 'ep_size', 'ep_rank',
'start_expert_id', 'num_experts_per_rank',
)
def __init__(
self,
num_experts: int = 64,
topk: int = 8,
hidden_size: int = 3584,
intermediate_size: int = 18944,
is_gated: bool = True,
renormalize: bool = True,
hidden_act: str = "silu",
scoring_func: str = "softmax",
tp_size: int = 4,
tp_rank: int = 0,
ep_size: int = 1,
ep_rank: int = 0,
):
self.num_experts = num_experts
self.topk = topk
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.is_gated = is_gated
self.renormalize = renormalize
self.hidden_act = hidden_act
self.scoring_func = scoring_func
self.tp_size = tp_size
self.tp_rank = tp_rank
self.ep_size = ep_size
self.ep_rank = ep_rank
self.num_experts_per_rank = num_experts // ep_size
self.start_expert_id = ep_rank * self.num_experts_per_rank
class MoEPipeline:
"""
7-step MoE pipeline matching xllm FusedMoEImpl::forward_experts().
Each step calls through the dispatch layer. The pipeline manages
intermediate tensor lifetimes to minimize GPU memory pressure,
matching xllm's explicit tensor release pattern:
- expand_hidden_states released after Step 6
- act_out released after Step 6
"""
def __init__(self, config: MoEPipelineConfig, dispatch_module=None):
self.config = config
# The dispatch module provides the per-op kernel calls
# At runtime this is ix_ops_dispatch or direct ixformer
if dispatch_module is None:
try:
from ex_engine.python import ix_ops_dispatch
self.dispatch = ix_ops_dispatch
except ImportError:
self.dispatch = None
logger.warning("ix_ops_dispatch not available, MoE pipeline "
"will use PyTorch fallbacks")
else:
self.dispatch = dispatch_module
# ===================================================================
# Step 1: Router — softmax + topk (36× per layer, 64 layers)
# ===================================================================
# Upstream: FusedMoEImpl::select_experts → kernel::ilu::moe_active_topk
# → infer::topk_softmax
# → cuda::moe_topk_softmax_kernels.cuh::topkGatingSoftmax
def step1_topk_route(
self,
router_logits: torch.Tensor, # (num_tokens, num_experts)
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Returns:
topk_weights: (num_tokens, topk) float32, renormalized
topk_ids: (num_tokens, topk) int32
"""
if self.dispatch is not None:
try:
return self.dispatch.topk_softmax(
router_logits, self.config.topk, self.config.renormalize)
except (RuntimeError, AttributeError) as e:
logger.debug("topk_softmax dispatch failed: %s", e)
# PyTorch fallback — matches xllm ilu::moe_active_topk
logits_f32 = router_logits.float()
probs = torch.softmax(logits_f32, dim=-1)
topk_weights, topk_ids = torch.topk(probs, self.config.topk, dim=-1)
if self.config.renormalize:
topk_weights = topk_weights / topk_weights.sum(
dim=-1, keepdim=True)
return topk_weights, topk_ids.to(torch.int32)
# ===================================================================
# Step 2: Generate expert indices (permutation maps)
# ===================================================================
# Upstream: kernel::ilu::moe_gen_idx
# → infer::moe_compute_token_index_api
# → cuda::moe_compute_index (3-phase: histogram, prefix_sum, place)
def step2_gen_idx(
self,
topk_ids: torch.Tensor, # (num_tokens, topk) int32
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Build bidirectional permutation maps for expert dispatch.
Returns:
src_to_dst: (num_tokens * topk,) int32 — original → sorted position
dst_to_src: (num_tokens * topk,) int32 — sorted → original position
expert_sizes: (num_experts,) int32 — tokens per expert
"""
flat_ids = topk_ids.view(-1)
num_elements = flat_ids.shape[0]
num_experts = self.config.num_experts
device = topk_ids.device
if self.dispatch is not None:
try:
return self.dispatch.moe_compute_token_index(
topk_ids, num_experts, self.config.start_expert_id)
except (RuntimeError, AttributeError):
pass
# PyTorch fallback — matches cuda::moe_compute_index 3-phase logic
# Phase 1: histogram
expert_sizes = torch.zeros(
num_experts, dtype=torch.int32, device=device)
for eid in range(num_experts):
expert_sizes[eid] = (flat_ids == eid).sum().to(torch.int32)
# Phase 2: exclusive prefix sum
expert_offsets = torch.zeros(
num_experts, dtype=torch.int32, device=device)
expert_offsets[1:] = torch.cumsum(expert_sizes[:-1], dim=0)
# Phase 3: place indices
dst_to_src = torch.empty(
num_elements, dtype=torch.int32, device=device)
src_to_dst = torch.empty(
num_elements, dtype=torch.int32, device=device)
offsets_scratch = expert_offsets.clone()
for i in range(num_elements):
eid = flat_ids[i].item()
if 0 <= eid < num_experts:
pos = offsets_scratch[eid].item()
offsets_scratch[eid] += 1
dst_to_src[pos] = i
src_to_dst[i] = pos
return src_to_dst, dst_to_src, expert_sizes
# ===================================================================
# Step 3: Expand input (gather tokens by expert ordering)
# ===================================================================
# Upstream: kernel::ilu::moe_expand_input
# → infer::moe_expand_input
def step3_expand_input(
self,
hidden_states: torch.Tensor, # (num_tokens, hidden_size)
dst_to_src: torch.Tensor, # (num_tokens * topk,) int32
) -> torch.Tensor:
"""
Reorder tokens into expert-grouped order for batched GEMM.
Returns:
expanded: (num_tokens * topk, hidden_size) same dtype as input
"""
if self.dispatch is not None:
try:
return self.dispatch.moe_expand_input(
hidden_states, dst_to_src, self.config.topk)
except (RuntimeError, AttributeError):
pass
# PyTorch fallback
src_indices = dst_to_src.long()
# Each entry in dst_to_src is a flat index into the expanded token list.
# The source token index is flat_idx // topk
token_indices = src_indices // self.config.topk
expanded = hidden_states[token_indices]
return expanded
# ===================================================================
# Step 4: Group GEMM 1 — gate_proj + up_proj (w13)
# ===================================================================
# Upstream: kernel::ilu::group_gemm
# → infer::moe_w16a16_group_gemm
# weight shape: (num_experts_per_rank, intermediate_size * 2, hidden_size)
# for gated MLP: w1 and w3 fused into one [2*inter, hidden] matrix
def step4_gemm1(
self,
expanded_input: torch.Tensor, # (total_tokens, hidden_size)
w13: torch.Tensor, # (E_local, inter*2, hidden) or flat
expert_sizes: torch.Tensor, # (num_experts,) int32
) -> torch.Tensor:
"""
Group GEMM: expanded_input × w13^T for each expert group.
Returns:
gemm1_out: (total_tokens, intermediate_size * 2)
"""
if self.dispatch is not None:
try:
inter2 = w13.shape[1] if w13.dim() == 3 else w13.shape[0]
return self.dispatch.moe_group_gemm(
expanded_input, w13, expert_sizes, inter2)
except (RuntimeError, AttributeError):
pass
# PyTorch fallback: loop over experts
total_tokens = expanded_input.shape[0]
out_dim = w13.shape[1] if w13.dim() == 3 else w13.shape[0]
output = torch.empty(
total_tokens, out_dim,
dtype=expanded_input.dtype, device=expanded_input.device)
offset = 0
for e in range(expert_sizes.shape[0]):
count = expert_sizes[e].item()
if count > 0:
local_e = e - self.config.start_expert_id
if 0 <= local_e < w13.shape[0]:
x_e = expanded_input[offset:offset + count]
w_e = w13[local_e] # (inter*2, hidden)
# matmul: (count, hidden) × (hidden, inter*2) = (count, inter*2)
output[offset:offset + count] = x_e @ w_e.t()
offset += count
return output
# ===================================================================
# Step 5: Activation — SiLU-and-mul for gated MLP
# ===================================================================
# Upstream: kernel::ilu::act_and_mul → infer::silu_and_mul
# Input: (total_tokens, intermediate_size * 2)
# Output: (total_tokens, intermediate_size)
# Split input in half: out = silu(input[:, :inter]) * input[:, inter:]
def step5_activation(
self,
gemm1_out: torch.Tensor, # (total_tokens, inter*2)
) -> torch.Tensor:
"""
Gated SiLU activation.
Returns:
act_out: (total_tokens, intermediate_size)
"""
if self.config.is_gated:
half_dim = gemm1_out.shape[-1] // 2
gate = gemm1_out[:, :half_dim]
up = gemm1_out[:, half_dim:]
if self.dispatch is not None:
try:
# ixformer expects concatenated input, produces half-width output
return self.dispatch.silu_and_mul(gemm1_out)
except (RuntimeError, AttributeError):
pass
# PyTorch fallback — explicit silu_and_mul
return torch.nn.functional.silu(gate) * up
else:
if self.config.hidden_act == "silu":
return torch.nn.functional.silu(gemm1_out)
elif self.config.hidden_act == "gelu":
return torch.nn.functional.gelu(gemm1_out)
else:
return gemm1_out
# ===================================================================
# Step 6: Group GEMM 2 — down_proj (w2)
# ===================================================================
# Upstream: kernel::ilu::group_gemm (same as Step 4, different weights)
# weight shape: (num_experts_per_rank, hidden_size, intermediate_size)
def step6_gemm2(
self,
act_out: torch.Tensor, # (total_tokens, intermediate_size)
w2: torch.Tensor, # (E_local, hidden, inter)
expert_sizes: torch.Tensor, # (num_experts,) int32
) -> torch.Tensor:
"""
Group GEMM: act_out × w2^T for each expert group.
Returns:
gemm2_out: (total_tokens, hidden_size)
"""
if self.dispatch is not None:
try:
return self.dispatch.moe_group_gemm(
act_out, w2, expert_sizes, self.config.hidden_size)
except (RuntimeError, AttributeError):
pass
# PyTorch fallback: loop over experts
total_tokens = act_out.shape[0]
output = torch.empty(
total_tokens, self.config.hidden_size,
dtype=act_out.dtype, device=act_out.device)
offset = 0
for e in range(expert_sizes.shape[0]):
count = expert_sizes[e].item()
if count > 0:
local_e = e - self.config.start_expert_id
if 0 <= local_e < w2.shape[0]:
x_e = act_out[offset:offset + count]
w_e = w2[local_e] # (hidden, inter)
output[offset:offset + count] = x_e @ w_e.t()
offset += count
return output
# ===================================================================
# Step 7: Combine — weighted reduce over topk experts
# ===================================================================
# Upstream: kernel::ilu::moe_combine_result
# → infer::moe_output_reduce_sum
# → cuda::moe_combine_kernel
# Reorder from expert-sorted back to token order, weighted sum.
def step7_combine(
self,
gemm2_out: torch.Tensor, # (total_tokens, hidden_size) sorted
topk_weights: torch.Tensor, # (num_tokens, topk) float32
src_to_dst: torch.Tensor, # (total_tokens,) int32
) -> torch.Tensor:
"""
Weighted combine of expert outputs back to token order.
Returns:
final: (num_tokens, hidden_size)
"""
num_tokens = topk_weights.shape[0]
topk = self.config.topk
hidden_size = gemm2_out.shape[-1]
if self.dispatch is not None:
try:
return self.dispatch.moe_output_reduce_sum(
gemm2_out, topk_weights, 1.0)
except (RuntimeError, AttributeError):
pass
# PyTorch fallback — matches cuda::moe_combine_kernel logic
output = torch.zeros(
num_tokens, hidden_size,
dtype=gemm2_out.dtype, device=gemm2_out.device)
for t in range(num_tokens):
for k in range(topk):
flat_idx = t * topk + k
dst_pos = src_to_dst[flat_idx].long().item()
w = topk_weights[t, k].item()
output[t] += w * gemm2_out[dst_pos].float()
return output.to(gemm2_out.dtype)
# ===================================================================
# Full forward — orchestrates all 7 steps
# ===================================================================
# Upstream: FusedMoEImpl::forward_experts (main orchestrator)
def forward(
self,
hidden_states: torch.Tensor, # (num_tokens, hidden_size)
router_logits: torch.Tensor, # (num_tokens, num_experts)
w13: torch.Tensor, # (E_local, inter*2, hidden)
w2: torch.Tensor, # (E_local, hidden, inter)
shared_expert_output: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
Full MoE forward pass.
Tensor lifetime management matches xllm:
- expand_hidden_states is released after step6
- act_out is released after step6
- gemm1_out can be released after step5
"""
# Step 1: Router
topk_weights, topk_ids = self.step1_topk_route(router_logits)
# Step 2: Generate permutation indices
src_to_dst, dst_to_src, expert_sizes = self.step2_gen_idx(topk_ids)
# Step 3: Expand input tokens into expert-sorted order
expand_hidden_states = self.step3_expand_input(
hidden_states, dst_to_src)
# Step 4: Group GEMM 1 (gate_proj + up_proj)
gemm1_out = self.step4_gemm1(
expand_hidden_states, w13, expert_sizes)
# Step 5: Activation (gated SiLU)
act_out = self.step5_activation(gemm1_out)
del gemm1_out # release intermediate
# Step 6: Group GEMM 2 (down_proj)
gemm2_out = self.step6_gemm2(act_out, w2, expert_sizes)
del expand_hidden_states, act_out # release intermediates
# Step 7: Weighted combine
final_hidden_states = self.step7_combine(
gemm2_out, topk_weights, src_to_dst)
# Add shared expert output if present (Qwen3.5 has shared experts)
if shared_expert_output is not None:
final_hidden_states = final_hidden_states + shared_expert_output
return final_hidden_states