From e5381b7312f71dbb5c4806d279c09dcc3811975e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Tue, 26 May 2026 15:08:40 +0000 Subject: [PATCH] diag: add baseline test (s_k=256 D3 mask, no sink bias) to isolate D5c issue --- tests/unit/test_d5c_fused.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/unit/test_d5c_fused.py b/tests/unit/test_d5c_fused.py index 3fc549a7..b72ac839 100644 --- a/tests/unit/test_d5c_fused.py +++ b/tests/unit/test_d5c_fused.py @@ -298,5 +298,36 @@ def test_d5c_with_causal(): if __name__ == '__main__': + # First: baseline test without sink bias (just multi-tile D3 masking) + print('=== Baseline: multi-tile attention with D3 mask (no D5c) ===\n') + _hd = 64; _m = 128; _s_k = 256; _swa_len = 192 # 192 valid, 64 masked + _scale = 1.0 / math.sqrt(_hd) + torch.manual_seed(42) + _q = torch.randn(_m, _hd, 1, dtype=torch.bfloat16, device='cuda') + _k = torch.randn(_s_k, _hd, 1, dtype=torch.bfloat16, device='cuda') + _v = torch.randn(_s_k, _hd, dtype=torch.bfloat16, device='cuda') + + # Reference + _qf = _q[:, :, 0].float(); _kf = _k[:, :, 0].float(); _vf = _v.float() + _scores = _qf @ _kf.T * _scale + _scores[:, _swa_len:] = float('-inf') + _ref = (torch.softmax(_scores, dim=-1) @ _vf).to(torch.bfloat16) + + # Kernel + _stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + _kernel = FmhaKernel(head_dim=_hd, s_k=_s_k, normalize=False, apply_swa_mask=True, is_causal=False, n_comp=0) + _c_out = torch.zeros(_m, _hd, 1, dtype=torch.bfloat16, device='cuda') + _lse = torch.zeros(_m, 1, 1, dtype=torch.float32, device='cuda') + _rs = torch.zeros(_m, 1, 1, dtype=torch.float32, device='cuda') + def _tc(t): return ct.from_dlpack(t).mark_layout_dynamic(leading_dim=ct.get_leading_dim(t)) + _mQ = _tc(_q); _mK = _tc(_k); _mV = _tc(_v.unsqueeze(-1).contiguous()) + _mC = _tc(_c_out); _mLSE = _tc(_lse); _mRS = _tc(_rs) + _compiled = cute.compile(_kernel, _mQ, _mK, _mV, _mC, _stream, _mLSE, swa_len=_swa_len, row_sums=_mRS) + _compiled(_mQ, _mK, _mV, _mC, _stream, _mLSE, swa_len=_swa_len, row_sums=_mRS) + torch.cuda.synchronize() + _o_k = _c_out[:, :, 0].float() / _rs[:, 0, 0].float().unsqueeze(1).clamp(min=1e-30) + _cos = torch.nn.functional.cosine_similarity(_o_k.flatten().unsqueeze(0), _ref.flatten().unsqueeze(0).float()).item() + print(f'Baseline (s_k=256, D3 mask, n_comp=0): cos {_cos:.6f} {"PASS" if _cos > 0.99 else "FAIL"}\n') + test_d5c_combined() test_d5c_with_causal()