From 04af94d150e5921801a6fdafeb042a7b983c250f Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 14 Jul 2026 22:30:27 +0800 Subject: [PATCH] fix: avoid tilelang cuda runtime pollution (#30870) --- .../device_communicators/cuda_wrapper.py | 35 ++++++++----- python/sglang/srt/models/deepseek_v4.py | 52 ++++++++++++++----- .../unit/distributed/test_cuda_wrapper.py | 38 ++++++++++++++ 3 files changed, 98 insertions(+), 27 deletions(-) create mode 100644 test/registered/unit/distributed/test_cuda_wrapper.py diff --git a/python/sglang/srt/distributed/device_communicators/cuda_wrapper.py b/python/sglang/srt/distributed/device_communicators/cuda_wrapper.py index eec8ebafa..794af8010 100644 --- a/python/sglang/srt/distributed/device_communicators/cuda_wrapper.py +++ b/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: - found = True - break - if not found: + if lib_name not in line or "/" not in line: + continue + path = line[line.index("/") :].strip() + 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 - start = line.index("/") - path = line[start:].strip() - filename = path.split("/")[-1] - assert filename.rpartition(".so")[0].startswith( - lib_name - ), f"Unexpected filename: {filename} for library {lib_name}" - return path + + # TileLang ships a ``libcudart_stub.so`` that is only sufficient for its + # JIT loader. It can precede the actual CUDA runtime in /proc/self/maps; + # choosing it makes CUDA IPC and FlashInfer all-reduce fail when resolving + # symbols such as cudaDeviceReset. + for path in candidates: + if "stub" not in os.path.basename(path): + return path + return candidates[0] class CudaRTLibrary: diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 00377253a..3fc207069 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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, diff --git a/test/registered/unit/distributed/test_cuda_wrapper.py b/test/registered/unit/distributed/test_cuda_wrapper.py new file mode 100644 index 000000000..e7383c63e --- /dev/null +++ b/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"]))