[faresobeid]

Sparse weight deltas for engines that quantize online to MXFP8 (vLLM quantization = "mxfp8"), Blackwell's native format. Stacked on #4085; until now such engines were refused (no _scale_inv, unknown layouts).

What changes

• Trainer. Engines report their tensors as "mxfp8". The trainer quantizes old and new weights with vLLM's own mxfp8_e4m3_quantize (FlashInfer on SM100+, torch elsewhere: the same bytes for bf16 inputs) and sends the fp8 byte diffs plus the changed e8m0 scales as {name}_mx_scale ([rows, cols / 32], uint8). 1x32 blocks never straddle row shards, so nothing is gathered. The fp8 encoder's block geometry is now per mode (fp8_block): 128x128 for fp32/ue8m0, 1x32 for mxfp8.
• Engine. fp8_layout finds each MXFP8 layer's uint8 weight_scale and tries candidate kernel layouts for every fp8 param, keeping the first whose recorded positions hold vLLM's MXFP8 quantization of the checkpoint byte for byte (weights and scales):
  • row-major scales;
  • FlashInfer's F8_128x4 scale swizzle (swizzle_mxfp8_scale; FlashInfer Cutlass/CuteDSL/torch linear kernels), as row/column offset tables;
  • for expert params, the FlashInfer TRT-LLM MoE layouts (w13 swapped to [up; gate], gate/up-interleaved and row-shuffled; w2 row-shuffled; scales through the same rows plus FlashInfer's 128x4 interleave). Their tables come from running vLLM's own swap_w13_to_w31 and _shuffle_mxfp8_moe_weights (FlashInfer's block_scale_interleave included) on byte-planes of index probes, so they follow the engine's code exactly.
    An unknown layout raises instead of patching the wrong bytes. load_rows quantizes anchors and full checkpoints the same way.

Validation

• RTX PRO 6000 (SM120), vLLM 0.33: Qwen3-0.6B reverse-text RL with an MXFP8 engine (FlashInferCutlassMxfp8LinearKernel, 196 tensors held in MXFP8 with swizzled scales), sparse deltas, anchor every 3 steps: 8/8 deltas applied and verified against the trainer's checksums, 0 mismatches; reward 0.24 -> 0.36-0.41.
• GPU unit test: the probed TRT-LLM tables reproduce vLLM's real swap_w13_to_w31 + _shuffle_mxfp8_moe_weights output for w13 and w2.
• CPU: the F8_128x4 offset tables match swizzle_mxfp8_scale (padded shapes included); _probe_permutation / _separable recover a known permuted, padded layout and reject non-permutations; 556 passed across the sparse test set; ruff clean.

Not covered

• The production SM100 MoE path itself (TRT-LLM MXFP8 MoE runs on SM100 only; this host has SM120, where vLLM picks Marlin for MXFP8 MoE): the tables are checked against vLLM's transform on GPU, and on a B200/B300 the engine's byte-for-byte check picks or refuses the layout at startup. Needs one B200/B300 run, e.g. of configs/advanced/deepseek-v4-flash/swe_mxfp8.toml.
• Marlin (SM120 MXFP8 MoE) and Humming (H100) layouts still raise; Marlin can use the same index-probe approach (gptq_marlin_repack is a byte permutation). NVFP4 (online NVFP4 is SM100-only; its per-expert global scale makes an expert dense when its amax moves) is next.

🤖 Generated with Claude Code