Author Login Chief Editor Login Reviewer Login Editor Login Remote Office

Computer Engineering

   

A Pipelined Inference Method for Transformers via Batch Partitioning

  

  • Published:2026-08-19

基于批次分割的Transformer流水推理方法

Abstract: With the widespread adoption of Transformer models in computer vision, natural language processing, and related domains, inference workloads have grown rapidly, placing increasingly stringent demands on inference throughput. Existing inference frameworks primarily improve local execution efficiency through operator fusion, kernel generation, and graph-level optimization. However, operators in Transformer models exhibit distinct resource preferences with respect to hardware resources such as Tensor Cores, CUDA Cores, and memory bandwidth. Constrained by data dependencies, these operators are typically executed in topological order, making it difficult to achieve sufficient overlap among heterogeneous hardware resources and thereby limiting further improvements in inference throughput. To address this issue, this paper proposes a batch-partitioned pipelined inference method for Transformer models. Transformer computation consists of diverse operators, including matrix multiplication, reduction, element-wise computation, and tensor reshaping. Based on operator resource-occupancy characteristics and offline performance profiling, we classify Transformer operators into three categories: Tensor Core-intensive, CUDA Core-intensive, and memory-bandwidth-intensive operators. To overcome the limited overlap among operators with different resource preferences under conventional execution, the original inference batch is partitioned into three sub-batches with adjustable sizes and scheduled concurrently while preserving the computational-graph dependencies within each sub-batch. By allowing different sub-batches to progress through different computation stages, operators with complementary resource preferences can be overlapped in a pipelined manner. This reduces contention for the same type of hardware resource, mitigates pipeline bubbles, and improves overall hardware utilization. Furthermore, we formulate the With the widespread adoption of Transformer models in computer vision, natural language processing, and related domains, inference workloads have grown rapidly, placing increasingly stringent demands on inference throughput. Existing inference frameworks primarily improve local execution efficiency through operator fusion, kernel generation, and graph-level optimization. However, operators in Transformer models exhibit distinct resource preferences with respect to hardware resources such as Tensor Cores, CUDA Cores, and memory bandwidth. Constrained by data dependencies, these operators are typically executed in topological order, making it difficult to achieve sufficient overlap among heterogeneous hardware resources and thereby limiting further improvements in inference throughput. To address this issue, this paper proposes a batch-partitioned pipelined inference method for Transformer models. Transformer computation consists of diverse operators, including matrix multiplication, reduction, element-wise computation, and tensor reshaping. Based on operator resource-occupancy characteristics and offline performance profiling, we classify Transformer operators into three categories: Tensor Core-intensive, CUDA Core-intensive, and memory-bandwidth-intensive operators. To overcome the limited overlap among operators with different resource preferences under conventional execution, the original inference batch is partitioned into three sub-batches with adjustable sizes and scheduled concurrently while preserving the computational-graph dependencies within each sub-batch. By allowing different sub-batches to progress through different computation stages, operators with complementary resource preferences can be overlapped in a pipelined manner. This reduces contention for the same type of hardware resource, mitigates pipeline bubbles, and improves overall hardware utilization. Furthermore, we formulate the

摘要: 随着Transformer模型在计算机视觉、自然语言处理等领域的广泛应用,推理请求数量呈现指数级增长趋势,任务规模的迅速扩张对模型推理吞吐率提出了更高的要求。现有推理框架主要通过算子融合、内核生成和图优化提升局部执行效率,但Transformer模型中的不同算子对Tensor Core、CUDA Core和内存带宽等硬件资源具有不同的需求偏好。受数据依赖限制,这些算子通常仍按拓扑顺序执行,导致不同资源之间难以形成充分重叠,限制了推理吞吐量的进一步提升。针对上述问题,提出一种基于批次分割的Transformer流水推理方法。Transformer模型由矩阵乘法、归约、逐元素计算和张量重排等多类算子组成。结合算子的资源占用特征及离线性能分析结果,将其划分为Tensor Core密集型、CUDA Core密集型和内存带宽密集型三类,为后续的跨子批次流水调度提供依据。随后,针对传统方法执行中严格数据依赖导致不同资源类型算子难以重叠的问题,将原始推理批次划分为三个规模可调的子批次,并在保持各子批次内部计算图拓扑依赖的前提下进行并行调度。通过使不同子批次处于不同的计算阶段,资源偏好不同的算子可在时间维度上形成流水重叠,从而降低同类资源竞争和减少流水线空泡,提高硬件资源利用效率。进一步地,将三个子批次的规模确定问题建模为离散优化问题,以计算图拓扑依赖和三类资源容量为约束,以三个子批次规模为决策变量,并以端到端推理吞吐量最大化为优化目标。考虑到不同子批次规模会改变各流水阶段的执行时长、资源占用和重叠程度,基于离线性能查表构建吞吐量评价函数。该函数综合考虑三个流水阶段的预计执行时长及其均衡程度:当某一阶段持续时间明显长于其余阶段时,其他阶段容易提前完成并进入等待,从而降低流水重叠效果。因此,算法优先选择能够缩短瓶颈阶段、平衡各阶段执行时长并减少资源空泡的候选划分方案。在此基础上,设计滑动窗口启发式搜索算法,在每轮迭代中评估当前窗口内的邻域候选划分方案,并根据评价结果调整子批次规模;同时采用窗口衰减、收益饱和和资源阈值等机制控制搜索过程,以低搜索开销获得吞吐量最优的批次划分方案。在NVIDIA GeForce RTX 4090 GPU上,在ViT-Base、Conformer-L、BERT-Base、GPT2-XL、BLOOM-7B1和LLaMA3-8B等六个具有代表性的Transformer模型上开展实验,覆盖视觉、语音、自然语言理解和生成式大语言模型等应用场景。实验结果表明,在不同模型和批次规模配置下,该方法相较于Welder、PyTorch Inductor、TensorRT、ONNX Runtime、MetaAttention和NanoFlow,吞吐量至多可提升1.74倍、6.26倍、1.79倍、8.94倍、3.56倍和3.68倍。结果验证了所提方法能够有效挖掘不同资源类型算子之间的互补性,并在多种Transformer推理负载下提升推理吞吐量。