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 13, 09:30 UTC
[Perf] Refactor cudagraph_support to enable full CUDA graphs for spec decoding with FlashInfer (#28479)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
Nov 12, 2025 at 11:56 UTC
304419576ae9dc2ecaa28c4506d3870f7c68bd85
18 files changed
+71
-41
docs/design/cuda_graphs.md
+2
-1
@@ -177,8 +177,9 @@ The following table lists backends that support full CUDA Graphs at the time of
177
| FlashAttention v3 | `ALWAYS` | has unified routine for both batches, so `FULL` mode is good |
178
| Triton Attention | `ALWAYS` | prefer `FULL_AND_PIECEWISE` since it has different kernels for prefill/mixed and pure decode batches |
179
| AITER FlashAttention | `UNIFORM_BATCH`| |
180
-| FlashInfer | `UNIFORM_SINGLE_TOKEN_DECODE` | |
180
+| FlashInfer | `UNIFORM_SINGLE_TOKEN_DECODE` | Will be set to `UNIFORM_BATCH` when using TRTLLM attention on Blackwell |
181
| FlashMLA | `UNIFORM_BATCH` | |
182
+| FlashInferMLA | `UNIFORM_BATCH` | |
183
| AITER MLA | `UNIFORM_SINGLE_TOKEN_DECODE` | |
184
| CUTLASS MLA | `UNIFORM_SINGLE_TOKEN_DECODE` | |
185
| Mamba attention| `UNIFORM_SINGLE_TOKEN_DECODE` | |
vllm/attention/layers/chunked_local_attention.py
+1
-1
@@ -32,7 +32,7 @@ def create_chunked_local_attention_backend(
32
underlying_builder = underlying_attn_backend.get_builder_cls()
33
34
class ChunkedLocalAttentionBuilder(underlying_builder): # type: ignore
35
- cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
35
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
36
37
def build(
38
self,
vllm/v1/attention/backends/flash_attn.py
+1
-1
@@ -207,7 +207,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
207
# to FULL_AND_PIECEWISE.
208
# TODO(luka, lucas): audit FA2 as part of:
209
# https://github.com/vllm-project/vllm/issues/22945
210
- cudagraph_support = (
210
+ _cudagraph_support = (
211
AttentionCGSupport.ALWAYS
212
if get_flash_attn_version() == 3
213
else AttentionCGSupport.UNIFORM_BATCH
vllm/v1/attention/backends/flashinfer.py
+23
-15
@@ -15,6 +15,7 @@ from flashinfer import (
15
from flashinfer.decode import _get_range_buf, trtllm_batch_decode_with_kv_cache
16
from flashinfer.prefill import trtllm_batch_context_with_kv_cache
17
from flashinfer.utils import FP4Tensor
18
+from typing_extensions import override
19
20
from vllm import envs
21
from vllm.attention.backends.abstract import (
@@ -274,10 +275,6 @@ class FlashInferMetadata:
275
276
277
class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
277
- cudagraph_support: ClassVar[AttentionCGSupport] = (
278
- AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
279
- )
280
-
278
reorder_batch_threshold: int = 1
279
280
def __init__(
@@ -355,6 +352,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
352
else:
353
self.q_data_type = self.model_config.dtype
354
355
+ # Prefer TRTLLM attention for decoding in all cases.
356
+ # This allows us to use AttentionCGSupport.UNIFORM_BATCH mode.
357
+ self.use_trtllm_decode_attention = can_use_trtllm
358
self._init_reorder_batch_threshold(1, supports_spec_as_decode=can_use_trtllm)
359
360
self._cascade_wrapper = None # Wrapper for cascade attention
@@ -412,6 +412,24 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
412
"passing --block-size 32 or --block-size 64."
413
)
414
415
+ @classmethod
416
+ @override
417
+ def get_cudagraph_support(
418
+ cls: type["FlashInferMetadataBuilder"],
419
+ vllm_config: VllmConfig,
420
+ kv_cache_spec: AttentionSpec,
421
+ ) -> AttentionCGSupport:
422
+ has_trtllm_support = can_use_trtllm_attention(
423
+ num_qo_heads=vllm_config.model_config.get_num_attention_heads(
424
+ vllm_config.parallel_config
425
+ ),
426
+ num_kv_heads=kv_cache_spec.num_kv_heads,
427
+ )
428
+ if has_trtllm_support:
429
+ return AttentionCGSupport.UNIFORM_BATCH
430
+ else:
431
+ return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
432
+
433
def _get_workspace_buffer(self):
434
if self._workspace_buffer is None:
435
buffer_size = envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE
@@ -573,17 +591,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
591
has_sinks=self.has_sinks,
592
has_spec=uses_spec_reorder,
593
)
576
- decode_use_trtllm = use_trtllm_attention(
577
- self.num_qo_heads,
578
- self.num_kv_heads,
579
- num_decode_tokens,
580
- max_seq_len,
581
- self.cache_dtype,
582
- self.q_data_type,
583
- is_prefill=False,
584
- has_sinks=self.has_sinks,
585
- has_spec=uses_spec_reorder,
586
- )
594
+ decode_use_trtllm = self.use_trtllm_decode_attention
595
596
if not (prefill_use_trtllm and decode_use_trtllm):
597
if self.has_sinks:
vllm/v1/attention/backends/gdn_attn.py
+1
-1
@@ -59,7 +59,7 @@ class GDNAttentionMetadata:
59
60
61
class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]):
62
- cudagraph_support = AttentionCGSupport.UNIFORM_BATCH
62
+ _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH
63
64
reorder_batch_threshold: int = 1
65
vllm/v1/attention/backends/mamba_attn.py
+1
-1
@@ -20,7 +20,7 @@ M = TypeVar("M")
20
21
class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
22
reorder_batch_threshold: int = 1
23
- cudagraph_support: ClassVar[AttentionCGSupport] = (
23
+ _cudagraph_support: ClassVar[AttentionCGSupport] = (
24
AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
25
)
26
vllm/v1/attention/backends/mla/cutlass_mla.py
+1
-1
@@ -29,7 +29,7 @@ logger = init_logger(__name__)
29
30
class CutlassMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
31
# enable full CUDA Graph support for decode-only capture
32
- cudagraph_support: ClassVar[AttentionCGSupport] = (
32
+ _cudagraph_support: ClassVar[AttentionCGSupport] = (
33
AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
34
)
35
vllm/v1/attention/backends/mla/flashattn_mla.py
+1
-1
@@ -92,7 +92,7 @@ class FlashAttnMLAMetadata(MLACommonMetadata[FlashAttnMLADecodeMetadata]):
92
93
94
class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]):
95
- cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
95
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
96
query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.VARLEN
97
reorder_batch_threshold: int = 512 # process small prefills with decode pathway
98
vllm/v1/attention/backends/mla/flashinfer_mla.py
+1
-1
@@ -29,7 +29,7 @@ FLASHINFER_MLA_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024
29
30
31
class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
32
- cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
32
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
33
query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
34
35
vllm/v1/attention/backends/mla/flashmla.py
+1
-1
@@ -96,7 +96,7 @@ class FlashMLAMetadata(MLACommonMetadata[FlashMLADecodeMetadata]):
96
97
98
class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
99
- cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
99
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
100
query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
101
reorder_batch_threshold: int = 128 # process small prefills with decode pathway
102
# ^ TODO(matt): tune this
vllm/v1/attention/backends/mla/flashmla_sparse.py
+1
-1
@@ -248,7 +248,7 @@ def triton_convert_req_index_to_global_index(
248
249
@dataclass
250
class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetadata]):
251
- cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
251
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
252
253
def __init__(
254
self,
vllm/v1/attention/backends/mla/indexer.py
+1
-1
@@ -206,7 +206,7 @@ def split_prefill_chunks(
206
207
208
class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
209
- cudagraph_support: ClassVar[AttentionCGSupport] = (
209
+ _cudagraph_support: ClassVar[AttentionCGSupport] = (
210
AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
211
)
212
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+1
-1
@@ -55,7 +55,7 @@ class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]):
55
class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
56
# TODO(luka, lucas): audit this as part of:
57
# https://github.com/vllm-project/vllm/issues/22945
58
- cudagraph_support: ClassVar[AttentionCGSupport] = (
58
+ _cudagraph_support: ClassVar[AttentionCGSupport] = (
59
AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
60
)
61
vllm/v1/attention/backends/rocm_aiter_fa.py
+1
-1
@@ -251,7 +251,7 @@ class AiterFlashAttentionMetadata:
251
class AiterFlashAttentionMetadataBuilder(
252
AttentionMetadataBuilder[AiterFlashAttentionMetadata]
253
):
254
- cudagraph_support = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
254
+ _cudagraph_support = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
255
reorder_batch_threshold: int = 1
256
257
def __init__(
vllm/v1/attention/backends/rocm_attn.py
+1
-1
@@ -63,7 +63,7 @@ class RocmAttentionMetadata:
63
64
65
class RocmAttentionMetadataBuilder(AttentionMetadataBuilder[RocmAttentionMetadata]):
66
- cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
66
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
67
68
def __init__(
69
self,
vllm/v1/attention/backends/triton_attn.py
+1
-1
@@ -67,7 +67,7 @@ class TritonAttentionMetadata:
67
68
69
class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMetadata]):
70
- cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
70
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
71
72
def __init__(
73
self,
vllm/v1/attention/backends/utils.py
+11
-1
@@ -244,7 +244,8 @@ class AttentionCGSupport(enum.Enum):
244
245
class AttentionMetadataBuilder(abc.ABC, Generic[M]):
246
# Does this backend/builder support CUDA Graphs for attention (default: no).
247
- cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
247
+ # Do not access directly. Call get_cudagraph_support() instead.
248
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.NEVER
249
# Does this backend/builder reorder the batch?
250
# If not, set this to None. Otherwise set it to the query
251
# length that will be pulled into the front of the batch.
@@ -263,6 +264,15 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):
264
self.vllm_config = vllm_config
265
self.device = device
266
267
+ @classmethod
268
+ def get_cudagraph_support(
269
+ cls: type["AttentionMetadataBuilder"],
270
+ vllm_config: VllmConfig,
271
+ kv_cache_spec: AttentionSpec,
272
+ ) -> AttentionCGSupport:
273
+ """Get the cudagraph support level of this builder class."""
274
+ return cls._cudagraph_support
275
+
276
def _init_reorder_batch_threshold(
277
self,
278
reorder_batch_threshold: int | None = 1,
vllm/v1/worker/gpu_model_runner.py
+21
-10
@@ -4167,14 +4167,16 @@ class GPUModelRunner(
4167
return attn_groups
4168
4169
attention_backend_maps = []
4170
- attention_backend_set: set[type[AttentionBackend]] = set()
4170
+ attention_backend_list = []
4171
for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
4172
attn_backends = get_attn_backends_for_group(kv_cache_group_spec)
4173
attention_backend_maps.append(attn_backends[0])
4174
- attention_backend_set.update(attn_backends[1])
4174
+ attention_backend_list.append(attn_backends[1])
4175
4176
# Resolve cudagraph_mode before actually initialize metadata_builders
4177
- self._check_and_update_cudagraph_mode(attention_backend_set)
4177
+ self._check_and_update_cudagraph_mode(
4178
+ attention_backend_list, kv_cache_config.kv_cache_groups
4179
+ )
4180
4181
for i, attn_backend_map in enumerate(attention_backend_maps):
4182
self.attn_groups.append(create_attn_groups(attn_backend_map, i))
@@ -4203,22 +4205,31 @@ class GPUModelRunner(
4205
self.calculate_reorder_batch_threshold()
4206
4207
def _check_and_update_cudagraph_mode(
4206
- self, attention_backends: set[type[AttentionBackend]]
4208
+ self,
4209
+ attention_backends: list[set[type[AttentionBackend]]],
4210
+ kv_cache_groups: list[KVCacheGroupSpec],
4211
) -> None:
4212
"""
4213
Resolve the cudagraph_mode when there are multiple attention
4210
- backends with potential conflicting CUDA graph support.
4214
+ groups with potential conflicting CUDA graph support.
4215
Then initialize the cudagraph_dispatcher based on the resolved
4216
cudagraph_mode.
4217
"""
4218
min_cg_support = AttentionCGSupport.ALWAYS
4219
min_cg_backend_name = None
4220
4217
- for attn_backend in attention_backends:
4218
- builder_cls = attn_backend.get_builder_cls()
4219
- if builder_cls.cudagraph_support.value < min_cg_support.value:
4220
- min_cg_support = builder_cls.cudagraph_support
4221
- min_cg_backend_name = attn_backend.__name__
4221
+ for attn_backend_set, kv_cache_group in zip(
4222
+ attention_backends, kv_cache_groups
4223
+ ):
4224
+ for attn_backend in attn_backend_set:
4225
+ builder_cls = attn_backend.get_builder_cls()
4226
+
4227
+ cg_support = builder_cls.get_cudagraph_support(
4228
+ self.vllm_config, kv_cache_group.kv_cache_spec
4229
+ )
4230
+ if cg_support.value < min_cg_support.value:
4231
+ min_cg_support = cg_support
4232
+ min_cg_backend_name = attn_backend.__name__
4233
# Flexible resolve the cudagraph mode
4234
cudagraph_mode = self.compilation_config.cudagraph_mode
4235
# check cudagraph for mixed batch is supported