From 6309fb9abb65724ff40a2aa393df17dab3aa11db Mon Sep 17 00:00:00 2001 From: kangwangamd <100359556+kangwangamd@users.noreply.github.com> Date: Thu, 18 Jun 2026 14:39:44 +0800 Subject: [PATCH] [AMD] Fix Always mask padded topk_ids on HIP to prevent garbage MoE routing (DeepSeek-R1-MXFP4 accuracy regression) (#28378) Co-authored-by: kangwangamd --- python/sglang/srt/layers/moe/topk.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 37964b346..dfc25064c 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -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 )