@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 23322431c8
Fork 0

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, 22:35 UTC

[V1][CUDA] Full cudagraph support for FlashInfer (#21367)

fhl2000 committed Aug 2, 2025 at 09:49 UTC 23322431c802bb1057426c7ca31b22e859b51644
8 files changed +377 -48
vllm/v1/attention/backends/flash_attn.py
+5 -2
@@ -25,7 +25,8 @@ if is_flash_attn_varlen_func_available():
25 from vllm.config import VllmConfig, get_layers_from_vllm_config
26 from vllm.logger import init_logger
27 from vllm.utils import cdiv
28 -from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
28 +from vllm.v1.attention.backends.utils import (AttentionCGSupport,
29 + AttentionMetadataBuilder,
30 CommonAttentionMetadata,
31 get_kv_cache_layout)
32 from vllm.v1.kv_cache_interface import AttentionSpec
@@ -153,7 +154,9 @@ def _get_sliding_window_configs(
154
155 class FlashAttentionMetadataBuilder(
156 AttentionMetadataBuilder[FlashAttentionMetadata]):
156 - full_cudagraph_supported: ClassVar[bool] = get_flash_attn_version() == 3
157 + attn_cudagraph_support: ClassVar[AttentionCGSupport] = \
158 + AttentionCGSupport.NEVER if get_flash_attn_version() == 2 \
159 + else AttentionCGSupport.ALWAYS
160
161 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
162 vllm_config: VllmConfig, device: torch.device):
vllm/v1/attention/backends/flashinfer.py
+323 -34
@@ -4,26 +4,28 @@
4 from __future__ import annotations
5
6 from dataclasses import dataclass
7 -from typing import TYPE_CHECKING, Optional
7 +from typing import TYPE_CHECKING, ClassVar, Optional, Union
8
9 import torch
10 from flashinfer import (BatchDecodeWithPagedKVCacheWrapper,
11 BatchPrefillWithPagedKVCacheWrapper,
12 MultiLevelCascadeAttentionWrapper)
13 -from flashinfer.decode import trtllm_batch_decode_with_kv_cache
13 +from flashinfer.decode import (_get_range_buf, get_seq_lens,
14 + trtllm_batch_decode_with_kv_cache)
15
16 import vllm.envs as envs
17 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
18 AttentionType)
19 from vllm.config import VllmConfig
20 from vllm.logger import init_logger
20 -from vllm.utils import cdiv
21 +from vllm.utils import cdiv, is_pin_memory_available
22 from vllm.utils.flashinfer import use_trtllm_decode_attention
23 from vllm.v1.attention.backends.flash_attn import use_cascade_attention
24 from vllm.v1.attention.backends.utils import (
24 - AttentionMetadataBuilder, CommonAttentionMetadata, get_kv_cache_layout,
25 - get_per_layer_parameters, infer_global_hyperparameters,
26 - reorder_batch_to_split_decodes_and_prefills, split_decodes_and_prefills)
25 + AttentionCGSupport, AttentionMetadataBuilder, CommonAttentionMetadata,
26 + get_kv_cache_layout, get_per_layer_parameters,
27 + infer_global_hyperparameters, reorder_batch_to_split_decodes_and_prefills,
28 + split_decodes_and_prefills)
29 from vllm.v1.kv_cache_interface import AttentionSpec
30
31 if TYPE_CHECKING:
@@ -174,26 +176,66 @@ class FlashInferMetadata:
176
177
178 class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
179 + attn_cudagraph_support: ClassVar[AttentionCGSupport] = \
180 + AttentionCGSupport.PURE_DECODE_ONLY
181
182 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
183 vllm_config: VllmConfig, device: torch.device):
184 self.device = device
185 + self.vllm_config = vllm_config
186 + self.cache_config = vllm_config.cache_config
187 + self.kv_cache_spec = kv_cache_spec
188 self._workspace_buffer = None
189 self._prefill_wrapper = None # Wrapper for prefill/append
183 - self._decode_wrapper = None # Wrapper for decode
190 + self._decode_wrapper = None # Wrapper for decode (general shape)
191 +
192 + self.compilation_config = vllm_config.compilation_config
193 + max_num_pages_per_req = cdiv(vllm_config.model_config.max_model_len,
194 + self.kv_cache_spec.block_size)
195 + max_num_reqs = vllm_config.scheduler_config.max_num_seqs
196 + max_num_pages = max_num_reqs * max_num_pages_per_req
197 + self.enable_cuda_graph = self.compilation_config.full_cuda_graph
198 + if self.enable_cuda_graph:
199 + # For full cudagraph capture, one `decode_wrapper` for each batch
200 + # size is needed for FlashInfer.
201 + self._decode_wrappers_cudagraph: dict[
202 + int, BatchDecodeWithPagedKVCacheWrapper] = {}
203 + self._decode_cudagraph_max_bs = min(
204 + max_num_reqs, self.compilation_config.max_capture_size)
205 +
206 self._cascade_wrapper = None # Wrapper for cascade attention
207
208 # Global hyperparameters shared by all attention layers
209 self.global_hyperparameters = infer_global_hyperparameters(
210 get_per_layer_parameters(vllm_config, layer_names, FlashInferImpl))
211
190 - self.vllm_config = vllm_config
191 - self.cache_config = vllm_config.cache_config
192 - self.kv_cache_spec = kv_cache_spec
193 - max_num_blocks_per_request = cdiv(
194 - vllm_config.model_config.max_model_len,
195 - self.kv_cache_spec.block_size)
196 - self.block_table_arange = torch.arange(max_num_blocks_per_request,
212 + # Preparing persistent buffers (device-side)
213 + self.paged_kv_indptr = torch.zeros(max_num_reqs + 1,
214 + dtype=torch.int32,
215 + device=self.device)
216 + self.paged_kv_indices = torch.zeros(
217 + max_num_pages, # max num pages possible
218 + dtype=torch.int32,
219 + device=self.device)
220 + self.paged_kv_last_page_len = torch.zeros(max_num_reqs,
221 + dtype=torch.int32,
222 + device=self.device)
223 + # host-side buffer
224 + pin_memory = is_pin_memory_available()
225 + self.paged_kv_indptr_cpu = torch.zeros(max_num_reqs + 1,
226 + dtype=torch.int32,
227 + device="cpu",
228 + pin_memory=pin_memory)
229 + self.paged_kv_indices_cpu = torch.zeros(max_num_pages,
230 + dtype=torch.int32,
231 + device="cpu",
232 + pin_memory=pin_memory)
233 + self.paged_kv_last_page_len_cpu = torch.zeros(max_num_reqs,
234 + dtype=torch.int32,
235 + device="cpu",
236 + pin_memory=pin_memory)
237 +
238 + self.block_table_arange = torch.arange(max_num_pages_per_req,
239 dtype=torch.int32,
240 device=self.device)
241
@@ -217,8 +259,16 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
259 self._get_workspace_buffer(), get_kv_cache_layout())
260 return self._prefill_wrapper
261
220 - def _get_decode_wrapper(self):
221 - if self._decode_wrapper is None:
262 + def _get_decode_wrapper(self,
263 + batch_size: int,
264 + use_cudagraph: bool = False):
265 + if use_cudagraph:
266 + decode_wrapper = self._decode_wrappers_cudagraph.get(
267 + batch_size, None)
268 + else:
269 + decode_wrapper = self._decode_wrapper
270 +
271 + if decode_wrapper is None:
272 num_qo_heads = (
273 self.vllm_config.model_config.get_num_attention_heads(
274 self.vllm_config.parallel_config))
@@ -226,11 +276,32 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
276 self.vllm_config.parallel_config)
277 use_tensor_cores = envs.VLLM_FLASHINFER_FORCE_TENSOR_CORES or (
278 num_qo_heads // num_kv_heads > 4)
229 - self._decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
279 +
280 + if use_cudagraph:
281 + paged_kv_indptr = self.paged_kv_indptr[:batch_size + 1]
282 + paged_kv_indices = self.paged_kv_indices
283 + paged_kv_last_page_len = self.paged_kv_last_page_len[:
284 + batch_size]
285 + else:
286 + paged_kv_indptr = None
287 + paged_kv_indices = None
288 + paged_kv_last_page_len = None
289 + decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
290 self._get_workspace_buffer(),
291 get_kv_cache_layout(),
292 + use_cuda_graph=use_cudagraph,
293 + paged_kv_indptr_buffer=paged_kv_indptr,
294 + paged_kv_indices_buffer=paged_kv_indices,
295 + paged_kv_last_page_len_buffer=paged_kv_last_page_len,
296 use_tensor_cores=use_tensor_cores)
233 - return self._decode_wrapper
297 +
298 + # save the decode wrapper
299 + if use_cudagraph:
300 + self._decode_wrappers_cudagraph[batch_size] = decode_wrapper
301 + else:
302 + self._decode_wrapper = decode_wrapper
303 +
304 + return decode_wrapper
305
306 def _get_cascade_wrapper(self):
307 if self._cascade_wrapper is None:
@@ -308,16 +379,44 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
379 )
380
381 if num_decodes > 0:
311 - attn_metadata.decode_wrapper = self._get_decode_wrapper()
382 + pure_decode = num_prefills == 0
383 + # possible required padding for cudagraph replay
384 + use_cudagraph = (self.enable_cuda_graph and pure_decode and
385 + num_decodes <= self._decode_cudagraph_max_bs)
386 + if use_cudagraph:
387 + num_input_tokens = (
388 + self.vllm_config.pad_for_cudagraph(num_decodes))
389 + # Carefully fulfill the padding region with reasonable value
390 + # on cpu.
391 + # Make sure paged_kv_indptr_cpu is not decreasing
392 + self.paged_kv_indptr_cpu[1 + num_decodes:1 +
393 + num_input_tokens].fill_(
394 + attn_metadata.
395 + paged_kv_indptr_cpu[-1])
396 + # Fill the remaining paged_kv_last_page_len_cpu with 1.
397 + # This is because flashinfer treats 0 as a full page
398 + # instead of empty.
399 + self.paged_kv_last_page_len_cpu[
400 + num_decodes:num_input_tokens].fill_(1)
401 +
402 + else:
403 + num_input_tokens = num_decodes
404 +
405 + attn_metadata.decode_wrapper = self._get_decode_wrapper(
406 + num_input_tokens, use_cudagraph)
407 if not use_trtllm_decode_attention(
408 num_decodes, attn_metadata.max_seq_len,
409 self.cache_config.cache_dtype,
410 attn_metadata.num_qo_heads, attn_metadata.num_kv_heads,
411 attn_metadata.head_dim):
317 - attn_metadata.decode_wrapper.plan(
318 - attn_metadata.paged_kv_indptr_cpu[:num_decodes + 1],
412 + # Use the persistent buffer with padding length,
413 + # instead of the same address but chunked version
414 + # in atten_metadata when using cudagraph.
415 + fast_plan_decode(
416 + attn_metadata.decode_wrapper,
417 + self.paged_kv_indptr_cpu[:num_input_tokens + 1],
418 attn_metadata.paged_kv_indices,
320 - attn_metadata.paged_kv_last_page_len_cpu[:num_decodes],
419 + self.paged_kv_last_page_len_cpu[:num_input_tokens],
420 attn_metadata.num_qo_heads,
421 attn_metadata.num_kv_heads,
422 attn_metadata.head_dim,
@@ -336,6 +435,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
435 common_prefix_len: int,
436 common_attn_metadata: CommonAttentionMetadata,
437 fast_build: bool = False) -> FlashInferMetadata:
438 + num_reqs = common_attn_metadata.num_reqs
439 num_actual_tokens = common_attn_metadata.num_actual_tokens
440 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens =\
441 split_decodes_and_prefills(common_attn_metadata)
@@ -381,18 +481,26 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
481 non_blocking=True)
482 mask = (self.block_table_arange[:max_num_blocks].unsqueeze(0)
483 < block_table_bounds.unsqueeze(1))
384 - paged_kv_indices = block_table_tensor[:, :max_num_blocks][mask]
385 -
386 - paged_kv_indptr_cpu = torch.zeros(len(block_table_bounds_cpu) + 1,
387 - dtype=torch.int32,
388 - device='cpu')
389 - paged_kv_indptr_cpu[1:] = block_table_bounds_cpu.cumsum(
390 - dim=0, dtype=torch.int32)
484 + # write self.paged_kv_indices inplace
485 + num_actual_pages = torch.sum(mask)
486 + paged_kv_indices = self.paged_kv_indices[:num_actual_pages]
487 + torch.masked_select(block_table_tensor[:, :max_num_blocks],
488 + mask,
489 + out=paged_kv_indices)
490 +
491 + # write self.paged_kv_indptr_cpu inplace (0-index is always 0)
492 + torch.cumsum(block_table_bounds_cpu,
493 + dim=0,
494 + dtype=torch.int32,
495 + out=self.paged_kv_indptr_cpu[1:1 + num_reqs])
496
497 paged_kv_last_page_len_cpu = seq_lens_cpu % page_size
393 - paged_kv_last_page_len_cpu = torch.where(
394 - paged_kv_last_page_len_cpu == 0, page_size,
395 - paged_kv_last_page_len_cpu)
498 + # write self.paged_kv_last_page_len_cpu inplace
499 + torch.where(paged_kv_last_page_len_cpu == 0,
500 + torch.tensor(page_size),
501 + paged_kv_last_page_len_cpu,
502 + out=self.paged_kv_last_page_len_cpu[:num_reqs])
503 +
504 cache_dtype = self.cache_config.cache_dtype
505 if cache_dtype.startswith("fp8"):
506 kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
@@ -402,9 +510,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
510 attn_metadata = FlashInferMetadata(
511 num_actual_tokens=num_actual_tokens,
512 qo_indptr_cpu=common_attn_metadata.query_start_loc_cpu,
405 - paged_kv_indptr_cpu=paged_kv_indptr_cpu,
513 + paged_kv_indptr_cpu=self.paged_kv_indptr_cpu[:1 + num_reqs],
514 paged_kv_indices=paged_kv_indices,
407 - paged_kv_last_page_len_cpu=paged_kv_last_page_len_cpu,
515 + paged_kv_last_page_len_cpu=self.
516 + paged_kv_last_page_len_cpu[:num_reqs],
517 num_qo_heads=self.vllm_config.model_config.get_num_attention_heads(
518 self.vllm_config.parallel_config),
519 num_kv_heads=self.kv_cache_spec.num_kv_heads,
@@ -431,6 +540,26 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
540
541 return attn_metadata
542
543 + def build_for_cudagraph_capture(
544 + self, common_attn_metadata: CommonAttentionMetadata):
545 + """
546 + This method builds the metadata for full cudagraph capture.
547 + Currently, only decode is supported for full cudagraphs with FlashInfer.
548 + """
549 + m = common_attn_metadata
550 +
551 + assert m.num_reqs == m.num_actual_tokens, \
552 + "FlashInfer only supports decode-only full CUDAGraph capture. " \
553 + "Make sure all cudagraph capture sizes <= max_num_seq."
554 +
555 + m.max_query_len = 1 # decode-only
556 +
557 + return self.build(0, m)
558 +
559 + def can_run_in_cudagraph(
560 + self, common_attn_metadata: CommonAttentionMetadata) -> bool:
561 + return common_attn_metadata.max_query_len == 1
562 +
563 def use_cascade_attention(self, *args, **kwargs) -> bool:
564 if self.kv_cache_spec.dtype != self.vllm_config.model_config.dtype:
565 # TODO: The cascade wrapper currently does not support setting
@@ -638,3 +767,163 @@ class FlashInferImpl(AttentionImpl):
767 out=output[:num_decode_tokens],
768 )
769 return output_padded
770 +
771 +
772 +def fast_plan_decode(
773 + self, # decode wrapper
774 + indptr_cpu: torch.Tensor,
775 + indices: torch.Tensor,
776 + last_page_len_cpu: torch.Tensor,
777 + num_qo_heads: int,
778 + num_kv_heads: int,
779 + head_dim: int,
780 + page_size: int,
781 + pos_encoding_mode: str = "NONE",
782 + window_left: int = -1,
783 + logits_soft_cap: Optional[float] = None,
784 + q_data_type: Optional[Union[str, torch.dtype]] = "float16",
785 + kv_data_type: Optional[Union[str, torch.dtype]] = None,
786 + data_type: Optional[Union[str, torch.dtype]] = None,
787 + sm_scale: Optional[float] = None,
788 + rope_scale: Optional[float] = None,
789 + rope_theta: Optional[float] = None,
790 + non_blocking: bool = True,
791 +) -> None:
792 + """
793 + A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for
794 + cudagraph capture/replay, while the no cudagraph version turns back
795 + to the original plan.
796 + using original plan after passing host-side buffers:
797 + - only host-to-device copy of indptr and last_page_len buffers
798 + Modifications for cudagraph:
799 + - only host-to-device copy of indptr and last_page_len buffers.
800 + - avoid device-to-device copy of indices buffer.
801 +
802 + Part of the code get inspiration from the original plan from FlashInfer repo
803 + and the implementation of fast_decode_plan for FlashInfer in SGlang repo.
804 + """
805 + # Warm up with the original plan if it is first call, and always run the
806 + # original plan if we run for dynamic shape. For fixed shape (cudagraph),
807 + # this warm up is to generate the _cached_module for the decode wrapper.
808 + if not self.is_cuda_graph_enabled or \
809 + getattr(self, "vllm_first_call", True):
810 + self.plan(
811 + indptr_cpu,
812 + indices,
813 + last_page_len_cpu,
814 + num_qo_heads,
815 + num_kv_heads,
816 + head_dim,
817 + page_size,
818 + pos_encoding_mode,
819 + window_left,
820 + logits_soft_cap,
821 + q_data_type,
822 + kv_data_type,
823 + data_type,
824 + sm_scale,
825 + rope_scale,
826 + rope_theta,
827 + non_blocking,
828 + )
829 + self.vllm_first_call = False
830 + return
831 +
832 + assert self.is_cuda_graph_enabled, "Should be cudagraph only here"
833 +
834 + batch_size = len(last_page_len_cpu)
835 + if logits_soft_cap is None:
836 + logits_soft_cap = 0.0
837 +
838 + # Handle data types consistently
839 + if data_type is not None:
840 + if q_data_type is None:
841 + q_data_type = data_type
842 + if kv_data_type is None:
843 + kv_data_type = data_type
844 + elif q_data_type is None:
845 + q_data_type = "float16"
846 +
847 + if kv_data_type is None:
848 + kv_data_type = q_data_type
849 + q_data_type = getattr(torch, q_data_type) if isinstance(
850 + q_data_type, str) else q_data_type
851 + kv_data_type = getattr(torch, kv_data_type) if isinstance(
852 + kv_data_type, str) else kv_data_type
853 +
854 + if self.use_tensor_cores:
855 + qo_indptr_host = _get_range_buf(batch_size + 1, "cpu")
856 +
857 + if batch_size != self._fixed_batch_size:
858 + raise ValueError(
859 + "The batch size should be fixed in cudagraph mode, the runtime "
860 + "batch size {} mismatches the batch size set during "
861 + "initialization {}".format(batch_size, self._fixed_batch_size))
862 + if len(indices) > len(self._paged_kv_indices_buf):
863 + raise ValueError(
864 + "The size of indices should be less than or equal to the "
865 + "allocated buffer")
866 +
867 + # host-to-device copy for the indptr buffer
868 + self._paged_kv_indptr_buf.copy_(indptr_cpu, non_blocking=True)
869 + # host-to-device copy for the last_page_len buffer
870 + self._paged_kv_last_page_len_buf.copy_(last_page_len_cpu,
871 + non_blocking=True)
872 +
873 + indptr_host = indptr_cpu
874 + last_page_len_host = last_page_len_cpu
875 +
876 + if self.use_tensor_cores:
877 + kv_lens_arr_host = get_seq_lens(indptr_host, last_page_len_host,
878 + page_size)
879 +
880 + try:
881 + # Make sure we pass exactly 15 arguments for tensor core version
882 + self._plan_info = self._cached_module.plan(
883 + self._float_workspace_buffer,
884 + self._int_workspace_buffer,
885 + self._pin_memory_int_workspace_buffer,
886 + qo_indptr_host,
887 + indptr_host,
888 + kv_lens_arr_host,
889 + batch_size, # total_num_rows
890 + batch_size,
891 + num_qo_heads,
892 + num_kv_heads,
893 + page_size,
894 + self.is_cuda_graph_enabled,
895 + head_dim,
896 + head_dim,
897 + False, # causal
898 + )
899 + except Exception as e:
900 + raise RuntimeError(f"Error in tensor core plan: {e}") from e
901 + else:
902 + try:
903 + # Make sure we pass exactly 15 arguments for standard version
904 + self._plan_info = self._cached_module.plan(
905 + self._float_workspace_buffer,
906 + self._int_workspace_buffer,
907 + self._pin_memory_int_workspace_buffer,
908 + indptr_host,
909 + batch_size,
910 + num_qo_heads,
911 + num_kv_heads,
912 + page_size,
913 + self.is_cuda_graph_enabled,
914 + window_left,
915 + logits_soft_cap,
916 + head_dim,
917 + head_dim,
918 + torch.empty(0, dtype=q_data_type),
919 + torch.empty(0, dtype=kv_data_type),
920 + )
921 + except Exception as e:
922 + raise RuntimeError(f"Error in standard plan: {e}") from e
923 +
924 + self._pos_encoding_mode = pos_encoding_mode
925 + self._window_left = window_left
926 + self._logits_soft_cap = logits_soft_cap
927 + self._sm_scale = sm_scale
928 + self._rope_scale = rope_scale
929 + self._rope_theta = rope_theta
vllm/v1/attention/backends/mla/flashmla.py
+3 -1
@@ -18,6 +18,7 @@ from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
18 MLACommonImpl,
19 MLACommonMetadata,
20 MLACommonMetadataBuilder)
21 +from vllm.v1.attention.backends.utils import AttentionCGSupport
22 from vllm.v1.kv_cache_interface import AttentionSpec
23
24 logger = init_logger(__name__)
@@ -54,7 +55,8 @@ class FlashMLAMetadata(MLACommonMetadata[FlashMLADecodeMetadata]):
55
56
57 class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
57 - full_cudagraph_supported: ClassVar[bool] = True # Decode-only
58 + attn_cudagraph_support: ClassVar[AttentionCGSupport] = \
59 + AttentionCGSupport.PURE_DECODE_ONLY
60
61 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
62 vllm_config: VllmConfig, device: torch.device):
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+3 -1
@@ -17,6 +17,7 @@ from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
17 MLACommonImpl,
18 MLACommonMetadata,
19 MLACommonMetadataBuilder)
20 +from vllm.v1.attention.backends.utils import AttentionCGSupport
21 from vllm.v1.kv_cache_interface import AttentionSpec
22
23 # yapf: enable
@@ -64,7 +65,8 @@ class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]):
65
66
67 class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
67 - full_cudagraph_supported: ClassVar[bool] = True # decode only
68 + attn_cudagraph_support: ClassVar[AttentionCGSupport] = \
69 + AttentionCGSupport.PURE_DECODE_ONLY
70
71 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
72 vllm_config: VllmConfig, device: torch.device):
vllm/v1/attention/backends/triton_attn.py
+4 -2
@@ -18,7 +18,8 @@ from vllm.config import VllmConfig
18 from vllm.logger import init_logger
19 from vllm.platforms import current_platform
20 from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
21 -from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
21 +from vllm.v1.attention.backends.utils import (AttentionCGSupport,
22 + AttentionMetadataBuilder,
23 CommonAttentionMetadata)
24 from vllm.v1.kv_cache_interface import AttentionSpec
25
@@ -57,7 +58,8 @@ class TritonAttentionMetadata:
58
59 class TritonAttentionMetadataBuilder(
60 AttentionMetadataBuilder[TritonAttentionMetadata]):
60 - full_cudagraph_supported: ClassVar[bool] = True
61 + attn_cudagraph_support: ClassVar[AttentionCGSupport] = \
62 + AttentionCGSupport.ALWAYS
63
64 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
65 vllm_config: VllmConfig, device: torch.device):
vllm/v1/attention/backends/utils.py
+17 -1
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import abc
4 +import enum
5 import functools
6 from abc import abstractmethod
7 from dataclasses import dataclass, make_dataclass
@@ -65,9 +66,24 @@ class CommonAttentionMetadata:
66 M = TypeVar("M")
67
68
69 +class AttentionCGSupport(enum.Enum):
70 + """ Constants for the cudagraph support of the attention backend
71 + Here we do not consider the cascade attention, as currently
72 + it is never cudagraph supported."""
73 +
74 + NEVER = 0
75 + """NO cudagraph support"""
76 + PURE_DECODE_ONLY = 1
77 + """Cudagraph supported for pure decode, need to run without
78 + cudagraph for mixed prefill-decode batches"""
79 + ALWAYS = 2
80 + """Cudagraph always supported"""
81 +
82 +
83 class AttentionMetadataBuilder(abc.ABC, Generic[M]):
84 # Does this backend/builder support CUDA Graphs for attention.
70 - full_cudagraph_supported: ClassVar[bool] = False
85 + attn_cudagraph_support: ClassVar[AttentionCGSupport] = \
86 + AttentionCGSupport.NEVER
87
88 @abstractmethod
89 def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
vllm/v1/worker/gpu_model_runner.py
+17 -7
@@ -47,7 +47,7 @@ from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
47 is_pin_memory_available, round_up, supports_dynamo)
48 from vllm.v1.attention.backends.mamba_selectors import get_mamba_attn_backend
49 from vllm.v1.attention.backends.utils import (
50 - AttentionMetadataBuilder, CommonAttentionMetadata,
50 + AttentionCGSupport, AttentionMetadataBuilder, CommonAttentionMetadata,
51 make_kv_sharing_fast_prefill_attention_metadata,
52 make_local_attention_virtual_batches)
53 from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
@@ -2619,12 +2619,22 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2619 self.device,
2620 )
2621
2622 - if (self.full_cuda_graph
2623 - and not attn_metadata_builder_i.full_cudagraph_supported):
2624 - raise ValueError(
2625 - f"Full CUDAGraph not supported for "
2626 - f"{attn_backend_i.__name__}. Turn off CompilationConfig."
2627 - f"full_cuda_graph or use a different attention backend.")
2622 + if self.full_cuda_graph:
2623 + if attn_metadata_builder_i.attn_cudagraph_support == \
2624 + AttentionCGSupport.NEVER:
2625 + raise ValueError(f"Full CUDAGraph not supported for "
2626 + f"{attn_backend_i.__name__}. Turn off "
2627 + f"CompilationConfig.full_cuda_graph or use a "
2628 + f" different attention backend.")
2629 + if attn_metadata_builder_i.attn_cudagraph_support == \
2630 + AttentionCGSupport.PURE_DECODE_ONLY:
2631 + # Limit the max cudagraph size to the max number of
2632 + # sequences for pure decode only cudagraph backend,
2633 + # whose max_query_len is 1.
2634 + self.cudagraph_batch_sizes = [
2635 + size for size in self.cudagraph_batch_sizes
2636 + if size <= self.scheduler_config.max_num_seqs
2637 + ]
2638 return attn_backend_i, attn_metadata_builder_i
2639
2640 def initialize_attn_backend(self, kv_cache_config: KVCacheConfig) -> None:
vllm/v1/worker/gpu_worker.py
+5
@@ -321,11 +321,16 @@ class Worker(WorkerBase):
321 if get_pp_group().is_last_rank:
322 max_num_reqs = min(self.scheduler_config.max_num_seqs,
323 self.scheduler_config.max_num_batched_tokens)
324 + # activate building attn_metadata for this dummy run to avoid
325 + # potential illegal memory access for full cudagraph relay.
326 + attn_cudagraph = self.compilation_config.full_cuda_graph and\
327 + not self.model_config.enforce_eager
328
329 # We skip EPLB here since we don't want to record dummy metrics
330 hidden_states, last_hidden_states = \
331 self.model_runner._dummy_run(
332 num_tokens=max_num_reqs,
333 + capture_attn_cudagraph=attn_cudagraph,
334 skip_eplb=True,
335 )
336 if self.model_runner.is_pooling_model: