feat: batched MoE expert GEMM — replaces Python for-loop

ixformer probe results:
  ✗ moe_w16a16_group_gemm NOT in ixformer .so
  ✗ CUTLASS grouped GEMM needs cuda/std (variadic function error on corex)
  ✓ ixformer_linear EXISTS (fused matmul)
  ✓ torch.mm works (uses corex cublas)

Solution: moe_batched_gemm.cu
  - C++ loop over experts (eliminates Python overhead)
  - torch::mm for GEMM (corex cublas, not F.linear Python)
  - Fused silu_and_mul CUDA kernel (not PyTorch ops)
  - Weighted scatter-add in C++
  - Skips empty experts (no wasted compute)

Integration in qwen3_5.py:
  _USE_XLLM_MOE_GEMM dispatches to moe_experts_forward()
  Falls back to Python for-loop if not available

Build: bash qwen3_6_scripts/build_xllm_kernels.sh
This commit is contained in:
claude
2026-08-14 11:43:46 +00:00
parent 93353a1414
commit 06d7713db6
3 changed files with 160 additions and 14 deletions

View File

@@ -67,5 +67,8 @@ build_kernel "xllm_cache" \
build_kernel "xllm_moe" \
"${CUDA_DIR}/moe/moe_fused_topk.cu" "${CUDA_DIR}/moe/moe_compute_index.cu" "${CUDA_DIR}/moe/moe_combine.cu" "${BIND_DIR}/xllm_moe_bind.cpp"
build_kernel "xllm_moe_gemm" \
"${CUDA_DIR}/moe_batched_gemm.cu"
echo "=== All kernels built ==="
ls -lh "${PREBUILT_DIR}"/xllm_*.so 2>/dev/null || echo "No .so files found"