Test softmax stats outputs in various permutations

[egilliam-nv]

Ensure that various layouts of softmax stats (and other softmax outputs) are tested, when in non-ragged mode.

Previously we didn't test this due to https://nvbugspro.nvidia.com/bug/6057616 , which apparently has already been fixed. This PR provides additional confirmation that the bug has been fixed and will help prevent it from reappearing in the future.

Summary by CodeRabbit

• Tests
  • Extended randomization logic in SDPA testing to provide more consistent and comprehensive coverage across configuration variants.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 2794f05b-33ed-4c31-bd91-7f8f532948e2

📥 Commits

Reviewing files that changed from the base of the PR and between 00a0b29 and e951831.

📒 Files selected for processing (1)

• test/python/sdpa/random_config.py

---

📝 Walkthrough

Walkthrough

In RandomizationContext, the non-ragged stride setup now generates stride_stats using randomized gaps and indexed stride computation, matching the pattern used for stride_q and stride_o. Previously, stride_stats was derived from a fixed "bhsd" layout without gap randomization.

Changes

Stride Stats Randomization

Layer / File(s)	Summary
Randomized stride stats computation test/python/sdpa/random_config.py	Introduces gaps_stats with 50% randomization chance and aligned padding, then computes stride_stats via get_strides_from_indices using the permuted indices and randomized gaps, consolidating stride randomization logic across stride_q, stride_o, and stride_stats.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Poem

A rabbit hops through strides so fine,
Where gaps now dance and indices align,
No more fixed layouts, dull and plain,
Randomization flows like morning rain! 🐰✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'Test softmax stats outputs in various permutations' is clearly related to the changeset, which adds permutation testing for softmax stats layouts in non-ragged mode.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[egilliam-nv]

@cudnn-ci-bot run

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-304-e951831
Pipeline: 54491145

[egilliam-nv]

These tests fail with the given changes, so apparently the underlying bug still hasn't been fully fixed:
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test1] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test3] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test5] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test7] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test8] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test11] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test15] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test16] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test17] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test18] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test19] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test20] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test22] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test24] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test30] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test31] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test33] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test37] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test38] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test43] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test44] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test45] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test46] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test47] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test48] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test49] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test53] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test55] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test58] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test59] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test60] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test62] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test64] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test65] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test67] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test74] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test75] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test76] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test77] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test79] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test80] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test82] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test85] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test87] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test88] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test89] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test90] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test92] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test93] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test94] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test95] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test96] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test102] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test103] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test105] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test108] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test110] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test111] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test112] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test113] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test117] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test118] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test123] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test124] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test126] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test127] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test129] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test132] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test136] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test137] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test138] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test141] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test142] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test145] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test146] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test149] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test150] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test152] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test153] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test154] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test155] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test157] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test159] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test160] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test162] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test163] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test165] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test167] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test168] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test169] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test173] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test177] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test181] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test182] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test184] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test188] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test189] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test190] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test194] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test198] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test199] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test200] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test201] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test202] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test204] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test205] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test206] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test207] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test210] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test214] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test217] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test220] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test221] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test225] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test227] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test228] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test229] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test234] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test235] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test236] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test239] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test244] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test246] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test248] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test249] - Failed: determinism check failed
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test250] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test251] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_L0[test254] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test1] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test3] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test4] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test6] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test7] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test8] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test9] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test14] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test18] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test22] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test23] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test25] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test26] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test28] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test29] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test31] - Failed: disallowed mismatches
FAILED test/python/test_mhas_v2.py::test_sdpa_random_bwd_bias_L0[test32] - Failed: disallowed mismatches

Marking this as draft again.