Advancing Mamba in PyTorch
summary
Mamba 2 is a drop-in replacement for attention that keeps a fixed-size recurrent state so inference is constant-time token generation and the cache does not grow with sequence length. IBM’s Bomba shows 2.5x inference speed and ~2x lower latency versus a comparable Llama 3 setup, training within 20–30% throughput of Llama 3. Padding-free supervised fine-tuning yields about 2x padding-free throughput, context-parallel training scales experiments up to 2M tokens and work adds a PyTorch associative-scan op, DTensor shards, and MOE support.
AI-written summary · about
chapters (11)
- 0:00Introduction
- 0:21What is Mamba
- 3:52Mamba vs attention
- 7:00IBM’s Mamba in PyTorch
- 8:01Long context training
- 10:06Padding-free fine-tuning
- 12:33Core PyTorch ops
- 14:18Mixture of experts
- 15:44Granite 4.0 hybrid
- 16:27Inference with VLLM
- 17:52Summary
related talks
27mAccelerating Generative AI
10mKeynote: How to Leverage PyTorch to Scale AI Training and Inferencing
25mHow Modern PyTorch Supercharges Multimodal Training and Inference at Luma AI
12mMemory-Efficient AI: How PEFT and PyTorch Enable Accessible LLM Fine-Tuning
11mLightning Talk: Lessons from Using Pytorch 2.0 Compile in IBM's Watsonx.AI Inference
30mOur Journey With TorchTitan
16mtorchtune: Easy and Accessible Finetuning in Native PyTorch
32mTaming the Machine: Basics of ML Models Training and Inference Optimization