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
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
The .so's torch_chunk_gated_delta_rule() only accepts positional args:
(Tensor, Tensor, Tensor, Tensor, Tensor, int, Optional[Tensor], bool, bool)
But Python calls it with keyword args:
(q, k, v, g, beta, initial_state=, output_final_state=, use_qk_l2norm_in_kernel=)
This causes 'incompatible function arguments' crash during profiling
(determine_num_available_blocks), killing the engine before it starts.
Fix: _HAS_COREX_GDN_CHUNK = False, forcing Python _torch_chunk_gated_delta_rule.
This is what a3c45d3b effectively did (its .so wasn't compiled), explaining
why a3c45d3b works but aa4b4992 crashes.
31 files had Windows line endings (\r\n) from merge commit. This causes
patch_ops.sh replace_once() to fail: anchor strings use \n but file
content has \r\n, so no match → patch fails → docker build fails.
Also added .gitattributes to force LF for all text files going forward.
- patch_ops.sh: build corex_gdn_chunk_recurrent.so alongside moe_index_combine
- qwen3_5.py: import corex_gdn_chunk_recurrent, use C++ version for prefill
chunks instead of Python _torch_chunk_gated_delta_rule
- C++ version from xllm upstream avoids Python loop overhead and has proper
fp32 accumulation (key for NaN prevention on BI-V100)
- Falls back to Python version if .so not available
Verified on real BI-V100:
moe_compute_index: 11.48x speedup (0.035ms vs 0.397ms)
moe_combine_result: 2.66x speedup (0.022ms vs 0.059ms)
Integration:
- qwen3_5.py: import corex_moe_index_combine, use in prefill path
with _USE_COREX_MOE_INDEX_COMBINE flag (env BI100_MOE_COREX_INDEX_COMBINE)
Falls back to PyTorch argsort+bincount if .so unavailable
- patch_ops.sh: compile corex_moe_index_combine.cu during docker build
Previous xformers.py fix only covered our attention backend. The crash
moved to qwen2_vl.py's Qwen2VisionAttention.forward (base image file)
which directly calls xops.memory_efficient_attention_forward during
profiling's _process_image_input → visual() → block.attn().
Fix: monkey-patch Qwen2VisionAttention.forward at import time to use
the same PyTorch F.scaled_dot_product_attention path that qwen2_vl.py
already has for CPU (is_cpu() branch). This is the exact same math,
just without xops dispatch to ixformer's broken varlen_fwd.
Also added try/except fallback in _process_image_input for safety.
Reverted the NO-FALLBACK rewrite of corex_gdn.py and qwen3_5.py.
Policy: do NOT rewrite modules that already exist in base image or
upstream_ref. If an interface doesn't match, fix the interface call
site — don't rewrite the entire module in pure PyTorch.
Base image has corex_gdn.py, corex_moe.py, corex_fa2.py with C++
backends. The right approach is to match their __init__ signatures,
not replace them with slower Python reimplementations.
Root cause from latest docker build log:
ValueError: You set image=0 in --limit-mm-per-prompt, but found 1 items
→ Engine background task crashes → AsyncEngineDeadError → all subsequent 503
Fixes:
1. computility-run.yaml: add --limit-mm-per-prompt image=1
Prevents multimodal ValueError from killing the engine process.
2. patch_ops.sh: DON'T overwrite base image's corex_gdn.py/corex_moe.py
Comp 168 log proves base image's corex modules work with libcorex_gdn.so.
Our overwrite broke CoreXGDN.__init__ (unexpected kwarg 'num_v_heads').
Only deploy ours if base has NO corex modules at all.
Also deploy corex_fa2.py if base lacks it.
3. qwen3_5.py: try multiple CoreXGDN init signatures
Base image CoreXGDN may accept different kwargs than ours.
Try kwargs form first, fall back to positional.
4. corex_gdn.py: accept both calling conventions in __init__
Future-proof for when we DO need to deploy ours.
5. Copied upstream_ref headers: ilu_layer_fused_moe.h, ilu_layer_attention.h
Last 2 missing ILU files from xllm. All 14/14 now present.
Previous: except ImportError: pass (silent failure)
Now: logs WHY import failed so we can diagnose from docker logs
Also includes the matmul dtype guard fix:
_ix_matmul only calls ixformer.matmul for float16 tensors
Prevents stderr spam from GDN float32 accumulation path
matmul.cu:149 'Expected input.dtype() == kHalf' error in competition log.
Root cause: _torch_chunk_gated_delta_rule returns fp32 core_out,
passed directly to self.norm() → self.out_proj() which calls ixformer matmul.
Fix: explicit .to(torch.float16) on core_out and z before norm.
Root cause from competition platform log:
/opt/apps/ixformer/functions/matmul.cu:149 'Expected input.dtype() == kHalf'
Repeats ~80 times — every GDN layer token pass calls _ix_matmul with float32
GDN chunked delta rule uses float32 accumulation (correct for precision).
_ix_matmul was calling ixformer.matmul on float32 tensors → stderr spam.
The try/except caught it and fell back to torch.matmul, but the stderr
output floods the log and may slow down inference.
Fix: check a.dtype == torch.float16 before calling ixformer.matmul.
Non-half tensors go directly to torch.matmul — zero stderr noise.
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.
Root cause from real machine test: gdn_forward.cu output abs mean = inf
- gate_raw can be positive → exp(gate) > 1 → state grows exponentially
- Over 64 tokens: exp(2.0)^64 = inf
- PyTorch ref clamps g ∈ [-5, 2] but CUDA kernel did not
Fix:
gdn_forward.cu: clamp gate_raw ∈ [-5, 2] before exp (both kernel variants)
gdn_forward.cu: clamp state ∈ [-65504, 65504] after update (fp16 safe range)
qwen3_5.py: clamp g_3d before passing to SM70 kernel (belt + suspenders)
qwen3_5.py: clamp temporal_state after decode update
