Stax (3): Scheduling and CPU Codegen
Buffer dependencies, the greedy fusion algorithm, buffer reuse, affine dependences and tiling, plus the full CPU codegen path from vectorization to wrappers.
Buffer dependencies, the greedy fusion algorithm, buffer reuse, affine dependences and tiling, plus the full CPU codegen path from vectorization to wrappers.
Reverse-mode autodiff over the operator-level graph, the minimum-cut model for save-versus-recompute, and lowering to a loop-level intermediate representation.
Just-in-time compilation as program specialization: guards, specialization keys, tracing-time proxy capture, data-dependent control flow, and graph rewrites.
A systematic tour of Stax, TensorPlay’s compiler backend: specialization, graph capture, transformations, autodiff, loop-level IR, scheduling, and codegen.
Introduction Since Vaswani et al. proposed the Transformer architecture in 2017, the attention mechanism based on Softmax Attention has become the core component of sequence modeling, supporting the rapid development of Large Language Models (LLMs). However, Softmax Attention inherently suffers from quadratic computational complexity: as the sequence length $L$ increases, the computational and memory overhead of $O(L^2d)$ grows quadratically, which has become a core bottleneck for long-sequence modeling and efficient inference and training of LLMs. ...