moe(b12x): zero-pad w4a8_mx expert shards to the QMMA 128-column layout - #80
voipmonitor wants to merge 6 commits into
Conversation
…d-MXFP8 dense support - store_dtype=nvfp4 now returns ModelOptNvFp4FusedMoE (registers w13_weight/w2_weight names the DeepSeek/GLM fused-experts loader expects; the earlier CompressedTensors method registers *_weight_packed and KeyErrors at load). - dense_format=mxfp8: FP8 checkpoints whose dense weights are serialized as MXFP8 (e4m3 values + per-32 ue8m0 uint8 scales, e.g. requantized offline from BF16 via mxfp8_e4m3_quantize) route LinearBase to a thin adapter over the compressed-tensors W8A8 MXFP8 scheme. ignored_layers keeps enumerated BF16 linears (routers, indexer heads, lm_head) on the unquantized path. Enables GLM-5.2-FP8-NVFP4experts and GLM-5.2-MXFP8dense-NVFP4experts. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> (cherry picked from commit 4b55617)
Align the MXFP4 online-overlay hook with the ModelOpt overlay semantics:
the `linear` spec covers dense linears only, and shared-expert
projections are quantized only when a `shared_experts` spec is given
explicitly. Before this, the same
--quantization-config '{"linear":{"weight":"mxfp8"}}' meant two
different things: shared experts stayed BF16 on ModelOpt/NVFP4
checkpoints but were silently quantized on MXFP4 checkpoints (the
OnlineQuantizationConfig fallback picks the linear spec for every
LinearBase). Quantized shared experts measured both lower quality
(mean|dlogprob| 0.156 vs 0.152) and lower decode throughput
(90.1 vs 92.5 tok/s) on GLM-5.2, so the divergence was strictly harmful.
Adds selection tests for the MXFP4 overlay path: shared experts skipped
without the spec, quantized with it, and `ignore` patterns honored
(e.g. re:.*kv_b_proj).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…kens The B12X PCIe DCP A2A path is a one-barrier, latency-optimized exchange: it wins small decode batches (+13% CC1 decode on GLM-5.2 DCP4) but loses ~21-27% standalone prefill throughput to AG+RS, because for sparse MLA every batch (prefill included) goes through forward_mqa and the one-shot staging copy is bandwidth-poor at large token counts. Add VLLM_DCP_A2A_MAX_TOKENS (0 = uncapped, current behavior): batches with more tokens than the cap skip the B12X path and take a pipelined NCCL collective instead — VLLM_DCP_A2A_LARGE_BACKEND selects AG+RS (default, the measured prefill winner) or the packed NCCL all-to-all. The cap also bounds the B12X IPC staging pool, which scales linearly with its token capacity (max_num_batched_tokens=8192 previously reserved ~0.9 GB/rank; a cap of 64-128 needs ~7-14 MB). Dispatch happens on the per-step token count, mirroring the deterministic size crossovers already used for the PCIe allreduce (oneshot/PyNCCL/DMA). CUDA-graph safe: capture bakes one path per padded graph size; eager prefill re-evaluates per step; all DCP ranks see the same batch so the choice is uniform across the group. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…d size The B12X PCIe DCP channel exists only for world sizes 2/4/8, and the runtime dispatchers already fall back to NCCL collectives per call. The dedicated DCP warmup added in 'Optimize B12X DCP collectives and warmup' turned that graceful fallback into a boot failure: with TP6 + DCP3/DCP6 (GLM head66 configs that worked on the v13 stack) warmup_b12x_dcp_a2a raised 'B12X PCIe DCP query all-gather is unavailable for the configured attention geometry' and EngineCore died. - warmup_b12x_dcp_a2a: log once and skip when world size is not 2/4/8 (still raises for genuinely broken geometries at supported sizes). - MLAAttention: don't set dcp_b12x for unsupported DCP sizes, so neither warmup nor the per-step dispatch attempts the PCIe channel at all. Verified on GLM-5.2 NVFP4 TP6/DCP6/MTP3 (v13 wiki shape: GMU 0.957, max_model_len 128000, max_num_batched_tokens 2048): boots clean, KV cache 989k tokens, test.py c0/c3000 coherent (CJK 0), 73-75 tok/s. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. 🗂️ Base branches to auto review (1)
Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
…layout Both B12X MoE kernel families reject expert shards whose intermediate size is not tile-aligned, which makes GLM-5.2 (moe_intermediate 2048) unbootable at TP6 (2048/6 -> virtual-TP padded 352 per rank) on MXFP4 checkpoints, for any DCP value: - w4a8_mx dies at weight prep: 'W4A8-MX QMMA layout requires hidden_size % 256 == 0 and intermediate_size % 128 == 0' - w4a16 (packed e8m0 path) dies at profile run: 'no valid W4A16 tile config for M/N/K=16384/6144/352, moe_block_size=64' Zero-pad the shard to the next 128 multiple at the vLLM/B12X weight handoff, before the preparation plan is built: gate/up halves and their e8m0 scale rows are extended per half, w2 columns (packed FP4 nibbles) plus its scale columns are extended to match. Padded gate/up rows produce silu(0) * 0 = 0 activations and the padded w2 columns only multiply those zeros, so MoE outputs are bit-exact; the plan, the repack, and the runtime problem size all derive dimensions from the padded tensors. Cost is proportional to the padding (352 -> 384 is ~9% extra expert GEMM work); source parameters are released after the repack as before. Guarded narrowly: e8m0 source format only, exact logical-shape match required, otherwise the original errors remain. Verified on GLM-5.2-BF16-AMDMXFP4experts, TP6/DCP1, online MXFP8 dense overlay, v13 memory shape (GMU 0.957, mml 128000, mnbt 2048), coding smoke c0/c3000 coherent (0 CJK) in both modes: | mode | KV cache | decode c0 | decode c3000 | prefill 8k | prefill 64k | |---|---|---|---|---|---| | force-A8 (w4a8_mx) | 335,168 | 72.2 | 71.2 | 4,770 | 5,225 | | A16 (w4a16) | 338,304 | 79.2 | 77.9 | 4,643 | 4,816 | (Previously both modes failed before ready.) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
c28cb30 to
49bed02
Compare
|
Amended ( Measured on GLM-5.2-BF16-AMDMXFP4experts TP6/DCP1 (both modes previously failed before ready):
Interesting inversion vs TP8: at TP6 A16 decodes ~9% faster while A8 prefills up to ~8% faster (64k) — the QMMA padding overhead lands on the decode side, prefill still benefits from the MXFP8×FP4 MMA. All outputs coherent (0 CJK). |
|
Follow-up on the TP6 A8-vs-A16 decode gap (72 vs 79 tok/s): the A8 decode parity kernel ( Options considered:
Until the kernel grows N=384 support, the practical TP6 guidance stands: A16 for decode-heavy, A8 for prefill-heavy (+8% @64k prefill). |
|
Final resolution: the tiny_decode kernel now supports odd FC2 K-tile counts — lukealonso/b12x#26 (FC2 K-tiles-per-task becomes a configure()-time value, 2→1 for odd counts; support gate relaxed to
The earlier per-mode guidance (A16 for decode) is obsolete once both land — A8 wins both axes at TP6. |
|
This one is no good - the whole point is to avoid as much of this padding as possible, and 128 is way too much. The other modes (A4, A16) support much smaller alignments to support this. |
2 participants
Problem
The W4A8-MX QMMA repack tiles the expert intermediate dimension in 128-column blocks, so TP shards that are not multiples of 128 fail weight preparation:
GLM-5.2 has
moe_intermediate_size=2048, so TP6 (2048/6 → virtual-TP padded to 352/rank, a multiple of 32 but not 128) could not boot withB12X_MOE_FORCE_A8at all — for any DCP value (the failure is TP6 itself, it dies at weight prep before KV sizing). TP8 (256/rank) and TP4 (512/rank) are unaffected.Fix
Zero-pad the shard to the next 128 multiple at the vLLM/B12X weight handoff in
b12x_moe.py, before the preparation plan is built:w13and their e8m0 scale rows are extended per half,w2columns (packed FP4 nibbles) and its scale columns are extended to match,moe_problem_sizeall derive dimensions from the padded tensors, so the whole pipeline stays consistent (prepared.intermediate_size= padded).Exactness: padded gate/up rows produce
silu(0) * 0 = 0activations, and the paddedw2columns only ever multiply those zeros — MoE outputs are bit-identical to an unpadded ideal. Cost is proportional to the padding: 352 → 384 is ~9% extra expert GEMM work on the padded rank, logged once as a warning. Source parameters are released after the repack as before, so the padded copy is the only surviving allocation (no double-resident weights).Guarded narrowly: only
quant_mode == "w4a8_mx", only whenintermediate % 128 != 0and% 32 == 0, and only when all four tensor shapes match the expected logical grid — anything else falls through to the existing error.Verification
GLM-5.2-BF16-AMDMXFP4experts (AMD MXFP4 experts, BF16 dense), TP6/DCP1,
MOE_MODE=force-a8-experimental+ online MXFP8 dense overlay, v13 memory shape (GMU 0.957, mml 128000, mnbt 2048), imagev4-vllmbbce67f-b12xe44cb77+ this file mounted:Boot log:
B12X w4a8_mx: expert intermediate size 352 is not a multiple of 128; zero-padding the shard to 384 for the QMMA layout (exact results, ~9% extra expert GEMM work).An alternative would be teaching the QMMA repack/kernels a 32-aligned tail — that is real kernel work for a niche shard shape; this padding unblocks the config with bounded, transparent cost and can be deleted if the kernel ever grows native tail support.
🤖 Generated with Claude Code