fix: EPLB dispatch OOB when shared experts fusion enabled under DeepEP (#22525)
This commit is contained in:
@@ -1017,9 +1017,21 @@ def _post_process_topk_ids(
|
|||||||
topk_ids=topk_ids,
|
topk_ids=topk_ids,
|
||||||
)
|
)
|
||||||
if _is_cuda:
|
if _is_cuda:
|
||||||
topk_ids = _biased_grouped_topk_postprocess(
|
# When shared experts are fused (appended as extra columns in topk_ids),
|
||||||
topk_ids, expert_location_dispatch_info, num_token_non_padded
|
# EPLB dispatch must only remap the routed expert columns.
|
||||||
)
|
# The shared expert column (value = n_routed_experts) would be out-of-bounds
|
||||||
|
# for the logical-to-physical dispatch table.
|
||||||
|
if num_fused_shared_experts > 0 and is_deepep_class_backend():
|
||||||
|
shared_cols = topk_ids[:, -num_fused_shared_experts:]
|
||||||
|
routed_cols = topk_ids[:, :-num_fused_shared_experts]
|
||||||
|
routed_cols = _biased_grouped_topk_postprocess(
|
||||||
|
routed_cols, expert_location_dispatch_info, num_token_non_padded
|
||||||
|
)
|
||||||
|
topk_ids = torch.cat([routed_cols, shared_cols], dim=-1)
|
||||||
|
else:
|
||||||
|
topk_ids = _biased_grouped_topk_postprocess(
|
||||||
|
topk_ids, expert_location_dispatch_info, num_token_non_padded
|
||||||
|
)
|
||||||
|
|
||||||
if num_fused_shared_experts > 0 and _use_aiter:
|
if num_fused_shared_experts > 0 and _use_aiter:
|
||||||
M, N = router_logits.shape
|
M, N = router_logits.shape
|
||||||
|
|||||||
Reference in New Issue
Block a user