diff --git a/tests/unit/test_fmha_v3_stage_d1.py b/tests/unit/test_fmha_v3_stage_d1.py index 33819e6f..2302271b 100644 --- a/tests/unit/test_fmha_v3_stage_d1.py +++ b/tests/unit/test_fmha_v3_stage_d1.py @@ -12,7 +12,7 @@ import cuda.bindings.driver as cuda from dsv4.kernels.attention.fmha import FmhaKernel -def test_head_dim(hd, n_kv): +def test_head_dim(hd, n_kv=128): m = 128 torch.manual_seed(42) @@ -21,31 +21,28 @@ def test_head_dim(hd, n_kv): v = torch.randn(n_kv, hd, dtype=torch.bfloat16, device='cuda') c = torch.zeros(m, hd, 1, dtype=torch.bfloat16, device='cuda') - # FP32 reference (normalized) + # FP32 reference qf = q[:, :, 0].float() kf = k[:, :, 0].float() scale = 1.0 / math.sqrt(hd) - attn = qf @ kf.T * scale - attn_max = attn.max(dim=-1, keepdim=True)[0] - attn_exp = torch.exp(attn - attn_max) + attn_max = (qf @ kf.T * scale).max(dim=-1, keepdim=True)[0] + attn_exp = torch.exp(qf @ kf.T * scale - attn_max) attn_sum = attn_exp.sum(dim=-1, keepdim=True) - attn_norm = attn_exp / attn_sum - ref_norm = attn_norm @ v.float() - - # FP32 reference (un-normalized): O_unnorm = sum(exp(S - max) * V) + ref_norm = (attn_exp / attn_sum) @ v.float() ref_unnorm = attn_exp @ v.float() - - # Reference LSE: lse = ln(row_sum) + max - ref_lse = torch.log(attn_sum.squeeze(-1)) + attn_max.squeeze(-1) # (m,) + ref_lse = (torch.log(attn_sum.squeeze(-1)) + attn_max.squeeze(-1))[0].item() lse_tensor = torch.zeros(m, 1, 1, dtype=torch.float32, device='cuda') - kernel = FmhaKernel(head_dim=hd, s_k=n_kv) + # Use TMEM-P for all hd (Saves SMEM, and works correctly after qk_mma_tiler fix) + # SMEM-P is available for hd>64 but requires more SMEM budget + kernel = FmhaKernel(head_dim=hd, s_k=n_kv, use_smem_p=False) pv_n_tile = kernel.pv_n_tile n_pv_tiles = kernel.n_pv_tiles stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + # Compile once v_tile = v[:, 0:pv_n_tile].contiguous() v_kernel = v_tile.unsqueeze(-1) c_tile = torch.zeros(m, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda') @@ -82,64 +79,34 @@ def test_head_dim(hd, n_kv): lse_val = lse_tensor[0, 0, 0].item() out_unnorm = c[:, :, 0].float() - - # Compare un-normalized O against reference + out_norm = out_unnorm / attn_sum cos_unnorm = torch.nn.functional.cosine_similarity( out_unnorm.flatten().unsqueeze(0), ref_unnorm.flatten().unsqueeze(0) ).item() - - # Normalize externally: O_norm = O_unnorm / row_sum - # row_sum = exp(lse - max) where max is already incorporated in O_unnorm - # For the D5 merge, we use exp(lse) directly. - # For standalone test: O_norm = O_unnorm * (1/row_sum) - # where row_sum per row = O_unnorm row doesn't work. We need lse. - # lse = ln(row_sum) + max, so row_sum = exp(lse - max) - # But max is the same max used in the softmax, and O_unnorm already has - # the exp(-max) scaling baked in. So: - # O_norm = O_unnorm / row_sum - # We can compute row_sum from O_unnorm by checking what row_sum should be. - # Since O_unnorm[i,j] = sum_k(P[i,k] * V[k,j]) where P = exp(S*s - max) - # and row_sum = sum_k(exp(S*s - max)), - # we can normalize: O_norm[i] = O_unnorm[i] / row_sum[i] - # But we can't easily get row_sum from O_unnorm alone. - # Use LSE instead: row_sum = exp(lse - max_in_nat) - # where max_in_nat = row_max * ln(2) but we only have lse. - # Actually for the merge: we just need exp(lse) * O_unnorm. - # For standalone: compute row_sum from attention explicitly. - # ref_row_sum = attn_sum.squeeze(-1) # (m,) - # O_norm = O_unnorm / ref_row_sum.unsqueeze(1) - # This uses the reference row_sum to normalize — verifies the O_unnorm is correct. - out_norm_using_ref = out_unnorm / attn_sum # (m, hd) cos_norm = torch.nn.functional.cosine_similarity( - out_norm_using_ref.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0) + out_norm.flatten().unsqueeze(0), ref_norm.flatten().unsqueeze(0) ).item() + lse_err = abs(lse_val - ref_lse) if lse_val is not None else float('inf') - # Verify LSE - ref_lse_val = ref_lse[0].item() - lse_err = abs(lse_val - ref_lse_val) if lse_val is not None else float('inf') - - status = "PASS" if cos_unnorm >= 0.99 else ("WARN" if cos_unnorm >= 0.97 else "FAIL") - print(f'hd={hd}, n={n_kv}: cos_unnorm {cos_unnorm:.6f} cos_norm(ref_sum) {cos_norm:.6f} lse_err {lse_err:.6f} {status}') + status = "PASS" if cos_unnorm >= 0.99 else "FAIL" + print(f'hd={hd}, n={n_kv}: cos_unnorm {cos_unnorm:.6f} cos_norm {cos_norm:.6f} lse_err {lse_err:.6f} {status}') return cos_unnorm def test(): print("=== Stage D1: Parameterized HEAD_DIM ===") - print("(Kernel outputs un-normalized O + LSE)\n") + print("(Kernel outputs un-normalized O + LSE; TMEM-P path)\n") - print("--- Regression: HEAD_DIM=64 ---") cos64 = test_head_dim(64, 128) - - print("\n--- HEAD_DIM=256 (single PV tile) ---") + cos128 = test_head_dim(128, 128) cos256 = test_head_dim(256, 128) - - print("\n--- HEAD_DIM=512 (2 PV tiles) ---") cos512 = test_head_dim(512, 128) print("\n=== Summary ===") - print(f"hd=64, n=128: cos_unnorm={cos64:.6f} {'PASS' if cos64 >= 0.99 else 'FAIL'}") - print(f"hd=256, n=128: cos_unnorm={cos256:.6f} {'PASS' if cos256 >= 0.99 else 'FAIL'}") - print(f"hd=512, n=128: cos_unnorm={cos512:.6f} {'PASS' if cos512 >= 0.99 else 'FAIL'}") + print(f"hd=64: cos={cos64:.6f} {'PASS' if cos64 >= 0.99 else 'FAIL'}") + print(f"hd=128: cos={cos128:.6f} {'PASS' if cos128 >= 0.99 else 'FAIL'}") + print(f"hd=256: cos={cos256:.6f} {'PASS' if cos256 >= 0.99 else 'FAIL'}") + print(f"hd=512: cos={cos512:.6f} {'PASS' if cos512 >= 0.99 else 'FAIL'}") if __name__ == '__main__':