Dao AI Lab

We are a research group at Princeton University led by professor Tri Dao.

Machine learning and systems, with a focus on efficient training and inference:

Oct 13, 2026	FLAT: fast, open-source TPU training kernels for linear attention and Mamba-2. TPU kernels face different bottlenecks than GPU kernels.
Oct 12, 2026	Building on our earlier blog post, Ze-Wei and NVIDIA collaborators released ReplaySSM. This makes inference optimizations like speculative decoding and cache quantization more efficient for hybrid models.
Sep 26, 2026	Rigel: a hybrid Mamba–attention MoE approaching Llama-3.2 quality with a fraction of the training compute. Trained on a mix of A100s, H100s, and TPUs, thanks to Mayank.
Sep 24, 2026	4 papers accepted to NeurIPS 2026: M²RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling GPU Hierarchy Meets Structured Matrices: Fast Algorithms for State-Space Models CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs NyoomFloat12: Accelerating LLM Inference via Lossless 12-bit Weight Compression

Jun 15, 2026	ReplaySSM: Cache SSM Inputs, Not State
Apr 22, 2026	SonicMoE: A Hardware-Efficient and Software-Extensible Blueprint for Fine-Grained MoEs
Mar 30, 2026	Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon
Mar 16, 2026	Mamba-3 Part 2 - Methodological Deep Dive
Mar 16, 2026	Mamba-3 Part 1