feat(MoE): verified CUDA topk_softmax kernel — zero fallback
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
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user