[Bugfix] Resolve B12X modules before Dynamo tracing (#53326)
Signed-off-by: Luke Alonso <lalonso@gmail.com>
Luke Alonso committed
Aug 25, 2026 at 09:20 UTC
3454335f7b64e31fd385e36eb590d76d08c84a92
2 files changed
+47
-9
tests/model_executor/kernels/test_b12x_linear.py
+21
@@ -130,6 +130,27 @@ def test_b12x_backend_registration_priority_and_selection(
130
assert isinstance(initializer(**kwargs), kernel_cls)
131
132
133
+def test_b12x_module_lookup_is_dynamo_safe(monkeypatch) -> None:
134
+ import vllm.utils.b12x as b12x_utils
135
+
136
+ module = types.ModuleType("b12x.gemm.blockscaled")
137
+ module.run = lambda x: x + 1 # type: ignore[attr-defined]
138
+ monkeypatch.setitem(
139
+ b12x_utils._B12X_SUBMODULES,
140
+ "b12x.gemm.blockscaled",
141
+ module,
142
+ )
143
+
144
+ @torch.compile(backend="eager", fullgraph=True)
145
+ def forward(x: torch.Tensor) -> torch.Tensor:
146
+ blockscaled = b12x_utils.get_b12x_blockscaled()
147
+ assert blockscaled is not None
148
+ return blockscaled.run(x) # type: ignore[attr-defined]
149
+
150
+ x = torch.ones(1)
151
+ torch.testing.assert_close(forward(x), x + 1)
152
+
153
+
154
def test_b12x_tensor_fp8_can_implement_supported_config() -> None:
155
config = FP8ScaledMMLinearLayerConfig(
156
activation_quant_key=kFp8StaticTensorSym,
vllm/utils/b12x.py
+26
-9
@@ -1,8 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Lazy accessors for the optional ``b12x`` package."""
3
+"""Accessors for the optional ``b12x`` package."""
4
5
-import functools
5
import importlib
6
import importlib.util
7
from collections.abc import Callable, Hashable, Iterable
@@ -20,15 +19,11 @@ class B12xWarmupUnit:
19
compile: Callable[[], None]
20
21
23
-@functools.cache
24
-def has_b12x() -> bool:
25
- """Return whether the B12X package is installed."""
26
- return importlib.util.find_spec("b12x") is not None
22
+_HAS_B12X = importlib.util.find_spec("b12x") is not None
23
24
29
-@functools.cache
30
-def _get_submodule(module_name: str) -> ModuleType | None:
31
- if not has_b12x():
25
+def _import_submodule(module_name: str) -> ModuleType | None:
26
+ if not _HAS_B12X:
27
return None
28
try:
29
return importlib.import_module(module_name)
@@ -36,6 +31,28 @@ def _get_submodule(module_name: str) -> ModuleType | None:
31
return None
32
33
34
+_B12X_SUBMODULES = {
35
+ module_name: _import_submodule(module_name)
36
+ for module_name in (
37
+ "b12x.gemm.blockscaled",
38
+ # TODO: Remove once B12X exposes the scale-swizzle API publicly.
39
+ "b12x._lib.intrinsics",
40
+ "b12x.gemm.mxfp8_linear",
41
+ "b12x.gemm.tensor_fp8_linear",
42
+ "b12x.moe.fused_moe",
43
+ )
44
+}
45
+
46
+
47
+def has_b12x() -> bool:
48
+ """Return whether the B12X package is installed."""
49
+ return _HAS_B12X
50
+
51
+
52
+def _get_submodule(module_name: str) -> ModuleType | None:
53
+ return _B12X_SUBMODULES.get(module_name)
54
+
55
+
56
def get_b12x_blockscaled() -> ModuleType | None:
57
return _get_submodule("b12x.gemm.blockscaled")
58