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.
Technical Blog
TensorPlay Stax (2): Autodiff and Lowering
Reverse-mode autodiff over the operator-level graph, the minimum-cut model for save-versus-recompute, and lowering to a loop-level intermediate representation.
Stax (1): Specialization & Graph Capture
Just-in-time compilation as program specialization: guards, specialization keys, tracing-time proxy capture, data-dependent control flow, and graph rewrites.
Stax: Graph Capture to Code Generation
A systematic tour of Stax, TensorPlay’s compiler backend: specialization, graph capture, transformations, autodiff, loop-level IR, scheduling, and codegen.
Talk Linear Attention
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. ...