move all get_stream in sgl_kernel to c++ to reduce the launch overhead (#12521)

python/sglang/srt/_custom_ops.py

@@ -4,32 +4,20 @@ from typing import List, Optional, Tuple
 
 import torch
 
-from sglang.srt.utils import get_bool_env_var, is_hip, is_hpu, is_npu
+from sglang.srt.utils import is_hip, is_hpu, is_npu
 
 logger = logging.getLogger(__name__)
-use_vllm_custom_allreduce = get_bool_env_var(
-    "USE_VLLM_CUSTOM_ALLREDUCE", default="false"
-)
+
 
 if not is_hpu():
-    # ROCm does not use vllm custom allreduce
-    if use_vllm_custom_allreduce and not is_hip():
-        try:
-            import vllm._C  # noqa: F401
-        except ImportError as e:
-            logger.warning("Failed to import from vllm._C with %r", e)
-    else:
-        try:
-            import sgl_kernel
-        except ImportError as e:
-            logger.warning("Failed to import from custom_ar with %r", e)
+    try:
+        import sgl_kernel
+    except ImportError as e:
+        logger.warning("Failed to import from custom_ar with %r", e)
 
 
 if not is_hip() and not is_npu():
-    if use_vllm_custom_allreduce:
-        custom_op = torch.ops._C_custom_ar
-    else:
-        custom_op = sgl_kernel.allreduce
+    custom_op = sgl_kernel.allreduce
 
     # custom allreduce
     def init_custom_ar(

python/sglang/srt/configs/falcon_h1.py

@@ -19,7 +19,6 @@ from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
 
 from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
-from sglang.srt.layers.dp_attention import get_tensor_model_parallel_world_size
 
 logger = logging.get_logger(__name__)
 
@@ -297,8 +296,10 @@ class FalconH1Config(PretrainedConfig):
 
     @property
     def mamba2_cache_params(self):
+        from sglang.srt.layers.dp_attention import get_attention_tp_size
+
         shape = Mamba2StateShape.create(
-            tp_world_size=get_tensor_model_parallel_world_size(),
+            tp_world_size=get_attention_tp_size(),
             intermediate_size=self.mamba_intermediate,
             n_groups=self.mamba_n_groups,
             num_heads=self.mamba_n_heads,

python/sglang/srt/configs/nemotron_h.py

@@ -20,7 +20,6 @@ from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
 
 from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 
 logger = logging.get_logger(__name__)
 
@@ -273,6 +272,8 @@ class NemotronHConfig(PretrainedConfig):
 
     @property
     def mamba2_cache_params(self) -> Mamba2CacheParams:
+        from sglang.srt.layers.dp_attention import get_attention_tp_size
+
         shape = Mamba2StateShape.create(
             tp_world_size=get_attention_tp_size(),
             intermediate_size=self.mamba_num_heads * self.mamba_head_dim,

python/sglang/srt/configs/qwen3_next.py

@@ -21,7 +21,6 @@ from transformers.modeling_rope_utils import rope_config_validation
 from transformers.utils import logging
 
 from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
-from sglang.srt.layers.dp_attention import get_attention_tp_size
 
 logger = logging.get_logger(__name__)
 
@@ -277,6 +276,8 @@ class Qwen3NextConfig(PretrainedConfig):
 
     @property
     def mamba2_cache_params(self) -> Mamba2CacheParams:
+        from sglang.srt.layers.dp_attention import get_attention_tp_size
+
         shape = Mamba2StateShape.create(
             tp_world_size=get_attention_tp_size(),
             intermediate_size=self.linear_value_head_dim * self.linear_num_value_heads,

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

@@ -21,24 +21,19 @@ from sglang.srt.distributed.parallel_state import in_the_same_node_as
 from sglang.srt.environ import envs
 from sglang.srt.utils import is_cuda, is_hip, log_info_on_rank0
 
-logger = logging.getLogger(__name__)
+try:
+    # Use custom allreduce from sgl kernel (ROCM and TRT-LLM)
+    import sgl_kernel  # noqa: F401
+
+    custom_ar = True
+except ImportError:
+    # For CPUs
+    custom_ar = False
+
 
 _is_cuda = is_cuda()
 _is_hip = is_hip()
 
-
-try:
-    if ops.use_vllm_custom_allreduce and not _is_hip:
-        # Use vLLM custom allreduce
-        ops.meta_size()
-    else:
-        # Use custom allreduce from sgl kernel (ROCM and TRT-LLM)
-        import sgl_kernel  # noqa: F401
-    custom_ar = True
-except Exception:
-    # For CPUs
-    custom_ar = False
-
 logger = logging.getLogger(__name__)
 
 

python/sglang/srt/environ.py

@@ -229,7 +229,6 @@ class Envs:
     SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False)
 
     # vLLM dependencies (TODO: they have been deprecated, we can remove them safely)
-    USE_VLLM_CUSTOM_ALLREDUCE = EnvBool(False)
     USE_VLLM_CUTLASS_W8A8_FP8_KERNEL = EnvBool(False)
 
     USE_TRITON_W8A8_FP8_KERNEL = EnvBool(False)

python/sglang/srt/utils/common.py

@@ -303,6 +303,7 @@ def xpu_has_xmx_support():
     return False
 
 
+@lru_cache(maxsize=1)
 def is_flashinfer_available():
     """
     Check whether flashinfer is available.