5. Parallel Associative Scan

Problem

The sequential scan from the previous problem (h_t = a_t * h_{t-1} + b_t via a Python for-loop) is unambiguously correct but launches one dependent operation per timestep — on a GPU that means length sequential kernel launches with no cross-timestep parallelism, which is exactly the bottleneck real SSM implementations (S4, S5, Mamba's selective_scan_cuda) work around by reformulating the recurrence as an associative scan.

The trick: represent each timestep as an affine map (a_t, b_t) meaning "x -> a_t * x + b_t", and define a combine operator that composes two such maps: applying (a1, b1) then (a2, b2) is equivalent to the single map (a1 * a2, a2 * b1 + b2). This combine operator is associative, so elements can be combined pairwise in any grouping — which is what lets a scan over length L run in O(log L) sequential steps (O(L log L) total work) instead of O(L) sequential steps, in the spirit of the Blelloch/Hillis-Steele family of parallel scan algorithms used to parallelize S4/S5/Mamba-style linear recurrences.

Source: src/5_parallel_associative_scan.py

def parallel_scan(a: torch.Tensor, b: torch.Tensor, h0: torch.Tensor | None = None) -> torch.Tensor:
    ...

Examples:
>>> a = torch.tensor([[[0.5], [0.5], [0.5]]])
>>> b = torch.tensor([[[1.0], [1.0], [1.0]]])
>>> parallel_scan(a, b)
tensor([[[1.0000], [1.5000], [1.7500]]])  # matches sequential_scan exactly

Step-by-Step Approach

1. Define the combine operator precisely: given an earlier map (a1,b1) and a later map (a2,b2), their composition is (a1*a2, a2*b1+b2). Get the order right — combine is not commutative, since applying gate 1 before gate 2 is not the same as the reverse.
2. Initialize a_cum = a.clone(), b_cum = b.clone() — these track, at each position, the composition of every map seen "so far" in the current doubling round.
3. Loop with offset = 1, 2, 4, ... while offset < length (a while loop over ceil(log2(length)) iterations, not over length).
4. At each round, build a shifted copy of (a_cum, b_cum) moved right by offset positions, padding the newly-exposed left edge with the identity map (a=1, b=0) (composing with the identity is a no-op, exactly like padding a cumsum with zero).
5. Combine the shifted (earlier) copy with the current (later) (a_cum, b_cum) using the operator from step 1, and overwrite a_cum, b_cum with the result.
6. After the loop, (a_cum, b_cum) at position t holds the full composition of maps 0..t. If h0 is given, the final output is a_cum * h0 + b_cum (applying the composed map to the initial state); otherwise it's just b_cum.
7. Verify against the sequential version on a non-power-of-2 length (e.g. 5) — the identity-padding logic is the most common source of bugs, and off-by-one errors there usually only show up on lengths that aren't a clean power of 2.

The key insight is that after each doubling round, every position holds the composition of twice as many original maps as before — so ceil(log2(length)) rounds suffice to cover the whole sequence, turning length sequential dependencies into O(log length) sequential rounds of large, fully parallel tensor ops.

Reference solution

def parallel_scan(a: torch.Tensor, b: torch.Tensor, h0: torch.Tensor | None = None) -> torch.Tensor:
    batch, length, dim = a.shape
    a_cum = a.clone()  # (batch, length, dim), running composed "a" of the affine map
    b_cum = b.clone()  # running composed "b" of the affine map
    offset = 1
    while offset < length:
        # identity affine map (a=1, b=0) pads positions with no partner `offset` behind them
        ident_a = torch.ones(batch, offset, dim, dtype=a.dtype, device=a.device)
        ident_b = torch.zeros(batch, offset, dim, dtype=b.dtype, device=b.device)
        # shift the running composition right by `offset`: this is the "earlier" operand
        a_shift = torch.cat([ident_a, a_cum[:, :-offset, :]], dim=1)
        b_shift = torch.cat([ident_b, b_cum[:, :-offset, :]], dim=1)
        # combine((a1,b1) earlier, (a2,b2) later) = (a1*a2, a2*b1 + b2)
        new_a = a_shift * a_cum
        new_b = a_cum * b_shift + b_cum
        a_cum, b_cum = new_a, new_b
        offset *= 2  # doubling => ceil(log2(length)) steps total, not `length` steps
    if h0 is not None:
        # h_t = A_t * h0 + B_t, where (A_t, B_t) is the fully composed affine map up to t
        return a_cum * h0.unsqueeze(1) + b_cum
    return b_cum

Key Functions & Tricks

• torch.cat([ident, a_cum[:, :-offset, :]], dim=1) — shift-with-identity-padding along the time axis
• while offset < length: ... offset *= 2 — log-depth doubling loop (Hillis-Steele style), the parallel-scan idiom
• Associative combine operator on (a, b) pairs — composing affine maps instead of applying the recurrence directly
• a_cum * h0.unsqueeze(1) + b_cum — applying the fully-composed affine map to the initial state in one shot
• .clone() — avoid mutating the input tensors in place before they're needed for the shift

How to Recognize This Pattern