sdpa: require FP32 Stats and default an unset Stats dtype to FP32

[vedaanta]

Problem

The SDPA kernels always compute and store Stats (logsumexp) as FP32, 4 bytes per row. A graph that declared Stats with a narrower dtype — explicitly, or implicitly by leaving it unset with a non-FP32 io_data_type (fill_from_context assigned the io dtype) — built and executed fine, and the kernel then wrote FP32 rows past the end of the buffer the caller sized from the declared dtype. Depending on what the stray writes hit, this surfaced as silent corruption of adjacent allocations, illegal memory accesses, or CUDNN_STATUS_EXECUTION_FAILED_CUDA_DRIVER launch failures.

Root-caused from a customer workload: an fp16 Stats buffer (half the required bytes) was overwritten by batch ≥ 1 stats rows, corrupting whichever allocation happened to be adjacent — verified with compute-sanitizer (invalid f32 global writes starting exactly at byte offset batch_idx * s_q * 4, one byte past the fp16-sized allocation).

Fix

• Graph::sdpa_internal sets the Stats output dtype to FLOAT at creation, so no later fill can assign the io dtype.
• SDPANodeBase::infer_properties_node defaults a still-unset Stats dtype to FLOAT (covers direct node-construction paths).
• SDPANodeBase::post_validate_node rejects an explicitly non-FP32 Stats with GRAPH_NOT_SUPPORTED ("The Stats output of sdpa must be an FP32 tensor.") instead of corrupting memory at execute.
• The python IR layer (nodes.py) applies the same FP32 default for SDPA / SDPA_FP8 / SDPA_MXFP8 Stats before its generic io-dtype fill.

Testing

• New test/python/test_sdpa_stats_fp32_required.py: explicit fp16/bf16 Stats rejected at validate() (per io dtype), unset Stats inferred FP32, FP32 Stats accepted — 7/7 pass on B200 + cuDNN 9.26.0.33.
• Equivalent C++ checks exercised via a local graph.validate() harness: FLOAT OK, HALF rejected, unset inferred FLOAT.
• Regression: test_sdpa_fp32_rejected.py, test_sdpa_with_caching.py, test_sdpa_custom_features.py — 23/23 pass (including native-engine paths).
• The customer repro's full shape suite passes once its wrapper binds an FP32 stats buffer; with this change the previous misuse fails loudly at graph build instead.

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes

  • SDPA statistics outputs now consistently use FP32, regardless of the graph’s input or output data types.
  • Invalid non-FP32 statistics output configurations are now rejected during validation.
  • Unspecified statistics output types are automatically inferred as FP32.

• Tests

  • Added coverage for FP16, BF16, explicitly defined FP32, and invalid statistics output configurations.

[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: ff011c16-5a59-4b70-ae0b-ddd63b9ee545

📥 Commits

Reviewing files that changed from the base of the PR and between d53cafc and 26b4467.

📒 Files selected for processing (4)

• include/cudnn_frontend/graph_interface.h
• include/cudnn_frontend/node/scaled_dot_product_flash_attention.h
• python/cudnn/nodes.py
• test/python/test_sdpa_stats_fp32_required.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Walkthrough

Walkthrough

SDPA Stats outputs now use FP32 when generated or inferred. Validation rejects explicitly declared non-FP32 Stats types. Python regression tests cover invalid, inferred, and explicit FP32 cases.

Changes

SDPA Stats dtype handling

Layer / File(s)	Summary
C++ Stats dtype creation and validation include/cudnn_frontend/graph_interface.h, include/cudnn_frontend/node/scaled_dot_product_flash_attention.h	Generated and unset SDPA Stats outputs use FP32. Validation rejects explicitly declared non-FP32 Stats outputs.
Python inference and regression coverage python/cudnn/nodes.py, test/python/test_sdpa_stats_fp32_required.py	Python SDPA-family nodes infer FP32 for unset Stats outputs. Tests cover rejected FP16/BF16 types, inferred FP32, and explicit FP32 acceptance.

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

Merge Risk: ⚪ Minimal · up to 26b44

The PR prevents undersized SDPA Stats buffers by defaulting Stats to FP32 and rejecting incompatible declarations; no actionable merge-blocking risk remains beyond normal checks and review.

Suggested reviewers: anerudhan

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely summarizes the main change: requiring FP32 Stats and defaulting unset Stats dtypes to FP32.
Description check	✅ Passed	The description clearly explains the problem, fix, compatibility behavior, and testing results, although some template sections are omitted.
Docstring Coverage	✅ Passed	Docstring coverage is 85.71% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 7 functions across 4 files.
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.}

[vedaanta]

@cudnn-ci-bot run frost

[cudnn-ci-bot]

🏁 Pipeline finished

SHA: 26b4467
Targets: frost
Branch: cudnn-gh/pr-690-26b4467
Pipeline: 63888491
Last updated: 2026-08-21 16:57 UTC