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 7, 14:00 UTC
[compile] Fix caching error over pytree slice node. (#35308)
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
Zhengxu Chen committed
Feb 27, 2026 at 14:34 UTC
29b35477b0661f527d2b951ff5125f5c58fce3fe
2 files changed
+40
-2
tests/compile/test_aot_compile.py
+21
@@ -16,6 +16,7 @@ import torch
16
import vllm.model_executor.layers.activation
17
from vllm.compilation.caching import (
18
StandaloneCompiledArtifacts,
19
+ VllmSerializableFunction,
20
)
21
from vllm.compilation.decorators import support_torch_compile
22
from vllm.config import (
@@ -156,6 +157,26 @@ def test_save_and_load(monkeypatch: pytest.MonkeyPatch):
157
assert torch.allclose(ret, expected)
158
159
160
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
161
+def test_save_and_load_slice(monkeypatch: pytest.MonkeyPatch):
162
+ def foo(x: torch.Tensor):
163
+ return x[slice(0, x.shape[0])]
164
+
165
+ vllm_config = make_vllm_config()
166
+
167
+ example_input = torch.randn(10, 10)
168
+ torch._dynamo.mark_dynamic(example_input, 0)
169
+ gm = torch.fx.symbolic_trace(foo)
170
+ assert "getitem_1 = x[slice(0, getitem, None)]" in gm.code
171
+ with use_vllm_config(vllm_config):
172
+ payload = VllmSerializableFunction.serialize_compile_artifacts(
173
+ VllmSerializableFunction(gm, (example_input,), "", foo)
174
+ )
175
+ fn = VllmSerializableFunction.deserialize_compile_artifacts(payload)
176
+
177
+ assert gm.code == fn.graph_module.code
178
+
179
+
180
@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
181
def test_cache_load_returns_tuple_consistency(monkeypatch: pytest.MonkeyPatch):
182
"""
vllm/compilation/caching.py
+19
-2
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import contextlib
5
import hashlib
6
import inspect
7
import os
@@ -144,6 +145,18 @@ class StandaloneCompiledArtifacts:
145
self.loaded_submodule_store = {}
146
147
148
+@contextlib.contextmanager
149
+def patch_pytree_map_over_slice():
150
+ pytree._private_register_pytree_node(
151
+ slice, lambda x: ([x.start, x.stop, x.step], None), lambda x, c: slice(*x)
152
+ )
153
+
154
+ try:
155
+ yield
156
+ finally:
157
+ pytree._deregister_pytree_node(slice)
158
+
159
+
160
class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
161
"""
162
A wrapper around a compiled function by vllm. It will forward the tensor
@@ -235,7 +248,10 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
248
lambda inp: torch.empty_like(inp, device="meta"),
249
state["example_inputs"],
250
)
238
- with patch.object(GraphPickler, "reducer_override", _graph_reducer_override):
251
+ with (
252
+ patch.object(GraphPickler, "reducer_override", _graph_reducer_override),
253
+ patch_pytree_map_over_slice(),
254
+ ):
255
state["graph_module"] = GraphPickler.dumps(
256
state["graph_module"], Options(ops_filter=None)
257
)
@@ -261,7 +277,8 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
277
278
state = pickle.loads(data)
279
fake_mode = FakeTensorMode(shape_env=ShapeEnv())
264
- state["graph_module"] = GraphPickler.loads(state["graph_module"], fake_mode)
280
+ with patch_pytree_map_over_slice():
281
+ state["graph_module"] = GraphPickler.loads(state["graph_module"], fake_mode)
282
state["graph_module"].recompile()
283
state["example_inputs"] = GraphPickler.loads(state["example_inputs"], fake_mode)
284