Reject FP32 I/O in the unified SDPA node - #449
Merged
Merged
Conversation
The unified SDPA node has no FP32 I/O kernel. mma_core_mode is auto-set to HALF in Graph::sdpa() regardless of the I/O dtype, so an FP32 graph would build and then dispatch a half-precision kernel onto 4-byte data, reading and writing shared memory out of bounds (see NVIDIA#424). Reject unsupported Q/K/V/O I/O dtypes in the UNIFIED case of verify_sdpa_support_surface_for_implementation so such graphs fail cleanly at build time with GRAPH_NOT_SUPPORTED. NOT_SET is allowed to pass (the dtype is resolved from the graph default before real validation re-runs this check), and FP8 is permitted for the sdpa_fp8 path. Under AUTO, auto-select now routes FP32 graphs to the composite implementation, which does support FP32 I/O. Signed-off-by: Emil Gilliam <egilliam@nvidia.com>
egilliam-nv
added
cat-bug
1 task
Add a python regression test that an FP32-I/O SDPA graph forced onto the UNIFIED implementation fails cleanly with cudnnGraphNotSupportedError instead of crashing, that the guard does not over-reject FP16/BF16, and that under AUTO an FP32 graph is routed to the composite implementation. Verified on A100 (sm80) and H100 (sm90, composite FP32 via forward compat). Signed-off-by: Emil Gilliam <egilliam@nvidia.com>
Contributor
📝 WalkthroughWalkthroughSDPA unified support validation now rejects unsupported Q/K/V/O data types, including FP32, while allowing supported FP16, BF16, and FP8 types. Python regression tests verify explicit unified rejection, supported unified types, and automatic FP32 composite routing. ChangesSDPA dtype validation
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
test/python/test_sdpa_fp32_rejected.py (1)
64-67: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover FP8 and gate supported configurations.
FP8 E4M3/E5M2 are part of the supported unified contract but are untested. Parameterize supported dtypes with architecture/backend capability checks so BF16/FP8-unavailable environments skip rather than fail.
As per coding guidelines, “Gate tests on supported capabilities and skip unsupported architecture, dtype, or backend-version combinations using support checks,
cudnn.backend_version(), andtorch.cuda.get_device_capability().”🤖 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_sdpa_fp32_rejected.py` around lines 64 - 67, Extend test_fp16_bf16_unified_supported to cover FP8 E4M3 and E5M2, and gate each supported dtype with the appropriate architecture/backend capability checks. Use the established support checks, cudnn.backend_version(), and torch.cuda.get_device_capability() so BF16 or FP8 unsupported environments skip the test while supported configurations still exercise _build_sdpa with UNIFIED.Source: Coding guidelines
🤖 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_sdpa_fp32_rejected.py`:
- Around line 70-78: Update test_fp32_auto_routes_to_composite to distinguish
unified-routing rejection from genuine composite-engine unavailability: assert
the error does not indicate unified-node support failure before skipping. Only
skip when the exception specifically reflects that no composite FP32 engine
exists, and fail the test for any unified-routing error.
- Around line 26-55: Update _build_sdpa and its callers to construct four
variants where exactly one of Q, K, V, or O uses FP32 and every other I/O tensor
uses the supported dtype. Preserve the shared graph setup while allowing O’s
dtype to be configured independently, and assert each variant is rejected
specifically during the intended validation stage.
---
Nitpick comments:
In `@test/python/test_sdpa_fp32_rejected.py`:
- Around line 64-67: Extend test_fp16_bf16_unified_supported to cover FP8 E4M3
and E5M2, and gate each supported dtype with the appropriate
architecture/backend capability checks. Use the established support checks,
cudnn.backend_version(), and torch.cuda.get_device_capability() so BF16 or FP8
unsupported environments skip the test while supported configurations still
exercise _build_sdpa with UNIFIED.
🪄 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: ad6133d3-54be-47f0-8ffb-157c98130628
📒 Files selected for processing (2)
include/cudnn_frontend/node/sdpa_support_surface.htest/python/test_sdpa_fp32_rejected.py
Comment on lines
+26
to
+55
| def _build_sdpa(io_dtype, implementation): | ||
| b, h, s, d = 2, 4, 128, 64 | ||
| cudnn_dtype = _DTYPE[io_dtype] | ||
| graph = cudnn.pygraph( | ||
| io_data_type=cudnn_dtype, | ||
| intermediate_data_type=cudnn.data_type.FLOAT, | ||
| compute_data_type=cudnn.data_type.FLOAT, | ||
| ) | ||
| dim = [b, h, s, d] | ||
| stride = [h * s * d, s * d, d, 1] | ||
| q = graph.tensor(name="q", dim=dim, stride=stride, data_type=cudnn_dtype) | ||
| k = graph.tensor(name="k", dim=dim, stride=stride, data_type=cudnn_dtype) | ||
| v = graph.tensor(name="v", dim=dim, stride=stride, data_type=cudnn_dtype) | ||
| o, _ = graph.sdpa( | ||
| name="sdpa", | ||
| q=q, | ||
| k=k, | ||
| v=v, | ||
| generate_stats=False, | ||
| attn_scale=1.0 / (d**0.5), | ||
| implementation=implementation, | ||
| ) | ||
| o.set_output(True).set_dim(dim).set_stride(stride) | ||
|
|
||
| graph.validate() | ||
| graph.build_operation_graph() | ||
| graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) | ||
| graph.check_support() | ||
| graph.build_plans() | ||
| return graph |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Test each rejected I/O port independently.
All four tensors use the same FP32 dtype, so the forced-unified test passes if only Q is checked. Build variants with exactly one of Q, K, V, or O as FP32 and the remaining I/O as a supported dtype; also assert the failure occurs at the intended validation stage.
Also applies to: 58-61
🤖 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_sdpa_fp32_rejected.py` around lines 26 - 55, Update
_build_sdpa and its callers to construct four variants where exactly one of Q,
K, V, or O uses FP32 and every other I/O tensor uses the supported dtype.
Preserve the shared graph setup while allowing O’s dtype to be configured
independently, and assert each variant is rejected specifically during the
intended validation stage.
Comment on lines
+70
to
+78
| def test_fp32_auto_routes_to_composite(): | ||
| """Under AUTO, FP32 must not be routed to the unified node; it builds via | ||
| the composite implementation where a real FP32 kernel exists.""" | ||
| try: | ||
| _build_sdpa(torch.float32, cudnn.attention_implementation.AUTO) | ||
| except cudnn.cudnnGraphNotSupportedError: | ||
| # No composite FP32 engine available on this architecture; the important | ||
| # invariant (FP32 is not silently sent to the unified node) still holds. | ||
| pytest.skip("no composite FP32 SDPA engine available on this GPU") |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Do not skip a unified-routing regression as a missing composite engine.
A wrongly routed AUTO FP32 graph raises the same cudnnGraphNotSupportedError and is currently skipped. Distinguish the unified rejection—e.g., by asserting the unified-support error is absent before skipping a genuine composite-engine unavailability.
🤖 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_sdpa_fp32_rejected.py` around lines 70 - 78, Update
test_fp32_auto_routes_to_composite to distinguish unified-routing rejection from
genuine composite-engine unavailability: assert the error does not indicate
unified-node support failure before skipping. Only skip when the exception
specifically reflects that no composite FP32 engine exists, and fail the test
for any unified-routing error.
Anerudhan
approved these changes
Jul 29, 2026
egilliam-nv
added a commit
to egilliam-nv/cudnn-frontend
that referenced
this pull request
Jul 30, 2026
Address CodeRabbit review comments on NVIDIA#449 and NVIDIA#454: - Exercise each Q/K/V/O port independently: override exactly one port with an unsupported dtype (rest FP16) so a check that only looked at Q would be caught. Assert via pytest.raises(match=...) that the rejection is our unified/composite dtype check rather than an incidental one (the intended validation stage). - Don't mask a unified-misroute as a missing composite engine: in the FP32 composite/AUTO test, fail if the rejection came from the unified node; skip only on genuine composite-engine unavailability. - Capability-guard the positive FP16/BF16 unified test: skip when unified SDPA is unsupported on this cuDNN/GPU combo, but still fail if our dtype guard wrongly rejects FP16/BF16. Not addressing the FP8-coverage nitpick: FP8 needs the sdpa_fp8 path with descale/scale tensors, out of scope for this dtype-rejection test. Verified on A100 (sm80) and H100 (sm90): 16 passed. Signed-off-by: Emil Gilliam <egilliam@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Reject unsupported Q/K/V/O I/O data types in the
UNIFIEDcase ofverify_sdpa_support_surface_for_implementation, so a graph built viagraph.sdpa(...)with FP32 I/O fails cleanly at build time withGRAPH_NOT_SUPPORTEDinstead of crashing at execution.Fixes #424.
Why
The unified SDPA node has no FP32 I/O kernel.
mma_core_modeis auto-set toHALFinGraph::sdpa()regardless of the I/O dtype, and the support surface only ever keyed onmma_core_mode— never on the raw I/O dtype. So an FP32 graph passed validation and then dispatched a half-precision kernel onto 4-byte data, reading/writing shared memory out of bounds (the CUDA "invalid memory reference" in #424).Behavior
NOT_SETis allowed to pass (the dtype is resolved from the graph default before real validation re-runs this check).sdpa_fp8path).AttentionImplementation_t::AUTO, auto-select now routes FP32 graphs to the composite implementation, which does support FP32 I/O — so FP32 attention keeps working where a real kernel exists; only an explicitset_implementation(UNIFIED)gets the graceful error.Testing
test/python/test_sdpa_fp32_rejected.py(L0): FP32+UNIFIED now fails cleanly, FP16/BF16+UNIFIED still build, and FP32+AUTO routes to composite. Verified on A100 (sm80) and H100 (sm90).Notes
A companion change will also be made in the cuDNN backend so that the backend likewise rejects FP32 on the unified engine and falls back to composite.
🤖 Generated with Claude Code