Reproduce “Implementing and Optimizing the Scaled Dot-Product Attention on Streaming Dataflow” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Gina Sohn, Nathan Zhang, Kunle Olukotun
Researchers at Stanford University developed and optimized the Scaled Dot-Product Attention for streaming dataflow accelerators, achieving constant ($O(1)$) intermediate memory complexity while maintaining full throughput. This was accomplished by carefully balancing dataflow latencies and utilizing algorithmic transformations like running maximums and dynamic rescaling.
Key insights
- A streaming execution model for SDPA can inherently reduce intermediate memory complexity from $O(N^2)$ (full matrix materialization) to $O(N)$ without explicit algorithmic modifications.
- Achieving constant ($O(1)$) intermediate memory in a streaming dataflow accelerator requires balancing the latencies of divergent data paths within the computation graph.
- Algorithmic modifications like dynamic rescaling of accumulated values based on running maximums enable online normalization and output accumulation, eliminating the need for large buffers.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.