Skip to content

fp8 support in unified fprop #301

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
egilliam-nv merged 1 commit into from
Jun 30, 2026
Merged

fp8 support in unified fprop #301

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 47 additions & 0 deletions include/cudnn_frontend/node/scaled_dot_product_flash_attention.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2641,6 +2641,53 @@ class UnifiedSDPANode : public SDPANodeBase<UnifiedSDPANode> {
#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
}

Comment thread
coderabbitai[bot] marked this conversation as resolved.
_CUDNN_CHECK_CUDNN_ERROR(detail::finalize(unified_sdpa_operation->get_backend_descriptor()));

raw_operations.push_back(unified_sdpa_operation);
Expand Down
31 changes: 22 additions & 9 deletions include/cudnn_frontend/node/sdpa_support_surface.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -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};
Expand Down Expand Up @@ -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;
}
Expand Down
7 changes: 5 additions & 2 deletions python/pygraph/pygraph.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -522,7 +522,9 @@ class PyGraph {
py::object const& generate_stats,
std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> score_max,
std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> score_sum_exp,
std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> sink_token);
std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> 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]
Expand All @@ -543,7 +545,8 @@ class PyGraph {
std::string const& name,
py::object const& generate_stats,
std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> 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
Expand Down
18 changes: 15 additions & 3 deletions python/pygraph/sdpa.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -540,7 +540,9 @@ PyGraph::sdpa_fp8(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& q,
py::object const& generate_stats,
std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> score_max,
std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> score_sum_exp,
std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> sink_token) {
std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> 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<cudnn_frontend::graph::Tensor_attributes> block_mask = nullptr;
// cu_seq_len_q/cu_seq_len_kv are not exposed via the fp8 path.
Expand Down Expand Up @@ -629,7 +631,9 @@ PyGraph::sdpa_fp8(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& 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};
Expand All @@ -653,7 +657,8 @@ PyGraph::sdpa_mxfp8(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& q
std::string const& name,
py::object const& generate_stats,
std::shared_ptr<cudnn_frontend::graph::Tensor_attributes> 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);

Expand Down Expand Up @@ -733,6 +738,7 @@ PyGraph::sdpa_mxfp8(std::shared_ptr<cudnn_frontend::graph::Tensor_attributes>& 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
Expand Down Expand Up @@ -1263,6 +1269,8 @@ init_pygraph_sdpa_submodule(py::class_<PyGraph>& 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.

Expand Down Expand Up @@ -1294,6 +1302,8 @@ init_pygraph_sdpa_submodule(py::class_<PyGraph>& 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.
Expand Down Expand Up @@ -1330,6 +1340,7 @@ init_pygraph_sdpa_submodule(py::class_<PyGraph>& 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.

Expand Down Expand Up @@ -1369,6 +1380,7 @@ init_pygraph_sdpa_submodule(py::class_<PyGraph>& 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.
Expand Down
5 changes: 3 additions & 2 deletions test/python/sdpa/fp8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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])
Expand Down
5 changes: 4 additions & 1 deletion test/python/sdpa/mxfp8.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down
2 changes: 2 additions & 0 deletions test/python/test_mhas_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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])
Expand Down