From 45a0fb997103da341672ebcf725bff0209dd56f3 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Fri, 22 May 2026 18:01:11 +0000 Subject: [PATCH] Add NaN/inf checking to stage C test --- tests/unit/test_fmha_v3_stage_c_full.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit/test_fmha_v3_stage_c_full.py b/tests/unit/test_fmha_v3_stage_c_full.py index a2e6cf4a..7b5707c0 100644 --- a/tests/unit/test_fmha_v3_stage_c_full.py +++ b/tests/unit/test_fmha_v3_stage_c_full.py @@ -449,6 +449,8 @@ def test(): torch.cuda.synchronize() out = c[:, :, 0].float() + has_nan = out.isnan().any().item() + has_inf = out.isinf().any().item() cos = torch.nn.functional.cosine_similarity( out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0) ).item() @@ -456,6 +458,7 @@ def test(): n_tiles = n // 128 print(f'FMHA Stage-C n={n} ({n_tiles} kv tiles): ' f'cos {cos:.6f} max_abs {max_abs:.4f} ' + f'nan={has_nan} inf={has_inf} ' f'{"PASS" if cos >= 0.99 else "FAIL"}') if cos < 0.99: print(f' out[0,:4]={out[0,:4].tolist()}')