Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache

Posted on 2024年4月12日2024年12月7日 by remaper in LLM 推理

Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache

1. 问题

模型的架构：

QKV Linear layer
Multi-Head Self-Attention Mechanism, or the attention module
FNN，Feed-Forward Neural Network

超长文推理的挑战，2k -> 256k ？

有2个：

不同的层对显存的需求差异很大（注意力层的显存和输入输出是成比例的，但是其他层不是），这就限制了并行的效率（类似木桶效应？？）
KV cache 的大小不可预估，显存容量不可规划

这2个问题，最终也极大的限制了超长文推理的效果

SpotServe – Serving Generative Large Language Models on Preemptible Instance

Posted on 2024年3月5日2024年12月5日 by remaper in LLM 推理

https://mp.weixin.qq.com/s/uK1owRF53FW4WI7nldYadg
https://arxiv.org/pdf/2311.15566.pdf
业界首个，在可抢占式实例上运行的分布式大语言模型（LLM）服务系统
目标：LLM 降本 -> 抢占式 GPU 实例 -> 怎么更好的用这些抢占实例
相比直接使用抢占式GPU实例，SpotServe 可以将推理引擎的长尾延迟降低 2.4x – 9.1x
LLM 特点：

高计算量
大内存占用

无论是哪一点，都意味着成本昂贵
把 LLM 运行在随时可抢占的 GPU 实例上
传统的方法，MArk、Cocktail，单实例多GPU卡，运行一个 small DNN 模型，通过请求重定向或者冗余计算来处理抢占，但是这种方式只是和数据并行这种小的DNN模型，不适合LLM
LLM 会同时使用数据并行、模型并行、流水线并行多种技术，单个实例抢占会影响整个多个实例的计算结果
所以需要有更有效的方法
亮点：第一个做推理容错的论文（其他的都是在做训练的容错）
SpotServe 的创新点：

动态配置并行度
实例迁移优化（复用模型参数、中间结果，减少迁移后的传输数据量）：aws 抢占GPU实例，迁移后冷启动需要2m
高效利用宽限期（30s，尽量不中断推理）

1. 背景

1.1. Generative LLM Inference

当前类似的系统：

FasterTransformer
Orca
FairSeq
Megatron-LM

单次推理的总耗时，分析

texe(Sin) 是输入序列的解码时间
Sout 是生成的 token 数量，texe(1) 是生成每个 token 的时间
KV cache 技术可以将每个 token 的生成时间优化到接近常量 (i.e., 𝑡𝑒𝑥𝑒 (1) in E.q.(2) and Figure 1a).

深入浅出 tvm – (16) Tensor Expression 详解

Posted on 2023年8月26日2024年9月22日 by remaper in 深入浅出 tvm Tagged: tvm 源码学习

https://homes.cs.washington.edu/~arvind/papers/tvm.pdf

TVM addresses these challenges with three key modules.

(1) We introduce a tensor expression languageto build operators and provide program transformation primitives that generate different versions of the program with various optimizations.

tensor expression 是 tvm 引入的一门面向张量计算的函数式编程语言。前面我们知道 topi 就是基于 te 实现了一系列常见的张量算子库

在 tvm 里面，每个算子都可以表达为一个 te 表达式，te 接受张量作为输入，并返回一个新的张量

这个语言本身是非常简单的：

Type：因为处理的输入输出，都是 tensor，所以只有一种数据类型（不像 relay ir，有 Var、Constant、Function、Call 等各种类型）
Op：操作符也很简单，最常见的就2个 placeholder_op，compute_op

te 的代码主要在 include/tvm/te/ src/te/ 这2个目录，全部加起来才1.3w行，所以说是并不复杂

这篇文章，我们从编译器的视角来理解一下 tensor expression

参考资料：

1. te 的设计目标

te 可能应该有其他更多的设计目标

但从我目前了解到的来看，te 最大的作用应该是简化tensor 算子的编写难度，并支持进一步的调度和优化

比如向量加法

1）TVMScript 版本

在没有 te 之前，比如只有 tir 的时候，用户只能用 TVM script 来编写tensor 算子。如下：

@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def main(a: T.handle, b: T.handle):
        # 我们通过 T.handle 进行数据交换，类似于内存指针
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # 通过 handle 创建 Buffer
        A = T.match_buffer(a, (8,), dtype="float32")
        B = T.match_buffer(b, (8,), dtype="float32")
        for i in range(8):
            # block 是针对计算的抽象
            with T.block("B"):
                # 定义一个空间（可并行）block 迭代器，并且将它的值绑定成 i
                vi = T.axis.spatial(8, i)
                B[vi] = A[vi] + 1.0

ir_module = MyModule

2）tensor expression 版本

但是如果用 te 来描述向量加法的话，只需要2行，并且通过script() 打印出来会发现和 TVM script 是一样的

from tvm import te

A = te.placeholder((8,), dtype="float32", name="A")
B = te.compute((8,), lambda *i: A(*i) + 1.0, name="B")

func = te.create_prim_func([A, B])
ir_module_from_te = IRModule({"main": func})
print(ir_module_from_te.script())

使用 ftrace 跟踪内核函数调用栈

Posted on 2023年8月15日2024年12月22日 by remaper in eBPF & ftrace

搞虚拟化的，经常会遇到很多和内核相关的线上问题。比如最近我们遇到一个 k8s 的 pod 在删除的时候 pvc 卸载不掉的情况。

umount 的时候提示挂载点已经 not mounted 了，但是在 /proc/mounts 里仍然存在。statvfs 看了一下，这个挂载点的device id和父目录的device id 是一样的，说明已经不是挂载点了（因为看 /proc/mounts，这个挂载点的src是在另外一个设备上）

特别注意：如果 /proc/mounts 里，src 和 dst 是同一个设备，那就是 –bind 的方式，这种情况下就很难判断挂载点是否已经真正的卸载掉了

针对这个问题，社区已经有了一个比较 hack 的解决方案：https://github.com/kubernetes/kubernetes/issues/114546，就是忽略 /proc/mounts 残留，允许 pod 删除流程继续往下走

但是根源还是内核的存在脏数据（大概率猜测可能是引用计数泄露导致），还是得从内核层面定位这种问题，而且这种问题，很多时候我们得知道内核代码的执行路径是什么样的。这个时候用 ftrace 来定位就很方便了，比 eBPF 好使，ftrace 能够直接跟踪一个内核函数的调用栈

Pathways – Asynchronous Distributed Dataflow for ML

Posted on 2023年7月1日2023年11月1日 by remaper in LLM 训练

论文：https://arxiv.org/pdf/2203.12533.pdf

google 认为 pathways 是下一代AI架构：https://blog.google/technology/ai/introducing-pathways-next-generation-ai-architecture/

其他参考：

1. 论文背景、动机、贡献

在探讨现有 AI 模型的局限时，Jeff Dean 曾经说过，今天的人工智能系统总是从头开始学习新问题，我们为数千个单独的任务开发了数千个单独的模型，这种方式需要大量的数据和算力，是非常低效的

理想的发展模式应该是训练1个模型来干多个事情

而导致这个现象的根本原因，是因为过去10年的机器学习，总是算法&硬件&系统一起演进的，这种 co-evolution 的方式有一定的隐患：系统过于针对当前的任务，而不能很好的适应未来的需求

未来 ML 的需求是什么？

模型会越来越大：
1. 多模态：复杂，参数多。纯粹的数据并行已经不行了，开始引入流水线并行、模型并行等新的并行训练方式
2. 计算稀疏：大模型特有的一个问题
硬件异构：SPMD 不适用了，因为我们需要在不同的“硬件”上运行不同的代码，需要转向 MPMD
训推一体：提升硬件利用率

为了实现这个愿景，论文提出了一种叫 [Pathways] 的通用AI架构。Pathways 旨在用一个架构同时处理多项任务，并且拥有快速学习新任务、更好地理解世界的能力

深入浅出 tvm – (15) TVM Operator Inventory (TOPI)

Posted on 2023年6月10日2024年9月22日 by remaper in 深入浅出 tvm Tagged: tvm 源码学习

topi 是 tvm 的一个张量算子库，提供了很多常见的算子操作，例如 conv2D，transpose，等等。

tvm 在编译计算图的时候，会先将计算图从 relay ir 翻译成 tir，再将 tir 翻译成目标设备代码（比如llvm，比如c，比如cuda），但是 relay ir 和 tir 之间还有一种中间语言叫 tensor expression，这是一种专门为张量计算而设计的语言。relay ir 的绝大多数部分 op 其实都是通过 tensor expression 这种语言来实现的

https://tvm.apache.org/docs/reference/langref/relay_expr.html

An operator is a primitive operation, such as add or conv2d, not defined in the Relay language. Operators are declared in the global operator registry in C++. Many common operators are backed by TVM’s Tensor Operator Inventory.

topi 和 te 的关系：

te( tensor expression ) 是一门函数式编程语言，面向张量计算的
topi 是基于 te 实现的张量算子库

类似于编程语言和标准库之间的关系

1. topi 算子列表

源码目录：src/topi

src/topi 只是把算子通过 TVM_REGISTER_GLOBAL 注册到 tvm 的全局函数中，并没有太多细节

函数的实现，全部是以 inline 函数的形式，实现在 include/tvm/topi 中，由于算子都是通过 tensor expression 来实现的，因此所有算子的实现最终都会调用 tvm::te::compute 来进行计算

最常见的有4类操作：broadcast，elemwise，nn，reduction

broadcast	elemwise	nn	reduction
topi.add topi.subtract topi.multiply topi.divide topi.floor_divide topi.mod topi.floor_mod topi.maximum topi.minimum topi.power topi.left_shift topi.logical_and topi.logical_or topi.logical_xor topi.bitwise_and topi.bitwise_or topi.bitwise_xor topi.right_shift topi.greater topi.less topi.equal topi.not_equal topi.greater_equal topi.less_equal topi.broadcast_to	topi.acos topi.acosh topi.asin topi.asinh topi.atanh topi.exp topi.fast_exp topi.erf topi.fast_erf topi.tan topi.cos topi.cosh topi.sin topi.sinh topi.tanh topi.fast_tanh topi.atan topi.sigmoid topi.sqrt topi.rsqrt topi.log topi.log2 topi.log10 topi.identity topi.negative topi.clip topi.cast topi.reinterpret topi.elemwise_sum topi.sign topi.full topi.full_like topi.logical_not topi.bitwise_not	topi.nn.relu topi.nn.leaky_relu topi.nn.prelu topi.nn.pad topi.nn.space_to_batch_nd topi.nn.batch_to_space_nd topi.nn.nll_loss topi.nn.dense topi.nn.bias_add topi.nn.dilate topi.nn.flatten topi.nn.scale_shift_nchw topi.nn.scale_shift_nhwc topi.nn.pool_grad topi.nn.global_pool topi.nn.adaptive_pool topi.nn.adaptive_pool3d topi.nn.pool1d topi.nn.pool2d topi.nn.pool3d topi.nn.softmax topi.nn.log_softmax topi.nn.lrn topi.nn.binarize_pack topi.nn.binary_dense	topi.sum topi.min topi.max topi.argmin topi.argmax topi.prod topi.all topi.any

2. topi 编程示例

由于 topi 实现的所有算子，最终都会注册到 tvm 的全局函数中，并已 relay.op 的形式存在

因此，我们在描述计算图的时候，可以直接调用相关的算子，相当于我们平时写 c 代码的时候，有很多编译器的 builtin 函数一样

比如 topi.add 对应的 relay.op 是 relay.add，其他类推

a1 = relay.var("a1", shape=(1,), dtype="float32")
c1 = relay.const(10, 'float32')
c2 = relay.add(c1, a1)

深入浅出 tvm – (13) Relay Pass 之常量折叠

Posted on 2023年6月10日2024年9月22日 by remaper in 深入浅出 tvm Tagged: tvm 源码学习

和算子融合一样，常量折叠是编译领域里最常见的一个优化，简单来说，就是把常量表达式前置计算，在编译阶段就计算好，然后以常量的形式翻译成底层机器码，以提高执行效率，减少计算量

实际上大部分的编译器，常量折叠一般包含2种优化技术：常量折叠和常量传播

1. 基本概念

1.1. 常量折叠

constant folding，常量折叠，编译器优化技术之一，通过对编译时常量或常量表达式进行计算来简化代码。以下面的代码为例：

i = 320 * 200 * 32;

上面的代码中，编译器通常会在编译过程中直接对表达式进行求值，计算出320 * 200 * 32的结果，而不会生成2个乘法指令。

还有一些更复杂（但不清楚tvm是否支持，后面验证下）。比如，在执行一些复杂表达式的计算时，我们可以将表达式内部一些常量运算合并，最终起到简化的效果，如下

优化前（左边）每个表达式的运算量是 8 flop，优化后（右边）的运算量是 2 flop，运算效率极大提升了

不过，实际上 tvm 并没有做的那么高级，tvm 的常亮折叠只用来处理一些比较简单的场景

1.2. 常量传播

constant propagation，常量传播，同样也是编译器最常见的优化技术之一，在编译的过程中，对常量依赖做一些基本的推演和提前计算，再使用常量折叠技术来简化代码。如下：

int x = 14;
int y = 7 - x / 2;
return y * (28 / x + 2);

//常量传播

int x = 14;
int y = 7 - 14 / 2;
return y * (28 / 14 + 2);

//常量折叠

int x = 14;
int y = 0;
return 0;

深入浅出 tvm – (12) Relay Pass 之布局转换

Posted on 2023年6月7日2024年9月22日 by remaper in 深入浅出 tvm Tagged: tvm 源码学习

深入浅出 tvm – (11) Relay Pass 之算子融合

Posted on 2023年6月5日2024年9月22日 by remaper in 深入浅出 tvm Tagged: tvm 源码学习

算子融合，FuseOps，是编译领域最常见的一个优化技术，在这里也算属于 relay 里面最复杂的一类优化了，整个优化的核心逻辑 1k+ 行

代码：src/relay/transforms/fuse_ops.cc

算子融合的目的最终是要解决 AI 处理器的内存墙、并行墙的问题，提升 Tensor 数据的访存局部性。目前算子融合的技术路线有比较多，这里不涉及，我们只需要知道 tvm 是基于支配树来实现算子融合的就行了

1. 基本概念

1.1. 算子融合

算子融合，即将多个算子组合在一起放到同一个核中，通过算子融合的方式，不需要将中间结果保存到全局内存，进而减少执行所需要的时间。

tvm 中将算子分为7种类型：

kElemWise：2个 tensor 之间按元素逐个操作的算子，实际上所有四则运算都是这种类型，https://deeplizard.com/learn/video/QscEWm0QTRY
kBroadcast：见上述链接，到操作两个不同形状的 tensor 时
kInjective：一对一映射函数，比如 add / sqrt / exp 等操作算子（operator）
kCommReduce：多到少的映射，输入到输出具有降维性质，如：sum / max / min等操作操作算子（operator）
kOutEWiseFusable：这是计算比较复杂的，如：conv2d / bn / relu等操作算子（operator）
kTuple：xx
kOpaque：无法被融合的算符，比如 sort

根据以上对算符的不同类型，TVM提供了三种融合规则（我看论文是这么写的，但是现在不止3种了？）：

原文：https://layman-n-ish.github.io/pdfs/TVM_Review_Report.pdf

从融合算子的内部视角看，这种融合实际上是数据计算pipeline化，即两次计算中间数据不再经历store-load过程，而是直接给到下一个计算单元完成计算。

成功，源于对美学的执著追求

作者： remaper

Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache

1. 问题

SpotServe – Serving Generative Large Language Models on Preemptible Instance

1. 背景

1.1. Generative LLM Inference

深入浅出 tvm – (16) Tensor Expression 详解

1. te 的设计目标

使用 ftrace 跟踪内核函数调用栈

Pathways – Asynchronous Distributed Dataflow for ML

1. 论文背景、动机、贡献

深入浅出 tvm – (15) TVM Operator Inventory (TOPI)

1. topi 算子列表

2. topi 编程示例

深入浅出 tvm – (13) Relay Pass 之常量折叠

1. 基本概念

1.1. 常量折叠

1.2. 常量传播

深入浅出 tvm – (12) Relay Pass 之布局转换

深入浅出 tvm – (11) Relay Pass 之算子融合

1. 基本概念

1.1. 算子融合

2025 年 5 月
一	二	三	四	五	六	日
			1	2	3	4
5	6	7	8	9	10	11
12	13	14	15	16	17	18
19	20	21	22	23	24	25
26	27	28	29	30	31

成功，源于对美学的执著追求

1. 问题

1. 背景

1.1. Generative LLM Inference

1. 系统架构

1. te 的设计目标

1. 论文背景、动机、贡献

1. topi 算子列表

2. topi 编程示例

1. 基本概念

1.1. 常量折叠

1.2. 常量传播

1. 基本概念

1.1. 算子融合