1ae398eeee
corex_moe.py: moe_forward now accepts both formats:
Format A: w1(E,I,H) + w2(E,H,I) + w3(E,I,H) — xllm style, separate gate/up
Format B: w13(E,2*I,H) + w2(E,H,I) + w3=None — vllm style, merged gate_up
Auto-detects by checking if w3 is None, splits w13 internally.
qwen3_5.py:
- Fix corex_moe call: use keyword args (w3=None, topk=self.top_k)
prevents topk integer going to w3 tensor position
- Remove silent fallback on corex_moe failure — raise RuntimeError
with full shape info for diagnosis. Zero score with no error log
is worse than a crash.
220 lines
7.6 KiB
Python
220 lines
7.6 KiB
Python
"""
|
|
corex_moe.py — Fused MoE dispatch for BI-V100
|
|
|
|
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
|
|
|
|
MoE expert GEMM: torch.matmul → cublas (libcublas.so in base image)
|
|
MoE activation: ixformer.silu_and_mul (confirmed working in base image)
|
|
"""
|
|
|
|
import os
|
|
import logging
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from typing import Optional, Tuple
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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",
|
|
# Deployed by patch_ops.sh into vllm models dir
|
|
os.path.join(os.path.dirname(os.path.abspath(__file__)), "moe_topk_softmax_v3.cu"),
|
|
]
|
|
# Also search in vllm model_executor/models/
|
|
try:
|
|
import vllm
|
|
vllm_models = os.path.join(os.path.dirname(vllm.__file__), "model_executor", "models")
|
|
cu_search.append(os.path.join(vllm_models, "moe_topk_softmax_v3.cu"))
|
|
except Exception:
|
|
pass
|
|
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:
|
|
import ixformer as _ix
|
|
except ImportError:
|
|
pass
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# topk_softmax — CUDA kernel (no fallback)
|
|
# ---------------------------------------------------------------------------
|
|
def topk_softmax(
|
|
gating_output: torch.Tensor,
|
|
topk: int,
|
|
renormalize: bool = True,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Fused softmax + top-k via CUDA kernel.
|
|
Returns: (topk_weights [num_tokens, topk], topk_ids [num_tokens, topk])
|
|
"""
|
|
if not _topk_cuda_available:
|
|
_load_topk_kernel()
|
|
|
|
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
|
|
|
|
# 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 — full pipeline
|
|
# ---------------------------------------------------------------------------
|
|
def moe_forward(
|
|
hidden_states: torch.Tensor,
|
|
gate_output: torch.Tensor,
|
|
w1_or_w13: torch.Tensor,
|
|
w2: torch.Tensor,
|
|
w3: Optional[torch.Tensor] = None,
|
|
topk: int = 8,
|
|
renormalize: bool = True,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Full MoE pipeline: CUDA topk → per-expert GEMM (cublas) → silu → GEMM → scatter-add.
|
|
|
|
Accepts two weight formats:
|
|
Format A (xllm style): w1=(E,I,H), w2=(E,H,I), w3=(E,I,H) — gate and up separate
|
|
Format B (vllm style): w13=(E,2*I,H), w2=(E,H,I), w3=None — gate_up merged
|
|
"""
|
|
num_tokens = hidden_states.shape[0]
|
|
hidden_size = hidden_states.shape[1]
|
|
dtype = hidden_states.dtype
|
|
|
|
# Detect weight format
|
|
if w3 is None:
|
|
# Format B: w13 merged — split into w1 (gate) and w3 (up)
|
|
w13 = w1_or_w13
|
|
inter2 = w13.shape[1]
|
|
w1 = w13[:, :inter2 // 2, :] # (E, I, H)
|
|
w3 = w13[:, inter2 // 2:, :] # (E, I, H)
|
|
else:
|
|
w1 = w1_or_w13
|
|
|
|
topk_weights, topk_ids = topk_softmax(gate_output, topk, renormalize)
|
|
|
|
num_experts = w1.shape[0]
|
|
flat_ids = topk_ids.view(-1)
|
|
flat_weights = topk_weights.view(-1)
|
|
|
|
expanded_hidden = hidden_states.unsqueeze(1).expand(
|
|
-1, topk, -1
|
|
).reshape(-1, hidden_size)
|
|
|
|
output = torch.zeros_like(expanded_hidden)
|
|
|
|
for expert_idx in range(num_experts):
|
|
mask = (flat_ids == expert_idx)
|
|
if not mask.any():
|
|
continue
|
|
|
|
expert_tokens = expanded_hidden[mask]
|
|
|
|
gate_out = expert_tokens @ w1[expert_idx].t()
|
|
up_out = expert_tokens @ w3[expert_idx].t()
|
|
|
|
# SiLU activation
|
|
if _ix is not None:
|
|
fused_input = torch.cat([gate_out, up_out], dim=-1)
|
|
activated = torch.empty_like(gate_out)
|
|
try:
|
|
_ix.silu_and_mul(fused_input, activated)
|
|
except Exception:
|
|
activated = F.silu(gate_out) * up_out
|
|
else:
|
|
activated = F.silu(gate_out) * up_out
|
|
|
|
expert_out = activated @ w2[expert_idx].t()
|
|
output[mask] = expert_out
|
|
|
|
output = output * flat_weights.unsqueeze(-1).to(output.dtype)
|
|
output = output.view(num_tokens, topk, hidden_size).sum(dim=1)
|
|
|
|
return output
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Logging wrappers
|
|
# ---------------------------------------------------------------------------
|
|
_prefill_logged = False
|
|
_decode_logged = False
|
|
|
|
def moe_prefill(hidden_states, gate_output, w1, w2, w3, topk=8, renormalize=True, **kw):
|
|
global _prefill_logged
|
|
if not _prefill_logged:
|
|
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, 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(hidden_states, gate_output, w1, w2, w3, topk, renormalize)
|