[AMD] fix: do not emit a shared-expert marker twice on the per-rank slot path (#36515)
This commit is contained in:
@@ -2313,6 +2313,18 @@ def select_experts(
|
|||||||
# DeepSeek V2/V3/R1 series models use grouped_top_k
|
# DeepSeek V2/V3/R1 series models use grouped_top_k
|
||||||
# remove num_fused_shared_experts from grouped_topk/biased_grouped_topk
|
# remove num_fused_shared_experts from grouped_topk/biased_grouped_topk
|
||||||
num_routed_topk = top_k - num_fused_shared_experts
|
num_routed_topk = top_k - num_fused_shared_experts
|
||||||
|
# On the per-rank shared-slot path the shared expert is appended afterwards by
|
||||||
|
# fused_append_remap_shared_experts_deepep, so the gate must not emit a shared
|
||||||
|
# marker as well. Doing both costs one routed expert (the gate is asked for
|
||||||
|
# K_routed = top_k - num_fused_shared_experts and then spends one of those
|
||||||
|
# slots on the marker) and places that marker at id num_experts, which the
|
||||||
|
# DeepEP remap shifts one past the end of the expert space -- 384 -> 392 for
|
||||||
|
# 384 routed experts on EP8, where the valid ids are 0..391.
|
||||||
|
num_fused_shared_experts_for_gate = (
|
||||||
|
0
|
||||||
|
if has_per_rank_fused_shared_slots(num_fused_shared_experts)
|
||||||
|
else num_fused_shared_experts
|
||||||
|
)
|
||||||
if use_grouped_topk:
|
if use_grouped_topk:
|
||||||
assert topk_group is not None
|
assert topk_group is not None
|
||||||
assert num_expert_group is not None
|
assert num_expert_group is not None
|
||||||
@@ -2374,7 +2386,7 @@ def select_experts(
|
|||||||
topk=num_routed_topk if _use_aiter else top_k,
|
topk=num_routed_topk if _use_aiter else top_k,
|
||||||
renormalize=renormalize,
|
renormalize=renormalize,
|
||||||
scoring_func=scoring_func,
|
scoring_func=scoring_func,
|
||||||
num_fused_shared_experts=num_fused_shared_experts,
|
num_fused_shared_experts=num_fused_shared_experts_for_gate,
|
||||||
routed_scaling_factor=routed_scaling_factor,
|
routed_scaling_factor=routed_scaling_factor,
|
||||||
num_token_non_padded=num_token_non_padded,
|
num_token_non_padded=num_token_non_padded,
|
||||||
expert_location_dispatch_info=expert_location_dispatch_info,
|
expert_location_dispatch_info=expert_location_dispatch_info,
|
||||||
@@ -2437,7 +2449,7 @@ def select_experts(
|
|||||||
renormalize=renormalize,
|
renormalize=renormalize,
|
||||||
correction_bias=correction_bias,
|
correction_bias=correction_bias,
|
||||||
scoring_func=scoring_func,
|
scoring_func=scoring_func,
|
||||||
num_fused_shared_experts=num_fused_shared_experts,
|
num_fused_shared_experts=num_fused_shared_experts_for_gate,
|
||||||
routed_scaling_factor=routed_scaling_factor,
|
routed_scaling_factor=routed_scaling_factor,
|
||||||
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
||||||
**_fused_topk_kwargs,
|
**_fused_topk_kwargs,
|
||||||
|
|||||||
Reference in New Issue
Block a user