Dao AI Lab

Affiliations. Address. Contacts. Motto. Etc.

prof_pic.jpg

555 your office number

123 your address street

Your City, State 12345

Write your biography here. Tell the world about yourself. Link to your favorite subreddit. You can put a picture in, too. The code is already in, just name your picture prof_pic.jpg and put it in the img/ folder.

Put your address / P.O. box / other info right below your picture. You can also disable any of these elements by editing profile property of the YAML header of your _pages/about.md. Edit _bibliography/papers.bib and Jekyll will render your publications page automatically.

Link to your social media connections, too. This theme is set up to use Font Awesome icons and Academicons, like the ones below. Add your Facebook, Twitter, LinkedIn, Google Scholar, or just disable all of them.

news

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

latest posts

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

selected publications

  1. FlashAttention-4: Algorithm and Kernel Pipelining Co-design for Asymmetric Hardware Scaling
    Ted Zadouri*, Jay Shah*, Markus Hohnerbach*, Timmy Liu, Vijay Thakkar, and Tri Dao
    In Machine Learning and Systems (MLSys), 2026
    Best Paper Honorable Mention
  2. Marconi: Prefix Caching for the Era of Hybrid LLMs
    Rui Pan, Zhuang Wang, Zhen Jia, Can Karakus, Luca Zancato, Tri Dao, Ravi Netravali, and Yida Wang
    In Machine Learning and Systems (MLSys), 2025
    Outstanding Paper Honorable Mention arXiv Code
  3. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
    Jay Shah*, Ganesh Bikshandi*, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao
    In Advances in Neural Information Processing Systems (NeurIPS), 2024
    arXiv Code
  4. Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality
    Tri Dao* and Albert Gu*
    In International Conference on Machine Learning (ICML), 2024
    arXiv Code
  5. Mamba: Linear-Time Sequence Modeling with Selective State Spaces
    Albert Gu* and Tri Dao*
    Conference on Language Modeling (COLM), 2024
    Outstanding Paper arXiv Code
  6. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
    Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré
    In Advances in Neural Information Processing Systems, 2022
    Awards arXiv Code
  7. Monarch: Expressive Structured Matrices for Efficient and Accurate Training
    Tri Dao, Beidi Chen, Nimit Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré
    In International Conference on Machine Learning (ICML), 2022
    Outstanding Paper runner-up arXiv Code