test_mhas_v2: extend bwd random head-dim coverage to d=256

[vedaanta]

What

Extends the randomized SDPA backward test coverage in test/python/test_mhas_v2.py to head dims up to 256:

• The three bwd suites (test_sdpa_random_bwd_L0, test_sdpa_random_bwd_ragged_L0, test_sdpa_random_bwd_bias_L0) already listed (256,256) in with_high_probability, but their random-draw branch was still capped at d_qk_max=192, d_v_max=128 — so backward configs with 128 < d <= 256 outside the fixed high-probability shapes were never generated. This raises the caps to 256/256, matching the fwd suites.
• Bumps test_sdpa_random_bwd_L0 from 256 to 384 seeds to keep sampling density over the larger shape space.

Why

A head_dim=256 dense backward plan-build failure recently escaped to PyTorch integration testing (test_cudnn_attention_dense_d256_sm90_sm10x) because the FE test suites never exercised d=256 backward: v2's random ranges were capped as above, and v1's test_sdpa_backward samples d only from [32, 56, 64, 128]. A cuDNN backend fix is in flight; this change closes the test-coverage side.

Validation

• With this change test_sdpa_random_bwd_L0 draws d_qk=d_v=256 in 66/384 configs and d>128 in 206/384 (previously 0 outside the three fixed shapes).
• Sampled (256,256) configs verified end-to-end (build + execute + reference and determinism checks) against a cuDNN build with d=256 bprop support; on backends without support the configs are waived by the existing cudnnGraphNotSupportedError handling, so no new failures are introduced on older cuDNN versions.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Tests
  • Expanded cuDNN SDPA attention backward test coverage.
  • Increased generated test scenarios from 256 to 384.
  • Added coverage for larger query/key and value hidden dimensions, up to 256.

[coderabbitai]

📝 Walkthrough

Walkthrough

The cuDNN SDPA backward tests now generate 384 cases instead of 256 and sample d_qk and d_v up to 256 across standard, ragged, and bias scenarios.

Changes

SDPA backward test coverage

Layer / File(s)	Summary
Expand backward test generation and dimensions test/python/test_mhas_v2.py	L0 backward test generation increases from 256 to 384 cases, while standard, ragged, and bias randomization contexts expand d_qk and d_v maximums to 256.

Estimated code review effort: 2 (Simple) | ~8 minutes

Suggested reviewers: anerudhan

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly states the test file and the main change: expanding randomized backward head-dim coverage to 256.
Description check	✅ Passed	The description covers the change, rationale, and validation, but it skips several template sections like affected area, related issues, and exact testing commands.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@test/python/test_mhas_v2.py`:
- Line 192: Reformat the three modified RandomHiddenDimSize(...) calls,
including the occurrences near lines 192, 481, and 654, with Black using a
160-character line length. Keep all arguments and values unchanged while
applying consistent Black formatting.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: a8d880d5-77fe-4677-8d61-b43f72f07fda

📥 Commits

Reviewing files that changed from the base of the PR and between fbc7136 and 5ade219.

📒 Files selected for processing (1)

• test/python/test_mhas_v2.py

[coderabbitai]

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Format the changed calls with Black’s 160-character limit.

These modified RandomHiddenDimSize(...) calls exceed the configured line length. Reformat the arguments consistently in all three locations.

Proposed formatting

-        d_qk_d_v=RandomHiddenDimSize(d_qk_min=1, d_qk_max=256, d_v_min=1, d_v_max=256, head_dim_distribution={"d_qk=d_v":5, "d_qk=random":1}, with_high_probability=[(64,64), (128,128), (192,128), (256,256)]),
+        d_qk_d_v=RandomHiddenDimSize(
+            d_qk_min=1,
+            d_qk_max=256,
+            d_v_min=1,
+            d_v_max=256,
+            head_dim_distribution={"d_qk=d_v": 5, "d_qk=random": 1},
+            with_high_probability=[(64, 64), (128, 128), (192, 128), (256, 256)],
+        ),

Apply the same formatting to Lines 481 and 654.

As per coding guidelines, **/*.{py,ipynb} files must be formatted with Black using a line length of 160.

Also applies to: 481-481, 654-654

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@test/python/test_mhas_v2.py` at line 192, Reformat the three modified
RandomHiddenDimSize(...) calls, including the occurrences near lines 192, 481,
and 654, with Black using a 160-character line length. Keep all arguments and
values unchanged while applying consistent Black formatting.

Source: Coding guidelines

[vedaanta]

@cudnn-ci-bot run

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-425-5ade219
Pipeline: 59274452

[Anerudhan]

Please fix the failing pipeline