Decode attention: FlashDecoding, split-K, paged kernels

FlashAttention parallelises over the query axis. In prefill that axis is 2,048 or 8,192 tokens long and there is more parallelism than the machine can absorb. In decode it is one token per sequence. Everything that made the prefill kernel fast stops working.

Here is the grid, computed from vLLM's real launcher rather than from a paper. The Triton unified-attention wrapper packs queries into blocks of BLOCK_Q tokens:

    num_queries_per_kv = num_query_heads // num_kv_heads
    head_size = q.shape[2]

    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

and launches a two-dimensional grid over (query blocks × KV heads):

    grid: tuple[Any, ...]
    if not use_3d:
        grid = (total_num_q_blocks, num_kv_heads)
        tile_size = TILE_SIZE_PREFILL
    else:
        grid = (total_num_q_blocks, num_kv_heads, num_par_softmax_segments)
        tile_size = TILE_SIZE_DECODE

Work it for Llama-3-8B ($h = 32$ query heads, $h_{kv} = 8$ KV heads, so $g = h/h_{kv} = 4$). BLOCK_M = 16, BLOCK_Q = 4, and total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs (vllm/v1/attention/ops/triton_unified_attention.py:L966). At decode q.shape[0] equals the batch size $B$, so the 2D grid holds $(\lfloor B/4 \rfloor + B) \times 8$ blocks. Every number below is derived from that expression — no GPU was involved.

Batch 1 uses 6% of the GPU. It gets worse under tensor parallelism, because TP shards the head dimension: Llama-3-70B on TP=8 gives each rank $h = 8$, $h_{kv} = 1$, so BLOCK_Q = 2 and the batch-1 2D grid is one thread block on 132 SMs — the same 0.76% figure §0.3 derived for RMSNorm.

There is only one query and it is tiny. There is an enormous amount of KV and it is read exactly once. So stop parallelising over queries and parallelise over keys: cut the sequence into $S$ contiguous slices, give each slice its own thread block, let each block run the ordinary online-softmax loop over only its slice, and have each block emit a partial answer — the unnormalised output $\tilde{O}_i$, the running max $m_i$, and the running sum $\ell_i$. A second, nearly free kernel then merges the $S$ partials into the true softmax using the rescale identity from §3.1. This is FlashDecoding; in GEMM vocabulary it is split-K, because the reduction axis is the one being split.

The trade is explicit: same FLOPs, same KV bytes, done 16× wider and 16× shorter, paid for with one extra round trip of partial state through HBM. The rest of the chapter sizes $S$, deals with the fact that the KV is not contiguous, and explains why vLLM deleted the kernel it wrote for exactly this job.

Symbols, all per attention layer: $B$ batch size, $s$ sequence length (KV positions attended), $h$ query heads, $h_{kv}$ KV heads, $g = h / h_{kv}$ the GQA group size, $d_h$ head dimension, $b$ bytes per KV element, $S$ the number of KV splits, $N_{SM} = 132$.

The grid

The 2D (unsplit) grid has $G_{2D} = (\lfloor B/\text{BLOCK\_Q}\rfloor+B)h_{kv}$ for the quoted vLLM wrapper blocks. Adding the split axis multiplies it:

and the smallest useful $S$ is the one that first fills the machine:

For Llama-3-8B at $B = 1$: $G_{2D} = 8$, so $S^{*} = \lceil 132/8 \rceil = 17$; vLLM's fixed 16 lands at 128 of 132 SMs. For 70B on TP=8 at $B=1$: $G_{2D} = 1$, $S^{*} = 132$, and the same fixed 16 gets you 16 blocks. That constant is not adaptive — precisely the axis on which SGLang differs (§4).

Splitting beyond one SM wave can still improve latency hiding or resource balance; it also increases merge traffic. S* is a first-wave coverage heuristic, not a universal optimal split count.

Segment length, and the early exit

vLLM does not divide the sequence into $S$ equal pieces. It computes a tile count per segment and lets the last segments fall off the end:

    if IS_3D:
        tiles_per_segment = cdiv_fn(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE)
        if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len:
            return
    else:
        tiles_per_segment = 0

TILE_SIZE for the decode path is 16 for bf16 KV (vllm/v1/attention/ops/triton_unified_attention.py:L798-L799), and NUM_SEGMENTS_PER_SEQ is 16. So each segment covers $16 \lceil s / 256 \rceil$ keys, and the number of segments that actually do work is $\lceil s / (16 \lceil s/256 \rceil) \rceil$. Derived:

Below $s = 256$ the split saturates at $s/16$ segments and the surplus blocks return immediately. That early return is how a fixed-shape grid stays legal under CUDA graph capture — a recording of the step's kernel launches that is replayed with fixed dimensions, §8.1 — where launch dimensions must not depend on the runtime sequence length.

What the merge costs

Stage 1 writes, and stage 2 reads, one fp32 output vector plus two fp32 scalars per (token, head, segment). Per (token, head) that is $S(d_h + 2) \cdot 4$ bytes $= 16 \times 130 \times 4 = 8{,}320$ B. For $B = 1$, $h = 32$: 266 KB read by the merge and 266 KB written by stage 1.

Compare against the KV that stage 1 must read anyway, per layer: $2 h_{kv} d_h s b = 2 \times 8 \times 128 \times 2560 \times 2 = 10.5$ MB.

The overhead is fixed per token while the KV read grows with $s$, which is the whole design in one sentence: split hard when $s$ is large. At $s = 512$ the same 532 KB of partial state sits against 2.1 MB of KV — a quarter of the traffic added to save an idle GPU. Whether that trades well depends on how idle the GPU was, which is why the decision is a threshold on batch size rather than a constant.

Where GQA enters

The tile the kernel actually computes is BLOCK_M = 16 query rows against TILE_SIZE = 16 keys. Those 16 rows are not 16 tokens — they are BLOCK_Q = 4 tokens × $g$ = 4 query heads, all sharing one KV head:

    query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv

    query_offset_0 = cur_batch_in_all_start_index + query_pos
    query_offset_1 = kv_head_idx * num_queries_per_kv + offs_m % num_queries_per_kv

Row offs_m of the tile belongs to token offs_m // g and query head kv_head_idx * g + offs_m % g. One K tile load therefore feeds $g$ query heads. Count it: the K and V tiles are $2 \times 16 \times 128 \times 2 = 8{,}192$ bytes; with a single sequence in flight only 4 of the 16 rows carry live data, and those 4 rows cost $2 \times (2 \times 4 \times 16 \times 128) = 32{,}768$ FLOP. The ratio is

exactly the decode-attention intensity §0.4 derived as $2h/(b\,h_{kv})$, now visible as a property of a single tile of a single kernel. Against the H100's ridge of $I^{*} = 295$ that is 74× memory-bound, and no scheduling trick changes it: the split exposes more parallel work but adds partial-state traffic, so actual operational intensity can decrease. What GQA buys is the factor $g$ itself — an MHA model ($g = 1$) runs the same kernel at $I = 1$.

Note the flip side: at batch 1 twelve of the sixteen BLOCK_M rows are masked off, so the tensor cores compute a full 16×16 tile and throw away 75% of it. The waste is invisible on the roofline — the kernel is bandwidth-bound anyway — which is why decode kernels are tuned for bytes moved, not FLOP efficiency.

Stage 1 is the ordinary kernel with one extra program_id and a different destination:

    q_block_global_idx = tl.program_id(0)
    kv_head_idx = tl.program_id(1)
    segm_idx = tl.program_id(2) if IS_3D else 0

Instead of dividing by $\ell$ and storing $O$, it stores $\tilde{O}$, $m$, $\ell$ into three scratch tensors sized at build time (vllm/v1/attention/backends/triton_attn.py:L156-L177). Stage 2, reduce_segments, is the rescale from §3.1 written out verbatim:

    segm_max = tl.load(segm_max_ptr + segm_offset, mask=segm_mask, other=float("-inf"))
    overall_max = tl.max(segm_max)

    # load and rescale segment exp sums
    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)
    overall_expsum = tl.sum(segm_expsum)
# ...
    segm_output *= tl.exp(segm_max - overall_max)[:, None]
    acc_sum = tl.sum(segm_output, axis=0)
    # safely divide by overall_expsum, returning 0.0 if overall_expsum is 0
    acc = tl.where(overall_expsum == 0.0, 0.0, acc_sum / overall_expsum)

Global max across segments, rescale each $\ell_i$ and $\tilde{O}_i$ by $e^{m_i - m}$, sum, divide once, no re-reading of K or V. Its grid is (q.shape[0], num_query_heads) — 32 blocks at $B = 1$, itself badly under-occupied, but it moves 266 KB, under 0.1 µs of bandwidth time. At that size the merge is a launch cost, not a memory cost.

Everything above assumed the KV slice is a contiguous range of memory. It is not. §2.2 established that KV lives in fixed-size blocks scattered across a pool, and a per-sequence block table maps logical block index to physical block index. The kernel has to perform that mapping itself, once per tile, in the inner loop:

        seq_offset = j * TILE_SIZE + offs_t
        tile_mask = seq_offset < max_seq_prefix_len

        physical_block_idx = tl.load(
            block_tables_ptr + block_table_offset + seq_offset // BLOCK_SIZE
        ).to(tl.int64)

and then folds the physical block index into the address arithmetic:

            k_offset = (
                physical_block_idx[None, :] * stride_k_cache_0
                + kv_head_idx * stride_k_cache_2
                + offs_d[:, None] * stride_k_cache_3
                + (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1
            )

That is the whole of PagedAttention's addressing, in four lines of Triton: position $p$ becomes $\text{blocktable}[\text{seq}][\lfloor p / \text{block\_size} \rfloor] \cdot \text{stride}_0 + (p \bmod \text{block\_size}) \cdot \text{stride}_1$ plus head and dim terms, with block_table_offset = seq_idx * block_table_stride (vllm/v1/attention/ops/triton_unified_attention.py:L375) selecting the row §2.2's slot mapping fills in.

The bandwidth cost of paging is negligible — one int32 per 16 keys against 512 bytes of K and V per key and KV head, and the whole block table for a 2,560-token sequence is 640 bytes, comfortably L2-resident. The latency cost is real but structural: the KV address depends on the result of a load, a serial dependency the memory pipeline cannot hide without prefetching block indices a tile ahead.

Coalescing survives as long as a tile does not straddle blocks. With TILE_SIZE = 16 and block_size = 16 the tile sits inside one physical block, physical_block_idx is uniform across the vector, and the addresses form one contiguous run. A smaller kernel block size still works — physical_block_idx becomes a genuine vector — but the load fragments into shorter runs. Hence MultipleOf(16) as the Triton backend's supported kernel block size (vllm/v1/attention/backends/triton_attn.py:L308-L310).

And it is why the block table the kernel sees may not be the allocator's. §2.2 found the kernel_block_size / block_size split; here is the function that performs the subdivision:

    def map_to_kernel_blocks(
        kv_manager_block_ids: np.ndarray,
        blocks_per_kv_block: int,
        kernel_block_arange: np.ndarray,
    ) -> np.ndarray:
        """Convert kv_manager_block_id IDs to kernel block IDs.

        Example:
            # kv_manager_block_ids: 32 tokens,
            # Kernel block size: 16 tokens
            # blocks_per_kv_block = 2
            >>> kv_manager_block_ids = np.array([0, 1, 2])
            >>> Result: [0, 1, 2, 3, 4, 5]
# ...
        kernel_block_ids = (
            kv_manager_block_ids.reshape(-1, 1) * blocks_per_kv_block
            + kernel_block_arange
        )

        return kernel_block_ids.reshape(-1)

The allocator wants large blocks (fewer entries to hash for prefix caching, cheaper free lists); the kernel wants a block size that divides its tile. Subdividing satisfies both, at the cost of a block table blocks_per_kv_block× wider — eight times as many int32s for a 128-token allocator block cut to 16. Still under 1% of KV traffic, and the width is computed once at build time rather than inferred in the kernel (vllm/v1/worker/block_table.py:L29-L49).

vLLM: a threshold, and a fixed split count

vLLM makes exactly one decision — 2D or 3D — and makes it on batch size:

# constants
MIN_LAUNCH_GRID_SIZE_2D = 128  # Minimum launch grid size of 2D kernel
NUM_PAR_SOFTMAX_SEGMENTS = 16  # Number of parallel tiled softmax segments
# ...
        # The launch grid for the 2D kernel is defined as (num_q_blocks, num_heads_kv).
        # A lower bound for num_q_blocks is the number of sequences.
        # To ensure the minimum launch grid size is achieved, the number of sequences
        # must be at least equal to the threshold below.
        # If this threshold is not reached (i.e., the batch size is not large enough),
        # the 3D kernel will be selected instead.
        self.seq_threshold_3D = MIN_LAUNCH_GRID_SIZE_2D // self.num_heads_kv

For Llama-3-8B, $128 / 8 = 16$: split for batches of 16 or fewer, do not split above. Under TP=8 on the 70B, $h_{kv} = 1$ per rank and the threshold becomes 128. The gate is evaluated in the wrapper, and note what else disqualifies the split path:

    # Launch the 2D kernel if
    # 1. No intermediate tiled softmax buffers for the 3D kernel have been allocated, or
    # 2. The batch includes at least one prefill request, or
    # 3. The number of sequences exceeds the configured threshold, or
    # 4. Batch invariance is enabled
    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 num_seqs > seq_threshold_3D
        or is_batch_invariant
    )

max_seqlen_q > 1 kills it for any mixed batch: one chunked-prefill request and the whole batch takes the unsplit path. That is the real cost of a unified kernel — the same code serves prefill and decode, so the decode optimisation is all-or-nothing per launch.

SGLang: a split count computed per sequence, on device

SGLang keeps separate prefill and decode kernels and can therefore be aggressive. Its split count is not a constant; it is a Triton kernel that reads the batch's sequence lengths and the GPU's SM count and solves for the count that fills the machine:

    # NOTE: this is a hack to let num_kv_split grows up with seqlen gradually
    ext_seq_len = tl.cast(max_seq_len, tl.float32) / 64.0
    ext_device_core_count = tl.cast(
        device_core_count * tl.maximum(tl.log2(ext_seq_len), 1.0), tl.int32
    )
    block_h, num_kv_group = 16, num_head // num_kv_head
    if num_kv_group == 1:
        token_grid = num_seq * num_group * num_head
    else:
        # from triton_ops/decode_attention.py:_decode_grouped_att_m_fwd
        block_h = tl.minimum(block_h, num_kv_group)
        token_grid = num_seq * num_group * tl.cdiv(num_head, block_h)
    max_kv_splits_2 = tl.minimum(
        tl.cdiv(ext_device_core_count, token_grid), max_kv_splits
    )
    kv_chunk_size_2 = tl.cdiv(max_seq_len, max_kv_splits_2)

    num_kv_splits = tl.maximum(
        tl.cdiv(seq_lens, kv_chunk_size_1), tl.cdiv(seq_lens, kv_chunk_size_2)
    )

token_grid is exactly the $G_{2D}$ of §3, and cdiv(ext_device_core_count, token_grid) is $S^{*}$ with a $\log_2(s/64)$ fudge that lets long contexts oversubscribe. Run it for Llama-3-8B, $s = 2560$, $N_{SM} = 132$ (so ext_device_core_count $= \lfloor 132 \times \log_2 40 \rfloor = 702$), block_h $= \min(16, 4) = 4$, head tiles $= \lceil 32/4 \rceil = 8$, and max_kv_splits at its default of 8 (python/sglang/srt/server_args.py:L1988-L1992):

SGLang adapts its active split count to batch and context while the quoted vLLM path uses a fixed segment count below a threshold. Distinguish allocated grid slots from useful work: extra query tiles can return early, and padded split rows can be masked. Data-dependent active splits can coexist with CUDA graphs when maximum launch dimensions and buffers stay fixed; adaptation does not inherently require a new device launch or recapture.

SGLang's stage-1 grid puts the GQA group on the head axis explicitly:

    # Blocks at or above the split count return immediately, so the grid shrinks too.
    grid = (batch, head_tiles, forced_kv_splits or MAX_KV_SPLITS)
# ...
    cur_head_id = tl.program_id(1)
    cur_kv_head = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H)
    split_kv_id = tl.program_id(2)

    if BLOCK_H < kv_group_num:
        VALID_BLOCK_H: tl.constexpr = BLOCK_H
    else:
        VALID_BLOCK_H: tl.constexpr = kv_group_num
    cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H)
    mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H

With _GROUPED_BLOCK_H = 16 (python/sglang/kernels/ops/attention/decode_attention.py:L40-L41) and $g = 4$, one block owns min(16, 4) = 4 query heads sharing one KV head — the same amortisation vLLM gets by packing rows of BLOCK_M, expressed as a grid dimension instead. For DeepSeek-style MLA where $g$ is 128, BLOCK_H caps at 16 and the group is tiled across $\lceil 128/16 \rceil = 8$ blocks.

SGLang's paging differs in kind too: kv_indices is a flat per-token index list, not a per-block table, so the kernel derives page and offset itself:

            # Page-aware KV address math (see _fwd_kernel_stage1).
            if PAGE_SIZE == 1:
                offs_buf_k = kv_loc[None, :] * stride_buf_kbs + base_offs_k
            else:
                page_id = kv_loc // PAGE_SIZE
                tok_in_p = kv_loc % PAGE_SIZE
                offs_buf_k = (
                    page_id[None, :] * stride_buf_kpage
                    + tok_in_p[None, :] * stride_buf_ktok
                    + base_offs_k
                )

At PAGE_SIZE = 1 — what RadixAttention's token-level sharing wants (§2.4) — the indirection is one int32 per key rather than per block: the same 0.78%, but with a fully scattered address stream instead of contiguous runs. The constant-folded branch exists to make that case as cheap as it can be.

The kernel vLLM deleted

The most instructive thing in this chapter is a file that is not there. At the pinned SHA, csrc/attention/ in vLLM contains six headers — attention_dtypes.h, attention_generic.cuh and four dtype_*.cuh — and no kernel at all. The original PagedAttention CUDA kernel — the one the 2023 paper is about, the one every blog post diagrams — was moved to csrc/libtorch_stable/attention/ in PR #43717 and then removed outright:

commit d715b3aa1ea6af3f663eb6d3cd8f5b6bb15770e9
Author: Michael Goin <mgoin64@gmail.com>
Date:   Thu Jul 2 15:31:26 2026 -0400

    Delete PagedAttention (#47361)

 CMakeLists.txt                                     |   2 -
 benchmarks/kernels/benchmark_paged_attention.py    | 133 ++--
 .../attention/attention_kernels.cuh                | 667 ---------------------
 .../attention/paged_attention_v1.cu                | 190 ------
 .../attention/paged_attention_v2.cu                | 202 -------
 csrc/libtorch_stable/ops.h                         |  26 -
 csrc/libtorch_stable/torch_bindings.cpp            |  30 -
 tests/kernels/attention/test_attention.py          | 218 ++-----
 vllm/_custom_ops.py                                |  94 ---
 9 files changed, 90 insertions(+), 1472 deletions(-)

What was in those 1,059 lines? Exactly this chapter. v1 was the unsplit kernel, dim3 grid(num_heads, num_seqs, 1). v2 was FlashDecoding: a PARTITION_SIZE of 512 keys, a grid of (num_heads, num_seqs, max_num_partitions), scratch tensors exp_sums, max_logits, tmp_out, and a paged_attention_v2_reduce_kernel on (num_heads, num_seqs). Rename exp_sums to segm_expsum and you have the Triton kernel of §4.1, three years later. (Recover it with git show d715b3aa1e~1:csrc/libtorch_stable/attention/paged_attention_v2.cu.)

Why did it become redundant? Because everything that made it special stopped being special:

What was lost is real but narrow: block-sparse attention (the IS_BLOCK_SPARSE and blocksparse_* template parameters), the odd head sizes 80/96/112/120 compiled by name — and a reference implementation that fit in one file and could be read end to end.

Two hand-written paged kernels survive in the tree. csrc/rocm/attention.cu is the AMD one — 3,717 lines, still very much alive, and its launcher is the clearest surviving statement of the split-K grid:

  // partition size is fixed at 256 since both mfma4 and mfma16 kernels support
  // it mfma4 kernel also supports partition size 512
  constexpr int PARTITION_SIZE = 256;
  const int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);
  const int gqa_ratio = num_heads / num_kv_heads;
  assert(num_heads % num_kv_heads == 0);
  assert(head_size == HEAD_SIZE);

  constexpr int NTHR = 256;
  dim3 grid(num_seqs, max_num_partitions, num_kv_heads);

Note the third grid dimension: num_kv_heads, not num_heads. The gqa_ratio is a template parameter, dispatched by a switch (csrc/rocm/attention.cu:L3295-L3300), so the whole query-head group lives in one thread block and reads the KV once. Same idea as vLLM's BLOCK_M packing and SGLang's BLOCK_H, this time baked into the launch.

The other survivor is csrc/libtorch_stable/attention/merge_attn_states.cu, the same LSE algebra extracted as a standalone op. Its header comment advertises the split-KV case, but read what it actually does — two operands, not N:

// Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
// can be used to combine partial attention results (in the split-KV case)
# ...
  p_lse = p_lse - max_lse;
  s_lse = s_lse - max_lse;
  const float p_se = expf(p_lse);
  const float s_se = expf(s_lse);
  const float out_se = p_se + s_se;
  const float p_scale = p_se / out_se;
  const float s_scale = s_se / out_se;

It merges exactly two partial states, carrying LSE rather than $(m, \ell)$ separately, and serves cases where the halves come from different launches — cascade attention's shared prefix plus per-sequence suffix, and decode context parallelism — not an $N$-way split inside one launch. Same algebra as reduce_segments, different caller.

Batch 1, Llama-3-8B, $s = 2560$, bf16 KV, block_size 16, --attention-backend TRITON_ATTN. Following one layer's attention from the model forward to the two kernel launches:

1. TritonAttentionMetadataBuilder.build() (vllm/v1/attention/backends/triton_attn.py:L197-L271) fills seq_lens = [2560], query_start_loc = [0, 1], max_query_len = 1, the block table row (160 int32), the build-time constants seq_threshold_3D = 16 and num_par_softmax_segments = 16, and the three preallocated softmax_segm_* scratch tensors.
2. TritonAttentionImpl.forward() (vllm/v1/attention/backends/triton_attn.py:L599-L749) writes the new K and V into the cache, then calls unified_attention(...) with max_seqlen_q=1.
3. unified_attention() (vllm/v1/attention/ops/triton_unified_attention.py:L802) computes num_queries_per_kv = 4, BLOCK_M = 16, BLOCK_Q = 4, total_num_q_blocks = 1 // 4 + 1 = 1, TILE_SIZE_DECODE = 16. The use_3d gate passes: max_seqlen_q is 1, num_seqs is 1 ≤ 16, batch invariance off. Grid becomes (1, 8, 16) = 128 blocks.
4. kernel_unified_attention[grid] — each block reads segm_idx = program_id(2), computes tiles_per_segment = cdiv(2560, 16*16) = 10, and compute_tile_loop_bounds (vllm/v1/attention/ops/triton_attention_helpers.py:L231-L236) narrows the tile loop to [segm_idx*10, (segm_idx+1)*10). Block 0 scans keys 0–159, block 15 scans 2400–2559.
5. Per tile: one tl.load from the block table gives the physical block, k_offset / v_offset address the 4 KB K and V runs, tl.dot produces a 16×16 score tile of which 4 rows are live, and the online softmax updates M, L, acc.
6. Each block then takes the IS_3D epilogue (vllm/v1/attention/ops/triton_unified_attention.py:L587-L642) and stores acc, M, L — without dividing by L. That omission is the point: normalisation cannot happen until every segment has reported.
7. reduce_segments[(1, 32)] (vllm/v1/attention/ops/triton_unified_attention.py:L1169-L1189) recomputes act_num_segments = cdiv(2560, 160) = 16, masks the surplus, takes overall_max, rescales, sums, divides once, and writes output[0, head, :] in bf16.

Two kernel launches, 128 + 32 thread blocks, 10.5 MB of KV read and 0.5 MB of partial state round-tripped, per layer. Times 32 layers: about 336 MB of KV plus 17 MB of partial-state round-trip traffic for one token, matching the 335 MB §0.2 derived from the KV size formula. The split does not reduce bytes. It only stops 124 SMs from watching.

Split-K breaks bitwise reproducibility

The merge sums $S$ fp32 partials, and $S$ depends on batch size and sequence length — that is, on who else is in the batch. Change $S$, change the summation order; fp32 addition is not associative; the logits move by a few ULP and eventually a different token is sampled. This is the mechanism behind "the same prompt gives different output at different load", and both engines ship an off switch.

vLLM disables the 3D path entirely under VLLM_BATCH_INVARIANT (vllm/v1/attention/ops/triton_unified_attention.py:L34 and the is_batch_invariant term in use_3d), and pins FA's split count to 1:

        max_num_splits = 0  # 0 means use FA3's heuristics, not CG compatible
        if (
            self.use_full_cuda_graph
            and self.max_cudagraph_size is not None
            and num_actual_tokens <= self.max_cudagraph_size
        ):
            # NOTE(woosuk): Setting num_splits > 1 may increase the memory
            # usage, because the intermediate buffers of size [num_splits,
            # num_heads, num_tokens, head_size] are allocated. Therefore,
            # we only set num_splits when using cuda graphs.
            max_num_splits = self.max_num_splits

        if envs.VLLM_BATCH_INVARIANT:
            max_num_splits = 1

SGLang instead fixes the tile so the count is a function of sequence length alone:

        self.prefill_split_tile_size = None
        self.decode_split_tile_size = None
        self.disable_cuda_graph_kv_split = False
        if self.enable_deterministic:
            self.decode_use_tensor_cores = True
            self.prefill_split_tile_size = get_int_env_var(
                "SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE", 4096
            )
            self.decode_split_tile_size = get_int_env_var(
                "SGLANG_FLASHINFER_DECODE_SPLIT_TILE_SIZE", 2048
            )
            self.disable_cuda_graph_kv_split = True

Its Triton backend does the same for its own kernel (python/sglang/srt/layers/attention/triton_backend.py:L376-L385). Full treatment in §10.4; the point here is that the non-determinism is caused by the optimisation in this chapter, and turning it off costs you the occupancy of Figure 1.

CUDA graphs need the split count pinned

A captured graph fixes launch geometry and buffer addresses, not necessarily every amount of useful work. A maximum grid with device-side masks can vary active split counts safely. Validity masks must agree between producer and reducer so unwritten partial rows are never consumed. Depending on the mistake, violations can produce stale results, an illegal access, or a runtime error.

Over-splitting short sequences

At s=128, TILE_SIZE=16 and 16 maximum segments give eight useful key segments. Inactive producers return and reducer loads can be masked, so 532 KB is an allocation/full-16-segment upper accounting, not guaranteed actual traffic. With eight written and read partials across 32 heads, useful partial traffic is about 266 KB against 524 KB of KV, still substantial. Measure cache traffic and masks instead of charging every padded row.

Kernel block size vs allocator block size

A backend that requires MultipleOf(16) and an allocator configured with --block-size 1 is a startup error, not a slow path. The two are reconciled by get_block_table_width / map_to_kernel_blocks, but only in the direction where kernel_block_size divides block_size:

    if block_size % kernel_block_size != 0:
        raise ValueError(
            f"kernel_block_size {kernel_block_size} must divide "
            f"block_size {block_size} evenly"
        )

Benchmarking a kernel that no longer runs

benchmarks/kernels/benchmark_paged_attention.py still exists and still says "paged attention". It now benchmarks only the ROCm kernel, and says so:

if __name__ == "__main__":
    logger.warning(
        "This script benchmarks the ROCm paged attention kernel. "
        "By default this is no longer used in vLLM inference."
    )
    if not current_platform.is_rocm():
        raise RuntimeError("This benchmark requires the ROCm platform.")

Numbers from an older checkout of it are not numbers about vLLM's CUDA decode path.

You need a GPU for the timings; the shape arithmetic you can do anywhere.

1. See the split decision flip. The gate is pure Python, so read it directly:

V=~/path/to/vllm
sed -n '54,56p;134,155p' $V/vllm/v1/attention/backends/triton_attn.py
sed -n '1036,1082p' $V/vllm/v1/attention/ops/triton_unified_attention.py

128 // num_kv_heads_per_rank is the batch size at which vLLM stops splitting. With threshold 2, batches above 2 use 2D; with threshold 128, batches up to 128 may use 3D when the other eligibility checks pass.

2. Watch the grid change. The cheapest observation is an Nsight Systems trace: two kernel names (kernel_unified_attention and reduce_segments) per layer on the split path, one on the unsplit path.

nsys profile -t cuda,nvtx -o decode_split \
  vllm bench latency --model meta-llama/Meta-Llama-3-8B-Instruct \
  --attention-backend TRITON_ATTN \
    --input-len 2560 --output-len 32 --batch-size 1 --num-iters 5

# then the same at --batch-size 32, which is above 128 // 8 = 16
# and should show reduce_segments disappear

3. Turn the split off. VLLM_BATCH_INVARIANT=1 forces the 2D path and pins FA's splits to 1. The TPOT difference at batch 1 and long context is the value of everything in §3 — plus whatever else batch invariance disables, so read §10.4 before attributing it all to split-K.

4. On the SGLang side the split count is a server flag: --triton-attention-num-kv-splits, default 8. Sweeping it from 1 to 16 at batch 1 with a long context isolates exactly the axis Figure 1 draws.

• FlashDecoding — Dao, Haziza, Massa, Sizov (2023). The original write-up of splitting the KV axis for batch-1 inference.
• Efficient Memory Management for Large Language Model Serving with PagedAttention — Kwon et al., SOSP 2023: the kernel §4.3 buries.
• FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving — section 2.2 is the merge that merge_attn_states.cu cites in its own header comment.
• vLLM PR #19152 — "[Kernel] Add Split-KV Support to Unified Triton Attention Kernel". Where the 3D path landed.
• vLLM PR #28306 — CUDA graphs for the 3D kernel, and #40631, which merged the separate 2D and 3D kernels into the IS_3D-parameterised one read above.
• vLLM PR #47361 — "Delete PagedAttention", and #43717, the torch-stable-ABI migration that moved it first.
• SGLang PR #4553 — "Optimize Triton decoding kernel for dynamic workload", which introduced the adaptive num_kv_splits computation.
• Next: §3.4 on backend selection and FlashInfer's plan/run split, and §3.5 on what happens to all of this when $g$ becomes 128.