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 14, 05:30 UTC
[Spec Decode] Enable efficient speculative decoding with FlashInfer-MLA (#25984)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
Oct 7, 2025 at 16:05 UTC
3d1f67616da88cbf0033bf5027cc0c6e5e9cacf6
3 files changed
+42
-5
vllm/v1/attention/backends/mla/common.py
+24
-2
@@ -190,7 +190,7 @@ return curr_o @ W_O
190
import functools
191
from abc import abstractmethod
192
from dataclasses import dataclass, field
193
-from typing import Generic, Optional, TypeVar, Union
193
+from typing import ClassVar, Generic, Optional, TypeVar, Union
194
195
import torch
196
from tqdm import tqdm
@@ -454,6 +454,20 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
454
understand this class
455
"""
456
457
+ # Whether the backend supports reordering the batch such that
458
+ # short sequences (i.e. verification for speculative decoding) are
459
+ # classified as decode requests.
460
+ # If True, this will increase `reorder_batch_threshold` (below) when
461
+ # speculative decoding is enabled, and set `require_uniform=True` when
462
+ # when reordering the batch. Non-uniform decode requests will
463
+ # fall back to prefill in this case.
464
+ supports_uniform_spec_as_decode: ClassVar[bool] = False
465
+
466
+ # The threshold for reordering the batch into decode and prefill requests.
467
+ # If > 1, the batch will be reordered such that requests with
468
+ # query length <= threshold are classified as decode requests.
469
+ # Use `supports_uniform_spec_as_decode` (above) to set this automatically
470
+ # when speculative decoding is enabled.
471
reorder_batch_threshold: int = 1
472
473
@staticmethod
@@ -503,6 +517,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
517
self.model_config = vllm_config.model_config
518
parallel_config = vllm_config.parallel_config
519
self.compilation_config = vllm_config.compilation_config
520
+ self.vllm_config = vllm_config
521
self.device = device
522
523
self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
@@ -578,6 +593,11 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
593
device=device,
594
)
595
596
+ supports_spec_as_decode = self.supports_uniform_spec_as_decode
597
+ self._init_reorder_batch_threshold(
598
+ self.reorder_batch_threshold, supports_spec_as_decode
599
+ )
600
+
601
def _build_fi_prefill_wrappers(self, prefill: FlashInferPrefillMetadata):
602
qo_indptr = prefill.query_start_loc
603
@@ -714,7 +734,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
734
735
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
736
split_decodes_and_prefills(
717
- common_attn_metadata, decode_threshold=self.reorder_batch_threshold
737
+ common_attn_metadata,
738
+ decode_threshold=self.reorder_batch_threshold,
739
+ require_uniform=self.supports_uniform_spec_as_decode,
740
)
741
)
742
vllm/v1/attention/backends/mla/flashinfer_mla.py
+15
-1
@@ -22,6 +22,9 @@ FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
22
23
24
class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
25
+ # enable spec-as-decode optimization
26
+ supports_uniform_spec_as_decode: ClassVar[bool] = True
27
+
28
# enable full CUDA Graph support for decode-only capture
29
cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
30
@@ -111,7 +114,15 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
114
q = torch.cat([q_nope, q_pe], dim=-1)
115
116
# trtllm API requires extra dimension q_len_per_request for MTP
114
- q = q.unsqueeze(1)
117
+ if attn_metadata.num_decode_tokens % attn_metadata.num_decodes != 0:
118
+ logger.warning_once(
119
+ """FlashInferMLAImpl got a query of uneven length.
120
+ This usually indicates an issue in batch reordering
121
+ or incorrect setup in dummy_run."""
122
+ )
123
+ q = q.unsqueeze(1)
124
+ else:
125
+ q = q.view(attn_metadata.num_decodes, -1, q.shape[-2], q.shape[-1])
126
127
if self.bmm1_scale is None:
128
self.bmm1_scale = layer._q_scale_float * layer._k_scale_float * self.scale
@@ -132,6 +143,9 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
143
bmm2_scale=self.bmm2_scale,
144
)
145
146
+ # Flatten the output for consistent shape
147
+ o = o.view(-1, o.shape[-2], o.shape[-1])
148
+
149
# TODO: Return LSE pending support from Flashinfer API:
150
# https://github.com/flashinfer-ai/flashinfer/pull/1566
151
return o, None
vllm/v1/attention/backends/utils.py
+3
-2
@@ -275,8 +275,9 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):
275
speculative_config is not None
276
and speculative_config.num_speculative_tokens is not None
277
):
278
- self.reorder_batch_threshold = (
279
- 1 + speculative_config.num_speculative_tokens
278
+ self.reorder_batch_threshold = max(
279
+ self.reorder_batch_threshold,
280
+ 1 + speculative_config.num_speculative_tokens,
281
)
282
283
@abstractmethod