关于《Streaming Tensor Programs: A Streaming Abstraction for Dynamic Parallelism》的学术研究报告
一、 主要作者、机构与发表信息
本研究的主要作者包括：Gina Sohn, Genghan Zhang, Konstantin Hossfeld, Jungwoo Kim, Nathan Sobotka, Nathan Zhang, Olivia Hsu, 以及 Kunle Olukotun。这些作者主要来自斯坦福大学（Stanford University），其中Nathan Zhang同时来自SambaNova Systems公司，Olivia Hsu则同时在斯坦福大学和卡内基梅隆大学（Carnegie Mellon University）任职。该研究论文发表于2026年3月22日至26日在美国匹兹堡举行的第31届ACM国际编程语言和操作系统架构支持会议（Proceedings of the 31st ACM International Conference on Architectural Support for Programming Languages and Operating Systems， ASPLOS ‘26）。这表明该工作属于计算机体系结构与编程系统交叉领域的高水平学术成果。
二、 学术背景与研究动机
本研究的主要科学领域是计算机体系结构，具体聚焦于针对张量应用（尤其是机器学习）的专用硬件加速器及其编程抽象。随着大型语言模型（Large Language Models, LLMs）等机器学习应用的普及，对高性能、高能效计算硬件的需求日益增长。空间数据流加速器（Spatial Dataflow Accelerators, SDAs）作为一种新兴的硬件架构，因其通过空间分布的、流水线化的计算与存储单元进行数据流式执行，能够避免传统处理器（如GPU）的控制开销，并实现激进的算子融合、流水线化和细粒度并行，从而在性能和能效上展现出巨大潜力，已被证明能超越当前最先进的GPU。
然而，当前SDAs的编程抽象在支持动态行为方面存在严重局限。现代机器学习模型（如MoE架构的LLM）越来越多地展现出数据依赖的张量形状（例如，序列长度、批处理大小可变）和数据依赖的控制流（例如，Mixture of Experts中的专家路由）。这些动态行为在现有的SDA编程抽象（无论是命令式如Spatial、Revet，还是数据流式如StreamIt、SAM、Ripple）中难以高效表达和优化。这些抽象要么缺乏对显式内存层次结构的支持，要么采用同步数据流模型限制了动态性，要么将数据速率隐藏在抽象之下，使得编译器难以分析和优化。这种不匹配导致动态工作负载在SDA上要么只能静态实现，要么性能未达最优，从而限制了SDA硬件潜力的充分发挥。
因此，本研究旨在解决这一核心矛盾。其目标是提出一种新的编程抽象，能够原生、高效地表达和优化动态张量工作负载在SDA上的执行。具体而言，该研究的目标是：1）设计一个支持动态数据速率和张量维度的流式抽象；2）通过该抽象暴露性能关键指标（如片外流量、片上内存需求）；3）基于该抽象实现并评估一系列此前无法表达的优化，如动态分块、动态并行化和配置时分复用。
三、 详细工作流程与方法
本研究的工作流程可以分为四个主要阶段：新抽象的设计与定义、符号化前端与性能模型的构建、基于新抽象的优化策略实现与评估、以及对未来编译与硬件支持的讨论。
第一阶段：Streaming Tensor Programs (STEP) 抽象的设计。 研究者首先分析了现有SDA编程抽象的不足（如表1所示），并针对动态张量应用的需求，设计了一套全新的异步数据流抽象——STEP。STEP的核心设计思想是“以流为中心”，将数据表示为流（Stream），其中可以包含具有动态形状的块（Tile）。STEP程序由异步数据流块组成，这些块通过硬件FIFO连接。其创新性体现在三个关键属性上： 1. 显式内存层次结构：提供了linearOffchipLoad、bufferize、streamify等操作符，明确区分和控制数据在片外内存、片上内存和流之间的移动，使程序员或编译器能够精确管理数据局部性和复用。 2. 符号化的数据消费与生产速率：STEP为每个流定义了“形状语义”，包括静态规则维度、动态规则维度和不规则维度。这使得即使在异步数据流中，也能在编译时和运行时推理数据的组织和流动模式。 3. 数据依赖的控制流操作符：引入了reassemble、eagerMerge、partition等操作符，专门用于高效表达动态路由和合并模式，这是支持MoE等动态控制流应用的关键。
此外，STEP还包含高阶操作符（如accum, map）和形状操作符（如flatten, reshape），共同构成一个完整的编程模型。论文通过一个简化的MoE示例（图6、图7和代码清单1）详细展示了如何使用STEP操作符组合来表达包含动态路由、动态形状张量计算和结果合并的完整工作流。
第二阶段：符号化前端与性能模型的开发。 为了验证STEP的可行性和进行设计空间探索，研究者实现了一个符号化的Python前端和一个周期近似模拟器。 * 符号化前端：该前端允许用户以类似PyTorch的风格编写STEP程序，但需要指定并行化、分块、内存放置等调度决策。其核心功能是基于STEP的形状语义，自动推导出程序的关键性能指标符号表达式，包括片外内存流量和片上内存需求。例如，对于linearOffchipLoad操作符，其片外流量可以表示为输出流基数 × 输出数据类型大小。当存在动态维度时，这些表达式会包含符号变量，通过代入具体运行时的形状值即可得到具体指标。这为快速分析和优化提供了强大工具。 * 周期近似模拟器：该模拟器基于数据流抽象机框架实现，用于评估性能（周期数）。它集成了Ramulator 2.0来模拟高带宽内存（HBM）的时序行为，并使用屋顶线模型（Roofline Model）来估算计算操作的延迟。模拟器允许为每个高阶操作符（如矩阵乘）指定计算带宽，从而更真实地反映硬件特性。 * 验证：为了确保抽象和模拟器的准确性，研究者将STEP程序映射到一个用Bluespec SystemVerilog实现的周期精确硬件描述语言模型上，并运行了一个SWiGLU层工作负载。通过对比不同分块大小下的周期数和片外流量，结果表明STEP模拟器的结果与HDL模拟器高度吻合（皮尔逊相关系数0.99），且符号化前端推导出的片外流量与模拟性能强相关，验证了STEP所捕获指标的有效性和模拟器的可靠性。
第三阶段：优化策略的实现与系统性评估。 利用STEP表达动态性的能力，研究者实现了三种此前在SDA抽象中难以或无法表达的优化，并进行了详尽的评估。 1. 动态分块：在MoE计算中，每个专家激活的令牌数量是动态的。静态分块需要填充到固定大小，导致内存浪费或频繁重载。STEP通过promote和accum操作符，允许分块大小在运行时根据实际令牌数动态决定。评估在Mixtral-8x7b和Qwen3-30b-a3b模型上进行，批处理大小分别为64和1024。结果：动态分块打破了静态分块的帕累托最优边界。例如，在批处理大小为64时，对于Qwen3-30b-a3b，动态分块在性能上与最佳静态分点（tile=64）持平时，减少了5.05倍的片上内存使用；在保持与次佳静态分点（tile=8）相同内存时，实现了1.69倍加速。帕累托改进距离达到1.33倍至2.11倍。 2. 配置时分复用：针对MoE中只有部分专家被激活的特点，该优化允许一个硬件配置（计算单元和内存）被多个专家在时间上复用。在STEP中，这通过在一组专家子图周围插入eagerMerge和randomOffchipLoad等控制流操作符来实现。结果：在Qwen3-30b-a3b的MoE层上，无论是静态分块还是动态分块，配置时分复用都能在性能开销小于5%的情况下，将计算利用率提升约2.5-2.6倍，从而释放出大量硬件资源用于处理更多并发请求或更大模型。 3. 动态并行化：在注意力机制等场景中，请求的KV缓存长度各不相同，导致负载不均。STEP的partition和eagerMerge操作符支持将工作动态分发到下游并行流水线，一旦某个流水线空闲就立即分配新任务，而不是静态分配固定数量的请求。结果：在注意力层评估中，与静态交错并行化相比，动态并行化在KV缓存长度变化较大时，能实现1.47倍至1.57倍的加速；与静态粗粒度并行化相比，在批处理大小较小时（batch=16），加速比高达2.72倍，有效缓解了负载不平衡问题。
此外，研究者还将这些优化整合到完整的Qwen3-30b-a3b和Mixtral-8x7b模型中进行端到端评估。结果显示，结合了动态分块和动态并行化的优化实现，相比内存匹配的静态实现，分别获得了1.27倍和1.15倍的端到端加速。对于Qwen3-30b-a3b，配置时分复用还额外减少了69%的片上内存和54%的计算资源占用。
第四阶段：对未来工作的讨论。 论文最后探讨了将STEP作为中间表示从高级框架（如PyTorch）编译而来的可能性，以及如何在SDA硬件上支持STEP的动态特性（如通过硬件状态机处理停止令牌、在片上网络或计算单元内实现动态路由、通过硬件管理的映射缓存支持非连续内存分配和虚拟化）。
四、 主要研究结果
本研究在各个阶段均获得了明确且相互支撑的结果： 1. 抽象设计可行性：成功设计了STEP抽象，其核心操作符（如动态路由、显式内存操作）能够简洁地表达复杂的动态张量工作负载（如MoE），并通过形状语义实现了对程序行为的可推理性。 2. 工具链有效性：开发的符号化前端能够准确推导性能指标的符号表达式，周期近似模拟器能够可靠地预测硬件性能，两者与周期精确HDL仿真的高度一致性（图8）强有力地证明了STEP抽象及其工具链的实用性和准确性。 3. 优化性能增益：通过STEP实现的三种优化均带来了显著的性能提升或资源节约。 * 动态分块的实验结果（图9，图10）直接证明了其能够突破静态分块的性能-内存帕累托边界，为解决内存受限应用中的分块权衡问题提供了新方案。 * 配置时分复用的结果（图12，图13）表明，通过共享硬件配置，可以大幅提高资源利用率，这对于专家数量多但激活稀疏的现代MoE模型至关重要。 * 动态并行化的结果（图14，图15）显示，它能有效应对工作负载大小的动态变化，改善负载均衡，从而提升整体吞吐量。 4. 端到端价值：完整的模型评估（图17）将上述优化置于真实场景中，证明了它们的组合能够为实际LLM推理带来可观的加速和资源节省，验证了STEP在整体系统层面的价值。
这些结果之间存在清晰的逻辑链条：STEP抽象的设计（结果1）为性能分析（结果2）和新型优化（结果3）的实现提供了基础；而各类优化的成功验证（结果3）又反过来证明了STEP抽象在解决动态工作负载挑战方面的有效性；最终，端到端的积极结果（结果4）综合证明了该研究对实际应用的贡献。
五、 研究结论与价值
本研究的核心结论是：Streaming Tensor Programs (STEP) 是一种有效的流式抽象，它首次为空间数据流加速器提供了对动态张量工作负载的一流支持，并解锁了此前无法实现的优化，从而显著提升了动态机器学习应用在SDA上的性能和资源效率。
其科学价值在于： 1. 提出了一个创新的编程模型：STEP融合了显式内存层次、符号形状语义和动态控制流原语，填补了SDA编程抽象在动态性支持方面的空白，为后续相关研究奠定了新的基础。 2. 提供了系统的分析工具和方法论：符号化前端和性能模型构成了一个完整的设计空间探索框架，使研究人员和工程师能够定量分析并优化动态数据流程序。 3. 揭示了新的优化机会：研究明确展示了动态分块、配置时分复用和动态并行化等优化策略的潜力和实现路径，为未来编译器和硬件协同设计指明了方向。
其应用价值在于： 1. 提升硬件利用率：使SDA能够更高效地运行日益流行的动态机器学习模型（如MoE LLM），充分发挥其硬件潜力。 2. 降低开发门槛：通过提供高级的、支持动态性的抽象，简化了为SDA编写高性能动态应用程序的复杂度。 3. 促进软硬件协同设计：STEP抽象为硬件设计者提供了明确的需求指引（如对动态路由、非连续存储的支持），有助于设计下一代更适应动态工作负载的SDA。
六、 研究亮点
七、 其他有价值内容
论文中关于可移植性的讨论也颇具价值。作者指出，STEP作为一个抽象，并不绑定于特定硬件实现。其符号化前端和模拟器可以通过定制成本函数、调整操作符延迟和间隔、替换内存模型等方式，适配不同特性的SDA硬件（如不同的计算瓦片大小、内存层次结构）。这种设计增强了STEP作为跨平台编程模型或中间表示的潜力。此外，与CUDA Graphs、动态任务并行框架（如TaskStream）等相关工作的对比，清晰地定位了STEP在技术谱系中的独特地位——专注于张量粒度、异步数据流、且与显式内存层次紧密集成的动态并行抽象。