@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3597b06a4f
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 15, 06:00 UTC

[CUDA] Enable full cudagraph for FlashMLA (#18581)

Signed-off-by: luka <luka@neuralmagic.com>

Luka Govedič committed Jun 13, 2025 at 14:12 UTC 3597b06a4faafae96f9696e736ebc198520dc8ec
17 files changed +453 -220
tests/compile/piecewise/test_full_cudagraph.py
+107 -51
@@ -2,15 +2,16 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import contextlib
4 import os
5 +import weakref
6 +from contextlib import ExitStack
7
8 import pytest
9
10 +from tests.utils import wait_for_gpu_memory_to_clear
11 from vllm import LLM, SamplingParams
12 from vllm.config import CompilationConfig
13 from vllm.platforms import current_platform
14
12 -MODEL = "Qwen/Qwen2-1.5B-Instruct"
13 -
15
16 @contextlib.contextmanager
17 def temporary_environ(env_vars):
@@ -31,64 +32,119 @@ def temporary_environ(env_vars):
32 os.environ[k] = v
33
34
34 -@pytest.fixture(scope="module")
35 -def full_cudagraph_llm():
36 - with temporary_environ({
37 - "VLLM_USE_V1": "1",
38 - "VLLM_FLASH_ATTN_VERSION": "3"
39 - }):
40 - return LLM(model=MODEL,
41 - gpu_memory_utilization=0.3,
42 - compilation_config=CompilationConfig(full_cuda_graph=True))
43 -
35 +@pytest.fixture(scope="class")
36 +def llm_pair(request):
37 + model = request.param
38
45 -@pytest.fixture(scope="module")
46 -def piecewise_llm():
39 with temporary_environ({
40 "VLLM_USE_V1": "1",
41 "VLLM_FLASH_ATTN_VERSION": "3"
42 }):
51 - return LLM(model=MODEL,
52 - gpu_memory_utilization=0.6,
53 - compilation_config=CompilationConfig())
54 -
55 -
56 -def generate_text(llm: LLM, batch_size: int, max_tokens: int):
57 - prompts = ["Hi my name is"] * batch_size
58 - sampling_params = SamplingParams(temperature=0.0,
59 - max_tokens=max_tokens,
60 - top_p=0.95)
61 -
62 - return llm.generate(prompts, sampling_params)
63 -
64 -
43 + full = LLM(
44 + model=model,
45 + gpu_memory_utilization=0.45,
46 + trust_remote_code=True,
47 + max_model_len=1024,
48 + compilation_config=CompilationConfig(full_cuda_graph=True),
49 + )
50 + piecewise = LLM(
51 + model=model,
52 + gpu_memory_utilization=0.45,
53 + trust_remote_code=True,
54 + max_model_len=1024,
55 + compilation_config=CompilationConfig(),
56 + )
57 +
58 + # PyTest caches the fixture values so we use weakref.proxy to enable GC
59 + yield weakref.proxy(full), weakref.proxy(piecewise)
60 + del full
61 + del piecewise
62 +
63 + wait_for_gpu_memory_to_clear(
64 + devices=[0],
65 + threshold_ratio=0.1,
66 + )
67 +
68 +
69 +@pytest.mark.parametrize(
70 + "llm_pair",
71 + [
72 + # Model names for the llm_pair fixture
73 + "deepseek-ai/DeepSeek-V2-Lite",
74 + "Qwen/Qwen2-1.5B-Instruct"
75 + ],
76 + indirect=True)
77 @pytest.mark.skipif(current_platform.get_device_capability() != (9, 0),
66 - reason="Only Hopper GPUs support FlashAttention 3")
67 -@pytest.mark.parametrize(("batch_size", "max_tokens"), [(1, 10), (7, 10),
68 - (16, 10), (25, 10),
69 - (32, 10), (45, 10),
70 - (64, 10), (8, 5),
71 - (8, 20), (8, 200)])
72 -def test_full_cudagraph(batch_size, max_tokens, full_cudagraph_llm,
73 - piecewise_llm):
78 + reason="Only Hopper GPUs support FA3 and FlashMLA")
79 +class TestFullCUDAGraph:
80 """
75 - Load full cudagraph model and piecewise model once, and at the same time to
76 - reuse them across various test cases.
81 + Use a class such that an llm pair is constructed once for all
82 + batch_size/max_tokens combinations and released immediately after.
83
78 - Test various batch sizes and max_tokens to ensure that the full cudagraph
79 - compilation works for padded cases too.
84 + Module-scope fixtures would stick around the whole time,
85 + meaning there would be multiple LLM instances hogging memory simultaneously.
86 """
81 - piecewise_responses = generate_text(piecewise_llm,
82 - batch_size=batch_size,
83 - max_tokens=max_tokens)
84 - full_cudagraph_responses = generate_text(full_cudagraph_llm,
85 - batch_size=batch_size,
86 - max_tokens=max_tokens)
87
88 - # Check that all responses are the same
89 - for i in range(len(piecewise_responses)):
90 - assert piecewise_responses[i].outputs[
91 - 0].text == full_cudagraph_responses[i].outputs[0].text
88 + @pytest.mark.parametrize(("batch_size", "max_tokens"), [
89 + (1, 10),
90 + (7, 10),
91 + (16, 10),
92 + (25, 10),
93 + (32, 10),
94 + (45, 10),
95 + (64, 10),
96 + (123, 10),
97 + (8, 5),
98 + (8, 30),
99 + ])
100 + def test_full_cudagraph(self, batch_size, max_tokens,
101 + llm_pair: tuple[LLM, LLM]):
102 + """
103 + Test various batch sizes and max_tokens to ensure that the
104 + full cudagraph compilation works for padded cases too.
105 + """
106 +
107 + piecewise_llm, full_cudagraph_llm = llm_pair
108 +
109 + prompts = ["Hello, my name is"] * batch_size
110 + sampling_params = SamplingParams(temperature=0.0,
111 + max_tokens=max_tokens,
112 + top_p=0.95)
113 +
114 + piecewise_responses = piecewise_llm.generate(prompts, sampling_params)
115 + full_responses = full_cudagraph_llm.generate(prompts, sampling_params)
116 +
117 + # Check that all responses are the same
118 + for piecewise_res, full_res in zip(piecewise_responses,
119 + full_responses):
120 + assert piecewise_res.outputs[0].text == full_res.outputs[0].text
121 +
122 +
123 +@pytest.mark.parametrize(
124 + "model, supported",
125 + [
126 + ("Qwen/Qwen2-1.5B-Instruct", True),
127 + # MLA does not support capturing CUDA Graphs with size > max_num_seqs
128 + ("deepseek-ai/DeepSeek-V2-Lite", False),
129 + ])
130 +@pytest.mark.skipif(current_platform.get_device_capability() != (9, 0),
131 + reason="Only Hopper GPUs support FA3 and FlashMLA")
132 +def test_lower_max_num_seqs(model, supported):
133 + with temporary_environ({
134 + "VLLM_USE_V1": "1",
135 + "VLLM_FLASH_ATTN_VERSION": "3"
136 + }), ExitStack() as stack:
137 + if not supported:
138 + stack.enter_context(pytest.raises(RuntimeError))
139 +
140 + llm = LLM(model=model,
141 + max_num_seqs=256,
142 + trust_remote_code=True,
143 + max_model_len=1024,
144 + compilation_config=CompilationConfig(
145 + full_cuda_graph=True,
146 + cudagraph_capture_sizes=[64, 256, 512]))
147 + llm.generate(["Hello, my name is"] * 10)
148
149
150 def test_full_cudagraph_with_invalid_backend():
@@ -97,5 +153,5 @@ def test_full_cudagraph_with_invalid_backend():
153 "VLLM_FLASH_ATTN_VERSION":
154 "2" #FA2 not supported with full_cuda_graph
155 }), pytest.raises(RuntimeError):
100 - LLM(model=MODEL,
156 + LLM(model="Qwen/Qwen2-1.5B-Instruct",
157 compilation_config=CompilationConfig(full_cuda_graph=True))
tests/compile/piecewise/test_simple.py
+5 -11
@@ -4,7 +4,7 @@
4 Test the piecewise compilation with a simple model so that we
5 can exactly calculate the expected output and side effects.
6 """
7 -
7 +import pytest
8 import torch
9 from torch import nn
10 from torch.library import Library
@@ -14,6 +14,7 @@ from vllm.compilation.decorators import support_torch_compile
14 from vllm.config import (CompilationConfig, CompilationLevel, VllmConfig,
15 set_current_vllm_config)
16 from vllm.envs import VLLM_USE_V1
17 +from vllm.forward_context import set_forward_context
18 from vllm.utils import direct_register_custom_op
19
20 global_counter = 0
@@ -76,7 +77,8 @@ class SillyModel(nn.Module):
77 return x
78
79
79 -def _test_simple_piecewise_compile(*, use_inductor):
80 +@pytest.mark.parametrize("use_inductor", [True, False])
81 +def test_simple_piecewise_compile(use_inductor):
82 assert VLLM_USE_V1
83
84 vllm_config = VllmConfig(compilation_config=CompilationConfig(
@@ -99,7 +101,7 @@ def _test_simple_piecewise_compile(*, use_inductor):
101 num_backend_compilations=3, # num_piecewise_capturable_graphs_seen
102 num_cudagraph_captured=
103 6, # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
102 - ):
104 + ), set_forward_context({}, vllm_config=vllm_config):
105
106 model(inputs)
107
@@ -112,11 +114,3 @@ def _test_simple_piecewise_compile(*, use_inductor):
114 output = model(input)
115 assert global_counter == 2
116 assert torch.allclose(output.cpu(), torch.tensor([3., 1.]))
115 -
116 -
117 -def test_simple_piecewise_compile_inductor():
118 - _test_simple_piecewise_compile(use_inductor=True)
119 -
120 -
121 -def test_simple_piecewise_compile_no_inductor():
122 - _test_simple_piecewise_compile(use_inductor=False)
tests/compile/piecewise/test_toy_llama.py
+21 -24
@@ -11,6 +11,7 @@ initialized randomly with a fixed seed.
11 from dataclasses import dataclass
12 from typing import Any, Optional
13
14 +import pytest
15 import torch
16 from torch import nn
17 from torch.library import Library
@@ -19,6 +20,7 @@ from vllm.compilation.counter import compilation_counter
20 from vllm.compilation.decorators import support_torch_compile
21 from vllm.config import (CompilationConfig, CompilationLevel, VllmConfig,
22 set_current_vllm_config)
23 +from vllm.forward_context import set_forward_context
24 from vllm.utils import direct_register_custom_op
25
26 # create a library to hold the custom op
@@ -285,29 +287,32 @@ def run_model(llama_config,
287 vllm_config=vllm_config,
288 prefix="").eval().cuda()
289
288 - B = 16 # max batch size
289 - input_ids = torch.randint(0, llama_config.vocab_size, (B, )).cuda()
290 - positions = torch.arange(B).cuda()
290 + with set_forward_context({}, vllm_config=vllm_config):
291 + B = 16 # max batch size
292 + input_ids = torch.randint(0, llama_config.vocab_size, (B, )).cuda()
293 + positions = torch.arange(B).cuda()
294
292 - model(input_ids, positions)
293 - model(input_ids[:2], positions[:2])
294 - model(input_ids[:1], positions[:1])
295 + model(input_ids, positions)
296 + model(input_ids[:2], positions[:2])
297 + model(input_ids[:1], positions[:1])
298
296 - input_ids[:2].zero_()
297 - output = model(input_ids[:2], positions[:2])
299 + input_ids[:2].zero_()
300 + output = model(input_ids[:2], positions[:2])
301
299 - output = output.cpu()
302 + output = output.cpu()
303
301 - if llama_config.tractable_init:
302 - expected_output = tractable_computation(input_ids[:2], positions[:2],
303 - llama_config).cpu()
304 + if llama_config.tractable_init:
305 + expected_output = tractable_computation(input_ids[:2],
306 + positions[:2],
307 + llama_config).cpu()
308
305 - assert torch.allclose(output, expected_output)
306 - else:
307 - return output.cpu()
309 + assert torch.allclose(output, expected_output)
310 + else:
311 + return output.cpu()
312
313
310 -def _test_toy_llama(*, use_inductor):
314 +@pytest.mark.parametrize("use_inductor", [True, False])
315 +def test_toy_llama(use_inductor: bool):
316 # compare output with and without piecewise compilation
317
318 llama_config = LlamaConfig(hidden_size=128,
@@ -379,14 +384,6 @@ def _test_toy_llama(*, use_inductor):
384 assert torch.allclose(outputs[0], outputs[i])
385
386
382 -def test_toy_llama_inductor():
383 - _test_toy_llama(use_inductor=True)
384 -
385 -
386 -def test_toy_no_inductor():
387 - _test_toy_llama(use_inductor=False)
388 -
389 -
387 @torch.inference_mode
388 def benchmark():
389 from triton.testing import do_bench
tests/utils.py
+21 -9
@@ -667,42 +667,54 @@ def get_physical_device_indices(devices):
667
668
669 @_nvml()
670 -def wait_for_gpu_memory_to_clear(devices: list[int],
671 - threshold_bytes: int,
670 +def wait_for_gpu_memory_to_clear(*,
671 + devices: list[int],
672 + threshold_bytes: Optional[int] = None,
673 + threshold_ratio: Optional[float] = None,
674 timeout_s: float = 120) -> None:
675 + assert threshold_bytes is not None or threshold_ratio is not None
676 # Use nvml instead of pytorch to reduce measurement error from torch cuda
677 # context.
678 devices = get_physical_device_indices(devices)
679 start_time = time.time()
680 while True:
681 output: dict[int, str] = {}
679 - output_raw: dict[int, float] = {}
682 + output_raw: dict[int, tuple[float, float]] = {}
683 for device in devices:
684 if current_platform.is_rocm():
685 dev_handle = amdsmi_get_processor_handles()[device]
686 mem_info = amdsmi_get_gpu_vram_usage(dev_handle)
687 gb_used = mem_info["vram_used"] / 2**10
688 + gb_total = mem_info["vram_total"] / 2**10
689 else:
690 dev_handle = nvmlDeviceGetHandleByIndex(device)
691 mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
692 gb_used = mem_info.used / 2**30
689 - output_raw[device] = gb_used
690 - output[device] = f'{gb_used:.02f}'
693 + gb_total = mem_info.total / 2**30
694 + output_raw[device] = (gb_used, gb_total)
695 + output[device] = f'{gb_used:.02f}/{gb_total:.02f}'
696
692 - print('gpu memory used (GB): ', end='')
697 + print('gpu memory used/total (GiB): ', end='')
698 for k, v in output.items():
699 print(f'{k}={v}; ', end='')
700 print('')
701
702 + if threshold_bytes is not None:
703 + is_free = lambda used, total: used <= threshold_bytes / 2**30
704 + threshold = f"{threshold_bytes/2**30} GiB"
705 + else:
706 + is_free = lambda used, total: used / total <= threshold_ratio
707 + threshold = f"{threshold_ratio:.2f}"
708 +
709 dur_s = time.time() - start_time
698 - if all(v <= (threshold_bytes / 2**30) for v in output_raw.values()):
710 + if all(is_free(used, total) for used, total in output_raw.values()):
711 print(f'Done waiting for free GPU memory on devices {devices=} '
700 - f'({threshold_bytes/2**30=}) {dur_s=:.02f}')
712 + f'({threshold=}) {dur_s=:.02f}')
713 break
714
715 if dur_s >= timeout_s:
716 raise ValueError(f'Memory of devices {devices=} not free after '
705 - f'{dur_s=:.02f} ({threshold_bytes/2**30=})')
717 + f'{dur_s=:.02f} ({threshold=})')
718
719 time.sleep(5)
720
vllm/compilation/cuda_piecewise_backend.py
+5 -1
@@ -14,6 +14,7 @@ from vllm.compilation.backends import VllmBackend
14 from vllm.compilation.counter import compilation_counter
15 from vllm.compilation.monitor import end_monitoring_torch_compile
16 from vllm.config import VllmConfig
17 +from vllm.forward_context import get_forward_context
18 from vllm.logger import init_logger
19 from vllm.utils import weak_ref_tensors
20
@@ -137,7 +138,10 @@ class CUDAPiecewiseBackend:
138 if self.is_last_graph and not self.to_be_compiled_sizes:
139 self.check_for_ending_compilation()
140
140 - if not entry.use_cudagraph:
141 + # Skip CUDA graphs if this entry doesn't use them OR
142 + # if we're supposed to skip them globally
143 + skip_cuda_graphs = get_forward_context().skip_cuda_graphs
144 + if not entry.use_cudagraph or skip_cuda_graphs:
145 return entry.runnable(*args)
146
147 if entry.cudagraph is None:
vllm/entrypoints/llm.py
+2 -1
@@ -179,7 +179,8 @@ class LLM:
179 hf_overrides: Optional[HfOverrides] = None,
180 mm_processor_kwargs: Optional[dict[str, Any]] = None,
181 override_pooler_config: Optional[PoolerConfig] = None,
182 - compilation_config: Optional[Union[int, dict[str, Any]]] = None,
182 + compilation_config: Optional[Union[int, dict[str, Any],
183 + CompilationConfig]] = None,
184 **kwargs,
185 ) -> None:
186 """LLM constructor."""
vllm/forward_context.py
+12 -6
@@ -94,6 +94,7 @@ class ForwardContext:
94 virtual_engine: int # set dynamically for each forward pass
95 # set dynamically for each forward pass
96 dp_metadata: Optional[DPMetadata] = None
97 + skip_cuda_graphs: bool = False
98
99
100 _forward_context: Optional[ForwardContext] = None
@@ -108,11 +109,14 @@ def get_forward_context() -> ForwardContext:
109
110
111 @contextmanager
111 -def set_forward_context(attn_metadata: Any,
112 - vllm_config: VllmConfig,
113 - virtual_engine: int = 0,
114 - num_tokens: Optional[int] = None,
115 - num_tokens_across_dp: Optional[torch.Tensor] = None):
112 +def set_forward_context(
113 + attn_metadata: Any,
114 + vllm_config: VllmConfig,
115 + virtual_engine: int = 0,
116 + num_tokens: Optional[int] = None,
117 + num_tokens_across_dp: Optional[torch.Tensor] = None,
118 + skip_cuda_graphs: bool = False,
119 +):
120 """A context manager that stores the current forward context,
121 can be attention metadata, etc.
122 Here we can inject common logic for every model forward pass.
@@ -135,7 +139,9 @@ def set_forward_context(attn_metadata: Any,
139 static_forward_context,
140 virtual_engine=virtual_engine,
141 attn_metadata=attn_metadata,
138 - dp_metadata=dp_metadata)
142 + dp_metadata=dp_metadata,
143 + skip_cuda_graphs=skip_cuda_graphs,
144 + )
145
146 try:
147 yield
vllm/v1/attention/backends/cpu_attn.py
+8 -4
@@ -7,7 +7,8 @@ from vllm.attention.backends.torch_sdpa import (TorchSDPABackendImpl,
7 TorchSDPAMetadata)
8 from vllm.attention.backends.utils import CommonAttentionState
9 from vllm.attention.ops.ipex_attn import PagedAttention
10 -from vllm.v1.attention.backends.utils import CommonAttentionMetadata
10 +from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
11 + CommonAttentionMetadata)
12 from vllm.v1.core.sched.output import SchedulerOutput
13 from vllm.v1.kv_cache_interface import AttentionSpec
14 from vllm.v1.worker.block_table import BlockTable
@@ -53,7 +54,7 @@ class TorchSDPABackend:
54 return False
55
56
56 -class TorchSDPAMetadataBuilderV1:
57 +class TorchSDPAMetadataBuilderV1(AttentionMetadataBuilder[TorchSDPAMetadata]):
58
59 def __init__(self, runner: CPUModelRunner, kv_cache_spec: AttentionSpec,
60 block_table: BlockTable) -> None:
@@ -118,9 +119,12 @@ class TorchSDPAMetadataBuilderV1:
119
120 return True
121
121 - def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
122 - common_prefix_len: int,
122 + def build(self, common_prefix_len: int,
123 common_attn_metadata: CommonAttentionMetadata):
124 + num_reqs = common_attn_metadata.num_reqs
125 + num_actual_tokens = common_attn_metadata.num_actual_tokens
126 + max_query_len = common_attn_metadata.max_query_len
127 +
128 runner = self.runner
129 block_table = self.block_table
130 seq_lens_np = runner.seq_lens_np[:num_reqs]
vllm/v1/attention/backends/flash_attn.py
+18 -11
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Attention layer with FlashAttention."""
4 from dataclasses import dataclass
5 -from typing import TYPE_CHECKING, Any, Optional
5 +from typing import TYPE_CHECKING, Any, ClassVar, Optional
6
7 import numpy as np
8 import torch
@@ -21,13 +21,12 @@ from vllm.distributed.kv_transfer.kv_connector.utils import (
21 from vllm.logger import init_logger
22 from vllm.platforms import current_platform
23 from vllm.utils import cdiv
24 -from vllm.v1.attention.backends.utils import CommonAttentionMetadata
24 +from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
25 + CommonAttentionMetadata)
26 from vllm.v1.kv_cache_interface import AttentionSpec
27 from vllm.v1.worker.block_table import BlockTable
28
29 if TYPE_CHECKING:
29 - from vllm.v1.core.sched.output import SchedulerOutput
30 - from vllm.v1.worker.gpu_input_batch import InputBatch
30 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
31
32 if current_platform.is_cuda():
@@ -306,7 +305,9 @@ def _get_sliding_window_configs(
305 return sliding_window_configs
306
307
309 -class FlashAttentionMetadataBuilder:
308 +class FlashAttentionMetadataBuilder(
309 + AttentionMetadataBuilder[FlashAttentionMetadata]):
310 + full_cudagraph_supported: ClassVar[bool] = get_flash_attn_version() == 3
311
312 def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,
313 block_table: BlockTable):
@@ -336,13 +337,14 @@ class FlashAttentionMetadataBuilder:
337 # populated on first build() call.
338 self.aot_sliding_window: Optional[tuple[int, int]] = None
339
339 - def reorder_batch(self, input_batch: "InputBatch",
340 - scheduler_output: "SchedulerOutput") -> bool:
341 - return False
340 + def build(
341 + self, common_prefix_len: int,
342 + common_attn_metadata: CommonAttentionMetadata
343 + ) -> FlashAttentionMetadata:
344 + num_reqs = common_attn_metadata.num_reqs
345 + num_actual_tokens = common_attn_metadata.num_actual_tokens
346 + max_query_len = common_attn_metadata.max_query_len
347
343 - def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
344 - common_prefix_len: int,
345 - common_attn_metadata: CommonAttentionMetadata):
348 max_seq_len = int(self.runner.seq_lens_np[:num_reqs].max())
349 query_start_loc = common_attn_metadata.query_start_loc
350 seq_lens = common_attn_metadata.seq_lens
@@ -496,6 +498,11 @@ class FlashAttentionMetadataBuilder:
498 )
499 return attn_metadata
500
501 + def can_run_in_cudagraph(
502 + self, common_attn_metadata: CommonAttentionMetadata) -> bool:
503 + # Full CUDA Graph always supported (FA2 support checked separately)
504 + return True
505 +
506 def use_cascade_attention(self, *args, **kwargs) -> bool:
507 return use_cascade_attention(*args, **kwargs)
508
vllm/v1/attention/backends/flashinfer.py
+7 -4
@@ -18,7 +18,8 @@ from vllm.attention.layer import Attention
18 from vllm.config import VllmConfig, get_layers_from_vllm_config
19 from vllm.logger import init_logger
20 from vllm.v1.attention.backends.flash_attn import use_cascade_attention
21 -from vllm.v1.attention.backends.utils import CommonAttentionMetadata
21 +from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
22 + CommonAttentionMetadata)
23 from vllm.v1.kv_cache_interface import AttentionSpec
24 from vllm.v1.worker.block_table import BlockTable
25
@@ -202,7 +203,7 @@ class FlashInferMetadata:
203 f" received {self.head_dim}.")
204
205
205 -class FlashInferMetadataBuilder:
206 +class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
207
208 def __init__(self, runner: GPUModelRunner, kv_cache_spec: AttentionSpec,
209 block_table: BlockTable):
@@ -399,9 +400,11 @@ class FlashInferMetadataBuilder:
400 kv_data_type=attn_metadata.data_type,
401 )
402
402 - def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
403 - common_prefix_len: int,
403 + def build(self, common_prefix_len: int,
404 common_attn_metadata: CommonAttentionMetadata):
405 + num_reqs = common_attn_metadata.num_reqs
406 + num_actual_tokens = common_attn_metadata.num_actual_tokens
407 +
408 assert self._num_decodes + self._num_prefills == num_reqs
409 assert (self._num_decode_tokens +
410 self._num_prefill_tokens == num_actual_tokens)
vllm/v1/attention/backends/flex_attention.py
+9 -13
@@ -15,7 +15,8 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
15 is_quantized_kv_cache)
16 from vllm.logger import init_logger
17 from vllm.platforms import current_platform
18 -from vllm.v1.attention.backends.utils import CommonAttentionMetadata
18 +from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
19 + CommonAttentionMetadata)
20 from vllm.v1.kv_cache_interface import AttentionSpec
21 from vllm.v1.worker.block_table import BlockTable
22
@@ -25,8 +26,6 @@ if current_platform.is_cuda():
26 logger = init_logger(__name__)
27
28 if TYPE_CHECKING:
28 - from vllm.v1.core.sched.output import SchedulerOutput
29 - from vllm.v1.worker.gpu_input_batch import InputBatch
29 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
30
31 create_block_mask_compiled = torch.compile(create_block_mask,
@@ -256,7 +255,8 @@ class FlexAttentionMetadata:
255 self.block_mask = self.build_block_mask()
256
257
259 -class FlexAttentionMetadataBuilder:
258 +class FlexAttentionMetadataBuilder(
259 + AttentionMetadataBuilder[FlexAttentionMetadata]):
260
261 def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,
262 block_table: BlockTable):
@@ -272,13 +272,12 @@ class FlexAttentionMetadataBuilder:
272 self.kv_cache_spec = kv_cache_spec
273 self.block_table = block_table
274
275 - def reorder_batch(self, input_batch: "InputBatch",
276 - scheduler_output: "SchedulerOutput") -> bool:
277 - return False
278 -
279 - def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
280 - common_prefix_len: int,
275 + def build(self, common_prefix_len: int,
276 common_attn_metadata: CommonAttentionMetadata):
277 + num_reqs = common_attn_metadata.num_reqs
278 + num_actual_tokens = common_attn_metadata.num_actual_tokens
279 + max_query_len = common_attn_metadata.max_query_len
280 +
281 max_seq_len = self.runner.seq_lens_np[:num_reqs].max()
282 query_start_loc = common_attn_metadata.query_start_loc
283 seq_lens = common_attn_metadata.seq_lens
@@ -332,9 +331,6 @@ class FlexAttentionMetadataBuilder:
331 )
332 return out
333
335 - def use_cascade_attention(self, *args, **kwargs) -> bool:
336 - return False
337 -
334
335 class FlexAttentionImpl(AttentionImpl):
336 sliding_window: Optional[tuple[int, int]]
vllm/v1/attention/backends/mla/common.py
+36 -8
@@ -207,7 +207,8 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
207 UnquantizedLinearMethod)
208 from vllm.platforms import current_platform
209 from vllm.utils import cdiv, round_down
210 -from vllm.v1.attention.backends.utils import CommonAttentionMetadata
210 +from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
211 + CommonAttentionMetadata)
212 from vllm.v1.kv_cache_interface import AttentionSpec
213 from vllm.v1.worker.block_table import BlockTable
214
@@ -329,7 +330,7 @@ class MLACommonMetadata(Generic[D]):
330 M = TypeVar("M", bound=MLACommonMetadata)
331
332
332 -class MLACommonMetadataBuilder(Generic[M]):
333 +class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
334 """
335 NOTE: Please read the comment at the top of the file before trying to
336 understand this class
@@ -450,9 +451,32 @@ class MLACommonMetadataBuilder(Generic[M]):
451 seq_lens=seq_lens,
452 )
453
453 - def build(self, num_reqs: int, num_actual_tokens: int, max_query_len: int,
454 - common_prefix_len: int,
454 + def build_for_cudagraph_capture(
455 + self, common_attn_metadata: CommonAttentionMetadata) -> M:
456 + """
457 + This method builds the metadata for full cudagraph capture.
458 + Currently, only decode is supported for full cudagraphs with MLA.
459 + """
460 + m = common_attn_metadata
461 + assert m.num_reqs == m.num_actual_tokens, \
462 + "MLA only supports decode-only full CUDAGraph capture. " \
463 + "Make sure all cudagraph capture sizes <= max_num_seq."
464 +
465 + m.max_query_len = 1 # decode-only
466 +
467 + # Update state usually set in reorder_batch.
468 + self._num_decodes = m.num_reqs
469 + self._num_decode_tokens = m.num_actual_tokens
470 + self._num_prefills = 0
471 + self._num_prefill_tokens = 0
472 + return self.build(0, m)
473 +
474 + def build(self, common_prefix_len: int,
475 common_attn_metadata: CommonAttentionMetadata) -> M:
476 + num_reqs = common_attn_metadata.num_reqs
477 + num_actual_tokens = common_attn_metadata.num_actual_tokens
478 + max_query_len = common_attn_metadata.max_query_len
479 +
480 assert self._num_decodes + self._num_prefills == num_reqs
481
482 # Note(simon): be careful about the CPU <> GPU memory movement in this
@@ -461,8 +485,11 @@ class MLACommonMetadataBuilder(Generic[M]):
485 device = self.runner.device
486 block_table = self.block_table
487 block_table_tensor = block_table.get_device_tensor()[:num_reqs]
464 - slot_mapping = block_table.slot_mapping_cpu[:num_actual_tokens].to(
465 - device, non_blocking=True).long()
488 + block_table.slot_mapping[:num_actual_tokens].copy_(
489 + block_table.slot_mapping_cpu[:num_actual_tokens],
490 + non_blocking=True)
491 + block_table.slot_mapping[num_actual_tokens:].fill_(-1)
492 + slot_mapping = block_table.slot_mapping[:num_actual_tokens]
493
494 query_start_loc = common_attn_metadata.query_start_loc
495 seq_lens = common_attn_metadata.seq_lens
@@ -564,8 +591,9 @@ class MLACommonMetadataBuilder(Generic[M]):
591 decode=decode_metadata,
592 )
593
567 - def use_cascade_attention(self, *args, **kwargs) -> bool:
568 - return False
594 + def can_run_in_cudagraph(
595 + self, common_attn_metadata: CommonAttentionMetadata) -> bool:
596 + return common_attn_metadata.max_query_len == 1
597
598
599 class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
vllm/v1/attention/backends/mla/flashmla.py
+31 -3
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass
5 -from typing import Any, Optional
5 +from typing import Any, ClassVar, Optional
6
7 import torch
8
@@ -44,7 +44,7 @@ class FlashMLABackend(MLACommonBackend):
44
45 @dataclass
46 class FlashMLADecodeMetadata(MLACommonDecodeMetadata):
47 - tile_scheduler_metadata: tuple[torch.Tensor, torch.Tensor]
47 + tile_scheduler_metadata: torch.Tensor
48 num_splits: torch.Tensor
49
50
@@ -54,14 +54,18 @@ class FlashMLAMetadata(MLACommonMetadata[FlashMLADecodeMetadata]):
54
55
56 class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
57 + full_cudagraph_supported: ClassVar[bool] = True # Decode-only
58
59 def __init__(self, runner, kv_cache_spec: AttentionSpec,
60 block_table: BlockTable):
60 - super().__init__(runner, kv_cache_spec, block_table)
61 + super().__init__(runner, kv_cache_spec, block_table, FlashMLAMetadata)
62
63 self.num_q_heads = self.runner.model_config.get_num_attention_heads(
64 self.runner.parallel_config)
65
66 + self.cg_buf_tile_scheduler_metadata = None
67 + self.cg_buf_num_splits = None
68 +
69 def _build_decode(self, block_table_tensor: torch.Tensor,
70 seq_lens: torch.Tensor) -> FlashMLADecodeMetadata:
71 tile_scheduler_metadata, num_splits = \
@@ -71,6 +75,30 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
75 1, # MQA for the decode path
76 )
77
78 + if self.runner.full_cuda_graph:
79 + # First time around (CUDAGraph capture), allocate the static buffer
80 + if self.cg_buf_tile_scheduler_metadata is None:
81 + self.cg_buf_tile_scheduler_metadata = tile_scheduler_metadata
82 + self.cg_buf_num_splits = num_splits
83 + else:
84 + assert self.cg_buf_num_splits is not None
85 +
86 + # Metadata per-SM, fixed size (#SMs, TileMetadataSize)
87 + assert (self.cg_buf_tile_scheduler_metadata.size() ==
88 + tile_scheduler_metadata.size())
89 + self.cg_buf_tile_scheduler_metadata.\
90 + copy_(tile_scheduler_metadata)
91 + tile_scheduler_metadata = self.cg_buf_tile_scheduler_metadata
92 +
93 + # Num splits is per-batch, varying size (batch_size,)
94 + n = num_splits.size(0)
95 + # make sure static buffer is large enough
96 + assert n <= self.cg_buf_num_splits.size(0)
97 + num_splits_view = self.cg_buf_num_splits[:n]
98 + num_splits_view.copy_(num_splits)
99 + self.cg_buf_num_splits[n:].fill_(0) # fill the rest with 0s
100 + num_splits = num_splits_view
101 +
102 return FlashMLADecodeMetadata(
103 block_table=block_table_tensor,
104 seq_lens=seq_lens,
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+1 -1
@@ -66,7 +66,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
66
67 def __init__(self, runner, kv_cache_spec: AttentionSpec,
68 block_table: BlockTable):
69 - super().__init__(runner, kv_cache_spec, block_table)
69 + super().__init__(runner, kv_cache_spec, block_table, AiterMLAMetadata)
70 assert self.kv_cache_spec.block_size == 1, "AITER MLA" \
71 "only supports block size 1."
72
vllm/v1/attention/backends/utils.py
+71 -2
@@ -1,15 +1,23 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import abc
4 +from abc import abstractmethod
5 from dataclasses import dataclass
6 +from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar
7
8 +import numpy as np
9 import torch
10
11 +if TYPE_CHECKING:
12 + from vllm.v1.core.sched.output import SchedulerOutput
13 + from vllm.v1.worker.gpu_input_batch import InputBatch
14 +
15
16 @dataclass
17 class CommonAttentionMetadata:
18 """
11 - Attention metadata attributes that can be shared by layers in different KV
12 - cache groups and thus having different block table.
19 + Per-batch attention metadata, shared across layers and backends.
20 + AttentionMetadataBuilder instances use it to construct per-layer metadata.
21 """
22
23 query_start_loc: torch.Tensor
@@ -18,6 +26,67 @@ class CommonAttentionMetadata:
26 """(batch_size,), the length of each request including both computed tokens
27 and newly scheduled tokens"""
28
29 + num_reqs: int
30 + """Number of requests"""
31 + num_actual_tokens: int
32 + """Total number of tokens in batch"""
33 + max_query_len: int
34 + """Longest query in batch"""
35 +
36 +
37 +M = TypeVar("M")
38 +
39 +
40 +class AttentionMetadataBuilder(abc.ABC, Generic[M]):
41 + # Does this backend/builder support CUDA Graphs for attention.
42 + full_cudagraph_supported: ClassVar[bool] = False
43 +
44 + @abstractmethod
45 + def build(self, common_prefix_len: int,
46 + common_attn_metadata: CommonAttentionMetadata) -> M:
47 + """
48 + Central method that builds attention metadata.
49 + Some builders (MLA) require reorder_batch to be called prior to build.
50 + """
51 + raise NotImplementedError
52 +
53 + def can_run_in_cudagraph(
54 + self, common_attn_metadata: CommonAttentionMetadata) -> bool:
55 + """
56 + Can this batch (with given metadata) use CUDA Graphs for attention.
57 + """
58 + return False
59 +
60 + def build_for_cudagraph_capture(
61 + self, common_attn_metadata: CommonAttentionMetadata) -> M:
62 + """
63 + Build attention metadata for CUDA graph capture. Uses build by default.
64 + Subclasses that override this method should call self.build or
65 + super().build_for_cudagraph_capture.
66 + """
67 + return self.build(common_prefix_len=0,
68 + common_attn_metadata=common_attn_metadata)
69 +
70 + def use_cascade_attention(
71 + self,
72 + common_prefix_len: int,
73 + query_lens: np.ndarray,
74 + num_query_heads: int,
75 + num_kv_heads: int,
76 + use_alibi: bool,
77 + use_sliding_window: bool,
78 + num_sms: int,
79 + ) -> bool:
80 + return False
81 +
82 + def reorder_batch(self, input_batch: "InputBatch",
83 + scheduler_output: "SchedulerOutput") -> bool:
84 + """
85 + This method can reorder the batch if desired by the backend.
86 + :return: Has the batch been reordered (default False).
87 + """
88 + return False
89 +
90
91 def validate_kv_sharing_target(current_layer_name, target_layer_name,
92 static_forward_context):
vllm/v1/spec_decode/eagle.py
+6 -4
@@ -138,15 +138,17 @@ class EagleProposer:
138 max_query_len = query_lens.max().item()
139
140 common_attn_metadata = CommonAttentionMetadata(
141 - query_start_loc=cu_num_tokens, seq_lens=seq_lens)
141 + query_start_loc=cu_num_tokens,
142 + seq_lens=seq_lens,
143 + num_reqs=batch_size,
144 + num_actual_tokens=num_tokens,
145 + max_query_len=max_query_len,
146 + )
147
148 assert self.runner is not None
149
150 # FIXME: need to consider multiple kv_cache_groups
151 attn_metadata = self.runner.attn_metadata_builder.build(
147 - num_reqs=batch_size,
148 - num_actual_tokens=num_tokens,
149 - max_query_len=max_query_len,
152 common_prefix_len=0,
153 common_attn_metadata=common_attn_metadata,
154 )
vllm/v1/worker/gpu_model_runner.py
+93 -67
@@ -16,10 +16,8 @@ from tqdm import tqdm
16
17 import vllm.envs as envs
18 from vllm.attention import AttentionType, get_attn_backend
19 -from vllm.attention.backends.abstract import (AttentionBackend,
20 - AttentionMetadataBuilder)
19 +from vllm.attention.backends.abstract import AttentionBackend
20 from vllm.attention.layer import Attention
22 -from vllm.attention.utils.fa_utils import get_flash_attn_version
21 from vllm.config import (CompilationLevel, VllmConfig,
22 get_layers_from_vllm_config)
23 from vllm.distributed.kv_transfer import (get_kv_transfer_group,
@@ -41,7 +39,8 @@ from vllm.sequence import IntermediateTensors
39 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
40 GiB_bytes, LazyLoader, async_tensor_h2d, cdiv,
41 check_use_alibi, is_pin_memory_available)
44 -from vllm.v1.attention.backends.utils import CommonAttentionMetadata
42 +from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
43 + CommonAttentionMetadata)
44 from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
45 from vllm.v1.kv_cache_interface import (AttentionSpec, FullAttentionSpec,
46 KVCacheConfig, KVCacheSpec,
@@ -89,6 +88,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
88 self.vllm_config = vllm_config
89 self.model_config = vllm_config.model_config
90 self.cache_config = vllm_config.cache_config
91 + self.compilation_config = vllm_config.compilation_config
92 self.lora_config = vllm_config.lora_config
93 self.load_config = vllm_config.load_config
94 self.parallel_config = vllm_config.parallel_config
@@ -197,7 +197,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
197 block_sizes=[self.cache_config.block_size],
198 )
199
200 - self.use_cuda_graph = (self.vllm_config.compilation_config.level
200 + self.use_cuda_graph = (self.compilation_config.level
201 == CompilationLevel.PIECEWISE
202 and not self.model_config.enforce_eager)
203 # TODO(woosuk): Provide an option to tune the max cudagraph batch size.
@@ -205,8 +205,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
205 # self.cudagraph_batch_sizes sorts in ascending order.
206 # The batch sizes in the config are in descending order.
207 self.cudagraph_batch_sizes = list(
208 - reversed(
209 - self.vllm_config.compilation_config.cudagraph_capture_sizes))
208 + reversed(self.compilation_config.cudagraph_capture_sizes))
209 +
210 + self.full_cuda_graph = self.compilation_config.full_cuda_graph
211
212 # Cache the device properties.
213 self._init_device_properties()
@@ -555,7 +556,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
556 def _prepare_inputs(
557 self,
558 scheduler_output: "SchedulerOutput",
558 - ) -> tuple[dict[str, Any], torch.Tensor, Optional[SpecDecodeMetadata]]:
559 + ) -> tuple[dict[str, Any], bool, torch.Tensor,
560 + Optional[SpecDecodeMetadata]]:
561 + """
562 + :return: tuple[
563 + attn_metadata: layer-to-attention_metadata mapping,
564 + attention_cuda_graphs: whether attention can run in cudagraph
565 + logits_indices, spec_decode_metadata
566 + ]
567 + """
568 total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
569 assert total_num_scheduled_tokens > 0
570 num_reqs = self.input_batch.num_reqs
@@ -669,7 +678,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
678 seq_lens = self.seq_lens[:num_reqs]
679
680 common_attn_metadata = CommonAttentionMetadata(
672 - query_start_loc=query_start_loc, seq_lens=seq_lens)
681 + query_start_loc=query_start_loc,
682 + seq_lens=seq_lens,
683 + num_reqs=num_reqs,
684 + num_actual_tokens=total_num_scheduled_tokens,
685 + max_query_len=max_num_scheduled_tokens,
686 + )
687
688 attn_metadata: dict[str, Any] = {}
689 # Prepare the attention metadata for each KV cache group and make layers
@@ -679,25 +693,28 @@ class GPUModelRunner(LoRAModelRunnerMixin):
693
694 # Prepare for cascade attention if enabled & beneficial.
695 common_prefix_len = 0
696 + builder = self.attn_metadata_builders[kv_cache_group_id]
697 if self.cascade_attn_enabled:
698 common_prefix_len = self._compute_cascade_attn_prefix_len(
699 num_scheduled_tokens,
700 scheduler_output.
701 num_common_prefix_blocks[kv_cache_group_id],
702 kv_cache_group_spec.kv_cache_spec,
688 - self.attn_metadata_builders[kv_cache_group_id],
703 + builder,
704 )
705
691 - attn_metadata_i = (
692 - self.attn_metadata_builders[kv_cache_group_id].build(
693 - num_reqs=num_reqs,
694 - num_actual_tokens=total_num_scheduled_tokens,
695 - max_query_len=max_num_scheduled_tokens,
696 - common_prefix_len=common_prefix_len,
697 - common_attn_metadata=common_attn_metadata))
706 + attn_metadata_i = (builder.build(
707 + common_prefix_len=common_prefix_len,
708 + common_attn_metadata=common_attn_metadata,
709 + ))
710 +
711 for layer_name in kv_cache_group_spec.layer_names:
712 attn_metadata[layer_name] = attn_metadata_i
713
714 + attention_cuda_graphs = all(
715 + b.can_run_in_cudagraph(common_attn_metadata)
716 + for b in self.attn_metadata_builders)
717 +
718 use_spec_decode = len(
719 scheduler_output.scheduled_spec_decode_tokens) > 0
720 if not use_spec_decode:
@@ -726,7 +743,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
743 if self.lora_config:
744 self.set_active_loras(self.input_batch, num_scheduled_tokens)
745
729 - return attn_metadata, logits_indices, spec_decode_metadata
746 + return (attn_metadata, attention_cuda_graphs, logits_indices,
747 + spec_decode_metadata)
748
749 def _compute_cascade_attn_prefix_len(
750 self,
@@ -1121,7 +1139,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1139 assert self.intermediate_tensors is not None
1140
1141 tp = self.vllm_config.parallel_config.tensor_parallel_size
1124 - enabled_sp = self.vllm_config.compilation_config.pass_config. \
1142 + enabled_sp = self.compilation_config.pass_config. \
1143 enable_sequence_parallelism
1144 if enabled_sp:
1145 # When sequence parallelism is enabled, we always pad num_tokens
@@ -1189,8 +1207,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1207 return self.kv_connector_no_forward(scheduler_output)
1208
1209 # Prepare the decoder inputs.
1192 - attn_metadata, logits_indices, spec_decode_metadata = (
1193 - self._prepare_inputs(scheduler_output))
1210 + (attn_metadata, attention_cuda_graphs, logits_indices,
1211 + spec_decode_metadata) = (self._prepare_inputs(scheduler_output))
1212 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
1213 if (self.use_cuda_graph
1214 and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
@@ -1203,7 +1221,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1221 # Pad tokens to multiple of tensor_parallel_size when
1222 # enabled collective fusion for SP
1223 tp_size = self.vllm_config.parallel_config.tensor_parallel_size
1206 - if self.vllm_config.compilation_config.pass_config. \
1224 + if self.compilation_config.pass_config. \
1225 enable_sequence_parallelism and tp_size > 1:
1226 from vllm.utils import round_up
1227 num_input_tokens = round_up(num_scheduled_tokens, tp_size)
@@ -1255,12 +1273,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1273 intermediate_tensors = self.sync_and_slice_intermediate_tensors(
1274 num_input_tokens, intermediate_tensors, True)
1275
1276 + # Some attention backends only support CUDA Graphs in pure decode.
1277 + # If attention doesn't support CUDA Graphs for this batch, but we
1278 + # compiled with full CUDA graphs, we have to skip them entirely.
1279 + skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
1280 +
1281 # Run the decoder.
1282 # Use persistent buffers for CUDA graphs.
1260 - with set_forward_context(attn_metadata,
1261 - self.vllm_config,
1262 - num_tokens=num_input_tokens,
1263 - num_tokens_across_dp=num_tokens_across_dp):
1283 + with set_forward_context(
1284 + attn_metadata,
1285 + self.vllm_config,
1286 + num_tokens=num_input_tokens,
1287 + num_tokens_across_dp=num_tokens_across_dp,
1288 + skip_cuda_graphs=skip_cuda_graphs,
1289 + ):
1290 self.maybe_setup_kv_connector(scheduler_output)
1291
1292 model_output = self.model(
@@ -1769,7 +1795,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1795 def _dummy_run(
1796 self,
1797 num_tokens: int,
1772 - skip_attn: bool = True,
1798 + capture_attn_cudagraph: bool = False,
1799 ) -> torch.Tensor:
1800
1801 # Padding for DP
@@ -1790,9 +1816,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1816 num_scheduled_tokens = np.array(num_scheduled_tokens_list,
1817 dtype=np.int32)
1818
1793 - if skip_attn:
1794 - attn_metadata: Optional[dict[str, Any]] = None
1795 - else:
1819 + attn_metadata: Optional[dict[str, Any]] = None
1820 + if capture_attn_cudagraph:
1821 + attn_metadata = {}
1822 +
1823 query_start_loc = self.query_start_loc[:num_reqs + 1]
1824 # Make sure max_model_len is used at the graph capture time.
1825 self.seq_lens_np[:num_reqs] = self.max_model_len
@@ -1802,19 +1829,19 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1829 seq_lens = self.seq_lens[:num_reqs]
1830
1831 common_attn_metadata = CommonAttentionMetadata(
1805 - query_start_loc=query_start_loc, seq_lens=seq_lens)
1832 + query_start_loc=query_start_loc,
1833 + seq_lens=seq_lens,
1834 + num_reqs=num_reqs,
1835 + num_actual_tokens=num_tokens,
1836 + max_query_len=num_tokens,
1837 + )
1838
1807 - attn_metadata = {}
1839 for kv_cache_group_id, kv_cache_group_spec in enumerate(
1840 self.kv_cache_config.kv_cache_groups):
1810 - attn_metadata_i = (
1811 - self.attn_metadata_builders[kv_cache_group_id].build(
1812 - num_reqs=num_reqs,
1813 - num_actual_tokens=num_tokens,
1814 - max_query_len=num_tokens,
1815 - common_prefix_len=0,
1816 - common_attn_metadata=common_attn_metadata,
1817 - ))
1841 +
1842 + attn_metadata_i = self.attn_metadata_builders[
1843 + kv_cache_group_id].build_for_cudagraph_capture(
1844 + common_attn_metadata)
1845 for layer_name in kv_cache_group_spec.layer_names:
1846 attn_metadata[layer_name] = attn_metadata_i
1847
@@ -2039,14 +2066,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2066 # Capture the large shapes first so that the smaller shapes
2067 # can reuse the memory pool allocated for the large shapes.
2068 with graph_capture(device=self.device):
2042 - skip_attn = not self.vllm_config.compilation_config.full_cuda_graph
2069 + full_cg = self.full_cuda_graph
2070 for num_tokens in tqdm(reversed(self.cudagraph_batch_sizes),
2071 desc="Capturing CUDA graphs",
2072 total=len(self.cudagraph_batch_sizes)):
2046 - for _ in range(self.vllm_config.compilation_config.
2047 - cudagraph_num_of_warmups):
2048 - self._dummy_run(num_tokens, skip_attn=skip_attn)
2049 - self._dummy_run(num_tokens, skip_attn=skip_attn)
2073 + for _ in range(
2074 + self.compilation_config.cudagraph_num_of_warmups):
2075 + self._dummy_run(num_tokens, capture_attn_cudagraph=full_cg)
2076 + self._dummy_run(num_tokens, capture_attn_cudagraph=full_cg)
2077
2078 end_time = time.perf_counter()
2079 end_free_gpu_memory = torch.cuda.mem_get_info()[0]
@@ -2089,20 +2116,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2116 "Non-Attention backend is not supported by V1 "
2117 "GPUModelRunner.")
2118
2092 - if self.vllm_config.compilation_config.full_cuda_graph:
2093 - attn_backend_name = attn_backend_i.__name__
2094 - flash_attn_version = get_flash_attn_version()
2095 - if attn_backend_name != "FlashAttentionBackend" or \
2096 - flash_attn_version != 3:
2097 - raise ValueError(
2098 - f"full_cuda_graph is only supported with "
2099 - f"FA3. Current attention backend is "
2100 - f"{attn_backend_name}, FlashAttention version is "
2101 - f"{flash_attn_version}.")
2102 -
2119 block_table_i = self.input_batch.block_table[i]
2120 attn_metadata_builder_i = attn_backend_i.get_builder_cls()(
2105 - weakref.proxy(self), kv_cache_spec, block_table_i)
2121 + weakref.proxy(self),
2122 + kv_cache_spec,
2123 + block_table_i,
2124 + )
2125 +
2126 + if (self.full_cuda_graph
2127 + and not attn_metadata_builder_i.full_cudagraph_supported):
2128 + raise ValueError(
2129 + f"Full CUDAGraph not supported for "
2130 + f"{attn_backend_i.__name__}. Turn off CompilationConfig."
2131 + f"full_cuda_graph or use a different attention backend.")
2132 +
2133 self.attn_backends.append(attn_backend_i)
2134 self.attn_metadata_builders.append(attn_metadata_builder_i)
2135
@@ -2142,9 +2169,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2169 to be reshaped to the desired shape before being used by the models.
2170
2171 Args:
2145 - kv_cache_config: The KV cache config
2172 + kv_cache_config: The KV cache config
2173 Returns:
2147 - dict[str, torch.Tensor]: A map between layer names to their
2174 + dict[str, torch.Tensor]: A map between layer names to their
2175 corresponding memory buffer for KV cache.
2176 """
2177 kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
@@ -2171,11 +2198,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2198 Reshape the KV cache tensors to the desired shape and dtype.
2199
2200 Args:
2174 - kv_cache_config: The KV cache config
2175 - kv_cache_raw_tensors: The KV cache buffer of each layer, with
2201 + kv_cache_config: The KV cache config
2202 + kv_cache_raw_tensors: The KV cache buffer of each layer, with
2203 correct size but uninitialized shape.
2204 Returns:
2178 - Dict[str, torch.Tensor]: A map between layer names to their
2205 + Dict[str, torch.Tensor]: A map between layer names to their
2206 corresponding memory buffer for KV cache.
2207 """
2208 kv_caches: dict[str, torch.Tensor] = {}
@@ -2227,7 +2254,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2254 Args:
2255 kv_cache_config: The KV cache config
2256 Returns:
2230 - Dict[str, torch.Tensor]: A map between layer names to their
2257 + Dict[str, torch.Tensor]: A map between layer names to their
2258 corresponding memory buffer for KV cache.
2259 """
2260 # Initialize the memory buffer for KV cache
@@ -2245,10 +2272,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2272 kv_caches,
2273 )
2274
2248 - bind_kv_cache(
2249 - kv_caches,
2250 - self.vllm_config.compilation_config.static_forward_context,
2251 - self.kv_caches)
2275 + bind_kv_cache(kv_caches,
2276 + self.compilation_config.static_forward_context,
2277 + self.kv_caches)
2278 return kv_caches
2279
2280 def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: