Making it parallel: the chunkwise form

This is the hard chapter. A recurrence whose write depends on its own read cannot be folded up with a prefix sum, and running 100,000 sequential steps on a GPU is not a training recipe. The way out reveals something better than a speed-up: attention and recurrence turn out to be the same computation at two settings of one knob.

S = torch.zeros(b, h, dh, dh) if cache is None else cache
outs = []
for i in range(t):                      # t = 131,072 for a long-context batch
    v_old = k[:,:,i:i+1] @ S
    u_i   = beta[:,:,i:i+1] * (v[:,:,i:i+1] - v_old)
    S     = S + k[:,:,i:i+1].transpose(-1,-2) @ u_i
    outs.append(q[:,:,i:i+1] @ S)
o = torch.cat(outs, dim=2)

Warm-up: chunking plain linear attention

Start with the easier case, chapter 2's purely additive state, and cut the sequence into blocks of \(C\) tokens. Write \(Q_r, K_r, V_r \in \mathbb{R}^{C\times d}\) for the rows belonging to block \(r\), and \(S_{r-1}\) for the state as it stood when the block began. Then for every token in the block at once:

Look at the two terms. They are the same triple product bracketed the two different ways from chapter 2, chosen independently for two different ranges:

• Inside the block we go score-first, \(q(k^{\top}v)\) evaluated as \((qk^{\top})v\) — a genuine \(C\times C\) attention matrix, causally masked. Exact, and cheap because \(C\) is small.
• Across blocks we go state-first, \((k^{\top}v)q\) — everything before this block has already been folded into \(S\), and one matmul reads it all back.

C is a hardware parameter, not a mathematical one

Both terms compute the same numbers regardless of \(C\). But their costs move in opposite directions. Per token, the work splits cleanly:

Set \(C = L\) and the second term becomes \(2L^2d\): full quadratic attention, recovered exactly. Set \(C = 1\) and it vanishes: the pure recurrence of chapter 2, with the fewest FLOPs of any setting. Fewer FLOPs, and yet nobody runs \(C=1\).

The reason is that a GPU's matrix units do not consume FLOPs, they consume tiles. Tensor-core instructions operate on fixed shapes — 16×8×16 on Ampere and Hopper, larger still with Blackwell's UMMA — and a matrix–vector product (\(C=1\)) fills almost none of that shape while reading the same operands. Its arithmetic intensity is about 1 FLOP per byte, the same pathology as chapter 1. At \(C = 64\) the intra-block matmul has intensity around \(C/2\), and the units are actually busy.

Now the delta twist

Equation 4.1 hoisted \(S_{r-1}\) out of the block because \(V_r\) did not depend on it. Under the delta rule it does. The two offending lines:

v_old = k_i @ S           # S here is the state *at position i*, not at the block start
u_i   = b_i * (v_i - v_old)

Every correction needs the state as it stood at its own position, and that state includes the corrections of every token before it inside the same block. There is no version of this you can hoist out.

The resolution: it is a triangular system

Unroll the recurrence within one block, holding \(S = S_{r-1}\) fixed at the block's entry. By definition \(S_{i-1} = S + \sum_{j

Read that carefully. Each \(u_i\) depends on the earlier \(u_j\) — but linearly, with known coefficients \(k_i \cdot k_j\). A set of linear equations in which each unknown depends only on earlier unknowns is a triangular system, and triangular systems are solved in closed form:

A = tril(K Kᵀ, −1), the strictly lower-triangular matrix of key–key dot products inside the block. Everything on the left is known before any \(u\) is computed.

Invert once and both dependencies fall out separately. Writing \(T = \big(I + \mathrm{diag}(\beta)A\big)^{-1}\), \(K_\beta = \mathrm{diag}(\beta)K\), \(V_\beta = \mathrm{diag}(\beta)V\):

This is the WY representation, borrowed from the numerical-linear-algebra literature on products of Householder matrices — which is exactly what \(\prod_j (I - \beta_j k_j^{\top}k_j)\) is. \(U_\beta\) and \(W\) depend only on \(K, V, \beta\), so they are computed for every block in the sequence in parallel, before any state exists. The only thing left that touches \(S\) is one matrix product per block.

The implementation

def chunk_delta_rule_forward(Q, K, V, beta, C):
    L, d = Q.shape
    Q, K, V = map(lambda x: x.reshape(-1, C, d), [Q, K, V])   # → (L/C, C, d)
    beta    = beta.reshape(-1, C)
    K_beta  = K * beta.unsqueeze(-1)
    V_beta  = V * beta.unsqueeze(-1)

    # ---- invert (I + diag(β)A) by forward substitution. C×C, all blocks at once ----
    T = -(K_beta @ K.t()).tril(-1)
    for i in range(1, C):
        T[i, :i] = T[i, :i] + (T[i, :, None] * T[:, :i]).sum(-2)
    T += torch.eye(C)
    W = T @ K_beta        # the part of U that depends on the incoming state
    U = T @ V_beta        # the part that does not

    # ---- now: exactly the chunked linear-attention loop, with U in place of V ----
    S = torch.zeros(d, d)
    O = torch.empty_like(V)
    for i in range(L // C):
        q_i, k_i, w_i = Q[i], K[i], W[i]
        u_i     = U[i] - w_i @ S           # every correction in the chunk, at once
        o_inter = q_i @ S                  # history before this chunk
        A_i     = (q_i @ k_i.t()).tril()   # masked attention inside the chunk
        o_intra = A_i @ u_i
        S      += k_i.t() @ u_i
        O[i]    = o_intra + o_inter
    return O.reshape(L, d)

Three things about this code are worth naming explicitly.

The for i in range(1, C) loop is sequential — and that is fine

It is forward substitution, and it terminates because a strictly lower-triangular \(C\times C\) matrix is nilpotent: the Neumann series \((I+L)^{-1} = I - L + L^2 - \cdots\) has at most \(C\) terms. The point is that the sequential depth is now \(C\), not \(L\). Sixty-four steps on a tiny matrix, not 131,072 steps on the full state. That is the entire trick.

The second loop is sequential too, over \(L/C\) blocks

Unavoidable — the state genuinely is a sequential object. But there are \(L/C\) iterations instead of \(L\), each doing dense matmuls that saturate the hardware, and the expensive parallel work (\(T\), \(W\), \(U\)) already happened for all blocks simultaneously.

Intermediate states are never materialised

The naive loop needs \(S_i\) at every position — \(O(d^2)\) live memory per token, which for backprop is ruinous. Equation 4.3 produces all the \(u_i\) using only \(O(d)\) per token, because the state dependence was factored into the single term \(W S_{r-1}\). The paper's memory-efficiency claim is this, not a kernel detail.

What the block actually looks like now