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 5, 07:15 UTC
[ZenCPU] Make PT Backport Patch Accessible to vLLM (#38205)
Signed-off-by: Lalithnarayan C <Lalithnarayan.C@amd.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Lalithnarayan C committed
Apr 10, 2026 at 13:59 UTC
8e8a3becd1e499a539bf1b9b0e4d135b82228727
3 files changed
+152
-43
tests/test_fxgraphcache_pickle_patch.py
new
+103
@@ -0,0 +1,103 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the FxGraphCachePickler.dumps ValueError patch in env_override.py.
4
+
5
+Validates that _apply_fxgraphcache_pickle_patch correctly wraps a pickler's
6
+dumps method to convert ValueError into a bypass exception, without affecting
7
+other exception types or normal return values.
8
+"""
9
+
10
+import pytest
11
+
12
+from vllm.env_override import _apply_fxgraphcache_pickle_patch
13
+
14
+
15
+class _BypassStub(Exception):
16
+ """Stand-in for BypassFxGraphCache in unit tests."""
17
+
18
+
19
+class TestApplyFxgraphcachePicklePatch:
20
+ def test_valueerror_converted_to_bypass(self):
21
+ class Pickler:
22
+ def dumps(self, obj):
23
+ raise ValueError("can't serialize blocked layout")
24
+
25
+ _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
26
+
27
+ with pytest.raises(_BypassStub, match="Failed to pickle cache key"):
28
+ Pickler().dumps(object())
29
+
30
+ def test_original_valueerror_chained(self):
31
+ class Pickler:
32
+ def dumps(self, obj):
33
+ raise ValueError("bad tensor layout")
34
+
35
+ _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
36
+
37
+ with pytest.raises(_BypassStub) as exc_info:
38
+ Pickler().dumps(object())
39
+
40
+ cause = exc_info.value.__cause__
41
+ assert isinstance(cause, ValueError)
42
+ assert str(cause) == "bad tensor layout"
43
+
44
+ def test_non_valueerror_propagates(self):
45
+ class Pickler:
46
+ def dumps(self, obj):
47
+ raise TypeError("unexpected type")
48
+
49
+ _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
50
+
51
+ with pytest.raises(TypeError, match="unexpected type"):
52
+ Pickler().dumps(object())
53
+
54
+ def test_normal_return_preserved(self):
55
+ sentinel = b"serialized-graph-key"
56
+
57
+ class Pickler:
58
+ def dumps(self, obj):
59
+ return sentinel
60
+
61
+ _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
62
+
63
+ assert Pickler().dumps(object()) is sentinel
64
+
65
+ def test_idempotent(self):
66
+ class Pickler:
67
+ def dumps(self, obj):
68
+ return b"ok"
69
+
70
+ _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
71
+ first_dumps = Pickler.dumps
72
+ _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
73
+
74
+ assert Pickler.dumps is first_dumps
75
+
76
+ def test_sentinel_attribute_set(self):
77
+ class Pickler:
78
+ def dumps(self, obj):
79
+ return b"ok"
80
+
81
+ assert not hasattr(Pickler.dumps, "_vllm_patched")
82
+ assert not getattr(Pickler, "_vllm_fxgraph_dumps_patched", False)
83
+
84
+ _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub)
85
+
86
+ assert Pickler.dumps._vllm_patched is True # type: ignore[attr-defined]
87
+ assert Pickler._vllm_fxgraph_dumps_patched is True # type: ignore[attr-defined]
88
+
89
+
90
+def test_patch_applied_in_current_environment():
91
+ """Integration: verify patch state matches current torch version."""
92
+ from torch._inductor.codecache import FxGraphCachePickler
93
+
94
+ from vllm.utils.torch_utils import is_torch_equal_or_newer
95
+
96
+ should_be_patched = is_torch_equal_or_newer(
97
+ "2.10.0"
98
+ ) and not is_torch_equal_or_newer("2.11.0")
99
+
100
+ assert getattr(FxGraphCachePickler, "_vllm_fxgraph_dumps_patched", False) == (
101
+ should_be_patched
102
+ )
103
+ assert hasattr(FxGraphCachePickler.dumps, "_vllm_patched") == should_be_patched
vllm/env_override.py
+49
@@ -586,3 +586,52 @@ if is_torch_equal_or_newer("2.10.0") and not is_torch_equal_or_newer("2.12.0"):
586
return runtime_env
587
588
GraphCaptureOutput.get_runtime_env = _patched_get_runtime_env
589
+
590
+# ===================================================
591
+# torch 2.10 FxGraphCachePickler.dumps ValueError fix
592
+# ===================================================
593
+# PyTorch 2.10's FxGraphCachePickler.dumps() doesn't catch ValueError,
594
+# causing torch.compile cache failures when tensors with non-standard
595
+# layouts (e.g. blocked-layout prepacked weights) are serialized.
596
+# PyTorch mainline fixed this in pytorch/pytorch#176557 (merged 2026-03-04).
597
+# This is a thin backport for 2.10 users; remove once 2.10 is dropped.
598
+
599
+
600
+def _apply_fxgraphcache_pickle_patch(pickler_cls, bypass_cls):
601
+ """Wrap pickler_cls.dumps to convert ValueError into bypass_cls.
602
+
603
+ Idempotent: sets `_vllm_fxgraph_dumps_patched` on the class after the
604
+ first apply to prevent re-application. The wrapper function is also
605
+ marked with `_vllm_patched` as an additional safeguard.
606
+ """
607
+ if getattr(pickler_cls, "_vllm_fxgraph_dumps_patched", False):
608
+ return
609
+
610
+ original_dumps = pickler_cls.dumps
611
+ if hasattr(original_dumps, "_vllm_patched"):
612
+ return
613
+
614
+ def patched_dumps(self, obj):
615
+ try:
616
+ return original_dumps(self, obj)
617
+ except ValueError as e:
618
+ raise bypass_cls("Failed to pickle cache key") from e
619
+
620
+ patched_dumps._vllm_patched = True # type: ignore[attr-defined]
621
+ pickler_cls.dumps = patched_dumps
622
+ pickler_cls._vllm_fxgraph_dumps_patched = True # type: ignore[attr-defined]
623
+
624
+
625
+def _patch_fxgraphcache_pickle_if_needed():
626
+ """Apply FxGraphCachePickler.dumps ValueError backport when on torch 2.10.x."""
627
+ from vllm.utils.torch_utils import is_torch_equal_or_newer
628
+
629
+ if not is_torch_equal_or_newer("2.10.0") or is_torch_equal_or_newer("2.11.0"):
630
+ return
631
+
632
+ from torch._inductor.codecache import BypassFxGraphCache, FxGraphCachePickler
633
+
634
+ _apply_fxgraphcache_pickle_patch(FxGraphCachePickler, BypassFxGraphCache)
635
+
636
+
637
+_patch_fxgraphcache_pickle_if_needed()
vllm/platforms/zen_cpu.py
-43
@@ -1,17 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import TYPE_CHECKING
5
-
4
from vllm.logger import init_logger
5
from vllm.platforms.cpu import CpuPlatform
8
-from vllm.utils.torch_utils import is_torch_equal_or_newer
6
7
logger = init_logger(__name__)
8
12
-if TYPE_CHECKING:
13
- from vllm.config import VllmConfig
14
-
9
10
class ZenCpuPlatform(CpuPlatform):
11
"""CPU platform with AMD Zen (ZenDNN/zentorch) optimizations.
@@ -28,40 +22,3 @@ class ZenCpuPlatform(CpuPlatform):
22
def is_zen_cpu(self) -> bool:
23
# is_cpu() also returns True for this platform (inherited from CpuPlatform).
24
return True
31
-
32
- @classmethod
33
- def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
34
- super().check_and_update_config(vllm_config)
35
- cls._apply_pytorch_backports()
36
-
37
- @classmethod
38
- def _apply_pytorch_backports(cls):
39
- """Backport PyTorch mainline fixes missing in 2.10.
40
-
41
- PyTorch 2.10 has a bug in FxGraphCachePickler.dumps that doesn't
42
- catch ValueError, causing torch.compile cache misses. Remove this
43
- once we drop PyTorch 2.10 support. PT mainline already has this fix.
44
- """
45
- if not is_torch_equal_or_newer("2.10.0") or is_torch_equal_or_newer("2.11.0"):
46
- return
47
-
48
- cls._patch_fxgraphcache_pickle()
49
-
50
- @classmethod
51
- def _patch_fxgraphcache_pickle(cls):
52
- """Backport mainline ValueError fix to FxGraphCachePickler.dumps()."""
53
- from torch._inductor.codecache import BypassFxGraphCache, FxGraphCachePickler
54
-
55
- original_dumps = FxGraphCachePickler.dumps
56
- if hasattr(original_dumps, "_zen_patched"):
57
- return
58
-
59
- def patched_dumps(self, obj):
60
- try:
61
- return original_dumps(self, obj)
62
- except ValueError as e:
63
- raise BypassFxGraphCache("Failed to pickle cache key") from e
64
-
65
- patched_dumps._zen_patched = True # type: ignore[attr-defined]
66
- FxGraphCachePickler.dumps = patched_dumps
67
- logger.info("[zen_cpu] Patched FxGraphCachePickler.dumps (ValueError fix)")