Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
github-actions
Bot
changed the title
Treat padded routing IDs as unrouted before expert-indexed accesses to prevent out-of-bounds writes in Opus and FlyDSL sorting.
5b61b6d to
e9a15cf
Compare
amd-weisun
left a comment
amd-weisun
left a comment
There was a problem hiding this comment.
LGTM, can you merge 387-388 into single line should_store, same for line 1035(valid_eid) and 1030(valid), and use valid/valid_eid whichever you prefer to replace line 1038 if valid & valid_eid
Thanks. Updated the code |
|
Apparently, the approval was dismissed after new commit was pushed (even if only style was changed). Could you approve again @coderfeli ? Thanks |
|
Or @amd-weisun could you review/approve? Thanks |
|
@coderfeli could you still have a look? Apparently an approval from team_aiter is required. This was approved already by @amd-weisun |
…2 histogram The fused router was pinned to a 4096 model dim by two constexpr in the entry: BlockSize 256 with TD 16, checked as cols == BlockSize * TD. The kernel body was never dim-specific -- it carries no static_assert, derives NWAVE from BlockSize and strides every loop by it -- so the pin was in the entry alone. TD is the quant's per-thread vector width and stays fixed; the block size now follows the model dim, so one thread still covers one vector of the row in a single pass. Instantiated dims are 4096 (256 threads) and 6144 (384, GLM), and a dim outside that set is refused rather than run on a mismatched block. Both are whole numbers of waves and satisfy E <= 2*BlockSize. 6144 needs 192 scale bytes per row against the existing 256-byte allowance, so the caller-owned workspace is unchanged. Selection is by shape alone -- no model name, no path registry, no environment threshold -- so GLM is reached the same way 4096 is. Phase 2's histogram also takes ROCm#5295's invalid-expert-id guard. Within the fused launch the ids are the kernel's own, but kPhase23 re-reads topk_ids from global memory, where a masked -1 would index an LDS counter out of bounds. GLM's gate needs no new code: at n_group == topk_group == 1 the group mask is all-ones, so biased grouped top-k is flat biased top-k, which is what phase 1 already computes. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
3 participants
Summary
topk_ids == -1is used as a padded/no-route sentinel, but several Opus and FlyDSL MoE sorting paths used the raw value as an expert mesh index. Negative IDs could write before the LDS/HBM mesh, and Opus local expert remapping could also read beforelocal_expert_mask, causing a GPU memory-access fault.This change consistently requires
0 <= eid < num_expertsbefore the remaining expert-indexed accesses. Valid routes continue through the same sorting logic, while invalid routes are skipped.This is a follow-up to ROCm/aiter#4839.
AI assistance was used for this PR.
Changes
vLLM workload impact
DeepSeek-V4 DSpark can produce all-padding routing during startup memory profiling. Fixed-K DSpark TP2/TP4 exposed this at M=16384 in Opus P0_v1; that exact site was fixed by ROCm/aiter#4839. Adaptive DSpark verification exposed additional unguarded sorting paths with an all-empty M=2048 fused-MoE input.
Together with ROCm/aiter#4839, this PR makes invalid
-1routing handling consistent across the remaining Opus and FlyDSL sorting paths. The fixed-K TP2/TP4 DSV4 Dspark startup configurations and the adaptive TP8 DSV4 Dspark startup configuration all completed successfully with the combined fixes.Validation
op_tests/test_moe_sorting.py: passed across Opus, CK, and FlyDSL, including CUDA graph capture/replay.