PipeThreader: Software-Defined Pipelining for Efficient DNN Execution

Yu Cheng†∗ Lei Wang† Yining Shi† Yuqing Xia⋄ Lingxiao Ma⋄ Jilong Xue⋄ Yang Wang⋄ Zhiwen Mo‡⋄ Feiyang Chen¶⋄ Fan Yang⋄ Mao Yang⋄ Zhi Yang† †School of Computer Science, Peking University ⋄Microsoft Research ‡Imperial College London ¶Shanghai Jiao Tong University

A1 主要贡献

本文提出PipeThreader，一个新的DNN编译器，旨在有效利用现代GPU中的异构专用硬件单元，如TensorCores和Tensor Memory Accelerators。通过将调度功能从硬件转移到软件，实现更高效和复杂的计算流水线，而无需大量手动努力。主要创新点包括：引入sTask-graph作为DNN计算抽象，用于表示细粒度任务和依赖；提出分层硬件抽象，捕捉专用单元的能力；以及新的调度原语（append、wait和propagate），支持软件定义的流水线。PipeThreader能为FlashAttention等已知DNN架构发现高效流水线调度，实现相当或更好的性能；同时为新兴模型如Mamba2揭示新型流水线方案，提供显著优于最先进手工优化的性能。代码开源于https://github.com/tile-ai/tilelang。

研究目标是解决现代GPU中异构硬件和操作融合带来的调度挑战：专用单元需要大张量粒度，导致并发线程减少，硬件调度器难以处理深计算流水线；操作融合加深流水线，手工优化难以泛化到新GPU类型、新模型或新张量形状。PipeThreader通过软件定义流水线，支持现代GPU架构上的高效DNN执行。

A3 背景知识/关键Observation/设计原则

硬件和软件复杂性的增长。 深度神经网络（DNN）模型的快速增长，特别是大型语言模型（LLM），推动硬件供应商开发专用异构硬件单元，如TensorCores和Tensor Memory Accelerators（TMA），以满足计算需求。同时，复杂的操作融合技术如FlashAttention被广泛用于减少内存开销并最大化数据局部性。这些趋势提升了计算密度和效率，但也引入了调度和执行的重大挑战，特别是在现代GPU上。

传统数据并行GPU执行中的利用率不足。 传统GPU编程模型如CUDA将线程块分派到流式多处理器（SM）上，将每个SM视为统一的孤立执行单元。这种抽象假设所有SM是可互换的，并隐藏其内部结构，这在早期架构如NVIDIA V100中有效。然而，现代GPU如NVIDIA H100在每个SM中集成异构组件，包括用于矩阵操作的TensorCores、用于通用计算的CUDA cores和用于内存移动的TMA。这些组件具有不同的角色和执行特性。如果不考虑这些内部差异而统一分派线程块，会导致资源使用效率低下。要充分利用此类架构，必须了解每个SM内的异构单元，并相应协调任务放置、调度和流水线。没有这种控制级别，会损失显著性能。

图1: NVIDIA H100上不同实现中代表性AI工作负载的每个单元硬件利用率，包括MatMul、FlashAttention和Mamba2的ChunkScan。每个条形显示特定工作负载实现的单个硬件单元利用率。请注意，FMA和XU单元未在MatMul中使用。
图1: NVIDIA H100上不同实现中代表性AI工作负载的每个单元硬件利用率，包括MatMul、FlashAttention和Mamba2的ChunkScan。每个条形显示特定工作负载实现的单个硬件单元利用率。请注意，FMA和XU单元未在MatMul中使用。

图1报告了H100上MatMul、FlashAttention和Mamba2的不同实现中每个单元的利用率。对于MatMul，由于没有流水线执行，内存移动成为瓶颈，TensorCore利用率仅为40%，而专家优化的cuBLAS达到97%。FlashAttention-3【[46], Flashattention-3: Fast and accurate attention with asynchrony and low-precision, 2024, arXiv preprint arXiv:2407.08608】通过手动优化将TensorCore利用率从40%提高到72%，相比基于Triton的FlashAttention-2【[26], Flashattention-2: Faster attention with better parallelism and work partitioning, 2023, arXiv preprint arXiv:2307.08691】。从FlashAttention-2到FlashAttention-3的演进花了近一年。然而，新兴的手动优化Mamba2【[28], Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, 2024, arXiv preprint arXiv:2405.21060】仍未充分利用单元，仅有15%的TensorCore利用率。因此，充分利用现代硬件对新兴DNN模型具有挑战性。

手动管理内核中流水线执行的挑战。 手动管理内核中流水线执行因巨大的硬件敏感设计空间而臭名昭著。开发者必须仔细平衡切片大小和流水线深度，同时遵守严格的芯片上资源约束。这些挑战因架构变异性而放大，包括内存层次和专用计算单元的差异。随着手动推理迅速变得不可行，自动化推理和调度对于实现性能和可移植性变得至关重要。然而，现有的编译器如TVM【[19], TVM: An automated end-to-end optimizing compiler for deep learning, 2018, 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)】和Triton【[49], Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations, 2019, Association for Computing Machinery】缺乏显式机制来表达流水线切片执行。通过抽象掉低级控制，它们限制了开发者指定执行顺序、资源分配和计算-通信重叠的能力，从而阻碍了充分利用性能潜力。实现高效流水线切片执行需要一个新编译器，能系统探索这个设计空间、生成优化调度并适应多样硬件平台。

观察和机会。 鉴于这些事实，我们观察到一个独特的机会来解决这个问题。由于新硬件单元处理大粒度数据，如张量切片，先前工作【[11], TensorIR, 讨论链接: https://discuss.tvm.apache.org/t/rfc-tensorir-a-schedulable-ir-for-tvm/7872】【[23], Cutlass: Cuda templates for linear algebra subroutines, 2024, https://github.com/NVIDIA/cutlass】【[43], Rammer: Enabling holistic deep learning compiler optimizations with rtasks, 2020, 14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20)】【[49], Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations, 2019, Association for Computing Machinery】【[53], Flextended tiles: A flexible extension of overlapped tiles for polyhedral compilation, 2019, ACM Transactions on Architecture and Code Optimization (TACO)】【[54], Ansor: Generating high-performance tensor programs for deep learning, 2020, 14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20)】【[56], ROLLER: Fast and efficient tensor compilation for deep learning, 2022, 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】已证明切片级执行可以在软件层高效调度，因为其在切片级的确定性性能。通过利用这一趋势，我们主张将流水线调度从隐式硬件行为转移到显式软件控制。这里，流水线调度不是指GPU的低级线程或warp分派，而是指软件引导的切片级操作映射到专用单元如TensorCores、CUDA cores或每个SM内的TMA。

图2: (a)现有方法中低效调度的说明；(b)利用专用执行单元之间流水线执行的优化调度。
图2: (a)现有方法中低效调度的说明；(b)利用专用执行单元之间流水线执行的优化调度。

图2(a)显示了融合MatMul-Sum的执行，其中MatMul在TensorCore上执行，Sum在CUDA core上执行。由于现有方法中的同构抽象，尽管TensorCores和CUDA cores之间存在固有并行性，但调度会序列化执行，导致低效。相反，图2(b)展示了利用专用执行单元的优化调度，实现流水线执行并充分利用异构硬件。

不幸的是，现有的DNN模型表示和现有GPU中的硬件接口都没有显式暴露切片级流水线执行所需的调度能力。

A2 方法细节

PipeThreader抽象概述。 §2中的观察激发了PipeThreader，一个整合切片基数据并行与流水线调度的DNN编译器框架。图3提供了系统概述。最先进的DNN编译器（如Triton【[49], Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations, 2019, Association for Computing Machinery】、Roller【[56], ROLLER: Fast and efficient tensor compilation for deep learning, 2022, 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】、Welder【[47], Welder: Scheduling deep learning memory access via tile-graph, 2023, 17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23)】）将硬件加速器抽象为同构执行单元（EU）的集合，用于SPMD风格的并行。这种方法忽略了现代GPU固有的硬件异构性。例如，SM内的TensorCores和CUDA cores针对不同工作负载进行优化，现有的编译器无法利用这种多样性。为了解决这些限制，PipeThreader引入两个关键抽象：专用任务（sTasks）和专用执行单元（sEUs）。从操作的数据流图（DFG）作为输入开始，PipeThreader将这些操作转换为sTasks，这些sTasks设计用于利用sEUs的异构能力，实现MPMD风格的并行。sTasks和sEUs的细节在§3.1中讨论。sTasks保留原始DFG的任务粒度数据依赖，形成sTask-graph。PipeThreader将sTask-graph到sEUs的映射组织成sTask-program，一个用于执行的结构化表示。通过sTask-program抽象，PipeThreader开辟了一个新的搜索空间。sTask-graph、sTask-program和搜索空间的细节在§3.2中讨论。

图3: PipeThreader的系统概述。
图3: PipeThreader的系统概述。

专用任务和执行单元。 sTask。 PipeThreader引入sTasks（专用任务的缩写）作为操作中的基本计算单元，在加速器设备的特定执行单元（sEU）上执行。sTask的概念自然与现代DNN加速器的异构专用处理器对齐，例如H100 GPU中的TMA、CUDA cores和TensorCores。为了最大化效率，此类加速器上的计算需要分为多个并行（异构）任务，每个类型针对专用处理器。每种类型的这些并行任务可以由一个sTask表示，从而暴露潜在的任务并行性，不仅对底层硬件的专用处理器，而且对PipeThreader编译器。

如图4所示，一个sTask处理从输入张量切片的切片数据，并在输出张量中产生一个数据切片，其计算由基于索引的张量表达式描述。sTask的形状（第8行）沿着张量表达式expr（第7行）的每个循环轴定义。此外，target_sEU（第9行）属性指定sTask可以执行的专用单元类型。相反，传统的切片基任务没有显式分类，限制了它们利用专用处理器并行性的能力。如图5所示，在特定EU上，传统的MatMul-Sum任务（FlashAttention中的融合操作）会顺序计算来自A的[2 × 2]数据切片和来自B的[2 × 2 × 2]数据切片，产生C的[2 × 2]输出切片。PipeThreader引入两种类型的sTasks：mma sTasks在TensorCores上执行矩阵乘累加，以及Sum sTasks在CUDA cores上运行。

图4: sTask和sEU的抽象。
图4: sTask和sEU的抽象。

图5: NVIDIA GPU上的sTask MatMul-Sum。
图5: NVIDIA GPU上的sTask MatMul-Sum。

这实现了流水线，因为第二个mma sTask，将A与分区B1相乘以在TensorCores上产生C1，可以与第一个Sum sTask重叠，后者在CUDA cores上将C0归约为S0。这种流水线执行允许任务并行，显著提高硬件利用率。

sEU。 现代加速器缺乏将sTask映射到特定执行单元的接口。为了解决这个问题，PipeThreader显式暴露GPU内的执行单元，并将它们抽象为分层执行阵列，捕捉并行性和支持带流水线的数据依赖执行的能力。

如图4所示，抽象设备由多个并行执行单元（EU）组成（第11行），每个包含几个异构专用执行单元（sEU）（第12行）。这些sEU作为PipeThreader有效调度带流水线的数据依赖任务的硬件基础。例如，在现代H100 GPU上，流式多处理器（SM）是一个EU，包括用于load sTasks的Tensor Memory Accelerators（TMA）和用于mma sTasks的TensorCores。一个sEU使用Execute接口（第15行）执行给定的sTask。is_async属性（第14行）指定sEU是同步操作（如CUDA cores）还是异步操作（如TMA）。异步sEU可以与异步或同步sEU并发执行。

空间分区以实现数据并行。 PipeThreader扩展支持空间和归约分区，解锁流水线执行的新机会。这种灵活性允许PipeThreader基于不同分区策略生成多样sTask-graph，实现更通用和高效的执行规划。

sProgram。 给定sTask-graph，PipeThreader以sTask-program（sProgram）的形式将其映射到硬件的sEUs。sProgram是一个二维数组sProg[sEU][order]，每个条目指定sTask到特定sEU的分配及其执行顺序。这种结构化表示促进了任务的有效调度和执行。为了维护依赖sTasks的正确执行顺序，PipeThreader引入barrier-sTasks，通过引用程序中标识的sTasks列表来同步执行。barrier-sTask等待所有引用的sTasks完成后再继续。

搜索空间。 PipeThreader的搜索空间结构化为sPrograms的集合，每个是二维数组sProg[sEU][order]，定义每个sTask-graph中sTask的切片大小和执行顺序（带同步障碍）。可能有许多sTask顺序和切片大小的组合。PipeThreader的搜索空间包括所有有效的sPrograms，其中操作可以在尊重数据依赖的情况下执行。例如，FlashAttention的搜索空间包括37,440个有效sPrograms。对于复杂融合操作，sTask调度通常对搜索空间贡献更大。随着sTask类型增多，可能有许多不同sTask执行顺序的有效sPrograms。例如，在FlashAttention中，有36个sTask大小（即切片）配置，但每个切片大小有1,040个sTask顺序配置。

从sTask-Graph到sProgram。 为了执行DNN，PipeThreader将输入DFG转换为针对现代异构硬件的专用表示。这个过程涉及两个关键步骤：构建捕捉计算和依赖的sTask-graph，以及将此图映射到专用执行单元（sEUs）作为sTask-program（sProgram），协调高效执行。

sTask-Graph。 如图3所示，输入DFG中的操作通过sTask-partition转换为sTasks，形成sTask-graph。此图保留原始DFG的计算和数据依赖，节点表示sTasks，边捕捉任务级别的细粒度依赖。sTask-partition过程通过配置sTasks的TileShape（即Map）来分区每个操作，指定其可分维度和大小。传统编译器主要关注空间分区以实现数据并行。PipeThreader扩展支持空间和归约分区，解锁流水线执行的新机会。这种灵活性允许PipeThreader基于不同分区策略生成多样sTask-graph，实现更通用和高效的执行规划。

运行示例。 Mamba2。 Mamba是一个流行的DNN模型，使用线性注意力机制通过块扫描处理序列。其线性注意力包括多个模块；这里，我们说明其关键ChunkScan操作。

前端。 对于ChunkScan函数，PipeThreader以图6(a)所示的简单IR作为输入。它将cb乘以dA和dt的指数乘积（第6行），然后将cb和x的矩阵乘法结果累加到acc_o（第8行）。PipeThreader将load_cb（第3行）、load_dA（第4行）、load_dt（第5行）、exp（第6行）、load_x（第7行）和mma（第8行）视为单独的sTasks。

sTask-graph。 PipeThreader基于它们的依赖构建相应的sTask-graph。sTask-graph可以沿着空间和归约维度分区。空间分区（即批大小）将图拆分为较小的子图，分发到EU上以实现切片基数据并行。归约分区（即序列长度）在EU内创建更细粒度的sTasks，暴露流水线执行的机会。给定acc_o（图6(a)第8行）的大小(M, N)和X的大小(K, N)，PipeThreader像往常一样分区空间维度(M, N)，为每个sTask分配(m, n)的切片。它还将归约维度(K)分区为loop_range迭代，允许计算在迭代间重叠。图6(b)显示了从归约维度分区派生的sTask-graph。

图6: Mamba2-ChunkScan的运行示例。(a)显示用户面向的前端。(b)呈现从(a)构建的sTask-graph，颜色表示不同迭代。(c)说明从(b)中sTask-graph派生的各种sPrograms，(d)比较它们的评估以识别最有效的sProgram。
图6: Mamba2-ChunkScan的运行示例。(a)显示用户面向的前端。(b)呈现从(a)构建的sTask-graph，颜色表示不同迭代。(c)说明从(b)中sTask-graph派生的各种sPrograms，(d)比较它们的评估以识别最有效的sProgram。

sPrograms。 给定sTask-graph，PipeThreader可以有不同的选择将sTasks映射到sEUs，形成多个sPrograms。图6(c)说明了从sTask-graph派生的三个sPrograms。在sProg-A中，load_x在其他load sTasks之前调度，而sProg-B和sProg-C以相反顺序调度这些操作。与sProg-B相比，sProg-C采用更大的切片大小。

评估。 图6(d)说明了三个sPrograms的评估。表格显示每个时间步运行的sTasks，以及相应的芯片上内存使用。这里，我们假设芯片上内存容量为1 KiB。在sProg-A中，load_x较早调度，但由于exp sTask依赖load_cb_dA_dt的完成，这导致exp调度延迟，与sProg-B相比整体效率较低。虽然sProg-C采用更大的sTasks，但这增加了其工作区的芯片上内存使用。我们观察到在时间步t4，工作区超过可用芯片上内存容量，使此sProgram无效。因此，选择sProg-B作为最终调度策略。

PipeThreader的任务分区遵循切片原则，但将新的归约切片提升为关键优化策略。传统上，切片优先考虑空间分区以实现数据重用，而归约切片通常被视为次要。PipeThreader积极利用归约切片来启用流水线，提高执行效率。这确保了高效流水线执行，同时保持数据重用。通过使归约切片成为一流优化，PipeThreader在传统切片策略之外解锁性能改进，特别是在流水线密集型工作负载中。此外，传统切片策略倾向于使用更大的切片大小，这可能与增加的流水线并行性冲突，因为两者都需要芯片上内存。PipeThreader进一步权衡切片和流水线。

FlashAttention-3。 FlashAttention是原始全注意力机制的高效实现，其中DNN操作图的输入节点是三个张量：Q、K和V。首先执行矩阵乘法MatMulQK以计算acc_s = QKT。接下来，acc_s通过Softmax操作产生P。最后，P和V用作第二个矩阵乘法MatMulPV的输入，以计算输出O。

图7: FlashAttention sTask-graph的伪代码。
图7: FlashAttention sTask-graph的伪代码。

在FlashAttention中，这些三个操作融合到一个内核中。PipeThreader注释模式以派生依赖sTasks，形成sTask-graph（图7）。分区后，PipeThreader将load_k（第4行）和load_v（第7行）分配到TMA，mma_qk（第5行）和mma_pv（第9行）到TensorCores，softmax（第6行）和rescale（第8行）到CUDA cores。PipeThreader应用其两层调度策略来探索空间，并生成sProgram，其中sTasks在各自sEUs上并行执行，利用sEUs的异步性和异构性。值得注意的是，我们的调度空间包括最新FlashAttention-3【[46], Flashattention-3: Fast and accurate attention with asynchrony and low-precision, 2024, arXiv preprint arXiv:2407.08608】的流水线计划。

PipeThreader调度概述。 sProgram抽象开辟了一个大型优化空间。PipeThreader旨在在此空间中生成高质量sPrograms。为此，PipeThreader将调度机制与其策略分离。在机制侧，它提供两个能力：(1) 调度接口，用于策略生成sProgram。(2) 分析器，提供调度策略请求的分析信息。在策略侧，PipeThreader提供一个两层策略，平衡切片和流水线并行。这个简单策略已经能显著优于最先进的方法。我们相信这个机制为未来研究更复杂的策略奠定基础，进一步利用sPrograms暴露的优化空间。

调度接口。 PipeThreader提供三个接口来在新空间中生成高质量sProgram，如图8所示。Append接口将特定sTask分配到EU内的特定sEU。Wait接口允许sTask s等待list中的sTasks完成，这隐式地在sTask s之前附加一个barrier-sTask。上述接口允许显式控制sTasks在sEUs上的放置和执行顺序（即sProgram），以探索并行化空间。

图8: 调度接口。
图8: 调度接口。

PipeThreader提供Propagate接口，通过自动推断sTask-graph中每个sTask的TileShape来探索sTask分区空间。从最终sTask的输出切片形状开始，Propagate通过图向后执行形状推断链，根据其张量表达式和输出切片形状确定每个sTask的依赖输入区域。例如，如果Softmax sTask需要[4 × 128]输出切片形状，Propagate推断其输入切片形状也必须是[4 × 128]。将此视为前一个mma sTask的输出切片形状，输入切片被推断为[4 × k]和[k × 128]，其中k是归约大小。

调度策略。 我们的调度策略受图3中描述的两层硬件抽象启发，其中同构EU启用SPMD风格的并行，每个EU内的异构sEU支持MPMD风格的并行。图9概述了PipeThreader中使用的两级调度算法。在EU间级别，策略通过将模型分区为sTask-subgraphs并均匀分发到EU上（第1-10行）来最小化延迟。在EU内级别，策略优化每个sTask-subgraph在给定EU上的执行成本，通过构建高效流水线计划（第11-32行）。EU间调度由每个分区的执行成本估计告知，由EU内调度提供。

最初，策略基于操作的计算阶段将每个操作表示为一个或多个sTasks（例如，MatMul拆分为load和mma sTasks）。在EU间传递中，调度器枚举函数GetsTaskPartitions（第2行）中输出sTask的不同分区。对于每个sTask分区，调度器使用Propagate（第4行）派生图中其他sTask分区，并将sTasks均匀分配到EU上，利用它们的等价计算能力。这种SPMD风格的方法显著降低了EU间并行计划的复杂性。策略调用EU内传递来优化EU内分配sTasks的执行（即sTask-subgraph）。在EU内传递期间，策略采用贪婪方法将sTasks调度到sEUs，迭代执行以下步骤直到所有分配到EU的sTasks被调度：1) 在get_complete_sTask（第15行）中选择endtime早于当前时间cur_time的sTask t；2) 识别其前驱已调度的就绪sTasks集合（第17-20行），并使用get_high_priority（第22行）出队具有最高优先级的sTask u；3) 使用Append()将选定的sTask附加到sEU（第23行），并通过调用Wait()确保sTask级依赖（第25行）。我们还调用Wait(u, t)处理u的调度必须等待t完成以释放内存的情况（第26-27行）。为了提高流水线效率，我们优先调度具有最小依赖于已调度任务的异步sTasks，以及高潜力解锁下游sTasks的sTasks。如图6(c)所示，早调度load_x会延迟exp的执行，而早调度load_y允许exp更快进行。因此，我们的算法为load_y分配更高优先级，自然偏好构建sProg-B和sProg-C而非sProg-A。

增加sTask重叠（即流水线并行）需要额外的芯片上内存（如GPU上的共享内存和寄存器）来缓冲阶段间的中间结果。然而，这种需求可能与使用更大切片大小冲突。我们的方法通过联合搜索策略平衡这些竞争需求，由分析反馈指导。为了确保内存可行性，我们调用check_valid基于当前sProgram和分析器验证选定的sTask是否符合目标sEU的内存约束（第24行）。超过限制的候选被跳过。例如，在图6中，sProg-C采用更大的切片大小，超过可用芯片上内存。check_valid步骤检测此违规并避免生成此类无效调度。

分析器。 PipeThreader引入分析器来指导搜索空间中高效sProgram的生成（图9）。分析器提供单个sTasks的信息以生成sProgram的有效执行时间线：(1) 单个sTasks在特定sEUs上的执行时间，(2) sTasks的资源使用，包括本地内存和寄存器消耗，(3) sProgram的整体执行时间。分析器通过利用现有编译器如TVM的代码生成后端自动处理新张量表达式，并测量孤立sTask的设备代码的执行时间和资源使用。在调度期间，PipeThreader使用这些分析结果估计何时启动任务以维持流水线效率并最小化空闲时间。在完成sTask调度后，分析器还测量整个生成调度的性能，提供真实延迟。这个分析数据告知调度策略并指导高效调度计划的生成。

图9: 调度算法。
图9: 调度算法。

PipeThreader实现概述。 PipeThreader使用8.5k行C++和Python代码实现，基于开源DNN编译器：TVM【[19], TVM: An automated end-to-end optimizing compiler for deep learning, 2018, 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)】和Ladder【[51], Ladder: Enabling efficient Low-Precision deep learning computing through hardware-aware tensor transformation, 2024, 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24)】。图10总结了PipeThreader的整体工作流。前端产生sTask-graph，然后由sTask-aware编译器（调度器）处理以生成sProgram。最后，映射优化器为sProgram生成设备代码。

图10: PipeThreader的实现。
图10: PipeThreader的实现。

前端。 PipeThreader前端包括sTask-IR，用于表达sTask级DNN计算，以及sTask-converter将DNN模型转换为sTask-graph。

sTask IR。 sTask中间表示（IR）为程序员和编译器提供灵活方式表达现有编译器IR（如表达式导向IR）无法轻易捕捉的sTask级计算。图6(a)和图7中的伪代码可以视为sTask IR的简化形式。伪代码说明了如何将复杂深度学习内核建模为数据流模式，包括内存操作（如在DRAM和SRAM之间移动sTasks）和sTasks上的计算序列。

带sEU的sTask转换器。 PipeThreader前端还可以将由sTask IR和ONNX图表达的DNN模型转换为sTask-graph。在此过程中，我们利用Ladder【[51], Ladder: Enabling efficient Low-Precision deep learning computing through hardware-aware tensor transformation, 2024, 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24)】，一个最先进的DNN编译器用于操作融合。Ladder的输出是一个切片图，以TVM的TIR作为中间表示。我们基于sEU信息为每个切片基任务注释target_sEU属性，将其转换为sTasks。例如，在NVIDIA H100 GPU上，我们将流式多处理器（SM）视为EU。每个SM内的sEU包括用于矩阵乘累加mma的TensorCores、用于通用浮点计算如reduce和parallel的CUDA cores，以及用于全局和共享内存间批量内存复制操作的TMA。这些基本操作可以组合以表示大多数常见深度学习内核中的数据操作。例如，数据移动转换为parallel操作，可以表示任意元素级切片操作。此外，用户可以定义自定义函数来描述其他sTasks。

虽然程序员需要编写简单IR来生成内核（例如，图7用于FlashAttention内核），但他们不必了解任务、如何将图（或IR）分为sTasks，以及哪些sTasks可以在哪些sEUs上运行。PipeThreader可以推断此信息，例如，在图4中的sTask类中设置属性如切片形状和target_sEU。sTask转换器注释每个操作（例如，mma）在哪些类型的sEU上执行（即target_sEU属性）。然后，调度器自动确定如何将其分为sTasks（例如，切片形状）和哪些sTasks在哪些sEUs上运行（例如，sProgram中的sEU分配）。因此，PipeThreader可以帮助减少手工实现所需的大量手动努力和强大领域专长。对于FlashAttention内核，PipeThreader只需68行Python代码，而手工实现FlashAttention-3的840行CUDA内核代码相比。

NVIDIA CUDA GPU上的sTask映射。 对于TMA和TensorCore sEUs，is_async属性设置为true，因为我们可以利用cp.async.bulk和wgmma.mma_async指令。相反，CUDA cores的is_async属性设置为false，因为它们不支持任何异步指令。CUDA Cores和TensorCores的指令可以同时分派的事实由NVIDIA官方确认【[8], Overlapping CUDA Cores and Tensor Cores, 论坛链接: https://forums.developer.nvidia.com/t/overlapping-cuda-cores-and-tensor-cores/288774/2】。干扰可能发生，因为两个单元使用同一组寄存器。我们通过在寄存器上实现双缓冲来缓解潜在干扰，以重叠TensorCore和CUDA Core的执行。我们使用PTX中的mbarrier对象【[9], PTX ISA, https://docs.nvidia.com/cuda/parallel-thread-execution/】实现barrier-sTask。为了高效实现sEU中的Execute函数，我们通过布局推断决定sTask操作和数据如何映射到不同线程和物理内存。我们还使用硬件特定指令（硬件内在函数）来加速操作。PipeThreader不需要大量工程努力来支持同一供应商的不同GPU模型。当针对不同架构如A100【[21], Nvidia a100 tensor core gpu architecture, 2020, Technical report, NVIDIA Corporation】、H100【[22], Nvidia h100 tensor core gpu architecture, 2023, Technical report, NVIDIA Corporation】或B100【[24], Nvidia blackwell architecture, 2025, Technical report, NVIDIA Corporation】时，只需轻量更新硬件特定配置，包括sEU布局、内在函数和资源限制。核心编译和调度逻辑保持完全可重用。

布局推断。 sTasks在专用执行单元（sEUs）上的高效执行需要遵守特定的布局和线程绑定约束。为了解决这个问题，PipeThreader引入Layout对象，描述sTasks的数据布局和线程绑定。Layout定义映射函数和迭代器域，指定逻辑数据元素如何转换为物理内存，并可选地分配到线程。

PipeThreader完全自动执行布局推断，无需手动指定。图11a显示了布局推断的简化版本。sTask mma分配到sEU TensorCore，具有严格布局约束。基于布局约束，我们可以派生相应的布局映射函数。这里，{T (m),n}表示映射到线程m的第n个位置的数据元素。

在sTask-graph中，连接sTasks的布局必须对齐以确保兼容性。使用特定sEUs的布局要求，PipeThreader推断sTasks的Layout并在整个图中传播这些要求。布局冲突通过基于优先级的推断算法解决，其中高优先级sTasks，如mma，决定依赖sTasks的布局。例如，在图11b中，mma和sum sTasks连接。鉴于mma sTask的布局已确定，我们可以相应推断sum sTask的布局。在此示例中，张量C_sum需要复制。

图11a: 将sTask的布局与sEU要求对齐。
图11a: 将sTask的布局与sEU要求对齐。

图11b: 在sTask-subgraph上推断连接sTasks的布局。
图11b: 在sTask-subgraph上推断连接sTasks的布局。

硬件内在函数。 对于需要在sEUs上进行批量操作的sTasks，我们将sTask降低为切片级函数模板。例如，矩阵乘累加操作使用CUTLASS/CuTe模板降低，这些模板集成硬件特定TensorCore内在函数。进一步的指令级优化，如寄存器分配，委托给低级编译器如LLVM【[38], Llvm: A compilation framework for lifelong program analysis & transformation, 2004, International symposium on code generation and optimization, 2004. CGO 2004】。对于NVIDIA H100，我们应用Warp Specialization【[17], Singe: Leveraging warp specialization for high performance on gpus, 2014, Proceedings of the 19th ACM SIGPLAN symposium on Principles and practice of parallel programming】【[25], Wasp: Exploiting gpu pipeline parallelism with hardware-accelerated automatic warp specialization, 2024, 2024 IEEE International Symposium on High-Performance Computer Architecture (HPCA)】来优化执行。这种技术将线程分为生产者和消费者warp，每个warp负责不同的流水线阶段。通过允许生产者warp释放未使用寄存器供消费者warp重用，Warp Specialization改善寄存器分配和效率。基于H100中TMA单元的特性，我们将load sTasks（在全局和共享内存之间复制数据）分配到生产者warp。其余sTasks，如mma和Softmax，由消费者warp处理。生产者和消费者之间的同步使用mbarrier实现的barrier-sTasks来实现，以确保正确数据依赖。

AMD ROCm GPU上的sTask映射。 我们还在AMD最新高性能GPU MI300X【[16], Amd cdna™ 3 architecture, 2023, Technical report, Advanced Micro Devices, Inc.】上实现了PipeThreader。MI300X GPU具有称为计算单元（CU）的并行执行单元，类似于NVIDIA的SM。每个CU包含多个sEU，包括用于矩阵乘累加的MatrixCores、算术逻辑单元（ALU）和异步复制单元。与CUDA GPU类似，PipeThreader在ROCm GPU上执行布局推断，将sTask数据映射到物理地址和线程。此外，我们显式利用lgkmcnt和s_waitcnt指令来管理异步障碍，实现对指令依赖和内存操作同步的精确控制。

讨论。 优于手工内核的优势。 PipeThreader在自动化流水线调度和跨架构可移植性上展示了基本优势。首先，它消除了专家级手动调优的需求。手动设计高效流水线调度容易出错、耗时，并对输入配置敏感。即使专家制作的内核如FlashAttention-3（FA3）最初缺乏对某些维度（如头大小256）的支持，说明了难度。PipeThreader自动化此过程，在硬件约束下系统探索调度空间。它实现了对FA3的最高2.18×加速，并优于vLLM的基于Triton的Mamba2达2.41×。其次，它在硬件间良好泛化。虽然手工调优内核通常紧密耦合到特定平台，特别是NVIDIA GPU，但PipeThreader可以在AMD硬件上实现显著收益。其抽象也自然映射到TPU-like架构【[34], Ten lessons from three generations shaped google’s tpuv4i : Industrial product, 2021, 2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA)】【[35], In-datacenter performance analysis of a tensor processing unit, 2017, Proceedings of the 44th Annual International Symposium on Computer Architecture, ISCA ’17】（例如，TPU核心和DMA引擎），启用高效流水线执行。最后，它降低了高性能的障碍。在Multihead Latent Attention (MLA)【[40], Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model, 2024, arXiv preprint arXiv:2405.04434】【[41], Deepseek-v3 technical report, 2024, arXiv preprint arXiv:2412.19437】上，PipeThreader仅用80行Python实现了对Triton的最高5×加速，匹配DeepSeek的500+行CUDA实现【[4], FlashMLA, https://github.com/deepseek-ai/FlashMLA】，但开发努力远少。

扩展到多GPU。 PipeThreader可以通过包括1) GPU间通信单元（如RDMA、NVLINK、IB）作为sEUs，和2) 集体通信作为sTasks，自然扩展到多GPU并与张量并行引入的集体通信互动。以此方式，PipeThreader可以重用策略在GPU内核级别搜索集体通信和计算间的有效流水线，并扩展到多GPU或多节点环境。当前结果显示，PipeThreader在常见通信模式上与最先进系统如TileLink【[55], Tilelink: Generating efficient compute-communication overlapping kernels using tile-centric primitives, 2025, arXiv preprint arXiv:2503.20313】相当。

支持新设备。 我们发现广泛使用的硬件（如NVIDIA/AMD GPU或TPU）与sEU抽象对齐，这些抽象包含所有sEUs的均匀集合。要将sTask-graphs编译到设备，编程模型仅需使用它们自己的load/store/compute指令实现每个sEU的Execute接口（图4）。设备虚拟化类似于Roller【[56], ROLLER: Fast and efficient tensor compilation for deep learning, 2022, 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】和Welder【[47], Welder: Scheduling deep learning memory access via tile-graph, 2023, 17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23)】的硬件抽象，但进一步暴露细粒度异构sEUs。

MoE FFN内核。 PipeThreader还可以支持MoE FFN内核中的分组MatMul。与批处理MatMul不同，每个组可以有不同形状。为了处理这个，PipeThreader可以将每个组分解为独立的sTask-subgraph，具有自己的输入形状，并分别应用策略，而不是共享单个调度。

相关工作。 深度学习编译器和框架。 大多数现有DNN编译器将硬件抽象为同构执行单元（EU）。Rammer【[43], Rammer: Enabling holistic deep learning compiler optimizations with rtasks, 2020, 14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20)】引入rTasks概念，设计用于EU间的并行执行，而Welder【[47], Welder: Scheduling deep learning memory access via tile-graph, 2023, 17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23)】关注通过垂直融合的整体内存优化。相反，PipeThreader引入sTasks和sEUs，显式暴露硬件异构性以启用流水线并行的优化和调度。

操作融合已在DNN编译器如TVM【[19], TVM: An automated end-to-end optimizing compiler for deep learning, 2018, 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18)】、Ansor【[54], Ansor: Generating high-performance tensor programs for deep learning, 2020, 14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20)】、XLA【[13], XLA, https://www.tensorflow.org/xla】和TensorRT【[6], NVIDIA TensorRT, https://developer.nvidia.com/tensorrt】中广泛采用，以减少内存开销，导致更深的计算阶段。编译器如Triton【[49], Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations, 2019, Association for Computing Machinery】、Welder【[47], Welder: Scheduling deep learning memory access via tile-graph, 2023, 17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23)】、Roller【[56], ROLLER: Fast and efficient tensor compilation for deep learning, 2022, 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)】、Cocktailer【[52], Cocktailer: Analyzing and optimizing dynamic control flow in deep learning, 2023, 17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23)】、TensorIR【[11], TensorIR, 讨论链接: https://discuss.tvm.apache.org/t/rfc-tensorir-a-schedulable-ir-for-tvm/7872】、ThunderKittens【[12], ThunderKittens, https://github.com/HazyResearch/ThunderKittens】、FractalTensor【[42], Uncovering nested data parallelism and data reuse in dnn computation with fractaltensor, 2024, Proceedings of the ACM SIGOPS 30th Symposium on Operating Systems Principles】和Ladder【[51], Ladder: Enabling efficient Low-Precision deep learning computing through hardware-aware tensor transformation, 2024, 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24)】基于切片抽象进行调度优化。然而，这些主要关注空间切片以提升数据局部性，实现EU间的数据并行，但忽略了sEU间流水线并行的机会。

虽然像Triton【[49], Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations, 2019, Association for Computing Machinery】和CUTLASS【[5], NVIDIA cutlass, https://github.com/NVIDIA/cutlass】这样的努力整合了流水线执行，但它们依赖于特定操作的临时规则，无法跨多样工作负载泛化。PipeThreader通过引入启用自动调度和流水线并行优化的抽象来解决这些限制。

框架如ALCOP【[33], Alcop: Automatic load-compute pipelining in deep learning compiler for aigpus, 2023, Proceedings of Machine Learning and Systems】关注数据加载和计算间的流水线，以优化内存层次利用。然而，它们未能充分利用现代计算单元的异构性，或探索具有深计算阶段工作负载如FlashAttention的流水线调度。PipeThreader通过引入细粒度抽象来解决这一差距，这些抽象启用跨异构硬件组件的全面流水线优化。

特定模式的优化。 由于现有编译器中缺乏sTask和sEU抽象，流水线并行优化通常手动针对特定模式。例如，FlashAttention【[46], Flashattention-3: Fast and accurate attention with asynchrony and low-precision, 2024, arXiv preprint arXiv:2407.08608】和CUTLASS中的Hopper MatMul【[5], NVIDIA cutlass, https://github.com/NVIDIA/cutlass】提供模式特定调度，但需要大量手动努力。此外，FlashAttention为不同输入提供单独调度，而CUTLASS【[5], NVIDIA cutlass, https://github.com/NVIDIA/cutlass】要求用户分析并选择最佳调度。相比之下，PipeThreader通过利用sTask和sEU抽象泛化流水线并行，启用跨广泛操作和配置的自动调度，而无需手动干预。

分布式深度学习框架。 Centauri【[18], Centauri: Enabling efficient scheduling for communication-computation overlap in large model training via communication partitioning, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】、PrimePar【[50], Primepar: Efficient spatial-temporal tensor partitioning for large transformer model training, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3】和TileLink【[55], Tilelink: Generating efficient compute-communication overlapping kernels using tile-centric primitives, 2025, arXiv preprint arXiv:2503.20313】通过分层调度、时序张量分区和切片基抽象分别改善通信-计算重叠。PipeThreader可以将通信和计算建模为单独sTasks，允许它表示这些工作中提出的调度策略，同时启用更广泛的调度优化。

汇总方法细节中的引用。
- [1]: A BLAS implementation on top of ROCm, https://rocmdocs.amd.com/en/latest/ROCm_Tools/rocblas.html, 引用在§6.1 Baselines中，描述为AMD ROCm上的BLAS实现。
- [2]: CUDA Basic Linear Algebra Subroutine library, https://docs.nvidia.com/cuda/cublas/index.html, 引用在§6.1 Baselines中，描述为NVIDIA CUDA上的BLAS库。
- [3]: CUDA Driver API, http://docs.nvidia.com/cuda/cuda-driver-api, 引用在§2中，描述传统GPU编程模型。
- [4]: FlashMLA, https://github.com/deepseek-ai/FlashMLA, 引用在§7中，描述为DeepSeek的MLA实现。
- [5]: NVIDIA cutlass, https://github.com/NVIDIA/cutlass, 引用多次，如§6.1和§8中，描述为CUDA线性代数子程序模板。
- [6]: NVIDIA TensorRT, https://developer.nvidia.com/tensorrt, 引用在§6.1和§8中，描述为NVIDIA特定推理库。
- [7]: ONNX Runtime, https://github.com/microsoft/onnxruntime, 引用在§6.1中，描述为DNN框架。
- [8]: Overlapping CUDA Cores and Tensor Cores, https://forums.developer.nvidia.com/t/overlapping-cuda-cores-and-tensor-cores/288774/2, 引用在§5.2中，描述NVIDIA确认的指令同时分派。
- [9]: PTX ISA, https://docs.nvidia.com/cuda/parallel-thread-execution/, 引用在§5.2中，描述为PTX指令集。
- [10]: PyTorch, https://pytorch.org/, 引用在§6.1中，描述为DNN框架。
- [11]: TensorIR, https://discuss.tvm.apache.org/t/rfc-tensorir-a-schedulable-ir-for-tvm/7872, 引用多次，如§1和§8中，描述为TVM的调度IR。
- [12]: ThunderKittens, https://github.com/HazyResearch/ThunderKittens, 引用在§8中，描述为深度学习编译器。
- [13]: XLA, https://www.tensorflow.org/xla, 引用在§8中，描述为TensorFlow编译器。
- [14]: Introducing amd cdna™ architecture, 2020, Technical report, Advanced Micro Devices, Inc., 引用在§1中，描述AMD架构。
- [15]: Introducing amd cdna™ 2 architecture, 2021, Technical report, Advanced Micro Devices, Inc., 引用在§1中。
- [16]: Amd cdna™ 3 architecture, 2023, Technical report, Advanced Micro Devices, Inc., 引用在§1和§6.1中，描述MI300X架构。
- [17]: Singe: Leveraging warp specialization for high performance on gpus, 2014, Proceedings of the 19th ACM SIGPLAN symposium on Principles and practice of parallel programming, 引用在§5.2中，描述Warp Specialization。
- [18]: Centauri: Enabling efficient scheduling for communication-computation overlap in large model training via communication partitioning, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, 引用在§8中。
- [19]: TVM: An automated end-to-end optimizing compiler for deep learning, 2018, 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18), 引用多次，如§1和§5中。
- [20]: Nvidia tesla v100 gpu architecture, 2017, Technical report, NVIDIA Corporation, 引用在§2中。
- [21]: Nvidia a100 tensor core gpu architecture, 2020, Technical report, NVIDIA Corporation, 引用在§5.2中。
- [22]: Nvidia h100 tensor core gpu architecture, 2023, Technical report, NVIDIA Corporation, 引用多次，如§1和§6.1中。
- [23]: Cutlass: Cuda templates for linear algebra subroutines, 2024, https://github.com/NVIDIA/cutlass, 引用在§1和§2中。
- [24]: Nvidia blackwell architecture, 2025, Technical report, NVIDIA Corporation, 引用在§5.2中。
- [25]: Wasp: Exploiting gpu pipeline parallelism with hardware-accelerated automatic warp specialization, 2024, 2024 IEEE International Symposium on High-Performance Computer Architecture (HPCA), 引用在§5.2中。
- [26]: Flashattention-2: Faster attention with better parallelism and work partitioning, 2023, arXiv preprint arXiv:2307.08691, 引用在§1和§2中。
- [27]: Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022, Advances in Neural Information Processing Systems, 引用在§1和§2中。
- [28]: Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, 2024, arXiv preprint arXiv:2405.21060, 引用在§1和§2中。
- [29]: The llama 3 herd of models, 2024, arXiv preprint arXiv:2407.21783, 引用在§6.1中，描述LLAMA3模型。
- [30]: bitsandbytes, 2024, 引用在§6.1中，描述低精度后端。
- [31]: Mamba: Linear-time sequence modeling with selective state spaces, 2023, arXiv preprint arXiv:2312.00752, 引用在§1中。
- [32]: Deep residual learning for image recognition, 2016, Proceedings of the IEEE conference on computer vision and pattern recognition, 引用在§6.1中，描述ResNet50。
- [33]: Alcop: Automatic load-compute pipelining in deep learning compiler for aigpus, 2023, Proceedings of Machine Learning and Systems, 引用在§8中。
- [34]: Ten lessons from three generations shaped google’s tpuv4i : Industrial product, 2021, 2021 ACM/IEEE 48th Annual International Symposium on Computer Architecture (ISCA), 引用在§7中。
- [35]: In-datacenter performance analysis of a tensor processing unit, 2017, Proceedings of the 44th Annual International Symposium on Computer Architecture, ISCA ’17, 引用在§7中。
- [36]: Miopen: An open source library for deep learning primitives, 2019, 引用在§6.1中，描述AMD库。
- [37]: Efficient memory management for large language model serving with pagedattention, 2023, Proceedings of the 29th Symposium on Operating Systems Principles, 引用在§6.1中，描述vLLM。
- [38]: Llvm: A compilation framework for lifelong program analysis & transformation, 2004, International symposium on code generation and optimization, 2004. CGO 2004, 引用在§5.2中。
- [39]: Performance analysis of gpu-based convolutional neural networks, 2016, 2016 45th International conference on parallel processing (ICPP), 引用在§6.2中，描述隐式GEMM。
- [40]: Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model, 2024, arXiv preprint arXiv:2405.04434, 引用在§7中。
- [41]: Deepseek-v3 technical report, 2024, arXiv preprint arXiv:2412.19437, 引用在§7中。
- [42]: Uncovering nested data parallelism and data reuse in dnn computation with fractaltensor, 2024, Proceedings of the ACM SIGOPS 30th Symposium on Operating Systems Principles, 引用在§8中。
- [43]: Rammer: Enabling holistic deep learning compiler optimizations with rtasks, 2020, 14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20), 引用在§1和§8中。
- [44]: Introducing chatgpt, 2022, https://openai.com/blog/chatgpt, 引用在§2中，描述LLM增长。
- [45]: Unet: Convolutional networks for biomedical image segmentation, 2015, Medical Image Computing and Computer-Assisted Intervention – MICCAI 2015, 引用在§6.1中，描述UNet。
- [46]: Flashattention-3: Fast and accurate attention with asynchrony and low-precision, 2024, arXiv preprint arXiv:2407.08608, 引用多次，如§1和§6.2中。
- [47]: Welder: Scheduling deep learning memory access via tile-graph, 2023, 17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23), 引用在§1和§3中。
- [48]: Retentive network: A successor to transformer for large language models, 2023, arXiv preprint arXiv:2307.08621, 引用在§6.1中，描述RetNet。
- [49]: Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations, 2019, Association for Computing Machinery, 引用多次，如§1和§2中。
- [50]: Primepar: Efficient spatial-temporal tensor partitioning for large transformer model training, 2024, Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, 引用在§8中。
- [51]: Ladder: Enabling efficient Low-Precision deep learning computing through hardware-aware tensor transformation, 2024, 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), 引用在§1和§5中。
- [52]: Cocktailer: Analyzing and optimizing dynamic control flow in deep learning, 2023, 17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23), 引用在§8中。
- [53]: Flextended tiles: A flexible extension of overlapped tiles for polyhedral compilation, 2019, ACM Transactions on Architecture and Code Optimization (TACO), 引用在§1和§2中。
- [54]: Ansor: Generating high-performance tensor programs for deep learning, 2020, 14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20), 引用在§1和§8中。
- [55]: Tilelink: Generating efficient compute-communication overlapping kernels using tile-centric primitives, 2025, arXiv preprint arXiv:2503.20313, 引用在§7和§8中。
- [56]: ROLLER: Fast and efficient tensor compilation for deep learning, 2022, 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), 引用在§1和§3中。

这些引用在方法细节中用于描述相关技术、硬件和先前工作，确保忠实于原文。

A4 实验环境

数据集名称、规模及用途。 评估使用六个典型DNN模型，包括LLAMA3-8B、LLAMA3-70B、Mamba2-1.3B、RetNet-65B、ResNet50和UNet。对于LLAMA3-8B、LLAMA3-70B和RetNet-65B，使用(BS, SEQ)配置为(1,1)、(32,1)和(1,4096)，覆盖在线和离线推理场景。对于Mamba，使用BS=1与序列长度1k、2k、4k、8k，以及BS=32或128与序列长度1。这些配置选择因为Mamba在长序列上的计算效率优于Transformer，并代表最常用场景。从每个模型中选择最频繁和昂贵的操作构建微基准。表1列出代表性操作、配置和缩写。

模型架构关键参数。 LLAMA3-8B和LLAMA3-70B为FP16和WFP4AFP16量化版本；Mamba2-1.3B为线性注意力模型；RetNet-65B为保留网络；ResNet50和UNet为CNN模型。微基准包括MatMul（从LLAMA3派生）、Conv2D（从ResNet-50派生）、低位MatMul（WFP4AFP16）、FlashAttention、FlashDecoding和线性注意力操作（ChunkScan和ChunkState从Mamba2派生）。

硬件配置。 在NVIDIA H100 (80GB) GPU和AMD Instinct MI300X (192GB) GPU上评估。H100使用CUDA 12.4，MI300X使用ROCm 6.1.0。两者在Ubuntu 20.04操作系统上评估。H100包括TensorCores、CUDA cores和TMA；MI300X包括MatrixCores、ALU和异步复制单元。

软件配置。 基于TVM和Ladder实现，8.5k行C++和Python代码。基线包括ONNXRuntime (v1.19.2)、Ladder、PyTorch-Inductor (v2.4.0 with Triton v3.0.0)、TensorRT (v10.0.1)、cuBLAS、rocBLAS、MIOpen、FlashAttention-3 (使用CUTLASS模板)、bitsandbytes和vLLM (v0.6.3)。所有评估从预热迭代开始，然后重复执行至少5秒。

A4 实验结果

实验内容、结果、分析结论，并附上图表的引用。

NVIDIA H100上的操作性能（图12）。 评估微基准中所有操作配置，y轴为相对于PipeThreader的归一化延迟。

图12: NVIDIA H100 GPU上的操作性能。
图12: NVIDIA H100 GPU上的操作性能。

NVIDIA H100上的端到端性能（图13，表2）。 评估八个模型的推理延迟，对于大型模型使用单个解码层作为代理。

图13: NVIDIA H100 GPU上的端到端性能。
图13: NVIDIA H100 GPU上的端到端性能。

表2: NVIDIA H100 GPU上LLAMA3-8B WFP16AFP16的延迟（毫秒）。
表2: NVIDIA H100 GPU上LLAMA3-8B WFP16AFP16的延迟（毫秒）。

调度策略评估（表3，表4）。 联合优化分区和流水线：在Mamba2-ChunkScan (BS=64, SEQ=8k)上，解耦变体 (PT-decouple) 选择更大切片 (64×128)，导致12.150ms；联合优化选择较小切片 (64×64)，实现6.981ms。分析：联合搜索平衡重用和并行。编译时间：MatMul 0.13min vs Triton 0.17min；FlashAttention 5.26min，探索大空间。

表3: 解耦和联合sTask-graph优化的延迟（毫秒）比较。
表3: 解耦和联合sTask-graph优化的延迟（毫秒）比较。

表4: H100上的编译时间（分钟）。
表4: H100上的编译时间（分钟）。

AMD MI300X上的操作性能（图14）。 子集操作：PipeThreader加速Triton 1.16×-5.42×、PyTorch最高6.21×、rocBLAS最高1.77×、MIOpen最高2.21×、FlashAttention-2最高2.82×、Ladder平均1.45×。分析：效率和可扩展性，尽管MI300X异步能力和共享内存较小。

图14: AMD MI300X GPU上的操作性能。
图14: AMD MI300X GPU上的操作性能。

AMD MI300X上的端到端性能（图15）。 对于FP16 LLAMA3，加速PyTorch-Inductor 1.48× (最高2.77×)、ONNXRuntime 6.33× (最高15.51×)、vLLM 1.02× (最高1.32×)、Ladder 1.07× (最高1.29×)；WFP4AFP16加速PyTorch-Inductor 3.97× (最高12.66×)、Ladder 1.12× (最高1.34×)。对于Mamba2，加速Ladder 32.93× (最高61.33×)、PyTorch-Inductor 1.31× (最高1.54×)。对于RetNet-65B，加速PyTorch-Inductor 1.03× (最高1.36×)、ONNXRuntime 4.75× (最高5.73×)、Ladder 1.01× (最高1.02×)。对于CNN，加速PyTorch-Inductor 2.74× (最高5.66×)、ONNXRuntime 5.84× (最高15.47×)、Ladder 2.14× (最高6.54×)。分析：PipeThreader在AMD上的收益，尽管硬件异步能力较低。

图15: AMD Instinct MI300X GPU上的端到端性能。
图15: AMD Instinct MI300X GPU上的端到端性能。

总体发现：PipeThreader为FlashAttention等发现高效调度，匹配或优于最先进；为Mamba2等新兴模型揭示新型调度，提供显著改进；抽象适应AMD等硬件，实现显著收益。

A5 结论

随着DNN模型变大和专用异构硬件单元出现，硬件调度器不再足以实现高效流水线执行。本文引入PipeThreader，一个DNN编译器，通过sTask-graph抽象和分层硬件能力（虚拟EU和专用sEU）启用软件定义流水线。借助关键调度原语，PipeThreader自动化流水线调度，在H100和AMD GPU上实现与FlashAttention等相当或优越的性能，同时泛化到新兴模型如Mamba2。我们相信PipeThreader为编译器基优化的进一步进步奠定基础，为利用演进GPU架构和DNN工作负载的高效利用铺平道路。未来工作展望包括扩展到多GPU、支持新设备和MoE FFN内核，以及更复杂的调度策略。

A6 附录

论文中无明确附录内容。