7. Sequential vs. Vectorized Scan: Runtime & Memory

Problem

Cartesia's round description explicitly calls out discussing engineering tradeoffs, and "why is the parallel/vectorized scan actually better than the loop" is exactly that kind of question — you should be able to both implement and instrument the comparison, not just assert it from memory. This problem builds the harness: run the same linear recurrence through a naive Python-loop scan and a vectorized cumsum/cumprod-based scan, confirm they agree numerically (a correctness regression check you'd want before trusting any "optimization"), and report wall-clock time for each alongside the number of tensor elements each strategy must retain as intermediate activations for backprop.

The activation-memory count is the same order, O(batch * length * dim), for both strategies here — chunking (a separate technique, the previous problem) is what actually reduces it. The practical difference this harness is built to expose is op-count: the loop issues length separate elementwise ops (each launching its own kernel(s) on a GPU), while the vectorized version issues O(1) large tensor ops. At the toy sizes used in a coding round that gap may not show up in wall-clock time — the value here is having built the instrumentation and being able to reason about why it matters at production sequence lengths (thousands of timesteps), not a specific number this test asserts.

Source: src/7_scan_runtime_memory_benchmark.py

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

Examples:
>>> a = torch.rand(1, 100, 8)
>>> b = torch.randn(1, 100, 8)
>>> stats = compare_scan_strategies(a, b)
>>> stats["activation_elements"]
800

Step-by-Step Approach

1. Write (or reuse) two independent scan implementations: a Python-loop _sequential_scan and a closed-form _vectorized_scan using torch.cumprod/torch.cumsum (the same trick as the within-chunk formula from the chunking problem, applied over the whole sequence).
2. Time each with time.perf_counter() bracketing the call — not time.time(), which has coarser resolution and is affected by system clock adjustments; perf_counter is the standard choice for measuring short elapsed intervals.
3. Run both on the exact same a, b, h0 inputs so any timing or output difference is attributable to the algorithm, not the data.
4. Compute activation_elements = batch * length * dim as the memory-footprint proxy — both strategies must retain one hidden state per timestep for backprop, so this count doesn't distinguish them (a deliberate detail: it isolates "compute pattern" from "memory," since conflating the two is a common reasoning mistake).
5. Package everything into a single results dict with clearly named keys (sequential_output, vectorized_output, *_time_sec, activation_elements) so the caller can assert correctness (outputs match) independently of the non-deterministic timing numbers.
6. Do not assert *which* strategy is faster in a unit test — wall-clock timing on tiny CPU tensors is noisy and can go either way; assert only that both times are non-negative floats, and separately assert the two outputs agree.

The key insight is separating what's deterministic and testable (do the two implementations agree? is the memory-footprint formula right?) from what's inherently noisy and environment-dependent (which one actually ran faster on this machine, right now) — a real interview conversation about this tradeoff should reason from first principles (kernel launch count, memory bandwidth) rather than from a single unreliable local timing.

Reference solution

def _sequential_scan(a, b, h0=None):
    batch, length, dim = a.shape
    h_prev = torch.zeros(batch, dim, dtype=a.dtype) if h0 is None else h0
    outs = []
    for t in range(length):
        # length separate elementwise ops -- each a distinct kernel launch on a GPU
        h_prev = a[:, t, :] * h_prev + b[:, t, :]
        outs.append(h_prev)
    return torch.stack(outs, dim=1)


def _vectorized_scan(a, b, h0=None):
    batch, length, dim = a.shape
    h0 = torch.zeros(batch, dim, dtype=a.dtype) if h0 is None else h0
    # closed form via cumulative product/sum: O(1) large ops instead of O(length) small ones
    a_cumprod = torch.cumprod(a, dim=1)
    b_scaled = b / a_cumprod
    cum_b_scaled = torch.cumsum(b_scaled, dim=1)
    return a_cumprod * (h0.unsqueeze(1) + cum_b_scaled)


def compare_scan_strategies(a, b, h0=None):
    batch, length, dim = a.shape
    t0 = time.perf_counter(); seq_out = _sequential_scan(a, b, h0); t1 = time.perf_counter()
    t2 = time.perf_counter(); vec_out = _vectorized_scan(a, b, h0); t3 = time.perf_counter()
    return {
        "sequential_output": seq_out,
        "vectorized_output": vec_out,
        "sequential_time_sec": t1 - t0,
        "vectorized_time_sec": t3 - t2,
        # O(batch * length * dim): every timestep's hidden state must be retained for backward,
        # regardless of which strategy computed it -- chunking, not vectorizing, reduces this
        "activation_elements": batch * length * dim,
    }

Key Functions & Tricks

• time.perf_counter() — high-resolution monotonic clock, the standard choice for micro-timing in Python
• torch.cumprod / torch.cumsum — vectorized closed-form scan, O(1) op count regardless of sequence length
• torch.testing.assert_close(out1, out2) — the correctness half of the comparison, independent of timing
• Separating correctness checks from performance checks — deterministic assertions on outputs/shapes, non-asserted (informational) timing
• Kernel-launch-count reasoning — O(length) small ops vs. O(1) large ops, the real mechanism behind the vectorized version's GPU advantage

How to Recognize This Pattern