Naive attention and the online softmax

Here is textbook scaled dot-product attention, for one layer of Llama-3-8B, with exact shapes. $T$ is the number of tokens in the sequence, $h = 32$ query heads, $h_{kv} = 8$ key/value heads (grouped-query, so each KV head serves four query heads), $d_h = 128$ the head dimension, all tensors fp16 ($b = 2$ bytes).

$Q$ is $[h, T, d_h]$; $K$ and $V$ are $[h_{kv}, T, d_h]$ and are broadcast four ways. The offending object is $S$. It is $[h, T, T]$ — quadratic in sequence length and linear in head count, and it does not depend on $d_h$ at all, so it does not shrink when the model gets narrower. At $T = 8192$:

Derived. Four gibibytes, for one layer, one sequence, one intermediate tensor that is thrown away microseconds later. And PyTorch's softmax is not in-place, so $P$ is a second 4.00 GiB allocation living at the same time. Double the context to 32,768 and the quadratic bites: $32 \times 32768^{2} \times 2 = 64.0$ GiB for $S$ alone. On an 80 GB H100 that already holds 15 GB of Llama-3-8B weights, a single 32k-token sequence cannot be prefilled by the textbook algorithm at all — not slowly, not at all.

Capacity is the loud failure. The quiet one is worse. Count the HBM traffic of the three kernels a naive implementation launches, at $T = 8192$, one layer, fp16. Compulsory traffic only — each distinct byte crossing the HBM boundary once, the convention fixed in §0.4.

Ninety-nine percent of the bytes are an intermediate. The FLOPs are unchanged by any of this: $F = 2hT^{2}d_h$ for $QK^{\top}$ plus the same again for $PV$, so $F = 4 \times 32 \times 8192^{2} \times 128 = 1.0995 \times 10^{12}$ FLOP per layer. Divide:

Two remarks on how these are counted, so the numbers reconcile with the ones §3.2 quotes per head. First, $I_{\text{materialised}}$ is $T d_h / \big(b(T + d_h)\big)$ once the $4|S|$ term dominates — the head count cancels, so the per-head figure (34.36 GFLOP over 545.3 MB) and the per-layer figure are the same 63. Second, $I_{\text{compulsory}} = 2hT / \big(b(h + h_{kv})\big)$ does depend on the GQA ratio, because K and V are read once and reused across four query heads; count per head without that sharing and you get $T/b = 4096$ instead of 6553.6. The GQA-aware number is the honest one for a whole layer.

Against the H100 SXM ridge point of $I^{*} = 295$ FLOP/byte fixed in §0.4: the second number is 22× past the ridge — attention is intrinsically one of the most compute-dense things a transformer does. The first is 4.7× left of it. Materialising $S$ costs a factor of 103 in arithmetic intensity and takes a comfortably compute-bound operation and makes it memory-bound.

In time: at 3.35 TB/s, 16.156 GiB per layer is 5.18 ms, or 165.7 ms for all 32 layers. The compute-bound floor is $F/\pi = 1.0995\times10^{12} \times 32 / 989.4\times10^{12} = 35.6$ ms. Derived, both. Textbook attention is 4.7× slower than the same arithmetic has any right to be, and that ratio grows with $T$, because $S$'s bytes grow quadratically while Q/K/V/O's grow linearly.

Keep this in proportion. §0.2 established that attention is only about 7% of decode FLOPs for Llama-3-8B at 2.5k context, and §0.4 gave decode attention a fixed intensity of $I = 2h/(b\,h_{kv}) = 4$ FLOP/byte, invariant in both context length and batch size. So this entire part of the book optimises a small share of the FLOPs and a large share of the latency. That is not a contradiction; it is the definition of a memory-bound bottleneck, and it is why attention kernels get more engineering attention than the GEMMs that do fourteen times the arithmetic.

The whole part rests on one observation. Every row of $S$ is consumed by exactly one row of $O$, through a softmax and a weighted sum. Nothing outside that row ever reads it. So the row never needs to exist in HBM — it needs to exist in registers or shared memory for as long as it takes to fold it into the output, and then it can be discarded. Hopper gives you 228 KB of shared memory per SM (§0.3); one row of $S$ at $T = 8192$ in fp16 is 16 KB, and a tile of 64 keys against 64 queries is 8 KB. There is room.

The obstacle is the softmax. Its denominator $\sum_j e^{s_j}$ is a reduction over the whole row, so the textbook algorithm cannot emit a single correct element of $P$ until it has seen the last element of $S$ — which is exactly why the row has to be parked somewhere in the meantime. Break that dependency and the round trip disappears.

The right-hand panel is what §3.2 builds. Everything it needs is the box marked m, l, O accumulators: a scheme for carrying a partially-computed softmax forward across KV blocks, correcting it as new information arrives, and arriving at exactly the answer the three-kernel version would have produced. That scheme is the rest of this chapter.

Fix one query row. Let $x_1, \dots, x_N \in \mathbb{R}$ be its attention logits — the row of $S$ after the $1/\sqrt{d_h}$ scale and any mask — and let $v_1, \dots, v_N \in \mathbb{R}^{d_h}$ be the corresponding value vectors. The target is

(1) Naive softmax, and why it is unusable

Evaluate that definition literally and you compute $e^{x_j}$ before dividing. The exponential overflows at

An fp16 softmax dies at a logit of 11.1. Idealised logits do not reach it: if $q$ and $k$ have independent unit-variance entries, the $1/\sqrt{d_h}$ scale makes $x_j$ roughly $\mathcal{N}(0,1)$, and the expected maximum of an 8192-long row of such samples is only $\approx \sqrt{2\ln N} - \tfrac{\ln\ln N + \ln 4\pi}{2\sqrt{2\ln N}} = 3.7$ — nine sigma of margin. That is exactly why the naive form is dangerous rather than obviously broken: it works on every toy tensor you test it with. Trained models are not $\mathcal{N}(0,1)$ (the callout below is the architectural admission of it), and the failure when it comes is not gradual. Once one entry is $+\infty$, the denominator is $+\infty$, overflowed entries yield $\infty/\infty=\mathrm{NaN}$ while finite exponentials divided by the infinite denominator yield zero. Later weighted reductions can propagate those NaNs throughout the output. And upcasting is no defence: $e^{100} = 2.688\times10^{43}$ overflows fp32 too.

(2) Safe softmax, and why it costs three passes

The exponential's translation property rescues it. For any constant $c$,

For finite logits, subtracting their maximum makes exponent arguments nonpositive and prevents exp overflow. Under nearest-even with subnormals, fp16 rounds exp(z) to zero below z=log(2^-25), approximately -17.33; flush-to-zero hardware can differ. Tiny terms are often negligible but many such terms or large associated values can make their aggregate relevant. Check the reduction's error budget rather than declaring every underflow harmless.

The price is passes over the data:

Pass 2 cannot start before pass 1 finishes, and pass 3 cannot start before pass 2 finishes. If the row does not fit in on-chip memory, each pass is a separate trip to HBM — which is precisely the round trip Figure 1's left panel is drawing, and precisely why $S$ has to be materialised.

(3) The online recurrence

Milakov and Gimelshein's observation (arXiv:1805.02867, Online normalizer calculation for softmax, 2018) is that passes 1 and 2 can be fused into one, at the cost of a correction term. Split the row into $n$ consecutive blocks $B_1, \dots, B_n$ — for a kernel, one block is one tile of keys, say 64 of them. Maintain two scalars, initialised to the identity of their respective reductions:

and after consuming block $j$:

Two symbols to be clear about. $m^{(j)}$ is the running maximum over everything seen so far. $\ell^{(j)}$ is the running sum of exponentials, expressed relative to $m^{(j)}$ — that qualifier is the whole idea. And $\alpha^{(j)} = e^{m^{(j-1)} - m^{(j)}}$ is the rescaling factor: when a block arrives containing a logit larger than anything seen before, the old $\ell^{(j-1)}$ was normalised against a stale, too-small maximum, and every term inside it is too large by exactly $e^{m^{(j)} - m^{(j-1)}}$. Multiplying by $\alpha^{(j)}$ retro-normalises all of them at once, in one scalar multiply, without revisiting a single term. Note $\alpha^{(j)} \in (0, 1]$ always, since $m^{(j)} \ge m^{(j-1)}$ — the correction only ever shrinks. It cannot overflow.

The proof

This is the load-bearing claim of the whole part, so it gets a real proof.

Proof, by induction on $j$.

Base case, $j = 0$. $X_0 = \varnothing$. The maximum over the empty set is $-\infty$ by convention, matching $m^{(0)}$; the sum over the empty set is $0$, matching $\ell^{(0)}$.

Inductive step. Assume the claim at $j-1$. For the maximum, associativity of $\max$ gives it immediately:

For the sum, substitute the induction hypothesis into the update and use the single algebraic fact the entire construction rests on, $e^{a-b}\,e^{b-c} = e^{a-c}$:

The blocks $B_j$ partition $\{1,\dots,N\}$, so $X_{j-1}$ and $B_j$ are disjoint and the two sums concatenate without double counting. $\square$

Corollary. At the final tile the state is the maximum, shifted exponential sum, and shifted weighted-value sum over all processed entries, so normalization yields the same attention function in exact arithmetic. Floating-point results can differ with tiling and reduction order. Conditioning belongs to the mathematical problem; algorithmic stability is separate. Individual exponentials lie in [0,1], but ell can grow to the number of keys and the weighted output accumulator need not be bounded by 1.

Verify it on numbers

Eight logits, four blocks of two. This is small enough to check by hand and large enough to show a rescale that matters.

Block 3 is the interesting one: the maximum jumps from 3 to 8, and $\alpha^{(3)} = e^{-5} = 0.0067379$ divides the accumulated state by 148. Everything banked before block 3 was, in hindsight, almost irrelevant — the recurrence discovers this and corrects for it in one multiply. Now the direct computation: $m = 8$, and

Identical to $\ell^{(4)}$. Derived.

Extending the recurrence to the output

Softmax alone is not the goal; $o = \sum_j y_j v_j$ is. Keep an unnormalised output accumulator $\tilde{O}^{(j)} \in \mathbb{R}^{d_h}$ alongside $m$ and $\ell$, and read the recurrence off the same argument:

The claim $\tilde{O}^{(j)} = \sum_{k \in X_j} e^{x_k - m^{(j)}} v_k$ has the identical induction — the only change is that each term carries a vector factor $v_k$ that is untouched by the rescale. That is exactly why $O$ needs the same $\alpha^{(j)}$ and no other: the stale normalisation $e^{-m^{(j-1)}}$ is a common scalar factor of every accumulated term, so correcting it distributes out of the sum. Had the correction depended on $k$, you would have to revisit every term, and the whole scheme would collapse back to storing the row.

This is also the step where the pass count finally reaches one. Milakov and Gimelshein's result is a two-pass softmax: the recurrence fuses the max sweep and the sum sweep, but you still need a second sweep over the row to divide each $e^{x_j - m}$ by $\ell$ — and that second sweep is what would force the row back into memory. Attention escapes it because the only consumer of $y_j$ is the weighted sum $\sum_j y_j v_j$, and division by a scalar is linear, so it commutes past the sum and can be deferred: accumulate $\tilde{O}$ unnormalised, divide once, at the very end, by one number. Attention gets a one-pass online softmax not because softmax admits one, but because $PV$ absorbs the normalisation. Every element of $S$ is now produced, consumed and discarded in registers, and never written anywhere.

Cost accounting: the rescale is $d_h$ multiplies for $\tilde{O}$ plus one for $\ell$, once per KV block, per query row. Against the $T \cdot d_h$ multiply-adds of the main accumulation that is a relative overhead of $1/B_c$ where $B_c$ is the block width in keys — 1.6% at $B_c = 64$. Derived. You spend under two percent more arithmetic to delete 99% of the memory traffic. That trade is the entire thesis of §3.2.

What has to be fp32

§0.5's rule was: low precision is safe for operands, dangerous for reductions. Every state variable here is a reduction, and each fails differently in fp16.

You will not find the recurrence written out as a recurrence in either engine, because inside FlashAttention it is compiled into a CUTLASS mainloop that §3.2 takes apart. But both engines expose the merge form of it as a standalone kernel — the operation that takes two partial attention results over disjoint key ranges and combines them. That kernel is the recurrence with the induction step written for two arbitrary states instead of state-plus-block, and reading it is the fastest way to see the mathematics as code.

The state that gets passed around is the log-sum-exp

Neither engine passes $(m, \ell)$ across a kernel boundary. They pass a single fp32 scalar per (token, head), the log-sum-exp:

One number instead of two, and it is the natural quantity: $e^{\mathrm{LSE}_A}$ is literally the total unnormalised mass of partition $A$. Given two disjoint ranges with normalised outputs $O_A, O_B$ and their LSEs, the merge is a convex combination weighted by those masses:

Which overflows, for exactly the reason §(1) gave. So you apply the safe-softmax shift one level up — subtract $M = \max(\mathrm{LSE}_A, \mathrm{LSE}_B)$ from both, which §(2) showed is free because softmax is shift-invariant, and which is legal for any common shift, not just the true maximum of the underlying logits. That gives the algorithm, and the algorithm is the kernel, line for line.

One more reconciliation before the code. The recurrence in §3 is asymmetric — a running state absorbing one new block — while the merge above is symmetric, treating both sides alike. They are the same operator. Take the merge and set side $A$ to the running state $(\tilde{O}^{(j-1)}/\ell^{(j-1)},\ m^{(j-1)} + \log \ell^{(j-1)})$ and side $B$ to a single block's own partial result; expand $e^{\mathrm{LSE}_A - M}$ and the $\alpha^{(j)}$ of §3 falls out, with $M$ playing the role of $m^{(j)}$. The streaming form is the merge specialised to "the second argument is one tile"; the merge is the streaming form with the privilege removed. Which is why the same algebra serves a FlashAttention mainloop, a split-KV reduction and a prefix/suffix combine — though, as §5 shows, those are three different kernels in vLLM, not three callers of one.

vLLM: merge_attn_states.cu

vLLM's CUDA implementation lives at csrc/libtorch_stable/attention/merge_attn_states.cu (362 lines at a556f3f). Its own header comment names the reference: // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 (csrc/libtorch_stable/attention/merge_attn_states.cu:L15). Here is the arithmetic, verbatim:

  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;

  if (pack_offset < head_size) {
    input_pack_t p_out_pack = reinterpret_cast<const input_pack_t*>(
        prefix_head_ptr)[pack_offset / pack_size];
    input_pack_t s_out_pack = reinterpret_cast<const input_pack_t*>(
        suffix_head_ptr)[pack_offset / pack_size];

    // Compute merged values in float32
    float o_out_f[pack_size];
#pragma unroll
    for (uint i = 0; i < pack_size; ++i) {
      const float p_out_f =
          vllm::to_float(reinterpret_cast<const scalar_t*>(&p_out_pack)[i]);
      const float s_out_f =
          vllm::to_float(reinterpret_cast<const scalar_t*>(&s_out_pack)[i]);
      o_out_f[i] = p_out_f * p_scale + (s_out_f * s_scale);
    }

Map it onto the equations. max_lse (set at csrc/libtorch_stable/attention/merge_attn_states.cu:L106 as fmaxf(p_lse, s_lse)) is M=max(LSE_A,LSE_B), not generally the maximum raw logit m. The two expf results are $e^{\mathrm{LSE}_A - M}$ and $e^{\mathrm{LSE}_B - M}$ — the rescaling factors $\alpha$, in their symmetric form where neither side is privileged. out_se is the merged $\ell$ relative to $M$; p_scale and s_scale are the normalised mixing weights, which sum to exactly 1 by construction, making o_out_f a convex combination — it cannot leave the range of its inputs. And the accumulation is fp32 even though scalar_t may be __nv_bfloat16: vllm::to_float on the way in, vllm::from_float on the way out, exactly the operand-vs-accumulator split of §0.5. The LSE pointers in the kernel signature are const float* unconditionally (csrc/libtorch_stable/attention/merge_attn_states.cu:L21-L22) — the running sum is never allowed to be narrow, whatever the tensors are.

Then the state for the next merge is written back in log space:

  // We only need to write to output_lse once per head.
  if (output_lse != nullptr && pack_idx == 0) {
    float out_lse = logf(out_se) + max_lse;
    output_lse[head_idx * output_lse_head_stride +
               token_idx * output_lse_token_stride] = out_lse;
  }

logf(out_se) + max_lse is $\log \ell + m$ — the LSE definition, undoing the shift. Because the output has the same type as the inputs, merges compose: this is what makes the operator associative, and associativity is what licenses everything built on it downstream. Formally, states $(O, \mathrm{LSE})$ under this merge form a commutative monoid with identity $(\mathbf{0}, -\infty)$. Streaming blocks left-to-right, reducing split-KV partials in a tree, or combining a cached prefix with a fresh suffix are all the same fold in different orders, and all give the same answer.

SGLang: the same thing, in Triton

SGLang dispatches through python/sglang/srt/layers/attention/merge_state.py (46 lines), preferring a compiled kernel and falling back to Triton:

    if (
        _is_cuda
        and _supported_dtypes(prefix_output)
        and _supported_headdim(prefix_output)
    ):
        return merge_state_v2(
            prefix_output, prefix_lse, suffix_output, suffix_lse, output, output_lse
        )
    else:
        # Fallback to Triton kernel
        return merge_state_triton(
            prefix_output, prefix_lse, suffix_output, suffix_lse, output, output_lse
        )

The Triton fallback is the clearest statement of the recurrence anywhere in either codebase — no packing, no dtype dispatch, just the algebra:

    p_lse = tl.load(prefix_lse + token_idx * num_heads + head_idx)
    s_lse = tl.load(suffix_lse + token_idx * num_heads + head_idx)
    p_lse = float("-inf") if p_lse == float("inf") else p_lse
    s_lse = float("-inf") if s_lse == float("inf") else s_lse

    max_lse = tl.maximum(p_lse, s_lse)
    p_lse = p_lse - max_lse
    s_lse = s_lse - max_lse
    out_se = tl.exp(p_lse) + tl.exp(s_lse)

    if OUTPUT_LSE:
        out_lse = tl.log(out_se) + max_lse
        tl.store(output_lse + token_idx * num_heads + head_idx, out_lse)

    p_scale = tl.exp(p_lse) / out_se
    s_scale = tl.exp(s_lse) / out_se
    out = p_out * p_scale + s_out * s_scale
    tl.store(
        output + token_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE + head_arange,
        out,
        mask=head_mask,
    )

Eleven lines of arithmetic, and every one of them appears in the CUDA kernel above. The engines differ only in packaging: vLLM keeps its CUDA kernel in-tree under csrc/ with a Triton fallback at vllm/v1/attention/ops/triton_merge_attn_states.py selected by vllm/v1/attention/ops/merge_attn_states.py:L80-L100; SGLang's compiled merge_state_v2 comes from the separately-built sgl_kernel package, with the in-tree Triton kernel as the AMD / odd-head-dim / fp8 path. The dispatch predicates are almost identical — head dimension a multiple of 8 for half types, 4 for fp32, because both CUDA kernels issue 128-bit vector loads (python/sglang/srt/layers/attention/merge_state.py:L19-L23, vllm/v1/attention/ops/merge_attn_states.py:L70-L78).

For the record on where this operation is not: at a556f3f vLLM has deleted its bespoke PagedAttention CUDA kernel (commit d715b3aa1e, Delete PagedAttention (#47361)). csrc/attention/ now contains only dtype headers. CUDA attention runs through FlashAttention, FlashInfer and Triton backends, which is where §3.2 and §3.4 pick the thread up.

vLLM's cascade-attention path is the shortest route from Python to the equations. It runs when a batch shares a long common prefix: attend to the shared prefix once for all queries, and to each query's private suffix separately, then merge. Trace it.

1. Two partial attentions. cascade_attention calls flash_attn_varlen_func twice, both with return_softmax_lse=True. The second call is quoted at vllm/v1/attention/backends/flash_attn.py:L1771-L1791 — note causal=True and block_table=block_table[:, num_common_kv_blocks:], the suffix being everything past the shared prefix. Each call returns $(O_A, \mathrm{LSE}_A)$: an already-normalised output over its own key range, plus the scalar that lets it be un-normalised later.
2. The merge call.
   vllm/v1/attention/backends/flash_attn.py:L1793-L1794 vLLM

       # Merge prefix and suffix outputs, and store the result in output.
       merge_attn_states(output, prefix_output, prefix_lse, suffix_output, suffix_lse)

3. Dispatch. vllm/v1/attention/ops/merge_attn_states.py:L80-L100 checks platform, dtype and head-dim divisibility, then calls either vllm._custom_ops.merge_attn_states (the CUDA kernel) or the Triton one. For Llama-3-8B, $d_h = 128$ is a multiple of 8 and bf16 is supported, so the CUDA path wins.
4. One thread's work. The kernel maps global_idx -> token_idx + head_idx + pack_idx (csrc/libtorch_stable/attention/merge_attn_states.cu:L44-L49) with threads_per_head = head_size / pack_size. For bf16, pack_size = 8, so $d_h = 128$ is covered by 16 threads, each handling 16 bytes — one 128-bit vector load. Every one of those 16 threads reloads the same p_lse and s_lse and recomputes the same two scales; only pack_idx == 0 writes the merged LSE back. Recomputing a scalar in 16 threads is cheaper than a shared-memory broadcast and a barrier.
5. The output. output[token, head, :] is written once, in the requested dtype, and output_lse (if requested) carries the merged state forward for a further merge.

Total HBM traffic for the merge, per token-head: two $d_h$-vectors and two fp32 scalars in, one $d_h$-vector out. At $d_h = 128$ in bf16 that is 776 bytes moved for roughly 390 FLOPs — $I \approx 0.5$ FLOP/byte, about as memory-bound as a kernel can be. Derived. It is tolerable only because it runs once per split, not once per key: the whole point of the design is that the expensive fold happened inside the attention kernel, in SRAM, and only the $(O, \mathrm{LSE})$ summary crossed the HBM boundary.

The base case is a NaN factory

The proof's base case sets $m^{(0)} = -\infty$ and $\ell^{(0)} = 0$, and the induction step computes $\alpha = e^{m^{(0)} - m^{(1)}}$. If $m^{(1)}$ is also $-\infty$ — a query row with no unmasked keys at all — then $\alpha = e^{-\infty + \infty} = e^{\mathrm{NaN}} = \mathrm{NaN}$, and the NaN propagates to every element of the output. In real arithmetic the empty case is harmless; in IEEE-754 it is a landmine, and both engines have stepped on it. vLLM's kernel carries the incident report inline:

  p_lse = std::isinf(p_lse) ? -std::numeric_limits<float>::infinity() : p_lse;
  s_lse = std::isinf(s_lse) ? -std::numeric_limits<float>::infinity() : s_lse;

  const float max_lse = fmaxf(p_lse, s_lse);

  /* In certain edge cases, MLA can produce p_lse = s_lse = -inf;
     continuing the pipeline then yields NaN. Root cause: with chunked prefill
     a batch may be split into two chunks; if a request in that batch has no
     prefix hit, every LSE entry for that request's position is -inf, and at
     this moment we merge cross-attention at first. For now we simply emit
     prefix_output (expected to be all zeros) and prefix_lse (-inf) to fix
     this problem.
  */
  if (std::isinf(max_lse)) {

Two separate defences in four lines. The first maps $+\infty$ to $-\infty$, because FA2 signals "empty range" with the wrong sign — vLLM's Triton kernel spells out the reasoning at vllm/v1/attention/ops/triton_merge_attn_states.py:L160-L164: "If we see an inf assume FA2 and convert inf to -inf for consistency and correctness." SGLang's Triton kernel carries the identical guard at python/sglang/kernels/ops/attention/merge_state.py:L27-L28. The second is the early-out for the both-empty case. vLLM's Triton version guards the same case with tl.where instead of a branch, at vllm/v1/attention/ops/triton_merge_attn_states.py:L176-L178 and :L207-L209, with the comment "If both sides are empty (max_lse == -inf) the scales are 0/0 = NaN; emit zeros rather than NaN."

The debugging signature: a NaN that appears only under chunked prefill, only for requests that missed the prefix cache, only in the first layer's output, and vanishes if you disable cascade or chunked prefill. If you write your own merge, write the empty-state guard first.

Scale first, then multiply

The order of operations in the merge is load-bearing, and vLLM's Triton kernel says so:

    # NOTE(woosuk): Be careful with the numerical stability.
    # We should compute the scale first, and then multiply it with the output.
    # Do not multiply the output with tl.exp(p_lse) or tl.exp(s_lse) directly.
    p_scale = p_se / out_se
    s_scale = s_se / out_se
    out = p_out * p_scale + s_out * s_scale

Both forms are algebraically identical. Numerically they are not: $p\_scale$ and $s\_scale$ are non-negative and sum to exactly 1, so out is a convex combination bounded by $\max(|p\_out|, |s\_out|)$ and cannot overflow whatever dtype it is stored in. Multiply by the raw $e^{\mathrm{LSE}-M}$ first and divide afterwards and you pass through an unbounded intermediate. The same instinct is why the running $\ell$ is carried normalised against $m$ rather than raw.

Bitwise reproducibility is a casualty

The merge is associative in real arithmetic and only approximately so in floating point. Change the number of KV splits — which both engines do adaptively, based on batch size and sequence length — and you change the fold order, and the last bits of the output move. vLLM has an explicit escape hatch for this: at vllm/v1/attention/backends/flash_attn.py:L1790 the split count is pinned when batch-invariant mode is on:

        num_splits=1 if envs.VLLM_BATCH_INVARIANT else max_num_splits,

num_splits=1 means one partial, so no merge, so no order-dependence. It also means no split-KV parallelism, which is the throughput lever §3.3 is built on. Determinism and decode throughput are in direct tension here, and the flag is where you choose.

Reproduce the proof numerically before trusting any kernel that claims it. No GPU required — this is CPU numpy and takes a second.

import numpy as np
rng = np.random.default_rng(0)
N, D, BC = 8192, 128, 64                 # keys, head dim, KV block width
x = rng.normal(0, 3, N).astype(np.float64)   # logits, sigma=3 so exp(x) overflows fp16
x[4000] = 40.0                           # a late outlier: forces a big rescale
v = rng.normal(0, 1, (N, D))

m, l, O = -np.inf, 0.0, np.zeros(D)
for j in range(0, N, BC):
    xb, vb = x[j:j+BC], v[j:j+BC]
    m_new = max(m, xb.max())
    a = np.exp(m - m_new) if np.isfinite(m) else 0.0   # the base-case guard
    p = np.exp(xb - m_new)
    l = a * l + p.sum()
    O = a * O + p @ vb
    m = m_new
online = O / l

ref = (np.exp(x - x.max()) / np.exp(x - x.max()).sum()) @ v   # safe 3-pass
np.testing.assert_allclose(online, ref, rtol=1e-12, atol=1e-12)
print(np.max(np.abs(online - ref)))      # rounding-scale difference

Then, on a machine with a GPU, exercise the real kernels. vLLM's merge has a dedicated benchmark and test:

# exercise the CUDA and Triton merge kernels against a torch reference
pytest -q tests/kernels/attention/test_merge_attn_states.py

# force the Triton fallback path by asking for an unsupported head dim,
# then diff the two implementations at your own shapes
python -c "import vllm.v1.attention.ops.merge_attn_states as m; print(m.__file__)"

# turn cascade attention off and on and watch the merge disappear from the profile
VLLM_BATCH_INVARIANT=1 vllm serve meta-llama/Meta-Llama-3-8B-Instruct

• Milakov & Gimelshein, Online normalizer calculation for softmax, arXiv:1805.02867 (2018). The two-pass softmax and the running-max recurrence, three years before FlashAttention used them. Short, and worth reading in full.
• Rabe & Staats, Self-attention Does Not Need $O(n^2)$ Memory, arXiv:2112.05682 (2021). The same recurrence applied to attention, with the memory argument made explicitly.
• arXiv:2501.01005, section 2.2 — the reference cited by name in vLLM's own source at csrc/libtorch_stable/attention/merge_attn_states.cu:L15, vllm/v1/attention/ops/merge_attn_states.py:L22 and csrc/cpu/mla_decode.cpp:L129. Read it as the engines' chosen statement of the merge formulation.
• vLLM PR #16173, [Kernel] support merge_attn_states CUDA kernel, 3x speedup — the commit that introduced the kernel dissected in §4, with the author's benchmark discussion.
• vLLM PR #47361, Delete PagedAttention (commit d715b3aa1e). Context for why there is no bespoke vLLM attention CUDA kernel to read at this SHA.
• The backend set this part will work through, enumerated for §3.4: SGLang's python/sglang/srt/layers/attention/ holds flashattention_backend.py, flashinfer_backend.py, flashinfer_mla_backend.py, triton_backend.py, torch_native_backend.py, torch_flex_backend.py, trtllm_mha_backend.py, trtllm_mla_backend.py, flashmla_backend.py, cutlass_mla_backend.py, cutedsl_mla_backend.py, aiter_backend.py, wave_backend.py, xpu_backend.py, intel_amx_backend.py, nsa_backend.py, dsa_backend.py, hybrid_attn_backend.py, hybrid_linear_attn_backend.py and tbo_backend.py, selected through python/sglang/srt/layers/attention/attention_registry.py.