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
[Core] Allow full cudagraph with separate attention routines and orthogonal to compilation, add support for FA2 and FlashInfer (#20059)
Signed-off-by: fhl <2410591650@qq.com> Signed-off-by: fhl2000 <63384265+fhl2000@users.noreply.github.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
fhl2000 committed
Aug 15, 2025 at 22:01 UTC
74f441f4b517a895ad12afd314a6f40caf657c4e
34 files changed
+1840
-598
tests/compile/piecewise/test_full_cudagraph.py
+125
-128
@@ -3,7 +3,8 @@
3
import contextlib
4
import os
5
import weakref
6
-from contextlib import ExitStack
6
+from dataclasses import dataclass
7
+from typing import Optional
8
9
import pytest
10
@@ -32,69 +33,133 @@ def temporary_environ(env_vars):
33
os.environ[k] = v
34
35
35
-@pytest.fixture(scope="class")
36
-def llm_pair(request):
37
- model = request.param
38
-
39
- with temporary_environ({
40
- "VLLM_USE_V1": "1",
41
- "VLLM_FLASH_ATTN_VERSION": "3"
42
- }):
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.fixture(scope="class")
70
-def cutlass_mla_llm_pair(request):
71
- model = request.param
72
-
73
- # force V1 engine and Cutlass MLA backend
74
- with temporary_environ({
36
+@dataclass
37
+class BackendConfig:
38
+ name: str
39
+ env_vars: dict
40
+ comp_config: dict
41
+ specific_gpu_arch: Optional[tuple] = None
42
+
43
+
44
+# Define all backend configurations of full cudagraph to be tested
45
+backend_configs = {
46
+ # FA3 on Hopper
47
+ "FA3":
48
+ BackendConfig(name="FA3",
49
+ env_vars={"VLLM_FLASH_ATTN_VERSION": "3"},
50
+ comp_config={
51
+ "cudagraph_mode": "FULL",
52
+ },
53
+ specific_gpu_arch=(9, 0)),
54
+ # FlashMLA on Hopper
55
+ "FlashMLA":
56
+ BackendConfig(name="FlashMLA",
57
+ env_vars={
58
+ "VLLM_ATTENTION_BACKEND": "FLASHMLA",
59
+ },
60
+ comp_config={
61
+ "cudagraph_mode": "FULL_AND_PIECEWISE",
62
+ },
63
+ specific_gpu_arch=(9, 0)),
64
+ # Cutlass MLA on Blackwell
65
+ "CutlassMLA":
66
+ BackendConfig(
67
+ name="CutlassMLA",
68
+ env_vars={
69
"VLLM_USE_V1": "1",
70
"VLLM_ATTENTION_BACKEND": "CUTLASS_MLA",
71
"FORCE_NUM_KV_SPLITS":
72
"1", # TODO: remove this when hang issue is fixed
79
- }):
73
+ },
74
+ comp_config={
75
+ "cudagraph_mode": "FULL_AND_PIECEWISE",
76
+ "cudagraph_capture_sizes": [16, 32, 64, 128, 256, 512],
77
+ },
78
+ specific_gpu_arch=(10, 0)),
79
+ # FA2
80
+ "FA2":
81
+ BackendConfig(name="FA2",
82
+ env_vars={"VLLM_FLASH_ATTN_VERSION": "2"},
83
+ comp_config={
84
+ "cudagraph_mode": "FULL",
85
+ }),
86
+ # Triton Attention
87
+ "TritonAttn":
88
+ BackendConfig(name="TritonAttn",
89
+ env_vars={"VLLM_ATTENTION_BACKEND": "TRITON_ATTN_VLLM_V1"},
90
+ comp_config={
91
+ "cudagraph_mode": "FULL",
92
+ }),
93
+ # FlashInfer
94
+ "FlashInfer":
95
+ BackendConfig(name="FlashInfer",
96
+ env_vars={"VLLM_ATTENTION_BACKEND": "FLASHINFER"},
97
+ comp_config={
98
+ "cudagraph_mode": "FULL_AND_PIECEWISE",
99
+ }),
100
+}
101
+
102
+test_params_full_cudagraph = []
103
+
104
+# deepseek-ai/DeepSeek-V2-Lite with MLA
105
+MLA_backends = ["FlashMLA", "CutlassMLA"]
106
+for mla_backend in MLA_backends:
107
+ test_params_full_cudagraph.append(
108
+ pytest.param(
109
+ ("deepseek-ai/DeepSeek-V2-Lite", backend_configs[mla_backend])))
110
+
111
+# Qwen/Qwen2-1.5B-Instruct with other backends
112
+other_backend_configs = [
113
+ backend_configs[c] for c in backend_configs if c not in MLA_backends
114
+]
115
+for backend_config in other_backend_configs:
116
+ test_params_full_cudagraph.append(
117
+ pytest.param(("Qwen/Qwen2-1.5B-Instruct", backend_config)))
118
+
119
+
120
+@pytest.fixture(scope="class")
121
+def llm_pair(request):
122
+ model, backend_config = request.param
123
+
124
+ # Dynamically skip test if GPU capability is not met
125
+ if backend_config.specific_gpu_arch and backend_config.specific_gpu_arch\
126
+ != current_platform.get_device_capability():
127
+ if backend_config.specific_gpu_arch == (9, 0):
128
+ pytest.skip("Only Hopper GPUs support FA3 and FlashMLA")
129
+ elif backend_config.specific_gpu_arch == (10, 0):
130
+ pytest.skip("Only Blackwell GPUs support Cutlass MLA")
131
+
132
+ env_vars = {
133
+ "VLLM_USE_V1": "1",
134
+ # Force native sampler to avoid potential nondeterminism in FlashInfer
135
+ # when per-request generators are not used in V1.
136
+ "VLLM_USE_FLASHINFER_SAMPLER": "0",
137
+ **backend_config.env_vars,
138
+ }
139
+ with temporary_environ(env_vars):
140
full = LLM(
141
model=model,
82
- gpu_memory_utilization=0.45,
142
+ gpu_memory_utilization=0.43,
143
trust_remote_code=True,
144
max_model_len=1024,
85
- compilation_config=CompilationConfig(
86
- full_cuda_graph=True,
87
- cudagraph_capture_sizes=[16, 32, 64, 128, 256, 512],
88
- ),
145
+ max_num_seqs=128,
146
+ compilation_config=\
147
+ CompilationConfig(**backend_config.comp_config),
148
+ generation_config="vllm",
149
+ seed=42,
150
)
151
piecewise = LLM(
152
model=model,
92
- gpu_memory_utilization=0.45,
153
+ gpu_memory_utilization=0.43,
154
trust_remote_code=True,
155
max_model_len=1024,
95
- compilation_config=CompilationConfig(),
156
+ max_num_seqs=128,
157
+ compilation_config=CompilationConfig(cudagraph_mode="PIECEWISE"),
158
+ generation_config="vllm",
159
+ seed=42,
160
)
161
162
+ # PyTest caches the fixture values so we use weakref.proxy to enable GC
163
yield weakref.proxy(full), weakref.proxy(piecewise)
164
del full
165
del piecewise
@@ -105,51 +170,7 @@ def cutlass_mla_llm_pair(request):
170
)
171
172
108
-@pytest.mark.parametrize(
109
- "cutlass_mla_llm_pair",
110
- [
111
- # use an MLA model
112
- "deepseek-ai/DeepSeek-V2-Lite",
113
- ],
114
- indirect=True)
115
-@pytest.mark.skipif(current_platform.get_device_capability() != (10, 0),
116
- reason="Only Blackwell GPUs support Cutlass MLA")
117
-class TestFullCUDAGraphCutlassMLA:
118
- """
119
- Validate full CUDA Graph with Cutlass MLA (decode-only capture).
120
- """
121
-
122
- @pytest.mark.parametrize(("batch_size", "max_tokens"), [
123
- (8, 8),
124
- ])
125
- def test_full_cudagraph_sm100_cutlass_mla(
126
- self, batch_size, max_tokens, cutlass_mla_llm_pair: tuple[LLM,
127
- LLM]):
128
- piecewise_llm, full_cudagraph_llm = cutlass_mla_llm_pair
129
-
130
- prompts = ["Hello, my name is"] * batch_size
131
- sampling_params = SamplingParams(temperature=0.0,
132
- max_tokens=max_tokens,
133
- top_p=0.95)
134
-
135
- piecewise_responses = piecewise_llm.generate(prompts, sampling_params)
136
- full_responses = full_cudagraph_llm.generate(prompts, sampling_params)
137
-
138
- for piecewise_res, full_res in zip(piecewise_responses,
139
- full_responses):
140
- assert piecewise_res.outputs[0].text == full_res.outputs[0].text
141
-
142
-
143
-@pytest.mark.parametrize(
144
- "llm_pair",
145
- [
146
- # Model names for the llm_pair fixture
147
- "deepseek-ai/DeepSeek-V2-Lite",
148
- "Qwen/Qwen2-1.5B-Instruct"
149
- ],
150
- indirect=True)
151
-@pytest.mark.skipif(current_platform.get_device_capability() != (9, 0),
152
- reason="Only Hopper GPUs support FA3 and FlashMLA")
173
+@pytest.mark.parametrize("llm_pair", test_params_full_cudagraph, indirect=True)
174
class TestFullCUDAGraph:
175
"""
176
Use a class such that an llm pair is constructed once for all
@@ -178,12 +199,14 @@ class TestFullCUDAGraph:
199
full cudagraph compilation works for padded cases too.
200
"""
201
181
- piecewise_llm, full_cudagraph_llm = llm_pair
202
+ full_cudagraph_llm, piecewise_llm = llm_pair
203
183
- prompts = ["Hello, my name is"] * batch_size
204
+ prompts = ["the quick brown fox"] * batch_size
205
+ # Use purely greedy decoding to avoid top-p truncation sensitivity
206
+ # that can amplify tiny numeric differences across runtimes.
207
sampling_params = SamplingParams(temperature=0.0,
208
max_tokens=max_tokens,
186
- top_p=0.95)
209
+ top_p=1.0)
210
211
piecewise_responses = piecewise_llm.generate(prompts, sampling_params)
212
full_responses = full_cudagraph_llm.generate(prompts, sampling_params)
@@ -191,42 +214,16 @@ class TestFullCUDAGraph:
214
# Check that all responses are the same
215
for piecewise_res, full_res in zip(piecewise_responses,
216
full_responses):
194
- assert piecewise_res.outputs[0].text == full_res.outputs[0].text
195
-
196
-
197
-@pytest.mark.parametrize(
198
- "model, supported",
199
- [
200
- ("Qwen/Qwen2-1.5B-Instruct", True),
201
- # MLA does not support capturing CUDA Graphs with size > max_num_seqs
202
- ("deepseek-ai/DeepSeek-V2-Lite", False),
203
- ])
204
-@pytest.mark.skipif(current_platform.get_device_capability() != (9, 0),
205
- reason="Only Hopper GPUs support FA3 and FlashMLA")
206
-def test_lower_max_num_seqs(model, supported):
207
- with temporary_environ({
208
- "VLLM_USE_V1": "1",
209
- "VLLM_FLASH_ATTN_VERSION": "3"
210
- }), ExitStack() as stack:
211
- if not supported:
212
- stack.enter_context(pytest.raises(RuntimeError))
213
-
214
- llm = LLM(model=model,
215
- max_num_seqs=256,
216
- trust_remote_code=True,
217
- max_model_len=1024,
218
- compilation_config=CompilationConfig(
219
- full_cuda_graph=True,
220
- cudagraph_capture_sizes=[64, 256, 512]))
221
- llm.generate(["Hello, my name is"] * 10)
217
+ assert piecewise_res.outputs[0].text.lower() == \
218
+ full_res.outputs[0].text.lower()
219
220
221
@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
222
def test_full_cudagraph_with_invalid_backend():
223
with temporary_environ({
224
"VLLM_USE_V1": "1",
228
- "VLLM_FLASH_ATTN_VERSION":
229
- "2" #FA2 not supported with full_cuda_graph
225
+ "VLLM_ATTENTION_BACKEND": "FLEX_ATTENTION"
226
+ # Flex_Attention is not supported with full cuda graph
227
}), pytest.raises(RuntimeError):
228
LLM(model="Qwen/Qwen2-1.5B-Instruct",
232
- compilation_config=CompilationConfig(full_cuda_graph=True))
229
+ compilation_config=CompilationConfig(cudagraph_mode="FULL"))
tests/compile/piecewise/test_simple.py
+25
-8
@@ -11,10 +11,10 @@ from torch.library import Library
11
12
from vllm.compilation.counter import compilation_counter
13
from vllm.compilation.decorators import support_torch_compile
14
-from vllm.config import (CompilationConfig, CompilationLevel, VllmConfig,
15
- set_current_vllm_config)
14
+from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,
15
+ VllmConfig, set_current_vllm_config)
16
from vllm.envs import VLLM_USE_V1
17
-from vllm.forward_context import set_forward_context
17
+from vllm.forward_context import BatchDescriptor, set_forward_context
18
from vllm.utils import direct_register_custom_op
19
20
global_counter = 0
@@ -101,16 +101,33 @@ 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
104
- ), set_forward_context({}, vllm_config=vllm_config):
105
-
104
+ ), set_forward_context(None,
105
+ vllm_config=vllm_config): # background context
106
+ # warm up with background context
107
model(inputs)
108
108
- model(torch.randn(2).cuda())
109
- model(torch.randn(1).cuda())
109
+ # capturing/replaying should under context of cudagraph dispatching
110
+ with set_forward_context(
111
+ None,
112
+ vllm_config=vllm_config,
113
+ cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
114
+ batch_descriptor=BatchDescriptor(num_tokens=2, )):
115
+ model(torch.randn(2).cuda())
116
+ with set_forward_context(
117
+ None,
118
+ vllm_config=vllm_config,
119
+ cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
120
+ batch_descriptor=BatchDescriptor(num_tokens=1, )):
121
+ model(torch.randn(1).cuda())
122
123
input = torch.zeros(2).cuda()
124
global global_counter
125
global_counter = 0
114
- output = model(input)
126
+ with set_forward_context(
127
+ None,
128
+ vllm_config=vllm_config,
129
+ cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
130
+ batch_descriptor=BatchDescriptor(num_tokens=2, )):
131
+ output = model(input)
132
assert global_counter == 2
133
assert torch.allclose(output.cpu(), torch.tensor([3., 1.]))
tests/compile/piecewise/test_toy_llama.py
+29
-7
@@ -18,9 +18,9 @@ from torch.library import Library
18
19
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
21
+from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,
22
+ VllmConfig, set_current_vllm_config)
23
+from vllm.forward_context import BatchDescriptor, set_forward_context
24
from vllm.utils import direct_register_custom_op
25
26
# create a library to hold the custom op
@@ -276,9 +276,11 @@ def run_model(llama_config,
276
)
277
if split_attn:
278
compilation_config.splitting_ops = ["silly.attention"]
279
+ cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
280
else:
281
compilation_config = CompilationConfig(
282
level=CompilationLevel.NO_COMPILATION, )
283
+ cudagraph_runtime_mode = CUDAGraphMode.NONE
284
285
vllm_config = VllmConfig(compilation_config=compilation_config,
286
additional_config=llama_config)
@@ -287,17 +289,37 @@ def run_model(llama_config,
289
vllm_config=vllm_config,
290
prefix="").eval().cuda()
291
290
- with set_forward_context({}, vllm_config=vllm_config):
292
+ with set_forward_context({},
293
+ vllm_config=vllm_config): # background context
294
B = 16 # max batch size
295
input_ids = torch.randint(0, llama_config.vocab_size, (B, )).cuda()
296
positions = torch.arange(B).cuda()
297
298
+ # warmup for the model with cudagraph_mode NONE
299
model(input_ids, positions)
296
- model(input_ids[:2], positions[:2])
297
- model(input_ids[:1], positions[:1])
300
+
301
+ # simulate cudagraphs capturing
302
+ with set_forward_context({},
303
+ vllm_config=vllm_config,
304
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
305
+ batch_descriptor=BatchDescriptor(
306
+ num_tokens=2, )):
307
+ model(input_ids[:2], positions[:2])
308
+ with set_forward_context({},
309
+ vllm_config=vllm_config,
310
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
311
+ batch_descriptor=BatchDescriptor(
312
+ num_tokens=1, )):
313
+ model(input_ids[:1], positions[:1])
314
315
input_ids[:2].zero_()
300
- output = model(input_ids[:2], positions[:2])
316
+ # simulate cudagraphs replay
317
+ with set_forward_context({},
318
+ vllm_config=vllm_config,
319
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
320
+ batch_descriptor=BatchDescriptor(
321
+ num_tokens=2, )):
322
+ output = model(input_ids[:2], positions[:2])
323
324
output = output.cpu()
325
tests/v1/cudagraph/__init__.py
tests/v1/cudagraph/test_cudagraph_dispatch.py
new
+406
@@ -0,0 +1,406 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from unittest.mock import MagicMock, patch
4
+
5
+import pytest
6
+import torch
7
+import torch.nn as nn
8
+
9
+from tests.utils import create_new_process_for_each_test
10
+from vllm.compilation.cuda_graph import CUDAGraphWrapper
11
+from vllm.compilation.monitor import set_cudagraph_capturing_enabled
12
+from vllm.config import (CompilationConfig, CompilationLevel, CUDAGraphMode,
13
+ ParallelConfig, SchedulerConfig, VllmConfig)
14
+from vllm.forward_context import BatchDescriptor, set_forward_context
15
+from vllm.platforms import current_platform
16
+from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
17
+
18
+
19
+# Helper MLP for testing
20
+class SimpleMLP(nn.Module):
21
+
22
+ def __init__(self):
23
+ super().__init__()
24
+ self.fc1 = nn.Linear(10, 10)
25
+ self.fc2 = nn.Linear(10, 10)
26
+
27
+ def forward(self, x):
28
+ return self.fc2(self.fc1(x))
29
+
30
+
31
+def _create_vllm_config(compilation_config: CompilationConfig,
32
+ max_num_seqs: int = 8) -> MagicMock:
33
+ mock_config = MagicMock(spec=VllmConfig)
34
+ mock_config.compilation_config = compilation_config
35
+ mock_config.scheduler_config = SchedulerConfig(max_num_seqs=max_num_seqs)
36
+ mock_config.parallel_config = ParallelConfig()
37
+
38
+ # Mimic the behavior of VllmConfig.__post_init__()
39
+ if compilation_config.level == CompilationLevel.PIECEWISE:
40
+ compilation_config.set_splitting_ops_for_v1()
41
+
42
+ return mock_config
43
+
44
+
45
+class TestCudagraphDispatcher:
46
+
47
+ @pytest.mark.parametrize(
48
+ "params",
49
+ [
50
+ # Test case 0: Full CG for mixed batches, no separate routine
51
+ {
52
+ "case_id": 0,
53
+ "cudagraph_mode": "FULL",
54
+ "compilation_level": CompilationLevel.NO_COMPILATION,
55
+ },
56
+ # Test case 1: Full CG for uniform batches, piecewise for mixed
57
+ {
58
+ "case_id": 1,
59
+ "cudagraph_mode": "FULL_AND_PIECEWISE",
60
+ "compilation_level": CompilationLevel.PIECEWISE,
61
+ },
62
+ # Test case 2: Full CG for uniform batches, no CG for mixed
63
+ {
64
+ "case_id": 2,
65
+ "cudagraph_mode": "FULL_DECODE_ONLY",
66
+ "compilation_level": CompilationLevel.NO_COMPILATION,
67
+ },
68
+ # Test case 3: Piecewise for all
69
+ {
70
+ "case_id": 3,
71
+ "cudagraph_mode": "PIECEWISE",
72
+ "compilation_level": CompilationLevel.PIECEWISE,
73
+ },
74
+ ])
75
+ def test_dispatcher(self, params):
76
+ # Setup dispatcher
77
+ comp_config = CompilationConfig(
78
+ cudagraph_mode=params["cudagraph_mode"],
79
+ level=params["compilation_level"],
80
+ cudagraph_capture_sizes=[1, 8])
81
+
82
+ config = _create_vllm_config(comp_config, max_num_seqs=8)
83
+ dispatcher = CudagraphDispatcher(config)
84
+ dispatcher.initialize_cudagraph_keys(
85
+ cudagraph_mode=comp_config.cudagraph_mode,
86
+ uniform_decode_query_len=1)
87
+
88
+ # Verify the key is initialized correctly
89
+ if params["cudagraph_mode"] in ["FULL_AND_PIECEWISE", "PIECEWISE"]:
90
+ assert len(dispatcher.cudagraph_keys[CUDAGraphMode.PIECEWISE]) == 2
91
+ else:
92
+ assert len(dispatcher.cudagraph_keys[CUDAGraphMode.PIECEWISE]) == 0
93
+ if params["cudagraph_mode"] not in ["NONE", "PIECEWISE"]:
94
+ assert len(dispatcher.cudagraph_keys[CUDAGraphMode.FULL]) == 2
95
+ else:
96
+ assert len(dispatcher.cudagraph_keys[CUDAGraphMode.FULL]) == 0
97
+
98
+ # Test dispatch logic
99
+ # 1. non-uniform batch, size in cudagraph size list
100
+ desc_full_exact = BatchDescriptor(num_tokens=8, uniform_decode=False)
101
+ rt_mode, key = dispatcher.dispatch(desc_full_exact)
102
+ if params["cudagraph_mode"] == "FULL":
103
+ assert rt_mode == CUDAGraphMode.FULL
104
+ assert key == desc_full_exact
105
+ elif params["cudagraph_mode"] in ["FULL_AND_PIECEWISE", "PIECEWISE"]:
106
+ assert rt_mode == CUDAGraphMode.PIECEWISE
107
+ assert key == desc_full_exact
108
+ else:
109
+ assert rt_mode == CUDAGraphMode.NONE
110
+
111
+ # 2. uniform decode batch, size in cudagraph size list
112
+ desc_uniform_exact = BatchDescriptor(num_tokens=8, uniform_decode=True)
113
+ rt_mode, key = dispatcher.dispatch(desc_uniform_exact)
114
+ if params["cudagraph_mode"] == "FULL":
115
+ assert rt_mode == CUDAGraphMode.FULL
116
+ assert key == desc_uniform_exact.non_uniform
117
+ elif params["cudagraph_mode"] in [
118
+ "FULL_DECODE_ONLY", "FULL_AND_PIECEWISE"
119
+ ]:
120
+ assert rt_mode == CUDAGraphMode.FULL
121
+ assert key == desc_uniform_exact
122
+ elif params["cudagraph_mode"] == "PIECEWISE":
123
+ assert rt_mode == CUDAGraphMode.PIECEWISE
124
+ assert key == desc_uniform_exact.non_uniform
125
+ else:
126
+ assert rt_mode == CUDAGraphMode.NONE
127
+
128
+ # 3. No key match
129
+ desc_no_match = BatchDescriptor(num_tokens=15, uniform_decode=False)
130
+ rt_mode, key = dispatcher.dispatch(desc_no_match)
131
+ assert rt_mode == CUDAGraphMode.NONE
132
+ assert key is None
133
+
134
+
135
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
136
+class TestCUDAGraphWrapper:
137
+
138
+ def setup_method(self):
139
+ self.vllm_config = _create_vllm_config(CompilationConfig())
140
+ self.model = SimpleMLP().to("cuda")
141
+ self.persistent_input_buffer = torch.zeros(1, 10, device="cuda")
142
+ self.input_tensor = torch.randn(1, 10, device="cuda")
143
+
144
+ @create_new_process_for_each_test("spawn")
145
+ def test_capture_and_replay(self):
146
+ wrapper = CUDAGraphWrapper(self.model,
147
+ self.vllm_config,
148
+ runtime_mode=CUDAGraphMode.FULL)
149
+ batch_descriptor = BatchDescriptor(num_tokens=10)
150
+
151
+ # 0. global warmup
152
+ with set_forward_context(attn_metadata=None,
153
+ vllm_config=self.vllm_config,
154
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
155
+ batch_descriptor=None):
156
+ wrapper(self.input_tensor)
157
+
158
+ # 1. Capture
159
+ with set_forward_context(
160
+ attn_metadata=None,
161
+ vllm_config=self.vllm_config,
162
+ cudagraph_runtime_mode=CUDAGraphMode.FULL,
163
+ batch_descriptor=batch_descriptor),\
164
+ patch("torch.cuda.graph",
165
+ wraps=torch.cuda.graph) as mock_cuda_graph:
166
+ output1 = wrapper(self.input_tensor)
167
+ # capturing phase should generate a zero output
168
+ assert torch.allclose(output1, torch.zeros_like(output1))
169
+ mock_cuda_graph.assert_called_once()
170
+
171
+ assert batch_descriptor in wrapper.concrete_cudagraph_entries
172
+ entry = wrapper.concrete_cudagraph_entries[batch_descriptor]
173
+ assert entry.cudagraph is not None
174
+
175
+ # 2. Replay
176
+ with set_forward_context(
177
+ attn_metadata=None,
178
+ vllm_config=self.vllm_config,
179
+ cudagraph_runtime_mode=CUDAGraphMode.FULL,
180
+ batch_descriptor=batch_descriptor),\
181
+ patch.object(entry.cudagraph, 'replay',
182
+ wraps=entry.cudagraph.replay) as mock_replay:
183
+ output2 = wrapper(self.input_tensor)
184
+ mock_replay.assert_called_once()
185
+
186
+ # Compare with eager output
187
+ eager_output = self.model(self.input_tensor)
188
+ torch.testing.assert_close(eager_output, output2)
189
+
190
+ @create_new_process_for_each_test("spawn")
191
+ def test_bypass_on_mode_mismatch(self):
192
+ wrapper = CUDAGraphWrapper(self.model,
193
+ self.vllm_config,
194
+ runtime_mode=CUDAGraphMode.FULL)
195
+ batch_descriptor = BatchDescriptor(num_tokens=10)
196
+
197
+ with set_forward_context(
198
+ attn_metadata=None,
199
+ vllm_config=self.vllm_config,
200
+ cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
201
+ batch_descriptor=batch_descriptor), \
202
+ patch('torch.cuda.graph',
203
+ wraps=torch.cuda.graph) as mock_cuda_graph, \
204
+ patch.object(self.model, 'forward',
205
+ wraps=self.model.forward) as mock_forward:
206
+ wrapper(self.input_tensor)
207
+ mock_cuda_graph.assert_not_called()
208
+ mock_forward.assert_called_once()
209
+ assert not wrapper.concrete_cudagraph_entries
210
+
211
+ @create_new_process_for_each_test("spawn")
212
+ def test_bypass_on_mode_none(self):
213
+ wrapper = CUDAGraphWrapper(self.model,
214
+ self.vllm_config,
215
+ runtime_mode=CUDAGraphMode.FULL)
216
+ batch_descriptor = BatchDescriptor(num_tokens=10)
217
+
218
+ with set_forward_context(
219
+ attn_metadata=None,
220
+ vllm_config=self.vllm_config,
221
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
222
+ batch_descriptor=batch_descriptor), \
223
+ patch('torch.cuda.graph',
224
+ wraps=torch.cuda.graph) as mock_cuda_graph:
225
+ wrapper(self.input_tensor)
226
+ mock_cuda_graph.assert_not_called()
227
+ assert not wrapper.concrete_cudagraph_entries
228
+
229
+
230
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
231
+class TestCudagraphIntegration:
232
+
233
+ def setup_method(self):
234
+ # only FULL mode for non-uniform batches
235
+ self.comp_config = CompilationConfig(level=CompilationLevel.PIECEWISE,
236
+ cudagraph_mode="FULL",
237
+ cudagraph_capture_sizes=[10, 20])
238
+ self.vllm_config = _create_vllm_config(self.comp_config)
239
+ self.dispatcher = CudagraphDispatcher(self.vllm_config)
240
+ self.dispatcher.initialize_cudagraph_keys(
241
+ self.comp_config.cudagraph_mode, uniform_decode_query_len=1)
242
+
243
+ def _run_and_monitor_call(self, wrapper, input_tensor, runtime_mode,
244
+ batch_descriptor):
245
+ """Helper to run a single call and monitor the action."""
246
+
247
+ with patch('torch.cuda.graph',
248
+ wraps=torch.cuda.graph) as mock_graph_context, \
249
+ patch.object(wrapper, 'runnable',
250
+ wraps=wrapper.runnable) as mock_runnable:
251
+
252
+ entry = wrapper.concrete_cudagraph_entries.get(
253
+ batch_descriptor, None)
254
+
255
+ context = set_forward_context(attn_metadata=None,
256
+ vllm_config=self.vllm_config,
257
+ cudagraph_runtime_mode=runtime_mode,
258
+ batch_descriptor=batch_descriptor)
259
+ mock_replay = MagicMock()
260
+ if entry and entry.cudagraph:
261
+ with context, \
262
+ patch.object(entry.cudagraph, 'replay',
263
+ new_callable=MagicMock) as mock_replay:
264
+ wrapper(input_tensor)
265
+ else:
266
+ with context:
267
+ wrapper(input_tensor)
268
+
269
+ if mock_graph_context.called:
270
+ # note that this is globally mocked, so it will be detected
271
+ # even whether called by the inner or outer wrapper
272
+ return "capture_global"
273
+ if mock_replay.called:
274
+ # only for outer wrapper
275
+ return "replay"
276
+ if mock_runnable.call_count > 0:
277
+ # only for outer wrapper
278
+ return "bypass"
279
+ return "unknown"
280
+
281
+ @create_new_process_for_each_test("spawn")
282
+ def test_capture_replay_bypass_logic(self):
283
+ model = SimpleMLP().to("cuda")
284
+ full_wrapper = CUDAGraphWrapper(model, self.vllm_config,
285
+ CUDAGraphMode.FULL)
286
+ max_bs = 16
287
+ persistent_input_buffer = torch.zeros(max_bs, 10, device="cuda")
288
+ input_1 = persistent_input_buffer[:1]
289
+ input_2 = persistent_input_buffer[:2]
290
+ input_3 = persistent_input_buffer[:3]
291
+
292
+ desc_1 = BatchDescriptor(num_tokens=1)
293
+ desc_2 = BatchDescriptor(num_tokens=2)
294
+ desc_3_unseen = BatchDescriptor(num_tokens=3)
295
+
296
+ # 0. global warmup
297
+ with set_forward_context(attn_metadata=None,
298
+ vllm_config=self.vllm_config,
299
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
300
+ batch_descriptor=None):
301
+ full_wrapper(input_1)
302
+
303
+ rt_mode, key = self.dispatcher.dispatch(desc_1)
304
+ # 1. Capture first shape
305
+ action = self._run_and_monitor_call(full_wrapper, input_1, rt_mode,
306
+ key)
307
+ assert action == "capture_global"
308
+
309
+ # 2. Replay first shape
310
+ action = self._run_and_monitor_call(full_wrapper, input_1, rt_mode,
311
+ key)
312
+ assert action == "replay"
313
+
314
+ rt_mode, key = self.dispatcher.dispatch(desc_2)
315
+ # 3. Capture second shape
316
+ action = self._run_and_monitor_call(full_wrapper, input_2, rt_mode,
317
+ key)
318
+ assert action == "capture_global"
319
+
320
+ # 4. Replay second shape
321
+ action = self._run_and_monitor_call(full_wrapper, input_2,
322
+ CUDAGraphMode.FULL, desc_2)
323
+ assert action == "replay"
324
+
325
+ # 5. Bypass if no key match
326
+ rt_mode, key = self.dispatcher.dispatch(desc_3_unseen)
327
+ assert rt_mode == CUDAGraphMode.NONE
328
+ action = self._run_and_monitor_call(full_wrapper, input_3, rt_mode,
329
+ key)
330
+ assert action == "bypass"
331
+
332
+ # capture unseen shape is not allowed after disable
333
+ set_cudagraph_capturing_enabled(False)
334
+ with pytest.raises(RuntimeError):
335
+ self._run_and_monitor_call(full_wrapper, input_3,
336
+ CUDAGraphMode.FULL, desc_3_unseen)
337
+ set_cudagraph_capturing_enabled(True)
338
+
339
+ @create_new_process_for_each_test("spawn")
340
+ def test_nested_wrappers(self):
341
+ """Tests a scenario with a PIECEWISE wrapper inside a FULL one."""
342
+ model = SimpleMLP().to("cuda")
343
+ full_wrapper = CUDAGraphWrapper(model, self.vllm_config,
344
+ CUDAGraphMode.FULL)
345
+ input_1 = torch.randn(1, 10, device="cuda")
346
+
347
+ # Setup: Inner model is wrapped with PIECEWISE, outer with FULL
348
+ inner_model = SimpleMLP().to("cuda")
349
+ piecewise_wrapper = CUDAGraphWrapper(inner_model, self.vllm_config,
350
+ CUDAGraphMode.PIECEWISE)
351
+ inner_model.forward = MagicMock(wraps=inner_model.forward)
352
+ outer_model = SimpleMLP().to("cuda")
353
+ # When outer model is called, it calls the piecewise_wrapper
354
+ outer_model.forward = MagicMock(wraps=outer_model.forward,
355
+ side_effect=piecewise_wrapper)
356
+ full_wrapper = CUDAGraphWrapper(outer_model, self.vllm_config,
357
+ CUDAGraphMode.FULL)
358
+
359
+ desc_1 = BatchDescriptor(num_tokens=1)
360
+
361
+ # 0. global warmup
362
+ with set_forward_context(attn_metadata=None,
363
+ vllm_config=self.vllm_config,
364
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
365
+ batch_descriptor=None):
366
+ full_wrapper(input_1)
367
+
368
+ # --- Test runtime mode FULL---
369
+ # Run with FULL mode context. Expect outer wrapper to capture.
370
+ # The inner mock should be called once inside the graph capture.
371
+ outer_model.forward.reset_mock()
372
+ inner_model.forward.reset_mock()
373
+ action = self._run_and_monitor_call(full_wrapper, input_1,
374
+ CUDAGraphMode.FULL, desc_1)
375
+ assert action == "capture_global"
376
+ assert outer_model.forward.call_count == 1
377
+ assert inner_model.forward.call_count == 1
378
+
379
+ # Run again. Expect outer wrapper to replay.
380
+ # The outer model should NOT be called because the whole graph
381
+ # is replayed.
382
+ action = self._run_and_monitor_call(full_wrapper, input_1,
383
+ CUDAGraphMode.FULL, desc_1)
384
+ assert action == "replay"
385
+ assert outer_model.forward.call_count == 1 # No new call
386
+ assert inner_model.forward.call_count == 1
387
+
388
+ # --- Test runtime mode PIECEWISE ---
389
+ outer_model.forward.reset_mock()
390
+ inner_model.forward.reset_mock()
391
+ # Run with PIECEWISE mode context.
392
+ # Expect outer wrapper to bypass and call inner wrapper.
393
+ # Inner wrapper should capture.
394
+ action = self._run_and_monitor_call(full_wrapper, input_1,
395
+ CUDAGraphMode.PIECEWISE, desc_1)
396
+ assert action == "capture_global"
397
+ assert outer_model.forward.call_count == 1
398
+ assert inner_model.forward.call_count == 1
399
+
400
+ # Run again with PIECEWISE.
401
+ # Outer bypasses, inner replays.
402
+ action = self._run_and_monitor_call(full_wrapper, input_1,
403
+ CUDAGraphMode.PIECEWISE, desc_1)
404
+ assert action == "bypass"
405
+ assert outer_model.forward.call_count == 2
406
+ assert inner_model.forward.call_count == 1
tests/v1/cudagraph/test_cudagraph_mode.py
new
+187
@@ -0,0 +1,187 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import contextlib
4
+import os
5
+import weakref
6
+from contextlib import ExitStack
7
+from dataclasses import dataclass
8
+from typing import Optional
9
+
10
+import pytest
11
+
12
+from tests.utils import wait_for_gpu_memory_to_clear
13
+from vllm import LLM
14
+from vllm.config import CompilationConfig
15
+from vllm.platforms import current_platform
16
+
17
+
18
+@contextlib.contextmanager
19
+def temporary_environ(env_vars):
20
+ """
21
+ Temporarily set environment variables and restore them afterward.
22
+ We have to do this vs monkeypatch because monkeypatch doesn't work
23
+ with "module" scoped fixtures.
24
+ """
25
+ original_env = {k: os.environ.get(k) for k in env_vars}
26
+ try:
27
+ os.environ.update(env_vars)
28
+ yield
29
+ finally:
30
+ for k, v in original_env.items():
31
+ if v is None:
32
+ os.environ.pop(k, None)
33
+ else:
34
+ os.environ[k] = v
35
+
36
+
37
+@dataclass
38
+class BackendConfig:
39
+ name: str
40
+ env_vars: dict
41
+ comp_config: dict
42
+ specific_gpu_arch: Optional[tuple] = None
43
+
44
+
45
+# Define all backend configurations of full cudagraph to be tested
46
+backend_configs = {
47
+ # FA3 on Hopper
48
+ "FA3":
49
+ BackendConfig(name="FA3",
50
+ env_vars={"VLLM_FLASH_ATTN_VERSION": "3"},
51
+ comp_config={
52
+ "cudagraph_mode": "FULL",
53
+ },
54
+ specific_gpu_arch=(9, 0)),
55
+ # FlashMLA on Hopper
56
+ "FlashMLA":
57
+ BackendConfig(name="FlashMLA",
58
+ env_vars={
59
+ "VLLM_ATTENTION_BACKEND": "FLASHMLA",
60
+ },
61
+ comp_config={
62
+ "cudagraph_mode": "FULL_AND_PIECEWISE",
63
+ },
64
+ specific_gpu_arch=(9, 0)),
65
+ # FA2
66
+ "FA2":
67
+ BackendConfig(name="FA2",
68
+ env_vars={"VLLM_FLASH_ATTN_VERSION": "2"},
69
+ comp_config={
70
+ "cudagraph_mode": "FULL_AND_PIECEWISE",
71
+ }),
72
+ # Triton Attention
73
+ "TritonAttn":
74
+ BackendConfig(name="TritonAttn",
75
+ env_vars={"VLLM_ATTENTION_BACKEND": "TRITON_ATTN_VLLM_V1"},
76
+ comp_config={
77
+ "cudagraph_mode": "FULL_AND_PIECEWISE",
78
+ }),
79
+ # FlashInfer
80
+ "FlashInfer":
81
+ BackendConfig(name="FlashInfer",
82
+ env_vars={"VLLM_ATTENTION_BACKEND": "FLASHINFER"},
83
+ comp_config={
84
+ "cudagraph_mode": "FULL_AND_PIECEWISE",
85
+ }),
86
+}
87
+
88
+# test attention backend and cudagraph_mode combo
89
+# (backend_name, cudagraph_mode, supported)
90
+combo_cases_1 = [
91
+ ("FA3", "FULL", True),
92
+ ("FA3", "FULL_AND_PIECEWISE", True),
93
+ ("FA2", "FULL", True), # Should fallback to FULL_AND_PIECEWISE
94
+ ("FA2", "FULL_AND_PIECEWISE", True),
95
+ ("FlashInfer", "FULL", True), # Should fallback to FULL_AND_PIECEWISE
96
+ ("FlashInfer", "FULL_AND_PIECEWISE", True),
97
+]
98
+
99
+
100
+@pytest.mark.parametrize("combo_case", combo_cases_1)
101
+def test_backend_and_cudagraph_mode_combo(combo_case):
102
+ backend_name, cudagraph_mode, supported = combo_case
103
+ if backend_name == "FlashInfer":
104
+ try:
105
+ import flashinfer # noqa: F401
106
+ except ImportError:
107
+ pytest.skip("FlashInfer is not installed")
108
+ backend_config = backend_configs[backend_name]
109
+ # Dynamically skip test if GPU capability is not met
110
+ if backend_config.specific_gpu_arch and backend_config.specific_gpu_arch\
111
+ != current_platform.get_device_capability():
112
+ pytest.skip("Only Hopper GPUs support FA3 and FlashMLA")
113
+
114
+ env_vars = {"VLLM_USE_V1": "1", **backend_configs[backend_name].env_vars}
115
+
116
+ with temporary_environ(env_vars), ExitStack() as stack:
117
+ if not supported:
118
+ stack.enter_context(pytest.raises(Exception))
119
+
120
+ llm = LLM(model="Qwen/Qwen2-1.5B-Instruct",
121
+ max_num_seqs=256,
122
+ trust_remote_code=True,
123
+ gpu_memory_utilization=0.45,
124
+ max_model_len=1024,
125
+ compilation_config=CompilationConfig(
126
+ level=3, cudagraph_mode=cudagraph_mode))
127
+ llm.generate(["Hello, my name is"] * 10)
128
+
129
+ try:
130
+ llm = weakref.proxy(llm)
131
+ del llm
132
+ except UnboundLocalError:
133
+ pass
134
+
135
+ wait_for_gpu_memory_to_clear(
136
+ devices=[0],
137
+ threshold_ratio=0.1,
138
+ )
139
+
140
+
141
+# test cudagraph_mode with different compilation level.
142
+# (backend_name, cudagraph_mode, compilation_level, supported)
143
+combo_cases_2 = [
144
+ ("FA2", "FULL", 0, True), # no compilation + full cudagraph
145
+ ("FA2", "FULL", 3, True), # piecewise compilation + full cudagraph
146
+ ("FA2", "PIECEWISE", 0, False), # no compilation + piecewise cudagraph
147
+ ("FA2", "PIECEWISE", 3,
148
+ True), # piecewise compilation + piecewise cudagraph
149
+ ("FA2", "FULL_AND_PIECEWISE", 0,
150
+ False), # piecewise cudagraph not supported without piecewise compilation
151
+ ("FA2", "FULL_AND_PIECEWISE", 3, True),
152
+ ("FA2", "FULL_DECODE_ONLY", 0, True),
153
+ ("FA2", "FULL_DECODE_ONLY", 3, True),
154
+ ("FA2", "NONE", 0, True), # no compilation + no cudagraph
155
+ ("FA2", "NONE", 3, True), # piecewise compilation + no cudagraph
156
+]
157
+
158
+
159
+@pytest.mark.parametrize("combo_case", combo_cases_2)
160
+def test_cudagraph_compilation_combo(combo_case):
161
+ backend_name, cudagraph_mode, compilation_level, supported\
162
+ = combo_case
163
+
164
+ env_vars = {"VLLM_USE_V1": "1", **backend_configs[backend_name].env_vars}
165
+
166
+ with temporary_environ(env_vars), ExitStack() as stack:
167
+ if not supported:
168
+ stack.enter_context(pytest.raises(Exception))
169
+
170
+ llm = LLM(model="Qwen/Qwen2-1.5B-Instruct",
171
+ max_num_seqs=256,
172
+ trust_remote_code=True,
173
+ gpu_memory_utilization=0.45,
174
+ max_model_len=1024,
175
+ compilation_config=CompilationConfig(
176
+ level=compilation_level, cudagraph_mode=cudagraph_mode))
177
+ llm.generate(["Hello, my name is"] * 10)
178
+ try:
179
+ llm = weakref.proxy(llm)
180
+ del llm
181
+ except UnboundLocalError:
182
+ pass
183
+ finally:
184
+ wait_for_gpu_memory_to_clear(
185
+ devices=[0],
186
+ threshold_ratio=0.1,
187
+ )
vllm/compilation/backends.py
+30
-12
@@ -15,7 +15,7 @@ import torch.fx as fx
15
from torch._dispatch.python import enable_python_dispatcher
16
17
import vllm.envs as envs
18
-from vllm.config import CompilationConfig, VllmConfig
18
+from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig
19
from vllm.logger import init_logger
20
from vllm.platforms import current_platform
21
from vllm.utils import is_torch_equal_or_newer, resolve_obj_by_qualname
@@ -277,9 +277,6 @@ def split_graph(graph: fx.GraphModule,
277
return split_gm, outputs
278
279
280
-# we share the global graph pool among all the backends
281
-global_graph_pool = None
282
-
280
compilation_start_time = 0.0
281
282
@@ -339,14 +336,37 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):
336
graph_index=index,
337
num_graphs=len(self.compile_submod_names),
338
runtime_shape=None)
339
+ # Lazy import here to avoid circular import
340
+ from .cuda_graph import CUDAGraphOptions
341
+ from .cuda_piecewise_backend import PiecewiseBackend
342
343
- piecewise_backend = resolve_obj_by_qualname(
344
- current_platform.get_piecewise_backend_cls())
345
- self.module.__dict__[target] = piecewise_backend(
346
- submod, self.vllm_config, self.graph_pool, index,
343
+ piecewise_backend = PiecewiseBackend(
344
+ submod, self.vllm_config, index,
345
len(self.compile_submod_names), sym_shape_indices,
346
compiled_graph_for_dynamic_shape, self.vllm_backend)
347
348
+ if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:
349
+ # resolve the static graph wrapper class (e.g. CUDAGraphWrapper
350
+ # class) as platform dependent.
351
+ static_graph_wrapper_class = resolve_obj_by_qualname(
352
+ current_platform.get_static_graph_wrapper_cls())
353
+
354
+ # Always assign PIECEWISE runtime mode to the
355
+ # CUDAGraphWrapper for piecewise_backend, to distinguish
356
+ # it from the FULL cudagraph runtime mode, no matter it
357
+ # is wrapped on a full or piecewise fx graph.
358
+ self.module.__dict__[target] = static_graph_wrapper_class(
359
+ runnable=piecewise_backend,
360
+ vllm_config=self.vllm_config,
361
+ runtime_mode=CUDAGraphMode.PIECEWISE,
362
+ graph_pool=self.graph_pool,
363
+ cudagraph_options=CUDAGraphOptions(
364
+ debug_log_enable=piecewise_backend.is_first_graph,
365
+ gc_disable=not piecewise_backend.is_first_graph,
366
+ weak_ref_output=piecewise_backend.is_last_graph))
367
+ else:
368
+ self.module.__dict__[target] = piecewise_backend
369
+
370
compilation_counter.num_piecewise_capturable_graphs_seen += 1
371
372
return output
@@ -413,9 +433,7 @@ class VllmBackend:
433
# them, e.g. backbone (default), eagle_head, etc.
434
self.prefix = prefix or model_tag
435
416
- global global_graph_pool
417
- if global_graph_pool is None:
418
- global_graph_pool = current_platform.graph_pool_handle()
436
+ global_graph_pool = current_platform.get_global_graph_pool()
437
438
# TODO: in the future, if we want to use multiple
439
# streams, it might not be safe to share a global pool.
@@ -585,7 +603,7 @@ class VllmBackend:
603
604
self._called = True
605
588
- if not self.compilation_config.use_cudagraph or \
606
+ if self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE or \
607
not self.compilation_config.cudagraph_copy_inputs:
608
return self.split_gm
609
vllm/compilation/base_piecewise_backend.py
deleted
-72
@@ -1,72 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from typing import Any, Callable, Protocol
5
-
6
-import torch.fx as fx
7
-
8
-from vllm.compilation.backends import VllmBackend
9
-from vllm.config import VllmConfig
10
-
11
-
12
-class AbstractPiecewiseBackend(Protocol):
13
- """
14
- PiecewiseBackend interface that allows platforms to extend
15
- piecewise static graph.
16
- """
17
-
18
- def __init__(self, graph: fx.GraphModule, vllm_config: VllmConfig,
19
- graph_pool: Any, piecewise_compile_index: int,
20
- total_piecewise_compiles: int, sym_shape_indices: list[int],
21
- compiled_graph_for_general_shape: Callable,
22
- vllm_backend: VllmBackend, **kwargs):
23
- """
24
- Initializes the PiecewiseBackend class with compilation and
25
- execution-related configurations.
26
-
27
- This class handles piecewise compilation, graph capturing,
28
- and dispatching for specific input shapes.
29
-
30
- Args:
31
- graph (fx.GraphModule): The graph represented in fx.
32
- vllm_config (VllmConfig): Global configuration for vLLM.
33
- graph_pool (Any):
34
- Graph memory pool handle, e.g.,
35
- `torch.cuda.graph_pool_handle()`.
36
- piecewise_compile_index (int):
37
- Index of the current piecewise subgraph.
38
- total_piecewise_compiles (int):
39
- Total number of piecewise-compiled graphs.
40
- sym_shape_indices (list[int]):
41
- Indices of symbolic shape.
42
- compiled_graph_for_general_shape (Callable):
43
- Callable that executes the graph compiled for general shapes.
44
- vllm_backend (VllmBackend):
45
- Backend compiler that manages compilation and graph runtime
46
- for vLLM.
47
-
48
- Keyword Args:
49
- kwargs: Additional keyword arguments reserved for future
50
- extensions or custom platforms.
51
- """
52
- raise NotImplementedError
53
-
54
- def __call__(self, *args) -> Any:
55
- """Executes the compiled graph for given input args.
56
-
57
- If this is the first invocation, executes the general compiled graph
58
- and initiates the compilation process tracking. For subsequent calls,
59
- dynamically dispatches execution to either a compiled graph or a static
60
- graph based on the input shape.
61
-
62
- Args:
63
- *args: Variable length input arguments to be passed into the
64
- graph. The symbolic shape is expected to be in position
65
- `sym_shape_indices[0]`.
66
-
67
- Returns:
68
- Any: Output of the executed graph. This can be from the general
69
- compiled graph, a specialized compiled version for the given shape,
70
- or a replayed static graph.
71
- """
72
- raise NotImplementedError
vllm/compilation/base_static_graph.py
new
+54
@@ -0,0 +1,54 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Any, Callable, Protocol
5
+
6
+from vllm.config import CUDAGraphMode, VllmConfig
7
+
8
+
9
+class AbstractStaticGraphWrapper(Protocol):
10
+ """
11
+ StaticGraphWrapper interface that allows platforms to wrap a callable
12
+ to be captured as a static graph.
13
+ """
14
+
15
+ def __init__(self, runnable: Callable, vllm_config: VllmConfig,
16
+ runtime_mode: CUDAGraphMode, graph_pool: Any, **kwargs):
17
+ """
18
+ Initializes the StaticGraphWrapper class with graph capturing and
19
+ execution-related configurations.
20
+
21
+ Args:
22
+ runnable (Callable): The callable to be wrapped and captured.
23
+ vllm_config (VllmConfig): Global configuration for vLLM.
24
+ runtime_mode (CUDAGraphMode): The style of the static
25
+ graph runtime. See CUDAGraphMode in vllm/config.py.
26
+ Note that only the subset enum `NONE`, `PIECEWISE` and `FULL`
27
+ are used as concrete runtime mode for cudagraph dispatching.
28
+ graph_pool (Any):
29
+ Graph memory pool handle, e.g.,
30
+ `torch.cuda.graph_pool_handle()`.
31
+ Keyword Args:
32
+ kwargs: Additional keyword arguments for platform-specific
33
+ configurations.
34
+ """
35
+ raise NotImplementedError
36
+
37
+ def __call__(self, *args, **kwargs) -> Any:
38
+ """
39
+ Executes the wrapped callable.
40
+
41
+ If the current runtime mode in the ForwardContext matches the runtime
42
+ mode of this instance, it replays the CUDAGraph or captures it using
43
+ the callable if it hasn't been captured yet. Otherwise, it calls the
44
+ original callable directly.
45
+
46
+ Args:
47
+ *args: Variable length input arguments to be passed into the
48
+ callable.
49
+ **kwargs: Keyword arguments to be passed into the callable.
50
+
51
+ Returns:
52
+ Any: Output of the executed callable.
53
+ """
54
+ raise NotImplementedError
vllm/compilation/cuda_graph.py
new
+193
@@ -0,0 +1,193 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import dataclasses
5
+from contextlib import ExitStack
6
+from typing import Any, Callable, Optional
7
+from unittest.mock import patch
8
+
9
+import torch
10
+
11
+import vllm.envs as envs
12
+from vllm.compilation.counter import compilation_counter
13
+from vllm.compilation.monitor import validate_cudagraph_capturing_enabled
14
+from vllm.config import CUDAGraphMode, VllmConfig
15
+from vllm.forward_context import BatchDescriptor, get_forward_context
16
+from vllm.logger import init_logger
17
+from vllm.platforms import current_platform
18
+from vllm.utils import weak_ref_tensors
19
+
20
+logger = init_logger(__name__)
21
+
22
+
23
+@dataclasses.dataclass
24
+class CUDAGraphEntry:
25
+ batch_descriptor: BatchDescriptor
26
+ cudagraph: Optional[torch.cuda.CUDAGraph] = None
27
+ output: Optional[Any] = None
28
+
29
+ # for cudagraph debugging, track the input addresses
30
+ # during capture, and check if they are the same during replay
31
+ input_addresses: Optional[list[int]] = None
32
+
33
+
34
+@dataclasses.dataclass
35
+class CUDAGraphOptions:
36
+ debug_log_enable: bool = True
37
+ gc_disable: bool = False
38
+ weak_ref_output: bool = True
39
+
40
+
41
+class CUDAGraphWrapper:
42
+ """Wraps a runnable to add CUDA graph capturing and replaying ability. And
43
+ provide attribute access to the underlying `runnable` via `__getattr__`.
44
+
45
+ The workflow of this wrapper in the cudagraph dispatching is as follows:
46
+ 1. At initialization, a runtime mode is assigned to the wrapper (FULL or
47
+ PIECEWISE).
48
+ 2. At runtime, the wrapper receives a runtime_mode and a
49
+ batch_descriptor(key) from the forward context and blindly trust them
50
+ for cudagraph dispatching.
51
+ 3. If runtime_mode is NONE or runtime_mode does not match the mode of the
52
+ wrapper, just call the runnable directly.
53
+ 4. Otherwise, i.e., the runtime_mode matches the mode of the wrapper,
54
+ the wrapper will perform cudagraph capture(if key does not exist, create
55
+ a new entry and cache it) or replay (if key exists in the cache).
56
+
57
+ Note: CUDAGraphWrapper does not store persistent buffers or copy any
58
+ runtime inputs into that buffers for replay. We assume implementing them
59
+ is done outside of the wrapper. That is because we do not make any
60
+ assumption on the dynamic shape (batch size) of the runtime inputs, as a
61
+ trade-off for staying orthogonal to compilation logic. Nevertheless,
62
+ tracing and checking the input addresses to be consistent during replay is
63
+ guaranteed when VLLM_LOGGING_LEVEL == "DEBUG".
64
+ """
65
+
66
+ def __init__(self,
67
+ runnable: Callable,
68
+ vllm_config: VllmConfig,
69
+ runtime_mode: CUDAGraphMode,
70
+ graph_pool: Any = None,
71
+ cudagraph_options: Optional[CUDAGraphOptions] = None):
72
+ self.runnable = runnable
73
+ self.vllm_config = vllm_config
74
+ self.graph_pool = graph_pool
75
+ self.runtime_mode = runtime_mode
76
+ self.compilation_config = vllm_config.compilation_config
77
+
78
+ self.first_run_finished = False
79
+ self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG"
80
+
81
+ # assert runtime_mode is not NONE(no cudagraph), otherwise, we don't
82
+ # need to initialize a CUDAGraphWrapper.
83
+ assert self.runtime_mode != CUDAGraphMode.NONE
84
+ if self.graph_pool is None:
85
+ self.graph_pool = current_platform.get_global_graph_pool()
86
+
87
+ if cudagraph_options is None:
88
+ cudagraph_options = CUDAGraphOptions()
89
+ self.cudagraph_options = cudagraph_options
90
+ # the entries for different batch descriptors that we need to capture
91
+ # cudagraphs for.
92
+ self.concrete_cudagraph_entries: dict[BatchDescriptor, CUDAGraphEntry]\
93
+ = {}
94
+
95
+ def __getattr__(self, key: str):
96
+ # allow accessing the attributes of the runnable.
97
+ if hasattr(self.runnable, key):
98
+ return getattr(self.runnable, key)
99
+ raise AttributeError(f"Attribute {key} not exists in the runnable of "
100
+ f"cudagraph wrapper: {self.runnable}")
101
+
102
+ def unwrap(self) -> Callable:
103
+ # in case we need to access the original runnable.
104
+ return self.runnable
105
+
106
+ def __call__(self, *args, **kwargs):
107
+ forward_context = get_forward_context()
108
+ batch_descriptor = forward_context.batch_descriptor
109
+ cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode
110
+
111
+ if cudagraph_runtime_mode == CUDAGraphMode.NONE or \
112
+ cudagraph_runtime_mode != self.runtime_mode:
113
+ # CUDAGraphMode.NONE could mean the profile run, a warmup run, or
114
+ # running without cudagraphs.
115
+ # We do not trigger capture/replay if the runtime mode is not
116
+ # matches. This enables properly dispatching to the correct
117
+ # CUDAGraphWrapper when nesting multiple instances with different
118
+ # runtime modes.
119
+ return self.runnable(*args, **kwargs)
120
+
121
+ if batch_descriptor not in self.concrete_cudagraph_entries:
122
+ # create a new entry for this batch descriptor
123
+ self.concrete_cudagraph_entries[batch_descriptor] = \
124
+ CUDAGraphEntry(batch_descriptor=batch_descriptor)
125
+
126
+ entry = self.concrete_cudagraph_entries[batch_descriptor]
127
+
128
+ if entry.cudagraph is None:
129
+ if self.cudagraph_options.debug_log_enable:
130
+ # Since we capture cudagraph for many different shapes and
131
+ # capturing is fast, we don't need to log it for every
132
+ # shape. E.g. we only log it for the first subgraph in
133
+ # piecewise mode.
134
+ logger.debug("Capturing a cudagraph on (%s,%s)",
135
+ self.runtime_mode.name, entry.batch_descriptor)
136
+ # validate that cudagraph capturing is legal at this point.
137
+ validate_cudagraph_capturing_enabled()
138
+
139
+ input_addresses = [
140
+ x.data_ptr() for x in args if isinstance(x, torch.Tensor)
141
+ ]
142
+ entry.input_addresses = input_addresses
143
+ cudagraph = torch.cuda.CUDAGraph()
144
+
145
+ with ExitStack() as stack:
146
+ if self.cudagraph_options.gc_disable:
147
+ # during every model forward for piecewise cudagraph
148
+ # mode, we will capture many pieces of cudagraphs
149
+ # (roughly one per layer). running gc again and again
150
+ # across layers will make the cudagraph capture very slow.
151
+ # therefore, we only run gc for the first graph,
152
+ # and disable gc for the rest of the graphs.
153
+ stack.enter_context(patch("gc.collect", lambda: None))
154
+ stack.enter_context(
155
+ patch("torch.cuda.empty_cache", lambda: None))
156
+
157
+ # mind-exploding: carefully manage the reference and memory.
158
+ with torch.cuda.graph(cudagraph, pool=self.graph_pool):
159
+ # `output` is managed by pytorch's cudagraph pool
160
+ output = self.runnable(*args, **kwargs)
161
+ if self.cudagraph_options.weak_ref_output:
162
+ # by converting it to weak ref,
163
+ # the original `output` will immediately be released
164
+ # to save memory. It is only safe to do this for
165
+ # the last graph in piecewise cuadgraph mode, because
166
+ # the output of the last graph will not be used by
167
+ # any other cuda graph.
168
+ output = weak_ref_tensors(output)
169
+
170
+ # here we always use weak ref for the output
171
+ # to save memory
172
+ entry.output = weak_ref_tensors(output)
173
+ entry.cudagraph = cudagraph
174
+
175
+ compilation_counter.num_cudagraph_captured += 1
176
+
177
+ # important: we need to return the output, rather than
178
+ # the weak ref of the output, so that pytorch can correctly
179
+ # manage the memory during cuda graph capture
180
+ return output
181
+
182
+ if self.is_debugging_mode:
183
+ # check if the input addresses are the same
184
+ new_input_addresses = [
185
+ x.data_ptr() for x in args if isinstance(x, torch.Tensor)
186
+ ]
187
+ assert new_input_addresses == entry.input_addresses, (
188
+ f"Input addresses for cudagraphs are different "
189
+ f"during replay. Expected {entry.input_addresses}, "
190
+ f"got {new_input_addresses}")
191
+
192
+ entry.cudagraph.replay()
193
+ return entry.output
vllm/compilation/cuda_piecewise_backend.py
+16
-117
@@ -2,21 +2,15 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import dataclasses
5
-from contextlib import ExitStack
6
-from typing import Any, Callable, Optional
7
-from unittest.mock import patch
5
+from typing import Any, Callable
6
9
-import torch
7
import torch.fx as fx
8
9
import vllm.envs as envs
10
from vllm.compilation.backends import VllmBackend
14
-from vllm.compilation.counter import compilation_counter
11
from vllm.compilation.monitor import end_monitoring_torch_compile
12
from vllm.config import VllmConfig
17
-from vllm.forward_context import get_forward_context
13
from vllm.logger import init_logger
19
-from vllm.utils import weak_ref_tensors
14
15
logger = init_logger(__name__)
16
@@ -24,44 +18,29 @@ logger = init_logger(__name__)
18
@dataclasses.dataclass
19
class ConcreteSizeEntry:
20
runtime_shape: int
27
- need_to_compile: bool # the size is in compile_sizes
28
- use_cudagraph: bool # the size is in cudagraph_capture_sizes
29
-
21
compiled: bool = False
22
runnable: Callable = None # type: ignore
32
- num_finished_warmup: int = 0
33
- cudagraph: Optional[torch.cuda.CUDAGraph] = None
34
- output: Optional[Any] = None
35
-
36
- # for cudagraph debugging, track the input addresses
37
- # during capture, and check if they are the same during replay
38
- input_addresses: Optional[list[int]] = None
23
24
41
-class CUDAPiecewiseBackend:
25
+class PiecewiseBackend:
26
27
def __init__(self, graph: fx.GraphModule, vllm_config: VllmConfig,
44
- graph_pool: Any, piecewise_compile_index: int,
45
- total_piecewise_compiles: int, sym_shape_indices: list[int],
28
+ piecewise_compile_index: int, total_piecewise_compiles: int,
29
+ sym_shape_indices: list[int],
30
compiled_graph_for_general_shape: Callable,
31
vllm_backend: VllmBackend):
32
"""
33
The backend for piecewise compilation.
50
- It mainly handles the compilation and cudagraph capturing.
34
+ It mainly handles the compilation of static shapes and
35
+ dispatching based on runtime shape.
36
37
We will compile `self.graph` once for the general shape,
38
and then compile for different shapes specified in
39
`compilation_config.compile_sizes`.
55
-
56
- Independently, we will capture cudagraph for different shapes.
57
-
58
- If a shape needs both compilation and cudagraph, we will
59
- compile it first, and then capture cudagraph.
40
"""
41
self.graph = graph
42
self.vllm_config = vllm_config
43
self.compilation_config = vllm_config.compilation_config
64
- self.graph_pool = graph_pool
44
self.piecewise_compile_index = piecewise_compile_index
45
self.total_piecewise_compiles = total_piecewise_compiles
46
self.vllm_backend = vllm_backend
@@ -70,11 +49,10 @@ class CUDAPiecewiseBackend:
49
self.is_last_graph = (
50
piecewise_compile_index == total_piecewise_compiles - 1)
51
52
+ self.is_full_graph = total_piecewise_compiles == 1
53
+
54
self.compile_sizes: set[int] = set(
55
self.compilation_config.compile_sizes)
75
- self.cudagraph_capture_sizes: set[int] = set(
76
- self.compilation_config.cudagraph_capture_sizes
77
- ) if self.compilation_config.use_cudagraph else set()
56
57
self.first_run_finished = False
58
@@ -84,18 +62,18 @@ class CUDAPiecewiseBackend:
62
63
self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG"
64
87
- # the entries for different shapes that we need to either
88
- # compile or capture cudagraph
65
+ # the entries for different shapes that we need to compile
66
self.concrete_size_entries: dict[int, ConcreteSizeEntry] = {}
67
68
# to_be_compiled_sizes tracks the remaining sizes to compile,
69
# and updates during the compilation process, so we need to copy it
70
self.to_be_compiled_sizes: set[int] = self.compile_sizes.copy()
94
- for shape in self.compile_sizes.union(self.cudagraph_capture_sizes):
71
+
72
+ # We only keep compilation management inside this class directly.
73
+ for shape in self.compile_sizes:
74
self.concrete_size_entries[shape] = ConcreteSizeEntry(
75
runtime_shape=shape,
97
- need_to_compile=shape in self.compile_sizes,
98
- use_cudagraph=shape in self.cudagraph_capture_sizes,
76
+ runnable=self.compiled_graph_for_general_shape,
77
)
78
79
def check_for_ending_compilation(self):
@@ -112,16 +90,14 @@ class CUDAPiecewiseBackend:
90
return self.compiled_graph_for_general_shape(*args)
91
92
runtime_shape = args[self.sym_shape_indices[0]]
93
+
94
if runtime_shape not in self.concrete_size_entries:
95
# we don't need to do anything for this shape
96
return self.compiled_graph_for_general_shape(*args)
97
98
entry = self.concrete_size_entries[runtime_shape]
99
121
- if entry.runnable is None:
122
- entry.runnable = self.compiled_graph_for_general_shape
123
-
124
- if entry.need_to_compile and not entry.compiled:
100
+ if not entry.compiled:
101
entry.compiled = True
102
self.to_be_compiled_sizes.remove(runtime_shape)
103
# args are real arguments
@@ -138,81 +114,4 @@ class CUDAPiecewiseBackend:
114
if self.is_last_graph and not self.to_be_compiled_sizes:
115
self.check_for_ending_compilation()
116
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:
148
- if entry.num_finished_warmup < self.compilation_config.cudagraph_num_of_warmups: # noqa
149
- entry.num_finished_warmup += 1
150
- if self.is_first_graph:
151
- logger.debug(
152
- "Warming up %s/%s for shape %s",
153
- entry.num_finished_warmup,
154
- self.compilation_config.cudagraph_num_of_warmups,
155
- runtime_shape)
156
- return entry.runnable(*args)
157
-
158
- if self.is_first_graph:
159
- # Since we capture cudagraph for many different shapes and
160
- # capturing is fast, we don't need to log it for every shape.
161
- # We only log it in the debug mode.
162
- logger.debug("Capturing a cudagraph for shape %s",
163
- runtime_shape)
164
-
165
- input_addresses = [
166
- x.data_ptr() for x in args if isinstance(x, torch.Tensor)
167
- ]
168
- entry.input_addresses = input_addresses
169
- cudagraph = torch.cuda.CUDAGraph()
170
-
171
- with ExitStack() as stack:
172
- if not self.is_first_graph:
173
- # during every model forward, we will capture
174
- # many pieces of cudagraphs (roughly one per layer).
175
- # running gc again and again across layers will
176
- # make the cudagraph capture very slow.
177
- # therefore, we only run gc for the first graph,
178
- # and disable gc for the rest of the graphs.
179
- stack.enter_context(patch("gc.collect", lambda: None))
180
- stack.enter_context(
181
- patch("torch.cuda.empty_cache", lambda: None))
182
-
183
- # mind-exploding: carefully manage the reference and memory.
184
- with torch.cuda.graph(cudagraph, pool=self.graph_pool):
185
- # `output` is managed by pytorch's cudagraph pool
186
- output = entry.runnable(*args)
187
- if self.is_last_graph:
188
- # by converting it to weak ref,
189
- # the original `output` will immediately be released
190
- # to save memory. It is only safe to do this for
191
- # the last graph, because the output of the last graph
192
- # will not be used by any other cuda graph.
193
- output = weak_ref_tensors(output)
194
-
195
- # here we always use weak ref for the output
196
- # to save memory
197
- entry.output = weak_ref_tensors(output)
198
- entry.cudagraph = cudagraph
199
-
200
- compilation_counter.num_cudagraph_captured += 1
201
-
202
- # important: we need to return the output, rather than
203
- # the weak ref of the output, so that pytorch can correctly
204
- # manage the memory during cuda graph capture
205
- return output
206
-
207
- if self.is_debugging_mode:
208
- # check if the input addresses are the same
209
- new_input_addresses = [
210
- x.data_ptr() for x in args if isinstance(x, torch.Tensor)
211
- ]
212
- assert new_input_addresses == entry.input_addresses, (
213
- "Input addresses for cudagraphs are different during replay."
214
- f" Expected {entry.input_addresses}, got {new_input_addresses}"
215
- )
216
-
217
- entry.cudagraph.replay()
218
- return entry.output
117
+ return entry.runnable(*args)
vllm/compilation/monitor.py
+18
@@ -37,3 +37,21 @@ def end_monitoring_torch_compile(vllm_config: VllmConfig):
37
if context_manager is not None:
38
context_manager.__exit__(None, None, None)
39
context_manager = None
40
+
41
+
42
+cudagraph_capturing_enabled: bool = True
43
+
44
+
45
+def validate_cudagraph_capturing_enabled():
46
+ # used to monitor whether an cudagraph capturing is legal at runtime.
47
+ # should be called before any cudagraph capturing.
48
+ # if an illegal cudagraph capturing happens, raise an error.
49
+ global cudagraph_capturing_enabled
50
+ if not cudagraph_capturing_enabled:
51
+ raise RuntimeError("CUDA graph capturing detected at an inappropriate "
52
+ "time. This operation is currently disabled.")
53
+
54
+
55
+def set_cudagraph_capturing_enabled(enabled: bool):
56
+ global cudagraph_capturing_enabled
57
+ cudagraph_capturing_enabled = enabled
vllm/compilation/wrapper.py
+4
-3
@@ -11,7 +11,8 @@ from typing import Callable, Optional
11
import torch
12
13
import vllm.envs as envs
14
-from vllm.config import CompilationLevel, get_current_vllm_config
14
+from vllm.config import (CompilationLevel, CUDAGraphMode,
15
+ get_current_vllm_config)
16
from vllm.logger import init_logger
17
18
logger = init_logger(__name__)
@@ -115,8 +116,8 @@ class TorchCompileWrapperWithCustomDispatcher:
116
except Exception:
117
pass
118
118
- if self.vllm_config.compilation_config.use_cudagraph and \
119
- "update" in new_code.co_names:
119
+ if self.vllm_config.compilation_config.cudagraph_mode != \
120
+ CUDAGraphMode.NONE and "update" in new_code.co_names:
121
import depyf
122
src = depyf.decompile(new_code)
123
msg = "Assigning / modifying buffers of nn.Module during forward pass is not allowed when using cudagraph inside the compiler because it will cause silent errors. Please use eager mode or fix the code. The following code contains clues about which buffer is being modified (please search for the usage of the function `update`):\n" + src # noqa
vllm/config/__init__.py
+40
-12
@@ -32,7 +32,7 @@ from vllm import version
32
from vllm.config.cache import (BlockSize, CacheConfig, CacheDType, MambaDType,
33
PrefixCachingHashAlgo)
34
from vllm.config.compilation import (CompilationConfig, CompilationLevel,
35
- PassConfig)
35
+ CUDAGraphMode, PassConfig)
36
from vllm.config.parallel import DistributedExecutorBackend, ParallelConfig
37
from vllm.config.scheduler import SchedulerConfig, SchedulerPolicy
38
from vllm.config.utils import ConfigType, config
@@ -3529,11 +3529,21 @@ class VllmConfig:
3529
else:
3530
self.compilation_config.level = \
3531
CompilationLevel.NO_COMPILATION
3532
+
3533
else:
3534
# NB: Passing both --enforce-eager and a compilation level
3535
# in V0 means the compilation level wins out.
3536
self.compilation_config.level = CompilationLevel.NO_COMPILATION
3537
3538
+ # if cudagraph_mode is not explicitly set by users, set default value
3539
+ if self.compilation_config.cudagraph_mode is None:
3540
+ if envs.VLLM_USE_V1 and self.compilation_config.level \
3541
+ == CompilationLevel.PIECEWISE:
3542
+ self.compilation_config.cudagraph_mode = \
3543
+ CUDAGraphMode.PIECEWISE
3544
+ else:
3545
+ self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
3546
+
3547
# async tp is built on top of sequence parallelism
3548
# and requires it to be enabled.
3549
if self.compilation_config.pass_config.enable_async_tp:
@@ -3541,12 +3551,13 @@ class VllmConfig:
3551
True
3552
if self.compilation_config.pass_config.enable_sequence_parallelism:
3553
self.compilation_config.custom_ops.append("+rms_norm")
3544
- if envs.VLLM_USE_V1 and self.model_config is not None and \
3545
- not self.model_config.enforce_eager:
3546
- # By default, V1 uses piecewise CUDA graphs. If full_cuda_graph
3547
- # is set to True, full CUDA graphs will be used.
3554
+
3555
+ # disable cudagraph when enforce eager execution
3556
+ if self.model_config is not None and self.model_config.enforce_eager:
3557
+ logger.info("Cudagraph is disabled under eager mode")
3558
+ self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
3559
+ elif envs.VLLM_USE_V1:
3560
self.compilation_config.cudagraph_num_of_warmups = 1
3549
- self.compilation_config.set_splitting_ops_for_v1()
3561
3562
self._set_cudagraph_sizes()
3563
@@ -3566,12 +3577,6 @@ class VllmConfig:
3577
"Disabling `torch.compile`.")
3578
self.compilation_config.level = CompilationLevel.NO_COMPILATION
3579
3569
- if self.compilation_config.full_cuda_graph and \
3570
- not self.model_config.disable_cascade_attn:
3571
- logger.info("full_cuda_graph is not supported with "
3572
- "cascade attention. Disabling cascade attention.")
3573
- self.model_config.disable_cascade_attn = True
3574
-
3580
disable_chunked_prefill_reasons: list[str] = []
3581
3582
if self.model_config and self.model_config.pooler_config:
@@ -3612,9 +3617,32 @@ class VllmConfig:
3617
"to True to enable.")
3618
current_platform.check_and_update_config(self)
3619
3620
+ # final check of cudagraph mode after platform-specific update
3621
+ if envs.VLLM_USE_V1:
3622
+ if self.compilation_config.cudagraph_mode == CUDAGraphMode.FULL \
3623
+ and self.model_config is not None and \
3624
+ not self.model_config.disable_cascade_attn:
3625
+ logger.info("CUDAGraphMode.FULL is not supported with "
3626
+ "cascade attention currently. Disabling cascade"
3627
+ "attention.")
3628
+ self.model_config.disable_cascade_attn = True
3629
+
3630
+ if self.compilation_config.cudagraph_mode\
3631
+ .requires_piecewise_compilation():
3632
+ assert self.compilation_config.level == \
3633
+ CompilationLevel.PIECEWISE, \
3634
+ "Compilation level should be CompilationLevel.PIECEWISE "\
3635
+ "when cudagraph_mode piecewise cudagraphs is used, "\
3636
+ f"cudagraph_mode={self.compilation_config.cudagraph_mode}"
3637
+
3638
if not self.instance_id:
3639
self.instance_id = random_uuid()[:5]
3640
3641
+ # Do this after all the updates to compilation_config.level
3642
+ if envs.VLLM_USE_V1 and \
3643
+ self.compilation_config.level == CompilationLevel.PIECEWISE:
3644
+ self.compilation_config.set_splitting_ops_for_v1()
3645
+
3646
if (envs.VLLM_USE_V1
3647
and not self.scheduler_config.disable_hybrid_kv_cache_manager):
3648
# logger should only print warning message for hybrid models. As we
vllm/config/compilation.py
+162
-26
@@ -1,12 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import enum
5
import hashlib
6
from collections import Counter
7
from dataclasses import asdict, field
7
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union
8
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, Union
9
9
-from pydantic import TypeAdapter
10
+from pydantic import TypeAdapter, field_validator
11
from pydantic.dataclasses import dataclass
12
13
import vllm.envs as envs
@@ -31,6 +32,40 @@ class CompilationLevel:
32
PIECEWISE = 3
33
34
35
+class CUDAGraphMode(enum.Enum):
36
+ """ Constants for the cudagraph mode in CompilationConfig.
37
+ Meanwhile, the subset enum `NONE`, `PIECEWISE` and `FULL` are also
38
+ treated as concrete runtime mode for cudagraph runtime dispatching.
39
+ """
40
+ NONE = 0
41
+ PIECEWISE = 1
42
+ FULL = 2
43
+ FULL_DECODE_ONLY = (FULL, NONE)
44
+ FULL_AND_PIECEWISE = (FULL, PIECEWISE)
45
+
46
+ def decode_mode(self) -> 'CUDAGraphMode':
47
+ return CUDAGraphMode(self.value[0]) if \
48
+ self.separate_routine() else self
49
+
50
+ def mixed_mode(self) -> 'CUDAGraphMode':
51
+ return CUDAGraphMode(self.value[1]) if \
52
+ self.separate_routine() else self
53
+
54
+ def requires_piecewise_compilation(self) -> bool:
55
+ return (self.decode_mode() == CUDAGraphMode.PIECEWISE
56
+ or self.mixed_mode() == CUDAGraphMode.PIECEWISE)
57
+
58
+ def max_cudagraph_mode(self) -> 'CUDAGraphMode':
59
+ return CUDAGraphMode(max(
60
+ self.value)) if self.separate_routine() else self
61
+
62
+ def has_full_cudagraphs(self) -> bool:
63
+ return self.max_cudagraph_mode() == CUDAGraphMode.FULL
64
+
65
+ def separate_routine(self) -> bool:
66
+ return isinstance(self.value, tuple)
67
+
68
+
69
@config
70
@dataclass
71
class PassConfig:
@@ -91,6 +126,7 @@ class CompilationConfig:
126
- [`splitting_ops`][vllm.config.CompilationConfig.splitting_ops]
127
- CudaGraph capture:
128
- [`use_cudagraph`][vllm.config.CompilationConfig.use_cudagraph]
129
+ - [`cudagraph_mode`][vllm.config.CompilationConfig.cudagraph_mode]
130
- [`cudagraph_capture_sizes`]
131
[vllm.config.CompilationConfig.cudagraph_capture_sizes]
132
- [`cudagraph_num_of_warmups`]
@@ -157,7 +193,7 @@ class CompilationConfig:
193
By default, all custom ops are enabled when running without Inductor and
194
disabled when running with Inductor: level>=PIECEWISE and use_inductor=True.
195
Inductor generates (fused) Triton kernels for disabled custom ops."""
160
- splitting_ops: list[str] = field(default_factory=list)
196
+ splitting_ops: Optional[list[str]] = None
197
"""A list of ops to split the full graph into subgraphs, used in piecewise
198
compilation."""
199
@@ -187,7 +223,43 @@ class CompilationConfig:
223
constructor, e.g. `CompilationConfig(inductor_passes={"a": func})`."""
224
225
# CudaGraph compilation
190
- use_cudagraph: bool = field(default_factory=lambda: envs.VLLM_USE_V1)
226
+ cudagraph_mode: Optional[CUDAGraphMode] = None
227
+ """
228
+ The mode of the cudagraph.
229
+ - NONE, no cudagraph capture.
230
+ - PIECEWISE. (v1 default)
231
+ - FULL.
232
+ - FULL_DECODE_ONLY.
233
+ - FULL_AND_PIECEWISE.
234
+
235
+ PIECEWISE mode build piecewise cudagraph only, keeping the cudagraph
236
+ incompatiable ops (i.e. some attention ops) outside the cudagraph
237
+ for general flexibility.
238
+ This is the default mode.
239
+
240
+ FULL mode: Capture full cudagraph for all batches. Can be good for small
241
+ models or workloads with small prompts; not supported by many backends.
242
+ Generally for performance FULL_AND_PIECEWISE is better.
243
+
244
+ FULL_DECODE_ONLY mode: Capture full cudagraph for decode batches only.
245
+ Mixed prefill-decode batches are run without cudagraphs. Can be good for
246
+ decode instances in a P/D setup where prefill is not as important so we
247
+ can save some memory.
248
+
249
+ FULL_AND_PIECEWISE mode: Capture full cudagraph for decode batches and
250
+ piecewise cudagraph for prefill and mixed prefill-decode batches.
251
+ This is like the most performant mode for most models.
252
+
253
+ Currently, the cudagraph mode is only used for the v1 engine.
254
+ Note that the cudagraph logic is generally orthogonal to the
255
+ compilation logic. While piecewise cudagraphs require piecewise
256
+ compilation (level=PIECEWISE and non-empty splitting_ops), full
257
+ cudagraphs are supported with and without compilation.
258
+
259
+ Warning: This flag is new and subject to change in addition
260
+ more modes may be added.
261
+ """
262
+ use_cudagraph: bool = True
263
"""Whether to use cudagraph inside compilation.
264
- False: cudagraph inside compilation is not used.
265
- True: cudagraph inside compilation is used. It requires
@@ -197,8 +269,9 @@ class CompilationConfig:
269
CompilationLevel.PIECEWISE (aka -O3).
270
Note that this is orthogonal to the cudagraph capture logic
271
outside of compilation.
200
- TODO: move outside cudagraph logic into compilation.
201
- torch.compile will handle cudagraph capture logic in the future."""
272
+ Warning: This flag is deprecated and will be removed in the next major or
273
+ minor release, i.e. v0.11.0 or v1.0.0. Please use cudagraph_mode instead.
274
+ """
275
cudagraph_num_of_warmups: int = 0
276
"""Number of warmup runs for cudagraph.
277
It means the first several runs will be treated as warmup runs.
@@ -213,12 +286,17 @@ class CompilationConfig:
286
cudagraph. If the caller can guarantee that the same input buffers
287
are always used, it can set this to False. Otherwise, it should
288
set this to True, and the compiler will copy the input to an
216
- internally managed buffer. Default is False."""
217
- full_cuda_graph: bool = False
289
+ internally managed buffer. Default is False.
290
+ Note that this flag is only effective when cudagraph_mode is PIECEWISE.
291
+ """
292
+ full_cuda_graph: Optional[bool] = False
293
"""whether to use a full cuda graph for the entire forward pass rather than
294
splitting certain operations such as attention into subgraphs. Thus this
295
flag cannot be used together with splitting_ops. This may provide
221
- performance benefits for smaller models."""
296
+ performance benefits for smaller models.
297
+ Warning: This flag is deprecated and will be removed in the next major or
298
+ minor release, i.e. v0.11.0 or v1.0.0. Please use cudagraph_mode instead.
299
+ """
300
301
pass_config: PassConfig = field(default_factory=PassConfig)
302
"""Custom inductor passes, see PassConfig for more details"""
@@ -253,6 +331,13 @@ class CompilationConfig:
331
Map from layer name to layer objects that need to be accessed outside
332
model code, e.g., Attention, FusedMOE when dp_size>1."""
333
334
+ # Attention ops; used for piecewise cudagraphs
335
+ _attention_ops: ClassVar[list[str]] = [
336
+ "vllm.unified_attention",
337
+ "vllm.unified_attention_with_output",
338
+ "vllm.mamba_mixer2",
339
+ ]
340
+
341
def compute_hash(self) -> str:
342
"""
343
WARNING: Whenever a new field is added to this config,
@@ -297,13 +382,26 @@ class CompilationConfig:
382
if pass_config_exclude:
383
exclude["pass_config"] = pass_config_exclude
384
300
- return TypeAdapter(CompilationConfig).dump_json(
301
- self,
302
- exclude=exclude, # type: ignore[arg-type]
303
- exclude_unset=True).decode()
385
+ # The cast to string is necessary because Pydantic is mocked in docs
386
+ # builds and sphinx-argparse doesn't know the return type of decode()
387
+ return str(
388
+ TypeAdapter(CompilationConfig).dump_json(
389
+ self,
390
+ exclude=exclude, # type: ignore[arg-type]
391
+ exclude_unset=True).decode())
392
393
__str__ = __repr__
394
395
+ @field_validator("cudagraph_mode", mode="before")
396
+ @classmethod
397
+ def validate_cudagraph_mode_before(cls, value: Any) -> Any:
398
+ """
399
+ enable parse the `cudagraph_mode` enum type from string
400
+ """
401
+ if isinstance(value, str):
402
+ return CUDAGraphMode[value.upper()]
403
+ return value
404
+
405
def __post_init__(self) -> None:
406
count_none = self.custom_ops.count("none")
407
count_all = self.custom_ops.count("all")
@@ -341,7 +439,26 @@ class CompilationConfig:
439
if isinstance(self.pass_config, dict):
440
self.pass_config = PassConfig(**self.pass_config)
441
344
- def init_backend(self, vllm_config: VllmConfig) -> Union[str, Callable]:
442
+ # migrate the deprecated flags
443
+ if not self.use_cudagraph:
444
+ logger.warning("use_cudagraph is deprecated, use "
445
+ "cudagraph_mode=NONE instead.")
446
+ if self.cudagraph_mode is not None:
447
+ raise ValueError(
448
+ "use_cudagraph and cudagraph_mode are mutually"
449
+ " exclusive, prefer cudagraph_mode since "
450
+ "use_cudagraph is deprecated.")
451
+ self.cudagraph_mode = CUDAGraphMode.NONE
452
+ if self.full_cuda_graph:
453
+ logger.warning("full_cuda_graph is deprecated, use "
454
+ "cudagraph_mode=FULL instead.")
455
+ if self.cudagraph_mode is not None:
456
+ raise ValueError("full_cuda_graph and cudagraph_mode are "
457
+ "mutually exclusive, prefer cudagraph_mode "
458
+ "since full_cuda_graph is deprecated.")
459
+ self.cudagraph_mode = CUDAGraphMode.FULL
460
+
461
+ def init_backend(self, vllm_config: "VllmConfig") -> Union[str, Callable]:
462
if self.level == CompilationLevel.NO_COMPILATION:
463
raise ValueError("No compilation level is set.")
464
@@ -414,15 +531,34 @@ class CompilationConfig:
531
self.max_capture_size] = self.max_capture_size
532
533
def set_splitting_ops_for_v1(self):
417
- # NOTE: this function needs to be called
418
- if self.splitting_ops and self.full_cuda_graph:
419
- raise ValueError("full_cuda_graph cannot be used together with "
420
- "splitting_ops, as Full CUDA graph will override "
421
- f"the splitting_ops: {self.splitting_ops}")
422
-
423
- if not self.splitting_ops:
424
- self.splitting_ops = [] if self.full_cuda_graph else [
425
- "vllm.unified_attention",
426
- "vllm.unified_attention_with_output",
427
- "vllm.mamba_mixer2",
428
- ]
534
+ # NOTE: this function needs to be called only when level is
535
+ # CompilationLevel.PIECEWISE
536
+ assert self.level == CompilationLevel.PIECEWISE, (
537
+ "set_splitting_ops_for_v1 should only be called when "
538
+ "level is CompilationLevel.PIECEWISE")
539
+
540
+ if self.splitting_ops is None:
541
+ # NOTE: When using full cudagraph, instead of setting an empty
542
+ # list and capture the full cudagraph inside the flattened fx
543
+ # graph, we keep the piecewise fx graph structure but capture the
544
+ # full cudagraph outside the fx graph. This reduces some cpu
545
+ # overhead when the runtime batch_size is not cudagraph captured.
546
+ # see https://github.com/vllm-project/vllm/pull/20059 for details.
547
+ self.splitting_ops = self._attention_ops
548
+ elif len(self.splitting_ops) == 0:
549
+ logger.warning_once("Using piecewise compilation with empty "
550
+ "splitting_ops.")
551
+ if self.cudagraph_mode == CUDAGraphMode.PIECEWISE:
552
+ logger.warning_once(
553
+ "When compilation level is piecewise with empty "
554
+ "splitting_ops, PIECEWISE cudagraph_mode will be "
555
+ "treated as FULL cudagraph_mode. Please ensure you are "
556
+ "using attention backends that support cudagraph or set "
557
+ "cudagraph_mode to NONE explicitly if encountering "
558
+ "any problems.")
559
+ self.cudagraph_mode = CUDAGraphMode.FULL
560
+ self.splitting_ops = []
561
+
562
+ def splitting_ops_contain_attention(self) -> bool:
563
+ return self.splitting_ops is not None and all(
564
+ op in self.splitting_ops for op in self._attention_ops)
vllm/forward_context.py
+41
-11
@@ -5,13 +5,13 @@ import time
5
from collections import defaultdict
6
from contextlib import contextmanager
7
from dataclasses import dataclass
8
-from typing import TYPE_CHECKING, Any, Optional, Union
8
+from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Union
9
10
import torch
11
import torch.distributed as dist
12
13
import vllm.envs as envs
14
-from vllm.config import ParallelConfig, VllmConfig
14
+from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig
15
from vllm.logger import init_logger
16
17
if TYPE_CHECKING:
@@ -26,6 +26,27 @@ batchsize_logging_interval: float = envs.VLLM_LOG_BATCHSIZE_INTERVAL
26
batchsize_forward_time: defaultdict = defaultdict(list)
27
28
29
+class BatchDescriptor(NamedTuple):
30
+ """
31
+ Batch descriptor for cudagraph dispatching. We should keep the num of
32
+ items as minimal as possible to properly and uniquely describe the padded
33
+ batch for cudagraph.
34
+ """
35
+ num_tokens: int
36
+ uniform_decode: bool = False
37
+ """
38
+ False can also be used for an uniform decode batch to dispatch to the
39
+ cudagraph supporting non-uniform batches.
40
+ """
41
+
42
+ @property
43
+ def non_uniform(self) -> "BatchDescriptor":
44
+ """
45
+ Return a non-uniform version of current batch descriptor.
46
+ """
47
+ return BatchDescriptor(self.num_tokens, uniform_decode=False)
48
+
49
+
50
def _compute_chunked_local_num_tokens(num_tokens_across_dp_cpu: list[int],
51
max_num_tokens: int,
52
chunk_idx: int) -> list[int]:
@@ -152,7 +173,15 @@ class ForwardContext:
173
virtual_engine: int # set dynamically for each forward pass
174
# set dynamically for each forward pass
175
dp_metadata: Optional[DPMetadata] = None
155
- skip_cuda_graphs: bool = False
176
+ # determine the cudagraph style at runtime to be FULL, PIECEWISE, or NONE.
177
+ # by default NONE, no cudagraph is used.
178
+ cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE
179
+ batch_descriptor: Optional[BatchDescriptor] = None
180
+
181
+ def __post_init__(self):
182
+ assert self.cudagraph_runtime_mode in [
183
+ CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL], \
184
+ f"Invalid cudagraph runtime mode: {self.cudagraph_runtime_mode}"
185
186
187
_forward_context: Optional[ForwardContext] = None
@@ -168,13 +197,13 @@ def get_forward_context() -> ForwardContext:
197
198
@contextmanager
199
def set_forward_context(
171
- attn_metadata: Any,
172
- vllm_config: VllmConfig,
173
- virtual_engine: int = 0,
174
- num_tokens: Optional[int] = None,
175
- num_tokens_across_dp: Optional[torch.Tensor] = None,
176
- skip_cuda_graphs: bool = False,
177
-):
200
+ attn_metadata: Any,
201
+ vllm_config: VllmConfig,
202
+ virtual_engine: int = 0,
203
+ num_tokens: Optional[int] = None,
204
+ num_tokens_across_dp: Optional[torch.Tensor] = None,
205
+ cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
206
+ batch_descriptor: Optional[BatchDescriptor] = None):
207
"""A context manager that stores the current forward context,
208
can be attention metadata, etc.
209
Here we can inject common logic for every model forward pass.
@@ -198,7 +227,8 @@ def set_forward_context(
227
virtual_engine=virtual_engine,
228
attn_metadata=attn_metadata,
229
dp_metadata=dp_metadata,
201
- skip_cuda_graphs=skip_cuda_graphs,
230
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
231
+ batch_descriptor=batch_descriptor,
232
)
233
234
try:
vllm/platforms/cuda.py
+8
-5
@@ -177,17 +177,20 @@ class CudaPlatformBase(Platform):
177
logger.info("Forcing kv cache block size to 128 for "
178
"CUTLASS_MLA backend.")
179
180
+ # lazy import to avoid circular import
181
+ from vllm.config import CUDAGraphMode
182
+
183
compilation_config = vllm_config.compilation_config
184
if (envs.VLLM_ALL2ALL_BACKEND == "deepep_high_throughput"
185
and parallel_config.data_parallel_size > 1
183
- and compilation_config.use_cudagraph):
186
+ and compilation_config.cudagraph_mode != CUDAGraphMode.NONE):
187
logger.info(
185
- "Data Parallel: Forcing enforce eager to be True since DP "
188
+ "Data Parallel: disabling cudagraphs since DP "
189
"with DeepEP high-throughput kernels are not CUDA Graph "
190
"compatible. The DeepEP low-latency kernels are CUDA Graph "
191
"compatible. Set the all_to_all backend to deepep_low_latency "
192
"to use those kernels instead.")
190
- compilation_config.use_cudagraph = False
193
+ compilation_config.cudagraph_mode = CUDAGraphMode.NONE
194
if model_config is not None:
195
model_config.enforce_eager = True
196
@@ -454,8 +457,8 @@ class CudaPlatformBase(Platform):
457
return True
458
459
@classmethod
457
- def get_piecewise_backend_cls(cls) -> str:
458
- return "vllm.compilation.cuda_piecewise_backend.CUDAPiecewiseBackend" # noqa
460
+ def get_static_graph_wrapper_cls(cls) -> str:
461
+ return "vllm.compilation.cuda_graph.CUDAGraphWrapper"
462
463
@classmethod
464
def stateless_init_device_torch_dist_pg(
vllm/platforms/interface.py
+15
-4
@@ -7,7 +7,7 @@ import random
7
import sys
8
from datetime import timedelta
9
from platform import uname
10
-from typing import TYPE_CHECKING, NamedTuple, Optional, Union
10
+from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Union
11
12
import numpy as np
13
import torch
@@ -137,6 +137,8 @@ class Platform:
137
138
additional_env_vars: list[str] = []
139
140
+ _global_graph_pool: Optional[Any] = None
141
+
142
@property
143
def supported_dtypes(self) -> list[torch.dtype]:
144
"""Returns the supported dtypes for the current platform."""
@@ -522,6 +524,15 @@ class Platform:
524
" attribute.", self.device_type, key)
525
return None
526
527
+ def get_global_graph_pool(self) -> Any:
528
+ """
529
+ Return the global graph pool for the this platform.
530
+ """
531
+ cls = self.__class__
532
+ if cls._global_graph_pool is None:
533
+ cls._global_graph_pool = self.graph_pool_handle()
534
+ return cls._global_graph_pool
535
+
536
@classmethod
537
def get_cu_count(cls, device_id: int = 0) -> int:
538
"""
@@ -530,11 +541,11 @@ class Platform:
541
raise NotImplementedError
542
543
@classmethod
533
- def get_piecewise_backend_cls(cls) -> str:
544
+ def get_static_graph_wrapper_cls(cls) -> str:
545
"""
535
- Get piecewise backend class for piecewise graph.
546
+ Get static graph wrapper class for static graph.
547
"""
537
- return "vllm.compilation.base_piecewise_backend.AbstractPiecewiseBackend" # noqa
548
+ return "vllm.compilation.base_static_graph.AbstractStaticGraphWrapper"
549
550
@classmethod
551
def stateless_init_device_torch_dist_pg(
vllm/platforms/rocm.py
+2
-2
@@ -421,8 +421,8 @@ class RocmPlatform(Platform):
421
return 'gfx1' in torch.cuda.get_device_properties(0).gcnArchName
422
423
@classmethod
424
- def get_piecewise_backend_cls(cls) -> str:
425
- return "vllm.compilation.cuda_piecewise_backend.CUDAPiecewiseBackend" # noqa
424
+ def get_static_graph_wrapper_cls(cls) -> str:
425
+ return "vllm.compilation.cuda_graph.CUDAGraphWrapper"
426
427
@classmethod
428
def stateless_init_device_torch_dist_pg(
vllm/platforms/tpu.py
+10
-2
@@ -99,7 +99,7 @@ class TpuPlatform(Platform):
99
100
@classmethod
101
def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
102
- from vllm.config import CompilationLevel
102
+ from vllm.config import CompilationLevel, CUDAGraphMode
103
104
cache_config = vllm_config.cache_config
105
# For v0, the default block size is 16.
@@ -109,9 +109,17 @@ class TpuPlatform(Platform):
109
110
# TPU only supports DYNAMO_ONCE compilation level
111
if compilation_config.level != CompilationLevel.DYNAMO_ONCE:
112
- logger.info("[TPU] Forcing DYNAMO_ONCE compilation level")
112
+ logger.info("[TPU] Forcing DYNAMO_ONCE compilation level, and "
113
+ "disabling cudagraph.")
114
compilation_config.level = CompilationLevel.DYNAMO_ONCE
115
116
+ if compilation_config.cudagraph_mode is None or \
117
+ compilation_config.cudagraph_mode.max_cudagraph_mode() \
118
+ != CUDAGraphMode.NONE:
119
+ logger.info("[TPU] CUDA graph is not supported on TPU, "
120
+ "disabling cudagraphs.")
121
+ compilation_config.cudagraph_mode = CUDAGraphMode.NONE
122
+
123
if compilation_config.backend == "":
124
compilation_config.backend = "openxla"
125
vllm/platforms/xpu.py
+12
-10
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Optional
7
import torch
8
9
import vllm.envs as envs
10
+from vllm.config import CUDAGraphMode
11
from vllm.logger import init_logger
12
from vllm.utils import DEFAULT_MAX_NUM_BATCHED_TOKENS
13
@@ -100,16 +101,17 @@ class XPUPlatform(Platform):
101
# Instances created using VllmConfig() typically have model_config as
102
# None by default. The modification involves adding a check to prevent
103
# potential null exceptions check and update model config.
103
- if model_config is not None:
104
- if model_config.dtype == torch.bfloat16:
105
- bf16_supported = cls.device_support_bf16()
106
- if not bf16_supported:
107
- model_config.dtype = torch.float16
108
- if not model_config.enforce_eager:
109
- logger.warning(
110
- "CUDA graph is not supported on XPU, fallback to the eager "
111
- "mode.")
112
- model_config.enforce_eager = True
104
+ if model_config is not None and model_config.dtype == torch.bfloat16 \
105
+ and not cls.device_support_bf16():
106
+ model_config.dtype = torch.float16
107
+
108
+ compilation_config = vllm_config.compilation_config
109
+ if compilation_config.cudagraph_mode is None or \
110
+ compilation_config.cudagraph_mode.max_cudagraph_mode() \
111
+ != CUDAGraphMode.NONE:
112
+ logger.info("[XPU] CUDA graph is not supported on XPU, "
113
+ "disabling cudagraphs.")
114
+ compilation_config.cudagraph_mode = CUDAGraphMode.NONE
115
116
# check and update parallel config
117
parallel_config = vllm_config.parallel_config
vllm/v1/attention/backends/flash_attn.py
+37
-31
@@ -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 ClassVar, Optional
5
+from typing import Optional
6
7
import numpy as np
8
import torch
@@ -154,9 +154,26 @@ def _get_sliding_window_configs(
154
155
class FlashAttentionMetadataBuilder(
156
AttentionMetadataBuilder[FlashAttentionMetadata]):
157
- attn_cudagraph_support: ClassVar[AttentionCGSupport] = \
158
- AttentionCGSupport.NEVER if get_flash_attn_version() == 2 \
159
- else AttentionCGSupport.ALWAYS
157
+ # FA3:
158
+ # Supports full cudagraphs for all cases.
159
+ #
160
+ # FA2:
161
+ # For FA2, a graph is captured with max_query_len=1, (which is what we
162
+ # capture by default for num_tokens <= max_num_seqs when there is no
163
+ # spec-decode) then these graphs will not work for mixed prefill-decode
164
+ # (unlike FA3). This is due to special max_query_len=1 packed-GQA handling
165
+ # in FA2.
166
+ # In summary if we are running with spec decodes the graphs would
167
+ # work for mixed prefill-decode and uniform-decode. But for non-spec decodes
168
+ # the graphs would not work for mixed prefill-decode; sorta the inverse
169
+ # of UNIFORM_SINGLE_TOKEN_DECODE.
170
+ # Theres probably a better way to describe this using `AttentionCGSupport`
171
+ # but for now just set it to `UNIFORM_BATCH` to get use to drop down
172
+ # to FULL_AND_PIECEWISE.
173
+ # TODO(luka, lucas): audit FA2 as part of:
174
+ # https://github.com/vllm-project/vllm/issues/22945
175
+ cudagraph_support = AttentionCGSupport.ALWAYS \
176
+ if get_flash_attn_version() == 3 else AttentionCGSupport.UNIFORM_BATCH
177
178
def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
179
vllm_config: VllmConfig, device: torch.device):
@@ -177,17 +194,13 @@ class FlashAttentionMetadataBuilder(
194
195
self.max_num_splits = 0 # No upper bound on the number of splits.
196
self.aot_schedule = (get_flash_attn_version() == 3)
180
- self.use_full_cuda_graph = self.compilation_config.full_cuda_graph
181
- if self.use_full_cuda_graph:
182
- if not self.aot_schedule:
183
- raise ValueError(
184
- "AoT scheduling is required for full cuda graph.")
185
- capture_sizes = self.compilation_config.cudagraph_capture_sizes
186
- if not capture_sizes:
187
- raise ValueError(
188
- "cudagraph_capture_sizes should not be None when "
189
- "full_cuda_graph is True.")
190
- self.max_cudagraph_size = max(capture_sizes)
197
+
198
+ self.use_full_cuda_graph = \
199
+ self.compilation_config.cudagraph_mode.has_full_cudagraphs()
200
+
201
+ if self.use_full_cuda_graph and self.aot_schedule:
202
+ self.max_cudagraph_size = self.compilation_config.max_capture_size
203
+
204
if self.max_cudagraph_size > 992:
205
# This condition derives from FA3's internal heuristic.
206
# TODO(woosuk): Support larger cudagraph sizes.
@@ -310,9 +323,9 @@ class FlashAttentionMetadataBuilder(
323
seqlens=seq_lens,
324
max_seq_len=max_seq_len,
325
causal=causal)
313
-
314
- if self.use_full_cuda_graph:
315
- assert scheduler_metadata is not None
326
+ # For FA3 + full cudagraph
327
+ max_num_splits = 0
328
+ if self.use_full_cuda_graph and scheduler_metadata is not None:
329
n = scheduler_metadata.shape[0]
330
self.scheduler_metadata[:n] = scheduler_metadata
331
# NOTE(woosuk): We should zero out the rest of the scheduler
@@ -322,14 +335,12 @@ class FlashAttentionMetadataBuilder(
335
self.scheduler_metadata[n:] = 0
336
scheduler_metadata = self.scheduler_metadata[:n]
337
325
- max_num_splits = 0
326
- if (self.use_full_cuda_graph
327
- and num_actual_tokens <= self.max_cudagraph_size):
328
- # NOTE(woosuk): Setting num_splits > 1 may increase the memory
329
- # usage, because the intermediate buffers of size [num_splits,
330
- # num_heads, num_tokens, head_size] are allocated. Therefore,
331
- # we only set num_splits when using cuda graphs.
332
- max_num_splits = self.max_num_splits
338
+ if num_actual_tokens <= self.max_cudagraph_size:
339
+ # NOTE(woosuk): Setting num_splits > 1 may increase the memory
340
+ # usage, because the intermediate buffers of size [num_splits,
341
+ # num_heads, num_tokens, head_size] are allocated. Therefore,
342
+ # we only set num_splits when using cuda graphs.
343
+ max_num_splits = self.max_num_splits
344
345
attn_metadata = FlashAttentionMetadata(
346
num_actual_tokens=num_actual_tokens,
@@ -350,11 +361,6 @@ class FlashAttentionMetadataBuilder(
361
causal=causal)
362
return attn_metadata
363
353
- def can_run_in_cudagraph(
354
- self, common_attn_metadata: CommonAttentionMetadata) -> bool:
355
- # Full CUDA Graph always supported (FA2 support checked separately)
356
- return True
357
-
364
def use_cascade_attention(self, *args, **kwargs) -> bool:
365
return use_cascade_attention(*args, **kwargs)
366
vllm/v1/attention/backends/flashinfer.py
+5
-8
@@ -17,7 +17,7 @@ from flashinfer.prefill import trtllm_batch_context_with_kv_cache
17
import vllm.envs as envs
18
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
19
AttentionType)
20
-from vllm.config import VllmConfig
20
+from vllm.config import CUDAGraphMode, VllmConfig
21
from vllm.logger import init_logger
22
from vllm.utils import cdiv, is_pin_memory_available
23
from vllm.utils.flashinfer import use_trtllm_attention
@@ -183,8 +183,8 @@ class FlashInferMetadata:
183
184
185
class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
186
- attn_cudagraph_support: ClassVar[AttentionCGSupport] = \
187
- AttentionCGSupport.PURE_DECODE_ONLY
186
+ cudagraph_support: ClassVar[AttentionCGSupport] = \
187
+ AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
188
189
reorder_batch_threshold: ClassVar[int] = 1
190
@@ -203,7 +203,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
203
self.kv_cache_spec.block_size)
204
max_num_reqs = vllm_config.scheduler_config.max_num_seqs
205
max_num_pages = max_num_reqs * max_num_pages_per_req
206
- self.enable_cuda_graph = self.compilation_config.full_cuda_graph
206
+ self.enable_cuda_graph = self.compilation_config.cudagraph_mode.\
207
+ decode_mode() == CUDAGraphMode.FULL
208
if self.enable_cuda_graph:
209
# For full cudagraph capture, one `decode_wrapper` for each batch
210
# size is needed for FlashInfer.
@@ -586,10 +587,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
587
588
return self.build(0, m)
589
589
- def can_run_in_cudagraph(
590
- self, common_attn_metadata: CommonAttentionMetadata) -> bool:
591
- return common_attn_metadata.max_query_len == 1
592
-
590
def use_cascade_attention(self, *args, **kwargs) -> bool:
591
if self.kv_cache_spec.dtype != self.vllm_config.model_config.dtype:
592
# TODO: The cascade wrapper currently does not support setting
vllm/v1/attention/backends/mamba2_attn.py
+2
-6
@@ -89,8 +89,8 @@ class Mamba2AttentionMetadata:
89
90
class Mamba2AttentionMetadataBuilder(
91
AttentionMetadataBuilder[Mamba2AttentionMetadata]):
92
- attn_cudagraph_support: ClassVar[AttentionCGSupport] = \
93
- AttentionCGSupport.PURE_DECODE_ONLY
92
+ cudagraph_support: ClassVar[AttentionCGSupport] = \
93
+ AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
94
95
reorder_batch_threshold: ClassVar[int] = 1
96
@@ -203,7 +203,3 @@ class Mamba2AttentionMetadataBuilder(
203
m.max_query_len = 1 # decode-only
204
205
return self.build(0, m)
206
-
207
- def can_run_in_cudagraph(
208
- self, common_attn_metadata: CommonAttentionMetadata) -> bool:
209
- return common_attn_metadata.max_query_len == 1
vllm/v1/attention/backends/mla/common.py
+1
-5
@@ -575,7 +575,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
575
"MLA only supports decode-only full CUDAGraph capture. " \
576
"Make sure all cudagraph capture sizes <= max_num_seq."
577
578
- m.max_query_len = 1 # decode-only
578
+ assert m.max_query_len == 1 # decode-only
579
580
return self.build(0, m)
581
@@ -728,10 +728,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
728
729
return attn_metadata
730
731
- def can_run_in_cudagraph(
732
- self, common_attn_metadata: CommonAttentionMetadata) -> bool:
733
- return common_attn_metadata.max_query_len == 1
734
-
731
732
class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
733
"""
vllm/v1/attention/backends/mla/cutlass_mla.py
+1
-1
@@ -22,7 +22,7 @@ logger = init_logger(__name__)
22
class CutlassMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
23
# enable full CUDA Graph support for decode-only capture
24
attn_cudagraph_support: ClassVar[
25
- AttentionCGSupport] = AttentionCGSupport.PURE_DECODE_ONLY
25
+ AttentionCGSupport] = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
26
27
28
class CutlassMLABackend(MLACommonBackend):
vllm/v1/attention/backends/mla/flashmla.py
+7
-4
@@ -55,8 +55,8 @@ class FlashMLAMetadata(MLACommonMetadata[FlashMLADecodeMetadata]):
55
56
57
class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
58
- attn_cudagraph_support: ClassVar[AttentionCGSupport] = \
59
- AttentionCGSupport.PURE_DECODE_ONLY
58
+ cudagraph_support: ClassVar[AttentionCGSupport] = \
59
+ AttentionCGSupport.UNIFORM_BATCH
60
61
def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
62
vllm_config: VllmConfig, device: torch.device):
@@ -73,7 +73,7 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
73
device_properties = torch.cuda.get_device_properties(self.device)
74
num_sms = device_properties.multi_processor_count
75
76
- if self.compilation_config.full_cuda_graph:
76
+ if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
77
self.cg_buf_tile_scheduler_metadata = torch.zeros(
78
# Upper bound on size (<= #SMs, TileSchedulerMetaDataSize)
79
# TileSchedulerMetaDataSize = 8
@@ -95,7 +95,10 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
95
1, # MQA for the decode path
96
)
97
98
- if self.compilation_config.full_cuda_graph:
98
+ # TODO: we can disambiguate between decode and mixed-prefill decode here
99
+ # so we can only use the persistent buffer if a cudagraph is actually
100
+ # being used.
101
+ if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
102
assert self.cg_buf_tile_scheduler_metadata is not None
103
assert self.cg_buf_num_splits is not None
104
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+9
-4
@@ -65,8 +65,10 @@ class AiterMLAMetadata(MLACommonMetadata[AiterMLADecodeMetadata]):
65
66
67
class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
68
- attn_cudagraph_support: ClassVar[AttentionCGSupport] = \
69
- AttentionCGSupport.PURE_DECODE_ONLY
68
+ # TODO(luka, lucas): audit this as part of:
69
+ # https://github.com/vllm-project/vllm/issues/22945
70
+ cudagraph_support: ClassVar[AttentionCGSupport] = \
71
+ AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
72
73
def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
74
vllm_config: VllmConfig, device: torch.device):
@@ -82,7 +84,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
84
max_num_pages = max_num_reqs * max_num_pages_per_req
85
86
# Preparing persistent buffers
85
- if vllm_config.compilation_config.full_cuda_graph:
87
+ # TODO: we can disambiguate between decode and mixed-prefill decode here
88
+ # so we can only use the persistent buffer if a cudagraph is actually
89
+ # being used.
90
+ if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
91
self.paged_kv_indptr = torch.zeros(max_num_reqs + 1,
92
dtype=torch.int32,
93
device=device)
@@ -120,7 +125,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
125
block_table_bounds.cumsum(dim=0, dtype=torch.int32)
126
])
127
123
- if self.compilation_config.full_cuda_graph:
128
+ if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
129
130
num_actual_pages = paged_kv_indices.size(0)
131
vllm/v1/attention/backends/rocm_aiter_fa.py
-5
@@ -311,11 +311,6 @@ class AiterFlashAttentionMetadataBuilder(
311
)
312
return attn_metadata
313
314
- def can_run_in_cudagraph(
315
- self, common_attn_metadata: CommonAttentionMetadata) -> bool:
316
- # Full CUDA Graph always supported (FA2 support checked separately)
317
- return True
318
-
314
def use_cascade_attention(self, *args, **kwargs) -> bool:
315
return False
316
vllm/v1/attention/backends/triton_attn.py
+1
-7
@@ -58,8 +58,7 @@ class TritonAttentionMetadata:
58
59
class TritonAttentionMetadataBuilder(
60
AttentionMetadataBuilder[TritonAttentionMetadata]):
61
- attn_cudagraph_support: ClassVar[AttentionCGSupport] = \
62
- AttentionCGSupport.ALWAYS
61
+ cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
62
63
def __init__(self, kv_cache_spec: AttentionSpec, layer_names: list[str],
64
vllm_config: VllmConfig, device: torch.device):
@@ -132,11 +131,6 @@ class TritonAttentionMetadataBuilder(
131
)
132
return attn_metadata
133
135
- def can_run_in_cudagraph(
136
- self, common_attn_metadata: CommonAttentionMetadata) -> bool:
137
- # Full CUDA Graph always supported
138
- return True
139
-
134
135
class TritonAttentionBackend(AttentionBackend):
136
vllm/v1/attention/backends/utils.py
+10
-14
@@ -158,18 +158,21 @@ class AttentionCGSupport(enum.Enum):
158
Here we do not consider the cascade attention, as currently
159
it is never cudagraph supported."""
160
161
+ ALWAYS = 3
162
+ """Cudagraph always supported; supports mixed-prefill-decode"""
163
+ UNIFORM_BATCH = 2
164
+ """Cudagraph supported for batches the only contain query lengths that are
165
+ the same, this can be used for spec-decode
166
+ i.e. "decodes" are 1 + num_speculative_tokens"""
167
+ UNIFORM_SINGLE_TOKEN_DECODE = 1
168
+ """Cudagraph supported for batches the only contain query_len==1 decodes"""
169
NEVER = 0
170
"""NO cudagraph support"""
163
- PURE_DECODE_ONLY = 1
164
- """Cudagraph supported for pure decode, need to run without
165
- cudagraph for mixed prefill-decode batches"""
166
- ALWAYS = 2
167
- """Cudagraph always supported"""
171
172
173
class AttentionMetadataBuilder(abc.ABC, Generic[M]):
171
- # Does this backend/builder support CUDA Graphs for attention.
172
- attn_cudagraph_support: ClassVar[AttentionCGSupport] = \
174
+ # Does this backend/builder support CUDA Graphs for attention (default: no).
175
+ cudagraph_support: ClassVar[AttentionCGSupport] = \
176
AttentionCGSupport.NEVER
177
# Does this backend/builder reorder the batch?
178
# If not, set this to None. Otherwise set it to the query
@@ -199,13 +202,6 @@ class AttentionMetadataBuilder(abc.ABC, Generic[M]):
202
"""
203
raise NotImplementedError
204
202
- def can_run_in_cudagraph(
203
- self, common_attn_metadata: CommonAttentionMetadata) -> bool:
204
- """
205
- Can this batch (with given metadata) use CUDA Graphs for attention.
206
- """
207
- return False
208
-
205
def build_for_cudagraph_capture(
206
self, common_attn_metadata: CommonAttentionMetadata) -> M:
207
"""
vllm/v1/cudagraph_dispatcher.py
new
+120
@@ -0,0 +1,120 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Optional
4
+
5
+from vllm.config import CompilationLevel, CUDAGraphMode, VllmConfig
6
+from vllm.forward_context import BatchDescriptor
7
+from vllm.logger import init_logger
8
+
9
+logger = init_logger(__name__)
10
+
11
+
12
+class CudagraphDispatcher:
13
+ """
14
+ Runtime cudagraph dispatcher to dispach keys for multiple set of cudagraphs.
15
+
16
+ The dispatcher stores two sets of dispatch keys, one for PIECEWISE and one
17
+ for FULL cudagraph runtime mode. The keys are initialized depending on
18
+ attention support and what cudagraph mode is set in CompilationConfig. The
19
+ keys stored in dispatcher are the only source of truth for valid
20
+ cudagraphs that can be dispatched at runtime.
21
+
22
+ At runtime, the dispatch method generates the runtime cudagraph mode (FULL,
23
+ PIECEWISE, or NONE for no cudagraph) and the valid key (batch descriptor)
24
+ based on the input key. After dispatching (commuicate via forward context),
25
+ the cudagraph wrappers will trust the dispatch key to do either capturing
26
+ or replaying (if mode matched), or pass through to the underlying runnable
27
+ without cudagraph (if mode no match or mode is NONE).
28
+ """
29
+
30
+ def __init__(self, vllm_config: VllmConfig):
31
+ self.vllm_config = vllm_config
32
+ self.compilation_config = vllm_config.compilation_config
33
+ self.cudagraph_mode = self.compilation_config.cudagraph_mode
34
+
35
+ # Dict to store valid cudagraph dispatching keys.
36
+ self.cudagraph_keys: dict[CUDAGraphMode, set[BatchDescriptor]] = {
37
+ CUDAGraphMode.PIECEWISE: set(),
38
+ CUDAGraphMode.FULL: set(),
39
+ }
40
+
41
+ assert not self.cudagraph_mode.requires_piecewise_compilation() or \
42
+ (self.compilation_config.level == CompilationLevel.PIECEWISE and
43
+ self.compilation_config.splitting_ops_contain_attention()), \
44
+ "Compilation level should be CompilationLevel.PIECEWISE when "\
45
+ "cudagraph_mode piecewise cudagraphs is used, "\
46
+ f"cudagraph_mode={self.cudagraph_mode}, "\
47
+ f"compilation_level={self.compilation_config.level}, "\
48
+ f"splitting_ops={self.compilation_config.splitting_ops}"
49
+
50
+ self.keys_initialized = False
51
+
52
+ def add_cudagraph_key(self, runtime_mode: CUDAGraphMode,
53
+ batch_descriptor: BatchDescriptor):
54
+ assert runtime_mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL], \
55
+ f"Invalid cudagraph runtime mode: {runtime_mode}"
56
+ self.cudagraph_keys[runtime_mode].add(batch_descriptor)
57
+
58
+ def initialize_cudagraph_keys(self, cudagraph_mode: CUDAGraphMode,
59
+ uniform_decode_query_len: int):
60
+ # This should be called only after attention backend is initialized.
61
+
62
+ # Note: we create all valid keys possible for cudagraph but do not
63
+ # guarantee all keys would be used. For example, we create keys for
64
+ # piecewise cudagraphs when it is piecewise compilation, which is always
65
+ # valid, but for attention backend support unified routine, we may not
66
+ # trigger capturing/replaying the piecewise cudagraphs depending on
67
+ # CompilationConfig.cudagraph_mode. In addition, if we allow lazy
68
+ # capturing in future PR, some keys may never be triggered.
69
+ if cudagraph_mode.mixed_mode() != CUDAGraphMode.NONE:
70
+ for bs in self.compilation_config.cudagraph_capture_sizes:
71
+ self.add_cudagraph_key(
72
+ cudagraph_mode.mixed_mode(),
73
+ BatchDescriptor(num_tokens=bs, uniform_decode=False))
74
+
75
+ # if decode cudagraph mode is FULL, and we don't already have mixed
76
+ # mode full cudagraphs then add them here.
77
+ if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL \
78
+ and cudagraph_mode.separate_routine():
79
+ max_num_tokens = uniform_decode_query_len * \
80
+ self.vllm_config.scheduler_config.max_num_seqs
81
+ cudagraph_capture_sizes_for_decode = [
82
+ x for x in self.compilation_config.cudagraph_capture_sizes
83
+ if x <= max_num_tokens and x >= uniform_decode_query_len
84
+ ]
85
+ for bs in cudagraph_capture_sizes_for_decode:
86
+ self.add_cudagraph_key(
87
+ CUDAGraphMode.FULL,
88
+ BatchDescriptor(num_tokens=bs, uniform_decode=True))
89
+ self.keys_initialized = True
90
+
91
+ def dispatch(
92
+ self, batch_descriptor: BatchDescriptor
93
+ ) -> tuple[CUDAGraphMode, Optional[BatchDescriptor]]:
94
+ """
95
+ Given a batch descriptor, dispatch to a cudagraph mode.
96
+ A new batch descriptor is returned as we might dispatch a uniform batch
97
+ to a graph that supports a more general batch (uniform to non-uniform).
98
+ """
99
+ # if not initialized, just skip dispatching.
100
+ if not self.keys_initialized:
101
+ logger.warning_once("cudagraph dispatching keys are not "
102
+ "initialized. No cudagraph will be used.")
103
+ return CUDAGraphMode.NONE, None
104
+
105
+ # check if key exists for full cudagraph
106
+ if batch_descriptor in self.cudagraph_keys[CUDAGraphMode.FULL]:
107
+ return CUDAGraphMode.FULL, batch_descriptor
108
+
109
+ # otherwise, check if non-uniform key exists
110
+ non_uniform_key = batch_descriptor.non_uniform
111
+ if non_uniform_key in self.cudagraph_keys[CUDAGraphMode.FULL]:
112
+ return CUDAGraphMode.FULL, non_uniform_key
113
+
114
+ # also check if non-uniform key exists for more "general"
115
+ # piecewise cudagraph
116
+ if non_uniform_key in self.cudagraph_keys[CUDAGraphMode.PIECEWISE]:
117
+ return CUDAGraphMode.PIECEWISE, non_uniform_key
118
+
119
+ # finally, just return no cudagraphs
120
+ return CUDAGraphMode.NONE, None
vllm/v1/worker/gpu_model_runner.py
+270
-89
@@ -21,7 +21,9 @@ from vllm.attention import Attention, AttentionType
21
from vllm.attention.backends.abstract import AttentionBackend
22
from vllm.attention.layers.chunked_local_attention import ChunkedLocalAttention
23
from vllm.compilation.counter import compilation_counter
24
-from vllm.config import (CompilationLevel, VllmConfig,
24
+from vllm.compilation.cuda_graph import CUDAGraphWrapper
25
+from vllm.compilation.monitor import set_cudagraph_capturing_enabled
26
+from vllm.config import (CompilationLevel, CUDAGraphMode, VllmConfig,
27
get_layers_from_vllm_config, update_config)
28
from vllm.distributed.eplb.eplb_state import EplbState
29
from vllm.distributed.kv_transfer import (get_kv_transfer_group,
@@ -29,7 +31,8 @@ from vllm.distributed.kv_transfer import (get_kv_transfer_group,
31
from vllm.distributed.parallel_state import (
32
get_pp_group, get_tp_group, graph_capture, is_global_first_rank,
33
prepare_communication_buffer_for_model)
32
-from vllm.forward_context import DPMetadata, set_forward_context
34
+from vllm.forward_context import (BatchDescriptor, DPMetadata,
35
+ set_forward_context)
36
from vllm.logger import init_logger
37
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaBase
38
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
@@ -48,13 +51,15 @@ from vllm.sampling_params import SamplingType
51
from vllm.sequence import IntermediateTensors, PoolerOutput
52
from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
53
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,
51
- GiB_bytes, LazyLoader, check_use_alibi, get_dtype_size,
52
- is_pin_memory_available, round_up, supports_dynamo)
54
+ GiB_bytes, LazyLoader, cdiv, check_use_alibi,
55
+ get_dtype_size, is_pin_memory_available, round_up,
56
+ supports_dynamo)
57
from vllm.v1.attention.backends.mamba_selectors import get_mamba_attn_backend
58
from vllm.v1.attention.backends.utils import (
59
AttentionCGSupport, AttentionMetadataBuilder, CommonAttentionMetadata,
60
make_kv_sharing_fast_prefill_attention_metadata,
61
reorder_batch_to_split_decodes_and_prefills)
62
+from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
63
from vllm.v1.kv_cache_interface import (AttentionSpec,
64
ChunkedLocalAttentionSpec,
65
FullAttentionSpec, KVCacheConfig,
@@ -218,11 +223,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
223
is_spec_decode=bool(self.vllm_config.speculative_config),
224
)
225
221
- self.use_cuda_graph = (
222
- self.vllm_config.compilation_config.level
223
- == CompilationLevel.PIECEWISE
224
- and self.vllm_config.compilation_config.use_cudagraph
225
- and not self.model_config.enforce_eager)
226
# TODO(woosuk): Provide an option to tune the max cudagraph batch size.
227
# The convention is different.
228
# self.cudagraph_batch_sizes sorts in ascending order.
@@ -230,8 +230,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
230
self.cudagraph_batch_sizes = list(
231
reversed(self.compilation_config.cudagraph_capture_sizes))
232
233
- self.full_cuda_graph = self.compilation_config.full_cuda_graph
234
-
233
# Cache the device properties.
234
self._init_device_properties()
235
@@ -326,6 +324,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
324
self.kv_sharing_fast_prefill_logits_indices = torch.zeros(
325
self.max_num_tokens, dtype=torch.int32, device=self.device)
326
327
+ self.uniform_decode_query_len = 1 if not self.speculative_config else \
328
+ 1 + self.speculative_config.num_speculative_tokens
329
+
330
+ # Cudagraph dispatcher for runtime cudagraph dispatching.
331
+ self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)
332
+
333
self.mm_budget = (MultiModalBudget(
334
self.model_config,
335
self.scheduler_config,
@@ -471,7 +475,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
475
assert (task := pooling_params.task) is not None, (
476
"You did not set `task` in the API")
477
474
- model = cast(VllmModelForPooling, self.model)
478
+ model = cast(VllmModelForPooling, self.get_model())
479
to_update = model.pooler.get_pooling_updates(task)
480
to_update.apply(pooling_params)
481
@@ -679,13 +683,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
683
def _prepare_inputs(
684
self,
685
scheduler_output: "SchedulerOutput",
682
- ) -> tuple[dict[str,
683
- Any], bool, torch.Tensor, Optional[SpecDecodeMetadata],
684
- np.ndarray, Optional[CommonAttentionMetadata]]:
686
+ ) -> tuple[dict[str, Any], torch.Tensor, Optional[SpecDecodeMetadata],
687
+ np.ndarray, Optional[CommonAttentionMetadata], int]:
688
"""
689
:return: tuple[
690
attn_metadata: layer-to-attention_metadata mapping,
688
- attention_cuda_graphs: whether attention can run in cudagraph
691
logits_indices, spec_decode_metadata
692
]
693
"""
@@ -820,7 +822,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
822
# valid, we fill the padded indices with the last index.
823
self.kv_sharing_fast_prefill_logits_indices[num_logits:].fill_(
824
logits_indices[-1].item())
823
- if (self.use_cuda_graph
825
+ if (self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
826
and num_logits <= self.cudagraph_batch_sizes[-1]):
827
# Use piecewise CUDA graphs.
828
# Add padding to the batch size.
@@ -925,17 +927,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
927
continue
928
attn_metadata[layer_name] = attn_metadata_i
929
928
- attention_cuda_graphs = all(
929
- g.metadata_builder.can_run_in_cudagraph(common_attn_metadata)
930
- for g in self._attn_group_iterator())
931
-
930
# Hot-Swap lora model
931
if self.lora_config:
932
self.set_active_loras(self.input_batch, num_scheduled_tokens)
933
936
- return (attn_metadata, attention_cuda_graphs, logits_indices,
937
- spec_decode_metadata, num_scheduled_tokens,
938
- spec_decode_common_attn_metadata)
934
+ return (attn_metadata, logits_indices, spec_decode_metadata,
935
+ num_scheduled_tokens, spec_decode_common_attn_metadata,
936
+ max_num_scheduled_tokens)
937
938
def _compute_cascade_attn_prefix_len(
939
self,
@@ -1259,6 +1257,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1257
return mm_embeds
1258
1259
def get_model(self) -> nn.Module:
1260
+ # get raw model out of the cudagraph wrapper.
1261
+ if isinstance(self.model, CUDAGraphWrapper):
1262
+ return self.model.unwrap()
1263
return self.model
1264
1265
def get_supported_generation_tasks(self) -> list[GenerationTask]:
@@ -1415,9 +1416,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1416
return
1417
1418
assert self.eplb_state is not None
1418
- assert is_mixture_of_experts(self.model)
1419
+ model = self.get_model()
1420
+ assert is_mixture_of_experts(model)
1421
self.eplb_state.step(
1420
- self.model,
1422
+ model,
1423
is_dummy,
1424
is_profile,
1425
log_stats=self.parallel_config.eplb_log_balancedness,
@@ -1507,15 +1509,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1509
self.vllm_config)
1510
1511
# Prepare the decoder inputs.
1510
- (attn_metadata, attention_cuda_graphs, logits_indices,
1511
- spec_decode_metadata, num_scheduled_tokens_np,
1512
- spec_decode_common_attn_metadata) = (
1513
- self._prepare_inputs(scheduler_output))
1512
+ (attn_metadata, logits_indices, spec_decode_metadata,
1513
+ num_scheduled_tokens_np, spec_decode_common_attn_metadata,
1514
+ max_query_len) = (self._prepare_inputs(scheduler_output))
1515
1516
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
1516
- if (self.use_cuda_graph
1517
+ if (self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
1518
and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
1518
- # Use piecewise CUDA graphs.
1519
+ # Use CUDA graphs.
1520
# Add padding to the batch size.
1521
num_input_tokens = self.vllm_config.pad_for_cudagraph(
1522
num_scheduled_tokens)
@@ -1581,10 +1582,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1582
intermediate_tensors = self.sync_and_slice_intermediate_tensors(
1583
num_input_tokens, intermediate_tensors, True)
1584
1584
- # Some attention backends only support CUDA Graphs in pure decode.
1585
- # If attention doesn't support CUDA Graphs for this batch, but we
1586
- # compiled with full CUDA graphs, we have to skip them entirely.
1587
- skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
1585
+ uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
1586
+ num_scheduled_tokens == self.input_batch.num_reqs * max_query_len)
1587
+ batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens,
1588
+ uniform_decode=uniform_decode)
1589
+ cudagraph_runtime_mode, batch_descriptor = \
1590
+ self.cudagraph_dispatcher.dispatch(batch_descriptor)
1591
1592
# Run the model.
1593
# Use persistent buffers for CUDA graphs.
@@ -1593,10 +1596,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1596
self.vllm_config,
1597
num_tokens=num_input_tokens,
1598
num_tokens_across_dp=num_tokens_across_dp,
1596
- skip_cuda_graphs=skip_cuda_graphs,
1599
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
1600
+ batch_descriptor=batch_descriptor,
1601
), self.maybe_get_kv_connector_output(
1602
scheduler_output) as kv_connector_output:
1599
-
1603
model_output = self.model(
1604
input_ids=input_ids,
1605
positions=positions,
@@ -2021,20 +2024,31 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2024
self.model.compile(
2025
fullgraph=envs.VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE,
2026
backend=backend)
2027
+ return
2028
+ # for other compilation levels, cudagraph behavior is controlled by
2029
+ # CudagraphWraper and CudagraphDispatcher of vllm.
2030
+
2031
+ # wrap the model with full cudagraph wrapper if needed.
2032
+ if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
2033
+ self.model = CUDAGraphWrapper(self.model,
2034
+ self.vllm_config,
2035
+ runtime_mode=CUDAGraphMode.FULL)
2036
2037
def reload_weights(self) -> None:
2038
assert getattr(self, "model", None) is not None, \
2039
"Cannot reload weights before model is loaded."
2040
model_loader = get_model_loader(self.load_config)
2041
logger.info("Reloading weights inplace...")
2030
- model_loader.load_weights(self.model, model_config=self.model_config)
2042
+ model = self.get_model()
2043
+ model_loader.load_weights(model, model_config=self.model_config)
2044
2045
def save_tensorized_model(
2046
self,
2047
tensorizer_config: "TensorizerConfig",
2048
) -> None:
2049
+ model = self.get_model()
2050
TensorizerLoader.save_model(
2037
- self.model,
2051
+ model,
2052
tensorizer_config=tensorizer_config,
2053
model_config=self.model_config,
2054
)
@@ -2210,31 +2224,82 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2224
def _dummy_run(
2225
self,
2226
num_tokens: int,
2213
- capture_attn_cudagraph: bool = False,
2227
+ cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
2228
+ force_attention: bool = False,
2229
+ uniform_decode: bool = False,
2230
skip_eplb: bool = False,
2231
is_profile: bool = False,
2232
) -> tuple[torch.Tensor, torch.Tensor]:
2233
+ """
2234
+ Run a dummy forward pass to warm up/profile run or capture the
2235
+ CUDA graph for the model.
2236
+
2237
+ Args:
2238
+ num_tokens: Number of tokens to run the dummy forward pass.
2239
+ cudagraph_runtime_mode: used to control the behavior.
2240
+ - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run
2241
+ - CUDAGraphMode.PIECEWISE: Piecewise cudagraph.
2242
+ - CUDAGraphMode.FULL: Full cudagraph, attention metadata is
2243
+ needed.
2244
+ force_attention: If True, always create attention metadata. Used to
2245
+ warm up attention backend when mode is NONE.
2246
+ uniform_decode: If True, the batch is a uniform decode batch.
2247
+ skip_eplb: If True, skip EPLB state update.
2248
+ is_profile: If True, this is a profile run.
2249
+ """
2250
+ assert cudagraph_runtime_mode in {
2251
+ CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL
2252
+ }
2253
2254
# Padding for DP
2255
num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
2256
num_tokens += num_pad
2257
2258
+ # If cudagraph_mode.decode_mode() == FULL and
2259
+ # cudagraph_mode.seperate_routine(). This means that we are using
2260
+ # different graphs and/or modes for mixed prefill-decode batches vs.
2261
+ # uniform decode batches. A uniform decode batch means that all
2262
+ # requests have identical query length, except a potential virtual
2263
+ # request (shorter) in the batch account for padding.
2264
+ # Uniform decode batch could either be common pure decode, where
2265
+ # max_query_len == 1, or speculative decode, where
2266
+ # max_query_len == 1 + num_spec_decode_tokens.
2267
+
2268
+ # When setting max_query_len = 1, we switch to and capture the optimized
2269
+ # routine of FA2 for pure decode, i.e., Flashdecode + an optimization
2270
+ # for GQA/MQA.
2271
+ max_query_len = self.uniform_decode_query_len if uniform_decode else \
2272
+ num_tokens
2273
+
2274
# Set num_scheduled_tokens based on num_tokens and max_num_seqs
2275
# for dummy run with LoRA so that the num_reqs collectively
2276
# has num_tokens in total.
2277
assert num_tokens <= self.scheduler_config.max_num_batched_tokens
2278
max_num_reqs = self.scheduler_config.max_num_seqs
2227
- num_reqs = min(num_tokens, max_num_reqs)
2228
- min_tokens_per_req = num_tokens // num_reqs
2229
- num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
2230
- num_scheduled_tokens_list[-1] += num_tokens % num_reqs
2279
+ if uniform_decode:
2280
+ num_reqs = cdiv(num_tokens, max_query_len)
2281
+ assert num_reqs <= max_num_reqs, \
2282
+ "Do not capture num_reqs > max_num_reqs for uniform batch"
2283
+ num_scheduled_tokens_list = [max_query_len] * num_reqs
2284
+ if num_tokens % max_query_len != 0:
2285
+ num_scheduled_tokens_list[-1] = num_tokens % max_query_len
2286
+ else:
2287
+ num_reqs = min(num_tokens, max_num_reqs)
2288
+ min_tokens_per_req = num_tokens // num_reqs
2289
+ num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
2290
+ num_scheduled_tokens_list[-1] += num_tokens % num_reqs
2291
+
2292
assert sum(num_scheduled_tokens_list) == num_tokens
2293
assert len(num_scheduled_tokens_list) == num_reqs
2294
num_scheduled_tokens = np.array(num_scheduled_tokens_list,
2295
dtype=np.int32)
2296
2297
attn_metadata: Optional[dict[str, Any]] = None
2237
- if capture_attn_cudagraph:
2298
+
2299
+ # If force_attention is True, we always capture attention. Otherwise,
2300
+ # it only happens for cudagraph_runtime_mode=FULL.
2301
+ if force_attention or cudagraph_runtime_mode == \
2302
+ CUDAGraphMode.FULL:
2303
attn_metadata = {}
2304
2305
# Make sure max_model_len is used at the graph capture time.
@@ -2255,7 +2320,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2320
num_computed_tokens_cpu_tensor[:num_reqs],
2321
num_reqs=num_reqs,
2322
num_actual_tokens=num_tokens,
2258
- max_query_len=num_tokens,
2323
+ max_query_len=max_query_len,
2324
block_table_tensor=self.input_batch.block_table[
2325
kv_cache_group_id].get_device_tensor()[:num_reqs],
2326
slot_mapping=self.input_batch.
@@ -2299,12 +2364,26 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2364
2365
intermediate_tensors = self.sync_and_slice_intermediate_tensors(
2366
num_tokens, None, False)
2367
+ if cudagraph_runtime_mode == CUDAGraphMode.NONE:
2368
+ batch_descriptor = None
2369
+ else:
2370
+ # filter out the valid batch descriptor
2371
+ _cg_mode, batch_descriptor = \
2372
+ self.cudagraph_dispatcher.dispatch(
2373
+ BatchDescriptor(num_tokens=num_tokens,
2374
+ uniform_decode=uniform_decode))
2375
+ # sanity check
2376
+ assert cudagraph_runtime_mode == _cg_mode, (
2377
+ f"Cudagraph runtime mode mismatch at dummy_run. "
2378
+ f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}.")
2379
2380
with self.maybe_randomize_inputs(input_ids), set_forward_context(
2381
attn_metadata,
2382
self.vllm_config,
2383
num_tokens=num_tokens,
2307
- num_tokens_across_dp=num_tokens_across_dp):
2384
+ num_tokens_across_dp=num_tokens_across_dp,
2385
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
2386
+ batch_descriptor=batch_descriptor):
2387
outputs = self.model(
2388
input_ids=input_ids,
2389
positions=positions,
@@ -2436,7 +2515,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2515
dtype=torch.int32,
2516
device=self.device)
2517
2439
- model = cast(VllmModelForPooling, self.model)
2518
+ model = cast(VllmModelForPooling, self.get_model())
2519
dummy_pooling_params = PoolingParams(task=task)
2520
to_update = model.pooler.get_pooling_updates(task)
2521
to_update.apply(dummy_pooling_params)
@@ -2546,12 +2625,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2625
gc.collect()
2626
2627
def capture_model(self) -> None:
2549
- if not self.use_cuda_graph:
2628
+ if self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE:
2629
logger.warning(
2630
"Skipping CUDA graph capture. To turn on CUDA graph capture, "
2552
- "set -O %s and ensure `use_cudagraph` was not manually set to "
2553
- "False", CompilationLevel.PIECEWISE)
2631
+ "ensure `cudagraph_mode` was not manually set to `NONE`")
2632
return
2633
+ else:
2634
+ self.initialize_cudagraph_capture()
2635
2636
compilation_counter.num_gpu_runner_capture_triggers += 1
2637
@@ -2576,25 +2656,41 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2656
# Trigger CUDA graph capture for specific shapes.
2657
# Capture the large shapes first so that the smaller shapes
2658
# can reuse the memory pool allocated for the large shapes.
2659
+ set_cudagraph_capturing_enabled(True)
2660
with freeze_gc(), graph_capture(device=self.device):
2580
- full_cg = self.full_cuda_graph
2581
- # Only rank 0 should print progress bar during capture
2582
- compilation_cases = reversed(self.cudagraph_batch_sizes)
2583
- if is_global_first_rank():
2584
- compilation_cases = tqdm(
2585
- list(compilation_cases),
2586
- disable=not self.load_config.use_tqdm_on_load,
2587
- desc="Capturing CUDA graph shapes")
2588
- for num_tokens in compilation_cases:
2589
- # We skip EPLB here since we don't want to record dummy metrics
2590
- for _ in range(
2591
- self.compilation_config.cudagraph_num_of_warmups):
2592
- self._dummy_run(num_tokens,
2593
- capture_attn_cudagraph=full_cg,
2594
- skip_eplb=True)
2595
- self._dummy_run(num_tokens,
2596
- capture_attn_cudagraph=full_cg,
2597
- skip_eplb=True)
2661
+ cudagraph_mode = self.compilation_config.cudagraph_mode
2662
+ if cudagraph_mode.mixed_mode() != CUDAGraphMode.NONE:
2663
+ cudagraph_runtime_mode = cudagraph_mode.mixed_mode()
2664
+
2665
+ compilation_cases = list(reversed(self.cudagraph_batch_sizes))
2666
+ self._capture_cudagraphs(
2667
+ compilation_cases,
2668
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
2669
+ uniform_decode=False)
2670
+
2671
+ # Capture full cudagraph for uniform decode batches if we have
2672
+ # dont already have full mixed prefill-decode cudagraphs
2673
+ if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL and \
2674
+ cudagraph_mode.separate_routine():
2675
+ max_num_tokens = self.scheduler_config.max_num_seqs * \
2676
+ self.uniform_decode_query_len
2677
+ decode_cudagraph_batch_sizes = [
2678
+ x for x in self.cudagraph_batch_sizes if
2679
+ x <= max_num_tokens and x >= self.uniform_decode_query_len
2680
+ ]
2681
+ compilation_cases_decode = list(
2682
+ reversed(decode_cudagraph_batch_sizes))
2683
+ self._capture_cudagraphs(
2684
+ compilation_cases=compilation_cases_decode,
2685
+ cudagraph_runtime_mode=CUDAGraphMode.FULL,
2686
+ uniform_decode=True)
2687
+
2688
+ # Disable cudagraph capturing globally, so any unexpected cudagraph
2689
+ # capturing will be detected and raise an error after here.
2690
+ # Note: We don't put it into graph_capture context manager because
2691
+ # we may doing lazy capturing in future that still allows capturing
2692
+ # after here.
2693
+ set_cudagraph_capturing_enabled(False)
2694
2695
end_time = time.perf_counter()
2696
end_free_gpu_memory = torch.cuda.mem_get_info()[0]
@@ -2604,6 +2700,41 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2700
logger.info("Graph capturing finished in %.0f secs, took %.2f GiB",
2701
elapsed_time, cuda_graph_size / (1 << 30))
2702
2703
+ def _capture_cudagraphs(self, compilation_cases: list[int],
2704
+ cudagraph_runtime_mode: CUDAGraphMode,
2705
+ uniform_decode: bool):
2706
+ assert cudagraph_runtime_mode != CUDAGraphMode.NONE and \
2707
+ cudagraph_runtime_mode in [CUDAGraphMode.FULL,
2708
+ CUDAGraphMode.PIECEWISE]
2709
+
2710
+ # Only rank 0 should print progress bar during capture
2711
+ if is_global_first_rank():
2712
+ compilation_cases = tqdm(
2713
+ compilation_cases,
2714
+ disable=not self.load_config.use_tqdm_on_load,
2715
+ desc="Capturing CUDA graphs ({}, {})".format(
2716
+ "decode" if uniform_decode else "mixed prefill-decode",
2717
+ cudagraph_runtime_mode.name))
2718
+ # We skip EPLB here since we don't want to record dummy metrics
2719
+ for num_tokens in compilation_cases:
2720
+ for _ in range(self.compilation_config.cudagraph_num_of_warmups):
2721
+ # Use CUDAGraphRuntimeStyle.NONE (default) for warmup.
2722
+ # But be careful, warm up with `NONE`is orthogonal to
2723
+ # if we want to warm up attention or not. This is
2724
+ # different from the case where `FULL` implies capture
2725
+ # attention while `PIECEWISE` implies no attention.
2726
+ force_attention = (
2727
+ cudagraph_runtime_mode == CUDAGraphMode.FULL)
2728
+ self._dummy_run(num_tokens,
2729
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
2730
+ force_attention=force_attention,
2731
+ uniform_decode=uniform_decode,
2732
+ skip_eplb=True)
2733
+ self._dummy_run(num_tokens,
2734
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
2735
+ uniform_decode=uniform_decode,
2736
+ skip_eplb=True)
2737
+
2738
def initialize_attn_backend(self, kv_cache_config: KVCacheConfig) -> None:
2739
"""
2740
Initialize the attention backends and attention metadata builders.
@@ -2648,25 +2779,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2779
attn_metadata_builder_i,
2780
layer_names)
2781
attn_groups.append(attn_group)
2651
-
2652
- if self.full_cuda_graph:
2653
- if attn_metadata_builder_i.attn_cudagraph_support == \
2654
- AttentionCGSupport.NEVER:
2655
- raise ValueError(
2656
- f"Full CUDAGraph not supported for "
2657
- f"{attn_backend.__name__}. Turn off "
2658
- f"CompilationConfig.full_cuda_graph or use a "
2659
- f" different attention backend.")
2660
- if attn_metadata_builder_i.attn_cudagraph_support == \
2661
- AttentionCGSupport.PURE_DECODE_ONLY:
2662
- # Limit the max cudagraph size to the max number of
2663
- # sequences for pure decode only cudagraph backend,
2664
- # whose max_query_len is 1.
2665
- self.cudagraph_batch_sizes = [
2666
- size for size in self.cudagraph_batch_sizes
2667
- if size <= self.scheduler_config.max_num_seqs
2668
- ]
2669
-
2782
return attn_groups
2783
2784
for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
@@ -2734,6 +2846,75 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2846
"All or none of the layers are expected to be encoder-only"
2847
self.is_encoder_only_model = True
2848
2849
+ def initialize_cudagraph_capture(self) -> None:
2850
+ min_cg_support = AttentionCGSupport.ALWAYS
2851
+ min_cg_builder_name = None
2852
+
2853
+ for attn_group in self._attn_group_iterator():
2854
+ builder = attn_group.metadata_builder
2855
+ if builder.cudagraph_support.value < min_cg_support.value:
2856
+ min_cg_support = builder.cudagraph_support
2857
+ min_cg_builder_name = builder.__class__.__name__
2858
+
2859
+ # Flexible resolve the cudagraph mode
2860
+ cudagraph_mode = self.compilation_config.cudagraph_mode
2861
+ # check cudagraph for mixed batch is supported
2862
+ if cudagraph_mode.mixed_mode() == CUDAGraphMode.FULL \
2863
+ and min_cg_support != AttentionCGSupport.ALWAYS:
2864
+ msg = (f"CUDAGraphMode.{cudagraph_mode.name} is not supported "
2865
+ f"with {min_cg_builder_name} backend (support: "
2866
+ f"{min_cg_support})")
2867
+ if min_cg_support == AttentionCGSupport.NEVER:
2868
+ # if not supported any full cudagraphs, just raise it.
2869
+ msg += "; please try cudagraph_mode=PIECEWISE, and "\
2870
+ "make sure compilation level is piecewise"
2871
+ raise ValueError(msg)
2872
+
2873
+ # attempt to resolve the full cudagraph related mode
2874
+ if self.compilation_config.splitting_ops_contain_attention():
2875
+ msg += "; setting cudagraph_mode=FULL_AND_PIECEWISE"
2876
+ cudagraph_mode = self.compilation_config.cudagraph_mode = \
2877
+ CUDAGraphMode.FULL_AND_PIECEWISE
2878
+ else:
2879
+ msg += "; setting cudagraph_mode=FULL_DECODE_ONLY"
2880
+ cudagraph_mode = self.compilation_config.cudagraph_mode = \
2881
+ CUDAGraphMode.FULL_DECODE_ONLY
2882
+ logger.warning(msg)
2883
+
2884
+ # check that if we are doing spec-decode + decode full-cudagraphs it is
2885
+ # supported
2886
+ if (cudagraph_mode.decode_mode() == CUDAGraphMode.FULL
2887
+ and self.uniform_decode_query_len > 1 and min_cg_support.value
2888
+ < AttentionCGSupport.UNIFORM_BATCH.value):
2889
+ msg = (f"CUDAGraphMode.{cudagraph_mode.name} is not supported"
2890
+ f" with spec-decode for attention backend "
2891
+ f"{min_cg_builder_name} (support: {min_cg_support})")
2892
+ if self.compilation_config.splitting_ops_contain_attention():
2893
+ msg += "; setting cudagraph_mode=PIECEWISE"
2894
+ cudagraph_mode = self.compilation_config.cudagraph_mode = \
2895
+ CUDAGraphMode.PIECEWISE
2896
+ else:
2897
+ msg += "; setting cudagraph_mode=NONE"
2898
+ cudagraph_mode = self.compilation_config.cudagraph_mode = \
2899
+ CUDAGraphMode.NONE
2900
+ logger.warning(msg)
2901
+
2902
+ # double check that we can support full cudagraph if they are requested
2903
+ # even after automatic downgrades
2904
+ if cudagraph_mode.has_full_cudagraphs() \
2905
+ and min_cg_support == AttentionCGSupport.NEVER:
2906
+ raise ValueError(f"CUDAGraphMode.{cudagraph_mode.name} is not "
2907
+ f"supported with {min_cg_builder_name} backend ("
2908
+ f"support:{min_cg_support}) "
2909
+ "; please try cudagraph_mode=PIECEWISE, "
2910
+ "and make sure compilation level is piecewise")
2911
+
2912
+ # Trigger cudagraph dispatching keys initialization here (after
2913
+ # initializing attn backends).
2914
+ self.cudagraph_dispatcher.initialize_cudagraph_keys(
2915
+ self.compilation_config.cudagraph_mode,
2916
+ self.uniform_decode_query_len)
2917
+
2918
def calculate_reorder_batch_threshold(self) -> None:
2919
"""
2920
Check that if any backends reorder batches; that the reordering
vllm/v1/worker/gpu_worker.py
-5
@@ -322,16 +322,11 @@ class Worker(WorkerBase):
322
if get_pp_group().is_last_rank:
323
max_num_reqs = min(self.scheduler_config.max_num_seqs,
324
self.scheduler_config.max_num_batched_tokens)
325
- # activate building attn_metadata for this dummy run to avoid
326
- # potential illegal memory access for full cudagraph relay.
327
- attn_cudagraph = self.compilation_config.full_cuda_graph and\
328
- not self.model_config.enforce_eager
325
326
# We skip EPLB here since we don't want to record dummy metrics
327
hidden_states, last_hidden_states = \
328
self.model_runner._dummy_run(
329
num_tokens=max_num_reqs,
334
- capture_attn_cudagraph=attn_cudagraph,
330
skip_eplb=True,
331
)
332
if self.model_runner.is_pooling_model: