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

Fix qwen3.5 mtp eplb related issues (#19767)

Browse Source
This commit is contained in:
luoyuyan
2026-03-09 16:05:32 +08:00
committed by GitHub
parent c76251f70c
commit cabe171b6c
5 changed files with 79 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/srt/models/qwen2_moe.py
View File

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

35
python/sglang/srt/models/qwen3_5.py
View File

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

25
python/sglang/srt/models/qwen3_5_mtp.py
View File

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

6
python/sglang/srt/models/qwen3_next.py
View File

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

19
python/sglang/srt/models/qwen3_next_mtp.py
View File

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