Reject FP32 I/O in the unified SDPA node

include/cudnn_frontend/node/sdpa_support_surface.h

@@ -399,6 +399,10 @@ SDPA_attributes::verify_sdpa_support_surface_for_implementation(const detail::Co
         auto const it = inputs.find(name);
         return it != inputs.end() && it->second != nullptr;
     };
+    auto const has_output = [this](output_names name) {
+        auto const it = outputs.find(name);
+        return it != outputs.end() && it->second != nullptr;
+    };
 
     switch (impl) {
         case AttentionImplementation_t::AUTO:
@@ -436,6 +440,27 @@ SDPA_attributes::verify_sdpa_support_surface_for_implementation(const detail::Co
                                            error_code_t::GRAPH_NOT_SUPPORTED,
                                            "Unified SDPA node doesn't yet support dynamic shape");
 
+            // The unified engine only implements FP16/BF16/FP8/MXFP8 inputs/outputs.
+            // NOT_SET is allowed to pass (the dtype is resolved from the graph default before
+            // real validation re-runs this check).
+            auto const io_dtype_supported = [](DataType_t dt) {
+                return dt == DataType_t::NOT_SET || dt == DataType_t::HALF || dt == DataType_t::BFLOAT16 ||
+                       dt == DataType_t::FP8_E4M3 || dt == DataType_t::FP8_E5M2;
+            };
+            bool bad_io = false;
+            for (auto const name : {input_names::Q, input_names::K, input_names::V}) {
+                if (has_input(name) && !io_dtype_supported(inputs.at(name)->get_data_type())) {
+                    bad_io = true;
+                }
+            }
+            if (has_output(output_names::O) && !io_dtype_supported(outputs.at(output_names::O)->get_data_type())) {
+                bad_io = true;
+            }
+            RETURN_CUDNN_FRONTEND_ERROR_IF(bad_io,
+                                           error_code_t::GRAPH_NOT_SUPPORTED,
+                                           "Unified SDPA node only supports FP16/BF16/FP8 Q/K/V/O I/O (FP32 I/O is not "
+                                           "supported); use the composite implementation for FP32.");
+
             // TODO: Provide smarter error messages that provide the required cuDNN version for each input.
             std::unordered_set<SDPA_attributes::input_names> allowed_input_names{
                 input_names::Q, input_names::K, input_names::V, input_names::Attn_scale};

test/python/test_sdpa_fp32_rejected.py

@@ -0,0 +1,78 @@
+"""Regression test for https://github.com/NVIDIA/cudnn-frontend/issues/424.
+
+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
+previously build and then dispatch a half-precision kernel onto 4-byte data,
+reading/writing shared memory out of bounds (a CUDA invalid memory reference).
+
+The support surface now rejects unsupported unified Q/K/V/O I/O dtypes, so such a
+graph fails cleanly at build time with GRAPH_NOT_SUPPORTED. Under AUTO, FP32 is
+routed to the composite implementation, which does support FP32 I/O.
+"""
+
+import cudnn
+import pytest
+import torch
+
+pytestmark = pytest.mark.L0
+
+_DTYPE = {
+    torch.float32: cudnn.data_type.FLOAT,
+    torch.float16: cudnn.data_type.HALF,
+    torch.bfloat16: cudnn.data_type.BFLOAT16,
+}
+
+
+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
+
+
+def test_fp32_unified_rejected():
+    """FP32 forced onto the unified node fails cleanly instead of crashing."""
+    with pytest.raises(cudnn.cudnnGraphNotSupportedError):
+        _build_sdpa(torch.float32, cudnn.attention_implementation.UNIFIED)
+
+
+@pytest.mark.parametrize("io_dtype", [torch.float16, torch.bfloat16])
+def test_fp16_bf16_unified_supported(io_dtype):
+    """The FP32 guard must not over-reject the supported FP16/BF16 I/O dtypes."""
+    _build_sdpa(io_dtype, cudnn.attention_implementation.UNIFIED)
+
+
+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")