FMHA + analytic sink bias correction using LSE

Instead of SDPA with virtual sink position, use the production FMHA
kernel and apply the sink bias as a post-hoc correction on the output.

The correction is: O_sink = O_raw * exp(lse) / (exp(lse) + exp(sink))

This simulates the attention sink (paper D5c) without modifying the
FMHA kernel. The sink absorbs some attention mass, reducing the
normalization constant and scaling down the output.

single_shot_inference.py

@@ -366,43 +366,36 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
     k_full, v_full = kv_cache.get()  # (1, seq_len, hd) each — RoPE'd, K=V
     seq_len = k_full.shape[1]
 
-    # -- FMHA: (n_h, T, hd) × (1, seq_len, hd) → (n_h, T, hd) --
+    # -- FMHA with sink bias correction (paper D5c) --
     q_input = q_heads.permute(1, 0, 2)  # (n_h, T, hd)
-    # Use PyTorch SDPA for correctness verification
-    USE_SDPA = False  # Use production FMHA kernel (better residual, no sinks)
-    if USE_SDPA:
-        # 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)
-    attn_out = attn_out.permute(1, 0, 2)  # (T, n_h, hd)
+    from dsv4.kernels.attention.production import dsv4_attention
+    # Get both output and LSE from the kernel
+    from dsv4.kernels.attention.fmha_multitile_op import fmha_multitile_decode_raw
+    
+    scale = 1.0 / math.sqrt(hd)
+    q_4d = q_input.unsqueeze(0).contiguous()  # (1, n_h, T, hd)
+    k_4d = k_full.unsqueeze(0).contiguous()  # (1, 1, seq_len, hd)
+    v_4d = v_full.unsqueeze(0).contiguous()  # (1, 1, seq_len, hd)
+    
+    o_4d, lse = fmha_multitile_decode_raw(q_4d, k_4d, v_4d, scale)  # (1, n_h, T, hd), (1, n_h, T)
+    attn_out = o_4d.squeeze(0).permute(1, 0, 2)  # (T, n_h, hd)
+    
+    # Apply sink bias correction: scale the output by softmax_normalizer / (normalizer + exp(sink))
+    # This simulates adding a virtual sink position with V=0 and logit=sink
+    # O_sink = O_raw * exp(lse) / (exp(lse) + exp(sink))
+    sink_key = f"{pre}.sinks"
+    if sink_key in w and seq_len > 0:
+        sinks = w[sink_key].to(device=device)  # (n_h,) BF16
+        # lse: (1, n_h, T) — log-sum-exp of attention scores per head
+        lse_3d = lse.squeeze(0).permute(1, 0)  # (T, n_h)
+        # For each head, compute the correction factor
+        sink_exp = torch.exp(sinks.float())  # (n_h,)
+        attn_exp = torch.exp(lse_3d.float())  # (T, n_h)
+        # Correction: attn_exp / (attn_exp + sink_exp)
+        correction = attn_exp / (attn_exp + sink_exp.unsqueeze(0) + 1e-10)  # (T, n_h)
+        # Apply per-head correction
+        attn_out = attn_out * correction.unsqueeze(-1)  # (T, n_h, hd) * (T, n_h, 1)
+    
 
     # -- Inverse RoPE on attention output (paper §2.3.3) --
     attn_out = apply_inverse_rope(attn_out, positions_dev, rope_cos, rope_sin, hd, rd)