@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 15f40b20aa
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 9, 05:00 UTC

[fix][torch.compile] Fix cold-start compilation time increase by adding kv cache update to splitting ops (#33441)

Signed-off-by: Luka Govedič <lgovedic@redhat.com> Co-authored-by: Richard Zou <zou3519@gmail.com>

Luka Govedič committed Jan 31, 2026 at 09:48 UTC 15f40b20aadf27af43bdae38bf4644b82e21634f
4 files changed +127 -1
tests/compile/test_cold_start.py new
+48
@@ -0,0 +1,48 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from torch._dynamo.utils import counters
5 +
6 +from vllm import LLM
7 +from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode
8 +
9 +
10 +def test_moe_compilation_cold_start(monkeypatch, use_fresh_inductor_cache):
11 + # Run in same process so we can access PyTorch's internal counters
12 + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
13 +
14 + # I'm not sure if this is going to affect the numbers
15 + monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "0")
16 +
17 + # Force cold compilation
18 + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
19 +
20 + compilation_config = CompilationConfig(
21 + mode=CompilationMode.VLLM_COMPILE,
22 + cudagraph_mode=CUDAGraphMode.NONE, # make the model loading faster
23 + )
24 +
25 + counters.clear()
26 +
27 + _ = LLM(
28 + model="microsoft/Phi-tiny-MoE-instruct",
29 + max_model_len=256,
30 + load_format="dummy", # make the model loading faster
31 + compilation_config=compilation_config,
32 + num_gpu_blocks_override=8, # make the model loading faster
33 + )
34 +
35 + # vLLM-compile cold start is special. By default, we do
36 + # one full dynamo capture of the entire forward pass.
37 + # The forward pass consists of 32 transformer layers.
38 + # Then, we split on the attention operation. This results in
39 + # 33 subgraphs (not including the attention operation).
40 + # The 33 subgraphs then get standalone_compile'd.
41 + #
42 + # There are actually only 3 unique subgraphs for this model
43 + # (all of its transformer layers are the same modulo weights);
44 + # this is true for most vLLM models.
45 + # So we test that during cold start, the aot_autograd cache
46 + # misses for 3 subgraphs and hits for the rest.
47 + assert counters["aot_autograd"]["autograd_cache_miss"] == 3
48 + assert counters["aot_autograd"]["autograd_cache_hit"] == 30
tests/compile/test_graph_partition.py
+62
@@ -8,6 +8,10 @@ import torch
8 from torch.fx.experimental.proxy_tensor import make_fx
9
10 from vllm.compilation.backends import split_graph
11 +from vllm.compilation.fx_utils import find_op_nodes
12 +
13 +# This import automatically registers `torch.ops.silly.attention`
14 +from . import silly_attention # noqa: F401
15
16
17 def test_getitem_moved_to_producer_subgraph():
@@ -122,3 +126,61 @@ def test_no_tuple_inputs_with_multiple_consumers():
126 output_split = split_gm(new_x)
127
128 assert torch.allclose(output_original, output_split), "Output mismatch after split"
129 +
130 +
131 +def test_consecutive_ops_in_split():
132 + """
133 + Test that consecutive splitting operations are grouped into the same subgraph
134 + """
135 +
136 + def model_fn(x: torch.Tensor) -> torch.Tensor:
137 + """
138 + Define a simple model where consecutive operations create opportunities
139 + for splitting subgraphs.
140 + """
141 + # Apply silly attention followed by consecutive operations
142 + intermediate = torch.relu(x)
143 + attn_inout = torch.sqrt(intermediate)
144 + torch.ops.silly.attention(intermediate, intermediate, attn_inout, attn_inout)
145 + final_result = torch.sigmoid(attn_inout)
146 + return final_result
147 +
148 + torch.set_default_device("cuda")
149 +
150 + # Create the traced FX graph for the model
151 + x = torch.randn(8, 4)
152 +
153 + gm = make_fx(model_fn)(x)
154 +
155 + # Assert presence of the expected operations in the setup
156 + assert (
157 + len(list(find_op_nodes(torch.ops.aten.relu, gm.graph))) == 1
158 + and len(list(find_op_nodes(torch.ops.aten.sqrt, gm.graph))) == 1
159 + ), "Test setup failed: Expected sqrt and relu operations in the graph."
160 +
161 + # Configure split operations to test
162 + splitting_ops = ["silly::attention", "aten::sqrt"]
163 + split_gm, split_items = split_graph(gm, splitting_ops)
164 +
165 + # Validate the number of partitions
166 + assert len(split_items) == 3, (
167 + "Consecutive splitting operations were not grouped correctly."
168 + )
169 +
170 + # Validate that correctness is preserved
171 + new_x = torch.randn(8, 4)
172 + output_original = gm(new_x)
173 + output_split = split_gm(new_x)
174 + assert torch.allclose(output_original, output_split), (
175 + "Output mismatch after splitting."
176 + )
177 +
178 + # Check the splitting item has 2 nodes exactly (relu and attn)
179 + splitting_items = list(s for s in split_items if s.is_splitting_graph)
180 + assert len(splitting_items) == 1, "Expecting a single splitting graph"
181 + print(splitting_items[0].graph.graph)
182 + splitting_gm = splitting_items[0].graph
183 + assert len(splitting_gm.graph.nodes) == 4, "Expecting 4 nodes in splitting graph"
184 + assert [node.op for node in splitting_gm.graph.nodes] == ["placeholder"] + 2 * [
185 + "call_function"
186 + ] + ["output"]
vllm/compilation/backends.py
+8 -1
@@ -359,7 +359,14 @@ def split_graph(
359 subgraph_id += 1
360 node_to_subgraph_id[node] = subgraph_id
361 split_op_graphs.append(subgraph_id)
362 - subgraph_id += 1
362 +
363 + # keep consecutive splitting ops together
364 + # (we know node.next exists because node isn't the last (output) node)
365 + if should_split(node.next, splitting_ops):
366 + # this will get incremented by the next node
367 + subgraph_id -= 1
368 + else:
369 + subgraph_id += 1
370 else:
371 node_to_subgraph_id[node] = subgraph_id
372
vllm/config/compilation.py
+9
@@ -943,6 +943,15 @@ class CompilationConfig:
943 # for details. Make a copy to avoid mutating the class-level
944 # list via reference.
945 self.splitting_ops = list(self._attention_ops)
946 +
947 + # unified_kv_cache_update has a string param that prevents Inductor
948 + # from reusing piecewise graphs. Remove it from the compiled graph.
949 + # This has the side-effect of excluding cache from cudagraphs but
950 + # that doesn't seem to affect performance.
951 + # https://github.com/vllm-project/vllm/issues/33267
952 + if not self.use_inductor_graph_partition:
953 + self.splitting_ops.append("vllm::unified_kv_cache_update")
954 +
955 elif len(self.splitting_ops) == 0:
956 if (
957 self.cudagraph_mode == CUDAGraphMode.PIECEWISE