Repository navigation
[ROCm][Perf] Fuse MXFP8 activation quant into the 1x32 dot_scaled linear for M <= 16 - #63915
Conversation
…ear for M <= 16 On gfx950, RocmDotScaledMxfp8LinearKernel quantized every activation with a separate mxfp8_e4m3_quantize launch before the tl.dot_scaled GEMM. For bf16 activations with 1-16 rows, the GEMM now loads the bf16 tile and quantizes it in registers with the same fp32 operations, so each linear is one launch and the output is bit-identical. Every N-tile requantizes its activation, so the in-kernel work grows with K while the saved launch does not. With BLOCK_K >= 256 the kernel quantizes only min(16, next_pow2(M)) rows and repeats them to fill the 16-row MFMA tile (rows >= M are never stored); at BLOCK_K=128 repeating the rows costs more than it saves. The fused path is taken for K <= 1024 at M <= 16 (1.14-1.64x on MI355X) and for K=2048 at M <= 4 (1.03-1.27x); larger M or K keep the separate quant. K becomes a constexpr, which the fused loop needs to be profitable. fp16 is not fused, because the in-kernel fp8 cast enables the FP16 overflow clamp for the whole kernel. VLLM_BATCH_INVARIANT keeps the unfused path. Adds test_rocm_mxfp8_fused_act_quant_bitwise_eq_unfused: the fused path is taken and its output is bitwise equal to mxfp8_e4m3_quantize plus the _select_cfg tile, for every quantized-row count, with and without padding rows, including K=2048 at M=1-4. Signed-off-by: RuibinCheung <ruibzhan@amd.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
This pull request has merge conflicts that must be resolved before it can be |
shen-shanshan
added
verified
|
✅ @RuibinCheung, CI is now available for this PR.
|
|
/ci run |
|
❌ This PR is 263 commits behind upstream |
Signed-off-by: RuibinCheung <ruibzhan@amd.com>
|
/ci run |
|
✅ Triggered Buildkite CI #97942 for commit ℹ️ This PR is 1 commit behind upstream |
benenzhu
left a comment
benenzhu
left a comment
There was a problem hiding this comment.
LGTM! Improvements under low concurrency GEMMs.
Overview
On gfx950,
RocmDotScaledMxfp8LinearKernel(1x32-scale MXFP8 dense layers: MiniMax-M3/M3.1, ModelOpt/INC MXFP8) runs a separatemxfp8_e4m3_quantizelaunch before everytl.dot_scaledGEMM. For bf16 activations with 1-16 rows, this PR quantizes the activation inside the GEMM instead, so each such linear is one launch and the output is bit-identical. Fused: K <= 1024 for M <= 16, and K = 2048 for M <= 4.The fused GEMM quantizes only
min(16, next_pow2(M))activation rows per K step and repeats them to fill the 16-row MFMA tile; that keeps the per-N-tile quant cost proportional to the real rows and is what makes K = 2048 profitable.Claims
_select_cfgtile), including where the fused tile has a differentBLOCK_K/BLOCK_N.VLLM_BATCH_INVARIANT, for fp16 activations, for M > 16, for K = 2048 at M > 4, or for K > 2048.Validation
Microbenchmark: "separate" is
mxfp8_e4m3_quantize+ GEMM (current behaviour); "fused" is this PR. Speedup over separate, MI355X:* not fused at this shape; the column times the fallback, i.e. the noise floor.
End to end: MiniMax-M3.1-Flash-Quark-MXFP4, TP4, MI355X, mono decode off, FULL_DECODE_ONLY CUDA graphs, on #63232 (
8d0f05a0c, imagevllm/vllm-openai-rocm:minimax-m31-flash) with this PR applied on top; baseline is the same image unchanged. AgentX trace replay (InferenceX), concurrency 1, 600 s per run,max-num-seqs2; run 1 in the order baseline, this PR, run 2 reversed. This PR / baseline (> 1 is better; TPOT as baseline / this PR; P90 interactivity is InferenceX'sintvty.p90= 1 / P90 TPOT):GSM8K (lm-eval, 5-shot, full test set, greedy), measured with K <= 1024 fused: baseline 96.51 / this PR 96.21 strict-match (stderr 0.51 / 0.53), within one stderr; vLLM is not batch invariant when requests are batched together, so greedy outputs vary between runs even with bit-identical kernels. The fused output is bitwise equal to the existing path for every tested shape, including K = 2048.
Details
BLOCK_M= 16 is the smallest MFMA tile, but at M < 16 only the first M rows are stored. The kernel takes a constexprQUANT_ROWS = min(16, next_pow2(M)), loads and quantizes that many rows, and repeats them withbroadcast_to+reshape. AtBLOCK_K= 128 the repeat costs more than the saved quant, so those tiles (K = 384, 640, ...) keep quantizing 16 rows. The unfused path passesQUANT_ROWS = BLOCK_Mand generates the same code as before.Kbecomes atl.constexpr. Without it the fused loop is not unrolled/pipelined and the fused GEMM is ~0.8x even at K = 512. [Perf][FlashAttention] FA4: pass no FP8 descale for a unit per-tensor scale #63232 makes the same change.== 1and% 16 == 0;QUANT_ROWSadds up to three more variants per fused shape (6 instead of 3 for M = 1-16 atBLOCK_K>= 256; 5 at K = 2048, where M = 5-16 stays unfused). For MiniMax-M3.1 at TP4 that is 17 instead of 12_mxfp8_linear_kernelcompiles over M = 1-16 (15 instead of 12 at TP8). They are compiled in the eager warm-up before CUDA graph capture; vLLM's capture sizes include every M from 1 to 16, so decode never compiles at runtime, and graph capture time was unchanged in the end-to-end runs.VLLM_BATCH_INVARIANT([Feature] Support batch invariance on ROCm #52231) every M stays on the unfused path.dot_scaledpath). Shapes are tuned on MI355X.Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.