fix(attention): zero empty decode rows in unified attention

aiter/ops/triton/attention/unified_attention.py

@@ -201,6 +201,29 @@ def use_2d_kernel(params: _UAParams, backend: str = "triton"):
     )
 
 
+def _zero_empty_decode_rows(output: torch.Tensor, params: _UAParams):
+    """Make zero-KV decode padding rows deterministic instead of NaN."""
+    if params.all_decode:
+        empty = params.seqused_k == 0
+        empty = empty.view(-1, *([1] * (output.ndim - 1)))
+        output.masked_fill_(empty, 0)
+    return output
+
+
+def _neutralize_empty_decode_partials(
+    segm_output: torch.Tensor,
+    segm_max: torch.Tensor,
+    segm_expsum: torch.Tensor,
+    params: _UAParams,
+):
+    if params.all_decode:
+        empty = params.seqused_k == 0
+        segm_output.masked_fill_(empty[:, None, None, None], 0)
+        segm_max.masked_fill_(empty[:, None, None], float("-inf"))
+        segm_expsum.masked_fill_(empty[:, None, None], 0)
+    return segm_output, segm_max, segm_expsum
+
+
 def unified_attention(
     q,
     k,
@@ -467,9 +490,11 @@ def unified_attention(
             )
 
         if NUM_SEGMENTS == 1:
-            return segm_output
+            return _zero_empty_decode_rows(segm_output, params)
         elif skip_reduce:
-            return segm_output, segm_max, segm_expsum
+            return _neutralize_empty_decode_partials(
+                segm_output, segm_max, segm_expsum, params
+            )
 
         # Gluon reduce (one workgroup/token, in-wave segment merge); valid for all-decode with small split counts, else the Triton reduce_segments.
         use_gluon_reduce = (
@@ -497,7 +522,7 @@ def unified_attention(
                 NUM_SEGMENTS,
                 TILE_SIZE,
             )
-    return out
+    return _zero_empty_decode_rows(out, params)
 
 
 def is_2d_gluon_available(params: _UAParams, backend: str):

op_tests/triton_tests/attention/test_unified_attention.py

@@ -548,6 +548,84 @@ def test_triton_unified_attn_3d(
     )
 
 
+@torch.inference_mode()
+def test_unified_attention_zero_kv_decode_rows() -> None:
+    if DEVICE_ARCH != "gfx1250":
+        pytest.skip("zero-KV padding regression covers the gfx1250 Gluon path")
+
+    seq_lens = [(1, 2048), (1, 0), (1, 0), (1, 0)]
+    query_lens = [query_len for query_len, _ in seq_lens]
+    (
+        query,
+        key_cache,
+        value_cache,
+        shuffled_key_cache,
+        shuffled_value_cache,
+        _sinks,
+        output,
+        cu_query_lens,
+        kv_lens,
+        max_query_len,
+        max_kv_len,
+        scale,
+        window_size,
+        block_tables,
+        _maybe_quant_query,
+        _query_scales,
+        q_descale,
+        k_descale,
+        v_descale,
+        output_scale,
+    ) = generate_data(
+        seq_lens=seq_lens,
+        num_blocks=128,
+        block_size=64,
+        head_size=64,
+        num_heads=(8, 1),
+        q_dtype=torch.bfloat16,
+        kv_dtype=torch.bfloat16,
+        out_dtype=torch.bfloat16,
+        shuffled_kv_cache=True,
+        device="cuda",
+    )
+
+    unified_attention(
+        q=query,
+        k=shuffled_key_cache,
+        v=shuffled_value_cache,
+        out=output,
+        cu_seqlens_q=cu_query_lens,
+        seqused_k=kv_lens,
+        max_seqlen_q=max_query_len,
+        max_seqlen_k=max_kv_len,
+        softmax_scale=scale,
+        causal=True,
+        window_size=window_size,
+        block_table=block_tables,
+        softcap=0,
+        q_descale=q_descale,
+        k_descale=k_descale,
+        v_descale=v_descale,
+        output_scale=output_scale,
+        shuffled_kv_cache=True,
+        backend="gluon",
+    )
+
+    ref_output = ref_paged_attn(
+        query=query,
+        key_cache=key_cache,
+        value_cache=value_cache,
+        query_lens=query_lens,
+        kv_lens=kv_lens,
+        block_tables=block_tables,
+        scale=scale,
+        out_dtype=torch.bfloat16,
+    )
+    torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1e-2)
+    assert torch.isfinite(output).all()
+    assert torch.count_nonzero(output[1:]) == 0
+
+
 @pytest.mark.parametrize(
     "seq_lens",
     [