2 posts tagged with "CUDA"

Mamba's (state-spaces/mamba) selective scan is written as a CUDA kernel and doesn't run on Apple Silicon out of the box. I rewrote it in Metal Shading Language (MSL), wired it up to load HuggingFace weights directly, and got end-to-end inference working — that's mamba-metal. This post is the design + results notebook.

Mamba (Gu & Dao 2023) is a sequence model called a Selective State Space Model. It's a candidate replacement for the Transformer's attention, with linear-time compute in context length and a per-token decode cost that is independent of context length. This post walks through the architecture, what "selective" actually means, why the recurrence can still be parallelised, and finally the official CUDA kernel csrc/selective_scan/selective_scan_fwd_kernel.cuh line by line.