[Deps] Upgrade CUDA PyTorch stack to 2.13 (#28836)

Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>

python/pyproject.toml

@@ -67,19 +67,19 @@ dependencies = [
   "scipy",
   "sentencepiece",
   "setproctitle",
-  "sgl-deep-gemm==0.1.5.post1",
-  "sglang-kernel==0.4.5",
+  "sgl-deep-gemm==0.1.5.post2",
+  "sglang-kernel==0.4.6.post1",
   "smg-grpc-servicer>=0.5.0",
   "soundfile==0.13.1",
   "tiktoken",
   "tilelang==0.1.11",
   "timm==1.0.16",
   "tokenspeed_mla==0.1.8",
-  "torch==2.11.0",
+  "torch==2.13.0",
   "torch_memory_saver>=0.0.9.post1",
   "torchao==0.17.0",
   "torchaudio==2.11.0",
-  "torchcodec==0.11.1 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # torchcodec 0.11.1 for torch 2.11.x (0.10 is ABI-incompatible: references the pre-2.11 c10::MessageLogger ctor signature). Not available on Linux ARM.
+  "torchcodec==0.15.0 ; sys_platform != 'linux' or (sys_platform == 'linux' and platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')", # Not available on Linux ARM.
   "torchvision",
   "tqdm",
   "transformers==5.12.1",

python/sglang/multimodal_gen/runtime/launch_server.py

@@ -766,6 +766,9 @@ def launch_disagg_role(server_args: ServerArgs):
 
 def dispatch_launch(server_args: ServerArgs):
     """Route to the correct launch function based on --disagg-role."""
+    if "NCCL_NVLS_ENABLE" not in os.environ or server_args.enable_nccl_nvls:
+        os.environ["NCCL_NVLS_ENABLE"] = str(int(server_args.enable_nccl_nvls))
+
     role = server_args.disagg_role
     if role == RoleType.MONOLITHIC:
         launch_server(server_args)

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -321,7 +321,11 @@ class RMSNormNoWeight(CustomOp):
         return F.rms_norm(x, normalized_shape=(x.shape[-1],), eps=eps)
 
     def forward_cuda(self, x: torch.Tensor, eps: float) -> torch.Tensor:
-        return self.forward_native(x, eps=eps)
+        # Torch 2.12+ runs rms_norm in fp32 under CUDA autocast. This operator
+        # historically preserved the activation dtype, and callers rely on
+        # that contract for both memory use and downstream kernel selection.
+        with torch.autocast(device_type="cuda", enabled=False):
+            return self.forward_native(x, eps=eps)
 
     def forward_npu(self, x: torch.Tensor, eps: float) -> torch.Tensor:
         return fused_rmsnorm_without_weight(x, eps)

python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py

@@ -847,8 +847,17 @@ class LTX2Attention(nn.Module):
             else:
                 if self.qk_norm:
                     assert self.q_norm is not None and self.k_norm is not None
-                    q = self.q_norm(q)
-                    k = self.k_norm(k)
+                    # Torch 2.12+ places rms_norm on the autocast fp32 list. A
+                    # cast after the norm preserves the attention contract but
+                    # still runs the much slower fp32 kernel. Torch 2.11 ran
+                    # this operation in the input dtype, so disable autocast
+                    # around Q/K norm to preserve both its precision path and
+                    # performance.
+                    q_dtype = q.dtype
+                    k_dtype = k.dtype
+                    with torch.autocast(device_type=q.device.type, enabled=False):
+                        q = self.q_norm(q).to(dtype=q_dtype)
+                        k = self.k_norm(k).to(dtype=k_dtype)
 
                 if pe is not None and cos.dim() == 3:
                     q = apply_interleaved_rotary_emb(q, (cos, sin))

python/sglang/multimodal_gen/runtime/server_args/server_args.py

@@ -203,6 +203,7 @@ class ServerArgs(DisaggServerArgsMixin):
 
     # Distributed executor backend
     nccl_port: Optional[int] = None
+    enable_nccl_nvls: bool = False
 
     # HuggingFace specific parameters
     trust_remote_code: bool = False
@@ -1382,6 +1383,12 @@ class ServerArgs(DisaggServerArgsMixin):
             ),
         )
         # Parallelism
+        parser.add_argument(
+            "--enable-nccl-nvls",
+            action=StoreBoolean,
+            default=ServerArgs.enable_nccl_nvls,
+            help="Enable NCCL NVLS when available.",
+        )
         parser.add_argument(
             "--num-gpus",
             type=int,

python/sglang/multimodal_gen/test/test_utils.py

@@ -39,7 +39,7 @@ logger = init_logger(__name__)
 # NPU/ascend) is read from sgl-project/ci-data-diffusion, where the GT-gen workflows
 # publish.
 SGL_TEST_FILES_CI_DATA_REPO = "sgl-project/ci-data-diffusion"
-SGL_TEST_FILES_CI_DATA_REVISION = "d05810e3ea3eff1d137dec723f6e66d9c11b470f"
+SGL_TEST_FILES_CI_DATA_REVISION = "dc0e1bb34f2776313a259bcfab3e30daed85160e"
 
 # The NPU pin is kept as a separate branch so ascend GT can be bumped independently
 # when it's regenerated on its own cadence.

python/sglang/multimodal_gen/test/unit/test_server_args.py

@@ -2114,5 +2114,23 @@ class TestDisaggTransferBackendArgs(unittest.TestCase):
         self.assertEqual(args.disagg_transfer_backend, "mock")
 
 
+class TestNcclNvlsArgs(unittest.TestCase):
+    def test_enable_nccl_nvls_cli_arg(self):
+        parser = FlexibleArgumentParser()
+        ServerArgs.add_cli_args(parser)
+
+        default_args, _ = parser.parse_known_args(["--model-path", "/fake"])
+        enabled_args, _ = parser.parse_known_args(
+            ["--model-path", "/fake", "--enable-nccl-nvls"]
+        )
+        disabled_args, _ = parser.parse_known_args(
+            ["--model-path", "/fake", "--enable-nccl-nvls", "false"]
+        )
+
+        self.assertFalse(default_args.enable_nccl_nvls)
+        self.assertTrue(enabled_args.enable_nccl_nvls)
+        self.assertFalse(disabled_args.enable_nccl_nvls)
+
+
 if __name__ == "__main__":
     unittest.main()

python/sglang/srt/entrypoints/engine.py

@@ -1651,7 +1651,7 @@ def _set_envs_and_config(server_args: ServerArgs):
         if _is_cuda:
             assert_pkg_version(
                 "sglang-kernel",
-                "0.4.5",
+                "0.4.6.post1",
                 "Please reinstall the latest version with `pip install sglang-kernel --force-reinstall`",
             )
 

python/sglang/srt/layers/moe/fused_moe_triton/triton_kernels_moe.py

@@ -7,17 +7,14 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Optional
 
 import torch
-from triton_kernels.matmul_ogs import (
+from triton_kernels.matmul import (
     FlexCtx,
     FnSpecs,
     FusedActivation,
-    GatherIndx,
     PrecisionConfig,
-    RoutingData,
-    ScatterIndx,
-    matmul_ogs,
+    matmul,
 )
-from triton_kernels.matmul_ogs_details.opt_flags import update_opt_flags_constraints
+from triton_kernels.matmul_details.opt_flags import update_opt_flags_constraints
 from triton_kernels.numerics import InFlexData
 from triton_kernels.swiglu import swiglu_fn
 from triton_kernels.tensor import FP4
@@ -35,6 +32,8 @@ else:
     from sgl_kernel import gelu_and_mul, silu_and_mul
 
 if TYPE_CHECKING:
+    from triton_kernels.tensor_details.ragged_tensor import RaggedTensorMetadata
+
     from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
     from sglang.srt.layers.moe.topk import TopKOutput
 
@@ -86,15 +85,17 @@ def triton_kernel_moe_forward(
 
     assert TopKOutputChecker.format_is_triton_kernels(topk_output)
 
-    routing_data, gather_idx, scatter_idx = topk_output
+    a_ragged_metadata, gather_idx, scatter_idx, gate_scal, n_expts_act = topk_output
 
     return triton_kernel_fused_experts(
         hidden_states,
         w1,
         w2,
-        routing_data,
+        a_ragged_metadata,
         gather_idx,
         scatter_idx,
+        gate_scal,
+        n_expts_act,
         inplace=False,  # triton kernel doesn't support inplace
         activation=moe_runner_config.activation,
         apply_router_weight_on_input=apply_router_weight_on_input,
@@ -115,9 +116,11 @@ def triton_kernel_fused_experts(
     hidden_states: torch.Tensor,
     w1: torch.Tensor,
     w2: torch.Tensor,
-    routing_data: RoutingData,
-    gather_indx: GatherIndx,
-    scatter_indx: ScatterIndx,
+    a_ragged_metadata: RaggedTensorMetadata,
+    gather_indx: torch.Tensor,
+    scatter_indx: Optional[torch.Tensor],
+    gate_scal: torch.Tensor,
+    n_expts_act: int,
     inplace: bool = False,
     activation: str = "silu",
     apply_router_weight_on_input: bool = False,
@@ -162,7 +165,6 @@ def triton_kernel_fused_experts(
 
     M, K = hidden_states.shape
     E, _, N = w1.shape
-    n_expts_act = routing_data.n_expts_act
     dtype = hidden_states.dtype
 
     if global_num_experts == -1:
@@ -170,16 +172,16 @@ def triton_kernel_fused_experts(
 
     # consistent with default implementation
     intermediate_cache2 = torch.empty(
-        (M * n_expts_act, N // 2), device="cuda", dtype=dtype
+        (M * n_expts_act, N // 2), device=hidden_states.device, dtype=dtype
     )
 
-    intermediate_cache1 = matmul_ogs(
+    intermediate_cache1 = matmul(
         hidden_states,
         w1,
         None,
-        routing_data,
+        a_ragged_metadata=a_ragged_metadata,
         gather_indx=gather_indx,
-        gammas=routing_data.gate_scal if apply_router_weight_on_input else None,
+        gammas=gate_scal if apply_router_weight_on_input else None,
     )
 
     if activation == "silu":
@@ -189,13 +191,13 @@ def triton_kernel_fused_experts(
     else:
         raise ValueError(f"Unsupported FusedMoe activation: {activation}")
 
-    intermediate_cache3 = matmul_ogs(
+    intermediate_cache3 = matmul(
         intermediate_cache2,
         w2,
         None,
-        routing_data,
+        a_ragged_metadata=a_ragged_metadata,
         scatter_indx=scatter_indx,
-        gammas=None if apply_router_weight_on_input else routing_data.gate_scal,
+        gammas=None if apply_router_weight_on_input else gate_scal,
     )
 
     return intermediate_cache3
@@ -226,7 +228,7 @@ def triton_kernel_moe_with_bias_forward(
 
     assert TopKOutputChecker.format_is_triton_kernels(topk_output)
 
-    routing_data, gather_idx, scatter_idx = topk_output
+    a_ragged_metadata, gather_idx, scatter_idx, gate_scal, n_expts_act = topk_output
 
     return triton_kernel_fused_experts_with_bias(
         hidden_states,
@@ -236,9 +238,11 @@ def triton_kernel_moe_with_bias_forward(
         w2=w2,
         w2_pcg=w2_pcg,
         b2=b2,
-        routing_data=routing_data,
+        a_ragged_metadata=a_ragged_metadata,
         gather_indx=gather_idx,
         scatter_indx=scatter_idx,
+        gate_scal=gate_scal,
+        n_expts_act=n_expts_act,
         inplace=False,  # triton kernel doesn't support inplace
         activation=moe_runner_config.activation,
         apply_router_weight_on_input=apply_router_weight_on_input,
@@ -264,9 +268,11 @@ def triton_kernel_fused_experts_with_bias(
     w2: torch.Tensor,
     w2_pcg,
     b2: torch.Tensor,
-    routing_data: RoutingData,
-    gather_indx: GatherIndx,
-    scatter_indx: ScatterIndx,
+    a_ragged_metadata: RaggedTensorMetadata,
+    gather_indx: torch.Tensor,
+    scatter_indx: Optional[torch.Tensor],
+    gate_scal: torch.Tensor,
+    n_expts_act: int,
     inplace: bool = False,
     activation: str = "silu",
     apply_router_weight_on_input: bool = False,
@@ -315,7 +321,6 @@ def triton_kernel_fused_experts_with_bias(
 
     M, K = hidden_states.shape
     E, _, N = w1.shape
-    n_expts_act = routing_data.n_expts_act
 
     if global_num_experts == -1:
         global_num_experts = E
@@ -335,35 +340,24 @@ def triton_kernel_fused_experts_with_bias(
         (gemm1_alpha, gemm1_clamp_limit),
     )
 
-    intermediate_cache = torch.empty(
-        (1, M * n_expts_act, N // 2),
-        device=hidden_states.device,
-        dtype=hidden_states.dtype,
-    )
-    output = torch.empty(
-        (1, M, K), device=hidden_states.device, dtype=hidden_states.dtype
-    )
-
-    matmul_ogs(
+    intermediate_cache = matmul(
         hidden_states,
         w1,
         b1,
-        routing_data,
+        a_ragged_metadata=a_ragged_metadata,
         gather_indx=gather_indx,
         precision_config=w1_pcg,
-        gammas=routing_data.gate_scal if apply_router_weight_on_input else None,
+        gammas=gate_scal if apply_router_weight_on_input else None,
         fused_activation=act,
-        y=intermediate_cache,
     )
 
-    matmul_ogs(
+    output = matmul(
         intermediate_cache.view(M * n_expts_act, N // 2),
         w2,
         b2,
-        routing_data,
+        a_ragged_metadata=a_ragged_metadata,
         scatter_indx=scatter_indx,
         precision_config=w2_pcg,
-        gammas=None if apply_router_weight_on_input else routing_data.gate_scal,
-        y=output,
+        gammas=None if apply_router_weight_on_input else gate_scal,
     )
-    return output.view(M, K)
+    return output.view(-1, K)

python/sglang/srt/layers/moe/moe_runner/triton_kernels.py

@@ -19,12 +19,8 @@ from sglang.srt.layers.moe.moe_runner.base import (
 from sglang.srt.layers.moe.utils import MoeRunnerBackend
 
 if TYPE_CHECKING:
-    from triton_kernels.matmul_ogs import (
-        GatherIndx,
-        PrecisionConfig,
-        RoutingData,
-        ScatterIndx,
-    )
+    from triton_kernels.matmul import PrecisionConfig
+    from triton_kernels.tensor_details.ragged_tensor import RaggedTensorMetadata
 
     from sglang.srt.layers.moe.token_dispatcher.standard import (
         StandardCombineInput,
@@ -42,9 +38,11 @@ class TritonKernelsRunnerInput(RunnerInput):
     """Input bundle passed to the triton-kernels runner core."""
 
     hidden_states: torch.Tensor
-    routing_data: RoutingData
-    gather_indx: GatherIndx
-    scatter_indx: ScatterIndx
+    a_ragged_metadata: RaggedTensorMetadata
+    gather_indx: torch.Tensor
+    scatter_indx: torch.Tensor
+    gate_scal: torch.Tensor
+    n_expts_act: int
 
     @property
     def runner_backend(self) -> MoeRunnerBackend:
@@ -102,9 +100,11 @@ class TritonKernelsRunnerCore(MoeRunnerCore):
         hidden_states = runner_input.hidden_states
 
         common_kwargs = dict(
-            routing_data=runner_input.routing_data,
+            a_ragged_metadata=runner_input.a_ragged_metadata,
             gather_indx=runner_input.gather_indx,
             scatter_indx=None if self.config.no_combine else runner_input.scatter_indx,
+            gate_scal=runner_input.gate_scal,
+            n_expts_act=runner_input.n_expts_act,
             inplace=False,
             activation=self.config.activation,
             apply_router_weight_on_input=self.config.apply_router_weight_on_input,
@@ -137,12 +137,14 @@ class TritonKernelsRunnerCore(MoeRunnerCore):
                 **common_kwargs,
             )
 
+        tokens = runner_input.hidden_states.shape[0]
+        hidden = runner_input.hidden_states.shape[-1]
+        top_k = runner_input.n_expts_act
+
         if self.config.no_combine:
-            tokens = runner_input.hidden_states.shape[0]
-            hidden = runner_input.hidden_states.shape[-1]
-            total_rows = output.shape[0]
-            top_k = total_rows // tokens
             output = output.view(tokens, top_k, hidden)
+        else:
+            output = output.view(tokens, top_k, hidden).sum(dim=1)
 
         return TritonKernelsRunnerOutput(hidden_states=output)
 
@@ -172,13 +174,15 @@ def pre_permute_standard_to_triton_kernels(
         topk_output
     ), "Triton-kernel runner expects TritonKernelTopKOutput"
 
-    routing_data, gather_indx, scatter_indx = topk_output
+    a_ragged_metadata, gather_indx, scatter_indx, gate_scal, n_expts_act = topk_output
 
     return TritonKernelsRunnerInput(
         hidden_states=hidden_states,
-        routing_data=routing_data,
+        a_ragged_metadata=a_ragged_metadata,
         gather_indx=gather_indx,
         scatter_indx=scatter_indx,
+        gate_scal=gate_scal,
+        n_expts_act=n_expts_act,
     )
 
 

python/sglang/srt/layers/moe/topk.py

@@ -32,10 +32,12 @@ from typing import (
 import torch
 import torch.nn.functional as F
 
+if TYPE_CHECKING:
+    from triton_kernels.tensor_details.ragged_tensor import RaggedTensorMetadata
+
 from sglang.srt.runtime_context import get_exec, get_lora, get_parallel
 
 try:
-    from triton_kernels.matmul_ogs import GatherIndx, RoutingData, ScatterIndx
     from triton_kernels.tensor import make_ragged_tensor_metadata
     from triton_kernels.topk import topk as triton_kernels_topk
 
@@ -49,7 +51,7 @@ try:
     ):
         if simulated_ep != 1:
             raise NotImplementedError(
-                "simulated_ep routing is not supported with triton_kernels 3.6.0"
+                "simulated_ep routing is not supported with triton_kernels 3.7.1"
             )
 
         if sm_first:
@@ -64,20 +66,13 @@ try:
         )
         dispatch_indx = sparse_logits.mask_metadata.row_sorted_indx
         combine_indx = sparse_logits.mask_metadata.col_sorted_indx
+        gather_indx = torch.div(combine_indx, n_expts_act, rounding_mode="trunc")
+        scatter_indx = combine_indx
         ragged_metadata = make_ragged_tensor_metadata(
             sparse_logits.mask_metadata.col_sum, dispatch_indx.shape[0]
         )
         gate_scal = sparse_logits.vals.flatten()[combine_indx]
-        routing_data = RoutingData(
-            gate_scal,
-            ragged_metadata.slice_sizes,
-            logits.shape[-1],
-            n_expts_act,
-            ragged_metadata,
-        )
-        gather_indx = GatherIndx(combine_indx, dispatch_indx)
-        scatter_indx = ScatterIndx(dispatch_indx, combine_indx)
-        return routing_data, gather_indx, scatter_indx
+        return ragged_metadata, gather_indx, scatter_indx, gate_scal, n_expts_act
 
 except ImportError:
     pass
@@ -319,9 +314,11 @@ class StandardTopKOutputPacked(NamedTuple):
 class TritonKernelTopKOutput(NamedTuple):
     """Triton kernel top-k output format."""
 
-    routing_data: RoutingData
-    gather_indx: GatherIndx
-    scatter_indx: ScatterIndx
+    a_ragged_metadata: RaggedTensorMetadata
+    gather_indx: torch.Tensor
+    scatter_indx: torch.Tensor
+    gate_scal: torch.Tensor
+    n_expts_act: int
 
     @property
     def format(self) -> TopKOutputFormat:
@@ -537,12 +534,24 @@ class TopK(BaseFusedOp):
 
         if output_format == TopKOutputFormat.TRITON_KERNEL:
             # renormalize=True is equivalent to sm_first=False
-            routing_data, gather_idx, scatter_idx = routing(
+            (
+                a_ragged_metadata,
+                gather_idx,
+                scatter_idx,
+                gate_scal,
+                n_expts_act,
+            ) = routing(
                 router_logits,
                 self.topk_config.top_k,
                 sm_first=not self.topk_config.renormalize,
             )
-            return TritonKernelTopKOutput(routing_data, gather_idx, scatter_idx)
+            return TritonKernelTopKOutput(
+                a_ragged_metadata,
+                gather_idx,
+                scatter_idx,
+                gate_scal,
+                n_expts_act,
+            )
         elif output_format == TopKOutputFormat.BYPASSED:
             return BypassedTopKOutput(
                 hidden_states=hidden_states,

python/sglang/srt/layers/quantization/fp8.py

@@ -263,7 +263,7 @@ class Fp8Config(QuantizationConfig):
         if weight_block_size is not None:
             if not is_checkpoint_fp8_serialized:
                 raise ValueError(
-                    f"The block-wise quantization only supports fp8-serialized checkpoint for now."
+                    "The block-wise quantization only supports fp8-serialized checkpoint for now."
                 )
             if len(weight_block_size) != 2:
                 raise ValueError(
@@ -1788,11 +1788,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
             from triton_kernels.tensor import convert_layout, wrap_torch_tensor
             from triton_kernels.tensor_details import layout
 
-            scale_layout, scale_layout_opts = (
-                layout.make_default_matmul_mxfp4_w_scale_layout(
-                    mx_axis=1, num_warps=num_warps
-                )
+            scale_layout = layout.make_default_matmul_mxfp4_w_scale_layout(
+                mx_axis=-2, num_warps=num_warps
             )
+            scale_layout_opts = {}
             scale = scale.transpose(-2, -1)
             scale = convert_layout(
                 wrap_torch_tensor(scale), scale_layout, **scale_layout_opts

python/sglang/srt/layers/quantization/mxfp4.py

@@ -169,17 +169,17 @@ if _is_hip:
 
 def _swizzle_mxfp4(quant_tensor, scale, num_warps):
     """weight swizzle for mxfp4 moe, used for OAI mxfp4 kernel"""
-    import triton_kernels.matmul_ogs_details.opt_flags as opt_flags
+    import triton_kernels.matmul_details.opt_flags as opt_flags
     from triton_kernels.numerics import InFlexData
     from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
     from triton_kernels.tensor_details import layout
 
-    value_layout, value_layout_opts = layout.make_default_matmul_mxfp4_w_layout(
-        mx_axis=1
-    )
-    scale_layout, scale_layout_opts = layout.make_default_matmul_mxfp4_w_scale_layout(
-        mx_axis=1, num_warps=num_warps
+    value_layout = layout.make_default_matmul_mxfp4_w_layout(mx_axis=-2)
+    value_layout_opts = {}
+    scale_layout = layout.make_default_matmul_mxfp4_w_scale_layout(
+        mx_axis=-2, num_warps=num_warps
     )
+    scale_layout_opts = {}
     if is_sm100_supported():
         constraints = {
             "is_persistent": True,
@@ -931,7 +931,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
 
         if self.use_triton_kernels:
 
-            from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
+            from triton_kernels.matmul import FlexCtx, PrecisionConfig
 
             w13_weight_bias = layer.w13_weight_bias.to(torch.float32)
             w2_weight_bias = layer.w2_weight_bias.to(torch.float32)
@@ -949,10 +949,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             )
 
             self.w13_precision_config = PrecisionConfig(
-                weight_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
+                b_mx_scale=w13_scale, flex_ctx=FlexCtx(rhs_data=w13_flex)
             )
             self.w2_precision_config = PrecisionConfig(
-                weight_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
+                b_mx_scale=w2_scale, flex_ctx=FlexCtx(rhs_data=w2_flex)
             )
 
             self.w13_weight_triton_tensor = w13_weight

python/sglang/srt/models/deepseek_v2.py

@@ -191,7 +191,6 @@ from sglang.srt.models.deepseek_common.utils import (
 from sglang.srt.runtime_context import (
     get_device,
     get_exec,
-    get_flags,
     get_forward,
     get_model,
     get_parallel,
@@ -910,12 +909,6 @@ class DeepseekV2MoE(nn.Module):
                 and self.num_fused_shared_experts == 0
                 and hidden_states.shape[0] > 0
                 and get_is_capture_mode()
-                and not (
-                    get_flags().capture.enable_torch_compile
-                    and hidden_states.shape[0]
-                    <= get_exec().graph.torch_compile_max_bs
-                    * (get_spec().speculative_num_draft_tokens or 1)
-                )
             ):
                 return self.forward_normal_dual_stream(
                     hidden_states,

python/sglang/srt/models/nemotron_h.py

@@ -279,17 +279,11 @@ class NemotronHMoE(nn.Module):
         self,
         hidden_states: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        overlap = _is_cuda and not torch.compiler.is_compiling()
-        if (
-            overlap
-            and get_moe_a2a_backend().is_flashinfer()
-            and not get_is_capture_mode()
+        if _is_cuda and (
+            not get_moe_a2a_backend().is_flashinfer() or get_is_capture_mode()
         ):
-            overlap = False
-        if overlap:
             return self._forward_core_shared_routed_overlap(hidden_states)
-        else:
-            return self._forward_core_normal(hidden_states)
+        return self._forward_core_normal(hidden_states)
 
     def _forward_core_normal(
         self,

python/sglang/srt/models/qwen2_moe.py

@@ -576,7 +576,11 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
             shared_output = None
             topk_output = self.topk.empty_topk_output(hidden_states.device)
             final_hidden_states = self.experts(hidden_states, topk_output)
-        elif self.alt_stream is not None and get_is_capture_mode():
+        elif (
+            self.alt_stream is not None
+            and get_is_capture_mode()
+            and not torch.compiler.is_compiling()
+        ):
             final_hidden_states, shared_output = self.forward_normal_dual_stream(
                 hidden_states, use_fused_gate=use_fused_gate
             )

python/sglang/srt/utils/common.py

@@ -4259,15 +4259,7 @@ class ConcurrentCounter:
 
 @lru_cache(maxsize=1)
 def is_triton_kernels_available() -> bool:
-    if importlib.util.find_spec("triton_kernels") is None:
-        return False
-    try:
-        ragged_metadata_spec = importlib.util.find_spec(
-            "triton_kernels.tensor_details.ragged_tensor"
-        )
-    except ModuleNotFoundError:
-        return False
-    return ragged_metadata_spec is not None
+    return importlib.util.find_spec("triton_kernels") is not None
 
 
 def json_list_type(value):