chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Feat/nemotron nano v3 support (#12690)

Browse Source
This commit is contained in:
roikoren755
2025-11-21 23:53:05 +02:00
committed by GitHub
parent a24aefe562
commit 1b48e1b974
13 changed files with 776 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
python/sglang/srt/configs/nemotron_h.py
View File

@@ -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,

146
python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=1856,device_name=NVIDIA_H100_80GB_HBM3.json Normal file
View File

@@ -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
}
}

146
python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=1856,device_name=NVIDIA_L40S.json Normal file
View File

@@ -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
}
}

146
python/sglang/srt/layers/moe/fused_moe_triton/configs/triton_3_4_0/E=128,N=928,device_name=NVIDIA_L40S.json Normal file
View File

@@ -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
}
}

23
python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
View File

@@ -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,

21
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
View File

@@ -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

12
python/sglang/srt/layers/moe/moe_runner/base.py
View File

@@ -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

9
python/sglang/srt/layers/moe/moe_runner/deep_gemm.py
View File

@@ -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

2
python/sglang/srt/layers/moe/moe_runner/triton.py
View File

@@ -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 = (

4
python/sglang/srt/layers/moe/moe_runner/triton_kernels.py
View File

@@ -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(

13
python/sglang/srt/layers/quantization/unquant.py
View File

@@ -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)

289
python/sglang/srt/models/nemotron_h.py
View File

@@ -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]