A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 11:00 UTC
[Core] Support full cuda graph in v1 (#16072)
Signed-off-by: Chanh Nguyen <cnguyen@linkedin.com> Co-authored-by: Chanh Nguyen <cnguyen@linkedin.com>
Chanh Nguyen committed
May 7, 2025 at 22:30 UTC
7ea2adb8026ec1213727a315a226b51b030b7af5
5 files changed
+190
-13
docs/source/design/v1/torch_compile.md
+6
@@ -137,3 +137,9 @@ By default, vLLM will try to determine a set of sizes to capture cudagraph. You
137
`vllm serve meta-llama/Llama-3.2-1B --compilation-config "{'cudagraph_capture_sizes': [1, 2, 4, 8]}"`
138
139
Then it will only capture cudagraph for the specified sizes. It can be useful to have fine-grained control over the cudagraph capture.
140
+
141
+### Full Cudagraph capture
142
+
143
+It is possible to include attention as part of the cudagraph if using an attention backend that is cudagraph compatible. This can improve performance in some cases such as decode speed for smaller models. Enable this using `--compilation-config "{'full_cuda_graph': True}"`
144
+
145
+Currently only FlashAttention 3 is compatible, and only when cascade attention is disabled.
tests/compile/piecewise/test_full_cudagraph.py
new
+97
@@ -0,0 +1,97 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import contextlib
3
+import os
4
+
5
+import pytest
6
+
7
+from vllm import LLM, SamplingParams
8
+from vllm.config import CompilationConfig
9
+
10
+MODEL = "Qwen/Qwen2-1.5B-Instruct"
11
+
12
+
13
+@contextlib.contextmanager
14
+def temporary_environ(env_vars):
15
+ """
16
+ Temporarily set environment variables and restore them afterward.
17
+ We have to do this vs monkeypatch because monkeypatch doesn't work
18
+ with "module" scoped fixtures.
19
+ """
20
+ original_env = {k: os.environ.get(k) for k in env_vars}
21
+ try:
22
+ os.environ.update(env_vars)
23
+ yield
24
+ finally:
25
+ for k, v in original_env.items():
26
+ if v is None:
27
+ os.environ.pop(k, None)
28
+ else:
29
+ os.environ[k] = v
30
+
31
+
32
+@pytest.fixture(scope="module")
33
+def full_cudagraph_llm():
34
+ with temporary_environ({
35
+ "VLLM_USE_V1": "1",
36
+ "VLLM_FLASH_ATTN_VERSION": "3"
37
+ }):
38
+ return LLM(model=MODEL,
39
+ gpu_memory_utilization=0.2,
40
+ compilation_config=CompilationConfig(full_cuda_graph=True))
41
+
42
+
43
+@pytest.fixture(scope="module")
44
+def piecewise_llm():
45
+ with temporary_environ({
46
+ "VLLM_USE_V1": "1",
47
+ "VLLM_FLASH_ATTN_VERSION": "3"
48
+ }):
49
+ return LLM(model=MODEL,
50
+ gpu_memory_utilization=0.5,
51
+ compilation_config=CompilationConfig())
52
+
53
+
54
+def generate_text(llm: LLM, batch_size: int, max_tokens: int):
55
+ prompts = ["Hi my name is"] * batch_size
56
+ sampling_params = SamplingParams(temperature=0.0,
57
+ max_tokens=max_tokens,
58
+ top_p=0.95)
59
+
60
+ return llm.generate(prompts, sampling_params)
61
+
62
+
63
+@pytest.mark.parametrize(("batch_size", "max_tokens"), [(1, 10), (7, 10),
64
+ (16, 10), (25, 10),
65
+ (32, 10), (45, 10),
66
+ (64, 10), (8, 5),
67
+ (8, 20), (8, 200)])
68
+def test_full_cudagraph(batch_size, max_tokens, full_cudagraph_llm,
69
+ piecewise_llm):
70
+ """
71
+ Load full cudagraph model and piecewise model once, and at the same time to
72
+ reuse them across various test cases.
73
+
74
+ Test various batch sizes and max_tokens to ensure that the full cudagraph
75
+ compilation works for padded cases too.
76
+ """
77
+ piecewise_responses = generate_text(piecewise_llm,
78
+ batch_size=batch_size,
79
+ max_tokens=max_tokens)
80
+ full_cudagraph_responses = generate_text(full_cudagraph_llm,
81
+ batch_size=batch_size,
82
+ max_tokens=max_tokens)
83
+
84
+ # Check that all responses are the same
85
+ for i in range(len(piecewise_responses)):
86
+ assert piecewise_responses[i].outputs[
87
+ 0].text == full_cudagraph_responses[i].outputs[0].text
88
+
89
+
90
+def test_full_cudagraph_with_invalid_backend():
91
+ with temporary_environ({
92
+ "VLLM_USE_V1": "1",
93
+ "VLLM_FLASH_ATTN_VERSION":
94
+ "2" #FA2 not supported with full_cuda_graph
95
+ }), pytest.raises(RuntimeError):
96
+ LLM(model=MODEL,
97
+ compilation_config=CompilationConfig(full_cuda_graph=True))
vllm/config.py
+17
-2
@@ -3605,6 +3605,10 @@ class CompilationConfig(BaseModel):
3605
are always used, it can set this to False. Otherwise, it should
3606
set this to True, and the compiler will copy the input to an
3607
internally managed buffer. Default is False.
3608
+ - full_cuda_graph: whether to use a full cuda graph for the entire forward
3609
+ pass rather than splitting certain operations such as attention into subgraphs.
3610
+ Thus this flag cannot be used together with splitting_ops. This may provide
3611
+ performance benefits for smaller models.
3612
- Inductor compilation:
3613
- use_inductor: whether to use inductor compilation.
3614
- False: inductor compilation is not used. graph runs in eager.
@@ -3649,6 +3653,7 @@ class CompilationConfig(BaseModel):
3653
cudagraph_num_of_warmups: int = 0
3654
cudagraph_capture_sizes: Optional[list[int]] = None
3655
cudagraph_copy_inputs: bool = False
3656
+ full_cuda_graph: bool = False
3657
3658
class PassConfig(BaseModel):
3659
"""
@@ -3871,10 +3876,14 @@ class CompilationConfig(BaseModel):
3876
self.max_capture_size] = self.max_capture_size
3877
3878
def set_splitting_ops_for_v1(self):
3874
- # If default, override splitting ops for piecewise cudagraph on V1.
3879
# NOTE: this function needs to be called
3880
+ if self.splitting_ops and self.full_cuda_graph:
3881
+ raise ValueError("full_cuda_graph cannot be used together with "
3882
+ "splitting_ops, as Full CUDA graph will override "
3883
+ f"the splitting_ops: {self.splitting_ops}")
3884
+
3885
if not self.splitting_ops:
3877
- self.splitting_ops = [
3886
+ self.splitting_ops = [] if self.full_cuda_graph else [
3887
"vllm.unified_attention",
3888
"vllm.unified_attention_with_output",
3889
]
@@ -4151,6 +4160,12 @@ class VllmConfig:
4160
"Disabling `torch.compile`.")
4161
self.compilation_config.level = CompilationLevel.NO_COMPILATION
4162
4163
+ if self.compilation_config.full_cuda_graph and \
4164
+ not self.model_config.disable_cascade_attn:
4165
+ logger.warning_once(
4166
+ "full_cuda_graph is not supported with "
4167
+ "cascade attention. Disabling cascade attention.")
4168
+ self.model_config.disable_cascade_attn = True
4169
4170
if self.model_config and self.model_config.use_mla and \
4171
not (current_platform.is_cuda() or current_platform.is_rocm()):
vllm/v1/attention/backends/flash_attn.py
+10
-3
@@ -291,6 +291,7 @@ class FlashAttentionMetadataBuilder:
291
292
def __init__(self, runner: "GPUModelRunner"):
293
model_config = runner.model_config
294
+ compilation_config = runner.vllm_config.compilation_config
295
296
self.runner = runner
297
self.num_heads_q = model_config.get_num_attention_heads(
@@ -300,7 +301,14 @@ class FlashAttentionMetadataBuilder:
301
self.headdim = model_config.get_head_size()
302
self.page_size = self.runner.block_size
303
303
- self.aot_schedule = (get_flash_attn_version() == 3)
304
+ if get_flash_attn_version() == 3:
305
+ self.aot_schedule = not compilation_config.full_cuda_graph
306
+ if not self.aot_schedule:
307
+ logger.warning(
308
+ "AOT Schedule is disabled when using full_cuda_graph")
309
+ else:
310
+ self.aot_schedule = False
311
+
312
# Sliding window size to be used with the AOT scheduler will be
313
# populated on first build() call.
314
self.aot_sliding_window: Optional[tuple[int, int]] = None
@@ -317,8 +325,7 @@ class FlashAttentionMetadataBuilder:
325
seq_lens = common_attn_metadata.seq_lens
326
block_table = (
327
self.runner.input_batch.block_table.get_device_tensor()[:num_reqs])
320
- slot_mapping = self.runner.slot_mapping_cpu[:num_actual_tokens].to(
321
- self.runner.device, non_blocking=True).long()
328
+ slot_mapping = self.runner.slot_mapping[:num_actual_tokens]
329
330
if self.aot_sliding_window is None:
331
self.aot_sliding_window = (-1, -1)
vllm/v1/worker/gpu_model_runner.py
+60
-8
@@ -12,6 +12,7 @@ import torch.nn as nn
12
13
from vllm.attention import AttentionType, get_attn_backend
14
from vllm.attention.layer import Attention
15
+from vllm.attention.utils.fa_utils import get_flash_attn_version
16
from vllm.config import (CompilationLevel, VllmConfig,
17
get_layers_from_vllm_config)
18
from vllm.distributed.kv_transfer import (get_kv_transfer_group,
@@ -139,6 +140,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
140
raise NotImplementedError(
141
"Non-Attention backend is not supported by V1 GPUModelRunner.")
142
143
+ if self.vllm_config.compilation_config.full_cuda_graph:
144
+ attn_backend_name = self.attn_backend.__name__
145
+ flash_attn_version = get_flash_attn_version()
146
+ if attn_backend_name != "FlashAttentionBackend" or \
147
+ flash_attn_version != 3:
148
+ raise ValueError(
149
+ f"full_cuda_graph is only supported with "
150
+ f"FA3. Current attention backend is {attn_backend_name}, "
151
+ f"FlashAttention version is {flash_attn_version}.")
152
+
153
self.attn_metadata_builder = self.attn_backend.get_builder_cls()(
154
weakref.proxy(self))
155
self.cascade_attn_enabled = not self.model_config.disable_cascade_attn
@@ -219,6 +230,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
230
self.positions = torch.zeros(self.max_num_tokens,
231
dtype=torch.int64,
232
device=self.device)
233
+ self.query_start_loc = torch.zeros(self.max_num_reqs + 1,
234
+ dtype=torch.int32,
235
+ device=self.device)
236
+ self.seq_lens = torch.zeros(self.max_num_reqs,
237
+ dtype=torch.int32,
238
+ device=self.device)
239
+ self.slot_mapping = torch.zeros(self.max_num_tokens,
240
+ dtype=torch.int64,
241
+ device=self.device)
242
+
243
# None in the first PP rank. The rest are set after load_model.
244
self.intermediate_tensors: Optional[IntermediateTensors] = None
245
@@ -271,7 +292,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
292
pin_memory=self.pin_memory)
293
self.positions_np = self.positions_cpu.numpy()
294
self.slot_mapping_cpu = torch.zeros(self.max_num_tokens,
274
- dtype=torch.int32,
295
+ dtype=torch.int64,
296
device="cpu",
297
pin_memory=self.pin_memory)
298
self.slot_mapping_np = self.slot_mapping_cpu.numpy()
@@ -589,10 +610,22 @@ class GPUModelRunner(LoRAModelRunnerMixin):
610
self.positions_cpu[:total_num_scheduled_tokens],
611
non_blocking=True)
612
592
- query_start_loc = self.query_start_loc_cpu[:num_reqs + 1].to(
593
- self.device, non_blocking=True)
594
- seq_lens = self.seq_lens_cpu[:num_reqs].to(self.device,
595
- non_blocking=True)
613
+ self.query_start_loc[:num_reqs + 1].copy_(
614
+ self.query_start_loc_cpu[:num_reqs + 1], non_blocking=True)
615
+ self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
616
+ non_blocking=True)
617
+ self.slot_mapping[:total_num_scheduled_tokens].copy_(
618
+ self.slot_mapping_cpu[:total_num_scheduled_tokens],
619
+ non_blocking=True)
620
+
621
+ # Fill unused with -1. Needed for reshape_and_cache
622
+ self.slot_mapping[total_num_scheduled_tokens:].fill_(-1)
623
+ self.seq_lens[num_reqs:].fill_(0)
624
+ self.query_start_loc[num_reqs + 1:].fill_(-1)
625
+
626
+ query_start_loc = self.query_start_loc[:num_reqs + 1]
627
+ seq_lens = self.seq_lens[:num_reqs]
628
+
629
common_attn_metadata = CommonAttentionMetadata(
630
query_start_loc=query_start_loc, seq_lens=seq_lens)
631
@@ -1478,6 +1511,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1511
def _dummy_run(
1512
self,
1513
num_tokens: int,
1514
+ skip_attn: bool = True,
1515
) -> torch.Tensor:
1516
1517
# Set num_scheduled_tokens based on num_tokens and max_num_seqs
@@ -1494,6 +1528,23 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1528
num_scheduled_tokens = np.array(num_scheduled_tokens_list,
1529
dtype=np.int32)
1530
1531
+ if skip_attn:
1532
+ attn_metadata = None
1533
+ else:
1534
+ query_start_loc = self.query_start_loc[:num_reqs + 1]
1535
+ seq_lens = self.seq_lens[:num_reqs]
1536
+
1537
+ common_attn_metadata = CommonAttentionMetadata(
1538
+ query_start_loc=query_start_loc, seq_lens=seq_lens)
1539
+
1540
+ attn_metadata = self.attn_metadata_builder.build(
1541
+ num_reqs=num_tokens,
1542
+ num_actual_tokens=num_tokens,
1543
+ max_query_len=num_tokens,
1544
+ common_prefix_len=0,
1545
+ common_attn_metadata=common_attn_metadata,
1546
+ )
1547
+
1548
with self.maybe_dummy_run_with_lora(self.lora_config,
1549
num_scheduled_tokens):
1550
model = self.model
@@ -1522,7 +1573,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1573
for k, v in self.intermediate_tensors.items()
1574
})
1575
1525
- with set_forward_context(None,
1576
+ with set_forward_context(attn_metadata,
1577
self.vllm_config,
1578
num_tokens=num_tokens):
1579
outputs = model(
@@ -1708,11 +1759,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1759
# Capture the large shapes first so that the smaller shapes
1760
# can reuse the memory pool allocated for the large shapes.
1761
with graph_capture(device=self.device):
1762
+ skip_attn = not self.vllm_config.compilation_config.full_cuda_graph
1763
for num_tokens in reversed(self.cudagraph_batch_sizes):
1764
for _ in range(self.vllm_config.compilation_config.
1765
cudagraph_num_of_warmups):
1714
- self._dummy_run(num_tokens)
1715
- self._dummy_run(num_tokens)
1766
+ self._dummy_run(num_tokens, skip_attn=skip_attn)
1767
+ self._dummy_run(num_tokens, skip_attn=skip_attn)
1768
1769
end_time = time.perf_counter()
1770
end_free_gpu_memory = torch.cuda.mem_get_info()[0]