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 10, 22:45 UTC
Fix/get raw stream patch #30905 (#30912)
Signed-off-by: baonudesifeizhai <baonudesifeizhai@gmail.com> Signed-off-by: baonudesifeizhai <85092850+baonudesifeizhai@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
baonudesifeizhai committed
Dec 26, 2025 at 23:08 UTC
8711b216766bb5d3cbe15161061c3a7d9fffe59c
3 files changed
+65
-1
tests/compile/fullgraph/test_full_graph.py
+14
@@ -156,6 +156,20 @@ def test_full_graph(
156
)
157
for model_info in models_list(all=False)
158
if is_torch_equal_or_newer("2.9.0.dev")
159
+ ]
160
+ + [
161
+ # Test get_raw_stream patch with compile_sizes
162
+ # This tests that TorchInductor autotune works correctly with get_raw_stream
163
+ # patch in torch 2.9 and without patch in torch 2.10+
164
+ (
165
+ CompilationConfig(
166
+ mode=CompilationMode.VLLM_COMPILE,
167
+ compile_sizes=[1, 2], # Triggers autotune which uses get_raw_stream
168
+ cudagraph_mode=CUDAGraphMode.NONE,
169
+ ),
170
+ "facebook/opt-125m",
171
+ {},
172
+ ),
173
],
174
)
175
# only test some of the models
tests/compile/test_config.py
+27
-1
@@ -13,7 +13,10 @@ from vllm.config import CompilationConfig, CUDAGraphMode, ParallelConfig, VllmCo
13
from vllm.config.compilation import CompilationMode, PassConfig
14
from vllm.engine.arg_utils import EngineArgs
15
from vllm.platforms import current_platform
16
-from vllm.utils.torch_utils import _is_torch_equal_or_newer
16
+from vllm.utils.torch_utils import (
17
+ _is_torch_equal_or_newer,
18
+ is_torch_equal,
19
+)
20
21
# This import automatically registers `torch.ops.silly.attention`
22
from . import silly_attention # noqa: F401
@@ -28,6 +31,29 @@ def test_version():
31
assert not _is_torch_equal_or_newer("2.7.1", "2.8.0.dev")
32
33
34
+def test_get_raw_stream_patch():
35
+ """Test that get_raw_stream patch is applied only for torch 2.9.0 or 2.9.1."""
36
+ import builtins
37
+
38
+ # Check if get_raw_stream exists in builtins
39
+ has_patch = hasattr(builtins, "get_raw_stream")
40
+
41
+ # Import torch to get actual version
42
+
43
+ is_torch_2_9 = is_torch_equal("2.9.0") or is_torch_equal("2.9.1")
44
+
45
+ if is_torch_2_9:
46
+ # For torch 2.9.x, the patch should be applied
47
+ assert has_patch, "get_raw_stream should be patched for torch 2.9.x"
48
+ # Verify it's callable (it should be the _cuda_getCurrentRawStream function)
49
+ get_raw_stream = builtins.get_raw_stream # type: ignore[attr-defined]
50
+ assert callable(get_raw_stream)
51
+ # Verify it's the correct function from torch._C
52
+ from torch._C import _cuda_getCurrentRawStream
53
+
54
+ assert get_raw_stream is _cuda_getCurrentRawStream
55
+
56
+
57
def test_copy_pass():
58
vllm_config = VllmConfig()
59
inductor_pass = FixFunctionalizationPass(vllm_config)
vllm/env_override.py
+24
@@ -363,6 +363,30 @@ def _update_scheduler_patched(self) -> None:
363
self.scheduler = Scheduler(self.operations)
364
365
366
+# ===================================================
367
+# torch 2.9 Inductor get_raw_stream workaround
368
+# ===================================================
369
+# Workaround for TorchInductor autotune using get_raw_stream() without defining it.
370
+# This occurs when compile_sizes > 1 in compilation_config.
371
+# For more context, see https://github.com/vllm-project/vllm/issues/30905.
372
+def _patch_get_raw_stream_if_needed():
373
+ """Workaround for TorchInductor autotune get_raw_stream() bug."""
374
+ from vllm.utils.torch_utils import is_torch_equal
375
+
376
+ # Only apply the patch for torch 2.9.0 or 2.9.1
377
+ if is_torch_equal("2.9.0") or is_torch_equal("2.9.1"):
378
+ import builtins
379
+
380
+ # Check if CUDA functionality is available without initializing CUDA
381
+ # _cuda_getCurrentRawStream only exists in CUDA builds of PyTorch
382
+ if hasattr(torch._C, "_cuda_getCurrentRawStream"):
383
+ from torch._C import _cuda_getCurrentRawStream as _get_raw_stream
384
+
385
+ builtins.get_raw_stream = _get_raw_stream
386
+
387
+
388
+_patch_get_raw_stream_if_needed()
389
+
390
if is_torch_equal("2.9.0"):
391
from torch._inductor.codegen.wrapper import PythonWrapperCodegen
392
from torch._inductor.graph import GraphLowering