[Kernel] Triton packed paged prefill and split-KV decode attention

tests/test_attention.py

@@ -105,3 +105,80 @@ def test_empty_batch_and_invalid_gqa():
     cache.write(0, ["a"], [0], [0], torch.ones(1, 2, 4), torch.ones(1, 2, 4))
     with pytest.raises(ValueError, match="positive multiple"):
         cache.attend(0, ["a"], [0], [0], torch.ones(1, 3, 4))
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("head_dim,heads,kv_heads", [(64, 4, 2), (96, 8, 8), (128, 16, 16)])
+def test_triton_paged_prefill_matches_per_query_reference(head_dim, heads, kv_heads):
+    from vllm_rlt.kernels.paged_attention import torch_paged_attention
+    from vllm_rlt.kernels.triton_attention import paged_prefill_attention
+
+    torch.manual_seed(0)
+    block, blocks = 16, 96
+    k = torch.randn(blocks, block, kv_heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn_like(k)
+    # (keys visible to the last query, queries in this chunk): prefixes of 32, 0, 0, 60, 9.
+    sequences = [(37, 5), (100, 100), (3, 3), (130, 70), (129, 120)]
+    order = torch.randperm(blocks).tolist()  # Pages deliberately out of order.
+    width = max(-(-length // block) for length, _ in sequences)
+    tables = torch.zeros(len(sequences), width, dtype=torch.int32, device="cuda")
+    for s, (length, _) in enumerate(sequences):
+        pages = [order.pop() for _ in range(-(-length // block))]
+        tables[s, : len(pages)] = torch.tensor(pages, dtype=torch.int32)
+    total = sum(count for _, count in sequences)
+    q = torch.randn(total, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    lengths = torch.tensor([length for length, _ in sequences], dtype=torch.int32, device="cuda")
+    starts = [0]
+    for _, count in sequences:
+        starts.append(starts[-1] + count)
+    cu = torch.tensor(starts, dtype=torch.int32, device="cuda")
+    actual = paged_prefill_attention(
+        q, k, v, tables, lengths, cu, max(count for _, count in sequences)
+    )
+    rows, row_tables, row_lengths = [], [], []
+    for s, (length, count) in enumerate(sequences):
+        for i in range(count):
+            rows.append(starts[s] + i)
+            row_tables.append(tables[s])
+            row_lengths.append(length - count + i + 1)  # Causal: keys up to own position.
+    expected = torch_paged_attention(
+        q[rows],
+        k,
+        v,
+        torch.stack(row_tables),
+        torch.tensor(row_lengths, dtype=torch.int32, device="cuda"),
+    )
+    torch.testing.assert_close(actual[rows], expected, rtol=2e-2, atol=2e-2)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_split_kv_decode_matches_reference_and_is_batch_invariant(dtype):
+    from vllm_rlt.kernels.paged_attention import torch_paged_attention
+    from vllm_rlt.kernels.triton_attention import paged_attention
+
+    torch.manual_seed(1)
+    block, heads, kv_heads, head_dim = 16, 8, 4, 128
+    # Empty, tile and split boundaries, and long rows that use every split.
+    lengths = [0, 1, 31, 32, 33, 511, 512, 513, 4097, 10000]
+    pages = sum(-(-n // block) for n in lengths)
+    k = torch.randn(pages, block, kv_heads, head_dim, device="cuda", dtype=dtype)
+    v = torch.randn_like(k)
+    order = torch.randperm(pages).tolist()  # Pages deliberately out of order.
+    width = max(-(-n // block) for n in lengths)
+    tables = torch.zeros(len(lengths), width, dtype=torch.int32, device="cuda")
+    for row, n in enumerate(lengths):
+        ids = [order.pop() for _ in range(-(-n // block))]
+        tables[row, : len(ids)] = torch.tensor(ids, dtype=torch.int32)
+    q = torch.randn(len(lengths), heads, head_dim, device="cuda", dtype=dtype)
+    context = torch.tensor(lengths, dtype=torch.int32, device="cuda")
+    actual = paged_attention(q, k, v, tables, context)
+    expected = torch_paged_attention(q, k, v, tables, context)
+    tolerance = 2e-2 if dtype == torch.bfloat16 else 1e-4
+    torch.testing.assert_close(actual, expected, rtol=tolerance, atol=tolerance)
+    assert torch.count_nonzero(actual[0]) == 0
+    # Speculative verification compares rows computed in different batches.
+    for row in (1, 5, 8, 9):
+        one = slice(row, row + 1)
+        alone = paged_attention(q[one], k, v, tables[one], context[one])
+        assert torch.equal(alone[0], actual[row])

tests/test_backend_selection.py

@@ -169,14 +169,42 @@ def test_flash_device_lengths_and_graph_requirements(packages, monkeypatch, sm,
     assert selected.info["selected_backend"] == backend
 
 
+@pytest.mark.parametrize("requested", ["triton", "auto"])
+def test_triton_packed_prefill_is_per_selection_opt_in(packages, monkeypatch, requested):
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (9, 0))
+    packages.difference_update(
+        {"flash_attn.cute", "flash_attn_3.flash_attn_interface", "flash_attn"}
+    )
+    args = (requested, torch.device("cuda"), torch.bfloat16, 128, 16)
+    default = create_backend(*args)
+    packed = create_backend(*args, required_capabilities=BackendCapabilities(packed_prefill=True))
+    later = create_backend(*args)
+    assert default.info["selected_backend"] == packed.info["selected_backend"] == "triton"
+    assert not default.capabilities.packed_prefill
+    assert packed.capabilities.packed_prefill
+    assert not packed.capabilities.packed_matches_decode
+    assert not later.capabilities.packed_prefill
+
+
 def test_packed_prefill_requires_declared_support(packages, monkeypatch):
     monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (9, 0))
     packages.remove("flash_attn.cute")
+    # Without FA4, auto skips FA3/FA2 and selects Triton's packed paged prefill.
+    backend = create_backend(
+        "auto",
+        torch.device("cuda"),
+        torch.bfloat16,
+        128,
+        16,
+        required_capabilities=BackendCapabilities(packed_prefill=True),
+    )
+    assert backend.info["selected_backend"] == "triton"
+    assert backend.capabilities.packed_prefill
     with pytest.raises(ValueError, match="packed_prefill"):
         create_backend(
-            "auto",
-            torch.device("cuda"),
-            torch.bfloat16,
+            "torch",
+            torch.device("cpu"),
+            torch.float32,
             128,
             16,
             required_capabilities=BackendCapabilities(packed_prefill=True),

tests/test_input_preparation.py

@@ -20,6 +20,7 @@
         ("cpu", "last_exited", "torch", False),
         pytest.param("cuda", "shared", "triton", False, marks=pytest.mark.gpu),
         pytest.param("cuda", "last_exited", "triton", False, marks=pytest.mark.gpu),
+        pytest.param("cuda", "last_exited", "triton", True, marks=pytest.mark.gpu),
         pytest.param("cuda", "last_exited", "flash_attn_4", False, marks=pytest.mark.gpu),
         pytest.param("cuda", "last_exited", "flash_attn_4", True, marks=pytest.mark.gpu),
     ],

tests/test_speculative.py

@@ -44,18 +44,27 @@ def drain(e):
     pytest.fail("speculative scheduler did not drain")
 
 
-@pytest.mark.parametrize("generation,supports_packed", [(2, True), (4, False)])
-def test_speculative_prefill_routes_by_declared_capabilities(generation, supports_packed):
+@pytest.mark.parametrize(
+    "generation,packed,matches_decode,expected",
+    # Triton packs prompt prefill, but its packed kernel is not bitwise equal to
+    # its decode kernel, so speculative verification keeps per-row attention.
+    [(2, True, True, True), (2, True, False, False), (4, False, False, False)],
+)
+def test_speculative_prefill_routes_by_declared_capabilities(
+    generation, packed, matches_decode, expected
+):
     from types import SimpleNamespace
 
     from vllm_rlt.attention import BackendCapabilities
 
     e = engine(model())
     e.cache_manager.attention = SimpleNamespace(
-        capabilities=BackendCapabilities(packed_prefill=supports_packed)
+        capabilities=BackendCapabilities(
+            packed_prefill=packed, packed_matches_decode=matches_decode
+        )
     )
     e.cache_manager.attention_info = {"generation": generation}
-    assert e.speculative_runner._packed(True) is supports_packed
+    assert e.speculative_runner._packed(True) is expected
     assert e.speculative_runner._packed(False) is False
 
 

vllm_rlt/attention/backend.py

@@ -3,12 +3,12 @@
 import importlib
 import logging
 from collections.abc import Callable
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 
 import torch
 
 from vllm_rlt.kernels.flash_attention import FLASH_BACKENDS, FlashPagedAttention
-from vllm_rlt.kernels.paged_attention import torch_paged_attention, triton_paged_attention
+from vllm_rlt.kernels.paged_attention import TritonPagedAttention, torch_paged_attention
 
 
 @dataclass(frozen=True)
@@ -18,6 +18,9 @@ class BackendCapabilities:
     packed_prefill: bool = False
     cuda_graphs: bool = False
     async_scheduling: bool = False
+    # Packed rows reproduce single-query decode exactly; speculative verification
+    # relies on this to match serial decoding token for token.
+    packed_matches_decode: bool = False
     # Consume GPU-owned context lengths without a host read between submissions.
     device_context_lengths: bool = False
 
@@ -50,7 +53,11 @@ def generation(self):
     2: BackendCapabilities(cuda_graphs=True, async_scheduling=True, device_context_lengths=True),
     3: BackendCapabilities(cuda_graphs=True, async_scheduling=True, device_context_lengths=True),
     4: BackendCapabilities(
-        packed_prefill=True, cuda_graphs=True, async_scheduling=True, device_context_lengths=True
+        packed_prefill=True,
+        cuda_graphs=True,
+        async_scheduling=True,
+        packed_matches_decode=True,
+        device_context_lengths=True,
     ),
 }
 _AUTO_PRIORITY = ("flash_attn_4", "flash_attn_3", "flash_attn_2", "triton", "torch")
@@ -75,7 +82,7 @@ def _create_explicit(backend, device, dtype, head_dim, block_size):
             raise ValueError("the Triton attention backend requires NVIDIA Compute Capability 8.0+")
         # Probe the actual module, including its dependency imports, before selection.
         importlib.import_module("vllm_rlt.kernels.triton_attention")
-        return AttentionBackend(triton_paged_attention, _TRITON_CAPABILITIES, {"backend": backend})
+        return AttentionBackend(TritonPagedAttention(), _TRITON_CAPABILITIES, {"backend": backend})
     if backend in FLASH_BACKENDS:
         kernel = FlashPagedAttention(device, dtype, head_dim, block_size, backend)
         return AttentionBackend(kernel, _FLASH_CAPABILITIES[kernel.generation], dict(kernel.info))
@@ -107,6 +114,13 @@ def create_backend(
     for candidate in candidates:
         try:
             implementation = _create_explicit(candidate, device, dtype, head_dim, block_size)
+            if candidate == "triton" and required_capabilities.packed_prefill:
+                # Packed BF16 prefill has different arithmetic from per-row decode.
+                # Opt in for prefill_uva without changing existing split/mixed or
+                # Huginn boundary attention when that option is disabled.
+                implementation.capabilities = replace(
+                    implementation.capabilities, packed_prefill=True
+                )
             missing = [
                 name
                 for name in (

vllm_rlt/kernels/paged_attention.py

@@ -52,3 +52,24 @@ def triton_paged_attention(
     from vllm_rlt.kernels.triton_attention import paged_attention
 
     return paged_attention(q, key_cache, value_cache, block_tables, context_lengths)
+
+
+def triton_paged_prefill(
+    q, key_cache, value_cache, block_tables, context_lengths, cu_seqlens_q, max_seqlen_q
+):
+    """Packed causal prefill with one block table per sequence (see triton_attention)."""
+    from vllm_rlt.kernels.triton_attention import paged_prefill_attention
+
+    return paged_prefill_attention(
+        q, key_cache, value_cache, block_tables, context_lengths, cu_seqlens_q, max_seqlen_q
+    )
+
+
+class TritonPagedAttention:
+    """Triton single-query decode attention plus packed causal prefill."""
+
+    def __call__(self, q, key_cache, value_cache, block_tables, context_lengths):
+        return triton_paged_attention(q, key_cache, value_cache, block_tables, context_lengths)
+
+    def prefill(self, *args, **kwargs):
+        return triton_paged_prefill(*args, **kwargs)