diff --git a/tests/unit/test_d5c_multitile.py b/tests/unit/test_d5c_multitile.py index 88fba532..b0a39544 100644 --- a/tests/unit/test_d5c_multitile.py +++ b/tests/unit/test_d5c_multitile.py @@ -247,9 +247,31 @@ def test_d5c_multitile(): if cos < 0.99: print(f' kernel[0,:4]={o_merged[0,:4].tolist()}') print(f' ref[0,:4]={ref[0,:4].tolist()}') - # Debug: check each segment individually + # Debug: check each segment individually against per-segment reference for i, (o_n, lse) in enumerate(segment_results): - print(f' Seg {i}: LSE[0]={lse[0].item():.4f} O_norm[0,:2]={o_n[0,:2].tolist()}') + k_start = i * seg_size + k_end = min(k_start + seg_size, n_total) + k_seg = k_combined[k_start:k_end] + v_seg = v_combined[k_start:k_end] + n_comp_local = min(max(n_comp - k_start, 0), seg_size) + # Per-segment reference + scores_ref = q.float() @ k_seg.float().T * scale + if k_start < n_comp: + scores_ref[:, n_comp - k_start:] += attn_sink_val + swa_end = n_comp + swa_len - k_start + if swa_end < seg_size: + scores_ref[:, swa_end:] = float('-inf') + else: + scores_ref += attn_sink_val + swa_len_in_seg = n_comp + swa_len - k_start + if swa_len_in_seg < seg_size: + scores_ref[:, swa_len_in_seg:] = float('-inf') + o_ref_seg = (torch.softmax(scores_ref, dim=-1) @ v_seg.float()).to(torch.bfloat16) + lse_ref = torch.logsumexp(scores_ref, dim=-1) + cos_seg = torch.nn.functional.cosine_similarity( + o_n.flatten().unsqueeze(0), o_ref_seg.flatten().unsqueeze(0).float() + ).item() + print(f' Seg {i}: LSE_kern={lse[0].item():.4f} LSE_ref={lse_ref[0].item():.4f} cos={cos_seg:.6f}') if __name__ == '__main__':