[Bugfix] Resolve B12X modules before Dynamo tracing

tests/model_executor/kernels/test_b12x_linear.py

@@ -130,6 +130,27 @@ def test_b12x_backend_registration_priority_and_selection(
     assert isinstance(initializer(**kwargs), kernel_cls)
 
 
+def test_b12x_module_lookup_is_dynamo_safe(monkeypatch) -> None:
+    import vllm.utils.b12x as b12x_utils
+
+    module = types.ModuleType("b12x.gemm.blockscaled")
+    module.run = lambda x: x + 1  # type: ignore[attr-defined]
+    monkeypatch.setitem(
+        b12x_utils._B12X_SUBMODULES,
+        "b12x.gemm.blockscaled",
+        module,
+    )
+
+    @torch.compile(backend="eager", fullgraph=True)
+    def forward(x: torch.Tensor) -> torch.Tensor:
+        blockscaled = b12x_utils.get_b12x_blockscaled()
+        assert blockscaled is not None
+        return blockscaled.run(x)  # type: ignore[attr-defined]
+
+    x = torch.ones(1)
+    torch.testing.assert_close(forward(x), x + 1)
+
+
 def test_b12x_tensor_fp8_can_implement_supported_config() -> None:
     config = FP8ScaledMMLinearLayerConfig(
         activation_quant_key=kFp8StaticTensorSym,

vllm/utils/b12x.py

@@ -1,8 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Lazy accessors for the optional ``b12x`` package."""
+"""Accessors for the optional ``b12x`` package."""
 
-import functools
 import importlib
 import importlib.util
 from collections.abc import Callable, Hashable, Iterable
@@ -20,22 +19,40 @@ class B12xWarmupUnit:
     compile: Callable[[], None]
 
 
-@functools.cache
-def has_b12x() -> bool:
-    """Return whether the B12X package is installed."""
-    return importlib.util.find_spec("b12x") is not None
+_HAS_B12X = importlib.util.find_spec("b12x") is not None
 
 
-@functools.cache
-def _get_submodule(module_name: str) -> ModuleType | None:
-    if not has_b12x():
+def _import_submodule(module_name: str) -> ModuleType | None:
+    if not _HAS_B12X:
         return None
     try:
         return importlib.import_module(module_name)
     except (ImportError, ModuleNotFoundError):
         return None
 
 
+_B12X_SUBMODULES = {
+    module_name: _import_submodule(module_name)
+    for module_name in (
+        "b12x.gemm.blockscaled",
+        # TODO: Remove once B12X exposes the scale-swizzle API publicly.
+        "b12x._lib.intrinsics",
+        "b12x.gemm.mxfp8_linear",
+        "b12x.gemm.tensor_fp8_linear",
+        "b12x.moe.fused_moe",
+    )
+}
+
+
+def has_b12x() -> bool:
+    """Return whether the B12X package is installed."""
+    return _HAS_B12X
+
+
+def _get_submodule(module_name: str) -> ModuleType | None:
+    return _B12X_SUBMODULES.get(module_name)
+
+
 def get_b12x_blockscaled() -> ModuleType | None:
     return _get_submodule("b12x.gemm.blockscaled")