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
[Bugfix] Remove VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE #2969 (#25090)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com>
Lucas Kabela committed
Sep 19, 2025 at 17:27 UTC
3da17c2cc2c2e1d750020e033535f942f156f64c
6 files changed
+11
-36
tests/compile/test_basic_correctness.py
+2
-14
@@ -20,7 +20,6 @@ class TestSetting:
20
tp_size: int
21
attn_backend: str
22
method: str
23
- fullgraph: bool
23
24
25
# we cannot afford testing the full Cartesian product
@@ -36,7 +35,6 @@ class TestSetting:
35
tp_size=2,
36
attn_backend="FLASH_ATTN",
37
method="generate",
39
- fullgraph=True,
38
),
39
# llama model with quantization
40
TestSetting(
@@ -46,7 +44,6 @@ class TestSetting:
44
tp_size=1,
45
attn_backend="FLASH_ATTN",
46
method="generate",
49
- fullgraph=True,
47
),
48
# MoE model
49
TestSetting(
@@ -56,7 +53,6 @@ class TestSetting:
53
tp_size=2,
54
attn_backend="FLASH_ATTN",
55
method="generate",
59
- fullgraph=True,
56
),
57
# embedding model
58
TestSetting(
@@ -73,7 +69,6 @@ class TestSetting:
69
tp_size=1,
70
attn_backend="FLASH_ATTN",
71
method="encode",
76
- fullgraph=True,
72
),
73
TestSetting(
74
model="BAAI/bge-base-en-v1.5",
@@ -82,7 +77,6 @@ class TestSetting:
77
tp_size=1,
78
attn_backend="FLASH_ATTN",
79
method="encode",
85
- fullgraph=True,
80
),
81
# vision language model
82
TestSetting(
@@ -92,7 +86,6 @@ class TestSetting:
86
tp_size=1,
87
attn_backend="FLASH_ATTN",
88
method="generate_with_image",
95
- fullgraph=False,
89
),
90
],
91
)
@@ -109,9 +102,8 @@ def test_compile_correctness(
102
tp_size = test_setting.tp_size
103
attn_backend = test_setting.attn_backend
104
method = test_setting.method
112
- fullgraph = test_setting.fullgraph
113
- if cuda_device_count_stateless() != pp_size * tp_size:
114
- pytest.skip(f"Need exactly {pp_size}*{tp_size} CUDA gpus but got "
105
+ if cuda_device_count_stateless() < pp_size * tp_size:
106
+ pytest.skip(f"Need at least {pp_size}*{tp_size} CUDA gpus but got "
107
f"{cuda_device_count_stateless()}")
108
109
with monkeypatch.context() as m:
@@ -149,9 +141,5 @@ def test_compile_correctness(
141
]:
142
all_args.append(final_args + [f"-O{level}"])
143
all_envs.append({})
152
- if level != CompilationLevel.DYNAMO_ONCE and not fullgraph:
153
- # "DYNAMO_ONCE" will always use fullgraph
154
- all_envs[-1][
155
- "VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE"] = "0" # type: ignore
144
145
compare_all_settings(model, all_args * 3, all_envs, method=method)
tests/compile/test_full_graph.py
+1
-3
@@ -79,9 +79,7 @@ def test_full_graph(
79
):
80
model, model_kwargs = model_info
81
82
- with monkeypatch.context() as m:
83
- # make sure these models can be captured in full graph mode
84
- m.setenv("VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1")
82
+ with monkeypatch.context():
83
print(f"MODEL={model}")
84
85
run_model(optimization_level, model, model_kwargs)
vllm/compilation/wrapper.py
+4
-6
@@ -10,7 +10,6 @@ from typing import Callable, Optional
10
11
import torch
12
13
-import vllm.envs as envs
13
from vllm.config import (CompilationLevel, CUDAGraphMode,
14
get_current_vllm_config)
15
from vllm.logger import init_logger
@@ -47,11 +46,10 @@ class TorchCompileWrapperWithCustomDispatcher:
46
options = get_current_vllm_config(
47
).compilation_config.inductor_compile_config
48
50
- compiled_callable = torch.compile(
51
- self.forward,
52
- fullgraph=envs.VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE,
53
- backend=backend,
54
- options=options)
49
+ compiled_callable = torch.compile(self.forward,
50
+ fullgraph=True,
51
+ backend=backend,
52
+ options=options)
53
54
self.compiled_callable = compiled_callable
55
self.original_code_object = self.__class__.forward.__code__
vllm/envs.py
-5
@@ -434,11 +434,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
434
"VLLM_FLASH_ATTN_VERSION":
435
lambda: maybe_convert_int(os.environ.get("VLLM_FLASH_ATTN_VERSION", None)),
436
437
- # Internal flag to enable Dynamo fullgraph capture
438
- "VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE":
439
- lambda: bool(
440
- os.environ.get("VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1") != "0"),
441
-
437
# Feature flag to enable/disable Inductor standalone compile.
438
# In torch <= 2.7 we ignore this flag; in torch >= 2.8 this is
439
# enabled by default.
vllm/v1/worker/gpu_model_runner.py
+1
-3
@@ -2602,9 +2602,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2602
backend = self.vllm_config.compilation_config.init_backend(
2603
self.vllm_config)
2604
compilation_counter.dynamo_as_is_count += 1
2605
- self.model.compile(
2606
- fullgraph=envs.VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE,
2607
- backend=backend)
2605
+ self.model.compile(fullgraph=True, backend=backend)
2606
return
2607
# for other compilation levels, cudagraph behavior is controlled by
2608
# CudagraphWraper and CudagraphDispatcher of vllm.
vllm/worker/model_runner.py
+3
-5
@@ -18,7 +18,6 @@ import torch.distributed
18
import torch.nn as nn
19
from tqdm.auto import tqdm
20
21
-import vllm.envs as envs
21
from vllm.attention import AttentionMetadata, get_attn_backend
22
from vllm.attention.backends.abstract import AttentionState
23
from vllm.attention.backends.utils import CommonAttentionState
@@ -1099,10 +1098,9 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):
1098
backend = self.vllm_config.compilation_config.init_backend(
1099
self.vllm_config)
1100
compilation_counter.dynamo_as_is_count += 1
1102
- self.model = torch.compile(
1103
- self.model,
1104
- fullgraph=envs.VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE,
1105
- backend=backend)
1101
+ self.model = torch.compile(self.model,
1102
+ fullgraph=True,
1103
+ backend=backend)
1104
1105
def get_model(self) -> nn.Module:
1106
return self.model