[Fangzhou-Ai]

Purpose

On ROCm, every MiniMax-M3.1 decode layer runs the NVFP4 sparse decode attention, then a separate _kernel launch (the MXFP8 activation quant of its bf16 output, K = 2048 per TP4 rank), then the o_proj GEMM. In the profile of a real decode step (rank-0 profile windows at N=4 and N=16) that quant is 60 launches per step, ~4.4-4.6 us each (profiled, inflated ~1.6x).

This PR adds an fp8_out variant of the FlyDSL sparse decode kernel (m3_nvfp4_sparse_decode_*_q8) whose final store writes E4M3 values and row-major E8M0 scales instead of bf16, and hands them to o_proj as a QuantizedActivation (same mechanism as #63637 for qkv_proj and the FFN fusion in #65478).

How it works

• Output layout: in the final store a lane owns head n, dims 32g..32g+31, i.e. exactly one 32-element MX block, so the quantization needs no cross-lane exchange. With split-K (splits > 1) the last workgroup of a row merges the partials and stores O; the quant runs there on the merged O, so every 32-block sees the final values.
• Arithmetic: it quantizes the bf16-rounded values the bf16 store would have written (pk_bf16), computes the block amax on their integer bit patterns and the E8M0 byte as clamp(exponent_field - 8 + (mantissa > 0.75), 0, 254) (448 = 1.75 x 2^8, so this is ceil(log2(amax/448)) + 127 exactly, no log2/division), scales by the exact power of two 2^(127-sb) and converts with v_cvt_pk_fp8_f32 (RNE). Hence the output is bit-identical to mxfp8_e4m3_quantize of the same kernel's bf16 output (not "more accurate": it is the same values, unlike the swiglu fusion in [ROCm][Perf][MiniMax-M3.1] Fuse the FFN-side MXFP8 activation quant into the post-attention norm and swiglu #65478).
• The scale formula was checked on CPU against _mxfp8_e4m3_quantize_torch on 20,000 random bf16 blocks over ~2^60 dynamic range plus edge blocks (zeros, subnormals, 447/448/449, the 1.75 boundary and the value just above it): 0 mismatches. The whole kernel compiles for gfx950 for splits 1/4/8 (FLYDSL compile-only on CPU). Bit-exactness on MI355X: see Tests.
• Gating: fp8_out is a compile-time variant; the bf16 kernel and all other callers are unchanged. The attention layer uses it only for decode-only batches (num_prefills == 0) in eager or full-CUDA-graph mode and when o_proj can take MXFP8 input (_can_fuse_input_norm_mxfp8); mixed batches, piecewise graphs and non-NVFP4 impls keep the bf16 path and the linear's own quant.

Files

ops/nvfp4_sparse_decode/kernel.py (epilogue), host.py (new optional out_q/out_scale; output may be None then), sparse_attention_nvfp4.py (forward_decode_mxfp8), model.py (MiniMaxM3SparseAttention NVFP4 branch + _run_attention keyword outputs), tests/kernels/attention/test_minimax_m3_nvfp4_sparse_decode.py. host.py is also touched by the open split-heuristic PR #65481 (different hunks: it patches cleanly with offsets).

Tests (MI355X, gfx950)

New in test_minimax_m3_nvfp4_sparse_decode.py: (1) N = 1,2,3,5,8,12,16,24 requests x 4/8 rows x 30K/150K/500K context at the default split choice, (2) every split count 1/2/4/8 x XCD on/off x 1/4 KV heads x 1/4/8 rows with padding requests and a wide magnitude range. Both assert bitwise equality of values and scales with mxfp8_e4m3_quantize of the bf16 kernel output and that rows past the batch are untouched. Result: the 96 new tests pass, and the whole file passes (176 passed), i.e. the existing decode tests are unchanged. The i8 scale store and the split-merge path were correct on the first run.

Measurements

Profile trace
Before:

After:

Graph-replay microbenchmark (a small benchmark script; 200 back-to-back launches per graph, median of 15 replays; 8 rows per request, 16 heads per rank, MI355X, 1 GPU), us per launch. "bf16 + quant" is the current path (decode kernel then mxfp8_e4m3_quantize). Context 150K:

N	bf16 kernel	bf16 + separate quant	MXFP8 epilogue kernel	saved	epilogue cost vs bf16
1	10.34	12.56	10.57	1.99	+0.24
2	10.74	12.78	10.93	1.85	+0.19
3	10.82	12.91	11.02	1.89	+0.20
4	10.86	12.99	11.14	1.85	+0.28
5	11.56	13.66	11.90	1.77	+0.34
6	10.07	12.15	10.29	1.87	+0.22
8	10.27	12.39	10.48	1.91	+0.21
12	13.67	15.91	14.13	1.78	+0.46
16	13.84	16.09	14.31	1.78	+0.47
24	18.39	20.72	18.64	2.09	+0.25

30K and 500K contexts give the same numbers within 0.1 us (saved 1.79-2.11 us at 30K, 1.81-2.04 us at 500K). The separate quant costs about 2.2 us in the graph; the epilogue adds 0.2-0.5 us to the attention kernel, so the net saving is 1.8-2.1 us per layer.

End to end (MI355X, TP4, DSpark 7 drafts, 150K context, DBO off, VLLM_ROCM_USE_SKINNY_GEMM=0)

Arms on the same installed tree, 5 repeats each: S0 = a tree that already contains the FFN quant fusions of #65478 and the routed-scale fold (so the baseline already has the other launch removals), S2 = S0 + this PR; S0 was also run again at the end to check drift (S0 start and end agree within 0.14 ms at every N).

Profile window (rank 0, 60 steps, all streams), N=4 and N=16: _kernel.kd (the MXFP8 activation quant) 82 -> 22 per decode step, total kernels per step 1399 -> 1339 (N=4) and 1456 -> 1396 (N=16): exactly the 60 launches this PR removes; nothing else changes.

Decode step p50 (ms), median of 5 repeats (noise of this benchmark about +-0.1-0.2 ms):

N	S0 start	S0 end	S2 (this PR)	vs mean(S0)
1	9.10	9.13	9.09	-0.3%
2	10.04	10.05	10.00	-0.4%
3	11.28	11.31	11.25	-0.4%
4	11.65	11.53	11.66	+0.6%
5	12.71	12.71	12.56	-1.2%
6	13.08	13.10	13.13	+0.3%
8	14.41	14.36	14.30	-0.6%
12	16.01	16.15	15.85	-1.4%
16	18.05	18.08	18.04	-0.1%

The step-level change is small: negative at 7 of 9 batch sizes, mean -0.4% (about 0.05 ms), average -0.8% at N >= 12, but it is not separable from the noise at any single N. I therefore claim the exact launch removal and the per-launch kernel saving above, not a specific end-to-end gain. AgentX replay and GSM8K were not run (the epilogue is bit-identical to quantizing the bf16 output in the unit tests).

Limitations

gfx950 FlyDSL decode only; decode-only batches; K = 2048 per rank assumes 16 query heads per KV head with TP4 (generalizes to any heads x 128 since each block is within one head).

AI assistance: this change was developed with Claude Code; the tests and benchmarks above were run on MI355X (gfx950).

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[Fangzhou-Ai]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99723 for commit a3ea9c6b3d5a.

[ChuanLi1101]

Looks correct on the fusion shape: quantizing after the split-K merge, one MX block per lane, and matching mxfp8_e4m3_quantize of the bf16-rounded store is the right contract for o_proj.

Two asks:

1. Please confirm _emits_o_proj_mxfp8 stays decode-only (prefill / mixed steps still take the bf16 → separate quant path), and that a graph-captured decode step cannot flip the cached _o_proj_mxfp8 decision mid-capture.
2. The bit-identical tests cover splits and heads — please also note (or test) the multi-XCD / multi-request path that #65481 retunes, so the fp8_out store is checked under the split counts we actually serve.

[Fangzhou-Ai]

Thanks for the review. On your two asks:

1. Decode-only and graph capture. _emits_o_proj_mxfp8 returns True only for a decode-only step: it needs num_prefills == 0 and num_decodes > 0 in the sparse metadata and a runtime mode other than PIECEWISE. Prefill and mixed steps keep the bf16 output and the separate quant path. The cached _o_proj_mxfp8 flag depends only on static module state (whether the impl has forward_decode_mxfp8, and whether o_proj's kernel/quant key accepts MXFP8 input), not on the batch or on capture state, so it cannot flip mid-capture. The per-step part is re-read from the forward context each forward. FULL graphs are only keyed on uniform-decode batch descriptors in the cudagraph dispatcher, so a graph captured with the MXFP8 hand-off is only replayed for decode-only steps; mixed batches run PIECEWISE or eager, which take the bf16 path. In the decode profile (graph replay) the fused path is active. I have not added a unit test for the dispatch decision itself; happy to add one if you'd like.

2. Multi-request / multi-XCD splits. Added test_decode_mxfp8_output_multi_request_xcd_splits (2d8d5e5): DSpark 7 rows (8 per request) on the 8-XCD grid, explicit splits of 2 and 4, with 1, 5, 8, 9, 16, 17 and 24 requests, checking the fp8 output and scales bit for bit against mxfp8_e4m3_quantize of the same kernel's bf16 output. The existing multi-request test (1-24 requests, 30K-500K context) runs at the default split choice, so it also covers whatever the heuristic picks once #65481 lands. #65481 is not in this branch, which is why the new test passes the split counts explicitly. All 110 MXFP8 epilogue tests pass on gfx950.

[ChuanLi1101]

Thanks, this addresses both points. Please add the focused dispatch test as well: decode-only should use the MXFP8 hand-off, while mixed and PIECEWISE batches keep the BF16 path.