diff --git a/tests/model_executor/kernels/test_b12x_linear.py b/tests/model_executor/kernels/test_b12x_linear.py index 1efc521bd8b4..36a0896ea609 100644 --- a/tests/model_executor/kernels/test_b12x_linear.py +++ b/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, diff --git a/vllm/utils/b12x.py b/vllm/utils/b12x.py index d7882fb6f3ab..dd4fa2940b48 100644 --- a/vllm/utils/b12x.py +++ b/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,15 +19,11 @@ 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) @@ -36,6 +31,28 @@ def _get_submodule(module_name: str) -> ModuleType | None: 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")