[Test] Handle tied top-k indices in graph-pool logprob regression

[ch-wan]

Motivation

test_input_logprobs_survive_replay_with_growing_chunks can fail when random float32 logprobs tie. The implementation computes topk(max_k) and truncates each sequence to its requested k, while the test computes topk(k) independently and requires identical token IDs. PyTorch does not guarantee which tied indices topk returns.

This caused the base-b failure on #40638. It also reproduces without that PR on public main commit 00986c81be687d82e66651fb3ce9a096070b06e6: seed 44 fails at rows=2600, choosing token 1783 versus 959 at the same logprob, -5.749078750610352. The failed test and its logprob, pool, and result-processing code remain identical on main at e332e1b84e643d02d6223542941030c03717f62c.

Modifications

Validate returned token IDs by gathering their expected scores and comparing those scores exactly against the expected top-k values. Require distinct IDs within each row and explicitly exercise tied scores in both sequences. All graph replay, pinned-copy, sampled-logit, and token-logprob checks remain in place. Only the existing test changes.

Accuracy Tests

On one H200 with PyTorch 2.13.0+cu130, using the public-main baseline above:

• Before the fix: the original test fails with seed 44.
• After the fix: the same seed passes, including the added explicit ties.
• Full test/registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py: 18 tests passed.
• Pre-commit checks passed, including registered-test validation.

No model output or inference behavior changes.

Speed Tests and Profiling

Not applicable; test-only change.

---

CI States

Latest PR Test (Base): ❌ Run #35677282078
Latest PR Test (Extra): ❌ Run #35677281936
Latest PR Test (AMD ROCm 10): ❌ Run #35677282031

[ch-wan]

/rerun-test test_graph_pool_borrow.py

[github-actions]

Results for /rerun-test test_graph_pool_borrow.py:

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py