Supporting Mixtral in gpt-fast
How gpt-fast added Mixtral MoE support using torch.compile — no custom kernels needed. Achieved fastest non-Groq inference speeds.
Key Claims
- MoE challenge: using a tensor to index into a Python list induces CUDA sync (CPU waits for GPU)
- Solution: GPU-side “gather” operation instead of Python indexing — torch.compile fuses gather + gemv into one kernel
- 98 tok/s on single A100 with int8, 280 tok/s with tensor parallelism (H100 node)
- No custom kernels — all codegen through torch.compile → Triton
- Effective bandwidth of 4.55 TB/s (above theoretical limit for a dense model — because MoE is sparse)
Takeaways
- Demonstrates the “gpt-fast ethos”: simple, native PyTorch, very fast
- torch.compile can handle MoE routing efficiently for BS=1 decode
- Doesn’t scale to larger batch sizes (latency-optimized, not throughput-optimized)
- Contrast with Cursor Warp Decode — different approach (custom CUDA kernels for Blackwell)
Connections
- Cursor Warp Decode — alternative MoE decode approach
- MoE Architectures — the architectures being optimized
- Kernel Engineering — torch.compile as alternative to manual kernels