feat(refit): support grouped MoE MXFP8 weights

[seonjinn]

Summary

Add MXFP8 rollout refit support for Qwen3.5 grouped MoE weights.

Split fused gate/up and down expert slabs into per-expert projections, quantize them to E4M3 values with E8M0 scales, and load them through vLLM's existing expert mapping. Scale names follow the legacy or native reload API as appropriate.

Changes

• Support grouped MoE MXFP8 refit while preserving BF16 passthrough for excluded experts.
• Handle vLLM kernel factories with or without the layer argument.
• Add focused conversion tests and a Qwen3.5 recipe; its end-to-end test remains disabled pending validation.

This PR contains only the original grouped-MoE feature. Trainer-side prequantization, cached loader routes, and IPC optimizations from #3294 are excluded. The prerequisite grouped-FP8 support from #2744 is already on main.

Validation

• After removing perf: Reduce MXFP8 refit weight-transfer overheads #3294, the code tree matches the original standalone head c22a14d29.
• CPU recipe checks: 3 passed, 1 skipped (vLLM unavailable).
• Ruff lint/format, shell syntax, and diff checks passed.
• The broader FP8 test suite could not be validated locally: 103 skipped and one existing test failed to import vLLM. GPU refit and end-to-end validation of this standalone revision remain pending.

Closes #3694.

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[seonjinn]

/ok to test 2c36da0

[yuki-97]

Thanks for supporting this. The PR lifts the grouped-MoE block on MXFP8 refit by expanding the fused expert slab per-expert into the gate_proj / up_proj / down_proj params vLLM actually registers, which lets the pre-refit rejection guard go away, and ships a Qwen3.5-35BA3B EP16/TP2 recipe on top of it.

Scope: the grouped-MoE MXFP8 expansion in quantization/fp8.py, the removed pre-refit guard, and the new recipe / driver / doc surface.

Checked:

• The expansion is anchored to what vLLM does rather than to what the layout looks like — the gate/up split at shape[1] // 2 matches vLLM's own chunk(2, dim=1) de-fusion with index 0 as w1/gate_proj, and both scale suffixes land on registered params through RoutedExperts' prefix rewrite.
• Keeping the new branch inside _is_fp8_grouped_moe_expert is the other half of that: experts excluded by num_{first,last}_layers_in_bf16 still pass through as the untouched bf16 slab, so the rewrite is scoped to the experts vLLM actually built as FP8.
• Every weight-transfer path the deleted assert_refit_unsupported_grouped_moe_params used to gate. Colocated IPC, collective legacy, collective native-reload and non-sharded checkpoint-engine all reach the new expander; sharded NIXL still rejects FP8 models outright; and NCCL reshard was never covered by that guard in the first place, since it implies non-colocated and therefore skips prepare_refit_info entirely.
• The vLLM calls against the pinned source at 0.25.1 and 0.26.0, and the kernel factory at every release through 0.30.0 — signature, kwarg names, return shape.

Remaining comments are on end-to-end evidence for the new recipe and the overlap with #4080.

nit, not blocking merging — the doc puts the grouped-MoE capability under ### Limitations, and the slab layout contract is now spelled out in two places.

[yuki-97]

nit: bullet 2 is a capability statement sitting under a ### Limitations heading, and next to bullet 1 — which scopes non-validation to Dense — it reads as though the grouped-MoE path is already validated.

Suggest fix: move it out of Limitations into the FP8 section body and write it as plain support, once the end-to-end run asked for in the disabled.txt comment has landed.

[yuki-97]

nit: lines 866-874 are byte-identical to the split in _expand_grouped_moe_expert_to_fp8 at :839-847, so the grouped-slab layout contract — gate first, split at shape[1] // 2 — is now asserted in two places with nothing tying them together.

Suggest fix: extract _split_grouped_moe_shards(key, weight) returning (base, shards) and call it from both; the per-expert loop, the quantizer choice, the .contiguous() decision and the scale suffix all stay at the call sites.

[yuki-97]

Nothing in the PR shows this path working end to end — this entry keeps the only e2e driver manual until its bounds exist, and the PR body lists GPU refit and end-to-end validation as pending.

Suggest fix:

1. Run the 20-step validation and set the bounds at ...mxfp8-trtllm.sh:36-37 from that run; today they are inherited verbatim from the 2n8g block-FP8 recipe.
2. Paste the run into the PR description — reward, loss and gen_kl_error — so the row added to the Example Recipes table is backed by evidence.
3. Move the driver out of disabled.txt into the recurring nightly suite once 1 and 2 are done.

[yuki-97]

#4080 touches both of this helper's call sites and resolves the same problem by dropping layer= unconditionally, with the pin at 0.29. Once this branch rebases onto that, the shim reproduces #4080's plain call through inspect.signature, and its if accepts_layer: arm (:1188) can no longer be taken on any pinned version.

Suggest fix: at rebase time, take #4080's call at :1255 and :1504, and drop this helper, the import inspect at :15 (its only other use), and test_make_fp8_moe_kernel_compat_matches_vllm_signature.