test: align test_fmha_v2_prefill SM gating with is_sm12x_supported

[leonardHONG]

📌 Description

The function-level skip in test_fmha_v2_prefill.py checks is_sm120a_supported, which only matches sm_120, but the actual FMHAv2 dispatch (prefill.py:4498, :4540) is gated on is_sm12x_supported —so SM121 ends up skipped from these tests even though the kernel path supports it.

This swaps the test gates to is_sm12x_supported. The top-level skip message already said "SM12x" (only the gate was lagging). The three sub-skips (FP8 / SEPARATE_Q_K_V / SLIDING_WINDOW) previously said "SM120+"; they're updated to "SM12x" to match what the gate actually means.

sm_120 behavior is unchanged: both helpers evaluate True for sm_120 on CUDA ≥ 12.8 (is_sm120a_supported: major == 12 and minor==0;is_sm12x_supported: major == 12, with a CUDA-version branch that collapses to ≥ 12.8 on minor 0).

Verified on RTX Pro 6000 (sm_120, CUDA 12.9): 233 passed, 2112 skipped,0 failed.

🔍 Related Issues

Refs #3170 (Action Item 1 / T1).

🚀 Pull Request Checklist

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Summary by CodeRabbit

• Tests
  • Updated GPU support detection logic in attention tests to use updated SM version targeting for test gating and conditional skips.

[coderabbitai]

📝 Walkthrough

Walkthrough

The test file updates GPU support checks from SM120a-specific to SM12x-family support, changing from is_sm120a_supported to is_sm12x_supported and re-keying multiple skip conditions to reflect SM12x rather than SM120+ requirements.

Changes

Cohort / File(s)	Summary
Test GPU Support Gating tests/attention/test_fmha_v2_prefill.py	Replaced is_sm120a_supported with is_sm12x_supported in availability checks and re-keyed skip conditions for FP8 FMHA v2, warp-specialization restrictions, and sliding-window mask restrictions from SM120+ to SM12x; condensed one sliding-window skip into a single conditional statement.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Possibly related PRs

• Enable sm120f compilation #2650: Introduces and uses is_sm12x_supported support predicate in tests and utilities that this PR now relies on.
• fix: guard CUTLASS FMHA against SM12x and fix fmha_v2 SM121a check #2560: Updates SM12x-related FMHA v2 acceptance logic that aligns with this PR's test gating changes.
• [Spark unit test debugging] Fix for tests/gemm/test_groupwise_scaled_gemm_fp8.py #2751: Modifies the same is_sm12x_supported device-support predicate used in this PR's gating logic.

Suggested labels

run-ci, op: attention

Suggested reviewers

• yzh119
• aleozlx
• jimmyzho
• cyx-6

Poem

🐰 A GPU tale we shall retell,
From SM120a's fading spell,
To SM12x's broader way,
We test with checks both sound and gay!
One sliding window now more neat,
Our gating logic is complete!

🚥 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
Title check	✅ Passed	The title clearly and concisely describes the main change: aligning test SM gating with is_sm12x_supported.
Description check	✅ Passed	The description covers all required template sections: a detailed explanation of the changes, related issue reference, and completed pre-commit and test checklist items.
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

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

🧹 Nitpick comments (2)

tests/attention/test_fmha_v2_prefill.py (2)

492-502: Optional: hoist the repeated is_sm12x_supported check.

is_sm12x = is_sm12x_supported(torch.device("cuda")) is computed twice per test invocation (Line 492 and Line 758). Since the device is fixed for the run, you could compute it once at the top of run_trtllm_fmha_v2_prefill_case and reuse it, eliminating the duplicate compute-capability lookup. Not a correctness concern.

♻️ Suggested refactor

-    if not is_sm90a_supported(torch.device("cuda")) and not is_sm12x_supported(
-        torch.device("cuda")
-    ):
+    is_sm12x = is_sm12x_supported(torch.device("cuda"))
+    if not is_sm90a_supported(torch.device("cuda")) and not is_sm12x:
         pytest.skip("FMHA v2 requires SM90+ (Hopper) or SM12x GPUs.")

     # Skip invalid combinations
-    is_sm12x = is_sm12x_supported(torch.device("cuda"))
     if dtype == torch.float8_e4m3fn and is_sm12x:
@@
-        is_sm12x = is_sm12x_supported(torch.device("cuda"))
         if logits_soft_cap > 0:

Also applies to: 756-767

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/attention/test_fmha_v2_prefill.py` around lines 492 - 502, The test
recomputes is_sm12x_supported(torch.device("cuda")) multiple times; hoist this
call to a single local variable at the start of run_trtllm_fmha_v2_prefill_case
and replace the duplicate calls with that variable (e.g., keep the existing
is_sm12x variable and remove the second call), referencing is_sm12x_supported
and is_sm12x so all conditional skips use the precomputed is_sm12x value.

---

740-755: Doc comment mentions SM90/SM120 explicitly while logic now covers SM12x.

The block comment (Lines 742-755) describes formats as "SM90 (Hopper)" vs "SM120 (Blackwell) tiled", but the runtime branch at Line 760 keys off is_sm12x (which includes SM121+). Consider updating the wording to "SM12x" for consistency with the gate. Minor doc nit; behavior is correct assuming all SM12x variants share the same Softmax_saver semantics.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/attention/test_fmha_v2_prefill.py` around lines 740 - 755, Update the
doc comment to match the runtime branch naming: replace occurrences of "SM120
(Blackwell)" with "SM12x (Blackwell family)" (and any other explicit "SM120"
mentions) so the description aligns with the gating variable is_sm12x; keep the
explanation of how max is stored and how lse is computed for the two cases (SM90
vs SM12x) and ensure phrases like "SM12x tiled" or "SM12x (Blackwell)" are used
consistently to reflect all SM12x variants.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@tests/attention/test_fmha_v2_prefill.py`:
- Around line 492-502: The test recomputes
is_sm12x_supported(torch.device("cuda")) multiple times; hoist this call to a
single local variable at the start of run_trtllm_fmha_v2_prefill_case and
replace the duplicate calls with that variable (e.g., keep the existing is_sm12x
variable and remove the second call), referencing is_sm12x_supported and
is_sm12x so all conditional skips use the precomputed is_sm12x value.
- Around line 740-755: Update the doc comment to match the runtime branch
naming: replace occurrences of "SM120 (Blackwell)" with "SM12x (Blackwell
family)" (and any other explicit "SM120" mentions) so the description aligns
with the gating variable is_sm12x; keep the explanation of how max is stored and
how lse is computed for the two cases (SM90 vs SM12x) and ensure phrases like
"SM12x tiled" or "SM12x (Blackwell)" are used consistently to reflect all SM12x
variants.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 8ba88841-b2d9-4214-a141-79f3f243c6ae

📥 Commits

Reviewing files that changed from the base of the PR and between 5e1318c and e170b97.

📒 Files selected for processing (1)

• tests/attention/test_fmha_v2_prefill.py

[gemini-code-assist]

Code Review

This pull request refactors the GPU architecture checks in the FMHA v2 prefill tests by replacing the specific is_sm120a_supported function with the more general is_sm12x_supported. Correspondingly, variable names and test skip messages have been updated to reference SM12x instead of SM120+. I have no feedback to provide.

[kahyunnam]

/bot run

[flashinfer-bot]

GitLab MR !608 has been created, and the CI pipeline #49610133 is currently running. I'll report back once the pipeline job completes.

[kahyunnam]

lgtm, thanks!