[AMD] Fix Always mask padded topk_ids on HIP to prevent garbage MoE routing (DeepSeek-R1-MXFP4 accuracy regression) (#28378)
Co-authored-by: kangwangamd <kangwangamd@users.noreply.github.com>
This commit is contained in:
co-authored by
kangwangamd
parent
3b61dc32c9
commit
6309fb9abb
@@ -1569,17 +1569,18 @@ def _post_process_topk_ids(
|
||||
topk_ids, expert_location_dispatch_info, num_token_non_padded
|
||||
)
|
||||
elif _is_hip:
|
||||
# On AMD HIP the aiter MoE kernels do not handle topk_ids=-1 safely
|
||||
# (negative indices cause illegal memory access). Always fill the padded
|
||||
# region with 0 so every kernel sees a valid in-range expert id.
|
||||
# Routing weights for padded tokens are zeroed below so their
|
||||
# contribution to the hidden state is still zero regardless of the id.
|
||||
# Regression: skipping this mask when EPLB is disabled caused garbage
|
||||
# MoE routing for models like DeepSeek-R1-MXFP4 (accuracy ~0.09 vs 0.94+).
|
||||
_mask_topk_ids_padded_region(topk_ids, num_token_non_padded, fill_value=0)
|
||||
# The logical->physical remap is only meaningful when a real
|
||||
# expert-location mapping exists. With a trivial placement and EPLB off
|
||||
# the map is identity, and indexing it here is both unnecessary and not
|
||||
# well-defined for the dp-attention padded region of topk_ids; skip it.
|
||||
# ``--ep-dispatch-algorithm fake`` is a routing benchmark artifact and
|
||||
# does not by itself require remapping.
|
||||
# the map is identity so the remap can be skipped safely.
|
||||
if _eplb_remap_enabled():
|
||||
# Mask the padded region to a valid in-range id (0) before the
|
||||
# gather so it never indexes the dispatch map out of bounds. -1 is
|
||||
# not used here: the aiter MoE kernels do not handle topk_ids=-1.
|
||||
_mask_topk_ids_padded_region(topk_ids, num_token_non_padded, fill_value=0)
|
||||
topk_ids = topk_ids_logical_to_physical(
|
||||
topk_ids, expert_location_dispatch_info
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user