add attention sinks to SDPA path (paper D5c)

single_shot_inference.py

@@ -487,22 +487,40 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
 
     # -- FMHA: (n_h, T, hd) × (1, seq_len, hd) → (n_h, T, hd) --
     q_input = q_heads.permute(1, 0, 2)  # (n_h, T, hd)
-    # Use PyTorch SDPA for correctness verification (production FMHA has no sink bias)
-    # dsv4_attention can be swapped back once sink bias is integrated into the kernel
-    USE_SDPA = True  # Set False to use production FMHA kernel
+    # Use PyTorch SDPA for correctness verification
+    USE_SDPA = True
     if USE_SDPA:
-        # PyTorch scaled_dot_product_attention
-        # q: (n_h, T, hd), k: (1, seq_len, hd), v: (1, seq_len, hd)
-        # Need to expand K/V for GQA: (1, seq_len, hd) → (n_h, seq_len, hd)
-        k_expanded = k_full.expand(n_h, -1, -1)  # (n_h, seq_len, hd)
-        v_expanded = v_full.expand(n_h, -1, -1)
-        scale = 1.0 / math.sqrt(hd)
-        attn_out = torch.nn.functional.scaled_dot_product_attention(
-            q_input, k_expanded, v_expanded,
-            scale=scale, is_causal=False)  # (n_h, T, hd)
+        # Expand K/V for GQA: (1, seq_len, hd) → (n_h, seq_len, hd)
+        k_expanded = k_full.expand(n_h, -1, -1).contiguous()  # (n_h, seq_len, hd)
+        v_expanded = v_full.expand(n_h, -1, -1).contiguous()
+
+        # Add attention sink (paper §2.3.3, D5c)
+        # The sink is a per-head logit bias added to a virtual position.
+        # We simulate it by appending a zero-valued KV position with the sink logit.
+        sink_key = f"{pre}.sinks"
+        if sink_key in w and seq_len > 0:
+            sinks = w[sink_key].to(device=device)  # (n_h,) BF16
+            # Append zero KV entry for the sink
+            sink_k = torch.zeros(n_h, 1, hd, dtype=torch.bfloat16, device=device)
+            sink_v = torch.zeros(n_h, 1, hd, dtype=torch.bfloat16, device=device)
+            k_with_sink = torch.cat([k_expanded, sink_k], dim=1)  # (n_h, seq_len+1, hd)
+            v_with_sink = torch.cat([v_expanded, sink_v], dim=1)
+            # Create attention bias: sink logit added to the last position for each head
+            # attn_mask shape: (n_h, T, seq_len+1)
+            sink_bias_mask = torch.zeros(n_h, T, seq_len + 1, dtype=torch.bfloat16, device=device)
+            for h in range(n_h):
+                sink_bias_mask[h, :, -1] = sinks[h]  # Add sink logit to sink position
+            attn_out = torch.nn.functional.scaled_dot_product_attention(
+                q_input, k_with_sink, v_with_sink,
+                attn_mask=sink_bias_mask,
+                scale=1.0 / math.sqrt(hd))
+        else:
+            attn_out = torch.nn.functional.scaled_dot_product_attention(
+                q_input, k_expanded, v_expanded,
+                scale=1.0 / math.sqrt(hd), is_causal=False)
     else:
         from dsv4.kernels.attention.production import dsv4_attention
-        attn_out = dsv4_attention(q_input, k_full, v_full)  # (n_h, T, hd)
+        attn_out = dsv4_attention(q_input, k_full, v_full)
     attn_out = attn_out.permute(1, 0, 2)  # (T, n_h, hd)
 
     # -- Inverse RoPE on attention output (paper §2.3.3) --