diff --git a/python/sglang/multimodal_gen/runtime/utils/common.py b/python/sglang/multimodal_gen/runtime/utils/common.py index a921130b5..b3826e685 100644 --- a/python/sglang/multimodal_gen/runtime/utils/common.py +++ b/python/sglang/multimodal_gen/runtime/utils/common.py @@ -398,14 +398,15 @@ def get_bool_env_var(name: str, default: str = "false") -> bool: return value in truthy_values -try: - import sgl_kernel # noqa: F401 +@lru_cache(maxsize=1) +def _is_intel_amx_backend_available(): + try: + import sgl_kernel # noqa: F401 + + return hasattr(torch.ops.sgl_kernel, "convert_weight_packed") + except Exception: + return False - is_intel_amx_backend_available = hasattr( - torch.ops.sgl_kernel, "convert_weight_packed" - ) -except: - is_intel_amx_backend_available = False try: # move torch.cpu._is_amx_tile_supported() from cpu_has_amx_support @@ -416,7 +417,7 @@ except: def cpu_has_amx_support(): - return is_amx_tile_supported and is_intel_amx_backend_available + return is_amx_tile_supported and _is_intel_amx_backend_available() def use_intel_amx_backend(layer): diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 690d4f629..b2a023c03 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -316,14 +316,15 @@ def is_sm121() -> bool: return is_cuda() and torch.cuda.get_device_capability() == (12, 1) -try: - import sgl_kernel # noqa: F401 +@lru_cache(maxsize=1) +def _is_intel_amx_backend_available(): + try: + import sgl_kernel # noqa: F401 + + return hasattr(torch.ops.sgl_kernel, "convert_weight_packed") + except Exception: + return False - is_intel_amx_backend_available = hasattr( - torch.ops.sgl_kernel, "convert_weight_packed" - ) -except: - is_intel_amx_backend_available = False try: # move torch.cpu._is_amx_tile_supported() from cpu_has_amx_support @@ -334,7 +335,7 @@ except: def cpu_has_amx_support(): - return is_amx_tile_supported and is_intel_amx_backend_available + return is_amx_tile_supported and _is_intel_amx_backend_available() def use_intel_amx_backend(layer): diff --git a/test/registered/kernels/ops/layernorm/test_kernels_namespace.py b/test/registered/kernels/ops/layernorm/test_kernels_namespace.py index 3c71dde2b..1dec217dc 100644 --- a/test/registered/kernels/ops/layernorm/test_kernels_namespace.py +++ b/test/registered/kernels/ops/layernorm/test_kernels_namespace.py @@ -181,6 +181,44 @@ def test_platform_detect_does_not_raise(): assert PlatformInfo.detect().device_type in ("cpu", "cuda", "hip", "npu") +@pytest.mark.parametrize( + "relative_path", + ( + "srt/utils/common.py", + "multimodal_gen/runtime/utils/common.py", + ), +) +def test_amx_backend_probe_is_lazy(relative_path): + loader = ( + "package = importlib.util.find_spec('sglang'); " + "path = pathlib.Path(next(iter(package.submodule_search_locations))) / " + f"{relative_path!r}; " + "spec = importlib.util.spec_from_file_location('_common_under_test', path); " + "module = importlib.util.module_from_spec(spec); " + "sys.modules[spec.name] = module; " + "spec.loader.exec_module(module)" + ) + code = "; ".join( + ( + "import builtins, importlib.util, pathlib, sys", + "from unittest import mock", + "real_import = builtins.__import__", + "import_mock = mock.Mock(wraps=real_import)", + "builtins.__import__ = import_mock", + loader, + "builtins.__import__ = real_import", + "attempted = any(call.args and call.args[0] == 'sgl_kernel' " + "for call in import_mock.call_args_list)", + "print('DIRTY' if attempted else 'CLEAN')", + ) + ) + result = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True + ) + assert result.returncode == 0, result.stderr + assert "CLEAN" in result.stdout + + def test_import_stays_metadata_only(): # Importing the namespace must not pull in the AOT backend (sgl_kernel) or # the JIT compilation infra (sglang.kernels.jit), which import torch / nvcc.