Tilus: A Virtual Machine for Arbitrary Low-Precision GPGPU Computation in LLM Serving

发表时间: 2025-04 · arXiv:2504.12984

原文: https://arxiv.org/abs/2504.12984

文章标题: Tilus：一种用于LLM服务中任意低精度GPGPU计算的虚拟机
作者/机构: Yaoyao Ding (University of Toronto), Bohan Hou (Carnegie Mellon University), Xiao Zhang (University of Toronto), Allan Lin (University of Waterloo), Tianqi Chen (Carnegie Mellon University), Cody Yu Hao (Anyscale), Yida Wang (Amazon), Gennady Pekhimenko (University of Toronto)

---

速读

一句话结论 本文提出了一个专为大语言模型服务设计的通用GPU虚拟机Tilus，通过原生地支持1到8位任意精度的张量计算，在保持可编程性的同时，实现了比现有最先进方案最高2.6倍的性能加速。

要解决什么问题 大语言模型推理的显存带宽和计算吞吐压力极大，量化是核心解法，但现有GPU编程抽象在处理低精度计算时存在严重的机制卡点。虽然4位量化能节省资源，但会带来不可忽视的精度损失；而能兼顾精度的5到7位等非2的幂次位宽，却缺乏高效的底层算子支持。现有的两类主流编译器都绕不开各自的死结：以Triton为代表的过程导向编译器没有暴露GPU的显存层级，且缺乏对子字节数据类型的原生支持，导致在处理低精度权重时，必须手动从大存储类型中解包，并依赖共享内存进行极其昂贵的张量布局转换才能满足张量核心的指令要求；以Ladder为代表的调度导向编译器虽然引入了低精度原语，但它采用的是类型级打包（将低精度数据硬塞进大存储类型中），这种方式不仅无法高效处理非2的幂次位宽，其原语式的调度还直接阻断了软件流水线优化，导致权重从全局内存加载到寄存器的过程无法与计算重叠，严重拖慢了执行效率。

怎么做的 核心思路是构建一个直接暴露GPU分层显存（全局内存、共享内存、寄存器）的通用GPU虚拟机，通过在寄存器级别对数据块进行重新解释，彻底绕开依赖共享内存的布局转换开销。该方法由三个关键部件构成。首先是代数布局系统，它精确定义了张量元素在一个数据块内如何分布到各个GPU线程上。布局被形式化为一个映射函数 $L(t, i)$，表示线程 $t$ 中第 $i$ 个局部元素的逻辑索引。系统提供了将元素存在单线程内的局部布局和分布在多线程间的空间布局作为基石，并通过组合算子构建复杂布局： $$h(t, i) = (f \circ g)(t, i)$$ 这种代数结构使得系统能够精确计算并匹配不同数据类型在寄存器中的位分布。其次是带分层显存的单指令多块编程模型，它在线程块粒度上维护标量、指针和张量变量，提供显式的内存搬运和计算指令，让开发者能精准控制数据在显存层级间的移动以实现软件流水线。最后是任意低精度处理机制，为了高效加载如6位整数这种非标准位宽，系统在全局内存中先将其重排并紧凑打包为8位无符号整数，随后通过异步指令将其流水线式地加载到共享内存再到寄存器中。此时，代数布局系统发挥关键作用：它直接在寄存器内将这块数据重新解释为目标低精度类型和布局，因为两者在跨线程的位分布上是完全兼容的（例如每个线程恰好持有24位数据，即3个8位整数或4个6位整数）。完成重新解释后，再利用特定目标的寄存器指令进行无线程间通信开销的向量化类型转换，转为高精度浮点数送入张量核心计算。这套机制完美支持了1到8位的任意位宽。

效果如何 实验在NVIDIA L40S、A100和H100硬件上展开，测试了Gemma-2-9B、QWen2.5-32B和Llama-3.3-70B等模型。对比基线涵盖了多条技术路线：厂商标准库cuBLAS、过程导向编译器Triton、调度导向编译器Ladder、手动优化的专用算子QuantLLM和Marlin，以及端到端推理框架vLLM。在算子级测试中，针对从1位到8位的全谱系量化数据类型（如uint8、f6e3m2、int4、uint2等），该方法全面超越所有基线。在端到端性能上，无论是批处理大小为1或16的解码阶段，还是提示词长度为2048的预填充阶段，该方法均持续优于Ladder，因为后者未能实现软件流水线和k维并行化，导致显存带宽利用率极低且执行了大量冗余的张量核心指令。在跨硬件测试中，该方法展现了极强的适应性，而基线方法Ladder在H100上出现了内核生成失败的运行时错误，vLLM在L40S上遭遇了显存溢出。不过，该方法也存在一定的工程代价，它需要一套包含领域特定语言、中间表示和多级优化遍的完整编译流水线来生成底层CUDA代码；同时，如果直接使用位运算来提取和存储低精度数据会非常低效，必须严格依赖其提出的全局内存布局预转换和寄存器重解释机制才能保证性能。

A1 主要贡献

核心问题：大型语言模型（LLM）服务对计算资源，特别是内存带宽和计算吞吐量，提出了巨大需求。低精度计算是提高效率的关键技术，但现有方法存在局限。用于生成低精度计算核心（kernel）的方法通常仅限于位宽为2的幂的权重，并且由于高级GPU编程抽象的限制，性能欠佳。这些抽象限制了精细的寄存器管理和优化的内存访问模式等关键优化。例如，4位量化虽能显著节省计算资源，但仍有不可忽视的精度损失；而5到7位量化能缓解精度损失，却缺乏高效的GPU核心支持，导致其应用受阻。

研究目标：为了解决现有低精度计算核心生成方法的覆盖范围和性能差距，本文旨在设计一个能够支持任意位宽（1-8位）低精度数据类型，同时保持GPU可编程性和高性能的系统。

创新点：
本文提出了Tilus，一个专为低精度计算设计的通用GPU（GPGPU）虚拟机，其核心贡献如下：
1. 提出一个专用于低精度计算的GPGPU虚拟机：该虚拟机旨在解决现有方法在支持任意位宽（如5-7位）量化方面的覆盖不足和性能不佳的问题。
2. 新颖的代数布局系统（Algebraic Layout System）：Tilus引入了一个代数布局系统，用于精确描述张量元素在一个tile（数据块）内如何在GPU线程间分布。该系统能够灵活地将寄存器中的低精度tile重新解释（reinterpret）为硬件友好的数据类型tile，从而实现高效处理。
3. 带分层内存空间的线程块级编程模型：Tilus提供了一个线程块级的编程模型，并显式暴露了GPU的分层内存空间。这使得开发者能够对数据在GPU内存层次结构中的移动、放置和计算进行精细控制。
4. 支持任意位宽的低精度数据类型：Tilus原生支持1到8位范围内的任意位宽数据类型，包括有符号整数、无符号整数和浮点数，极大地扩展了低精度计算的应用范围。

通过广泛评估，Tilus不仅能生成覆盖整个低精度范围的高效核心，并且在其支持的核心上，性能超越了当前最先进的解决方案，最高可达2.6倍。

---

A3 背景知识

2.1 LLM服务与量化

• LLM推理的两个阶段与优化重点：大型语言模型（LLM）服务的推理过程包括预填充（prefill）和解码（decode）两个阶段。预填充阶段处理输入提示以建立上下文，解码阶段则基于先前的token迭代生成输出token。在LLM的所有层中，矩阵乘法主导了计算时间和内存消耗，因此其优化对于高效的LLM服务至关重要。

• 量化的作用与局限性：量化【索引11, LLM.int8(): 8-bit matrix multiplication for transformers at scale, 2024, NIPS '22】【索引17, Gptq: Accurate post-training quantization for generative pre-trained transformers, 2022, arXiv】通过将模型权重和激活值降低到低精度格式（如8位或4位整数），来提高效率。它减少了内存使用、带宽需求和推理延迟，同时试图保持模型精度。尽管4位量化能带来显著的计算节省，但最先进的方法【索引7, QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs, 2024, NeurIPS】【索引9, QuIP: 2-bit quantization of large language models with guarantees, 2023, NIPS '23】【索引29, SpinQuant: LLM Quantization with Learned Rotations, 2025, ICLR】仍存在精度下降问题。

• 对任意位宽支持的需求：将精度提高到5位、6位或7位【索引3, eXmY: A Data Type and Technique for Arbitrary Bit Precision Quantization, 2024, arXiv】【索引54, Quant-LLM: Accelerating the Serving of Large Language Models via FP6-Centric Algorithm-System Co-Design on Modern GPUs, 2024, USENIX ATC】有助于在保持效率的同时保护精度，但这些位宽缺乏优化的GPU支持，限制了它们的采用。当前GPU架构和软件栈主要为2的幂的位宽（如4位和8位）进行优化，使得任意位宽的计算效率低下。然而，对灵活量化的需求正在增长，因为4位对某些模型可能过于激进，而8位则浪费资源。支持更广泛的位宽可以在LLM服务中实现更好的精度-效率权衡，从而推动了对能够高效处理非标准低精度格式（如3、5、6、7位）的新核心生成技术的需求。

2.2 GPGPU编程

• GPGPU的执行与内存层次结构：通用GPU（GPGPU）编程通过一个结构化的执行和内存层次结构来组织并行计算【索引32, The GPU Computing Era, 2010, IEEE Micro】。执行层次结构始于线程（thread），这是最小的执行单位，它使用自己的寄存器和本地内存独立执行指令。线程被组织成线程块（thread blocks），通过共享内存实现数据共享并支持同步执行。一个网格（grid）由多个独立的线程块组成，通过组织成千上万的线程来实现大规模并行。GPU内存层次结构包括寄存器（registers）、共享内存（shared memory）和全局内存（global memory）。寄存器提供最快且线程私有的存储。共享内存可由一个线程块内的所有线程访问，速度快于全局内存。全局内存可被整个网格访问，但延迟较高。这种结构通过利用执行和内存层次结构，实现了高效的并行执行。

2.3 GPGPU语言与编译器

• 2.3.1 GPGPU虚拟机和语言：GPGPU编程涉及多种语言和编译器，它们在硬件抽象和控制之间取得平衡。像SASS【索引37, SASS: Streaming Assembler for NVIDIA GPUs, 2024, NVIDIA】和CDNA3【索引5, CDNA 3 Architecture for Accelerated Computing, 2024, AMD】这样的低级语言提供直接的硬件访问以进行精细优化，但需要深入的架构知识。在抽象层次上稍高一点的是NVIDIA的PTX【索引36, Parallel Thread Execution ISA Version 12.0, 2024, NVIDIA】，它作为连接CUDA【索引35, CUDA C++ Programming Guide, 2024, NVIDIA】等高级语言与GPU特定指令的中间表示，同时保留了优化的灵活性。像CUDA【索引35, CUDA C++ Programming Guide, 2024, NVIDIA】和HIP【索引6, HIP: Heterogeneous-Compute Interface for Portability, 2024, AMD】这样的高级语言通过扩展C编程语言简化了编程。尽管有这些语言，GPGPU编程仍然复杂，受限于硬件特定的内存和计算层次结构，并需要针对特定工作负载进行优化。为应对这些挑战，研究人员引入了更高级的语言和编译器，分为两类：过程导向编译器，通过超越CUDA的抽象来简化编程；以及调度导向编译器，通过声明式调度原语来优化计算-硬件映射。

• 2.3.2 过程导向编译器：这类编译器【索引12, Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs, 2023, ASPLOS】【索引20, Graphene: An IR for Optimized Tensor Computations on GPUs, 2023, ASPLOS】【索引26, MLIR: scaling compiler infrastructure for domain specific computation, 2021, CGO】【索引47, Triton: an intermediate language and compiler for tiled neural network computations, 2019, MAPL】让程序员能够直接编写核心，并提供抽象来简化过程。例如，Triton【索引47, Triton: an intermediate language and compiler for tiled neural network computations, 2019, MAPL】引入了tile编程模型，其中线程块的行为是程序化定义的，tile取代标量成为基本数据类型。这种方法结合了编程的简便性和高性能核心的生成，使Triton被广泛采用。然而，Triton缺乏对像uint4这样的低精度数据类型的原生支持。处理这些类型需要手动从更大的存储类型（如uint32）中解包子字节数据【索引21, Accelerating a Triton Fused Kernel for W4A16 Quantized Inference with SplitK work decomposition, 2024, arXiv】。此外，Triton没有暴露GPU内存层次结构，限制了程序员对数据加载和内存范围使用的控制，这使得低精度核心的性能优化变得复杂。这些限制导致了低精度核心执行效率低下。图1(a)展示了Triton生成的低精度核心中的低效之处，以一个uint4权重加载流水线为例。该过程包括四个步骤：1) 权重使用流水线化的cp.async指令【索引22, Alcop: Automatic load-compute pipelining in deep learning compiler for ai-gpus, 2023, MLSys】从全局内存异步复制到共享内存；2) 共享内存数据加载到寄存器；3) 执行解包和类型转换操作；4) 将寄存器张量布局转换为满足张量核心指令的要求。其中，步骤4是主要瓶颈，因为它依赖共享内存进行布局转换，这会产生巨大的开销。

• 2.3.3 调度导向编译器：调度导向编译器将计算与调度分离，以优化计算到硬件的映射。Halide【索引40, Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines, 2013, PLDI】开创了这种方法，后来被TVM【索引10, TVM: An Automated End-to-End Optimizing Compiler for Deep Learning, 2018, OSDI】及后续工作【索引16, TensorIR: An Abstraction for Automatic Tensorized Program Optimization, 2023, ASPLOS】【索引22, Alcop: Automatic load-compute pipelining in deep learning compiler for ai-gpus, 2023, MLSys】【索引43, Tensor program optimization with probabilistic programs, 2024, NIPS '22】【索引51, Attention is All you Need, 2017, NeurIPS】【索引52, Ladder: Enabling Efficient Low-Precision Deep Learning Computing through Hardware-aware Tensor Transformation, 2024, OSDI】【索引60, Ansor: Generating High-Performance Tensor Programs for Deep Learning, 2020, OSDI】【索引61, AMOS: enabling automatic mapping for tensor computations on spatial accelerators with hardware abstraction, 2022, ISCA】扩展到深度学习领域。其中，Ladder【索引52, Ladder: Enabling Efficient Low-Precision Deep Learning Computing through Hardware-aware Tensor Transformation, 2024, OSDI】是第一个通过引入专用原语来支持低精度计算的，它将低精度数据（如4位整数）打包到更大的类型（如8位整数）中。然而，Ladder有两个局限性。首先，由于其类型级打包（将低精度类型打包到存储类型中）的方式，它无法高效处理非2的幂的位宽。其次，其原语式的调度方式阻碍了像软件流水线【索引22, Alcop: Automatic load-compute pipelining in deep learning compiler for ai-gpus, 2023, MLSys】这样的优化，导致性能不佳。图1(b)展示了Ladder低精度核心中的权重加载过程。该过程包括：1) 无流水线地将权重从全局内存加载到寄存器；2) 向量化类型转换；3) 将转换结果存储在共享内存中；最后4) 使用ldmatrix指令将权重从共享内存加载到寄存器以进行后续的张量核心操作。权重加载和计算之间缺乏流水线严重影响了性能。
  

  

---

A2 方法细节

3. 系统概述

3.1 关键思想

• Tilus的GPGPU虚拟机设计：本文引入了一种新颖的GPGPU虚拟机（VM），专门用于克服编写高效低精度深度学习核心的挑战。该VM原生支持1到8位任意位宽的低精度数据类型，从而实现高效的权重加载和计算。图1(c)展示了VM的权重加载流水线，以uint4为例。它始于1) 从全局内存到共享内存的流水线化异步内存复制，接着是2) 从共享内存加载寄存器张量。然后，它3) 无开销地将寄存器张量重新解释为不同的数据类型和布局，最后4) 执行向量化类型转换。与图1中的其他方法相比，该流水线实现了更高的效率，因为它消除了布局转换（不像Triton【索引47, Triton: an intermediate language and compiler for tiled neural network computations, 2019, MAPL】）并集成了流水线（不像Ladder【索引52, Ladder: Enabling Efficient Low-Precision Deep Learning Computing through Hardware-aware Tensor Transformation, 2024, OSDI】）。更重要的是，该流水线具有通用性，使得本文的工作成为第一个无缝支持1位到8位任意低精度数据类型的工作。

• 核心设计理念：为实现此效率，Tilus的设计基于几个关键思想：

  1. 用于灵活优化的GPGPU虚拟机：选择实现GPGPU VM是出于对GPU编程更大灵活性的需求。与标准的基于循环的转换不同，我们的VM允许程序员直接实现和微调超越传统循环转换的优化。这种灵活性对于低精度计算至关重要，因为对执行策略的精细控制可以带来显著的性能提升。
  2. 带分层内存空间的线程块级编程模型：我们的VM明确暴露了GPU内存层次结构，包括寄存器、共享内存和全局内存，而这些在Triton【索引47, Triton: an intermediate language and compiler for tiled neural network computations, 2019, MAPL】等现有解决方案中是被抽象掉的。通过赋予程序员对数据放置和移动的精细控制，我们的方法实现了内存流水线并消除了不必要的布局转换，如图1所示。
  3. 代数布局系统：我们引入了一个代数布局系统，它精确定义了寄存器张量内的元素如何在线程间分布。这种结构化的表示简化了张量布局的构建、分析和解释。值得注意的是，它能够将低精度寄存器张量无缝地重新解释为标准数据类型，如步骤3在图1(c)中所示。
  4. 对任意低精度数据类型的原生支持：我们的VM为广泛的低精度数据类型提供了内置支持，包括1到8位宽度的有符号和无符号整数以及浮点数。支持的类型包括int2到int8、uint1到uint8，以及float3到float8，浮点类型具有任意的指数和尾数分布。
     这些创新共同提升了现代GPU上低精度核心开发的可编程性、效率和灵活性。我们没有选择扩展Triton【索引47, Triton: an intermediate language and compiler for tiled neural network computations, 2019, MAPL】，因为它的编程模型内在地抽象了张量布局，这与我们显式控制布局的方法不兼容。类似地，Ladder【索引52, Ladder: Enabling Efficient Low-Precision Deep Learning Computing through Hardware-aware Tensor Transformation, 2024, OSDI】依赖于类型级打包，而Tilus采用tile级重解释，这使得两者在根本上不兼容。

3.2 虚拟机程序示例

• FP16 x INT6矩阵乘法示例：图2展示了虚拟机中一个低精度矩阵乘法的例子。矩阵乘法定义为 C_M,N = A_M,K × B_K,N，其中A和B的类型分别为float16（一种16位浮点数【索引23, IEEE Standard for Floating-Point Arithmetic (IEEE 754-2019), 2019, IEEE】）和int6（一种6位有符号整数）。该核心执行给定M、N和K的矩阵乘法，每个线程块计算C矩阵的一个BM × BN大小的tile（第1行）。因此，需要启动一个(M / BM, N / BN)的线程块网格（第2行）。

• 程序执行流程：在核心内部，BlockIndices指令获取线程块索引bi和bj（第3行），这决定了计算相应C tile的偏移量(bi * BM, bj * BN)。通过指定输入输出张量在全局内存中的地址和形状，创建了三个张量视图（第4-6行）。然后，创建了一个类型为f16[16, 8]的寄存器张量，其布局为 local(2, 1).spatial(8, 4).local(1, 2)。它将16 × 8 = 128个元素分布在32个线程中，每个线程存储4个元素（第7行）。这个布局由三个基本布局（第4节）组成，并与PTX【索引36, Parallel Thread Execution ISA Version 12.0, 2024, NVIDIA】中mma.m16n8k16张量核心指令使用的C矩阵布局对齐。

• 核心计算循环与低精度数据处理：对k维度的归约循环（第8-13行）重复地从全局内存加载A和B的tile到寄存器，并执行矩阵乘累加（mma）。在每次迭代中，我们首先使用LoadGlobal指令将一个f16[16, 16]的tile从全局内存加载到寄存器（第9行）。加载的寄存器tile的布局由张量核心指令指定和要求。offset参数指定了加载的tile在全局张量中的位置。加载数据类型为int6的张量B涉及一个更复杂的过程，详见第6节。我们在这里总结其高级思想。作为启动核心前的预处理步骤，权重张量在全局内存中的布局从i6[K, N]转换为u8[K / BK, N / BN, BK * BN * 6 / 8]，以便通过LoadGlobal指令高效加载（图2(b)中的“改变布局”步骤）。接下来，在核心中，转换后的tile被加载到一个寄存器张量中（第10行），然后被重新解释为一个具有不同数据类型和布局的张量（第11行）。这种重解释是有效的，因为两个张量都分布在相同数量的线程（32个）上，每个线程恰好持有24位（即3个u8或4个i6），如图2(c)所示。之后，i6张量被转换为f16张量（第12行），然后送入张量核心执行矩阵乘累加（mma）（第13行）。最后，累加张量从f32转换为f16并存储到全局内存中（第14-15行）。为简单起见，该程序未使用共享内存，并省略了软件流水线【索引22, Alcop: Automatic load-compute pipelining in deep learning compiler for ai-gpus, 2023, MLSys】等优化；此外，每个k-迭代仅执行一个张量核心指令【索引35, CUDA C++ Programming Guide, 2024, NVIDIA】。优化后的实现可在附录B中找到。
  

  

4. 代数布局系统

• 布局的定义与重要性：虚拟机向程序员暴露了带有全局内存、共享内存和寄存器的分层内存空间。我们需要一种方法来为所有三个内存范围建模张量元素的逻辑索引与相应元素在内存中位置之间的映射。这种映射通常被称为张量的布局。图3展示了一个张量核心指令 mma.m16n8k8.f32.f16.f16.f32 D, A, B, C 所使用的布局示例。它执行计算：D_16,8 = A_16,8 * B_8,8 + C_16,8，其中A、B、C、D是存储在线程寄存器中并分布在一个warp的32个线程中的张量。由于元素分布在不同的线程中，我们称这种布局为分布式布局【索引47, Triton: an intermediate language and compiler for tiled neural network computations, 2019, MAPL】。这种布局可以定义为一个函数 L，它将一个线程索引 t 和该线程内的一个局部索引 i 映射到相应张量元素的逻辑索引 L(t, i)。例如，图3中的布局可以表示为：
  

  

  
  在这里，t 的范围是0到31，i 的范围是0到3。函数 L(t, i) 表示线程 t 中元素 i 的逻辑索引。由于线程块中的所有线程都可以访问共享内存和全局内存，这种形式化也可以在单线程假设下用于描述它们的布局。也就是说，通过设置 t = 0，我们将 L(0, i) 定义为共享或全局内存中地址 i 处元素的逻辑索引。
  

  

4.1 参数化的基本布局

• 两种基本布局：Local 和 Spatial：在形式化定义布局之后，我们引入了参数化的基本布局，它们作为我们布局代数的基本构建块。给定一个形状为 (n1, n2) 的tile，主要有两种存储方式：(1) 将所有n1n2个元素存储在单个线程中，或 (2) 将所有元素分布在n1n2个线程中，每个线程持有一个元素。我们将第一种类型称为局部布局(local layouts)，表示为local(n1, n2)，第二种称为空间布局(spatial layouts)，表示为spatial(n1, n2)。这个概念可以自然地扩展到任意维度的tile。图4展示了这两种基本布局。local(2, 3)布局将线程t的第i个局部元素映射到逻辑索引 (i/3, i % 3)，而spatial(2, 3)布局将其映射到 (t/3, t % 3)。我们观察到，现代深度学习工作负载和GPU中使用的所有复杂布局都可以用这两种基本布局来构建。
  

  

4.2 布局组合

• 通过组合构建复杂布局：现代深度学习工作负载以及硬件指令定义的布局通常呈现出一种层次结构。以图5中的布局(c)为例。这个布局形状为(4, 6)，将24个元素存储在6个线程中，每个线程持有4个元素。我们将每个线程中存储的4个元素表示为e0, e1, e2, e3。比较它的前两行和后两行，我们观察到相似的结构，只是后两行存储的是e2和e3中的元素，而不是e0和e1。为了模拟这种结构不变性，布局(c)可以看作是布局(a)和(b)的组合，其中布局(a)中的每个元素代表一个具有布局(b)的tile。实际上，布局(a)和(b)可以组合起来表示布局(c)，如下所示：
  

  

  
  其中 0 ≤ t < 6，0 ≤ i < 4，⊙ 表示逐元素乘积，(2, 6) 表示布局(c)的形状。

• 布局组合的泛化与性质：这个组合原则可以被泛化。给定两个具有相同维度数量的布局 f 和 g，我们定义它们的组合 h = f ◦ g 为：
  

  

  
  其中 T_x, N_x, S_x 分别代表布局x的线程数、每个线程的局部元素数和形状。我们可以证明组合是结合的，即对于任意三个布局f, g, h，等式 f ◦(g◦h) = (f ◦g)◦h 成立。然而，组合是不可交换的，即通常情况下 f ◦g ≠ g◦f。空间布局和局部布局分别对线程和局部元素遵循行主序。使用组合，我们可以构造它们的列主序对应物，column_spatial(...)和column_repeat(...)，如布局(e)在图5中所示。回到图3中的张量核心指令布局，它可以表示为一个组合布局 local(2, 1).spatial(8, 4).local(1, 2)。使用布局组合，我们也可以定义其逆操作。如果 h = f ◦ g，我们定义 g = h/f 为布局h除以布局f的结果。例如，将local(2, 4)除以local(1, 2)得到local(2, 2)。在附录A中，我们正式定义了布局，并证明了在组合算子下，它们形成一个带有单位元的半群，使其成为一个幺半群(monoid)。
  

  

5. 线程块级编程模型

• 单指令多块（SIMB）模型：现代GPU编程模型，如PTX【索引36, Parallel Thread Execution ISA Version 12.0, 2024, NVIDIA】和CUDA【索引35, CUDA C++ Programming Guide, 2024, NVIDIA】，在线程级别定义操作，遵循单指令多线程（SIMT）范式【索引1, The SIMT Core: Instruction and Register Data Flow, 2018, General-Purpose Graphics Processor Architectures】。为了简化GPU编程，我们采用线程块级编程模型，在线程块的粒度上定义操作，而不是单个线程。此外，基于先前介绍的布局系统，我们提议明确暴露现代GPU中的分层内存结构，从而在降低编程复杂性的同时实现精细的内存控制。我们将此模型称为单指令多块（SIMB）。

5.1 状态空间和类型系统

• 虚拟机支持的变量类型：虚拟机支持三种类型的变量。标量变量存储单个值，如整数（例如，int32）或浮点数（例如，float16）。指针变量存储内存地址而不是直接的数据值。张量变量表示多维数组，其类型指定了它们的形状、元素类型、内存范围和布局。张量存在于不同的内存范围中，包括全局内存、共享内存和寄存器。张量布局决定了高维张量元素如何映射到线性内存。虚拟机中的所有变量都在线程块级别上操作，意味着一个块内的所有线程协同维护它们。

5.2 程序结构和控制流

• 虚拟机程序结构：图6展示了一个VM程序的结构。每个程序由一个程序名、一个网格形状、一个参数列表和一个程序主体组成。网格形状被指定为由<...>括起的一系列表达式，其中每个表达式要么是一个正整数，要么是基于程序参数的整数表达式。如果网格形状包含基于参数的表达式，其维度将在运行时根据程序的启动参数确定。程序主体由一系列语句组成，包括if-else语句、基于范围的for循环和while循环。与其他低级虚拟机【索引36, Parallel Thread Execution ISA Version 12.0, 2024, NVIDIA】或指令集架构（ISA）【索引37, SASS: Streaming Assembler for NVIDIA GPUs, 2024, NVIDIA】不同，我们的虚拟机不将控制流语句抽象为跳转指令。相反，它保留了高级控制结构，以提高人类开发者的可读性和编程便利性。除了控制流语句，单个指令也可以作为语句。虚拟机提供的大部分功能都是作为线程块级指令集中的指令实现的。
  

  

5.3 线程块级指令集

• 指令集特性：VM指令集中的每条指令都在线程块级别而不是线程级别上操作。表1显示了指令集中的指令列表，包括每条指令的签名和指令语义的简要描述。这些指令在指定的内存空间（如全局内存、共享内存、寄存器）中分配具有特定数据类型、形状和布局的张量，在内存空间之间传输张量，并对寄存器张量执行计算或转换。现代处理器会乱序执行指令【索引48, An efficient algorithm for exploiting multiple arithmetic units, 1967, IBM J. Res. Dev.】，这意味着在没有依赖关系的情况下，后续指令可能在当前指令完成前开始执行。同样，我们虚拟机中指令的执行也表现出这种行为：某些后续指令可能在当前指令完成前开始执行。通常，这种行为不会造成重大问题。然而，当两条指令访问共享或全局内存的同一区域，并且第二条指令依赖于第一条指令的完成时，就会出现异常。在这种情况下，必须插入一个Synchronize指令，以确保所有前面的指令在后续指令执行前完成。
  

  

6. 任意低精度数据类型

• 低精度数据类型的重要性：现代处理器使用字节（每个8位）作为最小处理单元。因此，现代编程语言中的标准数据类型通常具有8的倍数的位宽。然而，LLM的高计算和内存需求使得低精度数据类型（小于8位）对于减少资源消耗至关重要。本节描述虚拟机如何高效地支持低精度数据类型。

6.1 低精度数据的存储

• 紧凑存储与位操作：由于包括CPU和GPU在内的现代处理器使用字节作为内存访问和计算的最小单位，我们将低精度数据（每个元素少于8位）紧凑地存储在字节内，如图7所示。紧凑存储消除了连续低精度值之间的位间隙，这可能导致单个值跨越两个uint8条目（例如，图7中的b[1]）。我们采用位运算来提取、操作和存储打包字节数组中的低精度值。要加载一个低精度值，我们首先使用按位与（AND）提取相关位，用按位移位（SHIFT）操作调整它们的位置，如果值跨越多个字节，则最后使用按位或（OR）组合分离的部分。同样，要存储一个低精度值，我们首先使用位掩码清除目标位位置，然后使用按位或（OR）插入新值，同时保留其他位。低精度数据在算术计算前被转换为标准数据类型，计算后再转换回来。虽然这些方法能够支持任意位宽的数据类型，但它们通常效率低下，仅作为一种后备机制。为LLM服务，需要更高效地处理低精度数据。
  

  

6.2 在LLM中高效支持低精度

• 低精度核心的关键步骤：LLM中的低精度核心在计算前通常遵循两个步骤：(1) 将权重从全局内存加载到片上内存（寄存器或共享内存），以及 (2) 将低精度权重转换为高精度（例如，float16），然后进行反量化。因此，高效的内存加载和类型转换对性能至关重要。

• 高效的低精度权重加载：借助前一小节讨论的低精度支持，我们的虚拟机可以使用LoadGlobal指令加载低精度张量。然而，由于多次位运算和非合并的内存访问【索引35, CUDA C++ Programming Guide, 2024, NVIDIA】，直接这样加载效率低下。为解决此问题，我们转换了全局内存中权重张量的布局，以实现更高效的加载。如果不进行转换，加载一个数据类型为i6且布局为local(2, 1).column_spatial(4, 8).local(2, 1)的寄存器张量会导致非连续的内存访问，从而引起多次内存访问事务【索引35, CUDA C++ Programming Guide, 2024, NVIDIA】。此外，提取低精度位需要额外的位运算。为了优化这一点，我们识别出一个兼容的张量类型，其数据类型为uint8，布局为local(3).spatial(32)，这保留了线程数和线程局部元素数，同时实现了高效的内存加载。如图8所示，我们将权重张量[K, N]划分为形状为[BK, BN]的tile。每个tile从i6[BK, BN]重新解释为u8[BK * BN * 6 / 8]（第19行）并连续存储（第20行）。这使我们能够使用图2（第10、11行）中硬件友好的指令高效加载tile，同时也能够像标准数据类型一样进行流水线异步内存传输，并避免任何依赖共享内存的布局转换。此方法可推广到加载任何具有任意布局的低精度张量。更正式地，给定一个每线程B字节、T个线程的张量，我们使用数据类型uint8和布局local(n2).spatial(T).local(n1)来重新解释它，其中n1 = gcd(B, 16)，n2 = B/gcd(n1, 16)。
  

  

• 高效的类型转换：加载后，权重必须从低精度转换为高精度（例如，float16）以进行计算，特别是当硬件缺乏对给定低精度格式的原生支持时。我们利用特定目标的指令进行高效的向量化类型转换。在CUDA上，我们使用PRMT（在32位寄存器中排列字节）、LOP3（对三个输入进行任意逻辑操作）和位运算指令来以最小的开销执行类型转换，因为所有操作都在寄存器内执行，不需要线程间的任何通信。

7. 实现

• Tilus系统组件：Tilus由五个主要部分组成：一个Python中的领域特定语言（DSL）、一个中间表示（IR）、优化遍（optimization passes）、一个代码生成器和一个运行时系统。DSL使开发者能够用Python编写Tilus程序，然后这些程序被翻译成VM的IR以进行进一步处理。优化遍通过消除冗余和简化算术表达式来优化IR。代码生成器将优化后的IR翻译成Hidet IR【索引12, Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs, 2023, ASPLOS】，一种类似CUDA C的中间表示。之后，我们应用第6节中的转换来实现低精度类型与标准精度类型的转换，同时保留原始语义。最终的CUDA C代码由Hidet IR生成，并使用nvcc编译器【索引35, CUDA C++ Programming Guide, 2024, NVIDIA】编译成硬件二进制文件。运行时系统管理动态加载的二进制文件并提供执行环境。整个系统由大约2万行Python和C++代码组成。关于编译流水线的更多细节可以在附录C中找到。

---

A4 实验

实验环境

• 工作负载：

  • 模型: Gemma-2-9B, QWen2.5-32B, Llama-3.3-70B-Instruct。
  • 阶段: 评估了预填充（prefill）和解码（decode）两个阶段。
  • 算子: 算子级分析聚焦于从这些模型中提取的矩阵乘法核心。

• 基线系统:

  • 库: 厂商库 cuBLAS【索引34, NVIDIA cuBLAS Library, 2023, NVIDIA】。
  • 编译器: Triton v3.1.0【索引47, Triton: an intermediate language and compiler for tiled neural network computations, 2019, MAPL】和 Ladder (bitblas v0.0.1.dev15)【索引52, Ladder: Enabling Efficient Low-Precision Deep Learning Computing through Hardware-aware Tensor Transformation, 2024, OSDI】。
  • 手动优化核心: QuantLLM (commit 9802c5a)【索引54, Quant-LLM: Accelerating the Serving of Large Language Models via FP6-Centric Algorithm-System Co-Design on Modern GPUs, 2024, USENIX ATC】和 Marlin v0.1.1【索引18, MARLIN: Mixed-Precision Auto-Regressive Parallel Inference on Large Language Models, 2024, arXiv】。
  • 端到端框架: vLLM v0.5.3【索引25, Efficient Memory Management for Large Language Model Serving with PagedAttention, 2023, SOSP】。

• 硬件配置:

  • 主平台: NVIDIA L40S GPU (48 GiB)。
  • 其他平台: NVIDIA A100 和 H100 GPUs。

• 软件配置:

  • 驱动与工具包: GPU driver 565.57.01, CUDA Toolkit 12.6.3。
  • 实验协议: 算子实验每个核心运行50次，模型实验每个模型运行10次，报告中位延迟。每次运行前清除L2缓存。

实验结果

8.2 低精度核心性能

• 实验内容：使用一个自动调优的虚拟机程序模板，比较Tilus与cuBLAS、Triton、Ladder、QuantLLM和Marlin在多种低精度矩阵乘法（u8, f6, u4, i4, u2, u1）上的性能。
• 实验结果：如图9所示，Tilus在所有测试的数据类型上均实现了对cuBLAS的加速，并且性能优于所有基线。Triton因昂贵的布局转换而性能不佳；Ladder虽避免了转换，但缺乏软件流水线等关键优化，导致内存带宽利用不足；QuantLLM和Marlin虽为特定场景优化，但缺乏灵活性。
• 分析结论：Tilus通过一个参数化的虚拟机模板，结合其抽象良好的编程模型，高效地支持了全方位的量化类型，全面超越了现有方法。
  

  

8.3 任意数据类型支持

• 实验内容：评估Tilus在支持全谱量化权重数据类型（uint1-8, int2-8, float3-8）时的性能，与cuBLAS FP16核心进行比较。
• 实验结果：如图10所示，对于从1位到8位的各种整数和浮点数据类型，Tilus均实现了显著的加速。
• 分析结论：这证明了Tilus的虚拟机能够高效地支持任意低精度类型，为现代GPU上的低精度计算提供了一个强大的解决方案。值得注意的是，所有核心都来自同一个程序模板，仅通过参数化tile大小实现，编程工作量有限。
  

  

8.4 端到端性能

• 实验内容：在Gemma-2-9B, QWen-2.5-32B, Llama-3.3-70B模型上，评估Tilus和Ladder在预填充和解码阶段的端到端性能。
• 实验结果：如图11所示，Tilus在所有模型和阶段上都持续优于Ladder，尤其是在解码阶段。
• 分析结论：Ladder生成的内核由于未能实现软件流水线和k维并行化等关键优化，性能不佳。对于预填充阶段，Tilus通过高效处理量化权重布局，在将权重解码为float16进行计算时开销极小，从而获得了优越性能。
  

  

8.5 案例研究

• 8.5.1 跨不同硬件的加速：

  • 实验内容：在NVIDIA A100、L40S和H100 GPU上评估QWen2.5-30B模型的端到端性能，对比vLLM(float16)、Ladder(uint4)和Tilus(uint4)。
  • 实验结果：如图12所示，Tilus在所有配置中均优于Ladder。Ladder在H100上生成内核失败（ERR），vLLM在L40S上出现内存不足（OOM）错误。
  • 分析结论：Tilus展现了跨不同GPU架构的强大性能和适应性。
    

    

• 8.5.2 跨不同批处理大小的加速：

  • 实验内容：在Llama-3.3-70B模型上，使用f6和u4量化数据类型，分析不同批处理大小下（解码阶段1-16，预填充阶段4096-12288）的矩阵乘法性能。
  • 实验结果：如图13所示，Tilus在解码和预填充阶段的所有批处理大小下，性能始终优于基线方法。
  • 分析结论：Tilus的性能优势在各种规模的工作负载中都得以保持。详细的性能分析见附录D。
    

    

---

A5 结论

本文介绍了Tilus，一个为高效低精度LLM服务设计的GPGPU虚拟机，它解决了现有解决方案的关键局限性。Tilus的核心特性包括：用于线程块寄存器内张量分布的代数布局系统，具有精细内存管理的线程块级编程模型，以及对1到8位任意精度的子字节数据类型的广泛支持。实验结果表明，与Triton和Ladder等最先进的框架相比，Tilus取得了显著的性能提升，展示了其方法的灵活性和可扩展性。这项工作为高效和可扩展的LLM推理奠定了基础，为在新兴硬件、先进量化技术和多样化低精度格式方面的进一步优化铺平了道路。

---

A6 附录

A. 布局形式化

• 布局的正式定义：本附录形式化了布局的概念。定义1 (布局)：布局L是一个函数，将线程索引t和局部元素索引i映射到张量中的逻辑索引L(t, i)。定义2 (局部布局) local(n1, ..., nr)：一个布局，其中所有元素按行主序存储在单个线程中。定义3 (空间布局) spatial(n1, ..., nr)：一个布局，其中所有元素分布在多个线程中，每个线程持有一个元素。

• 布局操作与代数结构：定义4 (布局组合) h = f ◦ g：定义了如何通过组合两个布局来构建更复杂的层次化布局。定理1 (布局组合的结合性)：证明了布局组合操作满足结合律 f ◦ (g ◦ h) = (f ◦ g) ◦ h。定义5 (布局广播)：定义了如何通过前置补零来增加布局的秩，以处理不同秩的布局组合。定义6 (布局除法)：定义了组合的逆操作。定理2 (布局的幺半群结构)：证明了布局集合在组合算子◦下构成一个幺半群，因为它满足封闭性、结合性并存在单位元。

B. 优化的矩阵乘法

• 优化策略：本节展示了一个优化的矩阵乘法程序（图14），它在图2的简化版基础上引入了软件流水线和更优的tiling大小。关键改进是使用共享内存来存储全局子张量，以增强数据复用，并作为结果写回全局内存前的暂存区，确保合并的全局内存访问。通过CopyAsync等指令实现从全局内存到共享内存的异步数据传输。

• 代码分析：

  1. 初始化：定义了工作负载大小、tiling大小和流水线阶段数，并指定了寄存器和共享内存的张量布局。引入了reduce(...)和swizzle(...)两种新的布局操作，swizzle用于避免共享内存的bank-conflict。
  2. 内存分配与预加载：初始化全局张量视图，为软件流水线分配共享内存和累加寄存器。在主循环开始前，预加载前STAGES - 1个阶段的共享内存缓冲区。
  3. 主循环：在每次迭代中，k维度被进一步tiling以减少寄存器使用。从共享内存加载数据到寄存器，对量化数据类型进行重解释，并转换为计算数据类型（如FP16）。在执行点积运算的同时，使用单独的缓冲区异步发起全局到共享的内存传输，从而使GPU的计算和内存传输单元并行工作。
  4. 结果存储：主循环完成后，将累加张量从float32转换为float16，先存入共享内存，再重新加载到新的寄存器张量中，最后以合并访问的方式写回全局内存。
     

     

C. 程序编译和运行时

• 编译流程：Tilus提供一个Python DSL，程序通过以下步骤编译为GPU可执行代码：

  1. 内存规划：共享内存规划器计算核心所需的共享内存大小，并将共享张量映射到核心的共享内存空间区域。全局内存规划器管理所有线程块共享的全局内存分配。
  2. 代码生成：逐条为VM指令生成低级GPU代码（使用Hidet IR【索引12, Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs, 2023, ASPLOS】）。此过程进行指令选择（如使用更高效的ldmatrix替代lds）和自动向量化（如使用cp.async.v4, lds128）。
  3. 低精度数据类型降级：在低级IR上应用第6.1节讨论的规则，将所有低精度操作转换为硬件友好类型上的相应操作。通常，由于布局形式化和寄存器张量重解释，内存加载会被标准类型替代，因此主要应用的是从低精度到标准类型的向量化类型转换。最后，从低级IR生成CUDA代码，并用nvcc编译成硬件二进制文件。

• 虚拟机运行时：编译后的二进制文件可由虚拟机运行时加载。运行时维护内部状态以服务核心执行，包括：1) 按需分配的工作区内存；2) 存储CUDA流的执行上下文；3) 缓存在内存中的核心。

D. 量化核心的性能分析

• 性能剖析：使用Nsight Compute对cuBLAS、Ladder和Tilus的核心进行性能分析，结果如图15所示，展示了内存单元的吞吐量利用率和计算单元的流水线利用率。

• Ladder的瓶颈：

  • 对于小批量，Ladder主要依赖CUDA Cores，但受限于DRAM带宽（96.6%利用率），因为每个线程独立从全局内存加载操作数A，没有利用共享内存进行数据复用。
  • 对于大批量（≥16），Ladder转向Tensor Core，但未使用cp.async进行异步内存传输或软件流水线来重叠计算和内存操作。此外，其Tensor Core执行策略存在缺陷，因批量维度填充导致执行的Tensor Core指令比必要的多达八倍。

• Tilus的优势：

  • Tilus利用代数布局系统，将量化权重解释为标准数据类型，并使用cp.async优化内存访问模式。
  • Tilus集成了软件流水线，以重叠内存传输和计算，从而实现卓越的性能和硬件资源利用率，如图15中D行所示。