From 461addf428b06caea47c3b17b88063e9c3af6d0f Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 17 Aug 2026 05:32:19 +0000 Subject: [PATCH] =?UTF-8?q?Revert=20"feat:=2010-file=20algorithm=20factor?= =?UTF-8?q?=20system=20=E2=80=94=20full=2010-layer=20AST=20call=20chain"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 3581dd5435646a4208c70b9f3bf84350dfc11fe4. --- ex_engine/factors/attn_dispatch.py | 270 ---------- ex_engine/factors/factor_moe_combine.cu | 154 ------ ex_engine/factors/factor_moe_compute_index.cu | 174 ------- ex_engine/factors/factor_topk_softmax.cu | 456 ----------------- ex_engine/factors/hw_config.h | 245 ---------- ex_engine/factors/ilu_ops_api.h | 182 ------- ex_engine/factors/ixformer_infer.h | 246 ---------- ex_engine/factors/kernel_elem_ops.cpp | 210 -------- ex_engine/factors/kernel_moe_ops.cpp | 155 ------ ex_engine/factors/moe_pipeline.py | 461 ------------------ 10 files changed, 2553 deletions(-) delete mode 100644 ex_engine/factors/attn_dispatch.py delete mode 100644 ex_engine/factors/factor_moe_combine.cu delete mode 100644 ex_engine/factors/factor_moe_compute_index.cu delete mode 100644 ex_engine/factors/factor_topk_softmax.cu delete mode 100644 ex_engine/factors/hw_config.h delete mode 100644 ex_engine/factors/ilu_ops_api.h delete mode 100644 ex_engine/factors/ixformer_infer.h delete mode 100644 ex_engine/factors/kernel_elem_ops.cpp delete mode 100644 ex_engine/factors/kernel_moe_ops.cpp delete mode 100644 ex_engine/factors/moe_pipeline.py diff --git a/ex_engine/factors/attn_dispatch.py b/ex_engine/factors/attn_dispatch.py deleted file mode 100644 index 958232a7..00000000 --- a/ex_engine/factors/attn_dispatch.py +++ /dev/null @@ -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 diff --git a/ex_engine/factors/factor_moe_combine.cu b/ex_engine/factors/factor_moe_combine.cu deleted file mode 100644 index 1096507b..00000000 --- a/ex_engine/factors/factor_moe_combine.cu +++ /dev/null @@ -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 -#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 deleted file mode 100644 index 20cd578f..00000000 --- a/ex_engine/factors/factor_moe_compute_index.cu +++ /dev/null @@ -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 -#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 deleted file mode 100644 index 9b88ee66..00000000 --- a/ex_engine/factors/factor_topk_softmax.cu +++ /dev/null @@ -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 -// → 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 deleted file mode 100644 index 4c9665b3..00000000 --- a/ex_engine/factors/hw_config.h +++ /dev/null @@ -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 - -#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 deleted file mode 100644 index 83dec0ad..00000000 --- a/ex_engine/factors/ilu_ops_api.h +++ /dev/null @@ -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 -#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 deleted file mode 100644 index d58db2cb..00000000 --- a/ex_engine/factors/ixformer_infer.h +++ /dev/null @@ -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 -#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 deleted file mode 100644 index f4e324d9..00000000 --- a/ex_engine/factors/kernel_elem_ops.cpp +++ /dev/null @@ -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 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 deleted file mode 100644 index 9db0d227..00000000 --- a/ex_engine/factors/kernel_moe_ops.cpp +++ /dev/null @@ -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 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 deleted file mode 100644 index 16f93570..00000000 --- a/ex_engine/factors/moe_pipeline.py +++ /dev/null @@ -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