A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 05:00 UTC
[Bugfix][MTP][Sparse MLA] Allow sparse MLA with MTP to run with FULL cudagraphs (#34457)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
Matthew Bonanni committed
Feb 17, 2026 at 14:01 UTC
dc5fa77a4eb6680339cb77abe713fb22d7795560
2 files changed
+17
-6
docs/design/cuda_graphs.md
+1
@@ -182,6 +182,7 @@ The following table lists backends that support full CUDA Graphs at the time of
182
| FlashInfer | `UNIFORM_SINGLE_TOKEN_DECODE` | Will be set to `UNIFORM_BATCH` when using TRTLLM attention on Blackwell |
183
| FlashMLA | `UNIFORM_BATCH` | |
184
| FlashInferMLA | `UNIFORM_BATCH` | |
185
+| FlashInferMLASparse | `UNIFORM_BATCH` | |
186
| AITER MLA | `UNIFORM_SINGLE_TOKEN_DECODE` | |
187
| CUTLASS MLA | `UNIFORM_SINGLE_TOKEN_DECODE` | |
188
| Mamba attention| `UNIFORM_SINGLE_TOKEN_DECODE` | |
vllm/v1/attention/backends/mla/indexer.py
+16
-6
@@ -196,9 +196,7 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig):
196
197
198
class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
199
- _cudagraph_support: ClassVar[AttentionCGSupport] = (
200
- AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
201
- )
199
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
200
201
reorder_batch_threshold: int = 1
202
@@ -212,8 +210,14 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
210
if self.vllm_config.speculative_config
211
else 0
212
)
215
- # Now deepgemm fp8_paged_mqa_logits does not support next_n > 2
216
- self.reorder_batch_threshold += min(self.num_speculative_tokens, 1)
213
+ if self.num_speculative_tokens > 1:
214
+ raise ValueError(
215
+ "Sparse MLA only supports "
216
+ "num_speculative_tokens <= 1 because the DeepGEMM "
217
+ "fp8_paged_mqa_logits kernel does not support next_n > 2. "
218
+ f"Got num_speculative_tokens={self.num_speculative_tokens}."
219
+ )
220
+ self.reorder_batch_threshold += self.num_speculative_tokens
221
222
props = torch.cuda.get_device_properties(self.device)
223
sm_count = props.multi_processor_count
@@ -342,8 +346,14 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
346
self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
347
seq_lens, self.kv_cache_spec.block_size, self.num_sms
348
)
349
+ block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...]
350
+ # Padded CUDA graph requests have block_table entries of -1.
351
+ # Clamp to 0 to prevent OOB access in the DeepGEMM kernel.
352
+ # This is safe because padded requests have seq_lens=0, so the
353
+ # kernel produces no meaningful output for those rows.
354
+ block_table.clamp_(min=0)
355
decode_metadata = DeepSeekV32IndexerDecodeMetadata(
346
- block_table=common_attn_metadata.block_table_tensor[:num_decodes, ...],
356
+ block_table=block_table,
357
seq_lens=common_attn_metadata.seq_lens[:num_decodes],
358
decode_lens=decode_lens,
359
requires_padding=requires_padding,