From 4d01e3714dfb6e3d8558a90ef8796c64588b84a5 Mon Sep 17 00:00:00 2001 From: Emil Gilliam Date: Tue, 30 Jun 2026 18:55:04 +0000 Subject: [PATCH] Add unified-engine FP8 and MXFP8 forward SDPA support Wire per-tensor FP8 and block-scaled MXFP8 (E8M0) forward attention through the unified SDPA runtime fusion engine: - scaled_dot_product_flash_attention.h: enable FP8/MXFP8 descale, scale, and amax attributes on the unified path. - sdpa_support_surface.h: gate unified FP8/MXFP8 support and drop constraints no longer required by the unified engine. - python bindings (pygraph.h, sdpa.cpp): expose the new descale/scale/amax inputs and outputs. - tests: extend fp8.py, mxfp8.py, and test_mhas_v2.py to cover the unified-engine path. --- .../node/scaled_dot_product_flash_attention.h | 47 +++++++++++++++++++ .../node/sdpa_support_surface.h | 31 ++++++++---- python/pygraph/pygraph.h | 7 ++- python/pygraph/sdpa.cpp | 18 +++++-- test/python/sdpa/fp8.py | 5 +- test/python/sdpa/mxfp8.py | 5 +- test/python/test_mhas_v2.py | 2 + 7 files changed, 98 insertions(+), 17 deletions(-) diff --git a/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h b/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h index 9195ac828..c3c37c56b 100644 --- a/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h +++ b/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h @@ -2641,6 +2641,53 @@ class UnifiedSDPANode : public SDPANodeBase { #endif } + // FP8 per-tensor scaling attributes (descale Q/K/V/S, scale S/O inputs; amax S/O outputs). + if (attributes.mma_core_mode == DataType_t::FP8_E4M3 || attributes.mma_core_mode == DataType_t::FP8_E5M2) { + auto fp8_cudnn_ver_error = + error_t{error_code_t::GRAPH_NOT_SUPPORTED, "FP8/MXFP8 for the unified SDPA node requires cuDNN 9.30.0"}; +#if (CUDNN_VERSION >= 93000) + NV_CUDNN_FE_DYNAMIC_CHECK_CUDNN_BACKEND_VERSION(93000, fp8_cudnn_ver_error); + + // Wire an optional FP8 scaling/descaling tensor (from either the input or output map) to + // its backend descriptor, no-op if the caller did not provide it. + auto set_fp8_desc = [&](auto const& tensor_map, auto name, cudnnBackendAttributeName_t attr) -> error_t { + auto it = tensor_map.find(name); + if (it != tensor_map.end() && it->second != nullptr) { + auto backend_desc = tensors[it->second->get_uid()]->get_desc()->get_backend_descriptor(); + _CUDNN_CHECK_CUDNN_ERROR(detail::set_attribute(unified_sdpa_operation->get_backend_descriptor(), + attr, + CUDNN_TYPE_BACKEND_DESCRIPTOR, + 1, + &backend_desc)); + } + return {error_code_t::OK, ""}; + }; + + CHECK_CUDNN_FRONTEND_ERROR(set_fp8_desc(attributes.inputs, + SDPA_attributes::input_names::Descale_Q, + CUDNN_ATTR_OPERATION_SDPA_FWD_DESCALE_QDESC)); + CHECK_CUDNN_FRONTEND_ERROR(set_fp8_desc(attributes.inputs, + SDPA_attributes::input_names::Descale_K, + CUDNN_ATTR_OPERATION_SDPA_FWD_DESCALE_KDESC)); + CHECK_CUDNN_FRONTEND_ERROR(set_fp8_desc(attributes.inputs, + SDPA_attributes::input_names::Descale_V, + CUDNN_ATTR_OPERATION_SDPA_FWD_DESCALE_VDESC)); + CHECK_CUDNN_FRONTEND_ERROR(set_fp8_desc(attributes.inputs, + SDPA_attributes::input_names::Descale_S, + CUDNN_ATTR_OPERATION_SDPA_FWD_DESCALE_SDESC)); + CHECK_CUDNN_FRONTEND_ERROR(set_fp8_desc( + attributes.inputs, SDPA_attributes::input_names::Scale_S, CUDNN_ATTR_OPERATION_SDPA_FWD_SCALE_SDESC)); + CHECK_CUDNN_FRONTEND_ERROR(set_fp8_desc( + attributes.inputs, SDPA_attributes::input_names::Scale_O, CUDNN_ATTR_OPERATION_SDPA_FWD_SCALE_ODESC)); + CHECK_CUDNN_FRONTEND_ERROR(set_fp8_desc( + attributes.outputs, SDPA_attributes::output_names::Amax_S, CUDNN_ATTR_OPERATION_SDPA_FWD_AMAX_SDESC)); + CHECK_CUDNN_FRONTEND_ERROR(set_fp8_desc( + attributes.outputs, SDPA_attributes::output_names::Amax_O, CUDNN_ATTR_OPERATION_SDPA_FWD_AMAX_ODESC)); +#else + return fp8_cudnn_ver_error; +#endif // CUDNN_VERSION >= 93000 + } + _CUDNN_CHECK_CUDNN_ERROR(detail::finalize(unified_sdpa_operation->get_backend_descriptor())); raw_operations.push_back(unified_sdpa_operation); diff --git a/include/cudnn_frontend/node/sdpa_support_surface.h b/include/cudnn_frontend/node/sdpa_support_surface.h index 05755a554..6b73f73e7 100644 --- a/include/cudnn_frontend/node/sdpa_support_surface.h +++ b/include/cudnn_frontend/node/sdpa_support_surface.h @@ -459,6 +459,16 @@ SDPA_attributes::verify_sdpa_support_surface_for_implementation(const detail::Co allowed_input_msg += ", CU_SEQ_LEN_Q, CU_SEQ_LEN_KV"; } + if (effective_cudnn_ver >= 93000) { + allowed_input_names.insert({input_names::Descale_Q, + input_names::Descale_K, + input_names::Descale_V, + input_names::Descale_S, + input_names::Scale_S, + input_names::Scale_O}); + allowed_input_msg += ", Descale_Q, Descale_K, Descale_V, Descale_S, Scale_S, Scale_O"; + } + for (const auto& [key, value] : inputs) { if (allowed_input_names.find(key) == allowed_input_names.end() && value != nullptr) { return {error_code_t::GRAPH_NOT_SUPPORTED, allowed_input_msg}; @@ -474,6 +484,11 @@ SDPA_attributes::verify_sdpa_support_surface_for_implementation(const detail::Co allowed_output_msg += ", RNG_DUMP, Max, Sum_exp"; } + if (effective_cudnn_ver >= 93000) { + allowed_output_names.insert({output_names::Amax_S, output_names::Amax_O}); + allowed_output_msg += ", Amax_S, Amax_O"; + } + for (const auto& [key, value] : outputs) { if (allowed_output_names.find(key) == allowed_output_names.end() && value != nullptr) { return {error_code_t::GRAPH_NOT_SUPPORTED, allowed_output_msg}; @@ -521,15 +536,13 @@ SDPA_attributes::verify_sdpa_support_surface_for_implementation(const detail::Co "Attention score modifier for unified SDPA node requires cuDNN 9.21.0 or above"}; } - if (mma_core_mode != DataType_t::HALF) { - return {error_code_t::GRAPH_NOT_SUPPORTED, - "Unified SDPA node doesn't yet support a data type other than fp16/bf16"}; - } - - if ((compute_data_type != DataType_t::NOT_SET && compute_data_type != DataType_t::FLOAT) || - context.get_compute_data_type() != DataType_t::FLOAT) { - return {error_code_t::GRAPH_NOT_SUPPORTED, - "Unified SDPA node doesn't yet support compute data type other than float"}; + if (mma_core_mode == DataType_t::FP8_E4M3 || mma_core_mode == DataType_t::FP8_E5M2) { + // Per-tensor FP8 and MXFP8 (block-scaled, E8M0 descales) are supported by the + // unified node starting from cuDNN 9.30.0. + if (effective_cudnn_ver < 93000) { + return {error_code_t::GRAPH_NOT_SUPPORTED, + "FP8/MXFP8 for the unified SDPA node requires cuDNN 9.30.0 or above"}; + } } } break; } diff --git a/python/pygraph/pygraph.h b/python/pygraph/pygraph.h index 2f6f2eea2..0eee32987 100644 --- a/python/pygraph/pygraph.h +++ b/python/pygraph/pygraph.h @@ -522,7 +522,9 @@ class PyGraph { py::object const& generate_stats, std::shared_ptr score_max, std::shared_ptr score_sum_exp, - std::shared_ptr sink_token); + std::shared_ptr sink_token, + bool const unfuse_fma, + cudnn_frontend::AttentionImplementation_t const& implementation); // MXFP8 SDPA forward - uses block-wise scale factors (E8M0 with F8_128x4 reordering) // return [o, stats, amax_o] @@ -543,7 +545,8 @@ class PyGraph { std::string const& name, py::object const& generate_stats, std::shared_ptr sink_token, - bool const unfuse_fma); + bool const unfuse_fma, + cudnn_frontend::AttentionImplementation_t const& implementation); // return [dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP] // dSink_token is an optional output set via set_dsink_token() attribute diff --git a/python/pygraph/sdpa.cpp b/python/pygraph/sdpa.cpp index e6278f575..dabad1782 100644 --- a/python/pygraph/sdpa.cpp +++ b/python/pygraph/sdpa.cpp @@ -540,7 +540,9 @@ PyGraph::sdpa_fp8(std::shared_ptr& q, py::object const& generate_stats, std::shared_ptr score_max, std::shared_ptr score_sum_exp, - std::shared_ptr sink_token) { + std::shared_ptr sink_token, + bool const unfuse_fma, + cudnn_frontend::AttentionImplementation_t const& implementation) { cudnn_frontend::DataType_t mma_core_mode = cudnn_frontend::DataType_t::FP8_E4M3; std::shared_ptr block_mask = nullptr; // cu_seq_len_q/cu_seq_len_kv are not exposed via the fp8 path. @@ -629,7 +631,9 @@ PyGraph::sdpa_fp8(std::shared_ptr& q, descale_v, descale_s, scale_s, - scale_o); + scale_o, + implementation, + unfuse_fma); // Return all 4 outputs as array for backward compatibility return {internal_result.O, internal_result.Stats, internal_result.Amax_S, internal_result.Amax_O}; @@ -653,7 +657,8 @@ PyGraph::sdpa_mxfp8(std::shared_ptr& q std::string const& name, py::object const& generate_stats, std::shared_ptr sink_token, - bool const unfuse_fma) { + bool const unfuse_fma, + cudnn_frontend::AttentionImplementation_t const& implementation) { auto attributes = cudnn_frontend::graph::SDPA_fp8_attributes().set_name(name).set_compute_data_type(compute_data_type); @@ -733,6 +738,7 @@ PyGraph::sdpa_mxfp8(std::shared_ptr& q attributes.set_sink_token(sink_token); } attributes.set_unfuse_fma(unfuse_fma); + attributes.set_implementation(implementation); // Call the MXFP8 6-parameter overload of sdpa_fp8 // This uses block scale factors (E8M0 + F8_128x4) instead of regular scalar descales @@ -1263,6 +1269,8 @@ init_pygraph_sdpa_submodule(py::class_& m) { py::arg_v("score_max", nullptr), py::arg_v("score_sum_exp", nullptr), py::arg_v("sink_token", nullptr), + py::arg_v("unfuse_fma", false), + py::arg_v("implementation", cudnn_frontend::AttentionImplementation_t::AUTO), R"pbdoc( Perform scaled dot product attention with fp8 datatype inputs and outputs. @@ -1294,6 +1302,8 @@ init_pygraph_sdpa_submodule(py::class_& m) { score_max (Optional[cudnn_tensor]): The max of attention score. score_sum_exp (Optional[cudnn_tensor]): The numerically stable sum of exponents using normalized values wrt max score. sink_token (Optional[cudnn_tensor]): Sink token bias for streaming attention. Default is None. + unfuse_fma (Optional[bool]): For SM100: use unfused __fmul_rn + __fadd_rn instead of ffma2 in softmax. Default is False. + implementation (Optional[cudnn.attention_implementation]): Which underlying implementation to use in the cuDNN backend. Default is AUTO (recommended). Preferred masking Args: diagonal_alignment (Optional[cudnn.diagonal_alignment]): One of {"TOP_LEFT", "BOTTOM_RIGHT"}. E.g., causal masking can be performed by setting diagonal_alignment=TOP_LEFT, and right_bound=0. Default is TOP_LEFT. left_bound (Optional[int]): An integer >= 1 specifying the offset to the left of the main diagonal to attend to. Default is None, implying +Inf. @@ -1330,6 +1340,7 @@ init_pygraph_sdpa_submodule(py::class_& m) { py::arg("generate_stats"), py::arg_v("sink_token", nullptr), py::arg_v("unfuse_fma", false), + py::arg_v("implementation", cudnn_frontend::AttentionImplementation_t::AUTO), R"pbdoc( Perform MXFP8 (Microscaling FP8) scaled dot product attention. @@ -1369,6 +1380,7 @@ init_pygraph_sdpa_submodule(py::class_& m) { generate_stats (bool): If true, compute and output softmax stats (required for training). sink_token (Optional[cudnn_tensor]): Sink token bias for streaming attention. Shape is (1, h_q, 1, 1), type is float32. Default is None. unfuse_fma (Optional[bool]): For SM100: use unfused __fmul_rn + __fadd_rn instead of ffma2 in softmax. Default is False. + implementation (Optional[cudnn.attention_implementation]): Which underlying implementation to use in the cuDNN backend. Default is AUTO (recommended). Returns: o (cudnn_tensor): The output data. diff --git a/test/python/sdpa/fp8.py b/test/python/sdpa/fp8.py index 92e954d44..0da4cc596 100644 --- a/test/python/sdpa/fp8.py +++ b/test/python/sdpa/fp8.py @@ -87,7 +87,7 @@ class GraphBwdUid(IntEnum): sink_token = 133 dSink_token = 134 -def generate_graph_fwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d_qk, d_vo, attn_scale, block_size, is_ragged=False, generate_stats=True, left_bound=None, right_bound=None, diag_align=None, with_sink_token=False): +def generate_graph_fwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d_qk, d_vo, attn_scale, block_size, is_ragged=False, generate_stats=True, left_bound=None, right_bound=None, diag_align=None, with_sink_token=False, implementation=cudnn.attention_implementation.AUTO): graph_fwd = cudnn.pygraph(io_data_type=cudnn_itype, intermediate_data_type=cudnn.data_type.FLOAT, compute_data_type=cudnn.data_type.FLOAT) use_padding_mask = None @@ -154,6 +154,7 @@ def generate_graph_fwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d paged_attention_max_seq_len_kv=s_kv, left_bound=left_bound, right_bound=right_bound, sink_token=sink_token, + implementation=implementation, ) # Only pass diagonal_alignment if it's not None (pybind11 doesn't accept None for enum types) if diag_align is not None: @@ -339,7 +340,7 @@ def exec_sdpa_fp8(cfg, request, cudnn_handle): # Build forward graph (always needed) try: - graph_fwd = generate_graph_fwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d_qk, d_vo, attn_scale, block_size, is_ragged=is_ragged, left_bound=left_bound, right_bound=right_bound, diag_align=diag_align, with_sink_token=with_sink_token) + graph_fwd = generate_graph_fwd(cudnn_itype, cudnn_otype, b, h_q, h_k, h_v, s_qo, s_kv, d_qk, d_vo, attn_scale, block_size, is_ragged=is_ragged, left_bound=left_bound, right_bound=right_bound, diag_align=diag_align, with_sink_token=with_sink_token, implementation=cfg.implementation) graph_fwd.validate() graph_fwd.build_operation_graph() graph_fwd.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) diff --git a/test/python/sdpa/mxfp8.py b/test/python/sdpa/mxfp8.py index 9284191ae..8064653e4 100644 --- a/test/python/sdpa/mxfp8.py +++ b/test/python/sdpa/mxfp8.py @@ -194,7 +194,8 @@ def generate_graph_fwd(b, h_q, h_k, h_v, cudnn_otype=cudnn.data_type.HALF, left_bound=None, right_bound=None, diag_align=None, with_sink_token=False, - with_unfuse_fma=False): + with_unfuse_fma=False, + implementation=cudnn.attention_implementation.AUTO): # Compute padded dimensions for F8_128x4 scale factors s_q_padded = ceil_div(s_qo, 128) * 128 s_kv_padded = ceil_div(s_kv, 128) * 128 @@ -280,6 +281,7 @@ def generate_graph_fwd(b, h_q, h_k, h_v, diagonal_band_right_bound=right_bound, sink_token=sink_token, unfuse_fma=with_unfuse_fma, + implementation=implementation, ) # Set output tensor properties @@ -541,6 +543,7 @@ def exec_sdpa_mxfp8(cfg, request, cudnn_handle): left_bound=left_bound, right_bound=right_bound, diag_align=diag_align, with_sink_token=with_sink_token, with_unfuse_fma=with_unfuse_fma, + implementation=cfg.implementation, ) graph_fwd.validate() graph_fwd.build_operation_graph() diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 191f205b1..625b88ae3 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -695,6 +695,7 @@ def test_sdpa_fp8_fwd_L0(env_info, test_no, request, cudnn_handle): with_sink_token=RandomChoice({True : 1, False : 2}), ) as randomization_ctx: test.cfg = randomization_ctx(rng, data_seed, geom_seed) + test.cfg.implementation = getattr(cudnn.attention_implementation, request.config.getoption("--implementation") or "", cudnn.attention_implementation.AUTO) test.showConfig(test_no, request) # Randomly enable unfuse_fma via environment variable for SM100 @@ -942,6 +943,7 @@ def test_sdpa_mxfp8_fwd_L0(env_info, test_no, request, cudnn_handle): test.cfg = randomization_ctx(rng, data_seed, geom_seed) test.cfg.is_mxfp8 = True + test.cfg.implementation = getattr(cudnn.attention_implementation, request.config.getoption("--implementation") or "", cudnn.attention_implementation.AUTO) # Randomly enable unfuse_fma via environment variable for SM100 unfuse_fma = rng.choice([True, False])