C.W.K.
Stream
Lesson 02 of 05 · published

Hardware-Aware Selective Scan

~14 min · scan, cuda, kernels, memory-hierarchy

Level 0Observer
0 XP0/50 lessons0/14 achievements
0/100 XP to next level100 XP to go0% complete

The convolution trick is gone — what now?

Selectivity broke the convolution view, so Mamba needed a new way to be parallel during training. The answer is the hardware-aware selective scan: a parallel prefix-scan algorithm written specifically for the GPU memory hierarchy, keeping intermediate state in fast on-chip SRAM and avoiding the trip to slow HBM that would otherwise dominate runtime.

This is the same engineering philosophy as FlashAttention. Both papers come from Tri Dao's group; both are answers to the question "how do you make a complex sequence operation actually utilize a modern GPU?" The selective scan kernel (in mamba_ssm.ops.selective_scan_interface) is several hundred lines of CUDA that essentially does this: tile the sequence, keep tiles in SRAM, perform the recurrence locally, write back only the final state per tile. The result is GPU utilization comparable to a well-implemented attention kernel, despite the operation being structurally a recurrence rather than a matmul.

Architecture-level simplification

Mamba also threw out a lot of architectural overhead. The standard Transformer block has both an attention sub-layer and an MLP sub-layer per block. Mamba blocks have only the SSM sub-layer — no separate MLP. The internal structure is two parallel branches per block: one branch goes through a Conv1D and then the selective SSM; the other branch is a SiLU-gated linear projection. The two branches are combined elementwise, then projected back to model dim. Fewer parameters per layer, fewer ops, simpler graph.

Real wall-clock gains

Mamba reports 5× inference throughput over Transformers at equivalent quality, with the gap widening at longer sequences (because Transformer attention is the part that's quadratic, while Mamba's per-token cost is constant in sequence length). At the 3B scale, Mamba matches or exceeds Transformer models roughly twice its parameter count. These aren't just FLOP-count claims; the hardware-aware kernels mean the FLOP advantage actually shows up on the wall clock.

The takeaway: efficient algorithms aren't enough — you need hardware-shaped implementations. SSMs got their second wind from the 2010s not because the math improved (it had been there for decades) but because someone finally wrote the kernels.

External links

Exercise

Read the selective_scan_interface.py source. You don't have to understand the CUDA in depth — but trace through the Python wrapper and identify the points where the code could fall back to a non-CUDA path, and what would be lost in that case. (Hint: the torch reference implementation in selective_scan_ref is mathematically equivalent and is hundreds of times slower at long sequences.) The point of the exercise is to make tangible the dependency Mamba has on its kernel.

Progress

Progress is local-only — sign in to sync across devices.
Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)
Sign inPlease sign in to comment.

No comments yet — be the first.