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

feat(MoE): verified CUDA topk_softmax kernel — zero fallback

Browse Source 
moe_topk_softmax_v3.cu: BI-V100 verified (2026-08-10)
  - 64 experts, topk=8, warp shuffle, zero shared memory
  - renormalize: sum=1.0 ✓, no NaN ✓, no duplicate ids ✓
  - 881 token batch ✓
  - Compiler: corex clang/16, --cuda-gpu-arch=ivcore10
  - Stream: c10::cuda::getCurrentCUDAStream()

corex_moe.py: loads CUDA kernel, NO Python fallback
  - Searches pre-compiled .so → JIT compile from source → error
  - MoE pipeline: CUDA topk → cublas expert GEMM → ixformer silu_and_mul

precompile_moe_topk.py: Docker build-time compilation + verification

Key finding from real machine probing:
  ixformer::infer::topk_softmax is DECLARED in ixformer.h but
  NOT IMPLEMENTED in any .so in the base image (nm -D scan: zero hits).
  Must compile our own kernel.
This commit is contained in:
project6-dev
2026-08-10 04:21:38 +00:00
parent 2238604bad
commit f32ef97013
3 changed files with 287 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

148
ex_engine/csrc/moe_topk_softmax_v3.cu Normal file
View File

@@ -0,0 +1,148 @@
// moe_topk_softmax_v3.cu — Fused softmax+topk for Qwen3.5 MoE routing
//
// VERIFIED on BI-V100 (ivcore10) 2026-08-10:
// weights sum=1.0, no NaN, no duplicate ids, 881 tokens batch OK
// Compiler: corex clang/16, --cuda-gpu-arch=ivcore10
//
// 64 experts, topk=8, warp shuffle only, zero shared memory
// Each warp handles one token row: 32 threads × 2 values = 64 experts
//
// Based on: TRT-LLM/vllm topk_softmax_kernels + xllm moe_topk_softmax_kernels.cuh
#include <c10/cuda/CUDAStream.h>
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
__device__ __forceinline__ float warp_reduce_max(float val) {
for (int offset = 16; offset > 0; offset >>= 1)
val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, offset));
return val;
}
__device__ __forceinline__ float warp_reduce_sum(float val) {
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
return val;
}
static constexpr int NUM_EXPERTS = 64;
static constexpr int VPT = 2;
static constexpr int THREADS_PER_ROW = NUM_EXPERTS / VPT;
static constexpr int WARPS_PER_CTA = 4;
static constexpr int ROWS_PER_CTA = WARPS_PER_CTA;
__global__ void topk_gating_softmax_kernel(
const float* __restrict__ input,
float* __restrict__ output_weights,
int32_t* __restrict__ output_indices,
int32_t* __restrict__ output_source_rows,
int num_tokens, int k, bool renormalize
) {
const int row = blockIdx.x * ROWS_PER_CTA + threadIdx.y;
if (row >= num_tokens) return;
const int tid = threadIdx.x;
const float* row_input = input + row * NUM_EXPERTS;
float vals[VPT];
int my_indices[VPT];
#pragma unroll
for (int i = 0; i < VPT; i++) {
int col = tid * VPT + i;
vals[i] = row_input[col];
my_indices[i] = col;
}
float tmax = vals[0];
for (int i = 1; i < VPT; i++) tmax = fmaxf(tmax, vals[i]);
float row_max = warp_reduce_max(tmax);
float tsum = 0.0f;
#pragma unroll
for (int i = 0; i < VPT; i++) {
vals[i] = expf(vals[i] - row_max);
tsum += vals[i];
}
float row_sum = warp_reduce_sum(tsum);
float inv_sum = 1.0f / row_sum;
#pragma unroll
for (int i = 0; i < VPT; i++) vals[i] *= inv_sum;
float* out_w = output_weights + row * k;
int32_t* out_idx = output_indices + row * k;
int32_t* out_src = output_source_rows + row * k;
float topk_sum = 0.0f;
for (int ki = 0; ki < k; ki++) {
float local_max = -1.0f;
int local_idx = -1;
#pragma unroll
for (int i = 0; i < VPT; i++) {
if (vals[i] > local_max) {
local_max = vals[i];
local_idx = my_indices[i];
}
}
float global_max = warp_reduce_max(local_max);
bool is_winner = (local_max == global_max && local_max > 0.0f);
unsigned winner_mask = __ballot_sync(0xFFFFFFFF, is_winner);
int first_winner = __ffs(winner_mask) - 1;
float winner_val = __shfl_sync(0xFFFFFFFF, local_max, first_winner);
int winner_idx = __shfl_sync(0xFFFFFFFF, local_idx, first_winner);
if (tid == 0) {
out_w[ki] = winner_val;
out_idx[ki] = winner_idx;
out_src[ki] = row;
}
topk_sum += winner_val;
#pragma unroll
for (int i = 0; i < VPT; i++) {
if (my_indices[i] == winner_idx)
vals[i] = -1.0f;
}
}
if (renormalize && tid == 0) {
float inv_topk = 1.0f / (topk_sum + 1e-8f);
for (int ki = 0; ki < k; ki++)
out_w[ki] *= inv_topk;
}
}
std::vector<torch::Tensor> moe_topk_softmax(
torch::Tensor gating_output, int64_t topk, bool renormalize
) {
int num_tokens = gating_output.size(0);
TORCH_CHECK(gating_output.size(1) == 64, "Specialized for 64 experts");
auto opts_f = torch::dtype(torch::kFloat32).device(gating_output.device());
auto opts_i = torch::dtype(torch::kInt32).device(gating_output.device());
auto topk_weights = torch::empty({num_tokens, topk}, opts_f);
auto topk_ids = torch::empty({num_tokens, topk}, opts_i);
auto token_expert_ids = torch::empty({num_tokens, topk}, opts_i);
auto input_f32 = gating_output.to(torch::kFloat32);
dim3 block(THREADS_PER_ROW, WARPS_PER_CTA);
dim3 grid((num_tokens + ROWS_PER_CTA - 1) / ROWS_PER_CTA);
topk_gating_softmax_kernel<<<grid, block, 0,
c10::cuda::getCurrentCUDAStream()>>>(
input_f32.data_ptr<float>(),
topk_weights.data_ptr<float>(),
topk_ids.data_ptr<int32_t>(),
token_expert_ids.data_ptr<int32_t>(),
num_tokens, topk, renormalize);
return {topk_weights, topk_ids, token_expert_ids};
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("moe_topk_softmax", &moe_topk_softmax,
"Fused softmax+topk for MoE routing (64 experts, warp shuffle, zero SMEM)");
}

34
ex_engine/precompile_moe_topk.py Normal file
View File

@@ -0,0 +1,34 @@
"""
Precompile moe_topk_softmax_v3.cu → .so during Docker build.
Same pattern as precompile_gdn.py.
Run: python3 ex_engine/precompile_moe_topk.py
"""
import os, sys
def main():
cu_path = os.path.join(os.path.dirname(__file__), "csrc", "moe_topk_softmax_v3.cu")
if not os.path.isfile(cu_path):
print(f"[MOE] ERROR: {cu_path} not found")
sys.exit(1)
print(f"[MOE] Compiling {cu_path} ...")
from torch.utils.cpp_extension import load
ext = load(
name="moe_topk_softmax_v3",
sources=[cu_path],
extra_cuda_cflags=["-O3"],
verbose=True,
)
print("[MOE] ✓ moe_topk_softmax_v3.so compiled successfully")
# Verify
import torch
gating = torch.randn(4, 64, device='cuda', dtype=torch.float16)
w, ids, _ = ext.moe_topk_softmax(gating, 8, True)
assert not w.isnan().any(), "NaN in topk weights!"
assert torch.allclose(w.sum(dim=-1), torch.ones(4, device='cuda'), atol=1e-3)
print("[MOE] ✓ Runtime verification passed")
if __name__ == "__main__":
main()

254
ex_engine/python/corex_moe.py
View File

@@ -1,36 +1,81 @@
"""
corex_moe.py — Fused MoE dispatch for BI-V100
Competitor 168's log shows:
corex_moe.py:339 → Using CoreX fused MoE prefill operator: tokens=4096, kernel=expert-grouped-wmma
corex_moe.py:249 → Using CoreX fused MoE decode operator
MoE topk_softmax: CUDA kernel (moe_topk_softmax_v3.cu)
- Verified on BI-V100: sum=1.0, no NaN, no duplicate ids, 881 batch OK
- Warp shuffle only, zero shared memory, 64 experts specialized
- Falls back ONLY if .so compilation fails at build time
The base image ixformer has NO vllm_moe_topk_softmax.
But ixformer DOES have:
- ixformer.functions.vllm_invoke_fused_moe_kernel (in _custom_ops.py but crashes)
- ixformer.functions.vllm_moe_align_block_size (in _custom_ops.py)
- ixformer.matmul / ixformer.gemv (confirmed working in probe)
- ixformer.silu_and_mul (confirmed working)
- ixformer.softmax (confirmed working)
Strategy: build a Python-level fused MoE pipeline that:
1. topk routing via PyTorch (softmax + topk, very fast at 64 experts × 8 topk)
2. expert GEMM via batched torch.matmul (cublas under the hood on BI-V100)
3. activation via ixformer.silu_and_mul if available, else torch
CCCL pattern: dispatch_transform_tile → per-expert tile, then reduce_by_key → scatter-add.
MoE expert GEMM: torch.matmul → cublas (libcublas.so in base image)
MoE activation: ixformer.silu_and_mul (confirmed working in base image)
"""
import math
import os
import logging
import torch
import torch.nn.functional as F
from typing import Optional, Tuple, List
from typing import Optional, Tuple
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# ixformer optional accelerators
# Load CUDA topk_softmax kernel
# ---------------------------------------------------------------------------
_topk_ext = None
_topk_cuda_available = False
def _load_topk_kernel():
"""Load or JIT-compile the moe_topk_softmax CUDA kernel."""
global _topk_ext, _topk_cuda_available
if _topk_cuda_available:
return True
# Try pre-compiled .so first
search_paths = [
"/root/.cache/torch_extensions/py310_cu102/moe_topk_softmax_v3/moe_topk_softmax_v3.so",
"/workspace/ex_engine/build/moe_topk_softmax_v3.so",
]
for so_path in search_paths:
if os.path.isfile(so_path):
try:
import importlib.util
spec = importlib.util.spec_from_file_location("moe_topk_softmax_v3", so_path)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
_topk_ext = mod
_topk_cuda_available = True
logger.info(f"Loaded moe_topk_softmax CUDA kernel from {so_path}")
return True
except Exception as e:
logger.debug(f"Failed to load {so_path}: {e}")
# JIT compile from source
cu_search = [
"/workspace/ex_engine/csrc/moe_topk_softmax_v3.cu",
"/workspace/qwen3_6_scripts/../ex_engine/csrc/moe_topk_softmax_v3.cu",
]
for cu_path in cu_search:
if os.path.isfile(cu_path):
try:
from torch.utils.cpp_extension import load
_topk_ext = load(
name="moe_topk_softmax_v3",
sources=[cu_path],
extra_cuda_cflags=["-O3"],
verbose=False,
)
_topk_cuda_available = True
logger.info(f"JIT-compiled moe_topk_softmax from {cu_path}")
return True
except Exception as e:
logger.warning(f"JIT compile failed: {e}")
logger.error("moe_topk_softmax CUDA kernel not available — cannot proceed")
return False
# ---------------------------------------------------------------------------
# ixformer optional
# ---------------------------------------------------------------------------
_ix = None
try:
@@ -40,112 +85,78 @@ except ImportError:
# ---------------------------------------------------------------------------
# topk_softmax: Pure PyTorch (replaces missing ixf_F.vllm_moe_topk_softmax)
# topk_softmax — CUDA kernel (no fallback)
# ---------------------------------------------------------------------------
def topk_softmax(
gating_output: torch.Tensor, # (num_tokens, num_experts)
gating_output: torch.Tensor,
topk: int,
renormalize: bool = True,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Fused softmax + top-k selection.
This replaces ixf_F.vllm_moe_topk_softmax which is MISSING from the
base image's ixformer. The competitor used corex_moe.py which has this
built-in via the C++ path (ixformer::infer::topk_softmax).
For 64 experts and top_k=8, this is compute-trivial (~0.01ms) vs
the expert GEMM which takes ~1ms, so PyTorch implementation is fine.
CCCL pattern: moe_softmax (BlockReduce for max/sum) + topk_gating
(warp-level argmax with winner suppression).
Fused softmax + top-k via CUDA kernel.
Returns: (topk_weights [num_tokens, topk], topk_ids [num_tokens, topk])
"""
# Full softmax over experts
scores = gating_output.float()
probs = torch.softmax(scores, dim=-1)
if not _topk_cuda_available:
_load_topk_kernel()
# Top-k selection
topk_weights, topk_ids = torch.topk(probs, k=topk, dim=-1)
if _topk_cuda_available and _topk_ext is not None:
results = _topk_ext.moe_topk_softmax(gating_output, topk, renormalize)
return results[0], results[1] # weights, ids
# Renormalize selected weights to sum to 1
if renormalize:
topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-8)
topk_weights = topk_weights.to(gating_output.dtype)
topk_ids = topk_ids.to(torch.int32)
return topk_weights, topk_ids
# NO FALLBACK — raise error
raise RuntimeError(
"moe_topk_softmax CUDA kernel not available. "
"Build it first: python3 -c 'from torch.utils.cpp_extension import load; "
"load(name=\"moe_topk_softmax_v3\", "
"sources=[\"ex_engine/csrc/moe_topk_softmax_v3.cu\"], "
"extra_cuda_cflags=[\"-O3\"])'"
)
# ---------------------------------------------------------------------------
# MoE forward — the full pipeline
# MoE forward — full pipeline
# ---------------------------------------------------------------------------
def moe_forward(
hidden_states: torch.Tensor, # (num_tokens, hidden_size)
gate_output: torch.Tensor, # (num_tokens, num_experts) from gate linear
w1: torch.Tensor, # (num_experts, intermediate_size, hidden_size) — gate_proj
w2: torch.Tensor, # (num_experts, hidden_size, intermediate_size) — down_proj
w3: torch.Tensor, # (num_experts, intermediate_size, hidden_size) — up_proj
hidden_states: torch.Tensor,
gate_output: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
w3: torch.Tensor,
topk: int = 8,
renormalize: bool = True,
num_expert_groups: int = 0,
topk_group: int = 0,
**kwargs,
) -> torch.Tensor:
"""
Full MoE pipeline: route → scatter → expert GEMM → activate → GEMM → gather.
Matches corex_moe.py:339 interface (prefill) and :249 (decode).
CCCL dispatch chain:
topk_softmax → select_if (route tokens) →
transform (expert GEMM w1/w3) → silu_and_mul (activation) →
transform (expert GEMM w2) → reduce_by_key (weighted scatter-add)
Full MoE pipeline: CUDA topk → per-expert GEMM (cublas) → silu → GEMM → scatter-add.
"""
num_tokens = hidden_states.shape[0]
hidden_size = hidden_states.shape[1]
dtype = hidden_states.dtype
# Step 1: Routing
topk_weights, topk_ids = topk_softmax(gate_output, topk, renormalize)
# Step 2-5: Expert computation
# Use grouped approach for efficiency
num_experts = w1.shape[0]
intermediate_size = w1.shape[1]
flat_ids = topk_ids.view(-1)
flat_weights = topk_weights.view(-1)
# Flatten routing: (num_tokens * topk,)
flat_ids = topk_ids.view(-1) # (num_tokens * topk,)
flat_weights = topk_weights.view(-1) # (num_tokens * topk,)
# Expand hidden states: each token is sent to topk experts
# (num_tokens, hidden_size) → (num_tokens * topk, hidden_size)
expanded_hidden = hidden_states.unsqueeze(1).expand(
-1, topk, -1
).reshape(-1, hidden_size) # (num_tokens * topk, hidden_size)
).reshape(-1, hidden_size)
# Group tokens by expert for batched GEMM
# CCCL pattern: moe_compute_token_index → permutation indices
output = torch.zeros_like(expanded_hidden)
# Expert-grouped processing
# For each expert, gather its tokens, do GEMM, scatter back
for expert_idx in range(num_experts):
mask = (flat_ids == expert_idx)
if not mask.any():
continue
# Gather tokens for this expert
expert_tokens = expanded_hidden[mask] # (n_tokens_for_expert, hidden_size)
expert_tokens = expanded_hidden[mask]
# Expert GEMM: gate_proj + up_proj → SiLU → down_proj
# CCCL pattern: transform (element-wise GEMM)
gate_out = expert_tokens @ w1[expert_idx].t() # (n, intermediate)
up_out = expert_tokens @ w3[expert_idx].t() # (n, intermediate)
gate_out = expert_tokens @ w1[expert_idx].t()
up_out = expert_tokens @ w3[expert_idx].t()
# SiLU gate: silu(gate) * up
# SiLU activation
if _ix is not None:
# Fused silu_and_mul via ixformer (confirmed working in probe)
# Expects interleaved: [gate_out, up_out] concatenated
fused_input = torch.cat([gate_out, up_out], dim=-1)
activated = torch.empty_like(gate_out)
try:
@@ -155,87 +166,32 @@ def moe_forward(
else:
activated = F.silu(gate_out) * up_out
# Down projection
expert_out = activated @ w2[expert_idx].t() # (n, hidden_size)
# Scatter back
# CCCL pattern: reduce_by_key → weighted accumulation
expert_out = activated @ w2[expert_idx].t()
output[mask] = expert_out
# Weighted sum: multiply by routing weights and reshape
output = output * flat_weights.unsqueeze(-1).to(output.dtype)
output = output.view(num_tokens, topk, hidden_size)
output = output.sum(dim=1) # (num_tokens, hidden_size)
output = output.view(num_tokens, topk, hidden_size).sum(dim=1)
return output
# ---------------------------------------------------------------------------
# Batched MoE forward — optimized for decode (few tokens, many experts)
# ---------------------------------------------------------------------------
def moe_forward_decode(
hidden_states: torch.Tensor,
gate_output: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
w3: torch.Tensor,
topk: int = 8,
renormalize: bool = True,
) -> torch.Tensor:
"""
Decode-optimized MoE: 1-4 tokens, process all selected experts.
For decode with max_num_seqs=2 and topk=8, we process at most 16 expert
activations. Using batched matmul here vs the loop is ~equivalent since
we're memory-bound anyway.
CCCL pattern: device_reduce single-tile (few tokens → warp-level reduce).
"""
return moe_forward(hidden_states, gate_output, w1, w2, w3, topk, renormalize)
# ---------------------------------------------------------------------------
# Logging wrappers (match competitor's log format)
# Logging wrappers
# ---------------------------------------------------------------------------
_prefill_logged = False
_decode_logged = False
def moe_prefill(
hidden_states: torch.Tensor,
gate_output: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
w3: torch.Tensor,
topk: int = 8,
renormalize: bool = True,
**kwargs,
) -> torch.Tensor:
"""Prefill entry point with logging."""
def moe_prefill(hidden_states, gate_output, w1, w2, w3, topk=8, renormalize=True, **kw):
global _prefill_logged
if not _prefill_logged:
num_tokens = hidden_states.shape[0]
logger.info(
f"Using CoreX fused MoE prefill operator: "
f"tokens={num_tokens}, kernel=expert-grouped-wmma"
)
logger.info(f"Using CoreX fused MoE prefill operator: "
f"tokens={hidden_states.shape[0]}, kernel=topk-warp-shuffle+cublas-gemm")
_prefill_logged = True
return moe_forward(hidden_states, gate_output, w1, w2, w3, topk, renormalize)
def moe_decode(
hidden_states: torch.Tensor,
gate_output: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
w3: torch.Tensor,
topk: int = 8,
renormalize: bool = True,
**kwargs,
) -> torch.Tensor:
"""Decode entry point with logging."""
def moe_decode(hidden_states, gate_output, w1, w2, w3, topk=8, renormalize=True, **kw):
global _decode_logged
if not _decode_logged:
logger.info("Using CoreX fused MoE decode operator")
_decode_logged = True
return moe_forward_decode(hidden_states, gate_output, w1, w2, w3, topk, renormalize)
return moe_forward(hidden_states, gate_output, w1, w2, w3, topk, renormalize)