fix(qwen3): support DeepEP-class backends and early EPLB state (#34810)

This commit is contained in:
Xun Sun
2026-08-15 01:27:39 +08:00
committed by GitHub
parent 7562e741e2
commit d8399af70c
+11 -3
View File
@@ -52,6 +52,7 @@ from sglang.srt.layers.moe.topk import TopK
from sglang.srt.layers.moe.utils import (
RoutingMethodType,
filter_moe_weight_param_global_expert,
is_deepep_class_backend,
)
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.radix_attention import RadixAttention
@@ -284,7 +285,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
prefix=add_prefix("gate", prefix),
)
if get_moe_a2a_backend().is_deepep():
if is_deepep_class_backend():
# TODO: we will support tp < ep in the future
self.ep_size = get_parallel().moe_ep_size
self.num_experts = (
@@ -299,8 +300,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
) -> torch.Tensor:
if (
not get_moe_a2a_backend().is_deepep()
and not get_moe_a2a_backend().is_mori()
not is_deepep_class_backend()
and not get_moe_a2a_backend().is_ascend_fuseep()
):
return self.forward_normal(hidden_states)
@@ -965,6 +965,14 @@ class Qwen3MoeForCausalLM(nn.Module):
)
self.logits_processor = LogitsProcessor(config)
self.capture_aux_hidden_states = False
# IPC loading bypasses load_weights(), so initialize the EPLB descriptor here.
self.routed_experts_weights_of_layer = LazyValue(
lambda: {
layer_id: self.model.layers[layer_id].mlp.get_moe_weights()
for layer_id in range(self.start_layer, self.end_layer)
if isinstance(self.model.layers[layer_id].mlp, Qwen3MoeSparseMoeBlock)
}
)
self.attn_cp_size = get_parallel().attn_cp_size
self.attn_cp_rank = get_parallel().attn_cp_rank