From 2ab1ba3736dc29509d3d2618dd2c8250533157bc Mon Sep 17 00:00:00 2001 From: 0z5a Date: Thu, 15 Oct 2026 01:44:56 +0200 Subject: [PATCH 1/3] [Kernel] Triton packed paged prefill attention Without FlashAttention, Triton prefill ran the single-query decode kernel once per prompt row: one program per (row, head) streaming its whole prefix. Long prompts therefore read every key once per query row. Add a varlen paged prefill kernel with the FA4 packed interface (q packed by cu_seqlens_q, one block table per sequence, bottom-right causal alignment). It tiles 64 queries x 64 keys per program with online softmax, gathers K/V through the block table, supports GQA and non-power-of-two head sizes, and uses exact IEEE products for FP32 inputs. The Triton backend now declares packed_prefill, so prompt prefill takes the packed path on hardware without FA4. Speculative verification keeps per-row attention: it compares against serial decoding token for token, which requires the packed kernel to match the decode kernel bitwise. A new capability, packed_matches_decode (FA4 only), gates that. Ouro-1.4B prefill on one H100 (BF16), packed vs per-row tokens/s: 1K tokens 11.9K vs 4.7K, 4K 15.7K vs 1.5K, 8K 14.3K vs 0.8K, 16K 11.9K vs 0.41K (2.5x to 29x). Co-Authored-By: Claude Opus 5.5 --- tests/test_attention.py | 44 +++++++ tests/test_backend_selection.py | 17 ++- tests/test_speculative.py | 17 ++- vllm_rlt/attention/backend.py | 15 ++- vllm_rlt/kernels/paged_attention.py | 21 ++++ vllm_rlt/kernels/triton_attention.py | 164 +++++++++++++++++++++++++++ vllm_rlt/worker/speculative.py | 3 +- 7 files changed, 269 insertions(+), 12 deletions(-) diff --git a/tests/test_attention.py b/tests/test_attention.py index 5e3fea14..ccb23364 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -105,3 +105,47 @@ 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) diff --git a/tests/test_backend_selection.py b/tests/test_backend_selection.py index f8624beb..58915153 100644 --- a/tests/test_backend_selection.py +++ b/tests/test_backend_selection.py @@ -153,11 +153,22 @@ def test_execution_requirements_prevent_unsafe_torch_fallback(packages, monkeypa 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), diff --git a/tests/test_speculative.py b/tests/test_speculative.py index b7fb44ab..93037b15 100644 --- a/tests/test_speculative.py +++ b/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 diff --git a/vllm_rlt/attention/backend.py b/vllm_rlt/attention/backend.py index 15bde763..b2c82087 100644 --- a/vllm_rlt/attention/backend.py +++ b/vllm_rlt/attention/backend.py @@ -8,7 +8,7 @@ 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 @dataclass @@ -41,11 +44,15 @@ def generation(self): _TORCH_CAPABILITIES = BackendCapabilities() -_TRITON_CAPABILITIES = BackendCapabilities(cuda_graphs=True, async_scheduling=True) +_TRITON_CAPABILITIES = BackendCapabilities( + packed_prefill=True, cuda_graphs=True, async_scheduling=True +) _FLASH_CAPABILITIES = { 2: BackendCapabilities(cuda_graphs=True, async_scheduling=True), 3: BackendCapabilities(cuda_graphs=True, async_scheduling=True), - 4: BackendCapabilities(packed_prefill=True, cuda_graphs=True, async_scheduling=True), + 4: BackendCapabilities( + packed_prefill=True, cuda_graphs=True, async_scheduling=True, packed_matches_decode=True + ), } _AUTO_PRIORITY = ("flash_attn_4", "flash_attn_3", "flash_attn_2", "triton", "torch") @@ -69,7 +76,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)) diff --git a/vllm_rlt/kernels/paged_attention.py b/vllm_rlt/kernels/paged_attention.py index 2d1f2428..8dc005d2 100644 --- a/vllm_rlt/kernels/paged_attention.py +++ b/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) diff --git a/vllm_rlt/kernels/triton_attention.py b/vllm_rlt/kernels/triton_attention.py index cf53e818..71deacf7 100644 --- a/vllm_rlt/kernels/triton_attention.py +++ b/vllm_rlt/kernels/triton_attention.py @@ -123,3 +123,167 @@ def paged_attention(q, key_cache, value_cache, block_tables, context_lengths): num_warps=4, ) return output + + +@triton.jit +def _paged_prefill_kernel( + Q, + K, + V, + TABLES, + LENGTHS, + CU_Q, + OUT, + q_token_stride, + q_head_stride, + k_block_stride, + k_token_stride, + k_head_stride, + v_block_stride, + v_token_stride, + v_head_stride, + table_stride, + out_token_stride, + out_head_stride, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + SCALE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + PRECISION: tl.constexpr, + max_q_blocks, +): + pid = tl.program_id(0) + head = tl.program_id(1) + seq = pid // max_q_blocks + q_block = pid % max_q_blocks + q_begin = tl.load(CU_Q + seq) + q_len = tl.load(CU_Q + seq + 1) - q_begin + if q_block * BLOCK_M >= q_len: + return + kv_len = tl.load(LENGTHS + seq) + prefix = kv_len - q_len + kv_head = head // HEAD_GROUPS + rows = q_block * BLOCK_M + tl.arange(0, BLOCK_M) + row_valid = rows < q_len + dims = tl.arange(0, BLOCK_D) + dim_valid = dims < HEAD_DIM + query = tl.load( + Q + + (q_begin + rows)[:, None].to(tl.int64) * q_token_stride + + head * q_head_stride + + dims[None, :], + mask=row_valid[:, None] & dim_valid[None, :], + other=0.0, + ) + positions = prefix + rows + maximum = tl.full((BLOCK_M,), -float("inf"), tl.float32) + normalizer = tl.zeros((BLOCK_M,), tl.float32) + accumulator = tl.zeros((BLOCK_M, BLOCK_D), tl.float32) + # Bottom-right causal alignment: the last query row sees every key. + end = tl.minimum(kv_len, prefix + (q_block + 1) * BLOCK_M) + for start in range(0, end, BLOCK_N): + keys_pos = start + tl.arange(0, BLOCK_N) + key_valid = keys_pos < end + blocks = tl.load( + TABLES + seq * table_stride + keys_pos // PAGE_SIZE, mask=key_valid, other=0 + ).to(tl.int64) + offsets = keys_pos % PAGE_SIZE + keys = tl.load( + K + + blocks[:, None] * k_block_stride + + offsets[:, None] * k_token_stride + + kv_head * k_head_stride + + dims[None, :], + mask=key_valid[:, None] & dim_valid[None, :], + other=0.0, + ) + scores = tl.dot(query, tl.trans(keys), input_precision=PRECISION) * SCALE + visible = key_valid[None, :] & (keys_pos[None, :] <= positions[:, None]) + scores = tl.where(visible, scores, -float("inf")) + next_maximum = tl.maximum(maximum, tl.max(scores, axis=1)) + # Rows with no visible key in this tile keep -inf; avoid inf - inf. + safe = tl.where(next_maximum == -float("inf"), 0.0, next_maximum) + correction = tl.exp(maximum - safe) + weights = tl.exp(scores - safe[:, None]) + values = tl.load( + V + + blocks[:, None] * v_block_stride + + offsets[:, None] * v_token_stride + + kv_head * v_head_stride + + dims[None, :], + mask=key_valid[:, None] & dim_valid[None, :], + other=0.0, + ) + accumulator = accumulator * correction[:, None] + tl.dot( + weights.to(values.dtype), values, input_precision=PRECISION + ) + normalizer = normalizer * correction + tl.sum(weights, axis=1) + maximum = next_maximum + out = accumulator / tl.maximum(normalizer, 1.0e-30)[:, None] + tl.store( + OUT + + (q_begin + rows)[:, None].to(tl.int64) * out_token_stride + + head * out_head_stride + + dims[None, :], + out, + mask=row_valid[:, None] & dim_valid[None, :], + ) + + +def paged_prefill_attention( + q, key_cache, value_cache, block_tables, context_lengths, cu_seqlens_q, max_seqlen_q +): + """Packed causal prefill over paged K/V, bottom-right aligned per sequence. + + q is [total query tokens, heads, dim]; sequence s owns rows + cu_seqlens_q[s]:cu_seqlens_q[s + 1] and keys [0, context_lengths[s]) through + block_tables[s]. Query i of a sequence sits at position + context_length - query_length + i and attends keys up to that position. + """ + output = torch.empty_like(q) + sequences = block_tables.shape[0] + if q.shape[0] == 0 or sequences == 0: + return output + block_m, block_n = 64, 64 + q_blocks = triton.cdiv(int(max_seqlen_q), block_m) + for tensor in (key_cache, value_cache): + if tensor.stride(3) != 1: + raise ValueError("paged prefill requires a contiguous head dimension") + if q.stride(2) != 1: + raise ValueError("paged prefill requires contiguous query rows") + _paged_prefill_kernel[(sequences * q_blocks, q.shape[1])]( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + cu_seqlens_q, + output, + q.stride(0), + q.stride(1), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + block_tables.stride(0), + output.stride(0), + output.stride(1), + HEAD_DIM=q.shape[-1], + HEAD_GROUPS=q.shape[1] // key_cache.shape[2], + PAGE_SIZE=key_cache.shape[1], + SCALE=q.shape[-1] ** -0.5, + BLOCK_D=max(16, triton.next_power_of_2(q.shape[-1])), + BLOCK_M=block_m, + BLOCK_N=block_n, + # FP32 callers (reference checks) need exact products, not TF32. + PRECISION="ieee" if q.dtype == torch.float32 else "tf32", + max_q_blocks=q_blocks, + num_warps=4, + num_stages=2, + ) + return output diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py index 089e708c..b29186b8 100644 --- a/vllm_rlt/worker/speculative.py +++ b/vllm_rlt/worker/speculative.py @@ -76,7 +76,8 @@ def __init__(self, model, cache, config, execution, rng: RngRegistry | None = No self.coda_graphs = CodaGraphs(model, execution) if execution.cuda_graphs else None def _packed(self, packed): - return packed and self.cache.attention_capabilities.packed_prefill + capabilities = self.cache.attention_capabilities + return packed and capabilities.packed_prefill and capabilities.packed_matches_decode def _cores(self, hidden, ids, positions, depths, *, packed=False): # Depth-independent metadata is prepared and copied once for all depths. From 43aa99eb808d66b545d5f6bc5c592d8041fb3183 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Thu, 15 Oct 2026 03:01:24 +0200 Subject: [PATCH 2/3] [Kernel] Split-KV Triton decode attention One program per (row, head) walked a row's whole context serially, so the longest row in a decode batch set every attention call's latency (about 0.3 ms per call at 10K tokens; Ouro decode makes 96 calls per token). Each row's context is now split into 16 equal tile-aligned parts computed in parallel and merged with their softmax statistics. The partition depends only on the row's own length, so a row's result is batch-invariant, which speculative verification relies on. Co-Authored-By: Claude Opus 5.5 --- tests/test_attention.py | 33 ++++++++++ vllm_rlt/kernels/triton_attention.py | 97 +++++++++++++++++++++++----- 2 files changed, 114 insertions(+), 16 deletions(-) diff --git a/tests/test_attention.py b/tests/test_attention.py index ccb23364..e56a2e26 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -149,3 +149,36 @@ def test_triton_paged_prefill_matches_per_query_reference(head_dim, heads, kv_he 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]) diff --git a/vllm_rlt/kernels/triton_attention.py b/vllm_rlt/kernels/triton_attention.py index 71deacf7..f1953a92 100644 --- a/vllm_rlt/kernels/triton_attention.py +++ b/vllm_rlt/kernels/triton_attention.py @@ -1,14 +1,21 @@ -"""Paged GQA kernel using bounded-memory online softmax. +"""Paged GQA kernels using bounded-memory online softmax. -One Triton program computes one query head. Unlike a fixed-depth batch, every -query row receives its own physical block table. This is a correctness-first -kernel without a performance claim or autotuning. +Every query row receives its own physical block table. Single-query attention +splits each row's context into NUM_SPLITS equal parts (split-KV): one program +per (row, head, split) produces a partial softmax state and a second kernel +merges the parts. A row's partition depends only on its own length, so its +result does not depend on the batch it runs in. Without splitting, one long +row serialized its whole context in one program and set the decode step time. """ import torch import triton import triton.language as tl +# Splits per row for single-query attention. A power of two (tl.arange) and a +# constant, so the grid is static under CUDA graphs and independent of batch. +NUM_SPLITS = 16 + @triton.jit def _paged_attention_kernel( @@ -17,7 +24,9 @@ def _paged_attention_kernel( V, TABLES, LENGTHS, - OUT, + PART_M, + PART_L, + PART_O, q_batch_stride: tl.constexpr, q_head_stride: tl.constexpr, q_dim_stride: tl.constexpr, @@ -30,20 +39,24 @@ def _paged_attention_kernel( v_head_stride: tl.constexpr, v_dim_stride: tl.constexpr, table_stride: tl.constexpr, - out_batch_stride: tl.constexpr, - out_head_stride: tl.constexpr, - out_dim_stride: tl.constexpr, HEAD_DIM: tl.constexpr, HEAD_GROUPS: tl.constexpr, + NUM_HEADS: tl.constexpr, PAGE_SIZE: tl.constexpr, SCALE: tl.constexpr, BLOCK_D: tl.constexpr, BLOCK_T: tl.constexpr, + SPLITS: tl.constexpr, ): row = tl.program_id(0) head = tl.program_id(1) + split = tl.program_id(2) kv_head = head // HEAD_GROUPS length = tl.load(LENGTHS + row) + # Equal tile-aligned parts of this row's context; trailing parts may be empty. + split_tokens = ((length + SPLITS - 1) // SPLITS + BLOCK_T - 1) // BLOCK_T * BLOCK_T + begin = split * split_tokens + end = tl.minimum(length, begin + split_tokens) dims = tl.arange(0, BLOCK_D) query = tl.load( Q + row * q_batch_stride + head * q_head_stride + dims * q_dim_stride, @@ -53,9 +66,9 @@ def _paged_attention_kernel( maximum = -float("inf") normalizer = 0.0 accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32) - for start in range(0, length, BLOCK_T): + for start in range(begin, end, BLOCK_T): positions = start + tl.arange(0, BLOCK_T) - valid_tokens = positions < length + valid_tokens = positions < end blocks = tl.load( TABLES + row * table_stride + positions // PAGE_SIZE, mask=valid_tokens, @@ -90,6 +103,38 @@ def _paged_attention_kernel( accumulator = accumulator * correction + tl.sum(weights[:, None] * values, axis=0) normalizer = normalizer * correction + tl.sum(weights, axis=0) maximum = next_maximum + # An empty part stores (-inf, 0, 0), which the merge weights by zero. + part = (row * NUM_HEADS + head) * SPLITS + split + tl.store(PART_M + part, maximum) + tl.store(PART_L + part, normalizer) + tl.store(PART_O + part * BLOCK_D + dims, accumulator) + + +@triton.jit +def _paged_attention_merge_kernel( + PART_M, + PART_L, + PART_O, + OUT, + out_batch_stride: tl.constexpr, + out_head_stride: tl.constexpr, + out_dim_stride: tl.constexpr, + HEAD_DIM: tl.constexpr, + NUM_HEADS: tl.constexpr, + BLOCK_D: tl.constexpr, + SPLITS: tl.constexpr, +): + row = tl.program_id(0) + head = tl.program_id(1) + parts = (row * NUM_HEADS + head) * SPLITS + tl.arange(0, SPLITS) + dims = tl.arange(0, BLOCK_D) + maxima = tl.load(PART_M + parts) + top = tl.max(maxima, axis=0) + # All parts empty (length 0) leaves top = -inf; zero weights give a zero row. + scale = tl.where(maxima > -float("inf"), tl.exp(maxima - top), 0.0) + normalizer = tl.sum(tl.load(PART_L + parts) * scale, axis=0) + partial = tl.load(PART_O + parts[:, None] * BLOCK_D + dims[None, :]) + accumulator = tl.sum(partial * scale[:, None], axis=0) tl.store( OUT + row * out_batch_stride + head * out_head_stride + dims * out_dim_stride, accumulator / tl.maximum(normalizer, 1.0e-30), @@ -98,28 +143,48 @@ def _paged_attention_kernel( def paged_attention(q, key_cache, value_cache, block_tables, context_lengths): - """Launch over [batch row, query head]; inputs are validated by the manager.""" + """Launch split-KV parts over [row, head, split], then merge per [row, head].""" output = torch.empty_like(q) if q.shape[0] == 0: return output - _paged_attention_kernel[(q.shape[0], q.shape[1])]( + rows, heads = q.shape[0], q.shape[1] + block_d = triton.next_power_of_2(q.shape[-1]) + part_m = torch.empty((rows, heads, NUM_SPLITS), device=q.device, dtype=torch.float32) + part_l = torch.empty_like(part_m) + part_o = torch.empty((rows, heads, NUM_SPLITS, block_d), device=q.device, dtype=torch.float32) + _paged_attention_kernel[(rows, heads, NUM_SPLITS)]( q, key_cache, value_cache, block_tables, context_lengths, - output, + part_m, + part_l, + part_o, *q.stride(), *key_cache.stride(), *value_cache.stride(), block_tables.stride(0), - *output.stride(), HEAD_DIM=q.shape[-1], - HEAD_GROUPS=q.shape[1] // key_cache.shape[2], + HEAD_GROUPS=heads // key_cache.shape[2], + NUM_HEADS=heads, PAGE_SIZE=key_cache.shape[1], SCALE=q.shape[-1] ** -0.5, - BLOCK_D=triton.next_power_of_2(q.shape[-1]), + BLOCK_D=block_d, BLOCK_T=32, + SPLITS=NUM_SPLITS, + num_warps=4, + ) + _paged_attention_merge_kernel[(rows, heads)]( + part_m, + part_l, + part_o, + output, + *output.stride(), + HEAD_DIM=q.shape[-1], + NUM_HEADS=heads, + BLOCK_D=block_d, + SPLITS=NUM_SPLITS, num_warps=4, ) return output From 6b02350a083a4c9c616b3f3e737a6d587916abb2 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sat, 24 Oct 2026 11:36:51 +0200 Subject: [PATCH 3/3] Make Triton packed prefill an explicit opt-in The shared Triton capabilities no longer declare packed_prefill, so split prefill, mixed execution and Huginn boundary attention keep per-row prefill. Only a backend selected with required packed_prefill (prefill_uva) gets it, on that instance. Same change as #130's 50e8f7e, applied to this branch. Co-authored-by: Codex Signed-off-by: 0z5a Co-Authored-By: Claude Opus 5.5 Signed-off-by: 0z5a --- tests/test_backend_selection.py | 17 +++++++++++++++++ tests/test_input_preparation.py | 1 + vllm_rlt/attention/backend.py | 11 +++++++++-- 3 files changed, 27 insertions(+), 2 deletions(-) diff --git a/tests/test_backend_selection.py b/tests/test_backend_selection.py index 8e835d9d..f9862214 100644 --- a/tests/test_backend_selection.py +++ b/tests/test_backend_selection.py @@ -169,6 +169,23 @@ 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") diff --git a/tests/test_input_preparation.py b/tests/test_input_preparation.py index 361e559f..f38a3c80 100644 --- a/tests/test_input_preparation.py +++ b/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), ], diff --git a/vllm_rlt/attention/backend.py b/vllm_rlt/attention/backend.py index 40f5c42c..210e766e 100644 --- a/vllm_rlt/attention/backend.py +++ b/vllm_rlt/attention/backend.py @@ -3,7 +3,7 @@ import importlib import logging from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, replace import torch @@ -47,7 +47,7 @@ def generation(self): _TORCH_CAPABILITIES = BackendCapabilities() _TRITON_CAPABILITIES = BackendCapabilities( - packed_prefill=True, cuda_graphs=True, async_scheduling=True, device_context_lengths=True + cuda_graphs=True, async_scheduling=True, device_context_lengths=True ) _FLASH_CAPABILITIES = { 2: BackendCapabilities(cuda_graphs=True, async_scheduling=True, device_context_lengths=True), @@ -114,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 (