diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 4006de801..ed94416e5 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -150,6 +150,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module): quant_config: Optional[QuantizationConfig] = None, alt_stream: Optional[torch.cuda.Stream] = None, prefix: str = "", + is_nextn: bool = False, ): super().__init__() self.tp_size = get_tensor_model_parallel_world_size() @@ -220,6 +221,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module): config.num_experts + get_global_server_args().ep_num_redundant_experts ) self.top_k = config.num_experts_per_tok + self.is_nextn = is_nextn def get_moe_weights(self): return [ @@ -262,8 +264,12 @@ class Qwen2MoeSparseMoeBlock(nn.Module): hidden_states, router_logits, num_token_non_padded=forward_batch.num_token_non_padded, - expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( - layer_id=self.layer_id, + expert_location_dispatch_info=( + ExpertLocationDispatchInfo.init_new( + layer_id=self.layer_id, + ) + if not self.is_nextn + else None ), ) else: diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index b7c3f1ec0..5b2128e68 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -73,7 +73,14 @@ from sglang.srt.models.qwen2_moe import Qwen2MoeMLP, Qwen2MoeSparseMoeBlock from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration # Utils -from sglang.srt.utils import add_prefix, is_cuda, is_npu, make_layers, set_weight_attrs +from sglang.srt.utils import ( + LazyValue, + add_prefix, + is_cuda, + is_npu, + make_layers, + set_weight_attrs, +) from sglang.srt.utils.hf_transformers_utils import get_processor logger = logging.getLogger(__name__) @@ -295,6 +302,7 @@ class Qwen3_5LinearDecoderLayer(nn.Module): quant_config: Optional[QuantizationConfig] = None, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, + is_nextn: bool = False, ) -> None: super().__init__() self.config = config @@ -318,6 +326,7 @@ class Qwen3_5LinearDecoderLayer(nn.Module): quant_config=quant_config, alt_stream=alt_stream, prefix=add_prefix("mlp", prefix.replace(".linear_attn", "")), + is_nextn=is_nextn, ) is_layer_sparse = True is_previous_layer_sparse = True @@ -414,6 +423,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): quant_config: Optional[QuantizationConfig] = None, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, + is_nextn: bool = False, ) -> None: super().__init__() self.config = config @@ -516,6 +526,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): quant_config=quant_config, alt_stream=alt_stream, prefix=add_prefix("mlp", prefix.replace(".self_attn", "")), + is_nextn=is_nextn, ) is_layer_sparse = True is_previous_layer_sparse = True @@ -665,6 +676,7 @@ class Qwen3_5ForCausalLM(nn.Module): config: Qwen3_5TextConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", + is_nextn: bool = False, ) -> None: super().__init__() self.config = config @@ -698,6 +710,7 @@ class Qwen3_5ForCausalLM(nn.Module): quant_config=quant_config, prefix=prefix, alt_stream=alt_stream, + is_nextn=is_nextn, ) self.layers = make_layers( @@ -858,6 +871,14 @@ class Qwen3_5ForCausalLM(nn.Module): loaded_params.add(name) return loaded_params + @classmethod + def get_model_config_for_expert_location(cls, config): + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.num_experts, + num_groups=None, + ) + class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM): def __init__( @@ -1385,8 +1406,20 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration): logger.warning(f"Parameter {name} not found in params_dict") loaded_params.add(name) + self._routed_experts_weights_of_layer = LazyValue( + lambda: { + layer_id: layer.mlp.get_moe_weights() + for layer_id, layer in enumerate(self.model.layers) + if isinstance(layer.mlp, Qwen2MoeSparseMoeBlock) + } + ) + return loaded_params + @property + def routed_experts_weights_of_layer(self): + return self._routed_experts_weights_of_layer.value + @classmethod def get_model_config_for_expert_location(cls, config): text_config = getattr(config, "text_config", config) diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index 97d02274b..3fa89fcda 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -22,6 +22,8 @@ from torch import nn from transformers import PretrainedConfig from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.layers.layernorm import GemmaRMSNorm from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE @@ -69,6 +71,7 @@ class Qwen3_5ForCausalLMMTP(nn.Module): config, quant_config, prefix=add_prefix("mtp", prefix), + is_nextn=True, ) if get_pp_group().is_last_rank: @@ -84,6 +87,15 @@ class Qwen3_5ForCausalLMMTP(nn.Module): self.logits_processor = LogitsProcessor(config) + @classmethod + def get_model_config_for_expert_location(cls, config): + text_config = getattr(config, "text_config", config) + return ModelConfigForExpertLocation( + num_layers=text_config.num_hidden_layers, + num_logical_experts=text_config.num_experts, + num_groups=None, + ) + def get_embed_and_head(self): return self.model.embed_tokens.weight, self.lm_head.weight @@ -130,12 +142,13 @@ class Qwen3_5ForCausalLMMTP(nn.Module): hidden_states = self.fc(hidden_states) - hidden_states = self.model( - input_ids, - positions, - forward_batch, - hidden_states, - ) + with get_global_expert_distribution_recorder().disable_this_region(): + hidden_states = self.model( + input_ids, + positions, + forward_batch, + hidden_states, + ) return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch diff --git a/python/sglang/srt/models/qwen3_next.py b/python/sglang/srt/models/qwen3_next.py index 42146e605..6fcdec281 100644 --- a/python/sglang/srt/models/qwen3_next.py +++ b/python/sglang/srt/models/qwen3_next.py @@ -485,6 +485,7 @@ class Qwen3HybridLinearDecoderLayer(nn.Module): quant_config: Optional[QuantizationConfig] = None, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, + is_nextn: bool = False, ) -> None: super().__init__() self.config = config @@ -513,6 +514,7 @@ class Qwen3HybridLinearDecoderLayer(nn.Module): quant_config=quant_config, alt_stream=alt_stream, prefix=add_prefix("mlp", prefix.replace(".linear_attn", "")), + is_nextn=is_nextn, ) else: self.mlp = Qwen2MoeMLP( @@ -582,6 +584,7 @@ class Qwen3HybridAttentionDecoderLayer(nn.Module): quant_config: Optional[QuantizationConfig] = None, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, + is_nextn: bool = False, ) -> None: super().__init__() self.config = config @@ -688,6 +691,7 @@ class Qwen3HybridAttentionDecoderLayer(nn.Module): quant_config=quant_config, alt_stream=alt_stream, prefix=add_prefix("mlp", prefix.replace(".self_attn", "")), + is_nextn=is_nextn, ) else: self.mlp = Qwen2MoeMLP( @@ -822,6 +826,7 @@ class Qwen3NextModel(nn.Module): config: Qwen3NextConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", + is_nextn: bool = False, ) -> None: super().__init__() self.config = config @@ -847,6 +852,7 @@ class Qwen3NextModel(nn.Module): quant_config=quant_config, prefix=prefix, alt_stream=alt_stream, + is_nextn=is_nextn, ) self.layers = make_layers( diff --git a/python/sglang/srt/models/qwen3_next_mtp.py b/python/sglang/srt/models/qwen3_next_mtp.py index 9722f7e4a..b2bdbbbe8 100644 --- a/python/sglang/srt/models/qwen3_next_mtp.py +++ b/python/sglang/srt/models/qwen3_next_mtp.py @@ -22,6 +22,7 @@ from torch import nn from transformers import PretrainedConfig from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.layers.layernorm import GemmaRMSNorm from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.quantization.base_config import QuantizationConfig @@ -62,7 +63,10 @@ class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM): config.num_hidden_layers = 1 config.full_attention_interval = 1 self.model = Qwen3NextModel( - config, quant_config, prefix=add_prefix("model", prefix) + config, + quant_config, + prefix=add_prefix("model", prefix), + is_nextn=True, ) self.lm_head = ParallelLMHead( config.vocab_size, @@ -92,12 +96,13 @@ class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM): hidden_states = self.pre_fc_norm_hidden(hidden_states) hidden_states = self.fc(torch.cat((input_embeds, hidden_states), dim=-1)) - hidden_states = self.model( - input_ids, - positions, - forward_batch, - hidden_states, - ) + with get_global_expert_distribution_recorder().disable_this_region(): + hidden_states = self.model( + input_ids, + positions, + forward_batch, + hidden_states, + ) return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch