diff --git a/ex_engine/factors/attn_dispatch.py b/ex_engine/factors/attn_dispatch.py new file mode 100644 index 00000000..958232a7 --- /dev/null +++ b/ex_engine/factors/attn_dispatch.py @@ -0,0 +1,270 @@ +""" +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 diff --git a/ex_engine/factors/factor_moe_combine.cu b/ex_engine/factors/factor_moe_combine.cu new file mode 100644 index 00000000..1096507b --- /dev/null +++ b/ex_engine/factors/factor_moe_combine.cu @@ -0,0 +1,154 @@ +// 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 +#include +#include + +// ========================================================================= +// Compile-time constants +// ========================================================================= +static constexpr int32_t kCombineBlockSize = 256; + +// ========================================================================= +// Device helpers: type conversion to/from float +// ========================================================================= +template +__device__ __forceinline__ float to_float(T val); + +template <> +__device__ __forceinline__ float to_float(float val) { return val; } + +template <> +__device__ __forceinline__ float to_float<__half>(__half val) { + return __half2float(val); +} + +template +__device__ __forceinline__ T from_float(float val); + +template <> +__device__ __forceinline__ float from_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 +__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(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> + <<>>( + reinterpret_cast(gemm2_ptr), + reduce_weight, + reinterpret_cast<__half*>(output_ptr), + residual_ptr ? reinterpret_cast(residual_ptr) : nullptr, + N, topk, H); + } else { + // fp32 path — for debugging or fp32 inference + moe_combine_kernel + <<>>( + reinterpret_cast(gemm2_ptr), + reduce_weight, + reinterpret_cast(output_ptr), + residual_ptr ? reinterpret_cast(residual_ptr) : nullptr, + N, topk, H); + } + + return 0; +} diff --git a/ex_engine/factors/factor_moe_compute_index.cu b/ex_engine/factors/factor_moe_compute_index.cu new file mode 100644 index 00000000..20cd578f --- /dev/null +++ b/ex_engine/factors/factor_moe_compute_index.cu @@ -0,0 +1,174 @@ +// 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 +#include + +// CUB for BlockScan (exclusive prefix sum in shared memory) +#include + +// ========================================================================= +// 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; + __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<<>>( + 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<<>>( + expert_id, expert_offsets, dst_to_src, src_to_dst, + num_elements, num_experts); + + cudaFree(expert_offsets); + + return 0; +} diff --git a/ex_engine/factors/factor_topk_softmax.cu b/ex_engine/factors/factor_topk_softmax.cu new file mode 100644 index 00000000..9b88ee66 --- /dev/null +++ b/ex_engine/factors/factor_topk_softmax.cu @@ -0,0 +1,456 @@ +// 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 +// → 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 +// → 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 +// → 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 +#include +#include +#include + +// CUB for BlockReduce (non-power-of-2 fallback path) +#include + +// ========================================================================= +// 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 +__device__ __forceinline__ float to_float(T val); + +template <> +__device__ __forceinline__ float to_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 +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 +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 +__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 row_chunk_raw[VPT]; + AccessType* vec_ptr = reinterpret_cast(&row_chunk_raw); + const AccessType* src_ptr = reinterpret_cast(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 +__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; + __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; + +template +__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; + __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 +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 + <<>>( + 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( \ + 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 + <<>>( + gating_output, workspace, num_experts); + + moe_topk_fast_kernel + <<>>( + workspace, topk_weights, topk_indices, + num_experts, topk, 0, num_experts, renormalize); + + cudaFree(workspace); + } +} diff --git a/ex_engine/factors/hw_config.h b/ex_engine/factors/hw_config.h new file mode 100644 index 00000000..4c9665b3 --- /dev/null +++ b/ex_engine/factors/hw_config.h @@ -0,0 +1,245 @@ +// 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 + +#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 diff --git a/ex_engine/factors/ilu_ops_api.h b/ex_engine/factors/ilu_ops_api.h new file mode 100644 index 00000000..83dec0ad --- /dev/null +++ b/ex_engine/factors/ilu_ops_api.h @@ -0,0 +1,182 @@ +// 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 +#include +#include +#include +#include +#include + +// ========================================================================= +// 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& value, + torch::Tensor& key_cache, + std::optional& value_cache, + torch::Tensor& slot_mapping); + +void batch_prefill( + torch::Tensor& query, + const torch::Tensor& key, + const std::optional& value, + torch::Tensor& output, + std::optional& output_lse, + const std::optional& q_cu_seq_lens, + const std::optional& kv_cu_seq_lens, + const std::optional& alibi_slope, + const std::optional& attn_bias, + const std::optional& q_quant_scale, + const std::optional& k_quant_scale, + const std::optional& 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& v_cache, + std::optional& output_lse, + const std::optional& q_quant_scale, + const std::optional& k_cache_quant_scale, + const std::optional& v_cache_quant_scale, + const std::optional& out_quant_scale, + const std::optional& alibi_slope, + const std::optional& 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& residual, + torch::Tensor& weight, + std::optional& bias, + std::optional& 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 bias); + +// ---- MoE ops (Layer 5: fused_moe.cpp) ---- + +// Step 1: Router — softmax + topk +std::tuple moe_active_topk( + const torch::Tensor& input, + int64_t topk, + int64_t num_expert_group, + int64_t topk_group, + bool normalize, + const std::optional& mask, + const std::string& normed_by, + const std::string& scoring_func, + double route_scale, + const std::optional& e_score_correction_bias); + +// Step 2: Generate permutation indices +std::vector 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& 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 diff --git a/ex_engine/factors/ixformer_infer.h b/ex_engine/factors/ixformer_infer.h new file mode 100644 index 00000000..d58db2cb --- /dev/null +++ b/ex_engine/factors/ixformer_infer.h @@ -0,0 +1,246 @@ +// 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 +#include +#include +#include + +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& alibi_slopes, + const std::optional& sinks, + std::optional& 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& 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& 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& bias, + const std::optional& out, + const std::optional persistent); + +// ixformer linear extended: simplified interface +// BI-V100 status: PRESENT +torch::Tensor ixformer_linear_ex( + torch::Tensor& input, + torch::Tensor& weight, + const c10::optional& bias, + const c10::optional& 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& 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& 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& expert_mask, + const c10::optional& expert_sizes_cpu, + const c10::optional& 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& 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& dst_to_src, + const c10::optional& 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& mul_weight, + const c10::optional& mask, + const c10::optional& extra_residual, + double scaling_factor); + +} // namespace infer +} // namespace ixformer + +#endif // EX_FACTORS_IXFORMER_INFER_H diff --git a/ex_engine/factors/kernel_elem_ops.cpp b/ex_engine/factors/kernel_elem_ops.cpp new file mode 100644 index 00000000..f4e324d9 --- /dev/null +++ b/ex_engine/factors/kernel_elem_ops.cpp @@ -0,0 +1,210 @@ +// 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 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& residual, + torch::Tensor& weight, + std::optional& bias, + std::optional& 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& 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()); + + 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& value, + torch::Tensor& key_cache, + std::optional& 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 bias) { + + return infer::ixformer_linear_ex(a, b, bias, std::nullopt); +} + +} // namespace ilu +} // namespace kernel +} // namespace xllm diff --git a/ex_engine/factors/kernel_moe_ops.cpp b/ex_engine/factors/kernel_moe_ops.cpp new file mode 100644 index 00000000..9db0d227 --- /dev/null +++ b/ex_engine/factors/kernel_moe_ops.cpp @@ -0,0 +1,155 @@ +// 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 moe_active_topk( + const torch::Tensor& input, + int64_t topk, + int64_t num_expert_group, + int64_t topk_group, + bool normalize, + const std::optional& mask, + const std::string& normed_by, + const std::string& scoring_func, + double route_scale, + const std::optional& 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 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 diff --git a/ex_engine/factors/moe_pipeline.py b/ex_engine/factors/moe_pipeline.py new file mode 100644 index 00000000..16f93570 --- /dev/null +++ b/ex_engine/factors/moe_pipeline.py @@ -0,0 +1,461 @@ +""" +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