Skip to content

[ROCm][CI] Stabilize the sqrt-softplus top-k tie oracle - #54403

Draft
AndreasKaratzas wants to merge 1 commit into
mainfrom
akaratza_topk_softplus_tie_break
Draft

[ROCm][CI] Stabilize the sqrt-softplus top-k tie oracle#54403
AndreasKaratzas wants to merge 1 commit into
mainfrom
akaratza_topk_softplus_tie_break

Conversation

@AndreasKaratzas

Copy link
Copy Markdown
Member
  • Replace the reference torch.topk selection with stable descending argsort.
  • Match the fused kernel's deterministic lower-expert-ID tie-break.
  • Add a focused unit test for an exact tie at the top-k boundary.
  • Keep production routing and kernel behavior unchanged.
  • Validate the change with the complete parameter matrix and repeated fresh-process runs.

The MI355 MoE Shard 3 group in AMD CI build 12440 reported one mismatched expert in the BF16 384-expert case. Formal bisection found no responsible PR in the nightly window because the relevant source blobs were identical at f956e1c343dc63720b50c01209bd5ada6d8589a2 and the failing tip, while the exact node also passed repeatedly. The failure was an oracle ambiguity: torch.topk does not specify which equal-scored index wins, whereas the fused kernel intentionally selects the lower expert ID. Making the reference ordering explicit removes that nondeterminism without changing production behavior.

@mergify mergify Bot added the rocm Related to AMD ROCm label Aug 30, 2026
@github-project-automation github-project-automation Bot moved this to Todo in AMD Aug 30, 2026
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
@AndreasKaratzas
AndreasKaratzas force-pushed the akaratza_topk_softplus_tie_break branch from f82a3aa to 7aae9a2 Compare August 30, 2026 08:13
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

rocm Related to AMD ROCm

Projects

Status: Todo

Development

Successfully merging this pull request may close these issues.

1 participant