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

feat: pybind wrapper for CUTLASS batched GEMM → MoE decode path

Browse Source 
Based on verified result (issue #68):
  CUTLASS Cu10 TensorOp batched: 2.462ms (8 experts, 1 launch)
  vs 8× torch.matmul: 4.6ms (8 launches)
  vs Python F.linear loop: 10.36ms

New files:
  ex_engine/xllm_kernels/cuda/bindings/corex_batched_gemm_bind.cpp
    pybind11 wrapper: batched_gemm_fp16() + moe_decode_fused()
  ex_engine/xllm_kernels/cuda/corex_batched_gemm_kernel.cu
    CUTLASS GemmBatched<half> kernel (from cat_files/batched_gemm.cu)
  qwen3_6_scripts/build_corex_batched_gemm.sh
    Build script for BI-V100 (ivcore10)

Modified:
  qwen3_6_scripts/qwen3_5.py
    import corex_batched_gemm + _USE_COREX_BATCHED_GEMM flag
    Tier 1.5 in MoE decode: after corex_direct_routed, before corex_gather

Build on device: bash qwen3_6_scripts/build_corex_batched_gemm.sh
Output: prebuilt/corex-3.2.3-ivcore10/corex_batched_gemm.so
This commit is contained in:
dylan
2026-08-15 11:54:26 +00:00
parent a875fa5d4c
commit ddcfbad431
4 changed files with 332 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
qwen3_6_scripts/qwen3_5.py
View File

@@ -138,6 +138,11 @@ try:
except ImportError:
_corex_moe_direct_routed = None
try:
from vllm import corex_batched_gemm as _corex_batched_gemm
except ImportError:
_corex_batched_gemm = None
try:
from vllm import corex_moe_topk_softmax as _corex_moe_topk_softmax
except ImportError:
@@ -204,6 +209,9 @@ _USE_COREX_MOE_WEIGHT_GATHER = (
_USE_COREX_MOE_DIRECT_ROUTED = (
_corex_moe_direct_routed is not None
and env_bool("BI100_MOE_COREX_DIRECT_ROUTED", False))
_USE_COREX_BATCHED_GEMM = (
_corex_batched_gemm is not None
and env_bool("BI100_MOE_BATCHED_GEMM", True))
_USE_COREX_MOE_TOPK_SOFTMAX = (
_corex_moe_topk_softmax is not None
and env_bool("BI100_MOE_COREX_TOPK_SOFTMAX", True))
@@ -1729,6 +1737,15 @@ class Qwen3_5MoeSparseBlock(nn.Module):
return _corex_moe_direct_routed.w2_reduce(
act, w2, eids, ws)
# Tier 1.5: CUTLASS batched GEMM (verified 2.462ms, issue #68)
# 1 launch for 8 experts vs 8 launches for F.linear loop
if (_USE_COREX_BATCHED_GEMM
and hidden_states.dtype == torch.float16
and w13.dtype == torch.float16
and w2.dtype == torch.float16):
return _corex_batched_gemm.moe_decode_fused(
hidden_states, w13[eids], w2[eids], ws)
use_corex_gather = (
_USE_COREX_MOE_WEIGHT_GATHER
and hidden_states.dtype == torch.float16