Compare commits

...
Author SHA1 Message Date
yewentao256 77230471c0 remove torch 2.9, 2.10 workround
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-04-30 21:10:50 +00:00
5 changed files with 3 additions and 602 deletions
+2 -4
View File
@@ -150,13 +150,11 @@ def test_full_graph(
if is_torch_equal_or_newer("2.9.0.dev")
]
+ [
# Test get_raw_stream patch with compile_sizes
# This tests that TorchInductor autotune works correctly with get_raw_stream
# patch in torch 2.9 and without patch in torch 2.10+
# Cover compile_sizes autotune path.
(
CompilationConfig(
mode=CompilationMode.VLLM_COMPILE,
compile_sizes=[1, 2], # Triggers autotune which uses get_raw_stream
compile_sizes=[1, 2], # Triggers the autotune path.
cudagraph_mode=CUDAGraphMode.NONE,
),
"facebook/opt-125m",
-24
View File
@@ -24,7 +24,6 @@ from vllm.engine.arg_utils import EngineArgs
from vllm.platforms import current_platform
from vllm.utils.torch_utils import (
_is_torch_equal_or_newer,
is_torch_equal,
)
from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
@@ -43,29 +42,6 @@ def test_version():
assert not _is_torch_equal_or_newer("2.7.1", "2.8.0.dev")
def test_get_raw_stream_patch():
"""Test that get_raw_stream patch is applied only for torch 2.9.0 or 2.9.1."""
import builtins
# Check if get_raw_stream exists in builtins
has_patch = hasattr(builtins, "get_raw_stream")
# Import torch to get actual version
is_torch_2_9 = is_torch_equal("2.9.0") or is_torch_equal("2.9.1")
if is_torch_2_9:
# For torch 2.9.x, the patch should be applied
assert has_patch, "get_raw_stream should be patched for torch 2.9.x"
# Verify it's callable (it should be the _cuda_getCurrentRawStream function)
get_raw_stream = builtins.get_raw_stream # type: ignore[attr-defined]
assert callable(get_raw_stream)
# Verify it's the correct function from torch._C
from torch._C import _cuda_getCurrentRawStream
assert get_raw_stream is _cuda_getCurrentRawStream
def test_copy_pass():
vllm_config = VllmConfig()
inductor_pass = FixFunctionalizationPass(vllm_config)
-103
View File
@@ -1,103 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the FxGraphCachePickler.dumps ValueError patch in env_override.py.
Validates that _apply_fxgraphcache_pickle_patch correctly wraps a pickler's
dumps method to convert ValueError into a bypass exception, without affecting
other exception types or normal return values.
"""
import pytest
from vllm.env_override import _apply_fxgraphcache_pickle_patch
class _BypassStub(Exception):
"""Stand-in for BypassFxGraphCache in unit tests."""
class TestApplyFxgraphcachePicklePatch:
def test_valueerror_converted_to_bypass(self):
class Pickler:
def dumps(self, obj):
raise ValueError("can't serialize blocked layout")
_apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
with pytest.raises(_BypassStub, match="Failed to pickle cache key"):
Pickler().dumps(object())
def test_original_valueerror_chained(self):
class Pickler:
def dumps(self, obj):
raise ValueError("bad tensor layout")
_apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
with pytest.raises(_BypassStub) as exc_info:
Pickler().dumps(object())
cause = exc_info.value.__cause__
assert isinstance(cause, ValueError)
assert str(cause) == "bad tensor layout"
def test_non_valueerror_propagates(self):
class Pickler:
def dumps(self, obj):
raise TypeError("unexpected type")
_apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
with pytest.raises(TypeError, match="unexpected type"):
Pickler().dumps(object())
def test_normal_return_preserved(self):
sentinel = b"serialized-graph-key"
class Pickler:
def dumps(self, obj):
return sentinel
_apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
assert Pickler().dumps(object()) is sentinel
def test_idempotent(self):
class Pickler:
def dumps(self, obj):
return b"ok"
_apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
first_dumps = Pickler.dumps
_apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
assert Pickler.dumps is first_dumps
def test_sentinel_attribute_set(self):
class Pickler:
def dumps(self, obj):
return b"ok"
assert not hasattr(Pickler.dumps, "_vllm_patched")
assert not getattr(Pickler, "_vllm_fxgraph_dumps_patched", False)
_apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
assert Pickler.dumps._vllm_patched is True # type: ignore[attr-defined]
assert Pickler._vllm_fxgraph_dumps_patched is True # type: ignore[attr-defined]
def test_patch_applied_in_current_environment():
"""Integration: verify patch state matches current torch version."""
from torch._inductor.codecache import FxGraphCachePickler
from vllm.utils.torch_utils import is_torch_equal_or_newer
should_be_patched = is_torch_equal_or_newer(
"2.10.0"
) and not is_torch_equal_or_newer("2.11.0")
assert getattr(FxGraphCachePickler, "_vllm_fxgraph_dumps_patched", False) == (
should_be_patched
)
assert hasattr(FxGraphCachePickler.dumps, "_vllm_patched") == should_be_patched
-43
View File
@@ -203,47 +203,6 @@ def is_compile_cache_enabled(
)
def _patch_standalone_compile_atomic_save() -> None:
"""Backport of pytorch/pytorch#162432 for torch < 2.10.0.
Patches CompiledArtifact.save() to use write_atomic for binary format,
preventing corrupt cache files when multiple processes compile
concurrently.
"""
from torch._inductor.codecache import write_atomic
from torch._inductor.standalone_compile import CompiledArtifact as cls
if getattr(cls.save, "_vllm_patched", False):
return
original_save = cls.save
def _save(
self: Any, *, path: str, format: Literal["binary", "unpacked"] = "binary"
) -> None:
if format != "binary":
return original_save(self, path=path, format=format)
from torch._dynamo.utils import dynamo_timed
from torch._inductor.codecache import torch_key
from torch.utils._appending_byte_serializer import BytesWriter
with dynamo_timed("CompiledArtifact.save"):
assert self._artifacts is not None
artifact_bytes, cache_info = self._artifacts
assert len(cache_info.aot_autograd_artifacts) == 1, cache_info
key = cache_info.aot_autograd_artifacts[0]
assert not os.path.isdir(path)
writer = BytesWriter()
writer.write_bytes(torch_key())
writer.write_str(key)
writer.write_bytes(artifact_bytes)
write_atomic(path, writer.to_bytes())
_save._vllm_patched = True # type: ignore[attr-defined]
cls.save = _save # type: ignore[assignment]
logger.debug("Patched %s.save for atomic writes (torch < 2.10)", cls.__name__)
class InductorStandaloneAdaptor(CompilerInterface):
"""
The adaptor for the Inductor compiler.
@@ -257,8 +216,6 @@ class InductorStandaloneAdaptor(CompilerInterface):
name = "inductor_standalone"
def __init__(self, save_format: Literal["binary", "unpacked"]) -> None:
if not is_torch_equal_or_newer("2.10.0"):
_patch_standalone_compile_atomic_save()
self.save_format = save_format
def compute_hash(self, vllm_config: VllmConfig) -> str:
+1 -428
View File
@@ -87,7 +87,7 @@ _maybe_set_cuda_compatibility_path()
import torch
from vllm.logger import init_logger
from vllm.utils.torch_utils import is_torch_equal, is_torch_equal_or_newer
from vllm.utils.torch_utils import is_torch_equal_or_newer
logger = init_logger(__name__)
@@ -112,384 +112,6 @@ os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
# in the environment.
os.environ.setdefault("TRITON_CACHE_AUTOTUNING", "1")
# ===================================================
# torch 2.9 Inductor PythonWrapperCodegen monkeypatch
# ===================================================
# This change monkeypatches memory_plan_reuse in pytorch 2.9.0 to work around
# a test failure for test_multi_graph_piecewise_compile_outputs_equal.
# For more context, see https://github.com/pytorch/pytorch/pull/165514.
def memory_plan_reuse_patched(self):
import torch._inductor.ir as ir
from torch._inductor.codegen.wrapper import (
EnterSubgraphLine,
ExitSubgraphLine,
MemoryPlanningLine,
MemoryPlanningState,
SubgraphPythonWrapperCodegen,
)
from torch._inductor.virtualized import V
def get_output_names(graph_outputs) -> list[str]:
import itertools
names = []
shape_counter = itertools.count(0)
none_counter = itertools.count(0)
for node in graph_outputs:
if isinstance(node, ir.NoneAsConstantBuffer):
names.append(f"{V.graph.name}_none{next(none_counter)}")
elif isinstance(node, ir.ShapeAsConstantBuffer):
names.append(f"{V.graph.name}_shape{next(shape_counter)}")
else:
names.append(node.get_name())
return names
if (
isinstance(V.graph.wrapper_code, SubgraphPythonWrapperCodegen)
and V.graph.wrapper_code.partition_signatures is not None
):
out_names = get_output_names(
V.graph.wrapper_code.partition_signatures.output_nodes
)
else:
out_names = V.graph.get_output_names()
while (
self.lines
and isinstance(self.lines[-1], MemoryPlanningLine)
and self.lines[-1].node.name not in out_names # type: ignore[attr-defined]
):
# these lines will be pointless
self.lines.pop()
# codegen allocations in two passes
planning_states = [MemoryPlanningState()]
past_planning_states = []
for i in range(len(self.lines)):
line = self.lines[i]
if isinstance(line, MemoryPlanningLine):
self.lines[i] = line.plan(planning_states[-1])
elif isinstance(line, EnterSubgraphLine):
planning_states.append(MemoryPlanningState())
elif isinstance(line, ExitSubgraphLine):
past_planning_states.append(planning_states.pop())
past_planning_states.append(planning_states.pop())
assert len(planning_states) == 0
# ===================================================
# torch 2.9 Inductor get_graph_partition_signature monkeypatch
# ===================================================
# This change monkeypatches get_graph_partition_signature in pytorch 2.9.0 to
# fix inductor partition + attention-nvfp4 quant fusion, tested in
# `tests/compile/test_fusion_attn.py::test_attn_quant`.
# For more context, see https://github.com/pytorch/pytorch/pull/165815.
def get_graph_partition_signature_patched(
self, partitions, skip_cudagraphs: list[bool]
):
"""
Gets signature for each graph partition, including input nodes, output nodes, and
whether deallocating an input within graph partition.
"""
from torch._inductor import dependencies
from torch._inductor.ir import GraphPartitionSignature, MutationOutput, NoneLayout
from torch._inductor.virtualized import V
from torch.utils._ordered_set import OrderedSet
signatures = []
unmet_output_names = OrderedSet(V.graph.get_output_names())
name_to_node = self.get_name_to_nodes()
def is_none_layout(buf_name: str) -> bool:
"""
Checks if buf_name is NoneLayout. Buffers with NoneLayout is not allocated
so graph partition should not take it as inputs or outputs.
"""
buf = self.name_to_buf.get(buf_name, None)
if buf is None:
return False
if isinstance(buf.node.layout, NoneLayout):
if isinstance(buf.node, MutationOutput) and (
real_name := self.mutation_real_name.get(buf_name, None)
):
return is_none_layout(real_name)
return True
return False
for partition, skip_cudagraph in zip(
reversed(partitions), reversed(skip_cudagraphs)
):
output_names: OrderedSet[str] = OrderedSet()
for node in partition:
output_names.update(node.outputs_by_name.keys())
returned_output_names = output_names.intersection(unmet_output_names)
# all reads/writes are partition inputs except those generated
# within the partition and tensor constants
read_writes = dependencies.ReadWrites.merge_list(
[node.read_writes for node in partition]
)
# WeakDep is fake dependency on unused buffer. It should not appear
# in partition_input_names for inputs that are actually read or written.
partition_input_names = (
OrderedSet(
[
x.name
for x in read_writes.reads | read_writes.writes
if not is_none_layout(x.name)
]
)
- output_names
)
partition_input_names = OrderedSet(
self.mutation_real_name.get(name, name) for name in partition_input_names
)
buffer_names_to_free: OrderedSet[str] = OrderedSet()
for node in partition:
buffer_names_to_free.update(node.last_usage)
# buffer_names_to_free may contain buffers allocated in previous
# graph partitions. These buffers should also be a partition
# input.
extra_input_names = [
name
for name in (buffer_names_to_free - output_names)
if name in name_to_node
]
partition_input_names.update(extra_input_names)
input_nodes = {
name: name_to_node[name]
for name in partition_input_names
if name in name_to_node
}
input_deallocation = {
name: name in buffer_names_to_free
for name in partition_input_names
if name in name_to_node
}
# if an input tensor is not freed in the partition function, it should
# also be returned as an output. This brings benefits to cudagraph
# since the returned output tensor is a cudagraph managed tensor with
# a static tensor address.
extra_output_names = [
name
for name in partition_input_names
if name in name_to_node and name not in buffer_names_to_free
]
returned_output_names.update(extra_output_names)
returned_output_names = OrderedSet(
self.mutation_real_name.get(name, name) for name in returned_output_names
)
output_nodes = [
name_to_node[name]
for name in returned_output_names
if not is_none_layout(name)
]
constant_names = [
name for name in partition_input_names if name in V.graph.constants
]
symbol_inputs = self.get_graph_partition_symbol_inputs(partition, input_nodes)
partition_signature = GraphPartitionSignature(
symbol_inputs,
input_nodes,
output_nodes,
input_deallocation,
skip_cudagraph,
constant_names,
)
signatures.append(partition_signature)
unmet_output_names = partition_input_names.union(
unmet_output_names - returned_output_names
)
return signatures[::-1]
# ========================================
# torch 2.9 Inductor Scheduler monkeypatch
# ========================================
# This change monkeypatches a function in Inductor to work around the following
# bug: https://github.com/vllm-project/vllm/issues/26678
#
# The bug occurs when `use_inductor_graph_partition` is turned on and there
# exists operators inside of `splitting_ops` that have an in-place mutation. In
# vllm, this specifically occurs on the operator
# vllm.unified_attention_with_output. In this case, inductor does not populate
# the inductor IR's `origin_node` field, causing an assertion error when trying
# to access the node's `origin_node` field.
#
# So, we will monkeypatch torch._inductor.scheduler.Scheduler.should_partition
# so that it does not access the inductor IR node's `origin_node` field and just
# returns True if a node is registered as having a custom partition function.
# This is ok for now since vllm's implementation of the custom partition
# functions just return True.
# ========================================
def should_partition_patched(self, node, should_log: bool = False) -> bool:
# This is a patched version of
# torch._inductor.scheduler.Scheduler.should_partition that modifies
# the following piece of code so that we always return True:
# https://github.com/pytorch/pytorch/blob/ecb53078faf86ca1b33277df33b82985675bb011/torch/_inductor/scheduler.py#L4712-L4724
"""Return True if we should partition the inductor graph on this node"""
import torch._inductor.ir as ir
from torch._inductor.scheduler import (
BaseSchedulerNode,
FusedSchedulerNode,
)
from torch._inductor.utils import (
_unstable_customized_partition_wrapper,
is_cudagraph_unsafe_op,
maybe_log_cudagraph_partition,
)
# Allow users to manually specify if a node should be partitioned
# Can only do this for FallbackKernels
ir_node = node.node
if isinstance(ir_node, torch._inductor.ir.FallbackKernel) and (
op := ir_node.op_overload
):
op_overload_packet_name = op.name()
op_overload_name = (
f"{op_overload_packet_name}.{op._overloadname}"
if isinstance(op, torch._ops.OpOverload)
else op_overload_packet_name
)
if (
op_overload_packet_name
in torch._inductor.config.custom_should_partition_ops
or op_overload_name in torch._inductor.config.custom_should_partition_ops
):
assert isinstance(op, torch._ops.OpOverload)
return True
# When not using cudagraphs, keep all kernels in the `call` function
# instead of graph partition functions, since graph partition only brings
# benefit to cudagraph
if (
not torch._inductor.config.triton.cudagraphs
and _unstable_customized_partition_wrapper.wrapper is None
):
return True
# avoid duplicating logs when should_partition is called multiple times
# on the same node
def noop_log(msg: str, node: BaseSchedulerNode | None) -> None:
return
log_partition_reason = maybe_log_cudagraph_partition if should_log else noop_log
if isinstance(node, FusedSchedulerNode):
return any(self.should_partition(snode) for snode in node.snodes)
assert node.node is not None
if not node.is_gpu():
log_partition_reason("non gpu ops", node=node)
return True
if isinstance(node.node, ir.DeviceCopy):
log_partition_reason("DeviceCopy ops", node=node)
return True
if isinstance(node.node, ir.Conditional):
log_partition_reason("Conditional ops", node=node)
return True
if getattr(node.node, "unbacked_bindings", None):
log_partition_reason("unbacked binding ops", node=node)
return True
if is_cudagraph_unsafe_op(node.node):
log_partition_reason("CUDAGraph-unsafe custom ops", node=node)
return True
return False
def _update_scheduler_patched(self) -> None:
# Copied from torch._inductor.graph.GrahLowering._update_scheduler. Patches
# this method so that we can patch Scheduler.should_partition with the
# function above
"""
(Re)initializes the scheduler member. When initializing the scheduler, no CUBIN
files should be generated (to avoid biasing any benchmarks and pessimizing
fusion decisions).
"""
import torch._inductor.config as config
from torch._inductor.scheduler import Scheduler
Scheduler.should_partition = should_partition_patched
Scheduler.get_graph_partition_signature = get_graph_partition_signature_patched
with config.patch("triton.store_cubin", False):
self.scheduler = Scheduler(self.operations)
# ===================================================
# torch 2.9 Inductor get_raw_stream workaround
# ===================================================
# Workaround for TorchInductor autotune using get_raw_stream() without defining it.
# This occurs when compile_sizes > 1 in compilation_config.
# For more context, see https://github.com/vllm-project/vllm/issues/30905.
def _patch_get_raw_stream_if_needed():
"""Workaround for TorchInductor autotune get_raw_stream() bug."""
from vllm.utils.torch_utils import is_torch_equal
# Only apply the patch for torch 2.9.0 or 2.9.1
if is_torch_equal("2.9.0") or is_torch_equal("2.9.1"):
import builtins
# Check if CUDA functionality is available without initializing CUDA
# _cuda_getCurrentRawStream only exists in CUDA builds of PyTorch
if hasattr(torch._C, "_cuda_getCurrentRawStream"):
from torch._C import _cuda_getCurrentRawStream as _get_raw_stream
builtins.get_raw_stream = _get_raw_stream # type: ignore[attr-defined]
_patch_get_raw_stream_if_needed()
if is_torch_equal("2.9.0"):
from torch._inductor.codegen.wrapper import PythonWrapperCodegen
from torch._inductor.graph import GraphLowering
from torch.utils._config_module import _Config, _ConfigEntry
# `custom_should_partition_ops` is a new config after 2.9.0. So this would
# not overwrite any user configs.
torch._inductor.config._config["custom_should_partition_ops"] = _ConfigEntry(
_Config(default=[])
)
PythonWrapperCodegen.memory_plan_reuse = memory_plan_reuse_patched
GraphLowering._update_scheduler = _update_scheduler_patched
# ===================================================
# torch <2.12 GraphCaptureOutput.get_runtime_env monkeypatch
# ===================================================
@@ -586,55 +208,6 @@ if is_torch_equal_or_newer("2.10.0") and not is_torch_equal_or_newer("2.12.0.dev
GraphCaptureOutput.get_runtime_env = _patched_get_runtime_env
# ===================================================
# torch 2.10 FxGraphCachePickler.dumps ValueError fix
# ===================================================
# PyTorch 2.10's FxGraphCachePickler.dumps() doesn't catch ValueError,
# causing torch.compile cache failures when tensors with non-standard
# layouts (e.g. blocked-layout prepacked weights) are serialized.
# PyTorch mainline fixed this in pytorch/pytorch#176557 (merged 2026-03-04).
# This is a thin backport for 2.10 users; remove once 2.10 is dropped.
def _apply_fxgraphcache_pickle_patch(pickler_cls, bypass_cls):
"""Wrap pickler_cls.dumps to convert ValueError into bypass_cls.
Idempotent: sets `_vllm_fxgraph_dumps_patched` on the class after the
first apply to prevent re-application. The wrapper function is also
marked with `_vllm_patched` as an additional safeguard.
"""
if getattr(pickler_cls, "_vllm_fxgraph_dumps_patched", False):
return
original_dumps = pickler_cls.dumps
if hasattr(original_dumps, "_vllm_patched"):
return
def patched_dumps(self, obj):
try:
return original_dumps(self, obj)
except ValueError as e:
raise bypass_cls("Failed to pickle cache key") from e
patched_dumps._vllm_patched = True # type: ignore[attr-defined]
pickler_cls.dumps = patched_dumps
pickler_cls._vllm_fxgraph_dumps_patched = True # type: ignore[attr-defined]
def _patch_fxgraphcache_pickle_if_needed():
"""Apply FxGraphCachePickler.dumps ValueError backport when on torch 2.10.x."""
from vllm.utils.torch_utils import is_torch_equal_or_newer
if not is_torch_equal_or_newer("2.10.0") or is_torch_equal_or_newer("2.11.0"):
return
from torch._inductor.codecache import BypassFxGraphCache, FxGraphCachePickler
_apply_fxgraphcache_pickle_patch(FxGraphCachePickler, BypassFxGraphCache)
_patch_fxgraphcache_pickle_if_needed()
# ===================================================
# torch 2.11 Inductor cpp codegen indirect_assert scalar-mask fix
# ===================================================