Skip to content

feat: Token-sharded sparse-MLA prefill for DeepSeek-V3 on SM120 - #64043

Closed
YuvrajSinghBhadoria2 wants to merge 2 commits into
vllm-project:mainfrom
YuvrajSinghBhadoria2:feat/token-sharded-sparse-mla-prefill
Closed

YuvrajSinghBhadoria2 wants to merge 2 commits into
vllm-project:mainfrom
YuvrajSinghBhadoria2:feat/token-sharded-sparse-mla-prefill

Conversation

@YuvrajSinghBhadoria2

Copy link
Copy Markdown
Contributor

Implements token-sharded sparse prefill for MLA attention used in DeepSeek-V3 on Blackwell (SM120).

Key features:

  • Token sharding distributes prefill tokens across SMs
  • Sparse attention using top-k indices from indexer
  • fp8_ds_mla KV cache layout required
  • SM120 (Blackwell) only, uses FlashInfer trtllm-gen sparse MLA kernel

AI usage disclosure: NO

…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

⚠️ The sha of the head commit of this PR conflicts with #63969. Mergify cannot evaluate rules on this PR. Once #63969 is merged or closed, Mergify will resume processing this PR. ⚠️

@YuvrajSinghBhadoria2

Copy link
Copy Markdown
Contributor Author

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

@mergify mergify Bot added deepseek Related to DeepSeek models rocm Related to AMD ROCm structured-output speculative-decoding labels Oct 28, 2026
@mergify mergify Bot added dflash mrv2 Model Runner V2 specific scheduler labels Oct 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

deepseek Related to DeepSeek models dflash mrv2 Model Runner V2 specific rocm Related to AMD ROCm scheduler speculative-decoding structured-output

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant