[Feature][Attention] Cover a whole spec-verify in one DiffKV split-KV program (#59085)

Signed-off-by: Giulio De Pasquale <git@depasquale.giugl.io>

tests/kernels/attention/test_triton_unified_attention_diffkv.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Unit tests for the Triton DiffKV unified-attention kernel."""
 
+from unittest.mock import patch
+
 import pytest
 import torch
 
@@ -16,6 +18,7 @@ from vllm.v1.attention.backends.fa_utils import (
     get_flash_attn_version,
     is_flash_attn_varlen_func_available,
 )
+from vllm.v1.attention.ops import triton_unified_attention_diffkv as diffkv_module
 from vllm.v1.attention.ops.triton_unified_attention_diffkv import (
     unified_attention_diffkv,
 )
@@ -258,3 +261,95 @@ def test_triton_unified_attn_diffkv_fp8_query(seq_lens, seq_threshold_3D):
         seq_threshold_3D=seq_threshold_3D,
         quant_query=True,
     )
+
+
+def _capture_diffkv_launch(
+    query_lens: list[int], max_seqlen_q: int, segm_first_dim: int
+):
+    """Run the launcher with both Triton kernels mocked; return grid/kwargs."""
+    num_query_heads, num_kv_heads = 8, 2
+    head_size_qk, head_size_v = 128, 128
+    num_seqs = len(query_lens)
+    num_tokens = sum(query_lens)
+
+    query = torch.zeros(
+        num_tokens, num_query_heads, head_size_qk, dtype=torch.bfloat16, device="cpu"
+    )
+    kv_cache = torch.zeros(
+        NUM_BLOCKS,
+        BLOCK_SIZES[0],
+        num_kv_heads,
+        head_size_qk + head_size_v,
+        dtype=torch.bfloat16,
+        device="cpu",
+    )
+    key_cache = kv_cache[..., :head_size_qk]
+    value_cache = kv_cache[..., head_size_qk:]
+    out = torch.zeros(
+        num_tokens, num_query_heads, head_size_v, dtype=torch.bfloat16, device="cpu"
+    )
+    cu_seqlens_q = torch.tensor(
+        [0] + query_lens, dtype=torch.int32, device="cpu"
+    ).cumsum(dim=0, dtype=torch.int32)
+    seqused_k = torch.full((num_seqs,), 512, dtype=torch.int32, device="cpu")
+    block_tables = torch.zeros(num_seqs, 32, dtype=torch.int32, device="cpu")
+    segm_output, segm_max, segm_expsum = _alloc_segm_buffers(
+        segm_first_dim, num_query_heads, head_size_v
+    )
+
+    with (
+        patch.object(diffkv_module, "kernel_unified_attention_diffkv") as mock_kernel,
+        patch.object(diffkv_module, "kernel_reduce_segments_diffkv") as mock_reduce,
+    ):
+        unified_attention_diffkv(
+            q=query,
+            k=key_cache,
+            v=value_cache,
+            out=out,
+            cu_seqlens_q=cu_seqlens_q,
+            seqused_k=seqused_k,
+            softmax_scale=head_size_qk**-0.5,
+            causal=True,
+            window_size=(-1, -1),
+            block_table=block_tables,
+            softcap=0,
+            max_seqlen_q=max_seqlen_q,
+            seq_threshold_3D=8,
+            num_par_softmax_segments=NUM_PAR_SOFTMAX_SEGMENTS,
+            softmax_segm_output=segm_output,
+            softmax_segm_max=segm_max,
+            softmax_segm_expsum=segm_expsum,
+        )
+
+    grid = mock_kernel.__getitem__.call_args.args[0]
+    kwargs = mock_kernel.__getitem__.return_value.call_args.kwargs
+    return grid, kwargs, mock_reduce.__getitem__.called
+
+
+def test_spec_verify_block_m_grouping_and_launch_kwargs() -> None:
+    """Pin program granularity for verify/decode without launching Triton."""
+    verify_grid, verify_kwargs, verify_reduce = _capture_diffkv_launch(
+        query_lens=[8, 8], max_seqlen_q=8, segm_first_dim=16
+    )
+    num_queries_per_kv = 4
+    expected_block_m = min(128, next_power_of_2(8 * num_queries_per_kv))
+    assert verify_grid == (4, 2, NUM_PAR_SOFTMAX_SEGMENTS)
+    assert verify_kwargs["IS_3D"] is True
+    assert verify_kwargs["BLOCK_M"] == expected_block_m
+    assert verify_kwargs["BLOCK_Q"] == expected_block_m // num_queries_per_kv
+    assert verify_kwargs["TILE_SIZE"] == 16
+    assert verify_kwargs["num_warps"] == 4
+    assert "num_stages" not in verify_kwargs
+    assert verify_reduce
+
+    decode_grid, decode_kwargs, decode_reduce = _capture_diffkv_launch(
+        query_lens=[1], max_seqlen_q=1, segm_first_dim=8
+    )
+    assert decode_grid == (1, 2, NUM_PAR_SOFTMAX_SEGMENTS)
+    assert decode_kwargs["IS_3D"] is True
+    assert decode_kwargs["BLOCK_M"] == 16
+    assert decode_kwargs["BLOCK_Q"] == 4
+    assert decode_kwargs["TILE_SIZE"] == 16
+    assert "num_warps" not in decode_kwargs
+    assert "num_stages" not in decode_kwargs
+    assert decode_reduce

vllm/v1/attention/ops/triton_unified_attention_diffkv.py

@@ -46,6 +46,13 @@ logger = init_logger(__name__)
 
 is_batch_invariant = envs.VLLM_BATCH_INVARIANT
 
+# Whole-verify split-KV grouping: one program covers all q tokens of a
+# spec-decode verify request (BLOCK_M = q_len * GQA group, capped at 128).
+_SPEC_3D_MAX_Q = 16
+_SPEC_3D_BLOCK_M = 128
+_SPEC_3D_NUM_WARPS = 8
+_SPEC_3D_TILE = 16
+
 
 @triton.jit
 def kernel_unified_attention_diffkv(
@@ -290,6 +297,9 @@ def kernel_unified_attention_diffkv(
             acc,
             mask=dim_mask_v[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
         )
+        # A row whose keys are all masked in this segment ends at M=0, L=0;
+        # -inf makes reduce_segments treat that segment as absent.
+        M = tl.where(L > 0.0, M, float("-inf"))
         store_segm_reduce_scalars(
             segm_max_ptr,
             segm_expsum_ptr,
@@ -365,6 +375,8 @@ def kernel_reduce_segments_diffkv(
     )
     segm_max = tl.load(segm_max_ptr + segm_offset, mask=segm_mask, other=float("-inf"))
     overall_max = tl.max(segm_max)
+    # Graph-capture dummy lengths can leave a query with no causal keys.
+    overall_max = tl.where(overall_max > float("-inf"), overall_max, 0.0)
 
     segm_expsum = tl.load(segm_expsum_ptr + segm_offset, mask=segm_mask, other=0.0)
     segm_expsum = segm_expsum * tl.exp(segm_max - overall_max)
@@ -439,29 +451,53 @@ def unified_attention_diffkv(
     BLOCK_M = (
         16 if num_queries_per_kv <= 16 else triton.next_power_of_2(num_queries_per_kv)
     )
-    BLOCK_Q = BLOCK_M // num_queries_per_kv
-
-    total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs
+    launch_kw: dict[str, int] = {}
 
     sliding_window_val = 1 + window_size[0] if window_size[0] >= 0 else 0
 
     # Decide between 2D and 3D launch.  Mirrors the standard launcher:
     # 3D requires preallocated softmax buffers, decode-only batches, and
     # a small number of sequences (otherwise 2D already saturates the SM).
+    # Spec-decode verify batches (q_len <= _SPEC_3D_MAX_Q, full attention)
+    # also take 3D: their 2D grid is only q_blocks x kv_heads programs, each
+    # walking the whole context.  Sliding-window layers stay 2D: their loop
+    # is already window-bounded.
+    spec_3d = (
+        1 < max_seqlen_q <= _SPEC_3D_MAX_Q
+        and sliding_window_val == 0
+        and softmax_segm_output is not None
+        and q.shape[0] <= softmax_segm_output.shape[0]
+    )
     use_3d = not (
         seq_threshold_3D is None
         or num_par_softmax_segments is None
         or softmax_segm_output is None
         or softmax_segm_max is None
         or softmax_segm_expsum is None
-        or max_seqlen_q > 1
+        or (max_seqlen_q > 1 and not spec_3d)
         or num_seqs > seq_threshold_3D
         or is_batch_invariant
     )
 
+    spec_tile: int | None = None
+    if use_3d and spec_3d and _SPEC_3D_BLOCK_M > BLOCK_M:
+        spec_bm = min(
+            _SPEC_3D_BLOCK_M,
+            triton.next_power_of_2(max_seqlen_q * num_queries_per_kv),
+        )
+        if spec_bm > BLOCK_M and spec_bm % num_queries_per_kv == 0:
+            BLOCK_M = spec_bm
+            launch_kw["num_warps"] = _SPEC_3D_NUM_WARPS if BLOCK_M >= 128 else 4
+            spec_tile = _SPEC_3D_TILE
+    BLOCK_Q = BLOCK_M // num_queries_per_kv
+
+    total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs
+
     # Tile size: 32 for prefill-class kernels.  Decode (small Q) prefers
     # smaller tiles to expose more parallelism along the KV dim.
     tile_size = 32 if not use_3d else (16 if q.element_size() >= 2 else 32)
+    if spec_tile is not None:
+        tile_size = spec_tile
 
     grid: tuple[Any, ...]
     if use_3d:
@@ -530,6 +566,7 @@ def unified_attention_diffkv(
         BLOCK_M=BLOCK_M,
         NUM_SEGMENTS_PER_SEQ=num_segments,
         IS_3D=use_3d,
+        **launch_kw,
     )
 
     if use_3d: