Nemotron perf changes (#26733)

Co-authored-by: Brayden Zhong <brayden@radixark.ai>

python/sglang/jit_kernel/activation.py

@@ -42,11 +42,16 @@ def _jit_activation_module(dtype: torch.dtype) -> Module:
                 "run_activation_filtered",
                 f"ActivationKernel<{args}>::run_activation_filtered",
             ),
+            (
+                "run_unary_activation",
+                f"ActivationKernel<{args}>::run_unary_activation",
+            ),
         ],
     )
 
 
 SUPPORTED_ACTIVATIONS = {"silu", "gelu", "gelu_tanh"}
+SUPPORTED_UNARY_ACTIVATIONS = {"relu2"}
 
 
 @register_custom_op(mutates_args=["out"])
@@ -100,6 +105,42 @@ def run_activation(
     return out
 
 
+@register_custom_op(mutates_args=["out"])
+def _run_unary_activation_inplace(
+    op_name: str, input: torch.Tensor, out: torch.Tensor
+) -> None:
+    last = input.shape[-1]
+    module = _jit_activation_module(input.dtype)
+    module.run_unary_activation(input.view(-1, last), out.view(-1, last), op_name)
+
+
+def run_unary_activation(
+    op_name: str,
+    input: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Apply a standalone (non-gated) element-wise activation: ``out = act(input)``.
+
+    Unlike :func:`run_activation`, there is no gate/up split — ``input`` and
+    ``out`` share the same shape.
+    """
+    assert (
+        op_name in SUPPORTED_UNARY_ACTIVATIONS
+    ), f"Unsupported unary activation: {op_name}"
+    if out is None:
+        out = torch.empty_like(input)
+    _run_unary_activation_inplace(op_name, input, out)
+    return out
+
+
+def relu2(
+    input: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Squared ReLU: ``out = max(0, input) ** 2`` (element-wise)."""
+    return run_unary_activation("relu2", input, out)
+
+
 def silu_and_mul(
     input: torch.Tensor,
     out: Optional[torch.Tensor] = None,

python/sglang/jit_kernel/benchmark/bench_activation.py

@@ -6,6 +6,7 @@ from sgl_kernel import silu_and_mul as silu_and_mul_aot
 
 from sglang.jit_kernel.activation import gelu_and_mul as gelu_and_mul_jit
 from sglang.jit_kernel.activation import gelu_tanh_and_mul as gelu_tanh_and_mul_jit
+from sglang.jit_kernel.activation import relu2 as relu2_jit
 from sglang.jit_kernel.activation import silu_and_mul as silu_and_mul_jit
 from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.benchmark.utils import create_random
@@ -87,6 +88,21 @@ def benchmark_filter(
     )
 
 
+@torch.compile
+def relu2_torch(input: torch.Tensor) -> torch.Tensor:
+    return F.relu(input).pow(2)
+
+
+@marker.parametrize("dim", [1024, 4096, 6144, 8192], [4096])
+@marker.parametrize("batch_size", [2**x for x in range(0, 15)], [8, 512])
+@marker.benchmark("impl", ["jit", "torch"])
+def benchmark_unary(dim: int, batch_size: int, impl: str):
+    x = create_random(batch_size, dim)
+    fn = {"jit": relu2_jit, "torch": relu2_torch}[impl]
+    return marker.do_bench(fn, input_args=(x,))
+
+
 if __name__ == "__main__":
     benchmark.run()
     benchmark_filter.run()
+    benchmark_unary.run()

python/sglang/jit_kernel/csrc/elementwise/activation.cuh

@@ -19,6 +19,7 @@ enum class ActivationKind : uint32_t {
   kSiLU,
   kGELU,
   kGELUTanh,
+  kReLU2,
 };
 
 template <ActivationKind kAct>
@@ -33,6 +34,9 @@ SGL_DEVICE float apply_activation_f32(float x_f32) {
     constexpr auto kGeluTanhBeta = 0.7978845608028654f;
     const float cdf = 0.5f * (1.0f + tanhf(kGeluTanhBeta * (x_f32 + kGeluTanhAlpha * x_f32 * x_f32 * x_f32)));
     return x_f32 * cdf;
+  } else if constexpr (kAct == ActivationKind::kReLU2) {
+    const float relu = x_f32 > 0.0f ? x_f32 : 0.0f;
+    return relu * relu;
   } else {
     static_assert(host::dependent_false_v<decltype(kAct)>, "unsupported activation kind");
     return 0.0f;
@@ -81,6 +85,30 @@ __global__ void act_and_mul_kernel(const __grid_constant__ ActivationParams para
   PDLTriggerSecondary<kUsePDL>();
 }
 
+struct UnaryActivationParams {
+  const void* __restrict__ input;
+  void* __restrict__ out;
+  uint32_t num_vecs;
+};
+
+template <typename T, ActivationKind kAct, bool kUsePDL>
+__global__ void act_kernel(const __grid_constant__ UnaryActivationParams params) {
+  using namespace device;
+  constexpr auto kVecSize = kMaxVecBytes / sizeof(T);
+  using vec_t = AlignedVector<T, kMaxVecBytes / sizeof(T)>;
+  const auto vec_id = blockIdx.x * blockDim.x + threadIdx.x;
+  if (vec_id >= params.num_vecs) return;
+  PDLWaitPrimary<kUsePDL>();
+  const auto in = device::load_as<vec_t>(params.input, vec_id);
+  vec_t out;
+#pragma unroll
+  for (int i = 0; i < kVecSize; ++i) {
+    out[i] = device::cast<T>(apply_activation_f32<kAct>(device::cast<fp32_t>(in[i])));
+  }
+  device::store_as<vec_t>(params.out, out, vec_id);
+  PDLTriggerSecondary<kUsePDL>();
+}
+
 template <typename T, bool kUsePDL>
 struct ActivationKernel {
   static constexpr auto kVecSize = device::kMaxVecBytes / sizeof(T);
@@ -174,6 +202,50 @@ struct ActivationKernel {
     RuntimeCheck(expert_step >= 1, "expert_step must be positive");
     launch(input, out, type, static_cast<const int32_t*>(expert_ids.data_ptr()), static_cast<uint32_t>(expert_step));
   }
+
+  template <ActivationKind kAct>
+  static constexpr auto unary_kernel = act_kernel<T, kAct, kUsePDL>;
+
+  static auto select_unary_kernel(const std::string& type)
+      -> decltype(ActivationKernel::template unary_kernel<ActivationKind::kReLU2>) {
+    using namespace host;
+    if (type == "relu2") {
+      return ActivationKernel::template unary_kernel<ActivationKind::kReLU2>;
+    } else {
+      Panic("unsupported unary activation type: ", type);
+    }
+    return nullptr;
+  }
+
+  static void run_unary_activation(const tvm::ffi::TensorView input, const tvm::ffi::TensorView out, std::string type) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto D = SymbolicSize{"hidden"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({N, D})  //
+        .with_dtype<T>()
+        .with_device(device_)
+        .verify(out)
+        .verify(input);
+
+    const auto num_elems = static_cast<int64_t>(N.unwrap()) * D.unwrap();
+    const auto device = device_.unwrap();
+    if (num_elems == 0) return;
+    RuntimeCheck(num_elems % kVecSize == 0, "num elements must be divisible by vector size");
+    const auto num_vecs = num_elems / kVecSize;
+    RuntimeCheck(num_vecs <= std::numeric_limits<uint32_t>::max(), "too many items for 32-bit indexing");
+    const auto num_blocks = div_ceil(static_cast<uint32_t>(num_vecs), kBlockSize);
+    const auto params = UnaryActivationParams{
+        .input = input.data_ptr(),
+        .out = out.data_ptr(),
+        .num_vecs = static_cast<uint32_t>(num_vecs),
+    };
+    const auto kernel = select_unary_kernel(type);
+    LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
+  }
 };
 
 }  // namespace

python/sglang/jit_kernel/tests/test_activation.py

@@ -4,7 +4,11 @@ import pytest
 import torch
 import torch.nn.functional as F
 
-from sglang.jit_kernel.activation import SUPPORTED_ACTIVATIONS, run_activation
+from sglang.jit_kernel.activation import (
+    SUPPORTED_ACTIVATIONS,
+    relu2,
+    run_activation,
+)
 from sglang.jit_kernel.utils import get_ci_test_range
 from sglang.test.ci.ci_register import register_cuda_ci
 
@@ -155,5 +159,47 @@ def test_activation_filter_expert_none_skipped(op_name: str) -> None:
     torch.testing.assert_close(out_filtered, out_unfiltered, atol=0.0, rtol=0.0)
 
 
+UNARY_SHAPES = get_ci_test_range(
+    full_range=[
+        (7, 16),
+        (83, 1024),
+        (3, 5, 16),
+        (2, 3, 512),
+        (1, 17, 4096),
+        *[(2**x, 2048) for x in range(0, 15, 2)],
+    ],
+    ci_range=[(7, 16), (2, 3, 512)],
+)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("shape", UNARY_SHAPES)
+def test_relu2_correctness(dtype: torch.dtype, shape: tuple[int, ...]) -> None:
+    x = torch.randn(shape, dtype=dtype, device="cuda")
+    out = relu2(x)
+    expected = F.relu(x.float()).pow(2).to(dtype=dtype)
+    atol, rtol = _tolerances(dtype)
+    torch.testing.assert_close(out, expected, atol=atol, rtol=rtol)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("shape", UNARY_SHAPES)
+def test_relu2_out_param(dtype: torch.dtype, shape: tuple[int, ...]) -> None:
+    x = torch.randn(shape, dtype=dtype, device="cuda")
+    out = torch.empty(shape, dtype=dtype, device="cuda")
+    result = relu2(x, out)
+    assert result is out
+    expected = F.relu(x.float()).pow(2).to(dtype=dtype)
+    atol, rtol = _tolerances(dtype)
+    torch.testing.assert_close(out, expected, atol=atol, rtol=rtol)
+
+
+def test_relu2_negative_inputs_zeroed() -> None:
+    """All-negative input must produce an all-zero output."""
+    x = -torch.rand((64, 512), dtype=torch.bfloat16, device="cuda") - 1e-3
+    out = relu2(x)
+    assert torch.count_nonzero(out) == 0
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v", "-s"]))

python/sglang/srt/arg_groups/nemotron_h_hook.py

@@ -12,12 +12,13 @@ logger = logging.getLogger(__name__)
 def apply_nemotron_h_defaults(server_args: "ServerArgs", model_arch: str) -> None:
     """Apply NemotronH model-specific server arg defaults and constraints."""
     model_config = server_args.get_model_config()
-    if model_config.quantization in [
+    is_modelopt = model_config.quantization in [
         "modelopt",
         "modelopt_fp8",
         "modelopt_fp4",
         "modelopt_mixed",
-    ]:
+    ]
+    if is_modelopt:
         assert model_config.hf_config.mlp_hidden_act == "relu2"
         if model_config.quantization == "modelopt":
             quant_algo = model_config.hf_config.quantization_config["quant_algo"]
@@ -29,28 +30,30 @@ def apply_nemotron_h_defaults(server_args: "ServerArgs", model_arch: str) -> Non
                 )
         else:
             server_args.quantization = model_config.quantization
-        if server_args.moe_runner_backend == "auto":
-            if is_sm100_supported() and server_args.moe_a2a_backend == "none":
-                server_args.moe_runner_backend = "flashinfer_trtllm"
-                logger.info(
-                    "Use flashinfer_trtllm as MoE runner backend on sm100 for "
-                    f"{model_arch}"
-                )
-            elif (
-                (
-                    model_config.quantization in ("modelopt_fp4", "modelopt_mixed")
-                    or server_args.quantization == "modelopt_fp4"
-                )
-                and is_cuda()
-                and (8, 0) <= get_device_capability() < (10, 0)
-            ):
-                server_args.moe_runner_backend = "marlin"
-                logger.info(
-                    "Use marlin as MoE runner backend on SM80-SM90 for "
-                    f"{model_arch} {model_config.quantization}"
-                )
-            else:
-                server_args.moe_runner_backend = "flashinfer_cutlass"
+
+    if (is_modelopt or model_config.quantization is None) and (
+        server_args.moe_runner_backend == "auto"
+    ):
+        if is_sm100_supported() and server_args.moe_a2a_backend == "none":
+            server_args.moe_runner_backend = "flashinfer_trtllm"
+            logger.info(
+                f"Use flashinfer_trtllm as MoE runner backend on sm100 for {model_arch}"
+            )
+        elif (
+            (
+                model_config.quantization in ("modelopt_fp4", "modelopt_mixed")
+                or server_args.quantization == "modelopt_fp4"
+            )
+            and is_cuda()
+            and (8, 0) <= get_device_capability() < (10, 0)
+        ):
+            server_args.moe_runner_backend = "marlin"
+            logger.info(
+                "Use marlin as MoE runner backend on SM80-SM90 for "
+                f"{model_arch} {model_config.quantization}"
+            )
+        else:
+            server_args.moe_runner_backend = "flashinfer_cutlass"
 
     server_args._handle_mamba_radix_cache(
         model_arch=model_arch,

python/sglang/srt/layers/activation.py

@@ -57,6 +57,7 @@ if _is_cuda:
     from sglang.jit_kernel.activation import (
         gelu_and_mul,
         gelu_tanh_and_mul,
+        relu2,
         silu_and_mul,
     )
 elif _is_xpu:
@@ -190,16 +191,19 @@ class NewGELU(MultiPlatformOp):
         return self.forward_native(x)
 
 
-class ReLU2(nn.Module):
+class ReLU2(MultiPlatformOp):
     """
     Applies the squared Rectified Linear Unit function.
     y = max(0, x)^2
     """
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         x = F.relu(x)
         return x * x
 
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        return relu2(x)
+
 
 class QuickGELU(MultiPlatformOp):
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:

python/sglang/srt/layers/attention/fla/layernorm_gated.py

@@ -14,6 +14,7 @@ import triton
 import triton.language as tl
 from einops import rearrange
 
+from sglang.jit_kernel.utils import is_arch_support_pdl
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     cdiv,
@@ -86,7 +87,11 @@ def _layer_norm_fwd_1pass_kernel(
     NORM_BEFORE_GATE: tl.constexpr,
     IS_RMS_NORM: tl.constexpr,
     ACTIVATION: tl.constexpr,
+    USE_GDC: tl.constexpr = False,
 ):
+    if USE_GDC:
+        tl.extra.cuda.gdc_wait()
+
     # Map the program id to the starting row of X and Y it should compute.
     row_start = tl.program_id(0) * ROWS_PER_BLOCK
     group = tl.program_id(1)
@@ -168,6 +173,9 @@ def _layer_norm_fwd_1pass_kernel(
     # Write output
     tl.store(Y_base, y, mask=mask)
 
+    if USE_GDC:
+        tl.extra.cuda.gdc_launch_dependents()
+
 
 @lru_cache
 def _get_sm_count(device: torch.device) -> int:
@@ -243,6 +251,7 @@ def _layer_norm_fwd(
     rows_per_block = calc_rows_per_block(M, x.device)
     # Update grid to use rows_per_block
     grid = (cdiv(M, rows_per_block), ngroups)
+    pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
     with device_context(x.device):
         _layer_norm_fwd_1pass_kernel[grid](
             x,
@@ -266,6 +275,7 @@ def _layer_norm_fwd(
             IS_RMS_NORM=is_rms_norm,
             num_warps=num_warps,
             ACTIVATION=activation,
+            **pdl_kwargs,
         )
     return out, mean, rstd
 

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

@@ -710,7 +710,7 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
         self,
         mixer: MambaMixer2,
         hidden_states: torch.Tensor,
-        output: torch.Tensor,
+        output: Optional[torch.Tensor],
         layer_id: int,
         forward_batch: ForwardBatch,
         mup_vector: Optional[torch.Tensor] = None,
@@ -718,7 +718,7 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
     ):
         assert isinstance(self.forward_metadata, Mamba2Metadata)
         layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer_id)
-        intermediate_states = mixer.forward(
+        mixer_out, intermediate_states = mixer.forward(
             hidden_states=hidden_states,
             output=output,
             layer_cache=layer_cache,
@@ -752,6 +752,8 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
                     num_decodes,
                 )
 
+        return mixer_out
+
     def forward_decode(self, *args, **kwargs):
         raise NotImplementedError(
             "Mamba2AttnBackend's forward is called directly instead of through HybridLinearAttnBackend, as it supports mixed prefill and decode"

python/sglang/srt/layers/attention/mamba/causal_conv1d_triton.py

@@ -10,6 +10,8 @@ import torch
 import triton
 import triton.language as tl
 
+from sglang.jit_kernel.utils import is_arch_support_pdl
+
 PAD_SLOT_ID = -1
 
 
@@ -629,8 +631,12 @@ def _causal_conv1d_update_kernel(
     BLOCK_N: tl.constexpr,
     SAVE_INTERMEDIATE: tl.constexpr,
     HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: tl.constexpr,
+    USE_GDC: tl.constexpr = False,
 ):
     # ruff: noqa: E501
+    if USE_GDC:
+        tl.extra.cuda.gdc_wait()
+
     idx_seq = tl.program_id(0)
     if idx_seq >= batch:
         return
@@ -978,6 +984,9 @@ def _causal_conv1d_update_kernel(
                 mask=mask_retrieve,
             )
 
+    if USE_GDC:
+        tl.extra.cuda.gdc_launch_dependents()
+
 
 def causal_conv1d_update(
     x: torch.Tensor,
@@ -1124,6 +1133,8 @@ def causal_conv1d_update(
     else:
         stride_retrieve_parent_token_seq = stride_retrieve_parent_token_token = 0
 
+    pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
+
     _causal_conv1d_update_kernel[grid](
         # Pointers to matrices
         x,
@@ -1183,6 +1194,7 @@ def causal_conv1d_update(
         BLOCK_N=256,
         SAVE_INTERMEDIATE=intermediate_conv_window is not None,
         HAS_EAGLE_TREE_CUSTOM_ATTN_MASK=retrieve_next_token is not None,
+        **pdl_kwargs,
     )
     if unsqueeze:
         out = out.squeeze(-1)

python/sglang/srt/layers/attention/mamba/mamba.py

@@ -408,13 +408,17 @@ class MambaMixer2(torch.nn.Module):
         self,
         *,
         hidden_states: torch.Tensor,
-        output: torch.Tensor,
+        output: Optional[torch.Tensor] = None,
         layer_cache: MambaPool.State,
         metadata: Mamba2Metadata,
         forward_batch: ForwardBatch,
         mup_vector: Optional[torch.Tensor] = None,
         use_triton_causal_conv: bool = False,
     ):
+        # Returns the projected result. When `output` is given it is also
+        # written into that buffer (required by the cuda-graph split ops, which
+        # need a stable buffer); otherwise the caller uses the return value and
+        # avoids a copy.
         # metadata contains metadata necessary for the mamba2 triton
         # kernels to operate in continuous batching and in chunked prefill
         # modes; they are computed at top-level model forward since they
@@ -718,9 +722,11 @@ class MambaMixer2(torch.nn.Module):
         hidden_states = self.norm(preallocated_ssm_out, gate[:num_actual_tokens])
 
         # 5. Final linear projection
-        output[:num_actual_tokens], _ = self.out_proj(hidden_states)
+        mixer_out, _ = self.out_proj(hidden_states)
+        if output is not None:
+            output[:num_actual_tokens].copy_(mixer_out)
 
-        return intermediate_states
+        return mixer_out, intermediate_states
 
     @property
     def mamba_type(self) -> str:

python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py

@@ -106,7 +106,7 @@ class Mixer2RMSNormGated(MultiPlatformOp):
             # Keep gate in float32 for numerical stability during silu
             return x * torch.nn.functional.silu(gate.to(torch.float32)).to(input_dtype)
 
-        if ((self.n_groups % self.tp_size) != 0) or self.n_groups != 1:
+        if (self.n_groups % self.tp_size) != 0:
             return self.forward_native(x, gate)
 
         return rms_norm_gated(
@@ -115,6 +115,7 @@ class Mixer2RMSNormGated(MultiPlatformOp):
             bias=None,
             z=gate,
             eps=self.variance_epsilon,
+            group_size=self.group_size,
             norm_before_gate=False,
             is_rms_norm=True,
         )

python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py

@@ -11,6 +11,8 @@ import triton
 import triton.language as tl
 from packaging import version
 
+from sglang.jit_kernel.utils import is_arch_support_pdl
+
 PAD_SLOT_ID = -1
 
 TRITON3 = version.parse(triton.__version__) >= version.parse("3.0.0")
@@ -141,7 +143,11 @@ def _selective_scan_update_kernel(
     HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: tl.constexpr,
     HAS_INTERMEDIATE_STATE_INDICES: tl.constexpr,
     BLOCK_SIZE_DSTATE: tl.constexpr,
+    USE_GDC: tl.constexpr = False,
 ):
+    if USE_GDC:
+        tl.extra.cuda.gdc_wait()
+
     pid_m = tl.program_id(axis=0)
     pid_b = tl.program_id(axis=1)
     pid_h = tl.program_id(axis=2)
@@ -296,6 +302,9 @@ def _selective_scan_update_kernel(
     if not DISABLE_STATE_UPDATE:
         tl.store(state_ptrs, state.to(state_ptrs.dtype.element_ty), mask=mask)
 
+    if USE_GDC:
+        tl.extra.cuda.gdc_launch_dependents()
+
 
 def selective_state_update(
     state,
@@ -427,6 +436,8 @@ def selective_state_update(
         else (0, 0)
     )
 
+    pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
+
     with torch.get_device_module(x.device).device(x.device.index):
         _selective_scan_update_kernel[grid](
             state,
@@ -491,4 +502,5 @@ def selective_state_update(
             BLOCK_SIZE_M,
             DISABLE_STATE_UPDATE=disable_state_update,
             num_warps=num_warps,
+            **pdl_kwargs,
         )

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

@@ -1145,6 +1145,12 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
                 activation_type=activation,
             )[0]
 
+            if (
+                not layer.should_fuse_routed_scaling_factor_in_topk
+                and self.moe_runner_config.routed_scaling_factor is not None
+            ):
+                output.mul_(self.moe_runner_config.routed_scaling_factor)
+
             return StandardCombineInput(hidden_states=output)
 
         quant_info = TritonMoeQuantInfo(

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

@@ -308,6 +308,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
                     self._cache_permute_indices,
                     layer.w13_weight.data[i].view(torch.uint8),
                     epilogue_tile_m,
+                    is_gated_act_gemm=layer.moe_runner_config.is_gated,
                 )
                 tmp_weights1 = (
                     layer.w13_weight.data[i]
@@ -509,6 +510,13 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
                     else ActivationType.Swiglu
                 ),
             )[0]
+
+            if (
+                not layer.should_fuse_routed_scaling_factor_in_topk
+                and moe_runner_config.routed_scaling_factor is not None
+            ):
+                output.mul_(moe_runner_config.routed_scaling_factor)
+
             return StandardCombineInput(hidden_states=output)
         elif self.use_flashinfer_trtllm_moe:
             from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import (

python/sglang/srt/models/nemotron_h.py

@@ -164,7 +164,6 @@ class NemotronHMoE(nn.Module):
             config.hidden_size,
             config.n_routed_experts,
             bias=False,
-            params_dtype=torch.float32,
             quant_config=None,
             prefix=f"{prefix}.gate",
         )
@@ -172,16 +171,6 @@ class NemotronHMoE(nn.Module):
             torch.empty(config.n_routed_experts, dtype=torch.float32)
         )
 
-        self.topk = TopK(
-            top_k=config.num_experts_per_tok,
-            use_grouped_topk=True,
-            topk_group=config.topk_group,
-            num_expert_group=config.n_group,
-            renormalize=config.norm_topk_prob,
-            scoring_func="sigmoid",
-            correction_bias=self.gate.e_score_correction_bias,
-            routed_scaling_factor=1.0,
-        )
         self.experts = get_moe_impl_class(quant_config)(
             num_experts=config.n_routed_experts
             + get_global_server_args().ep_num_redundant_experts,
@@ -195,6 +184,18 @@ class NemotronHMoE(nn.Module):
             layer_id=layer_idx,
             is_gated=False,
             routing_method_type=RoutingMethodType.DeepSeekV3,
+            routed_scaling_factor=self.routed_scaling_factor,
+        )
+        self.topk = TopK(
+            top_k=config.num_experts_per_tok,
+            use_grouped_topk=True,
+            topk_group=config.topk_group,
+            num_expert_group=config.n_group,
+            renormalize=config.norm_topk_prob,
+            scoring_func="sigmoid",
+            correction_bias=self.gate.e_score_correction_bias,
+            routed_scaling_factor=self.routed_scaling_factor,
+            apply_routed_scaling_factor_on_output=self.experts.should_fuse_routed_scaling_factor_in_topk,
         )
         if config.n_shared_experts:
             self.shared_experts = NemotronHMLP(
@@ -243,7 +244,10 @@ class NemotronHMoE(nn.Module):
         hidden_states: torch.Tensor,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         # router_scores: [num_tokens, num_experts]
-        router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32))
+        # bf16 gemm on tensor cores with fp32 accumulation/output for sigmoid/topk.
+        router_logits = torch.mm(
+            hidden_states, self.gate.weight.t(), out_dtype=torch.float32
+        )
         if self.shared_experts is not None:
             shared_output = self.shared_experts(hidden_states)
         else:
@@ -269,7 +273,10 @@ class NemotronHMoE(nn.Module):
 
         with self.device_module.stream(alt_stream):
             # router_scores: [num_tokens, num_experts]
-            router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32))
+            # bf16 gemm on tensor cores with fp32 accumulation/output for sigmoid/topk.
+            router_logits = torch.mm(
+                hidden_states, self.gate.weight.t(), out_dtype=torch.float32
+            )
             topk_output = self.topk(hidden_states, router_logits)
             if self.use_latent_moe:
                 hidden_states, _ = self.fc1_latent_proj(hidden_states)
@@ -280,15 +287,10 @@ class NemotronHMoE(nn.Module):
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         num_tokens, hidden_dim = hidden_states.shape
+        # routed_scaling_factor is fused into the experts call (applied by the
+        # MoE runner / topk), so final_hidden_states is already scaled.
         final_hidden_states, shared_output = self._forward_core(hidden_states)
 
-        # Fix FP16 overflow
-        if hidden_states.dtype != torch.float16:
-            final_hidden_states *= self.routed_scaling_factor
-        elif self.shared_experts is not None:
-            assert shared_output is not None
-            shared_output *= 1.0 / self.routed_scaling_factor
-
         if self.use_latent_moe:
             final_hidden_states, _ = self.fc2_latent_proj(final_hidden_states)
 
@@ -414,20 +416,18 @@ class NemotronHMambaDecoderLayer(nn.Module):
     def _forward_mamba(
         self, hidden_states: torch.Tensor, forward_batch: ForwardBatch
     ) -> torch.Tensor:
-        """Core Mamba forward logic, called directly or via split op."""
-        output = torch.empty_like(hidden_states)
+        """Core Mamba forward logic for the eager path; returns the result."""
         attn_backend = get_attn_backend()
         assert isinstance(attn_backend, HybridLinearAttnBackend)
         assert isinstance(attn_backend.linear_attn_backend, Mamba2AttnBackend)
-        attn_backend.linear_attn_backend.forward(
+        return attn_backend.linear_attn_backend.forward(
             mixer=self.mixer,
             layer_id=self.layer_id,
             hidden_states=hidden_states,
-            output=output,
+            output=None,
             forward_batch=forward_batch,
             use_triton_causal_conv=True,
         )
-        return output
 
     def forward(
         self,