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

2025-11-02 13:15:05 -08:00
parent c9db79117f
commit 20315697f4
20 changed files with 61 additions and 94 deletions
@@ -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,