diff --git a/python/sglang/jit_kernel/activation.py b/python/sglang/jit_kernel/activation.py index 89b28bbf6..e94aecb4c 100644 --- a/python/sglang/jit_kernel/activation.py +++ b/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, diff --git a/python/sglang/jit_kernel/benchmark/bench_activation.py b/python/sglang/jit_kernel/benchmark/bench_activation.py index 5b0e7e332..9277ec132 100644 --- a/python/sglang/jit_kernel/benchmark/bench_activation.py +++ b/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() diff --git a/python/sglang/jit_kernel/csrc/elementwise/activation.cuh b/python/sglang/jit_kernel/csrc/elementwise/activation.cuh index f6626ee1e..2dbf4cc78 100644 --- a/python/sglang/jit_kernel/csrc/elementwise/activation.cuh +++ b/python/sglang/jit_kernel/csrc/elementwise/activation.cuh @@ -19,6 +19,7 @@ enum class ActivationKind : uint32_t { kSiLU, kGELU, kGELUTanh, + kReLU2, }; template @@ -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, "unsupported activation kind"); return 0.0f; @@ -81,6 +85,30 @@ __global__ void act_and_mul_kernel(const __grid_constant__ ActivationParams para PDLTriggerSecondary(); } +struct UnaryActivationParams { + const void* __restrict__ input; + void* __restrict__ out; + uint32_t num_vecs; +}; + +template +__global__ void act_kernel(const __grid_constant__ UnaryActivationParams params) { + using namespace device; + constexpr auto kVecSize = kMaxVecBytes / sizeof(T); + using vec_t = AlignedVector; + const auto vec_id = blockIdx.x * blockDim.x + threadIdx.x; + if (vec_id >= params.num_vecs) return; + PDLWaitPrimary(); + const auto in = device::load_as(params.input, vec_id); + vec_t out; +#pragma unroll + for (int i = 0; i < kVecSize; ++i) { + out[i] = device::cast(apply_activation_f32(device::cast(in[i]))); + } + device::store_as(params.out, out, vec_id); + PDLTriggerSecondary(); +} + template 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(expert_ids.data_ptr()), static_cast(expert_step)); } + + template + static constexpr auto unary_kernel = act_kernel; + + static auto select_unary_kernel(const std::string& type) + -> decltype(ActivationKernel::template unary_kernel) { + using namespace host; + if (type == "relu2") { + return ActivationKernel::template unary_kernel; + } 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(); + + TensorMatcher({N, D}) // + .with_dtype() + .with_device(device_) + .verify(out) + .verify(input); + + const auto num_elems = static_cast(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::max(), "too many items for 32-bit indexing"); + const auto num_blocks = div_ceil(static_cast(num_vecs), kBlockSize); + const auto params = UnaryActivationParams{ + .input = input.data_ptr(), + .out = out.data_ptr(), + .num_vecs = static_cast(num_vecs), + }; + const auto kernel = select_unary_kernel(type); + LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params); + } }; } // namespace diff --git a/python/sglang/jit_kernel/tests/test_activation.py b/python/sglang/jit_kernel/tests/test_activation.py index e384f31f3..832d73944 100644 --- a/python/sglang/jit_kernel/tests/test_activation.py +++ b/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"])) diff --git a/python/sglang/srt/arg_groups/nemotron_h_hook.py b/python/sglang/srt/arg_groups/nemotron_h_hook.py index 8a31d04a2..a4e17670f 100644 --- a/python/sglang/srt/arg_groups/nemotron_h_hook.py +++ b/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, diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index a76c454fa..27f3c297c 100644 --- a/python/sglang/srt/layers/activation.py +++ b/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: diff --git a/python/sglang/srt/layers/attention/fla/layernorm_gated.py b/python/sglang/srt/layers/attention/fla/layernorm_gated.py index 38c226033..5a8fda41b 100644 --- a/python/sglang/srt/layers/attention/fla/layernorm_gated.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index af3395999..d0a8ff7a9 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/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" diff --git a/python/sglang/srt/layers/attention/mamba/causal_conv1d_triton.py b/python/sglang/srt/layers/attention/mamba/causal_conv1d_triton.py index 3f9b89eb5..d70f176e8 100644 --- a/python/sglang/srt/layers/attention/mamba/causal_conv1d_triton.py +++ b/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) diff --git a/python/sglang/srt/layers/attention/mamba/mamba.py b/python/sglang/srt/layers/attention/mamba/mamba.py index 1c46a35d4..042cf43dd 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba.py +++ b/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: diff --git a/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py b/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py index 547b4aa99..8fa5b0850 100644 --- a/python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py +++ b/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, ) diff --git a/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py b/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py index c89a4f86b..6aba2737e 100644 --- a/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py +++ b/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, ) diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index af525600a..20a699661 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/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( diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 4c252d093..9aab8f72c 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/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 ( diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 3b903c3b3..3d50a20cf 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/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,