Skip to content

[ROCm][Perf] Fuse MXFP8 activation quant into the 1x32 dot_scaled linear for M≤16 - #63969

Closed
YuvrajSinghBhadoria2 wants to merge 2 commits into
vllm-project:mainfrom
YuvrajSinghBhadoria2:fused-mxfp8-act-quant
Closed

YuvrajSinghBhadoria2 wants to merge 2 commits into
vllm-project:mainfrom
YuvrajSinghBhadoria2:fused-mxfp8-act-quant

Conversation

@YuvrajSinghBhadoria2

Copy link
Copy Markdown
Contributor

This PR fuses MXFP8 activation quantization into the ROCm dot_scaled GEMM kernel for M≤16 on gfx950 (MI355X).

Key changes:

  • Eliminates separate mxfp8_e4m3_quantize kernel launch by quantizing bf16 activations in-register during GEMM
  • Uses QUANT_ROWS = min(16, next_pow2(M)) rows per K-tile, repeated to fill 16-row MFMA tile
  • K loop uses constexpr BLOCK_K for loop unrolling/pipelining
  • Fused for K≤1024 (M≤16) and K=2048 (M≤4); larger shapes use existing path

Performance (MI355X):

  • 1.14x–1.64x speedup for K≤1024 at M=1–16
  • 1.03x–1.27x for K=2048 at M=1–4

Constraints:

  • bf16 only (fp16 triggers FP16 overflow clamp)
  • Disabled under VLLM_BATCH_INVARIANT
  • Bitwise identical output validated via new test test_rocm_mxfp8_fused_act_quant_bitwise_eq_unfused

The optimization targets small-batch decode where separate quant launch overhead dominates.

Closes #63915

…s available

- Remove draft_grammar_bitmask from SchedulerOutput
- Add compute_draft_grammar_bitmask() to SchedulerInterface
- EngineCore computes bitmask after take_draft_token_ids()
- Pass _next_draft_grammar_bitmask to next execute_model()
- Works with both sync and async scheduling

Signed-off-by: Yuvraj Singh Bhadoria <yuvrajsinghbhado2030@gmail.com>
…≤16 on ROCm

This adds a fused kernel that quantizes bf16 activations in-register during
the GEMM, eliminating the separate mxfp8_e4m3_quantize kernel launch for
small M (M≤16) on gfx950 (MI355X).

Key changes:
- Added _mxfp8_fused_quant_linear_kernel: Fused quantization + dot_scaled kernel
- Added _mxfp8_fused_quant_dot_scaled_linear: Host function to launch fused kernel
- Modified RocmDotScaledMxfp8LinearKernel.apply_weights to use fused kernel
  for M≤16, K≤1024 (and M≤4, K=2048)
- Disabled under VLLM_BATCH_INVARIANT
- Added test_rocm_mxfp8_fused_act_quant_bitwise_eq_unfused for bitwise validation
- Added test for VLLM_BATCH_INVARIANT disable behavior

Performance targets (from issue):
- 1.14x–1.64x speedup for K≤1024 at M=1–16
- 1.03x–1.27x for K=2048 at M=1–4
- bf16 only (fp16 triggers FP16 overflow clamp)

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@mergify

mergify Bot commented Oct 27, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @YuvrajSinghBhadoria2.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@YuvrajSinghBhadoria2

Copy link
Copy Markdown
Contributor Author

Consolidating focus to priority PRs (#63479, #63798, #63919). Will reopen when ready.

@github-project-automation github-project-automation Bot moved this from Backlog to Done in Sprint - DFlash Oct 28, 2026
@github-project-automation github-project-automation Bot moved this from Todo to Done in AMD Oct 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

Status: Done
Status: Done
Status: Done

Development

Successfully merging this pull request may close these issues.

1 participant