Probe data (probe_moe_fused_breakdown.sh on BI-V100):
F.linear loop 8 experts: 8.060 ms
bmm pre-transposed full MoE: 6.918 ms ← 14% faster
transpose+contiguous runtime: 22.219 ms ← why CUTLASS was 27ms
Changes:
- Lazy-cache w13_t (E,H,2I) and w2_t (E,I,H) on first decode call
- FC1: torch.bmm(x_expand, w13_t_sel) replaces F.linear(x, w13_sel.reshape)
- FC2: torch.bmm(act, w2_t_sel) replaces torch.bmm(w2_sel, act^T)
- Zero runtime transpose cost after first call
Root cause of 25ms (vs expected 2.5ms):
1. ElementAccumulator was half_t → now float (FP32 accumulation)
2. Missing OpClassTensorOp → was defaulting to OpClassSimt (CUDA cores only)
3. Missing arch::Cu10 → was defaulting to arch::Sm61
With these fixes it should use __ivcorex_matrix_mad_f32x4_f16x4 (TCU)
same as moe_cutlass_batched.cu which benchmarked at 2.462ms.
Sources:
siboehm/SGEMM_CUDA → upstream_ref/sgemm_siboehm/ (25 files)
wangzyon/NVIDIA_SGEMM_PRACTICE → upstream_ref/nvidia_sgemm_practice/ (23 files, filled gaps)
edtallison/sgemm-cuda → upstream_ref/sgemm_edtallison/ (41 files)
All files cat'd one by one from git clone (no --depth).
These are the 3 public SGEMM repos that can compile on CUDA 10.2 + CoreX ivcore10.
Key files for BI-V100 porting:
kernel 10 (warp tiling) — already proven on device with WARPSIZE=64
kernel 11/12 (double buffering) — next optimization target
sgemm.cu + runner.cu — complete build+benchmark harness
CMakeLists.txt — build system reference
Tests xllm_norm, xllm_activation, xllm_rope, xllm_moe, ix_full_bridge,
and 11 corex_*.so modules with correctness checks against PyTorch reference.
Run: python3 test_ex_engine_cuda.py
torch profiler confirmed: torch.mm launches Gemm_tcu_bi_kernel::gemm_h_h_tcu_25
which is BI-V100 TCU (Tensor Compute Unit) hardware-accelerated GEMM.
0.58ms per call vs our custom kernel 7.7ms — TCU is 13x faster.
Python for-loop overhead measured: 0.892 ms/expert = 7.1 ms for 8 experts.
This C++ dispatch eliminates that overhead while using the same TCU kernel.
Three entry points:
- moe_decode: full MoE forward (FC1 + SiLU*mul + FC2) for decode
- moe_prefill: group-by-expert MoE forward for prefill
- moe_expert_gemm_tcu: raw GEMM loop for benchmarking
Only 3 changes from upstream_ref/sgemm_cuda/6_kernel_vectorize.cuh:
1. float → __half for A/B/C data and shared memory
2. float4 vectorized load → 4 scalar half loads (float4 needs 16-byte align)
3. threadResults accumulator stays float (FP32 accumulation)
Everything else identical: same shared mem layout, same indexing,
same A-transpose-while-loading, same thread tile computation.
No WARPSIZE. No cooperative_groups. No cuda::barrier.
Key difference from the reverted batched approach:
- Does NOT use torch::mm in a C++ loop (that was the reverted commit)
- Uses ixformer_torch_ext::ixformer_linear — the base image's optimized GEMM
- Same kernel the competitor (sub 168) uses via corex_moe.py
- Eliminates Python interpreter + dispatcher overhead per expert
- Links against _ixformer_torch.cpython-310.so (already in base image)
Decode: 1 Python call → 8 C++ ixformer_linear (vs 8 Python F.linear)
Prefill: 1 Python call → 64 C++ ixformer_linear (vs 64 Python F.linear)
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
CUTLASS grouped GEMM (example 24) requires SM80 Tensor Core + cuda/std headers.
Cannot compile on corex (same issue as CCCL 3.6 variadic functions).
Alternative path: ix_moe_bridge.so calls ixformer::infer::moe_w16a16_group_gemm
which is BI-V100 optimized grouped GEMM already in the base image.
This probe script checks if the MoE functions exist in ixformer .so
before attempting to build ix_moe_bridge.so.
Run: bash qwen3_6_scripts/probe_ixformer_symbols.sh
