Optional ``FlexMoEInference(compile_walker=True)`` / env var
``UNSLOTH_FLEX_COMPILE_WALKER=1`` wraps the decode walker
(``call_moe_model_with_flex_kwargs``) with
``torch.compile(fullgraph=False, dynamic=False)`` before the CUDA
graph capture kicks in. Inductor fuses the layernorm + residual +
router pointwise ops, and the compiled kernels end up recorded
inside the captured graph. Net: ~2x decode tok/s on the grouped_mm
path with no change in VRAM, no correctness regression, and no
user-facing API change unless the flag is set.
Numbers on Qwen3-30B-A3B-Instruct-2507, B200, 128 new tokens,
bs sweep, median of 2 timed rounds after 1 warmup:
| precision | bs | baseline (v2) | + compile_walker | speedup |
|-----------|---:|--------------:|-----------------:|--------:|
| 4bit | 16 | 699 | 1347.8 | 1.93x |
| 4bit | 32 | 1243 | 2423.9 | 1.95x |
| 4bit | 48 | 1735 | 3383.2 | 1.95x |
| 4bit | 64 | 1523 | 2981.9 | 1.96x |
| bf16 | 16 | — | 1401.1 | — |
| bf16 | 32 | — | 2864.0 | — |
| bf16 | 48 | — | **3911.1** | — |
Peak throughput: 3911 tok/s at bf16 bs=48 — 39x the pure-HF naive
baseline on the same workload (101.1 tok/s with
``AutoModelForCausalLM`` + eager attn, left-padded, no unsloth).
At 4bit bs=48, 51x the pure-HF naive baseline (66.7 tok/s).
GRPO 5-step validation (Qwen3_MoE_GRPO.py --backend flex
--max_steps 5 on DAPO-Math-17k):
| precision | baseline (v2) | + compile_walker | speedup |
|-----------|--------------:|-----------------:|--------:|
| 4bit | 548.3s | 434.5s | 1.26x |
| bf16 | 451.8s | **407.4s** | 1.10x |
Peak VRAM unchanged (130-133 GB). Loss / KL stable on both, no
NaN, rewards pegged at -7.5 (base-model artifact; orthogonal).
Parity (greedy 32 tokens × 3 prompts at bf16 and 4bit via
``FLEX_MOE_COMPILE_WALKER=1 tests/flex_moe_parity.py``):
flex-captured with the compile wrap matches flex-captured without
the compile wrap on 6/6 prompts with no gibberish, and matches pure
``transformers.AutoModelForCausalLM`` 32/32 on 5 of 6 (prompt ×
precision) pairs (the one divergence is a tie-break logit boundary
on an open-ended continuation — both coherent English).
Bisection of a few torch.compile flag sets against the default at
bs=32 4bit (max_batch_size=32):
| config | tok/s |
|------------------------------------------------------------|-------:|
| default (``torch.compile(fullgraph=False, dynamic=False)``)| 1581.5 |
| + max_autotune + coord_descent + aggressive_fusion | 1704.9 |
| + ``freezing=True`` | 935.7 |
``freezing=True`` is a regression on this path; shipping with the
default config only. The other flags are +7.8% at this size but
at large bs (48+) the max_autotune variant timed out during
compile (>40 min) so the default stays the ship-target for now.
Other attention backends don't help on B200 today:
- pure HF with ``attn_implementation="sdpa"``: cuDNN Frontend error
("No valid execution plans built") on sm_100 + torch 2.11.
- ``flash_attention_2`` 2.8.3: works, but kernels compiled for
sm_80/sm_90 only — slower than eager on B200 (46.7 / 67.9 tok/s
vs eager 66.7 / 101.1 at 4bit / bf16).
- ``flash_attention_3``: ``no kernel image for execution on the
device`` — sm_100 kernels not yet in flash_attn_interface.
- FA4 / ``flash_attention_4``: works standalone but transformers'
integration hard-codes ``flash_attn_with_kvcache = None`` for it,
so it can't service decode. Prefill-only, out of scope here.
New tests:
- ``tests/flex_moe_micro_bench.py``: tight probe that loads the
model once, sweeps batch sizes, prints a sample completion per
bucket (catches gibberish early). Supports ``--compile_mode
{off, walker, walker_fullgraph}`` and ``--compile_opts
{stock, unsloth_O3, inference_freeze}``.
- ``tests/flex_moe_bench.py``: add ``--backend hf_naive`` which
imports pure ``transformers`` (no ``import unsloth``) for the
reference HF baseline, with ``HF_ATTN_IMPL`` env var to switch
between eager / sdpa / flash_attention_{2,3,4}.
- ``tests/flex_moe_parity.py``: add ``FLEX_MOE_COMPILE_WALKER=1``
env var to exercise the compile wrap through the parity harness.
cc4832b605