5. Chunked Scan for Long Sequences

Problem

For a real-time voice model streaming minutes of audio, materializing the full-length scan output all at once — or holding every timestep's activations for backprop — doesn't scale: memory grows linearly with sequence length, which is exactly the cost SSMs are supposed to avoid relative to attention. The standard fix is a chunked (blocked) scan: split the sequence into fixed-size chunks, compute each chunk's local scan with a closed-form vectorized formula, and carry only a single (batch, dim) hidden state across chunk boundaries — so peak intermediate memory is bounded by one chunk, not the whole sequence, and the only sequential Python-level loop is over num_chunks = length / chunk_size (far fewer iterations than a per-timestep loop).

Within a chunk, the closed form comes from unrolling h_t = a_t * h_{t-1} + b_t: writing A_t = prod(a_1..a_t) as the chunk-local cumulative product, h_t = A_t * h_prev + A_t * cumsum(b_k / A_k for k=1..t). This uses torch.cumprod and torch.cumsum instead of a timestep loop within the chunk. Note the division by A_k means very long chunks re-invite numerical trouble (A_k can underflow toward zero for decaying gates) — chunking bounds both memory and the numerical error of this cumsum trick, which is part of why production scan kernels keep chunk sizes modest rather than relying on the trick over arbitrarily long spans.

Source: src/5_chunked_blocked_scan.py

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

Examples:
>>> a = torch.tensor([[[0.5], [0.6], [0.7], [0.8]]])
>>> b = torch.tensor([[[1.0], [1.0], [1.0], [1.0]]])
>>> chunked_scan(a, b, chunk_size=2)
tensor([[[1.0000], [1.6000], [2.1200], [2.6960]]])

Step-by-Step Approach

1. Assert length % chunk_size == 0 for simplicity (a production version would pad or handle a ragged final chunk; call this assumption out explicitly rather than silently mishandling it).
2. Initialize the carried state h to h0 or zeros, shape (batch, dim) — this is the only thing that crosses a chunk boundary.
3. Loop over chunk start indices 0, chunk_size, 2*chunk_size, ... (only num_chunks iterations). Slice out a_chunk, b_chunk of shape (batch, chunk_size, dim).
4. Within the chunk, compute a_cumprod = torch.cumprod(a_chunk, dim=1) — this is A_t, the product of gates from the start of the chunk through position t.
5. Compute b_scaled = b_chunk / a_cumprod, then cum_b_scaled = torch.cumsum(b_scaled, dim=1) — together these realize sum_k b_k / A_k via a running sum instead of a nested loop.
6. Combine: h_chunk = a_cumprod * (h.unsqueeze(1) + cum_b_scaled) gives every hidden state within the chunk in one vectorized expression, correctly incorporating the carried-in h.
7. Update the carried state to the chunk's last position, h = h_chunk[:, -1, :], and append h_chunk to the output list; after the loop, torch.cat the chunks along the time axis.

The key insight is that the within-chunk formula does real work vectorized (cumprod/cumsum over chunk_size, not a Python loop), while the only genuinely sequential Python-level iteration is over the much smaller num_chunks — this is the same "trade some parallelism for bounded memory" idea a production streaming decoder uses, just expressed as training-time chunking here instead of true incremental inference.

Reference solution

def chunked_scan(a: torch.Tensor, b: torch.Tensor, chunk_size: int, h0: torch.Tensor | None = None) -> torch.Tensor:
    batch, length, dim = a.shape
    assert length % chunk_size == 0, "length must be divisible by chunk_size"
    h = torch.zeros(batch, dim, dtype=a.dtype, device=a.device) if h0 is None else h0
    chunks = []
    for start in range(0, length, chunk_size):  # only num_chunks iterations, not `length`
        a_chunk = a[:, start:start + chunk_size, :]
        b_chunk = b[:, start:start + chunk_size, :]
        # closed form within the chunk: A_t = cumulative product, no per-timestep loop
        a_cumprod = torch.cumprod(a_chunk, dim=1)  # (batch, chunk_size, dim)
        b_scaled = b_chunk / a_cumprod             # b_k / A_k
        cum_b_scaled = torch.cumsum(b_scaled, dim=1)
        # h_t = A_t * (h_prev_from_last_chunk + cumsum(b_k / A_k))
        h_chunk = a_cumprod * (h.unsqueeze(1) + cum_b_scaled)
        chunks.append(h_chunk)
        h = h_chunk[:, -1, :]  # carry only the last state across the chunk boundary
    return torch.cat(chunks, dim=1)

Key Functions & Tricks

• torch.cumprod(a_chunk, dim=1) — vectorized running product within a chunk, replaces a per-timestep multiply loop
• torch.cumsum(b_scaled, dim=1) — vectorized running sum, paired with cumprod to realize the closed-form scan solution
• a[:, start:start + chunk_size, :] — chunk slicing along the time axis
• Bounded-memory streaming pattern — carry only a fixed-size state across an outer loop, vectorize the inner work
• h_chunk[:, -1, :] — extract the last timestep of a chunk to hand off as the next chunk's initial state

How to Recognize This Pattern