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") 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 # # ===================