Split single-segment (normalized) and multi-segment (KV merge) paths

dsv4/kernels/attention/production.py

@@ -123,6 +123,58 @@ def dsv4_attention(
     return output
 
 
+def _attention_single_head_normalized(
+    q: torch.Tensor,      # (1, T, hd)
+    k: torch.Tensor,      # (1, N, hd)
+    v: torch.Tensor,      # (1, N, hd)
+    scale: float,
+    swa_len: int = None,
+    is_causal: bool = False,
+    n_comp: int = 0,
+    sink_bias: torch.Tensor = None,
+    use_smem_p: bool = False,
+) -> torch.Tensor:
+    """Run FMHA for a single head with in-kernel normalization (single KV tile)."""
+    _, T, hd = q.shape
+    N = k.shape[1]
+    apply_swa_mask = swa_len is not None
+    apply_sink_bias = sink_bias is not None
+
+    compiled, kernel = _get_or_compile_kernel(
+        head_dim=hd, s_k=N, use_smem_p=use_smem_p,
+        normalize=True, apply_swa_mask=apply_swa_mask,
+        is_causal=is_causal, n_comp=n_comp if n_comp > 0 else None,
+        apply_sink_bias=apply_sink_bias,
+    )
+    pv_n_tile = kernel.pv_n_tile
+    n_pv_tiles = kernel.n_pv_tiles
+
+    output = torch.zeros(T, hd, dtype=torch.bfloat16, device='cuda')
+
+    for nt in range(n_pv_tiles):
+        v_start = nt * pv_n_tile
+        v_end = v_start + pv_n_tile
+        v_tile = v[0, :, v_start:v_end].contiguous()
+        v_kernel = v_tile.unsqueeze(-1)
+        c_tile = torch.zeros(T, pv_n_tile, 1, dtype=torch.bfloat16, device='cuda')
+
+        q_input = q[0].contiguous().unsqueeze(-1)
+        k_input = k[0].contiguous().unsqueeze(-1)
+        stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+
+        mQ = ct.from_dlpack(q_input).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q_input))
+        mK = ct.from_dlpack(k_input).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k_input))
+        mV = ct.from_dlpack(v_kernel).mark_layout_dynamic(leading_dim=ct.get_leading_dim(v_kernel))
+        mC = ct.from_dlpack(c_tile).mark_layout_dynamic(leading_dim=ct.get_leading_dim(c_tile))
+
+        compiled(mQ, mK, mV, mC, stream)
+        torch.cuda.synchronize()
+
+        output[:, v_start:v_end] = c_tile[:, :, 0]
+
+    return output.unsqueeze(0)  # (1, T, hd)
+
+
 def _attention_single_head(
     q: torch.Tensor,      # (1, T, hd)
     k: torch.Tensor,      # (1, N, hd)
@@ -149,12 +201,15 @@ def _attention_single_head(
     attn_sum = attn_exp.sum(dim=-1, keepdim=True)
     ref_norm = (attn_exp / attn_sum) @ vf
 
-    # Kernel execution with Python KV merge for multi-KV-tile
     # Each segment is s_k=128 (one KV tile)
     s_k_per_seg = 128
     n_segments = (N + s_k_per_seg - 1) // s_k_per_seg
 
-    # Use un-normalized output + LSE for Python KV merge
+    if n_segments == 1:
+        # Single segment: use kernel with normalize=True (in-kernel normalization)
+        return _attention_single_head_normalized(q, k, v, scale, swa_len, is_causal, n_comp, sink_bias, use_smem_p)
+
+    # Multi-segment: use un-normalized output + Python KV merge
     compiled, kernel = _get_or_compile_kernel(
         head_dim=hd, s_k=s_k_per_seg, use_smem_p=use_smem_p,
         normalize=False, apply_swa_mask=apply_swa_mask,