Qwen3-30B-A3B decode on bs=8/64-new-tokens jumps from 55.6 tok/s
(eager) to 374.3 tok/s at 4bit / 421.6 tok/s at bf16 — 4-7x over pure
Hugging Face generate (66.7 / 101.1 tok/s) on B200. Output is bit-exact
against eager decode on all 32 greedy tokens across 3 prompts.
FlexMoEInference.capture_decode_cudagraph now mirrors the dense
FlexInference capture pattern: reserve a scratch page per batch_idx
slot, allocate static input_ids / batch_idx / outputs buffers, capture
one graph per bucket in [1, 2, 4, 8] + range(16, max_bs+1, 16), share
the CUDA memory pool across buckets, erase scratch pages at the end.
_decode_step replays on subsequent calls: zero graph buffers (except
outputs), copy live input_ids / batch_idx into the pinned buffers,
graph.replay(), slice outputs. Matches flex_qwen3_llama.py:725-801
verbatim.
Capture is gated on select_moe_backend() == "grouped_mm". Other
backends (unsloth_triton, native_torch) print a warning and leave
cudagraph_captured = False so generate() silently falls back to eager.
The legacy NotImplementedError rationale — torch.where + Python
for-loop over experts — only applied to native_torch; grouped_mm's
bincount + cumsum + argsort + torch._grouped_mm + index_add_ path is
fully capture-friendly on H100/B200.
FlexEngine.__init__ drops the blanket capture_cudagraph=False force
for qwen3_moe arch since capture correctness is now backend-aware.
tests/flex_moe_bench.py gains a hf_naive backend that imports pure
transformers (no unsloth patches) for a clean "fast inference vs naive
HF" comparison: flex 374.3 vs 66.7 tok/s at 4bit = 5.6x, flex 421.6
vs 101.1 tok/s at bf16 = 4.2x.
tests/flex_moe_parity.py adds greedy-decode parity across three
prompts with flex-captured / flex-eager / HF-pure. Results at bf16:
flex-captured matches flex-eager 32/32 on all three prompts (CUDA
graph is numerically identical to eager), and matches HF pure
transformers 32/32 on 2/3 prompts with the remaining divergence at a
tie-break logit boundary.
Pairs with an unsloth-zoo PR dropping @torch.compiler.disable on the
sparse MoE block and replacing torch.bincount with a capture-safe
scatter_add_ — without those, capture trips on a CPU→CUDA scalar copy
inside bincount(minlength=python_int).
80aaf1121b