[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:
kangwangamd
2026-06-17 23:39:44 -07:00
committed by GitHub
co-authored by kangwangamd
parent 3b61dc32c9
commit 6309fb9abb
+9 -8
View File
@@ -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
)