[ROCm][CI] Stabilize the sqrt-softplus top-k tie oracle (#54403)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Andreas Karatzas committed
Sep 1, 2026 at 00:34 UTC
40b2f62061575905aaac8bc360eaea62a4baeb67
1 file changed
+19
-1
tests/kernels/moe/test_topk_softplus_sqrt.py
+19
-1
@@ -38,7 +38,12 @@ def _torch_topk_softplus_sqrt(
38
assert input_ids is not None
39
topk_ids = hash_indices_table[input_ids.long()]
40
else:
41
- topk_ids = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=True)[1]
41
+ # Match the fused kernel's deterministic tie-break: lower expert ids
42
+ # win when scores are equal. torch.topk does not guarantee which tied
43
+ # index it selects at the k-th boundary.
44
+ topk_ids = torch.argsort(
45
+ scores_for_choice, dim=-1, descending=True, stable=True
46
+ )[:, :topk]
47
48
topk_weights = original_scores.gather(1, topk_ids.long())
49
if renormalize:
@@ -48,6 +53,19 @@ def _torch_topk_softplus_sqrt(
53
return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
54
55
56
+def test_torch_topk_softplus_sqrt_breaks_ties_by_expert_id():
57
+ gating_output = torch.tensor([[2.0, 1.0, 1.0, 0.0]])
58
+
59
+ _, topk_ids = _torch_topk_softplus_sqrt(
60
+ gating_output,
61
+ topk=2,
62
+ renormalize=False,
63
+ routed_scaling_factor=1.0,
64
+ )
65
+
66
+ assert topk_ids.tolist() == [[0, 1]]
67
+
68
+
69
def test_sqrtsoftplus_bias_uses_deepseek_v4_routing_method():
70
assert (
71
get_routing_method_type(