fix: avoid tilelang cuda runtime pollution (#30870)

python/sglang/srt/distributed/device_communicators/cuda_wrapper.py

@@ -9,6 +9,7 @@ convenient for use when we just need to call a few functions.
 
 import ctypes
 import logging
+import os
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional
 
@@ -47,24 +48,30 @@ def find_loaded_library(lib_name) -> Optional[str]:
     shared libraries loaded by the process. We can use this file to find the path of the
     a loaded library.
     """  # noqa
-    found = False
+    candidates = []
     with open("/proc/self/maps") as f:
         for line in f:
-            if lib_name in line:
+            if lib_name not in line or "/" not in line:
-                found = True
+                continue
-                break
+            path = line[line.index("/") :].strip()
-    if not found:
+            if path.endswith(" (deleted)"):
+                path = path[: -len(" (deleted)")]
+            filename = os.path.basename(path)
+            if filename.rpartition(".so")[0].startswith(lib_name):
+                candidates.append(path)
+
+    if not candidates:
         # the library is not loaded in the current process
         return None
-    # if lib_name is libcudart, we need to match a line with:
+
-    # address /path/to/libcudart-hash.so.11.0
+    # TileLang ships a ``libcudart_stub.so`` that is only sufficient for its
-    start = line.index("/")
+    # JIT loader.  It can precede the actual CUDA runtime in /proc/self/maps;
-    path = line[start:].strip()
+    # choosing it makes CUDA IPC and FlashInfer all-reduce fail when resolving
-    filename = path.split("/")[-1]
+    # symbols such as cudaDeviceReset.
-    assert filename.rpartition(".so")[0].startswith(
+    for path in candidates:
-        lib_name
+        if "stub" not in os.path.basename(path):
-    ), f"Unexpected filename: {filename} for library {lib_name}"
+            return path
-    return path
+    return candidates[0]
 
 
 class CudaRTLibrary:

python/sglang/srt/models/deepseek_v4.py

@@ -1,13 +1,17 @@
 from __future__ import annotations
 
 import concurrent.futures
+import functools
 import logging
 import time
 from contextlib import nullcontext
 from typing import (
     TYPE_CHECKING,
+    Any,
+    Callable,
     Iterable,
     List,
+    NamedTuple,
     Optional,
     Set,
     Tuple,
@@ -129,15 +133,6 @@ if not _is_hip:
         prepare_context_parallel_metadata,
     )
 
-if _is_xpu:
-    from sgl_kernel import hc_split_sinkhorn
-else:
-    from sglang.kernels.ops.layernorm.mhc import (
-        hc_split_sinkhorn,
-        mhc_fused_post_pre,
-        npu_hc_pre,
-    )
-
 from sglang.srt.utils import (
     LazyValue,
     add_prefix,
@@ -155,6 +150,37 @@ from sglang.srt.utils.hf_transformers_utils import get_rope_config
 if _is_npu:
     import torch_npu
 
+
+class MhcOps(NamedTuple):
+    hc_split_sinkhorn: Callable[..., Any]
+    mhc_fused_post_pre: Optional[Callable[..., Any]]
+    npu_hc_pre: Optional[Callable[..., Any]]
+
+
+@functools.cache
+def _get_mhc_ops() -> MhcOps:
+    """Load MHC kernels only when a DeepSeek-V4 layer needs them.
+
+    Model modules are imported eagerly by the registry.  Importing
+    ``sglang.kernels.ops.layernorm.mhc`` owns TileLang-backed MHC kernels.
+    Import it only when a DeepSeek-V4 layer executes so registry discovery
+    cannot initialize an optional CUDA runtime before unrelated models set up
+    their communication workspaces.  DeepSeek-V4 is the sole consumer here.
+    """
+    if _is_xpu:
+        from sgl_kernel import hc_split_sinkhorn
+
+        return MhcOps(hc_split_sinkhorn, None, None)
+
+    from sglang.kernels.ops.layernorm.mhc import (
+        hc_split_sinkhorn,
+        mhc_fused_post_pre,
+        npu_hc_pre,
+    )
+
+    return MhcOps(hc_split_sinkhorn, mhc_fused_post_pre, npu_hc_pre)
+
+
 logger = logging.getLogger(__name__)
 
 _FP8_WO_A_GEMM = envs.SGLANG_OPT_FP8_WO_A_GEMM.get()
@@ -1349,7 +1375,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         shape, dtype = x.size(), x.dtype
 
         if _is_npu:
-            return npu_hc_pre(
+            return _get_mhc_ops().npu_hc_pre(
                 x,
                 hc_fn,
                 hc_scale,
@@ -1426,7 +1452,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         else:
             x_flat, mixes = hc_pre_torch_impl(x, hc_fn)
 
-        pre, post, comb = hc_split_sinkhorn(
+        pre, post, comb = _get_mhc_ops().hc_split_sinkhorn(
             mixes,
             hc_scale,
             hc_base,
@@ -1497,7 +1523,7 @@ class DeepseekV4DecoderLayer(nn.Module):
         use_fused = self.use_fused_mhc_post_pre
 
         if prev_residual is not None and use_fused:
-            residual, post, comb, hidden_states = mhc_fused_post_pre(
+            residual, post, comb, hidden_states = _get_mhc_ops().mhc_fused_post_pre(
                 hidden_states,
                 prev_residual,
                 prev_post,
@@ -1567,7 +1593,7 @@ class DeepseekV4DecoderLayer(nn.Module):
             if fused_mhc is not None:
                 residual, hidden_states, post, comb, norm_fused = fused_mhc
             else:
-                residual, post, comb, hidden_states = mhc_fused_post_pre(
+                residual, post, comb, hidden_states = _get_mhc_ops().mhc_fused_post_pre(
                     hidden_states,
                     residual,
                     post.unsqueeze(-1) if post.ndim == 2 else post,

test/registered/unit/distributed/test_cuda_wrapper.py

@@ -0,0 +1,38 @@
+import io
+
+import pytest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=2, stage="base-b", runner_config="1-gpu-small")
+
+from sglang.srt.distributed.device_communicators import cuda_wrapper
+
+
+def test_find_loaded_library_prefers_real_cudart_over_tilelang_stub(monkeypatch):
+    maps = """\
+7f000000-7f010000 r-xp 00000000 00:00 0 /site-packages/tilelang/lib/libcudart_stub.so
+7f020000-7f030000 r-xp 00000000 00:00 0 /cuda/lib64/libcudart.so.13
+"""
+
+    monkeypatch.setattr("builtins.open", lambda *args, **kwargs: io.StringIO(maps))
+
+    assert (
+        cuda_wrapper.find_loaded_library("libcudart") == "/cuda/lib64/libcudart.so.13"
+    )
+
+
+def test_find_loaded_library_strips_deleted_suffix(monkeypatch):
+    maps = """\
+7f020000-7f030000 r-xp 00000000 00:00 0 /cuda/lib64/libcudart.so.13 (deleted)
+"""
+
+    monkeypatch.setattr("builtins.open", lambda *args, **kwargs: io.StringIO(maps))
+
+    assert (
+        cuda_wrapper.find_loaded_library("libcudart") == "/cuda/lib64/libcudart.so.13"
+    )
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))