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, 07:15 UTC
[Bugfix] Re-enable use_cudagraph in vLLM v1 (#19299)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Jun 7, 2025 at 20:56 UTC
eaa2e51088d4daf36d47e566ad90e812f80e91b8
6 files changed
+52
-8
tests/compile/piecewise/test_simple.py
+1
-1
@@ -95,7 +95,7 @@ def _test_simple_piecewise_compile(*, use_inductor):
95
num_piecewise_graphs_seen=5, # 2 * num_layers + 1
96
num_piecewise_capturable_graphs_seen=3, # 1 + num_layers
97
num_backend_compilations=3, # num_piecewise_capturable_graphs_seen
98
- num_cudagraph_caputured=
98
+ num_cudagraph_captured=
99
6, # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
100
):
101
tests/compile/piecewise/test_toy_llama.py
+3
-3
@@ -327,7 +327,7 @@ def _test_toy_llama(*, use_inductor):
327
num_piecewise_graphs_seen=0,
328
num_piecewise_capturable_graphs_seen=0,
329
num_backend_compilations=0,
330
- num_cudagraph_caputured=0,
330
+ num_cudagraph_captured=0,
331
):
332
outputs.append(
333
run_model(llama_config, use_inductor=False, use_compile=False))
@@ -343,7 +343,7 @@ def _test_toy_llama(*, use_inductor):
343
num_piecewise_graphs_seen=1,
344
num_piecewise_capturable_graphs_seen=1,
345
num_backend_compilations=1, # num_piecewise_capturable_graphs_seen
346
- num_cudagraph_caputured=
346
+ num_cudagraph_captured=
347
2, # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
348
**kwargs,
349
):
@@ -361,7 +361,7 @@ def _test_toy_llama(*, use_inductor):
361
llama_config.num_layers, # 1 + num_layers
362
num_backend_compilations=1 +
363
llama_config.num_layers, # num_piecewise_capturable_graphs_seen
364
- num_cudagraph_caputured=2 *
364
+ num_cudagraph_captured=2 *
365
(1 + llama_config.num_layers
366
), # num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
367
):
tests/compile/test_config.py
new
+43
@@ -0,0 +1,43 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+import torch
5
+
6
+import vllm
7
+from vllm.compilation.counter import compilation_counter
8
+from vllm.config import (CompilationConfig, CompilationLevel, VllmConfig,
9
+ set_current_vllm_config)
10
+
11
+from .piecewise.test_simple import SillyModel
12
+
13
+
14
+@pytest.fixture(scope="function", autouse=True)
15
+def use_v1(monkeypatch):
16
+ """
17
+ TODO(rzou): The rest of tests/compile runs VLLM_USE_V1=0 right now,
18
+ I'll switch them over later.
19
+ """
20
+ monkeypatch.setenv('VLLM_USE_V1', '1')
21
+
22
+
23
+@pytest.mark.parametrize("enabled", [True, False])
24
+def test_use_cudagraphs(enabled):
25
+ assert vllm.envs.VLLM_USE_V1
26
+ vllm_config = VllmConfig(compilation_config=CompilationConfig(
27
+ level=CompilationLevel.PIECEWISE,
28
+ use_cudagraph=enabled,
29
+ cudagraph_capture_sizes=[100],
30
+ ))
31
+ with set_current_vllm_config(vllm_config):
32
+ model = SillyModel(vllm_config=vllm_config, prefix='')
33
+
34
+ inputs = torch.randn(100, device="cuda")
35
+
36
+ with compilation_counter.expect(
37
+ num_graphs_seen=1, # one graph for the model
38
+ num_cudagraph_captured=1 if enabled else 0,
39
+ ):
40
+ # first run is warmup
41
+ model(inputs)
42
+ # second run does CUDAGraphs recording (if enabled)
43
+ model(inputs)
vllm/compilation/counter.py
+1
-1
@@ -15,7 +15,7 @@ class CompilationCounter:
15
# not including the splitting ops
16
num_piecewise_capturable_graphs_seen: int = 0
17
num_backend_compilations: int = 0
18
- num_cudagraph_caputured: int = 0
18
+ num_cudagraph_captured: int = 0
19
# InductorAdapter.compile calls
20
num_inductor_compiles: int = 0
21
# EagerAdapter.compile calls
vllm/compilation/cuda_piecewise_backend.py
+1
-1
@@ -193,7 +193,7 @@ class CUDAPiecewiseBackend:
193
entry.output = weak_ref_tensors(output)
194
entry.cudagraph = cudagraph
195
196
- compilation_counter.num_cudagraph_caputured += 1
196
+ compilation_counter.num_cudagraph_captured += 1
197
198
# important: we need to return the output, rather than
199
# the weak ref of the output, so that pytorch can correctly
vllm/config.py
+3
-2
@@ -3918,12 +3918,14 @@ class CompilationConfig:
3918
constructor, e.g. `CompilationConfig(inductor_passes={"a": func})`."""
3919
3920
# CudaGraph compilation
3921
- use_cudagraph: bool = False
3921
+ use_cudagraph: bool = envs.VLLM_USE_V1
3922
"""Whether to use cudagraph inside compilation.
3923
- False: cudagraph inside compilation is not used.
3924
- True: cudagraph inside compilation is used. It requires
3925
that all input buffers have fixed addresses, and all
3926
splitting ops write their outputs to input buffers.
3927
+ In the vLLM V1 Engine, this flag only applies for
3928
+ CompilationLevel.PIECEWISE (aka -O3).
3929
Note that this is orthogonal to the cudagraph capture logic
3930
outside of compilation.
3931
TODO: move outside cudagraph logic into compilation.
@@ -4425,7 +4427,6 @@ class VllmConfig:
4427
# FIXME(rob): Add function to set all of these.
4428
if not self.compilation_config.custom_ops:
4429
self.compilation_config.custom_ops = ["none"]
4428
- self.compilation_config.use_cudagraph = True
4430
self.compilation_config.cudagraph_num_of_warmups = 1
4431
self.compilation_config.pass_config.enable_fusion = False
4432
self.compilation_config.pass_config.enable_noop = False