Implementation of Mamba in one file of PyTorch

Original Article ↗ Hacker News Discussion ↗

A compact, single-file PyTorch implementation of the new Mamba neural network architecture is being praised as a clear, educational alternative to the heavily optimized official code. Contributors explain how Mamba’s structured state-space models aim to rival Transformers by offering linear-time sequence modeling, lower VRAM use, and faster inference, while also debating trade-offs in readability, performance tricks, and tooling (from Fortran to Julia and Mojo). The exchange highlights how quickly ML jargon and model families are evolving, and shares resources for practitioners who want to understand where architectures like Mamba fit in the broader landscape of large language models.

Educational one-file implementation

  • Many appreciate the single-file PyTorch Mamba as a learning tool.
  • Compared to the highly optimized official CUDA code, this version is seen as far easier to understand.
  • Line count and file count are viewed as important here because the goal is conceptual clarity, not speed.

Alternative implementations & languages

  • A Fortran-based Mamba/SSM inference implementation is mentioned; Fortran is praised for NumPy-like array syntax, compiled speed, easy parallelism, and good linear algebra support, but criticized for clumsy non-numerical tasks.
  • Some suggest considering newer languages like Julia or Mojo, but the Fortran proponent accepts the “uphill battle.”
  • Another library is highlighted that factors out shared transformer code so models like BERT, LLaMA, and MPT fit in ~100 lines.

Performance and parallelization

  • Commenters point out that the sequential selective_scan loop can be parallelized using a couple of PyTorch calls and that all outputs can be computed with a single einsum.
  • Others note this might hurt readability, which is the main purpose of the repo; suggestions include optional parallel paths or keeping the simple version as comments.

What Mamba / SSMs are and why interesting

  • Mamba is framed as a state-space model (SSM)–based architecture aiming to compete with transformers: linear-time sequence modeling instead of quadratic attention.
  • Key idea: make some SSM parameters input-dependent (a “selection” mechanism) while keeping the latent transition linear, enabling both expressiveness and efficient parallel algorithms.
  • It is described as a continuation of earlier SSM/long-convolution work (HIPPO, S4, Hyena, etc.), adding missing ingredients like input-dependent gating.

Jargon, naming, and field velocity

  • Several posters struggle with rapid-fire jargon and new model names (Mamba, RetNet, RWKV, etc.), calling the space fad-driven and heavily branded.
  • Others argue short names are useful and that understanding comes from “overhearing” community conversations (HN, subreddits, newsletters).
  • Resources such as glossaries, references in the README, blog primers, and explainer videos are recommended.

Capabilities, efficiency, and open questions

  • Claims include:
    • Mamba trains faster and does inference much faster than similarly sized transformers, with lower VRAM usage per token.
    • It may be more compute/sample efficient and well-suited for very long context and other sequential domains.
  • Some report Mamba using about 60% of the VRAM of RetNet (in a specific test) and being cheaper per token than transformers.
  • Others note that:
    • Downstream applications are still immature compared to attention-based models.
    • It’s unclear how Mamba compares to RetNet and RWKV at scale; context-length testing and efficient implementations are still evolving.
    • One person claims current Mamba models are “less coherent than GPT-2,” another flatly disagrees and says they are better; no resolution is given.

Relation to RNNs / LSTMs

  • A recurring question is how Mamba differs from RNNs/LSTMs.
  • Explanation offered: it’s essentially a linear RNN in the latent space (for parallelizability and stability), with input-dependent transitions providing the expressive power traditionally given by nonlinearities.
  • A practical limitation mentioned: unlike transformers, packing different-length sequences efficiently for training may be harder.

Code style, tools, and research process

  • Opinions on einops/einsum are mixed: some find them elegant and efficient; others say they hurt readability.
  • There’s enthusiasm for minimal, self-contained “code sketches” as a way to speed research and lower the “Kolmogorov complexity” of experimentation.
  • Single-file implementations are praised as a counterweight to sprawling, interdependent ML codebases.