feat PD: add eagle3 support for DeepSeek V3 in EP mode (#14280)
Signed-off-by: Michael Qiu <qiudayu.qdy@antgroup.com> Co-authored-by: Mike_Qiu <qiudayu.qdy@antgroup.com>
This commit is contained in:
@@ -47,6 +47,7 @@ from sglang.srt.distributed import (
|
|||||||
get_moe_expert_parallel_world_size,
|
get_moe_expert_parallel_world_size,
|
||||||
get_pp_group,
|
get_pp_group,
|
||||||
get_tensor_model_parallel_world_size,
|
get_tensor_model_parallel_world_size,
|
||||||
|
tensor_model_parallel_all_gather,
|
||||||
tensor_model_parallel_all_reduce,
|
tensor_model_parallel_all_reduce,
|
||||||
)
|
)
|
||||||
from sglang.srt.environ import envs
|
from sglang.srt.environ import envs
|
||||||
@@ -3127,6 +3128,10 @@ class DeepseekV2Model(nn.Module):
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
self.layers_to_capture = []
|
self.layers_to_capture = []
|
||||||
|
if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
|
||||||
|
self.enable_a2a_moe = True
|
||||||
|
else:
|
||||||
|
self.enable_a2a_moe = False
|
||||||
|
|
||||||
# llama_4_scaling: for supporting Mistral-Large-3 model
|
# llama_4_scaling: for supporting Mistral-Large-3 model
|
||||||
self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None)
|
self.llama_4_scaling_config = getattr(config, "llama_4_scaling", None)
|
||||||
@@ -3211,6 +3216,12 @@ class DeepseekV2Model(nn.Module):
|
|||||||
)
|
)
|
||||||
with ctx:
|
with ctx:
|
||||||
if i in self.layers_to_capture:
|
if i in self.layers_to_capture:
|
||||||
|
if self.enable_a2a_moe and i > self.first_k_dense_replace:
|
||||||
|
aux_hidden_state = tensor_model_parallel_all_gather(
|
||||||
|
hidden_states + residual, dim=0
|
||||||
|
)
|
||||||
|
aux_hidden_states.append(aux_hidden_state)
|
||||||
|
else:
|
||||||
aux_hidden_states.append(hidden_states + residual)
|
aux_hidden_states.append(hidden_states + residual)
|
||||||
layer = self.layers[i]
|
layer = self.layers[i]
|
||||||
hidden_states, residual = layer(
|
hidden_states, residual = layer(
|
||||||
|
|||||||
Reference in New Issue
Block a user