diff --git a/include/cudnn_frontend/node/sdpa_support_surface.h b/include/cudnn_frontend/node/sdpa_support_surface.h index 9d7cddd5d..5fde69780 100644 --- a/include/cudnn_frontend/node/sdpa_support_surface.h +++ b/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 allowed_input_names{ input_names::Q, input_names::K, input_names::V, input_names::Attn_scale}; diff --git a/test/python/test_sdpa_fp32_rejected.py b/test/python/test_sdpa_fp32_rejected.py new file mode 100644 index 000000000..d2b04a319 --- /dev/null +++ b/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")