Stax 编译栈(三):调度、融合与 CPU 生成
缓冲区的依赖与生命周期、贪心融合算法、缓冲区复用、仿射依赖与分块合法性,以及向量化、缓存分块、内核形态、并行化与包装函数的 CPU 代码生成全流程。
缓冲区的依赖与生命周期、贪心融合算法、缓冲区复用、仿射依赖与分块合法性,以及向量化、缓存分块、内核形态、并行化与包装函数的 CPU 代码生成全流程。
算子级图上的反向模式自动微分、保存与重算的最小割模型、联合图,以及降级到循环级中间表示的规则、张量即函数的语义、规约的代数性质与实体化决策。
从「即时编译是一种程序特化」出发,讲清 Stax 的守卫与特化键、执行期代理追踪、数据相关控制流的处理,以及改写系统下的代数恒等式、常量折叠、公共子表达式消除与算子级分解。
TensorPlay 默认编译后端 Stax 的系统讲解:六个部分依次覆盖特化与缓存、图捕获、图变换、自动微分与重物化、循环级中间表示、调度与融合、CPU 与 GPU 代码生成、动态形状,最后给出测量的统计模型与习题。
引言 自2017年Vaswani等人提出Transformer架构以来,基于Softmax Attention的注意力机制成为序列建模的核心组件,支撑了大语言模型的快速发展。但Softmax Attention存在二次计算复杂度的固有问题,随着序列长度$L$的增加,$O(L^2d)$的计算量和内存开销呈平方级增长,成为长序列建模和大模型高效推理、训练的核心瓶颈。 ...