Adds tests/flex_moe_bench.py: 2-round median decode throughput bench
comparing flex (FlexMoEInference) and HF generate on the same
(n_prompts, max_new_tokens, precision) workload. Writes
async_task_outputs/qwen3_moe_grpo_bench/bench_decode_{backend}_{precision}.json.
Also defensively unpacks self.mlp(hidden_states) in
Qwen3MoeDecoderLayer_fast_forward's training branch:
unsloth_zoo.temporary_patches.qwen3_moe.sparse_moe_block_forward
returns a plain tensor for transformers 5.x stacked experts, but the
decoder wrapper unpacked a 2-tuple. The inference branch was already
fixed in the previous commit; the training branch hit the same
ValueError under plain HF generate (no _flag_for_generation).
Bench numbers (Qwen3-30B-A3B, 4bit, rank 16 LoRA, bs=8, 64 new tokens,
B200):
| backend | median tok/s | peak VRAM (GB) | median wall (s) |
|---------|--------------|----------------|------------------|
| HF | 80.5 | 57.2 | 6.36 |
| Flex | 55.6 | 116.0 | 9.21 |
Flex is correctness-complete but not yet performance-competitive on
MoE decode at bs=8. Two structural reasons:
- MoE decode runs eager (forward_moe_backend uses bincount + Python
expert loops which are not CUDA-graph capturable), so flex loses
its main dense-model advantage.
- FlexEngine deep-copies the HF model for the rollout copy, doubling
weight residency. For Qwen3-30B-A3B at bf16 that is ~60 GB extra.
The pristine-base third copy is skipped for Qwen3 MoE (see the
first commit of this series) but the inference deep-copy remains.
Follow-ups (not blockers for correctness):
- torch.compile(dynamic=True) on call_moe_model_with_flex_kwargs to
recover some of the CUDA-graph throughput without requiring graph
capture.
- Evaluate flex's scaling vs HF generate at bs=32 / bs=64, where
paged-KV reuse should dominate per-prompt cost.
- A quantised-only inference copy (4bit forward, fp32 LoRA injection)
so the flex path fits inside 2x 4bit weight residency (~34 GB)
instead of the current post-dequantisation footprint.
4f2bfe7f69