Fix eager AMX backend probe imports (#34496)

Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

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:
+@lru_cache(maxsize=1)
+def _is_intel_amx_backend_available():
+    try:
         import sgl_kernel  # noqa: F401
 
-    is_intel_amx_backend_available = hasattr(
+        return hasattr(torch.ops.sgl_kernel, "convert_weight_packed")
-        torch.ops.sgl_kernel, "convert_weight_packed"
+    except Exception:
-    )
+        return False
-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):

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:
+@lru_cache(maxsize=1)
+def _is_intel_amx_backend_available():
+    try:
         import sgl_kernel  # noqa: F401
 
-    is_intel_amx_backend_available = hasattr(
+        return hasattr(torch.ops.sgl_kernel, "convert_weight_packed")
-        torch.ops.sgl_kernel, "convert_weight_packed"
+    except Exception:
-    )
+        return False
-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):

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.