parallel_scan down-sweep is incorrect — produces wrong outputs for power-of-2 lengths

[eeshsaxena]

Bug: parallel_scan Falls Back to Sequential for Non-Power-of-2 Lengths — But Most Real Sequences Are Not Powers of 2

In mamba/selective_scan.py:

def parallel_scan(A, X):
    if length & (length - 1) != 0:
        return sequential_scan(A, X)  # silent fallback
    ...

For real training sequences (e.g. length 512, 1024, 2048), this actually works. But for any sequence that is not a power-of-2 — including padded batches where the last batch has a different length, or inference with variable-length inputs — the code silently falls back to the O(N) sequential scan without warning the caller.

Additionally, the down-sweep phase of parallel_scan has a bug: it iterates but only sets result[:, :, right_idx] when left_idx > 0, which means result[:, :, 0] is never set for intermediate positions. This makes the parallel scan output incorrect even for power-of-2 lengths, so callers get wrong gradients during training.

Impact

• Incorrect outputs from parallel_scan corrupt training gradients.
• The fallback to sequential_scan is silent — users have no idea they're getting O(N) instead of O(N log N) behavior.

Suggested Fix

1. Raise a warning (or error) when falling back to sequential mode.
2. Fix the down-sweep to correctly initialize all output positions, or replace with a validated reference implementation (e.g., using torch.cumprod).
3. Add a test that validates parallel_scan(A, X) == sequential_scan(A, X) for random inputs of power-of-2 length.