diff --git a/benchmark/kernels/fused_moe_triton/common_utils.py b/benchmark/kernels/fused_moe_triton/common_utils.py index 630efd73e..d87350f9f 100644 --- a/benchmark/kernels/fused_moe_triton/common_utils.py +++ b/benchmark/kernels/fused_moe_triton/common_utils.py @@ -104,7 +104,7 @@ def get_model_config( E = config.num_experts // ep_size topk = config.num_experts_per_tok intermediate_size = config.moe_intermediate_size - elif architecture in ["Glm4MoeForCausalLM"]: + elif architecture in ["Glm4MoeForCausalLM", "NemotronHForCausalLM"]: E = config.n_routed_experts // ep_size topk = config.num_experts_per_tok intermediate_size = config.moe_intermediate_size diff --git a/python/sglang/srt/configs/nemotron_h.py b/python/sglang/srt/configs/nemotron_h.py index b73b146fe..b3526ded5 100644 --- a/python/sglang/srt/configs/nemotron_h.py +++ b/python/sglang/srt/configs/nemotron_h.py @@ -26,6 +26,7 @@ logger = logging.get_logger(__name__) MAMBA = "M" ATTENTION = "*" MLP = "-" +MOE = "E" class NemotronHConfig(PretrainedConfig): @@ -189,6 +190,15 @@ class NemotronHConfig(PretrainedConfig): mamba_proj_bias=False, mamba_chunk_size=256, rescale_prenorm_residual=True, + n_routed_experts=8, + n_shared_experts=1, + moe_intermediate_size=7688, + moe_shared_expert_intermediate_size=7688, + num_experts_per_tok=2, + routed_scaling_factor=1.0, + n_group=1, + topk_group=1, + norm_topk_prob=True, **kwargs, ): self.vocab_size = vocab_size @@ -206,12 +216,12 @@ class NemotronHConfig(PretrainedConfig): # Validate hybrid_override_pattern # M: Mamba2, *: Attention, -: MLP - assert len(self.hybrid_override_pattern) == self.num_hidden_layers, ( - "hybrid_override_pattern must have same length as " "num_hidden_layers" - ) - assert re.match(r"^[*-M]+$", self.hybrid_override_pattern), ( - "hybrid_override_pattern must only contain characters " "'M', '*', or '-'" - ) + assert ( + len(self.hybrid_override_pattern) == self.num_hidden_layers + ), "hybrid_override_pattern must have same length as num_hidden_layers" + assert re.match( + r"^[*\-ME]+$", self.hybrid_override_pattern + ), "hybrid_override_pattern must only contain characters 'M', '*', '-' or 'E'" # for backward compatibility if num_key_value_heads is None: @@ -245,6 +255,15 @@ class NemotronHConfig(PretrainedConfig): self.mamba_proj_bias = mamba_proj_bias self.mamba_chunk_size = mamba_chunk_size self.rescale_prenorm_residual = rescale_prenorm_residual + self.n_routed_experts = n_routed_experts + self.n_shared_experts = n_shared_experts + self.moe_intermediate_size = moe_intermediate_size + self.moe_shared_expert_intermediate_size = moe_shared_expert_intermediate_size + self.num_experts_per_tok = num_experts_per_tok + self.routed_scaling_factor = routed_scaling_factor + self.n_group = n_group + self.topk_group = topk_group + self.norm_topk_prob = norm_topk_prob super().__init__( pad_token_id=pad_token_id, diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=1856,device_name=NVIDIA_H100_80GB_HBM3.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=1856,device_name=NVIDIA_H100_80GB_HBM3.json new file mode 100644 index 000000000..ea50fa9ec --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=1856,device_name=NVIDIA_H100_80GB_HBM3.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "48": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "1536": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "2048": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "3072": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 4 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=1856,device_name=NVIDIA_L40S.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=1856,device_name=NVIDIA_L40S.json new file mode 100644 index 000000000..fbf1adb13 --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=1856,device_name=NVIDIA_L40S.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 5 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 2 + }, + "8": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 2 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "512": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 2 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 3 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=928,device_name=NVIDIA_L40S.json b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=928,device_name=NVIDIA_L40S.json new file mode 100644 index 000000000..b3ccd09bb --- /dev/null +++ b/python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=928,device_name=NVIDIA_L40S.json @@ -0,0 +1,146 @@ +{ + "1": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 3 + }, + "2": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2 + }, + "4": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 3 + }, + "8": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 5 + }, + "16": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "24": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 4 + }, + "32": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4 + }, + "48": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 3 + }, + "64": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "96": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 5 + }, + "128": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "256": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "512": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 5 + }, + "1024": { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 3 + }, + "1536": { + "BLOCK_SIZE_M": 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 256, + "GROUP_SIZE_M": 16, + "num_warps": 8, + "num_stages": 3 + }, + "2048": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 32, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 4, + "num_stages": 5 + }, + "3072": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 16, + "num_warps": 4, + "num_stages": 4 + }, + "4096": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 64, + "num_warps": 8, + "num_stages": 3 + } +} diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py index b92f2159e..2a88703bf 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py @@ -10,6 +10,7 @@ import os from typing import TYPE_CHECKING, List, Optional import torch +import torch.nn.functional as F import triton.language as tl from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig @@ -66,6 +67,7 @@ def inplace_fused_experts( b1: Optional[torch.Tensor] = None, b2: Optional[torch.Tensor] = None, activation: str = "silu", + is_gated: bool = True, apply_router_weight_on_input: bool = False, use_fp8_w8a8: bool = False, use_int8_w8a8: bool = False, @@ -94,6 +96,7 @@ def inplace_fused_experts( b2, True, activation, + is_gated, apply_router_weight_on_input, use_fp8_w8a8, use_int8_w8a8, @@ -124,6 +127,7 @@ def inplace_fused_experts_fake( b1: Optional[torch.Tensor] = None, b2: Optional[torch.Tensor] = None, activation: str = "silu", + is_gated: bool = True, apply_router_weight_on_input: bool = False, use_fp8_w8a8: bool = False, use_int8_w8a8: bool = False, @@ -162,6 +166,7 @@ def outplace_fused_experts( b1: Optional[torch.Tensor] = None, b2: Optional[torch.Tensor] = None, activation: str = "silu", + is_gated: bool = True, apply_router_weight_on_input: bool = False, use_fp8_w8a8: bool = False, use_int8_w8a8: bool = False, @@ -191,6 +196,7 @@ def outplace_fused_experts( b2, False, activation, + is_gated, apply_router_weight_on_input, use_fp8_w8a8, use_int8_w8a8, @@ -221,6 +227,7 @@ def outplace_fused_experts_fake( b1: Optional[torch.Tensor] = None, b2: Optional[torch.Tensor] = None, activation: str = "silu", + is_gated: bool = True, apply_router_weight_on_input: bool = False, use_fp8_w8a8: bool = False, use_int8_w8a8: bool = False, @@ -288,6 +295,7 @@ def fused_experts( b1, b2, moe_runner_config.activation, + moe_runner_config.is_gated, moe_runner_config.apply_router_weight_on_input, use_fp8_w8a8, use_int8_w8a8, @@ -367,6 +375,7 @@ def fused_experts_impl( b2: Optional[torch.Tensor] = None, inplace: bool = False, activation: str = "silu", + is_gated: bool = True, apply_router_weight_on_input: bool = False, use_fp8_w8a8: bool = False, use_int8_w8a8: bool = False, @@ -533,7 +542,8 @@ def fused_experts_impl( c_sorted=down_moe_use_tma, filter_expert=filter_expert, ) - if activation == "silu": + # Activation function with multiplication + if activation == "silu" and is_gated: if gemm1_alpha is not None: assert gemm1_limit is not None intermediate_cache2 = swiglu_with_alpha_and_limit( @@ -547,7 +557,7 @@ def fused_experts_impl( vllm_ops.silu_and_mul( intermediate_cache2, intermediate_cache1.view(-1, N) ) - elif activation == "gelu": + elif activation == "gelu" and is_gated: assert gemm1_alpha is None, "gemm1_alpha is not supported for gelu" assert gemm1_limit is None, "gemm1_limit is not supported for gelu" if _is_cuda or _is_hip: @@ -556,8 +566,15 @@ def fused_experts_impl( vllm_ops.gelu_and_mul( intermediate_cache2, intermediate_cache1.view(-1, N) ) + # Activation function without multiplication + elif activation == "silu" and not is_gated: + intermediate_cache2 = F.silu(intermediate_cache1.view(-1, N)) + elif activation == "gelu" and not is_gated: + intermediate_cache2 = F.gelu(intermediate_cache1.view(-1, N)) + elif activation == "relu2" and not is_gated: + intermediate_cache2 = torch.square(F.relu(intermediate_cache1.view(-1, N))) else: - raise ValueError(f"Unsupported activation: {activation=}") + raise ValueError(f"Unsupported activation: {activation=}, with {is_gated=}") invoke_fused_moe_kernel( intermediate_cache2, diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index aae387734..47a9da781 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -148,6 +148,7 @@ class FusedMoE(torch.nn.Module): use_weight_loader_fused: bool = False, with_bias=False, routing_method_type: Optional[RoutingMethodType] = None, + is_gated: bool = True, ): super().__init__() if params_dtype is None: @@ -211,6 +212,7 @@ class FusedMoE(torch.nn.Module): routed_scaling_factor=routed_scaling_factor, gemm1_alpha=gemm1_alpha, gemm1_clamp_limit=gemm1_clamp_limit, + is_gated=is_gated, ) self.quant_method: Optional[FusedMoEMethodBase] = None @@ -344,10 +346,12 @@ class FusedMoE(torch.nn.Module): # if this weight is a bias, the last dimension must be the sharded dimension shard_dim = -1 - if shard_id in {"w1", "w3"}: + if shard_id in {"w1", "w3"} and self.moe_runner_config.is_gated: # non-fused version shard_size = expert_data.shape[shard_dim] // 2 - elif shard_id in {"w13"}: + elif shard_id in {"w13"} or ( + shard_id in {"w1", "w3"} and not self.moe_runner_config.is_gated + ): # fused version shard_size = expert_data.shape[shard_dim] else: @@ -620,9 +624,7 @@ class FusedMoE(torch.nn.Module): ) if shard_id not in ("w1", "w2", "w3"): - raise ValueError( - f"shard_id must be ['w1','w2','w3'] but " f"got {shard_id}." - ) + raise ValueError(f"shard_id must be ['w1','w2','w3'] but got {shard_id}.") # Flashinfer assumes w31 format for w13_weight. Same for the scales. if get_moe_runner_backend().is_flashinfer_trtllm() and ( @@ -821,7 +823,7 @@ class FusedMoE(torch.nn.Module): ) if shard_id not in ("w13", "w2"): - raise ValueError(f"shard_id must be ['w13','w2'] but " f"got {shard_id}.") + raise ValueError(f"shard_id must be ['w13','w2'] but got {shard_id}.") # Fetch the dim to shard the parameter/loaded weight # based on the shard id. This will be whatever @@ -1022,6 +1024,9 @@ class FlashInferFusedMoE(FusedMoE): assert ( self.num_fused_shared_experts == 0 ), "Fused shared experts are not supported for flashinfer blockscale fp8 moe" + assert ( + self.moe_runner_config.is_gated + ), "Only gated MoEs are supported for flashinfer blockscale fp8 moe" assert TopKOutputChecker.format_is_bypassed(topk_output) @@ -1089,6 +1094,10 @@ class FlashInferFP4MoE(FusedMoE): """ assert isinstance(self.quant_method, ModelOptNvFp4FusedMoEMethod) + assert ( + self.moe_runner_config.is_gated + ), "Only gated MoEs are supported for flashinfer fp4 moe" + assert TopKOutputChecker.format_is_bypassed(topk_output) router_logits = topk_output.router_logits diff --git a/python/sglang/srt/layers/moe/moe_runner/base.py b/python/sglang/srt/layers/moe/moe_runner/base.py index 4d95540e6..118206a90 100644 --- a/python/sglang/srt/layers/moe/moe_runner/base.py +++ b/python/sglang/srt/layers/moe/moe_runner/base.py @@ -24,7 +24,6 @@ if TYPE_CHECKING: @dataclass class MoeRunnerConfig: - # MoE parameters num_experts: Optional[int] = None num_local_experts: Optional[int] = None @@ -37,6 +36,7 @@ class MoeRunnerConfig: # Runner configuration activation: str = "silu" + is_gated: bool = True apply_router_weight_on_input: bool = False inplace: bool = True no_combine: bool = False @@ -47,7 +47,6 @@ class MoeRunnerConfig: @dataclass class RunnerInput(ABC): - @property @abstractmethod def runner_backend(self) -> MoeRunnerBackend: ... @@ -57,7 +56,6 @@ class RunnerInput(ABC): class RunnerOutput(ABC): - @property @abstractmethod def runner_backend(self) -> MoeRunnerBackend: ... @@ -74,7 +72,6 @@ class MoeQuantInfo(ABC): class MoeRunnerCore(ABC): - def __init__(self, config: MoeRunnerConfig): self.config = config @@ -93,7 +90,6 @@ class MoeRunnerCore(ABC): class FusedOpPool: - _fused_funcs: dict[str, Callable] = {} @classmethod @@ -121,7 +117,6 @@ class FusedOpPool: class PermuteMethodPool: - _pre_permute_methods: dict[ Tuple[DispatchOutputFormat, MoeRunnerBackend], Callable ] = {} @@ -250,9 +245,8 @@ def register_pre_permute( def decorator( permute_func: Callable[ [DispatchOutput, MoeQuantInfo, MoeRunnerConfig, dict], RunnerInput - ] + ], ) -> Callable: - PermuteMethodPool.register_pre_permute( dispatch_output_name, runner_backend_name, permute_func ) @@ -276,7 +270,7 @@ def register_post_permute( def decorator( permute_func: Callable[ [RunnerOutput, MoeQuantInfo, MoeRunnerConfig, dict], CombineInput - ] + ], ) -> Callable: PermuteMethodPool.register_post_permute( runner_backend_name, combine_input_name, permute_func diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index b74560910..436f0364a 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -104,6 +104,7 @@ class DeepGemmRunnerCore(MoeRunnerCore): def __init__(self, config: MoeRunnerConfig): super().__init__(config) assert self.config.activation == "silu" + assert self.config.is_gated def run( self, @@ -111,7 +112,6 @@ class DeepGemmRunnerCore(MoeRunnerCore): quant_info: DeepGemmMoeQuantInfo, running_state: dict, ) -> DeepGemmRunnerOutput: - if not runner_input.use_masked_gemm: hidden_states = self._run_contiguous_gemm( runner_input, quant_info, running_state @@ -128,7 +128,6 @@ class DeepGemmRunnerCore(MoeRunnerCore): quant_info: DeepGemmMoeQuantInfo, running_state: dict, ) -> torch.Tensor: - from sglang.srt.layers.moe.ep_moe.kernels import tma_align_input_scale from sglang.srt.layers.quantization.fp8_kernel import ( sglang_per_token_group_quant_fp8, @@ -212,7 +211,6 @@ class DeepGemmRunnerCore(MoeRunnerCore): quant_info: DeepGemmMoeQuantInfo, running_state: dict, ) -> torch.Tensor: - from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.moe.ep_moe.kernels import ( silu_and_mul_masked_post_quant_fwd, @@ -339,7 +337,6 @@ def pre_permute_standard_to_deep_gemm( runner_config: MoeRunnerConfig, running_state: dict, ) -> DeepGemmRunnerInput: - from sglang.srt.layers.moe.ep_moe.kernels import moe_ep_deepgemm_preprocess hidden_states, topk_output = ( @@ -432,7 +429,6 @@ def pre_permute_deepep_ll_to_deep_gemm( runner_config: MoeRunnerConfig, running_state: dict, ) -> DeepGemmRunnerInput: - hidden_states, hidden_states_scale, topk_ids, topk_weights, masked_m, expected_m = ( dispatch_output ) @@ -459,7 +455,6 @@ def post_permute_deep_gemm_to_deepep_ll( runner_config: MoeRunnerConfig, running_state: dict, ) -> DeepEPLLCombineInput: - from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPLLCombineInput return DeepEPLLCombineInput( @@ -476,7 +471,6 @@ def pre_permute_deepep_normal_to_deep_gemm( runner_config: MoeRunnerConfig, running_state: dict, ) -> DeepGemmRunnerInput: - from sglang.srt.layers.moe.ep_moe.kernels import ep_scatter ( @@ -569,7 +563,6 @@ def post_permute_deep_gemm_to_deepep_normal( runner_config: MoeRunnerConfig, running_state: dict, ) -> DeepEPNormalCombineInput: - from sglang.srt.layers.moe.ep_moe.kernels import ep_gather from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPNormalCombineInput diff --git a/python/sglang/srt/layers/moe/moe_runner/triton.py b/python/sglang/srt/layers/moe/moe_runner/triton.py index 25fbcf770..7de28f482 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton.py @@ -153,6 +153,8 @@ class TritonRunnerCore(MoeRunnerCore): routed_scaling_factor = self.config.routed_scaling_factor apply_router_weight_on_input = self.config.apply_router_weight_on_input + assert self.config.is_gated, "Only gated MoEs are supported for Triton runner" + M = hidden_states.shape[0] E, N, _ = w13.shape compute_type = ( diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py b/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py index 9d088faa8..b13cd2759 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_kernels.py @@ -90,6 +90,10 @@ class TritonKernelsRunnerCore(MoeRunnerCore): triton_kernel_fused_experts_with_bias, ) + assert ( + self.config.is_gated + ), "Only gated MoEs are supported for Triton Kernels runner" + hidden_states = runner_input.hidden_states common_kwargs = dict( diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 6c6c728e2..c0f35905a 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -153,7 +153,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): self.with_bias = with_bias # Fused gate_up_proj (column parallel) - w13_weight_n, w13_weight_k = 2 * intermediate_size_per_partition, hidden_size + w13_up_dim = ( + 2 * intermediate_size_per_partition + if layer.moe_runner_config.is_gated + else intermediate_size_per_partition + ) + w13_weight_n, w13_weight_k = (w13_up_dim, hidden_size) if self.use_triton_kernels: w13_weight_n, w13_weight_k = w13_weight_k, w13_weight_n w13_weight = torch.nn.Parameter( @@ -165,11 +170,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): if self.with_bias: w13_weight_bias = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - dtype=torch.float32, - ), + torch.empty(num_experts, w13_up_dim, dtype=torch.float32), requires_grad=False, ) layer.register_parameter("w13_weight_bias", w13_weight_bias) diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index eadff130f..8a2075ea9 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/python/sglang/srt/models/nemotron_h.py @@ -22,8 +22,13 @@ import torch from torch import nn from sglang.srt.configs import NemotronHConfig -from sglang.srt.configs.nemotron_h import ATTENTION, MAMBA, MLP -from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size +from sglang.srt.configs.nemotron_h import ATTENTION, MAMBA, MLP, MOE +from sglang.srt.distributed import ( + get_moe_ep_group, + get_pp_group, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) from sglang.srt.layers.activation import ReLU2 from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( HybridLinearAttnBackend, @@ -34,9 +39,13 @@ from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( ColumnParallelLinear, QKVParallelLinear, + ReplicatedLinear, RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.moe.topk import TopK from sglang.srt.layers.quantization import QuantizationConfig from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.vocab_parallel_embedding import ( @@ -51,31 +60,30 @@ from sglang.srt.model_loader.weight_utils import ( replace_prefix, replace_substrings, ) -from sglang.srt.utils import add_prefix, make_layers_non_pp +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import ( + add_prefix, + get_current_device_stream_fast, + is_cuda, + make_layers_non_pp, +) from sglang.utils import logger +_is_cuda = is_cuda() + class NemotronHMLP(nn.Module): def __init__( self, config: NemotronHConfig, - layer_idx: int, + intermediate_size: int, quant_config: Optional[QuantizationConfig] = None, bias: bool = False, + reduce_results: bool = True, prefix: str = "", ) -> None: super().__init__() - hybrid_override_pattern = config.hybrid_override_pattern - mlp_index = hybrid_override_pattern[: layer_idx + 1].count("-") - 1 - if isinstance(config.intermediate_size, list): - if len(config.intermediate_size) == 1: - intermediate_size = config.intermediate_size[0] - else: - intermediate_size = config.intermediate_size[mlp_index] - else: - intermediate_size = config.intermediate_size - self.up_proj = ColumnParallelLinear( input_size=config.hidden_size, output_size=intermediate_size, @@ -88,6 +96,7 @@ class NemotronHMLP(nn.Module): output_size=config.hidden_size, bias=bias, quant_config=quant_config, + reduce_results=reduce_results, prefix=f"{prefix}.down_proj", ) self.act_fn = ReLU2() @@ -99,6 +108,148 @@ class NemotronHMLP(nn.Module): return x +_alt_stream = None + + +def _get_or_create_alt_stream(device_module): + global _alt_stream + if _alt_stream is None: + _alt_stream = device_module.Stream() + return _alt_stream + + +class NemotronHMoE(nn.Module): + def __init__( + self, + config: NemotronHConfig, + layer_idx: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.tp_size = get_tensor_model_parallel_world_size() + self.routed_scaling_factor = config.routed_scaling_factor + self.device_module = torch.get_device_module() + + self.ep_group = get_moe_ep_group().device_group + self.ep_rank = self.ep_group.rank() + self.ep_size = self.ep_group.size() + self.n_routed_experts = config.n_routed_experts + self.n_shared_experts = config.n_shared_experts + + self.gate = ReplicatedLinear( + config.hidden_size, + config.n_routed_experts, + bias=False, + params_dtype=torch.float32, + quant_config=None, + prefix=f"{prefix}.gate", + ) + self.gate.e_score_correction_bias = nn.Parameter( + 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, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + reduce_results=False, + quant_config=quant_config, + prefix=f"{prefix}.experts", + activation=config.mlp_hidden_act, + layer_id=layer_idx, + is_gated=False, + ) + if config.n_shared_experts: + self.shared_experts = NemotronHMLP( + config, + intermediate_size=config.moe_shared_expert_intermediate_size + * config.n_shared_experts, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_experts", + ) + else: + self.shared_experts = None + + def _forward_core( + self, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if _is_cuda: + return self._forward_core_shared_routed_overlap(hidden_states) + else: + return self._forward_core_normal(hidden_states) + + def _forward_core_normal( + self, + 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)) + if self.shared_experts is not None: + shared_output = self.shared_experts(hidden_states) + else: + shared_output = None + topk_output = self.topk(hidden_states, router_logits) + final_hidden_states = self.experts(hidden_states, topk_output) + return final_hidden_states, shared_output + + def _forward_core_shared_routed_overlap( + self, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + alt_stream = _get_or_create_alt_stream(self.device_module) + + alt_stream.wait_stream(get_current_device_stream_fast()) + + if self.shared_experts is not None: + shared_output = self.shared_experts(hidden_states) + else: + shared_output = None + + with self.device_module.stream(alt_stream): + # router_scores: [num_tokens, num_experts] + router_logits, _ = self.gate(hidden_states.to(dtype=torch.float32)) + topk_output = self.topk(hidden_states, router_logits) + final_hidden_states = self.experts(hidden_states, topk_output) + get_current_device_stream_fast().wait_stream(alt_stream) + + return final_hidden_states, shared_output + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_dim = hidden_states.shape + 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 shared_output is not None: + final_hidden_states += shared_output + + if self.tp_size > 1: + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + + return final_hidden_states.view(num_tokens, hidden_dim) + + class NemotronHMLPDecoderLayer(nn.Module): def __init__( self, @@ -110,15 +261,61 @@ class NemotronHMLPDecoderLayer(nn.Module): super().__init__() self.config = config + hybrid_override_pattern = config.hybrid_override_pattern + mlp_index = hybrid_override_pattern[: layer_idx + 1].count("-") - 1 + if isinstance(config.intermediate_size, list): + if len(config.intermediate_size) == 1: + intermediate_size = config.intermediate_size[0] + else: + intermediate_size = config.intermediate_size[mlp_index] + else: + intermediate_size = config.intermediate_size + self.mixer = NemotronHMLP( config, + intermediate_size=intermediate_size, quant_config=quant_config, bias=config.mlp_bias, prefix=f"{prefix}.mixer", - layer_idx=layer_idx, ) - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) + + def forward( + self, + *, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor], + forward_batch: ForwardBatch, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.norm(hidden_states) + else: + hidden_states, residual = self.norm(hidden_states, residual) + + hidden_states = self.mixer.forward(hidden_states) + return hidden_states, residual + + +class NemotronHMoEDecoderLayer(nn.Module): + def __init__( + self, + config: NemotronHConfig, + layer_idx: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.mixer = NemotronHMoE( + config, + layer_idx=layer_idx, + quant_config=quant_config, + prefix=f"{prefix}.mixer", + ) + + self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) def forward( self, @@ -154,13 +351,13 @@ class NemotronHMambaDecoderLayer(nn.Module): use_conv_bias=config.use_conv_bias, use_bias=config.use_bias, n_groups=config.mamba_n_groups, - rms_norm_eps=config.rms_norm_eps, + rms_norm_eps=config.layer_norm_epsilon, activation=config.mamba_hidden_act, quant_config=quant_config, prefix=f"{prefix}.mixer", ) - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) def forward( self, @@ -275,7 +472,7 @@ class NemotronHAttentionDecoderLayer(nn.Module): prefix=f"{prefix}.mixer", ) - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) def forward( self, @@ -300,11 +497,13 @@ Layers = ( NemotronHAttentionDecoderLayer | NemotronHMLPDecoderLayer | NemotronHMambaDecoderLayer + | NemotronHMoEDecoderLayer ) ALL_DECODER_LAYER_TYPES: dict[str, type[Layers]] = { ATTENTION: NemotronHAttentionDecoderLayer, MLP: NemotronHMLPDecoderLayer, MAMBA: NemotronHMambaDecoderLayer, + MOE: NemotronHMoEDecoderLayer, } @@ -341,7 +540,7 @@ class NemotronHModel(nn.Module): self.layers = make_layers_non_pp( len(config.hybrid_override_pattern), get_layer, prefix=f"{prefix}.layers" ) - self.norm_f = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norm_f = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -473,6 +672,18 @@ class NemotronHForCausalLM(nn.Module): name = replace_prefix(name, self.remap_prefix) name = replace_substrings(name, self.remap_substr) updated_weights.append((name, loaded_weight)) + + # - FusedMoe.w1 (aka gate_proj) should be up_proj since that's + # what the activation is applied to + # - FusedMoe.w3 (aka up_proj) should be ignored since we're + # using non-gated MoE + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="up_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="", + num_experts=self.config.n_routed_experts, + ) + params_dict = dict(self.named_parameters()) for name, loaded_weight in updated_weights: @@ -495,17 +706,37 @@ class NemotronHForCausalLM(nn.Module): weight_loader(param, loaded_weight, shard_id) break else: - # Skip loading extra bias for GPTQ models. - if name.endswith(".bias") and name not in params_dict: - continue - if name in params_dict.keys(): - param = params_dict[name] - weight_loader = getattr( - param, "weight_loader", default_weight_loader + is_expert_weight = False + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + is_expert_weight = True + name_mapped = name.replace(weight_name, param_name) + param = params_dict[name_mapped] + param.weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, ) - weight_loader(param, loaded_weight) + name = name_mapped + break else: - logger.warning(f"Parameter {name} not found in params_dict") + if is_expert_weight: + continue + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if name in params_dict.keys(): + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + else: + logger.warning(f"Parameter {name} not found in params_dict") EntryClass = [NemotronHForCausalLM]