diag: per-segment reference comparison for D5c multi-tile

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__':