From 71b4779980d61ead0d6889bbc480a6bd3609602e Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Mon, 13 Jul 2026 08:09:55 -0700 Subject: [PATCH 1/9] Restore SDPA repro tensor dumps --- include/cudnn_frontend/graph_interface.h | 20 +++++++++ .../node/scaled_dot_product_flash_attention.h | 42 +++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/include/cudnn_frontend/graph_interface.h b/include/cudnn_frontend/graph_interface.h index 0f6ba3456..2d28b928c 100644 --- a/include/cudnn_frontend/graph_interface.h +++ b/include/cudnn_frontend/graph_interface.h @@ -120,6 +120,25 @@ class Graph : public ICudnn, public INode { return {error_code_t::OK, ""}; } + error_t + log_tensors_to_dump_(cudnnHandle_t handle, + std::vector const &tensor_uids, + void *const *tensor_ptrs) const { + if (!isLoggingTensorDumpEnabled()) { + return {error_code_t::OK, ""}; + } + + std::unordered_map tensor_uid_to_pointer_map; + tensor_uid_to_pointer_map.reserve(tensor_uids.size()); + for (size_t i = 0; i < tensor_uids.size(); i++) { + if (tensor_ptrs[i] != nullptr) { + tensor_uid_to_pointer_map.emplace(tensor_uids[i], tensor_ptrs[i]); + } + } + + return log_tensors_to_dump_(handle, tensor_uid_to_pointer_map); + } + error_t pre_validate_node() const override final { RETURN_CUDNN_FRONTEND_ERROR_IF( @@ -1425,6 +1444,7 @@ class Graph : public ICudnn, public INode { CHECK_CUDNN_FRONTEND_ERROR(run_auxiliary_kernels(handle, workspace, cached_workspace_modifications)); CUDNN_FE_LOG_LABEL_ENDL("INFO: Executing graph_uid " << graph_uid); + CHECK_CUDNN_FRONTEND_ERROR(log_tensors_to_dump_(handle, varpack_template.all_uids, ptrs)); // 5. Dispatch void *engine_workspace = static_cast(workspace) + fe_workspace_size; 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 55e67e910..e46dfde4e 100644 --- a/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h +++ b/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h @@ -2182,6 +2182,48 @@ class CompositeSDPABackwardNode : public NodeCRTP { return {error_code_t::OK, ""}; } + error_t + collect_tensors_to_dump_node( + std::vector, char>>& tensors_to_dump) const override final { + std::unordered_set seen_uids; + auto add_tensor = [&tensors_to_dump, &seen_uids](std::shared_ptr const& tensor) { + if (tensor != nullptr && seen_uids.insert(tensor->get_uid()).second) { + tensors_to_dump.emplace_back(tensor, 'd'); + } + }; + auto add_input = [&](input_names name) { + auto it = attributes.inputs.find(name); + if (it != attributes.inputs.end()) { + add_tensor(it->second); + } + }; + auto add_input_offset = [&](input_names name) { + auto it = attributes.inputs.find(name); + if (it != attributes.inputs.end() && it->second != nullptr) { + add_tensor(it->second->get_ragged_offset()); + } + }; + auto add_output_offset = [&](output_names name) { + auto it = attributes.outputs.find(name); + if (it != attributes.outputs.end() && it->second != nullptr) { + add_tensor(it->second->get_ragged_offset()); + } + }; + + add_input(input_names::SEQ_LEN_Q); + add_input(input_names::SEQ_LEN_KV); + + for (auto name : + {input_names::Q, input_names::K, input_names::V, input_names::O, input_names::dO, input_names::Stats}) { + add_input_offset(name); + } + for (auto name : {output_names::dQ, output_names::dK, output_names::dV}) { + add_output_offset(name); + } + + return {error_code_t::OK, ""}; + } + #ifndef CUDNN_FRONTEND_SKIP_JSON_LIB virtual void serialize(json& j) const override final { From 7a367aee31c4d39176fb30be958beae0a6f2dff8 Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Wed, 10 Jun 2026 18:45:41 -0700 Subject: [PATCH 2/9] Skip FP8 reference checks in perf mode --- test/python/sdpa/fp8.py | 158 +++++++++++++++++++++------------------- 1 file changed, 85 insertions(+), 73 deletions(-) diff --git a/test/python/sdpa/fp8.py b/test/python/sdpa/fp8.py index c9802845d..500ef4aaf 100644 --- a/test/python/sdpa/fp8.py +++ b/test/python/sdpa/fp8.py @@ -311,6 +311,7 @@ def create_paged_container_and_block_table(tensor, block_size): def exec_sdpa_fp8(cfg, request, cudnn_handle): if request.config.option.dryrun: pytest.skip("dryrun") + perf = request.config.getoption("--perf") cudnn_version = LooseVersion(cudnn.backend_version_string()) if cudnn_version < "9.14.0": @@ -420,12 +421,15 @@ def exec_sdpa_fp8(cfg, request, cudnn_handle): else: padding = None - o_ref, stats_ref, o_amax = compute_ref(q_fp8, k_fp8, v_fp8, attn_scale=attn_scale, - q_descale=q_descale_gpu, k_descale=k_descale_gpu, v_descale=v_descale_gpu, - s_scale=s_scale_gpu, s_descale=s_descale_gpu, torch_itype=torch_itype, - torch_otype=torch_otype, padding=padding, - left_bound=left_bound, right_bound=right_bound, diag_align=diag_align, - sink_token=sink_token_gpu, rescale_threshold=rescale_threshold) + if perf: + o_amax = 1.0 + else: + o_ref, stats_ref, o_amax = compute_ref(q_fp8, k_fp8, v_fp8, attn_scale=attn_scale, + q_descale=q_descale_gpu, k_descale=k_descale_gpu, v_descale=v_descale_gpu, + s_scale=s_scale_gpu, s_descale=s_descale_gpu, torch_itype=torch_itype, + torch_otype=torch_otype, padding=padding, + left_bound=left_bound, right_bound=right_bound, diag_align=diag_align, + sink_token=sink_token_gpu, rescale_threshold=rescale_threshold) o_scale_gpu = torch.tensor([get_fp8_scale_factor(o_amax, torch_otype)], dtype=torch.float, device="cuda") @@ -497,7 +501,7 @@ def exec_sdpa_fp8(cfg, request, cudnn_handle): variant_pack[int(GraphFwdUid.sink_token)] = sink_token_gpu workspace = torch.empty(graph_fwd.get_workspace_size(), dtype=torch.uint8, device="cuda") - if request.config.getoption("--perf"): + if perf: times_ms = time_execution(graph_fwd.execute, variant_pack, workspace, cudnn_handle) print(f"@@@@ FP8 Fwd graph_fwd.execute avg_time_ms={times_ms.mean().item():.3f}") profile_execution(graph_fwd.execute, variant_pack, workspace, cudnn_handle) @@ -505,21 +509,22 @@ def exec_sdpa_fp8(cfg, request, cudnn_handle): torch.cuda.synchronize() # Compare forward output - if is_ragged: - o_ref_comp = convert_uniform_to_packed(torch.einsum("bshd->bhsd", o_ref), seq_len_q_ref, max_t_q) - else: - o_ref_comp = o_ref + if not perf: + if is_ragged: + o_ref_comp = convert_uniform_to_packed(torch.einsum("bshd->bhsd", o_ref), seq_len_q_ref, max_t_q) + else: + o_ref_comp = o_ref - o_gpu_float = o_gpu.detach().float() * get_fp8_descale_factor(o_amax, torch_otype) - o_ref_float = o_ref_comp.detach().float() * get_fp8_descale_factor(o_amax, torch_otype) + o_gpu_float = o_gpu.detach().float() * get_fp8_descale_factor(o_amax, torch_otype) + o_ref_float = o_ref_comp.detach().float() * get_fp8_descale_factor(o_amax, torch_otype) - if is_ragged: - t_idx = sum(seq_len_q_list) - o_gpu_float[t_idx:] = 0 - o_ref_float[t_idx:] = 0 + if is_ragged: + t_idx = sum(seq_len_q_list) + o_gpu_float[t_idx:] = 0 + o_ref_float[t_idx:] = 0 - atol, rtol = 0.08, 0.2 - torch.testing.assert_close(o_gpu_float, o_ref_float, atol=atol, rtol=rtol) + atol, rtol = 0.08, 0.2 + torch.testing.assert_close(o_gpu_float, o_ref_float, atol=atol, rtol=rtol) # Backward pass if not cfg.is_infer: @@ -530,30 +535,36 @@ def exec_sdpa_fp8(cfg, request, cudnn_handle): o_descale_gpu = torch.tensor([get_fp8_descale_factor(o_amax, torch_otype)], dtype=torch.float, device="cuda") dO_descale_gpu = torch.tensor([get_fp8_descale_factor(dO_amax, torch_itype)], dtype=torch.float, device="cuda") - # Get unpacked BSHD references for backward - if is_ragged: - q_ref_bwd = torch.einsum("bhsd->bshd", convert_packed_to_uniform(q_gpu, seq_len_q_ref, s_qo)) - k_ref_bwd = torch.einsum("bhsd->bshd", convert_packed_to_uniform(k_gpu, seq_len_kv_ref, s_kv)) - v_ref_bwd = torch.einsum("bhsd->bshd", convert_packed_to_uniform(v_gpu, seq_len_kv_ref, s_kv)) - o_ref_bwd = torch.einsum("bhsd->bshd", convert_packed_to_uniform(o_gpu, seq_len_q_ref, s_qo)) - dO_ref_bwd = dO_fp8 + if perf: + dP_amax = 1.0 + dQ_amax = 1.0 + dK_amax = 1.0 + dV_amax = 1.0 else: - q_ref_bwd = q_gpu - k_ref_bwd = k_gpu - v_ref_bwd = v_gpu - o_ref_bwd = o_gpu - dO_ref_bwd = dO_fp8 - - padding_bwd = (seq_len_q_ref, seq_len_kv_ref) if is_ragged else None - dQ_ref, dK_ref, dV_ref, dSink_token_ref, dP_amax, dQ_amax, dK_amax, dV_amax = compute_ref_backward( - q_ref_bwd, k_ref_bwd, v_ref_bwd, o_ref_bwd, dO_ref_bwd, attn_scale=attn_scale, - q_descale=q_descale_gpu, k_descale=k_descale_gpu, v_descale=v_descale_gpu, - s_scale=s_scale_gpu, s_descale=s_descale_gpu, torch_itype=torch_itype, - o_descale=o_descale_gpu, dO_descale=dO_descale_gpu, - torch_otype=torch_otype, padding=padding_bwd, - left_bound=left_bound, right_bound=right_bound, diag_align=diag_align, - sink_token=sink_token_gpu - ) + # Get unpacked BSHD references for backward + if is_ragged: + q_ref_bwd = torch.einsum("bhsd->bshd", convert_packed_to_uniform(q_gpu, seq_len_q_ref, s_qo)) + k_ref_bwd = torch.einsum("bhsd->bshd", convert_packed_to_uniform(k_gpu, seq_len_kv_ref, s_kv)) + v_ref_bwd = torch.einsum("bhsd->bshd", convert_packed_to_uniform(v_gpu, seq_len_kv_ref, s_kv)) + o_ref_bwd = torch.einsum("bhsd->bshd", convert_packed_to_uniform(o_gpu, seq_len_q_ref, s_qo)) + dO_ref_bwd = dO_fp8 + else: + q_ref_bwd = q_gpu + k_ref_bwd = k_gpu + v_ref_bwd = v_gpu + o_ref_bwd = o_gpu + dO_ref_bwd = dO_fp8 + + padding_bwd = (seq_len_q_ref, seq_len_kv_ref) if is_ragged else None + dQ_ref, dK_ref, dV_ref, dSink_token_ref, dP_amax, dQ_amax, dK_amax, dV_amax = compute_ref_backward( + q_ref_bwd, k_ref_bwd, v_ref_bwd, o_ref_bwd, dO_ref_bwd, attn_scale=attn_scale, + q_descale=q_descale_gpu, k_descale=k_descale_gpu, v_descale=v_descale_gpu, + s_scale=s_scale_gpu, s_descale=s_descale_gpu, torch_itype=torch_itype, + o_descale=o_descale_gpu, dO_descale=dO_descale_gpu, + torch_otype=torch_otype, padding=padding_bwd, + left_bound=left_bound, right_bound=right_bound, diag_align=diag_align, + sink_token=sink_token_gpu + ) dP_descale_gpu = torch.tensor([get_fp8_descale_factor(dP_amax, torch_itype)], dtype=torch.float, device="cuda") dQ_scale_gpu = torch.tensor([get_fp8_scale_factor(dQ_amax, torch_otype)], dtype=torch.float, device="cuda") @@ -623,7 +634,7 @@ def exec_sdpa_fp8(cfg, request, cudnn_handle): variant_pack_bwd[int(GraphBwdUid.dSink_token)] = dSink_token_gpu workspace_bwd = torch.empty(graph_bwd.get_workspace_size(), dtype=torch.uint8, device="cuda") - if request.config.getoption("--perf"): + if perf: times_ms = time_execution(graph_bwd.execute, variant_pack_bwd, workspace_bwd, cudnn_handle) print(f"@@@@ FP8 Bwd graph.execute avg_time_ms={times_ms.mean().item():.3f}") profile_execution(graph_bwd.execute, variant_pack_bwd, workspace_bwd, cudnn_handle) @@ -652,36 +663,37 @@ def exec_sdpa_fp8(cfg, request, cudnn_handle): pytest.fail("determinism check failed", pytrace=False) print("@@@@ Determinism check: PASSED, dQ, dK, dV bitwise match between runs.") - if is_ragged: - dQ_ref = convert_uniform_to_packed(torch.einsum("bshd->bhsd", dQ_ref), seq_len_q_ref, max_t_q) - dK_ref = convert_uniform_to_packed(torch.einsum("bshd->bhsd", dK_ref), seq_len_kv_ref, max_t_kv) - dV_ref = convert_uniform_to_packed(torch.einsum("bshd->bhsd", dV_ref), seq_len_kv_ref, max_t_kv) - - dQ_out = dQ_gpu.detach().float() * get_fp8_descale_factor(dQ_amax, torch_otype) - dK_out = dK_gpu.detach().float() * get_fp8_descale_factor(dK_amax, torch_otype) - dV_out = dV_gpu.detach().float() * get_fp8_descale_factor(dV_amax, torch_otype) - - dQ_ref_float = dQ_ref.detach().float() * get_fp8_descale_factor(dQ_amax, torch_otype) - dK_ref_float = dK_ref.detach().float() * get_fp8_descale_factor(dK_amax, torch_otype) - dV_ref_float = dV_ref.detach().float() * get_fp8_descale_factor(dV_amax, torch_otype) - - if is_ragged: - t_idx_q = sum(seq_len_q_list) - dQ_out[t_idx_q:] = 0 - dQ_ref_float[t_idx_q:] = 0 - t_idx_kv = sum(seq_len_kv_list) - dK_out[t_idx_kv:] = 0 - dK_ref_float[t_idx_kv:] = 0 - dV_out[t_idx_kv:] = 0 - dV_ref_float[t_idx_kv:] = 0 - - atol, rtol = 0.04, 0.2 - torch.testing.assert_close(dQ_out, dQ_ref_float, atol=atol, rtol=rtol) - torch.testing.assert_close(dK_out, dK_ref_float, atol=atol, rtol=rtol) - torch.testing.assert_close(dV_out, dV_ref_float, atol=atol, rtol=rtol) - - if with_sink_token: - torch.testing.assert_close(dSink_token_gpu, dSink_token_ref, atol=0.02, rtol=0.2) + if not perf: + if is_ragged: + dQ_ref = convert_uniform_to_packed(torch.einsum("bshd->bhsd", dQ_ref), seq_len_q_ref, max_t_q) + dK_ref = convert_uniform_to_packed(torch.einsum("bshd->bhsd", dK_ref), seq_len_kv_ref, max_t_kv) + dV_ref = convert_uniform_to_packed(torch.einsum("bshd->bhsd", dV_ref), seq_len_kv_ref, max_t_kv) + + dQ_out = dQ_gpu.detach().float() * get_fp8_descale_factor(dQ_amax, torch_otype) + dK_out = dK_gpu.detach().float() * get_fp8_descale_factor(dK_amax, torch_otype) + dV_out = dV_gpu.detach().float() * get_fp8_descale_factor(dV_amax, torch_otype) + + dQ_ref_float = dQ_ref.detach().float() * get_fp8_descale_factor(dQ_amax, torch_otype) + dK_ref_float = dK_ref.detach().float() * get_fp8_descale_factor(dK_amax, torch_otype) + dV_ref_float = dV_ref.detach().float() * get_fp8_descale_factor(dV_amax, torch_otype) + + if is_ragged: + t_idx_q = sum(seq_len_q_list) + dQ_out[t_idx_q:] = 0 + dQ_ref_float[t_idx_q:] = 0 + t_idx_kv = sum(seq_len_kv_list) + dK_out[t_idx_kv:] = 0 + dK_ref_float[t_idx_kv:] = 0 + dV_out[t_idx_kv:] = 0 + dV_ref_float[t_idx_kv:] = 0 + + atol, rtol = 0.04, 0.2 + torch.testing.assert_close(dQ_out, dQ_ref_float, atol=atol, rtol=rtol) + torch.testing.assert_close(dK_out, dK_ref_float, atol=atol, rtol=rtol) + torch.testing.assert_close(dV_out, dV_ref_float, atol=atol, rtol=rtol) + + if with_sink_token: + torch.testing.assert_close(dSink_token_gpu, dSink_token_ref, atol=0.02, rtol=0.2) # Print hash and stats for determinism verification print_tensor_stats(o_gpu, tag="o_gpu") From f4e08f532fa921d0b679299ef864f280ea2a3be1 Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Mon, 29 Jun 2026 15:10:04 -0700 Subject: [PATCH 3/9] Fix large MXFP8 performance repros --- test/python/sdpa/mxfp8.py | 180 ++++++++++++++++++++---------------- test/python/test_mhas_v2.py | 33 +++++++ 2 files changed, 135 insertions(+), 78 deletions(-) diff --git a/test/python/sdpa/mxfp8.py b/test/python/sdpa/mxfp8.py index 8064653e4..c5fd42b29 100644 --- a/test/python/sdpa/mxfp8.py +++ b/test/python/sdpa/mxfp8.py @@ -7,7 +7,7 @@ if LooseVersion(transformer_engine.__version__) < LooseVersion("2.12.0"): raise ImportError(f"TransformerEngine >= 2.12.0 required, found {transformer_engine.__version__}") - from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor import transformer_engine_torch as tex HAS_TE = True @@ -15,6 +15,7 @@ HAS_TE = False tex = None MXFP8Quantizer = None + MXFP8Tensor = None import cudnn import pytest @@ -31,6 +32,29 @@ def ceil_div(a: int, b: int) -> int: return (a + b - 1) // b +MXFP8_QUANTIZER_MAX_ROWS = 65535 * 64 + + +def quantize_mxfp8_compact(tensor_2d, quantizer): + if tensor_2d.shape[0] <= MXFP8_QUANTIZER_MAX_ROWS: + return quantizer(tensor_2d) + + chunk_rows = MXFP8_QUANTIZER_MAX_ROWS // 128 * 128 + results = [quantizer(chunk) for chunk in tensor_2d.split(chunk_rows, dim=0)] + return MXFP8Tensor( + shape=tensor_2d.shape, + dtype=tensor_2d.dtype, + fp8_dtype=quantizer.dtype, + rowwise_data=torch.cat([result._rowwise_data for result in results], dim=0), + rowwise_scale_inv=torch.cat([result._rowwise_scale_inv for result in results], dim=0), + columnwise_data=torch.cat([result._columnwise_data for result in results], dim=0), + columnwise_scale_inv=torch.cat([result._columnwise_scale_inv for result in results], dim=0), + quantizer=quantizer, + requires_grad=False, + with_gemm_swizzled_scales=False, + ) + + class GraphFwdUid(IntEnum): q = 0 k = 1 @@ -134,7 +158,16 @@ def compute_mxfp8_scale_dims(s, d, block_size=32): -def quantize_to_mxfp8(tensor, b, h, s, d, block_size=32, fp8_dtype=torch.float8_e4m3fn): +def quantize_to_mxfp8( + tensor, + b, + h, + s, + d, + block_size=32, + fp8_dtype=torch.float8_e4m3fn, + with_ref=True, +): l = b * h te_dtype = tex.DType.kFloat8E4M3 if fp8_dtype == torch.float8_e4m3fn else tex.DType.kFloat8E5M2 @@ -155,34 +188,36 @@ def quantize_to_mxfp8(tensor, b, h, s, d, block_size=32, fp8_dtype=torch.float8_ # without swizzle quantizer = MXFP8Quantizer(fp8_dtype=te_dtype, rowwise=True, columnwise=True) - quantizer_swizzle = quantizer.copy() - mxfp8_result = quantizer(tensor_2d) + mxfp8_result = quantize_mxfp8_compact(tensor_2d, quantizer) # --- Rowwise results (quantized along D dimension) --- fp8_data_d_flat = mxfp8_result._rowwise_data fp8_data_d = fp8_data_d_flat.reshape(l, s_padded, d_padded)[:, :s, :d].contiguous() fp8_data_d = fp8_data_d.view(fp8_dtype).reshape(b, h, s, d) - scale_inv_d = mxfp8_result._rowwise_scale_inv - scale_inv_d_f32 = scale_inv_d.view(torch.float8_e8m0fnu).float() - sf_d_ref = torch.repeat_interleave(scale_inv_d_f32.reshape(l, s_padded, d_scale_padded), repeats=32, dim=2)[:, :s, :d].contiguous() + sf_d_ref = None + if with_ref: + scale_inv_d = mxfp8_result._rowwise_scale_inv + scale_inv_d_f32 = scale_inv_d.view(torch.float8_e8m0fnu).float() + sf_d_ref = torch.repeat_interleave(scale_inv_d_f32.reshape(l, s_padded, d_scale_padded), repeats=32, dim=2)[:, :s, :d].contiguous() # --- Columnwise results (quantized along S dimension) --- fp8_data_s_flat = mxfp8_result._columnwise_data fp8_data_s = fp8_data_s_flat.reshape(l, s_padded, d_padded)[:, :s, :d].contiguous() fp8_data_s = fp8_data_s.view(fp8_dtype).reshape(b, h, s, d) - scale_inv_s = mxfp8_result._columnwise_scale_inv - scale_inv_s_f32 = scale_inv_s.view(torch.float8_e8m0fnu).float() - sf_s_ref = torch.repeat_interleave(scale_inv_s_f32.reshape(l, s_scale_padded, d_padded), repeats=32, dim=1)[:, :s, :d].contiguous() + sf_s_ref = None + if with_ref: + scale_inv_s = mxfp8_result._columnwise_scale_inv + scale_inv_s_f32 = scale_inv_s.view(torch.float8_e8m0fnu).float() + sf_s_ref = torch.repeat_interleave(scale_inv_s_f32.reshape(l, s_scale_padded, d_padded), repeats=32, dim=1)[:, :s, :d].contiguous() # with swizzle - quantizer_swizzle.optimize_for_gemm = True - mxfp8_result_swizzle = quantizer_swizzle(tensor_2d) + tex.swizzle_scales_for_gemm_(mxfp8_result) # --- Rowwise results (quantized along D dimension) --- - sf_d_swizzle = mxfp8_result_swizzle._rowwise_scale_inv + sf_d_swizzle = mxfp8_result._rowwise_scale_inv # --- Columnwise results (quantized along S dimension) --- - sf_s_swizzle = mxfp8_result_swizzle._columnwise_scale_inv + sf_s_swizzle = mxfp8_result._columnwise_scale_inv return fp8_data_d, sf_d_ref, sf_d_swizzle, fp8_data_s, sf_s_ref, sf_s_swizzle @@ -498,6 +533,7 @@ def exec_sdpa_mxfp8(cfg, request, cudnn_handle): """Execute MXFP8 SDPA test.""" if request.config.option.dryrun: pytest.skip("dry run mode") + perf = request.config.getoption("--perf") cudnn_version = LooseVersion(cudnn.backend_version_string()) @@ -563,9 +599,9 @@ def exec_sdpa_mxfp8(cfg, request, cudnn_handle): v_f32 = torch.empty(b, h_v, s_kv, d_vo, dtype=torch.float32, device="cuda") fill_sparse_small_int(v_f32, rng_data, sparsity=0.8, abs_max=2) - q_fp8_d, sf_q_d_ref, sf_q_d_swizzle, q_fp8_s, sf_q_s_ref, sf_q_s_swizzle = quantize_to_mxfp8(q_f32, b, h_q, s_qo, d_qk, block_size, torch_itype) - k_fp8_d, sf_k_d_ref, sf_k_d_swizzle, k_fp8_s, sf_k_s_ref, sf_k_s_swizzle = quantize_to_mxfp8(k_f32, b, h_k, s_kv, d_qk, block_size, torch_itype) - v_fp8_d, sf_v_d_ref, sf_v_d_swizzle, v_fp8_s, sf_v_s_ref, sf_v_s_swizzle = quantize_to_mxfp8(v_f32, b, h_v, s_kv, d_vo, block_size, torch_itype) + q_fp8_d, sf_q_d_ref, sf_q_d_swizzle, q_fp8_s, sf_q_s_ref, sf_q_s_swizzle = quantize_to_mxfp8(q_f32, b, h_q, s_qo, d_qk, block_size, torch_itype, with_ref=not perf) + k_fp8_d, sf_k_d_ref, sf_k_d_swizzle, k_fp8_s, sf_k_s_ref, sf_k_s_swizzle = quantize_to_mxfp8(k_f32, b, h_k, s_kv, d_qk, block_size, torch_itype, with_ref=not perf) + v_fp8_d, sf_v_d_ref, sf_v_d_swizzle, v_fp8_s, sf_v_s_ref, sf_v_s_swizzle = quantize_to_mxfp8(v_f32, b, h_v, s_kv, d_vo, block_size, torch_itype, with_ref=not perf) # Generate sink_token if needed sink_token_gpu = None @@ -596,41 +632,35 @@ def exec_sdpa_mxfp8(cfg, request, cudnn_handle): # Execute workspace = torch.empty(graph_fwd.get_workspace_size(), dtype=torch.uint8, device="cuda") torch.cuda.synchronize() - if request.config.getoption("--perf"): + if perf: times_ms = time_execution(graph_fwd.execute, variant_pack, workspace, cudnn_handle) print(f"@@@@ MXFP8 Fwd graph_fwd.execute avg_time_ms={times_ms.mean().item():.3f}") profile_execution(graph_fwd.execute, variant_pack, workspace, cudnn_handle) graph_fwd.execute(variant_pack, workspace, handle=cudnn_handle) torch.cuda.synchronize() - # Compute reference - o_ref, stats_ref = compute_ref(q_fp8_d, k_fp8_d, v_fp8_s, sf_q_d_ref, sf_k_d_ref, sf_v_s_ref, attn_scale, - torch_itype=torch_itype, output_type=torch_otype, - left_bound=left_bound, right_bound=right_bound, diag_align=diag_align, - sink_token=sink_token_gpu, rescale_threshold=rescale_threshold) - - # Compare output - o_atol, o_rtol = 0.12, 0.20 - o_err = compare_tensors(o_gpu, o_ref, o_atol, o_rtol, "output") - - # Compare stats (logsumexp) - tight tolerance - stats_atol, stats_rtol = 0.05, 0.05 - stats_err = compare_tensors(stats_gpu, stats_ref, stats_atol, stats_rtol, "stats") - - # Compare amax - amax_err = compare_amax(o_gpu, o_ref, rtol=0.05, tag="amax") - - # Assert all checks pass - assert o_err == 0, f"Output mismatch: {o_err} elements differ" - assert stats_err == 0, f"Stats mismatch: {stats_err} elements differ" - assert amax_err, "Amax mismatch: 1 element differs" + o_f16 = o_gpu + stats_bwd = stats_gpu + if not perf: + o_ref, stats_ref = compute_ref(q_fp8_d, k_fp8_d, v_fp8_s, sf_q_d_ref, sf_k_d_ref, sf_v_s_ref, attn_scale, + torch_itype=torch_itype, output_type=torch_otype, + left_bound=left_bound, right_bound=right_bound, diag_align=diag_align, + sink_token=sink_token_gpu, rescale_threshold=rescale_threshold) + o_f16 = o_ref.to(torch.bfloat16) + stats_bwd = stats_ref + for actual, expected, atol, rtol, name in ( + (o_gpu, o_ref, 0.12, 0.20, "output"), + (stats_gpu, stats_ref, 0.05, 0.05, "stats"), + ): + error = compare_tensors(actual, expected, atol, rtol, name) + assert error == 0, f"{name} mismatch: {error} elements differ" + assert compare_amax(o_gpu, o_ref, rtol=0.05, tag="amax"), "Amax mismatch: 1 element differs" if not cfg.is_infer: dO_f32 = torch.empty(b, h_q, s_qo, d_vo, dtype=torch.float32, device="cuda") fill_sparse_small_int(dO_f32, rng_data, sparsity=0.8, abs_max=2) - dO_fp8_d, sf_dO_d_ref, sf_dO_d_swizzle, dO_fp8_s, sf_dO_s_ref, sf_dO_s_swizzle = quantize_to_mxfp8(dO_f32, b, h_q, s_qo, d_vo, block_size, torch_itype) + dO_fp8_d, sf_dO_d_ref, sf_dO_d_swizzle, dO_fp8_s, sf_dO_s_ref, sf_dO_s_swizzle = quantize_to_mxfp8(dO_f32, b, h_q, s_qo, d_vo, block_size, torch_itype, with_ref=not perf) - o_f16 = o_ref.to(torch.bfloat16) dO_f16 = dO_f32.to(torch.bfloat16) # Build backward graph @@ -675,7 +705,7 @@ def exec_sdpa_mxfp8(cfg, request, cudnn_handle): int(GraphBwdUid.dO): dO_fp8_d, int(GraphBwdUid.dO_t): dO_fp8_s, int(GraphBwdUid.dO_f16): dO_f16, - int(GraphBwdUid.stats): stats_ref, + int(GraphBwdUid.stats): stats_bwd, int(GraphBwdUid.sf_q): sf_q_d_swizzle, int(GraphBwdUid.sf_q_t): sf_q_s_swizzle, int(GraphBwdUid.sf_k): sf_k_d_swizzle, @@ -696,7 +726,7 @@ def exec_sdpa_mxfp8(cfg, request, cudnn_handle): # Execute backward graph workspace_bwd = torch.empty(graph_bwd.get_workspace_size(), dtype=torch.uint8, device="cuda") - if request.config.getoption("--perf"): + if perf: times_ms = time_execution(graph_bwd.execute, variant_pack_bwd, workspace_bwd, cudnn_handle) print(f"@@@@ MXFP8 Bwd graph_bwd.execute avg_time_ms={times_ms.mean().item():.3f}") profile_execution(graph_bwd.execute, variant_pack_bwd, workspace_bwd, cudnn_handle) @@ -735,39 +765,33 @@ def exec_sdpa_mxfp8(cfg, request, cudnn_handle): pytest.fail("determinism check failed", pytrace=False) - # Compute reference backward - dQ_ref, dK_ref, dV_ref, dSink_token_ref = compute_ref_backward( - q_fp8_d, q_fp8_s, k_fp8_d, k_fp8_s, v_fp8_d, - o_f16, dO_f16, dO_fp8_d, dO_fp8_s, - attn_scale, - sf_q_d_ref, sf_q_s_ref, sf_k_d_ref, sf_k_s_ref, sf_v_d_ref, - sf_dO_d_ref, sf_dO_s_ref, - torch_itype=torch_itype, torch_otype=torch_otype, - left_bound=left_bound, right_bound=right_bound, diag_align=diag_align, - sink_token=sink_token_gpu, - ) + if not perf: + dQ_ref, dK_ref, dV_ref, dSink_token_ref = compute_ref_backward( + q_fp8_d, q_fp8_s, k_fp8_d, k_fp8_s, v_fp8_d, + o_f16, dO_f16, dO_fp8_d, dO_fp8_s, + attn_scale, + sf_q_d_ref, sf_q_s_ref, sf_k_d_ref, sf_k_s_ref, sf_v_d_ref, + sf_dO_d_ref, sf_dO_s_ref, + torch_itype=torch_itype, torch_otype=torch_otype, + left_bound=left_bound, right_bound=right_bound, diag_align=diag_align, + sink_token=sink_token_gpu, + ) - # Compare output - grad_atol, grad_rtol = 0.08, 0.20 - - dQ_err = compare_tensors(dQ_gpu, dQ_ref, grad_atol, grad_rtol, "dQ") - dK_err = compare_tensors(dK_gpu, dK_ref, grad_atol, grad_rtol, "dK") - dV_err = compare_tensors(dV_gpu, dV_ref, grad_atol, grad_rtol, "dV") - - assert dQ_err == 0, f"dQ mismatch: {dQ_err} elements differ" - assert dK_err == 0, f"dK mismatch: {dK_err} elements differ" - assert dV_err == 0, f"dV mismatch: {dV_err} elements differ" - - # Compare dSink_token if using sink_token - if with_sink_token and dSink_token_ref is not None: - dSink_err = compare_tensors(dSink_token_gpu, dSink_token_ref, grad_atol, grad_rtol, "dSink_token") - assert dSink_err == 0, f"dSink_token mismatch: {dSink_err} elements differ" - - # Compare amax - dQ_amax_err = compare_amax(dQ_gpu, dQ_ref, rtol=0.04, tag="dQ") - dK_amax_err = compare_amax(dK_gpu, dK_ref, rtol=0.04, tag="dK") - dV_amax_err = compare_amax(dV_gpu, dV_ref, rtol=0.04, tag="dV") - - assert dQ_amax_err, "dQ amax mismatch: 1 element differs" - assert dK_amax_err, "dK amax mismatch: 1 element differs" - assert dV_amax_err, "dV amax mismatch: 1 element differs" + for actual, expected, name in ( + (dQ_gpu, dQ_ref, "dQ"), + (dK_gpu, dK_ref, "dK"), + (dV_gpu, dV_ref, "dV"), + ): + error = compare_tensors(actual, expected, 0.08, 0.20, name) + assert error == 0, f"{name} mismatch: {error} elements differ" + + if with_sink_token and dSink_token_ref is not None: + dSink_err = compare_tensors(dSink_token_gpu, dSink_token_ref, 0.08, 0.20, "dSink_token") + assert dSink_err == 0, f"dSink_token mismatch: {dSink_err} elements differ" + + for actual, expected, name in ( + (dQ_gpu, dQ_ref, "dQ"), + (dK_gpu, dK_ref, "dK"), + (dV_gpu, dV_ref, "dV"), + ): + assert compare_amax(actual, expected, rtol=0.04, tag=name), f"{name} amax mismatch: 1 element differs" diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 746cf2d1d..6095e3c64 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -1020,6 +1020,39 @@ def test_sdpa_mxfp8_bwd_L0(env_info, test_no, request, cudnn_handle): if "CUDNN_RESCALE_THRESHOLD" in os.environ: del os.environ["CUDNN_RESCALE_THRESHOLD"] + +@pytest.mark.L0 +def test_sdpa_mxfp8_perf_L0(monkeypatch, request, cudnn_handle): + def fail_reference(*args, **kwargs): + raise AssertionError("reference called in perf mode") + + getoption = request.config.getoption + def getoption_perf(name, *args, **kwargs): + return True if name == "--perf" else getoption(name, *args, **kwargs) + + monkeypatch.setattr(request.config, "getoption", getoption_perf) + monkeypatch.setattr("sdpa.mxfp8.compute_ref", fail_reference) + monkeypatch.setattr("sdpa.mxfp8.compute_ref_backward", fail_reference) + + cfg = ExecConfig( + data_type=torch.float8_e4m3fn, + output_type=torch.bfloat16, + rng_data_seed=1, + is_infer=False, + is_determin=False, + is_mxfp8=True, + batches=1, + h_q=1, + h_k=1, + h_v=1, + s_q=128, + s_kv=128, + d_qk=64, + d_v=64, + ) + cfg.fill_derived_fields() + exec_sdpa_mxfp8(cfg, request, cudnn_handle) + # # =================== # # Single repro test # # =================== From 40a2a01adc0d315fbde326a60c759c293d8d6ab4 Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Mon, 13 Jul 2026 11:29:54 -0700 Subject: [PATCH 4/9] Address SDPA repro review comments --- include/cudnn_frontend/graph_interface.h | 8 ++++++++ test/python/sdpa/mxfp8.py | 23 +++++++++++++---------- 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/include/cudnn_frontend/graph_interface.h b/include/cudnn_frontend/graph_interface.h index 2d28b928c..1c6556498 100644 --- a/include/cudnn_frontend/graph_interface.h +++ b/include/cudnn_frontend/graph_interface.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -975,6 +976,13 @@ class Graph : public ICudnn, public INode { CHECK_CUDNN_FRONTEND_ERROR(create_cudnn_tensors_subtree(uid_to_tensors, start_uid, used_uids)); tensors_to_dump.clear(); CHECK_CUDNN_FRONTEND_ERROR(collect_tensors_to_dump_subtree(tensors_to_dump)); + std::unordered_set dumped_uids; + tensors_to_dump.erase(std::remove_if(tensors_to_dump.begin(), + tensors_to_dump.end(), + [&dumped_uids](auto const &entry) { + return !dumped_uids.insert(entry.first->get_uid()).second; + }), + tensors_to_dump.end()); CUDNN_FE_LOG_BANNER(" 3/4 CREATE OPERATIONS "); // INode keeps track of all uids that an operation graph uses. diff --git a/test/python/sdpa/mxfp8.py b/test/python/sdpa/mxfp8.py index c5fd42b29..ef64d143d 100644 --- a/test/python/sdpa/mxfp8.py +++ b/test/python/sdpa/mxfp8.py @@ -41,14 +41,15 @@ def quantize_mxfp8_compact(tensor_2d, quantizer): chunk_rows = MXFP8_QUANTIZER_MAX_ROWS // 128 * 128 results = [quantizer(chunk) for chunk in tensor_2d.split(chunk_rows, dim=0)] + metadata = [result.get_metadata() for result in results] return MXFP8Tensor( shape=tensor_2d.shape, dtype=tensor_2d.dtype, fp8_dtype=quantizer.dtype, - rowwise_data=torch.cat([result._rowwise_data for result in results], dim=0), - rowwise_scale_inv=torch.cat([result._rowwise_scale_inv for result in results], dim=0), - columnwise_data=torch.cat([result._columnwise_data for result in results], dim=0), - columnwise_scale_inv=torch.cat([result._columnwise_scale_inv for result in results], dim=0), + rowwise_data=torch.cat([item["rowwise_data"] for item in metadata], dim=0), + rowwise_scale_inv=torch.cat([item["rowwise_scale_inv"] for item in metadata], dim=0), + columnwise_data=torch.cat([item["columnwise_data"] for item in metadata], dim=0), + columnwise_scale_inv=torch.cat([item["columnwise_scale_inv"] for item in metadata], dim=0), quantizer=quantizer, requires_grad=False, with_gemm_swizzled_scales=False, @@ -189,35 +190,37 @@ def quantize_to_mxfp8( # without swizzle quantizer = MXFP8Quantizer(fp8_dtype=te_dtype, rowwise=True, columnwise=True) mxfp8_result = quantize_mxfp8_compact(tensor_2d, quantizer) + metadata = mxfp8_result.get_metadata() # --- Rowwise results (quantized along D dimension) --- - fp8_data_d_flat = mxfp8_result._rowwise_data + fp8_data_d_flat = metadata["rowwise_data"] fp8_data_d = fp8_data_d_flat.reshape(l, s_padded, d_padded)[:, :s, :d].contiguous() fp8_data_d = fp8_data_d.view(fp8_dtype).reshape(b, h, s, d) sf_d_ref = None if with_ref: - scale_inv_d = mxfp8_result._rowwise_scale_inv + scale_inv_d = metadata["rowwise_scale_inv"] scale_inv_d_f32 = scale_inv_d.view(torch.float8_e8m0fnu).float() sf_d_ref = torch.repeat_interleave(scale_inv_d_f32.reshape(l, s_padded, d_scale_padded), repeats=32, dim=2)[:, :s, :d].contiguous() # --- Columnwise results (quantized along S dimension) --- - fp8_data_s_flat = mxfp8_result._columnwise_data + fp8_data_s_flat = metadata["columnwise_data"] fp8_data_s = fp8_data_s_flat.reshape(l, s_padded, d_padded)[:, :s, :d].contiguous() fp8_data_s = fp8_data_s.view(fp8_dtype).reshape(b, h, s, d) sf_s_ref = None if with_ref: - scale_inv_s = mxfp8_result._columnwise_scale_inv + scale_inv_s = metadata["columnwise_scale_inv"] scale_inv_s_f32 = scale_inv_s.view(torch.float8_e8m0fnu).float() sf_s_ref = torch.repeat_interleave(scale_inv_s_f32.reshape(l, s_scale_padded, d_padded), repeats=32, dim=1)[:, :s, :d].contiguous() # with swizzle tex.swizzle_scales_for_gemm_(mxfp8_result) + metadata = mxfp8_result.get_metadata() # --- Rowwise results (quantized along D dimension) --- - sf_d_swizzle = mxfp8_result._rowwise_scale_inv + sf_d_swizzle = metadata["rowwise_scale_inv"] # --- Columnwise results (quantized along S dimension) --- - sf_s_swizzle = mxfp8_result._columnwise_scale_inv + sf_s_swizzle = metadata["columnwise_scale_inv"] return fp8_data_d, sf_d_ref, sf_d_swizzle, fp8_data_s, sf_s_ref, sf_s_swizzle From 3e4266b1a7f223b4de6a6f9c457214a65fcd330b Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Mon, 13 Jul 2026 11:40:40 -0700 Subject: [PATCH 5/9] Keep MXFP8 storage access direct --- test/python/sdpa/mxfp8.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/test/python/sdpa/mxfp8.py b/test/python/sdpa/mxfp8.py index ef64d143d..c5fd42b29 100644 --- a/test/python/sdpa/mxfp8.py +++ b/test/python/sdpa/mxfp8.py @@ -41,15 +41,14 @@ def quantize_mxfp8_compact(tensor_2d, quantizer): chunk_rows = MXFP8_QUANTIZER_MAX_ROWS // 128 * 128 results = [quantizer(chunk) for chunk in tensor_2d.split(chunk_rows, dim=0)] - metadata = [result.get_metadata() for result in results] return MXFP8Tensor( shape=tensor_2d.shape, dtype=tensor_2d.dtype, fp8_dtype=quantizer.dtype, - rowwise_data=torch.cat([item["rowwise_data"] for item in metadata], dim=0), - rowwise_scale_inv=torch.cat([item["rowwise_scale_inv"] for item in metadata], dim=0), - columnwise_data=torch.cat([item["columnwise_data"] for item in metadata], dim=0), - columnwise_scale_inv=torch.cat([item["columnwise_scale_inv"] for item in metadata], dim=0), + rowwise_data=torch.cat([result._rowwise_data for result in results], dim=0), + rowwise_scale_inv=torch.cat([result._rowwise_scale_inv for result in results], dim=0), + columnwise_data=torch.cat([result._columnwise_data for result in results], dim=0), + columnwise_scale_inv=torch.cat([result._columnwise_scale_inv for result in results], dim=0), quantizer=quantizer, requires_grad=False, with_gemm_swizzled_scales=False, @@ -190,37 +189,35 @@ def quantize_to_mxfp8( # without swizzle quantizer = MXFP8Quantizer(fp8_dtype=te_dtype, rowwise=True, columnwise=True) mxfp8_result = quantize_mxfp8_compact(tensor_2d, quantizer) - metadata = mxfp8_result.get_metadata() # --- Rowwise results (quantized along D dimension) --- - fp8_data_d_flat = metadata["rowwise_data"] + fp8_data_d_flat = mxfp8_result._rowwise_data fp8_data_d = fp8_data_d_flat.reshape(l, s_padded, d_padded)[:, :s, :d].contiguous() fp8_data_d = fp8_data_d.view(fp8_dtype).reshape(b, h, s, d) sf_d_ref = None if with_ref: - scale_inv_d = metadata["rowwise_scale_inv"] + scale_inv_d = mxfp8_result._rowwise_scale_inv scale_inv_d_f32 = scale_inv_d.view(torch.float8_e8m0fnu).float() sf_d_ref = torch.repeat_interleave(scale_inv_d_f32.reshape(l, s_padded, d_scale_padded), repeats=32, dim=2)[:, :s, :d].contiguous() # --- Columnwise results (quantized along S dimension) --- - fp8_data_s_flat = metadata["columnwise_data"] + fp8_data_s_flat = mxfp8_result._columnwise_data fp8_data_s = fp8_data_s_flat.reshape(l, s_padded, d_padded)[:, :s, :d].contiguous() fp8_data_s = fp8_data_s.view(fp8_dtype).reshape(b, h, s, d) sf_s_ref = None if with_ref: - scale_inv_s = metadata["columnwise_scale_inv"] + scale_inv_s = mxfp8_result._columnwise_scale_inv scale_inv_s_f32 = scale_inv_s.view(torch.float8_e8m0fnu).float() sf_s_ref = torch.repeat_interleave(scale_inv_s_f32.reshape(l, s_scale_padded, d_padded), repeats=32, dim=1)[:, :s, :d].contiguous() # with swizzle tex.swizzle_scales_for_gemm_(mxfp8_result) - metadata = mxfp8_result.get_metadata() # --- Rowwise results (quantized along D dimension) --- - sf_d_swizzle = metadata["rowwise_scale_inv"] + sf_d_swizzle = mxfp8_result._rowwise_scale_inv # --- Columnwise results (quantized along S dimension) --- - sf_s_swizzle = metadata["columnwise_scale_inv"] + sf_s_swizzle = mxfp8_result._columnwise_scale_inv return fp8_data_d, sf_d_ref, sf_d_swizzle, fp8_data_s, sf_s_ref, sf_s_swizzle From ae27d0be7ca90e535470c6d68ebc03f23ebf8693 Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Mon, 13 Jul 2026 11:54:45 -0700 Subject: [PATCH 6/9] Use UID map for tensor dump collection --- include/cudnn_frontend/graph_interface.h | 21 +++++++------------ .../node/scaled_dot_product_flash_attention.h | 20 ++++++++---------- include/cudnn_frontend/node_interface.h | 7 +++++-- 3 files changed, 21 insertions(+), 27 deletions(-) diff --git a/include/cudnn_frontend/graph_interface.h b/include/cudnn_frontend/graph_interface.h index 1c6556498..048eb0cb5 100644 --- a/include/cudnn_frontend/graph_interface.h +++ b/include/cudnn_frontend/graph_interface.h @@ -1,6 +1,5 @@ #pragma once -#include #include #include #include @@ -62,7 +61,7 @@ class Graph : public ICudnn, public INode { mutable std::unordered_map>> cached_workspace_modifications; // char: 'x'=hex, 'd'=decimal, 'b'=base64 - std::vector, char>> tensors_to_dump; + std::map, char>> tensors_to_dump; error_t get_pre_assigned_uids(std::unordered_set &used_uids) { @@ -106,8 +105,9 @@ class Graph : public ICudnn, public INode { cudaStream_t stream; _CUDNN_CHECK_CUDNN_ERROR(detail::get_stream(handle, &stream)); - for (auto const &[tensor, fmt] : tensors_to_dump) { - auto it = tensor_uid_to_pointer_map.find(tensor->get_uid()); + for (auto const &[uid, tensor_and_format] : tensors_to_dump) { + auto const &[tensor, fmt] = tensor_and_format; + auto it = tensor_uid_to_pointer_map.find(uid); if (it != tensor_uid_to_pointer_map.end()) { auto const &dims = tensor->get_dim(); size_t num_elements = 1; @@ -976,13 +976,6 @@ class Graph : public ICudnn, public INode { CHECK_CUDNN_FRONTEND_ERROR(create_cudnn_tensors_subtree(uid_to_tensors, start_uid, used_uids)); tensors_to_dump.clear(); CHECK_CUDNN_FRONTEND_ERROR(collect_tensors_to_dump_subtree(tensors_to_dump)); - std::unordered_set dumped_uids; - tensors_to_dump.erase(std::remove_if(tensors_to_dump.begin(), - tensors_to_dump.end(), - [&dumped_uids](auto const &entry) { - return !dumped_uids.insert(entry.first->get_uid()).second; - }), - tensors_to_dump.end()); CUDNN_FE_LOG_BANNER(" 3/4 CREATE OPERATIONS "); // INode keeps track of all uids that an operation graph uses. @@ -1671,8 +1664,8 @@ class Graph : public ICudnn, public INode { j["fe_workspace_size"] = fe_workspace_size; std::vector> tensors_to_dump_uids; - for (auto const &[tensor, fmt] : tensors_to_dump) { - tensors_to_dump_uids.emplace_back(tensor->get_uid(), fmt); + for (auto const &[uid, tensor_and_format] : tensors_to_dump) { + tensors_to_dump_uids.emplace_back(uid, tensor_and_format.second); } j["tensors_to_dump"] = tensors_to_dump_uids; @@ -1809,7 +1802,7 @@ class Graph : public ICudnn, public INode { for (auto const &[uid, fmt] : dump_uids) { for (auto const &tensor : deserialized_tensor_properties) { if (tensor->get_uid() == uid) { - tensors_to_dump.emplace_back(tensor, fmt); + tensors_to_dump.emplace(uid, std::make_pair(tensor, fmt)); break; } } 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 e46dfde4e..636ff5863 100644 --- a/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h +++ b/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h @@ -522,15 +522,13 @@ class SDPANodeBase : public NodeCRTP { error_t collect_tensors_to_dump_node( - std::vector, char>>& tensors_to_dump) const override final { - std::unordered_set seen_uids; - auto add_tensor = [&tensors_to_dump, &seen_uids](std::shared_ptr const& tensor) { + std::map, char>>& tensors_to_dump) + const override final { + auto add_tensor = [&tensors_to_dump](std::shared_ptr const& tensor) { if (tensor == nullptr) { return; } - if (seen_uids.insert(tensor->get_uid()).second) { - tensors_to_dump.emplace_back(tensor, 'd'); - } + tensors_to_dump.emplace(tensor->get_uid(), std::make_pair(tensor, 'd')); }; auto const seq_len_q_it = attributes.inputs.find(input_names::SEQ_LEN_Q); @@ -2184,11 +2182,11 @@ class CompositeSDPABackwardNode : public NodeCRTP { error_t collect_tensors_to_dump_node( - std::vector, char>>& tensors_to_dump) const override final { - std::unordered_set seen_uids; - auto add_tensor = [&tensors_to_dump, &seen_uids](std::shared_ptr const& tensor) { - if (tensor != nullptr && seen_uids.insert(tensor->get_uid()).second) { - tensors_to_dump.emplace_back(tensor, 'd'); + std::map, char>>& tensors_to_dump) + const override final { + auto add_tensor = [&tensors_to_dump](std::shared_ptr const& tensor) { + if (tensor != nullptr) { + tensors_to_dump.emplace(tensor->get_uid(), std::make_pair(tensor, 'd')); } }; auto add_input = [&](input_names name) { diff --git a/include/cudnn_frontend/node_interface.h b/include/cudnn_frontend/node_interface.h index 8d46ece33..ab72a9030 100644 --- a/include/cudnn_frontend/node_interface.h +++ b/include/cudnn_frontend/node_interface.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -96,7 +97,8 @@ class INode { }; virtual error_t - collect_tensors_to_dump_node(std::vector, char>>&) const { + collect_tensors_to_dump_node( + std::map, char>>&) const { return {error_code_t::OK, ""}; }; @@ -351,7 +353,8 @@ class INode { error_t collect_tensors_to_dump_subtree( - std::vector, char>>& tensors_to_dump) const { + std::map, char>>& tensors_to_dump) + const { CHECK_CUDNN_FRONTEND_ERROR(collect_tensors_to_dump_node(tensors_to_dump)); for (auto const& sub_node : sub_nodes) { CHECK_CUDNN_FRONTEND_ERROR(sub_node->collect_tensors_to_dump_subtree(tensors_to_dump)); From 749a3e3a7c1e2a7a4d4c1705625a2f0a65b2377a Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Mon, 13 Jul 2026 12:34:56 -0700 Subject: [PATCH 7/9] Avoid monkeypatching MXFP8 performance test --- test/python/sdpa/mxfp8.py | 3 +-- test/python/test_mhas_v2.py | 21 +++++---------------- 2 files changed, 6 insertions(+), 18 deletions(-) diff --git a/test/python/sdpa/mxfp8.py b/test/python/sdpa/mxfp8.py index c5fd42b29..a5b4b8f8e 100644 --- a/test/python/sdpa/mxfp8.py +++ b/test/python/sdpa/mxfp8.py @@ -529,11 +529,10 @@ def generate_graph_bwd(b, h_q, h_k, h_v, return graph_bwd -def exec_sdpa_mxfp8(cfg, request, cudnn_handle): +def exec_sdpa_mxfp8(cfg, request, cudnn_handle, perf): """Execute MXFP8 SDPA test.""" if request.config.option.dryrun: pytest.skip("dry run mode") - perf = request.config.getoption("--perf") cudnn_version = LooseVersion(cudnn.backend_version_string()) diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 6095e3c64..68ea4192b 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -966,7 +966,7 @@ def test_sdpa_mxfp8_fwd_L0(env_info, test_no, request, cudnn_handle): if request.node.name in test.blocked_tests: pytest.skip(f"blocked test: {request.node.name}") try: - exec_sdpa_mxfp8(test.cfg, request, cudnn_handle) + exec_sdpa_mxfp8(test.cfg, request, cudnn_handle, request.config.getoption("--perf")) finally: if "CUDNN_UNFUSE_FMA" in os.environ: del os.environ["CUDNN_UNFUSE_FMA"] @@ -1015,25 +1015,14 @@ def test_sdpa_mxfp8_bwd_L0(env_info, test_no, request, cudnn_handle): if request.node.name in test.blocked_tests: pytest.skip(f"blocked test: {request.node.name}") try: - exec_sdpa_mxfp8(test.cfg, request, cudnn_handle) + exec_sdpa_mxfp8(test.cfg, request, cudnn_handle, request.config.getoption("--perf")) finally: if "CUDNN_RESCALE_THRESHOLD" in os.environ: del os.environ["CUDNN_RESCALE_THRESHOLD"] @pytest.mark.L0 -def test_sdpa_mxfp8_perf_L0(monkeypatch, request, cudnn_handle): - def fail_reference(*args, **kwargs): - raise AssertionError("reference called in perf mode") - - getoption = request.config.getoption - def getoption_perf(name, *args, **kwargs): - return True if name == "--perf" else getoption(name, *args, **kwargs) - - monkeypatch.setattr(request.config, "getoption", getoption_perf) - monkeypatch.setattr("sdpa.mxfp8.compute_ref", fail_reference) - monkeypatch.setattr("sdpa.mxfp8.compute_ref_backward", fail_reference) - +def test_sdpa_mxfp8_perf_L0(request, cudnn_handle): cfg = ExecConfig( data_type=torch.float8_e4m3fn, output_type=torch.bfloat16, @@ -1051,7 +1040,7 @@ def getoption_perf(name, *args, **kwargs): d_v=64, ) cfg.fill_derived_fields() - exec_sdpa_mxfp8(cfg, request, cudnn_handle) + exec_sdpa_mxfp8(cfg, request, cudnn_handle, perf=True) # # =================== # # Single repro test @@ -1083,7 +1072,7 @@ def test_repro(env_info, request, cudnn_handle): try: if cfg.cfg.is_mxfp8: - exec_sdpa_mxfp8(cfg.cfg, request, cudnn_handle) + exec_sdpa_mxfp8(cfg.cfg, request, cudnn_handle, request.config.getoption("--perf")) elif cfg.cfg.data_type in (torch.float8_e4m3fn, torch.float8_e5m2): exec_sdpa_fp8(cfg.cfg, request, cudnn_handle) else: From 848ad26003347f6413e5129f283e5389588fbfe9 Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Mon, 13 Jul 2026 12:51:44 -0700 Subject: [PATCH 8/9] Use vector for tensor dump collection --- include/cudnn_frontend/graph_interface.h | 22 +++++++++++++------ .../node/scaled_dot_product_flash_attention.h | 10 ++++----- include/cudnn_frontend/node_interface.h | 7 ++---- 3 files changed, 21 insertions(+), 18 deletions(-) diff --git a/include/cudnn_frontend/graph_interface.h b/include/cudnn_frontend/graph_interface.h index 048eb0cb5..8e92853aa 100644 --- a/include/cudnn_frontend/graph_interface.h +++ b/include/cudnn_frontend/graph_interface.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -61,7 +62,7 @@ class Graph : public ICudnn, public INode { mutable std::unordered_map>> cached_workspace_modifications; // char: 'x'=hex, 'd'=decimal, 'b'=base64 - std::map, char>> tensors_to_dump; + std::vector, char>> tensors_to_dump; error_t get_pre_assigned_uids(std::unordered_set &used_uids) { @@ -105,9 +106,8 @@ class Graph : public ICudnn, public INode { cudaStream_t stream; _CUDNN_CHECK_CUDNN_ERROR(detail::get_stream(handle, &stream)); - for (auto const &[uid, tensor_and_format] : tensors_to_dump) { - auto const &[tensor, fmt] = tensor_and_format; - auto it = tensor_uid_to_pointer_map.find(uid); + for (auto const &[tensor, fmt] : tensors_to_dump) { + auto it = tensor_uid_to_pointer_map.find(tensor->get_uid()); if (it != tensor_uid_to_pointer_map.end()) { auto const &dims = tensor->get_dim(); size_t num_elements = 1; @@ -976,6 +976,14 @@ class Graph : public ICudnn, public INode { CHECK_CUDNN_FRONTEND_ERROR(create_cudnn_tensors_subtree(uid_to_tensors, start_uid, used_uids)); tensors_to_dump.clear(); CHECK_CUDNN_FRONTEND_ERROR(collect_tensors_to_dump_subtree(tensors_to_dump)); + std::stable_sort(tensors_to_dump.begin(), tensors_to_dump.end(), [](auto const &lhs, auto const &rhs) { + return lhs.first->get_uid() < rhs.first->get_uid(); + }); + tensors_to_dump.erase( + std::unique(tensors_to_dump.begin(), + tensors_to_dump.end(), + [](auto const &lhs, auto const &rhs) { return lhs.first->get_uid() == rhs.first->get_uid(); }), + tensors_to_dump.end()); CUDNN_FE_LOG_BANNER(" 3/4 CREATE OPERATIONS "); // INode keeps track of all uids that an operation graph uses. @@ -1664,8 +1672,8 @@ class Graph : public ICudnn, public INode { j["fe_workspace_size"] = fe_workspace_size; std::vector> tensors_to_dump_uids; - for (auto const &[uid, tensor_and_format] : tensors_to_dump) { - tensors_to_dump_uids.emplace_back(uid, tensor_and_format.second); + for (auto const &[tensor, fmt] : tensors_to_dump) { + tensors_to_dump_uids.emplace_back(tensor->get_uid(), fmt); } j["tensors_to_dump"] = tensors_to_dump_uids; @@ -1802,7 +1810,7 @@ class Graph : public ICudnn, public INode { for (auto const &[uid, fmt] : dump_uids) { for (auto const &tensor : deserialized_tensor_properties) { if (tensor->get_uid() == uid) { - tensors_to_dump.emplace(uid, std::make_pair(tensor, fmt)); + tensors_to_dump.emplace_back(tensor, fmt); break; } } 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 636ff5863..ae25c1d18 100644 --- a/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h +++ b/include/cudnn_frontend/node/scaled_dot_product_flash_attention.h @@ -522,13 +522,12 @@ class SDPANodeBase : public NodeCRTP { error_t collect_tensors_to_dump_node( - std::map, char>>& tensors_to_dump) - const override final { + std::vector, char>>& tensors_to_dump) const override final { auto add_tensor = [&tensors_to_dump](std::shared_ptr const& tensor) { if (tensor == nullptr) { return; } - tensors_to_dump.emplace(tensor->get_uid(), std::make_pair(tensor, 'd')); + tensors_to_dump.emplace_back(tensor, 'd'); }; auto const seq_len_q_it = attributes.inputs.find(input_names::SEQ_LEN_Q); @@ -2182,11 +2181,10 @@ class CompositeSDPABackwardNode : public NodeCRTP { error_t collect_tensors_to_dump_node( - std::map, char>>& tensors_to_dump) - const override final { + std::vector, char>>& tensors_to_dump) const override final { auto add_tensor = [&tensors_to_dump](std::shared_ptr const& tensor) { if (tensor != nullptr) { - tensors_to_dump.emplace(tensor->get_uid(), std::make_pair(tensor, 'd')); + tensors_to_dump.emplace_back(tensor, 'd'); } }; auto add_input = [&](input_names name) { diff --git a/include/cudnn_frontend/node_interface.h b/include/cudnn_frontend/node_interface.h index ab72a9030..8d46ece33 100644 --- a/include/cudnn_frontend/node_interface.h +++ b/include/cudnn_frontend/node_interface.h @@ -1,6 +1,5 @@ #pragma once -#include #include #include #include @@ -97,8 +96,7 @@ class INode { }; virtual error_t - collect_tensors_to_dump_node( - std::map, char>>&) const { + collect_tensors_to_dump_node(std::vector, char>>&) const { return {error_code_t::OK, ""}; }; @@ -353,8 +351,7 @@ class INode { error_t collect_tensors_to_dump_subtree( - std::map, char>>& tensors_to_dump) - const { + std::vector, char>>& tensors_to_dump) const { CHECK_CUDNN_FRONTEND_ERROR(collect_tensors_to_dump_node(tensors_to_dump)); for (auto const& sub_node : sub_nodes) { CHECK_CUDNN_FRONTEND_ERROR(sub_node->collect_tensors_to_dump_subtree(tensors_to_dump)); From a0cbef7bbb05389d9356181cdc7c779f89fc085f Mon Sep 17 00:00:00 2001 From: Hwanseo Choi Date: Mon, 13 Jul 2026 13:15:39 -0700 Subject: [PATCH 9/9] Remove MXFP8 performance smoke test --- test/python/sdpa/mxfp8.py | 3 ++- test/python/test_mhas_v2.py | 28 +++------------------------- 2 files changed, 5 insertions(+), 26 deletions(-) diff --git a/test/python/sdpa/mxfp8.py b/test/python/sdpa/mxfp8.py index a5b4b8f8e..c5fd42b29 100644 --- a/test/python/sdpa/mxfp8.py +++ b/test/python/sdpa/mxfp8.py @@ -529,10 +529,11 @@ def generate_graph_bwd(b, h_q, h_k, h_v, return graph_bwd -def exec_sdpa_mxfp8(cfg, request, cudnn_handle, perf): +def exec_sdpa_mxfp8(cfg, request, cudnn_handle): """Execute MXFP8 SDPA test.""" if request.config.option.dryrun: pytest.skip("dry run mode") + perf = request.config.getoption("--perf") cudnn_version = LooseVersion(cudnn.backend_version_string()) diff --git a/test/python/test_mhas_v2.py b/test/python/test_mhas_v2.py index 68ea4192b..746cf2d1d 100644 --- a/test/python/test_mhas_v2.py +++ b/test/python/test_mhas_v2.py @@ -966,7 +966,7 @@ def test_sdpa_mxfp8_fwd_L0(env_info, test_no, request, cudnn_handle): if request.node.name in test.blocked_tests: pytest.skip(f"blocked test: {request.node.name}") try: - exec_sdpa_mxfp8(test.cfg, request, cudnn_handle, request.config.getoption("--perf")) + exec_sdpa_mxfp8(test.cfg, request, cudnn_handle) finally: if "CUDNN_UNFUSE_FMA" in os.environ: del os.environ["CUDNN_UNFUSE_FMA"] @@ -1015,33 +1015,11 @@ def test_sdpa_mxfp8_bwd_L0(env_info, test_no, request, cudnn_handle): if request.node.name in test.blocked_tests: pytest.skip(f"blocked test: {request.node.name}") try: - exec_sdpa_mxfp8(test.cfg, request, cudnn_handle, request.config.getoption("--perf")) + exec_sdpa_mxfp8(test.cfg, request, cudnn_handle) finally: if "CUDNN_RESCALE_THRESHOLD" in os.environ: del os.environ["CUDNN_RESCALE_THRESHOLD"] - -@pytest.mark.L0 -def test_sdpa_mxfp8_perf_L0(request, cudnn_handle): - cfg = ExecConfig( - data_type=torch.float8_e4m3fn, - output_type=torch.bfloat16, - rng_data_seed=1, - is_infer=False, - is_determin=False, - is_mxfp8=True, - batches=1, - h_q=1, - h_k=1, - h_v=1, - s_q=128, - s_kv=128, - d_qk=64, - d_v=64, - ) - cfg.fill_derived_fields() - exec_sdpa_mxfp8(cfg, request, cudnn_handle, perf=True) - # # =================== # # Single repro test # # =================== @@ -1072,7 +1050,7 @@ def test_repro(env_info, request, cudnn_handle): try: if cfg.cfg.is_mxfp8: - exec_sdpa_mxfp8(cfg.cfg, request, cudnn_handle, request.config.getoption("--perf")) + exec_sdpa_mxfp8(cfg.cfg, request, cudnn_handle) elif cfg.cfg.data_type in (torch.float8_e4m3fn, torch.float8_e5m2): exec_sdpa_fp8(cfg.cfg, request, cudnn_handle) else: