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

Hardware Advantage — GEMM, Tensor Cores, and Why Scans Are Hard

~12 min · hardware, gemm, tensor-cores, kernels

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

GPUs are matmul machines

Modern GPUs — H100, B200, and beyond — are designed first and foremost as dense matrix-multiply engines. Tensor cores saturate efficiently with large matmul operations. The entire memory hierarchy (HBM → L2 → SRAM) is tuned for the access patterns that matmul produces.

Standard attention is essentially three matmuls: Q·Kᵀ, the softmax (a non-matmul interlude), and weights·V. FlashAttention-3 fuses these elegantly into a tile-based pattern that maps perfectly onto Hopper's hardware. The result is ~85% of H100 theoretical peak — close to the upper bound of what's possible on the chip.

Parallel scans are structurally different

SSMs require parallel scan algorithms for training-time parallelism. A scan is structurally different from a matmul: it has dependencies between sequential elements, and parallelizing it requires clever algorithmic tricks (the Blelloch scan, the Brent-Kung adder pattern). These tricks work, but the result doesn't map onto tensor cores as cleanly as matmul does. Mapping a scan to GPU SRAM tile-by-tile is harder, and the constants are worse.

The TFLA kernel (NeurIPS 2025) is a major breakthrough here — it finally matched FlashAttention-3 speed for linear RNNs on H100. But TFLA took years of dedicated kernel engineering to write, and there's nothing comparable to FlashAttention's level of cross-architecture maturity yet. Every new scan-based architecture has to fight for similar kernels.

Hardware-software co-design is starting

Custom ASICs are beginning to target SSM-friendly designs. Google Ironwood (TPU v5p generation) and AWS Trainium3 have features that better support parallel scan and recurrent compute. Cerebras WSE and Groq LPU have always had different sweet spots than NVIDIA GPUs and are more SSM-friendly in places.

But the GPU ecosystem's head start is measured in years and billions of dollars of optimization. Hardware-software co-design for alternatives is in early stages — promising but not yet decisive. NVIDIA's commitment to hybrid SSM-attention via Nemotron is itself the strongest single signal that the GPU side will eventually have first-class kernel support for hybrid stacks.

External links

Exercise

Profile a Mamba-130M and a Transformer of similar parameter count on the same single GPU using torch.profiler or nsys. Look at the kernel trace: what fraction of time is spent in compute vs. memory access? You should see the Mamba kernel spending more relative time on memory-access patterns than the FlashAttention call does. That's the scan-vs-matmul gap, made visible in profiler output.

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.