ggaaooppeenngg

为什么计算机科学是无限的但生命是有限的

发表于 Disqus：

本文讨论了稀疏注意力的计算方法，其核心思想是通过选择矩阵的分块，将重要的矩阵挑选出来参与计算。这是因为注意力矩阵具有稀疏性，而如何选择这些重点矩阵是各类算法需要解决的主要问题。

稀疏性模式总结

MInference 总结了三种稀疏性模式，这些模式具有动态特性，分别是：

  • A-shape：注意力集中在初始词元及其相近词元上。
  • Vertical-Slash：注意力集中在一些重点词元及其相近词元上。
  • Block-Sparse：注意力具有明显的分块特性。

以下是稀疏性模式的示意图：

稀疏性的动态性从左到右逐渐增强：

MInference 算法

MInference 使用最后 64 个 Q 进行计算，选出 top-k 的垂线和斜线作为重点块的索引。对于 Block-Sparse 模式，使用 mean pool 方法选出 top-k。

FlexPrefill 算法

FlexPrefill 同样使用最后的 64 个 Q 进行 query-aware index selection，计算 QK^T 的 pool。通过 Jensen-Shannon divergence 计算分布距离，如果距离不满足条件，则回退到匹配垂线和斜线模式。

Xattention 算法

Xattention 采用分块和反斜对角线的形式选择块的索引。与仅使用最后一块 Q 进行选择的方式相比，Xattention 允许所有词元参与计算，不依赖垂线和斜线模式的连续性。

反斜对角线的构造

首先对矩阵进行分块，并按 stride 构造反斜对角线。与其说是反斜对角线，更像是一种形似反斜对角线的纹理构造方式：

反斜对角线的优势在于，它可以与垂线和斜线交叉，从而让相关词元参与计算：

选择方式就是通过纹理匹配到的词元注意力进行求和，根据求和结果选择重要矩阵。

发表于 Disqus：

Dynamo 发布以后，我大概速览了一些设计文档，并且提取了一些关键点，并对比一些其他方案的异同点。

Smart Router

worker 上有 KVPublisher 负责发送 kvcache 的创建和删除事件，同时 KvMetricsPublisher 用于发送监控指标（如队列排队长度等）。
router 上则包含 KVIndexer，用于收集 kvcache 的事件并建立前缀树，同时 KvMetricsAggregator 用于收集监控指标。

路由策略基于 KV match rate - Load 的最大值，旨在平衡负载与 kvcache 的匹配度。

KVPublisher 应该是侵入式实现，需要给vLLM打这个patch才能实现，需要修改代码才能捕获这些事件。所以光从他的依赖来看，应该是只支持了vLLM，其他的支持估计还没开源出来。

sgl-router 完全不依赖 worker 的信息，仅通过路由自身的请求实现可过期的前缀匹配。虽然这种方式的匹配精度不如直接获取信息，但实现上更为解耦。

vllm-router 则基于 vLLM 的 Prometheus 接口，通过 /metrics 获取监控指标，其前缀匹配是通过 block hash 的近似度实现的。

llumnix 支持请求的重调度功能，可以将排队中的请求重新分配。

aibrix gateway 同时支持基于树和哈希的匹配方式，并且支持用tokenizer使用 token 进行前缀匹配，而不像 sgl-router 基于字符的匹配。

从 Dynamo 的 Indexer 实现来看，其基于 block 级别的 radix tree，事件通过 Component 的 publish 机制进行分发然后触发radix tree的更新。

条件 PD 分离

并非所有请求的 prefill 阶段都需要在 prefill instance 中计算。如果 prefill 很短，或者 decode instance 的 KV 缓存命中率较高，通常在 decode instance 中直接完成 prefill 更为高效。Dynamo 的分解设计充分考虑了这些场景，并提供了一个灵活的框架，能够在多种条件下实现卓越性能。

在 Decode Instance（在 Dynamo 中称为普通的 worker）上，需要决定是否执行分离操作。如果需要PD分离，则将 prefill 请求交给 prefill worker，通过 prefill queue 进行处理。当 prefill queue 完成后，再通过 prefill queue 将结果传回 worker，开始 decode 阶段。

具体而言，只有在满足以下两个条件时，才会向远程 prefill instance 发送请求：

  1. 没有前缀缓存命中的 prompt 长度超过设定阈值。
  2. Prefill queue 的长度小于设定阈值。

这种条件化的 PD 分离设计，使得 Dynamo 能够在动态工作负载下实现高性能。

Prefill Queue

Prefill Queue 是一个基于 NATS Stream 的全局消息队列。

在这一部分中，最具挑战性的是 KV Cache 的传输。Mooncake 开源了其 TransferEngine，而 vLLM 提供了一些 KV Connector 和 KVStore 的抽象。可以推测 Dynamo 也在 vLLM 的基础上实现了相关功能，可以看到在这个patch中，给vLLM的kv connector实现了一个DynamoNixlConnector。

The key to high-performance disaggregation is efficient KV transfer. Dynamo leverages NIXL to transfer KV cache directly from the VRAM of the prefill engine to the VRAM of the decode engine.

Dynamo 的 KV Cache 传输是通过直接 RDMA（远程直接内存访问）实现的。

为了减少 Memory Descriptors（RDMA 的描述对象）的大小，Dynamo 采用了以下两种优化：

  1. Memory Descriptors 缓存
    每个 Worker（对应传统的 Decode Instance，但在 Prefill 较短时也会执行 Prefill）在初始化并分配所有 KV 缓存池后，会将所有块的 Memory Descriptors（也称为 NIXL 元数据）存储在分布式键值存储 ETCD 中。当 Prefill Worker 第一次服务来自 Worker 的远程预填充请求时，会从 ETCD 加载这些 Memory Descriptors 并缓存到该 Worker 中。因此，在发出 Prefill 请求时，只需要传递 KV 块 ID，而无需传递完整的 Memory Descriptors。这一优化的具体作用可能需要进一步分析 NIXL 的传输过程才能完全理解。

  2. 显存分配优化
    Dynamo 在 Prefill 过程中提升了显存分配能力，通过分配连续的内存块并将其合并为更大的块，从而减少 KV 块的总数。这种合并的具体效果需要结合实现NIXL细节进一步评估。

此外，对于不同 KV 布局（例如由于不同的 TP 导致的 Decode 和 Prefill 布局差异），Dynamo 使用了一个高性能内核。在 NIXL 读取之后和写入之前，该内核会将 KV 块转置为 KV Receiver 中的匹配布局。这可能是为了将 KV Cache 分块传输到不同的 TP 上。

由于引入了 ETCD，Dynamo 支持动态调整 Worker 和 Prefill Worker 的数量。

和其他方案对比

Mooncake 的设计在架构上更加分离，主要通过一个调度器（scheduler）来负责 kvcache 的传输调度，并直接决定 P 和 D 之间的 P2P 传输，基于其 TransferEngine 实现了以下功能：

  1. 基于 kvcache 的前缀匹配分配 prefill 请求
    如果 prefill 节点上缓存了足够的前缀（由 kvcache_balancing_threshold 控制），则选择预估 TTFT（Time to First Token）最小的实例：
    TTFT = min(T_queue + T_prefill)。
    如果 prefill 节点上缓存不足，则选择：
    TTFT = min(T_queue + T_prefill + T_transfer)，
    其中 T_transfer 指的是将最长匹配的 KVCache 从其他实例拷贝到当前实例的预估时间。

  2. 高频使用的 kvcache P2P 传输
    Scheduler 负责 kvcache 的传输调度，例如从一个 prefill 节点传输到另一个 prefill 节点，或者从 prefill 节点传输到 decode 节点。

  3. 基于负载均衡的 decode 请求分配
    通过负载均衡的方式预估 TBT（Time to Best Throughput），从而优化 decode instance 的请求分配。

Mooncake 的设计在模块划分上更加清晰，调度器（scheduler）与各个组件的职责分离明确。

相比之下，Dynamo 的入口在 worker（相当于 Mooncake 中的 decode instance），由 worker 决定是否将 prefill 请求交给 prefill instance。Dynamo 的特点包括：

  • Worker 也可以执行 prefill 操作（即 decode instance 有时也会承担 prefill 的职责）。
  • 引入了全局队列（queue）来处理 kvcache 的计算和计算就绪信息。
  • 提供了 NIXL 传输引擎，但仅支持 P 到 D 的 kvcache 传输，相对实现更为直白。

AIBrix 的现状

AIBrix 目前尚未实现 PD 分离功能，相关文档和白皮书中未提及此功能。

依赖与工程复杂度

从 Dynamo 的依赖项来看，其使用了 ai-dynamo-vllm v0.7.2，这是对 vLLM v0.7.2 的定制化补丁版本，需修改 vLLM 以支持 Publisher 功能。

Dynamo 的工程栈相对复杂，依赖消息队列和 ETCD，但其 PD 分离设计较为直白，例如仅支持 P 到 D 的传输。相比之下，Mooncake 的设计更注重架构分离，尽管目前未实现 offload 功能，但其 P2P kvcache pool 的设计为未来扩展提供了可能性。

关键问题

俗话说得好，关键问题是问题的关键。无论是 Mooncake 还是 Dynamo，其核心目标都是提高传输效率和 kvcache 的利用率。Dynamo 的实现更简化，而 Mooncake 则在架构设计上更具层次感。

KVCache 管理

KVCache Offload

当显存不足时，可以将 KVCache 卸载到更低级别的存储中，例如内存、磁盘，甚至对象存储。
管理器的核心在于结合驱逐策略，在以下两种情况之间取得平衡：

  • 过度缓存：可能引入查找延迟。
  • 缓存不足：导致查找失败和 KV 缓存的重新计算。

V1 单机版本

V1 版本支持将 KVCache 卸载到磁盘，同时使用 CPU 的内存作为缓存。在需要加载时，从磁盘读取数据回显存。

V2 分布式版本

V2 版本将扩展为分布式架构，形成一个全局的 KVCache 池。

Mooncake 的实现

Mooncake 的 KVCache Pool 完全基于显存的 P2P 传输，不涉及 offload 操作。它通过开源的 TransferEngine，将缓存节点上的 KVCache 调度到需要缓存的节点上。

AIBrix 的实现

AIBrix 提供了一个分布式 KVCache Pool，基于 Vineyard 的分布式内存存储。通过 Vineyard 实现 KVCache 的共享，但与专门的传输引擎相比，其传输效率可能稍逊一筹。

NIXL

NIXL 通过简化的同步和批处理以及简化的源和目标抽象简化了数据搬迁。
NIXL 能够在不同类型的内存和快速存储中抽象数据搬迁，而其他数据搬迁库通常只支持一层内存。
这些增强带来了显着的性能提升，加速了第一个词元的时间（TTFT）和整体吞吐量。

NIXL的地位应该是和Mooncake的TransferEngine相当的，至于两者谁的效果更好可能要具体看一下。

总结

看设计的话，感觉还是Mooncake更漂亮一点，层次分得较清楚，不额外依赖什么中间件，kvcache pool的这个设计虽然是纯P2P的，应该后面也可以去做offload之类的。
dynamo就显得更具有工程具体性，并且实现相对来说是要更简单一些，毕竟依赖了message queue又依赖了etcd，把一些复杂度转移给了中间件，入口从worker（or decode instance）可以自己直接prefill短prompt肯定也是做了很多tradeoff才给出了一个不完全分离的条件PD分离的实现。

发表于 Disqus：

参考 OpenAI Triton 主页

参考 Triton论文

参考 GPU MODE Lecture 14: Practitioners Guide to Triton

从Trinton主页引用的话

现代 GPU 的架构大致可以分为三个主要组件 ——DRAM、SRAM 和 ALU—— 在优化 CUDA 代码时必须考虑每个组件：

  1. 来自 DRAM 的内存传输必须合并为更大的事务，以利用现代内存接口的大型总线宽度。
  2. 数据在被再次使用之前，必须手动存储到SRAM中，并且要对数据进行管理，以便在检索数据时尽量减少共享内存存储体冲突的情况。
  3. 计算必须在流式多处理器（SM）之间和内部仔细分区和调度，以促进指令 / 线程级并行性并利用专用 ALU（例如Tensor Core）。

这几句话可能比较抽象，下面给一下这几个组件的指标可能感受更直观，参考Which GPU(s) to Get for Deep Learning: My Experience and Advice for Using GPUs in Deep Learning。
其中指出：

  • 全局内存访问（高达 80GB）：~380 个周期
  • L2 缓存：~200 个周期
  • L1 缓存或共享内存访问（每个流式多处理器高达 128 kb）：~34 个周期
  • 融合乘法和加法，ab+c（FFMA）： 4 个周期
  • Tensor Core 矩阵乘法：1 个周期

每个操作总是由 32 个线程组成的Warp执行，Warp中的线程必须相互等待。GPU 上的所有内存操作都针对warp进行了优化。

根据Simplifying CUDA kernels with Triton: A Pythonic Approach to GPU Programming的说法，GPU中的HBM（High Bandwidth Memory）等价于我们讲的Global Memory，SRAM对应的是L1和L2 Cache对应的是Shared Memory，这几个词在一些文档中可能会有不同的叫法，但是意思是一样的。

A100中的内存带宽约为 2TB/s，L1 缓存带宽：~100-200 TB/s 理论带宽，L2 缓存带宽：~4-7 TB/s 理论带宽。

再看OpenAI的三条说明的意思就是：

  1. 因为DRAM很大，比较容易占满总线带宽，所以尽量合并传输的事务可以减少传输的时间，让高速公路跑满。
  2. 如果数据要重复利用，反复参与计算，尽量让他们在SRAM当中能够缓存住，比如L1的读取只要34个cyle，能比从L2中快6到7倍。
  3. 尽量跑满并行度，并且利用更高效的计算单元，比如Tensor Core。

这个是OpenAI给出的GPU架构的简图，我们需要明确不同内存，缓存，和执行单元的周期之间的关系就比较好理解GPU计算当中的性能瓶颈。

Triton的目标其实就是优化 HBM -> SRAM -> 寄存器 的带宽，这在Torch里面直接实现不了，通过一些融合算子是可以减少写回到HBM的。

Triton的文档给出的很多实现的代码，可能都不太奏效了，笔者自己测试下来并没有超过torch本身的实现，
可能torch本身也再不断改进吧，这些差别很快就超越了，但是在一些写自定义融合算子方面应该还是比较有优势的。

Triton的使用

和CUDA对应的关系：

  • 程序（Program）：处理数据块的kernel实例。
  • PID（程序 ID）：等同于 CUDA 中的块 ID。
  • 向量化操作：在多个数据点上同时操作（triton不需要用户关心向量操作的并行化）。

先给出变量和修饰器的解释，大部分文档都混在注释里面不是很好阅读，我觉得先介绍一些简单概念再看代码会比较好一点。

@triton.jit 装饰器表示这个函数会被编译。
tl.constexpr 代表常量表达式，可以让编译器在编译期间直接求值，可以当作常量使用了。
BLOCK_SIZE对于GPU来说是比较固定的，因为一个block是有threads数上限的。
通过执行cuda-samples中的deviceQuery
可以发现L40的显卡BLOCK_SIZE最大是1024，大部分显卡应该都是这个固定大小。 

1
2
3
4
5
6
7
8
9
10
Total number of registers available per block: 65536  
Warp size:                                     32  
L2 Cache Size:                                 100663296 bytes（96MB）
Maximum number of threads per multiprocessor:  1536  
Maximum number of threads per block:           1024  
Max dimension size of a thread block (x,y,z): (1024, 1024, 64)  
Max dimension size of a grid size    (x,y,z): (2147483647, 65535, 65535)  
Total amount of shared memory per block:       49152 bytes（48KB）
Total shared memory per multiprocessor:        102400 bytes (100KB)
Total number of registers available per block: 65536

pid = tl.program_id(axis=0) 应该是对应的CUDA中的threadIdx.x的作用，对应block的一维下标，
pid = tl.program_id(axis=1) 应该是对应的CUDA中的threadIdx.y的作用，对应block的二维下标。

autotune是一个黑盒优化，通过内部的小benchmark的方式去基于key的变量，优化configs里面的参数。
下面是一个只有两个配置的搜索空间，当n_elements的值发生变化的时候，会自动选择最优的配置。

1
2
3
4
5
@triton.autotune(configs = [
    triton.Config({'BLOCK_SIZE': 128}, num_warps = 4, pre_hook = clone_inplace_updated_params),
    triton.Config({'BLOCK_SIZE': 1024}, num_warps = 8, pre_hook = clone_inplace_updated_params),
], key = ['n_elements'])
@triton.jit

BLOCK_SIZE表示的是一个program负责的BLOCK大小，放在triton的语境下更像是L2 Cache的大小。
但是cuda当中的block是包含n个thread的，表示的是并行线程的大小。
笔者的这个说法可能不太精准，但是这两种风格导致Cuda写一些element wise的操作比较合适.
每个element wise的操作都是一个thread，这样可以充分利用GPU的并行性。
而triton比较适合一些Reduce操作，例如对数据（也就是矩阵）切BLOCK，然后每个kernel去负责一个block，
他的好处就是比如softmax这样的在行上做reduce操作会比较直观，而矩阵乘法也可以沿着MxK,KxN的维度，沿着不同的维度切块。
Triton能够帮你把矩阵乘法优化得很不错，虽然可能还比不上精准手写的Cuda算子。

Triton的范式和CUDA的Single Instruction, Multiple Thread (SIMT)不一样，官网给出了一个简化的例子。

这是CUDA like的写法，每个threadId.x代表的线程只算一个element。

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
BLOCK = 512

# This is a GPU kernel in Numba.
# Different instances of this
# function may run in parallel.
@jit
def add(X, Y, Z, N):
   # In Numba/CUDA, each kernel 
   # instance itself uses an SIMT execution
   # model, where instructions are executed in
   # parallel for different values of threadIdx
   tid = threadIdx.x
   bid = blockIdx.x
   # scalar index
   idx = bid * BLOCK + tid
   if id < N:
     # There is no pointer in Numba.
     # Z,X,Y are dense tensors
     Z[idx] = X[idx] + Y[idx]


...
grid = (ceil_div(N, BLOCK),)
block = (BLOCK,)
add[grid, block](x, y, z, x.shape[0])

Triton文档中的Matrix乘法简化来说就是并行计算M*N个block（沿着K所代表的维度）。
这是Triton的写法，每个Program负责了一个block，他就少了一个block的切分维度：

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
BLOCK = 512

# This is a GPU kernel in Triton.
# Different instances of this
# function may run in parallel.
@jit
def add(X, Y, Z, N):
   # In Triton, each kernel instance
   # executes block operations on a
   # single thread: there is no construct
   # analogous to threadIdx
   pid = program_id(0)
   # block of indices
   idx = pid * BLOCK + arange(BLOCK)
   mask = idx < N
   # Triton uses pointer arithmetics  
   # rather than indexing operators
   x = load(X + idx, mask=mask)
   y = load(Y + idx, mask=mask)
   store(Z + idx, x + y, mask=mask)


...
grid = (ceil_div(N, BLOCK),)
# no thread-block
add[grid](x, y, z, x.shape[0])

笔者比较疑惑，单就这两个代码他们的并行度貌似是不一样的，难道是把block那一层隐式的放在了load和store当中，他的load和store其实是隐含了并行能力的。
援引知乎的文章，Triton中一直是以Block为中心来计算，直到Lowering到LLVM和PTX才会转为Thread为中心的计算，而这些对于使用Block抽象进行编程的用户来说都是无感的。是符合笔者预期的，triton简化了CUDA的写法，block具体的线程数的，每个线程处理多少元素，triton自己会去帮你处理。

当使用 triton 的时候,x = tl.load(x_ptr + offsets, mask=mask)时，我们正在加载到 L2 缓存 或者叫 SRAM 中。

根据Torch的blog，以及参考 OpenAI/Triton MLIR 迁移工作简介，Triton编译的过程是@triton.jit装饰器通过遍历提供的 Python 函数的抽象语法树（AST）来工作，以便使用通用的 SSA 构造算法即时生成 Triton-IR。
然后，生成的 IR 代码被我们的编译器后端简化、优化和自动并行化，最后被转换成高质量的 LLVM IR，最终是 PTX（Nvidia GPU的汇编），可以在最近的Nvidia GPU上执行。

矩阵乘法

这段代码是基于K切BLOCK，比上面的代码要好理解一点。

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
32
33
34
35
36
@triton.jit
def simple_mm(a, b, o, k, n,
              K_BLOCK_SIZE: tl.constexpr = 64,
              ) -> None:
    # a -> Matrix of size M x K and b -> Matrix of size K x N
    # K is the common inner dimension
    num_blocks = k//K_BLOCK_SIZE + 1
    row_id = tl.program_id(0)
    col_id = tl.program_id(1)

    # Lets pick one column and one row and do a dot product
    # Like the 1-D example we dont want to look at the entire row/column
    # We are making use of the fact that each row/column will be of the size
    # 'k' which is the inner common dimension of these matrices
    # But this will only be a part of the dot product so we have to keep track of many to cover the entire column or row.
    
    # What we are going to do is to access block size elements from the column
    # and the row and compute the dot product and keep adding to a value till
    # we run out of numbers
    value = 0.
    for k_id in range(num_blocks):
        row_start = row_id * k + k_id * K_BLOCK_SIZE
        row_offsets = tl.arange(0, K_BLOCK_SIZE) + row_start
        # The masks are a little more trickier as we cant just see if its
        # less than 'k'. We need to account for the row we are in
        row_masks = row_offsets < (row_id + 1) * k
        row = tl.load(a + row_offsets, mask=row_masks) # Load this into the GPU SRAM

        col_start = (K_BLOCK_SIZE * k_id)
        col_offsets = n * (tl.arange(0, K_BLOCK_SIZE) + col_start) + col_id # 0, n, 2n || 3n, 4n, 5n for a block size of 3 for eg
        col_masks = col_offsets/n < k
        col = tl.load(b + col_offsets, mask=col_masks)
        value += tl.sum(row * col)
    
    output_offset = row_id * n + col_id
    tl.store(o + output_offset, value)

BLOCK_SIZE 和 GROUP_SIZE 的优化。

一次计算的时候尽量用满L2 Cache，所以可以把多个BLOCK放到一个GROUP里面，这个GROUP变成了grid的切分，但在GROUP里面我们
再去做BLOCK级别的计算，要计算好对应的线性空间中的stride。

Softmax

下面的代码只能在n_cols小于BLOCK_SIZE的数据上运行。

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
32
33
34
35
36
import triton
import triton.language as tl

@triton.jit
def softmax(Y, stride_ym, stride_yn, X, stride_xm, stride_xn, M, N):
    # row index
    m = tl.program_id(0)
    # col indices
    # this specific kernel only works for matrices that 
    # have less than BLOCK_SIZE columns
    BLOCK_SIZE = 1024
    n = tl.arange(0, BLOCK_SIZE)
    # the memory address of all the elements
    # that we want to load can be computed as follows
    X = X + m * stride_xm + n * stride_xn
    # load input data; pad out-of-bounds elements with 0 
    x = tl.load(X, mask=n < N, other=-float('inf'))
    # compute numerically-stable softmax
    z = x - tl.max(x, axis=0)
    num = tl.exp(z)
    denom = tl.sum(num, axis=0)
    y = num / denom
    # write back to Y
    Y = Y + m * stride_ym + n * stride_yn
    tl.store(Y, y, mask=n < N)

import torch
# Allocate input/output tensors
X = torch.normal(0, 1, size=(583, 931), device='cuda')
Y = torch.empty_like(X)
# SPMD launch grid
grid = (X.shape[0], )
# enqueue GPU kernel
softmax[grid](Y, Y.stride(0), Y.stride(1), 
              X, X.stride(0), X.stride(1),
              X.shape[0]    , X.shape[1])

发表于 Disqus：

从一些大模型的训练技术报告来看有一些比较有代表性的挑战，比如 Meta 的 Research Super Compute (RSC) 和 X 的 Grok Infra。这些技术报告中提到了一些关键的技术挑战和解决方案，包括 GPU 架构与互联、存储系统、训练的稳定性等。

X Grok Infra

从 Grok-1.5 Infra 的技术报告中可以窥见，Grok-1.5 在基础设施方面具有以下核心优势：

  1. 先进的分布式训练框架：基于 JAX、Rust 和 Kubernetes 的技术栈，不仅确保了高性能，还能快速适配和训练新的模型架构。
  2. 卓越的可靠性和可用性：通过自研的训练协调器，系统能够智能地检测并隔离故障节点，大幅降低训练任务中断的风险。
  3. 高效的存储与数据处理：在检查点存储、数据加载和训练作业重启等环节都进行了深度优化，将训练过程中的停机时间降至最低。

Meta Reasearch Super Compute

另一个典型案例是 Meta 的 Research Super Compute (RSC) 超算集群，在这上面训练了Llama3.2，有一份92页的技术报告，RSC的相关Talk，以及里面用到的MAST论文调度器：

算力规模

已升级至 16,000 张 H100 GPU，算力获得质的飞跃。每个服务器配备了 8 块 GPU 和 2 块 CPU。在服务器内部，八块 GPU 通过 NVLink 连接。

网络互联

采用双网络方案：

  • NVIDIA Quantum InfiniBand，带宽高达 1600 Gb/s，RoCE（RDMA over Converged Ethernet）作为补充互联方案。

网络拓扑

  • 底层网络（第一个层）：每个机架（rack）包含 16 块 GPU，分散在两个服务器上，并通过一个 Minipack2 顶层网络（ToR）交换机连接。
  • 中间网络（第二层）：192 个这样的机架通过 Cluster Switches 连接，形成一个包含 3,072 块 GPU 的 Pod。这种设计确保了从任何两个 GPU 之间的通信都有满速带宽，没有过度订阅。
  • 顶层网络（第三层）：八个这样的 Pod 通过 Aggregation Switches 连接，形成一个包含 24,000 块 GPU 的集群。然而，顶层网络的连接没有保持满速带宽，而是存在过度订阅比例为 1：7。

负载均衡

  • Collective library 将 16 个网络流中的两个 GPU 之间的数据传输从一个流变为 16 个流。
  • Enhanced-ECMP（E-ECMP）协议 通过在 RoCE（Rdma over Converged Ethernet）报头中添加额外的字段，进行 hash 计算，从而有效地在不同网络路径上平衡 16 个流。

拥塞控制

使用深度缓冲区（deep-buffer switches）来解决在 Spine（Gangidi et al., 2024）中由于集体通信模式引起的暂时拥堵和缓冲问题。

存储系统

采用自研的 Tectonic 文件系统，通过 FUSE 提供标准的 Linux 文件系统接口，确保高效的数据访问。

  • 存储容量：240 PB，基于 7,500 台 SSD servers
  • 支持的最大吞吐量：7 TB/s
  • 支持的可持续吞吐量：2 TB/s
  • 检查点写入：非常时断时续，导致存储网络饱和
  • 检查点的目标：因为 checkpoint 非常大，最小化 GPU 停顿时间，加快检查点频率也变得非常重要

总结

从这些实践可以看出，现代 AI 基础设施主要围绕三大核心要素展开：

  • 计算能力（以 GPU 为核心）
  • 网络互联（RoCE 或 InfiniBand）
  • 存储系统

而在上层的编排调度领域，系统的容错能力和可靠性则成为关键考量因素。

GPU 架构与互联

在当前AI训练领域，主流的GPU型号主要是NVIDIA的A100、H100和H200系列，它们按照发布时间依次提供了更强大的算力和更优化的架构设计。关于GPU的详细架构，特别是其拓扑结构，可以参考这篇深度解析文章。

GPU互联拓扑

GPU之间的互联拓扑结构主要取决于不同总线间的传输特性，GPU之间可以通过NVIDIA专有的NVLink高速互联技术直接通信。在现代GPU集群中，主要有以下几种互联方式：

  1. NVSwitch架构：通过NVIDIA的交换架构实现所有GPU之间的全互联
  2. 走网卡，如果卡之间没有NVSwitch的话，可以绕过CPU走网卡：
    1
    GPU0 -> PCIe -> IB(InfiniBand) -> PCIe -> GPU1
    这种通信模式由NCCL（NVIDIA Collective Communications Library）负责协调和优化。

GPU分配策略

NVIDIA开源的go-gpuallocator库提供了一系列基于拓扑关系的GPU分配策略。例如，其中的NewStaticDGX1Policy专门针对DGX-1标准配置优化。考虑到单机环境下GPU组合的可能性有限，这种基于静态规则的分配策略已经能够很好地满足需求。

这些分配策略的核心目标是最小化跨总线和跨NUMA节点的通信开销，确保GPU间通信尽可能利用最高带宽的数据通路，从而提供最优的训练性能。

跨节点的通信

在分布式训练场景下，跨节点通信需要经过更长的数据传输路径：

1
GPU -> NIC -> 叶层交换机 -> 核心交换机 -> NIC -> GPU

这种通信模式面临两个主要的优化方向：

  1. 本地化优化：尽可能将相关联的GPU任务分配在物理位置相近的节点上，以减少网络延迟。

  2. 负载均衡：避免将所有任务集中在同一交换机下，防止出现网络拥塞。过度集中可能导致局部带宽饱和，反而降低整体训练效率。

这种权衡本质上是一个网络流优化问题。通过图论中的网络流算法，可以在通信延迟和带宽利用率之间找到最优平衡点，从而实现更高效的跨节点通信。

一个分布式训练的带宽瓶颈来源于带宽最低的那条路径。

利用 Kubernetes Pod 亲和性优化网络拓扑

在 Kubernetes 环境下，我们可以通过 Pod 亲和性（Affinity）和规则来优化 GPU 任务的分配。主要可以从以下几个方面入手：

拓扑感知调度：使用 topologyKey 确保相关联的 Pod 被调度到网络拓扑上接近的节点：
例如同一个分布式训练任务(training-group = group1)尽让分配在一个机架上，同交换机，同核心交换机也是类似的。

1
2
3
4
5
6
7
8
9
10
11
12
affinity:
  podAffinity:
    preferredDuringSchedulingIgnoredDuringExecution:
    - weight: 50
      podAffinityTerm:
        labelSelector:
          matchExpressions:
          - key: training-group
            operator: In
            values:
            - group1
        topologyKey: topology.kubernetes.io/rack  # 同机架优先

这种方案的优势在于：

  • 配置简单，易于理解和维护
  • 充分利用 Kubernetes 原生能力，无需额外组件
  • 可以根据实际需求灵活调整权重和策略

存储系统

AI训练中的存储系统面临着两个主要挑战：

1. 海量小文件问题

AI训练数据集通常包含大量的小文件，这对传统文件系统的性能和管理造成了巨大压力。一些现代分布式文件系统提供了很好的解决方案，例如 Meta 的 Tectonic 和与其架构类似的 JuiceFS，它们采用了以下优化方案：

元数据管理优化：

  • 使用元数据库管理文件结构，将 ls 命令转化为简单的字符串前缀匹配操作
  • 避免了传统 Linux 文件系统依赖 inode 管理的方式
  • 解决了 inode 臃肿问题（在传统系统中，一个 inode 的大小可能与文件本身相当）

2. Checkpoint 存储挑战

分布式训练中的 checkpoint 文件体积巨大，这在大语言模型训练中尤为明显：

  • 以 LLaMA-2-70B 为例，单个完整的 checkpoint 就需要 140GB 存储空间（FP16格式）
  • 训练过程中需要定期保存 checkpoint，累积存储需求可能达到 TB 甚至 PB 级别
  • 需要存储系统能够提供高带宽和低延迟的读写性能，同时保证数据的可靠性

这些挑战要求存储系统具备：

  • 强大的扩展性
  • 高效的数据压缩能力
  • 智能的数据分层存储机制
  • 可靠的数据备份和恢复能力

训练的稳定性

在大规模 AI 训练中，硬件故障是一个常见问题。特别是新型号显卡往往会有较高的故障率，再加上传统的硬件错误，这些都可能导致训练中断。因此，快速识别错误并恢复训练成为了一个关键挑战。目前主流的解决方案主要有以下两种：

基于 torchrun 的弹性训练

torchrun 提供了两种容错机制：简单重试和弹性训练。

  1. 简单重试模式：
    通过 --max-restarts 参数配置重试次数：

    1
    2
    3
    4
    5
    6
    7
    8
    torchrun \
        --nnodes=$NUM_NODES \
        --nproc-per-node=$NUM_TRAINERS \
        --max-restarts=3 \
        --rdzv-id=$JOB_ID \
        --rdzv-backend=c10d \
        --rdzv-endpoint=$HOST_NODE_ADDR \
        YOUR_TRAINING_SCRIPT.py [script args...]

  2. 弹性训练模式：
    通过设置 nnodes 的范围来支持动态节点数：

    1
    2
    3
    4
    5
    6
    7
    8
    torchrun \
        --nnodes=1:4 \  # 支持1-4个节点的动态伸缩
        --nproc-per-node=$NUM_TRAINERS \
        --max-restarts=3 \
        --rdzv-id=$JOB_ID \
        --rdzv-backend=c10d \
        --rdzv-endpoint=$HOST_NODE_ADDR \
        YOUR_TRAINING_SCRIPT.py [script args...]

弹性训练模式需要配置服务发现机制，默认使用 c10d 作为内置的节点发现服务，也支持使用 etcd 等外部服务。

当节点发生变化时，系统会自动处理以下场景：

  • 节点离开（缩容）：系统通知 agent，停止现有 workers，重新组建 WorkerGroup，使用新的 RANK 和 WORLD_SIZE 启动所有 workers
  • 节点加入（扩容）：接纳新节点，按照相同流程重组 WorkerGroup

基于 DeepSpeed 的弹性训练

DeepSpeed 提供了更细粒度的弹性训练配置：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
{
  "elasticity": {
    "enabled": true,
    "max_train_batch_size": "seqlen",
    "micro_batch_sizes": 8,
    "min_gpus": 1024,
    "max_gpus": "fixed_linear",
    "min_time": "seqlen",
    "version": 8,
    "ignore_non_elastic_batch_info": 1024,
    "num_gpus_per_node": "fixed_linear",
    "model_parallel_size": MODEL_PARALLEL_SIZE
  }
}

DeepSpeed 的特点是：

  • 支持动态调整 batch size
  • 以 GPU 为粒度进行弹性伸缩（而不是节点级别）
  • 提供更丰富的训练参数配置

弹性训练控制器

要实现完整的弹性训练支持，控制器需要：

  1. 依赖服务发现机制进行节点注册和健康检查
  2. 动态调整弹性策略（如 min_nodes、max_nodes 等参数）

对于简单的降级场景，通过静态配置即可实现：

  • 将 max_nodes 设置为总资源规格
  • 将 min_nodes 设置为最小运行要求（如设置为 1:4 表示支持 1-4 张显卡的动态伸缩）

节点的问题发现

在大规模语言模型（LLM）预训练过程中，常见的硬件异常包括：

  1. GPU ECC 错误：当 GPU 发生不可纠正的显存 ECC（Error Correcting Code）错误时，通常需要重置 GPU 或重启节点来清除这个错误。

  2. Infiniband（IB）/NCCL 问题：这类问题通常源于硬件故障，如网卡损坏或网络抖动，可能导致训练速度下降或任务异常中断。

  3. 任务挂起（Hang）：通常与 IB/NCCL 问题相关，需要人工检测和处理。

  4. GPU 掉卡：此时一般会触发 CUDA 错误或程序异常退出，可能需要重置 GPU 或重启节点来解决。

  5. 机器异常：包括 GPU 之外的硬件异常，如硬盘、CPU 等，甚至整机故障，可能需要更换硬件或进行系统维护。

  6. 机器配置异常：例如，某台机器意外启用了 MIG（多实例 GPU），可能影响训练任务的正常运行。

  7. 集群维护：集群中的其他任务或系统维护、升级，可能需要暂停当前训练任务。

可以使用node-promblem-detector，
node-problem-detector 是一个用于在集群管理栈的上游层次中使各个节点问题可见的守护进程。它在每个节点上运行，检测节点问题并将其报告给 apiserver。

监控和容错是一个比较难的问题，需要结合硬件和软件的特性，以及业务需求，进行综合考量。
特别是万卡集群，MFU 只有 50%左右。

在训练 OPT-175B 模型的过程中，Meta团队使用了 992 个 80GB 的 A100 GPU，每个 GPU 实现了约 147 TFLOP/s 的性能，对应的机器浮点利用率（MFU）约为 47%（147/312）。

为了应对可能的硬件故障，团队额外准备了 12 台备用机器，以便在出现问题时进行替换。在训练期间，平均每天约有 2 台机器发生故障，即每台机器每天发生故障的概率约为 1.61%。

整个训练过程持续了约 2 个多月，包括从 2021 年 10 月 20 日到 2021 年 11 月 11 日的测试阶段，以及从 2021 年 11 月 11 日到 2022 年 1 月 6 日的正式训练阶段，正式训练约 57 天。

根据预估，实际训练时间应为约 25 天，但由于各种问题，实际有效训练时间仅占总时间的约 44%。在前期，由于各种问题，团队至少手动重启了 35 次任务。为减少人工干预，后续引入了自动重启机制，但由于硬件故障，仍触发了超过 70 次重启，平均每天需要重启一次任务。

这些经验表明，在大规模模型训练中，硬件故障和其他问题会显著影响训练效率。为此，团队采取了多种措施来应对这些挑战，包括准备备用硬件、引入自动重启机制等，以确保训练过程的顺利进行。

这个问题在用新的卡的时候会有更多问题。

总结

To train our largest Llama 3 models, we combined three types of parallelization: data parallelization, model parallelization, and pipeline parallelization. Our most efficient implementation achieves a compute utilization of over 400 TFLOPS per GPU when trained on 16K GPUs simultaneously. We performed training runs on two custom-built 24K GPU clusters. To maximize GPU uptime, we developed an advanced new training stack that automates error detection, handling, and maintenance. We also greatly improved our hardware reliability and detection mechanisms for silent data corruption, and we developed new scalable storage systems that reduce overheads of checkpointing and rollback. Those improvements resulted in an overall effective training time of more than 95%. Combined, these improvements increased the efficiency of Llama 3 training by ~three times compared to Llama 2.

LLAMA3 的技术博客揭示了许多令人振奋的优化成果，这些优化背后蕴含着大量值得深入研究的技术细节。虽然我们可能难以直接接触如此大规模的训练集群及其面临的挑战，但这些技术进展仍然为整个 AI 基础设施领域提供了宝贵的参考和启发。

发表于 更新于 Disqus：

参考 Benchmarking Text Generation Inference。

参考 SGLang issue 364。

参考 LLM inference server performances comparison llama.cpp / TGI / vLLM。

相关代码：

sglang bench。

vLLM bech prefix cache。

vLLM bench serving。

TokenAttention和PagedAttention，感觉TokenAttention是个很离谱的设计，而Radix的话和PagedAttention的颗粒度不是完全对应的。

vLLM 的默认block size最多是32，虽然这个32对应的字符串长度不是固定的，一般一个Token平均对应4个字母，所以有效前缀大概120比较合适。

最近的vLLM production stack的RFC准备采用分block的trie树或者直接使用simHash，这些可能更匹配vLLM本身的实现

包括最近字节推出了aibrix也提供了相关能力。

前缀重复度

为了能够测试不同数据集的前缀重复度，需要一种方法衡量对话的前缀重复度，如果前缀的重复度不高，可能测试结果不太能体现前缀缓存的优势。

对于所有的对话构造一个Radix树，每个树节点保存一个计数器记录经过该节点的字符串的数量。

计数重复前缀的数量，比如W这个前缀是比较多的因为很多英文问句都是Wh-开头的，而中文的话是比较随机的。

对于每个节点，在进行计数器过滤的时候，要一直遍历到某个节点的子节点都小于计数器N才结束，这样防止过滤出多个公共前缀的前缀，
因为较短的前缀肯定是被较长的前缀包含的。相当于对这棵树做剪枝，删除所有计数器小于过滤值的节点。

再从满足要求的所有被剪枝完的叶子结点中选择长度大于L的前缀。

对话数据集的前缀重复度 = 基于N剪枝的所有长度大于L的叶子前缀节点数 / 所有对话数量

压力测试数据集

  • databricks-dolly-15k 这个数据集的前缀重复度不高。
    只有两个前缀长度超过00，重复次数大于1，因为里面都是单轮的对话。
    (‘Extract all of the dates mentioned in this paragraph and list them using bullets in the format {Date} - {Description}’, 11) (‘Extract all of the names of people mentioned in this paragraph and list them using bullets in the format {Name}’, 15)
  • LMSYS-CHAT-1M
    一个parquet有16W个对话。前缀重复比较高的是30~40次。这样的对话有9483条，也就占总数的5%，重复前缀的平均长度只有300左右。
  • ShareGPT这是vLLM官方使用的一个压测数据集。压测脚本在这。这个的比重也只有2%，重复前缀的平均长度是4K。

以上数据集可能对于前缀缓存的优势体现不太明显。

  • 测试工具

    • sglang inference benchmark

  • 测试参数

    • batch_size: 30
    • max_length: 4096
    • num_samples: 1000

  • 测试结果

    • TTFT
    • TBT
    • Throughput

构造数据集

用实际的数据集结果不是特别好，差异度不是很高，因为这些数据集的前缀重复度比重都不是很高。
没有特别好的现成的数据集，需要使用人工构造的方式去构造数据集。

sglang 的benchmark提供了 generated-shared-prefix dataset arguments相关的参数。
他是通过随机生成一个系统提示词再组合问题，但是Prompt是随机的。语言不是很明朗。但可能并不
影响测试效果。

比较理想的应该是认为构造一些长度的系统提示词加一些问题进行组合，这个可读性会更高一点，但是没那么灵活
不太好按要求生成指定上下文长度的提示词。

测试结果

结果来看，在batch size更大的情况下，TTFT会变得特别长，而TBT也会相应的增加一些但没有TTFT恐怖。
batch size变大以后，TTFT从300s变成了900s，而ITL则从0.2s变成了0.3s。
这和MoonCacke的论文是一致的。

测试一下PD分离的效果，使用vLLM的1P1D。
PD分离以后TTFT可以降低一个数量级，这个效果还是很明显的，直接降了一个数量级。

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
============ Serving Benchmark Result ============
Backend:                                 vllm
Traffic request rate:                    inf
Max reqeuest concurrency:                not set
Successful requests:                     47
Benchmark duration (s):                  127.03
Total input tokens:                      14545
Total generated tokens:                  2993
Total generated tokens (retokenized):    2992
Request throughput (req/s):              0.37
Input token throughput (tok/s):          114.50
Output token throughput (tok/s):         23.56
Total token throughput (tok/s):          138.06
Concurrency:                             24.49
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   66177.90
Median E2E Latency (ms):                 61336.75
---------------Time to First Token----------------
Mean TTFT (ms):                          39888.70
Median TTFT (ms):                        22421.85
P99 TTFT (ms):                           116090.20
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          491.86
Median TPOT (ms):                        394.97
P99 TPOT (ms):                           1917.39
---------------Inter-token Latency----------------
Mean ITL (ms):                           419.69
Median ITL (ms):                         275.52
P99 ITL (ms):                            1766.40
==================================================

双v100 LLAMA3.2:11b

python -m sglang_router.launch_router –worker-urls http://127.0.0.1:8081 http://127.0.0.1:8082

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
============ Serving Benchmark Result ============
Backend:                                 vllm
Traffic request rate:                    inf
Max reqeuest concurrency:                not set
Successful requests:                     1000
Benchmark duration (s):                  1247.16
Total input tokens:                      289255
Total generated tokens:                  184429
Total generated tokens (retokenized):    184388
Request throughput (req/s):              0.80
Input token throughput (tok/s):          231.93
Output token throughput (tok/s):         147.88
Total token throughput (tok/s):          379.81
Concurrency:                             470.04
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   586218.50
Median E2E Latency (ms):                 596155.97
---------------Time to First Token----------------
Mean TTFT (ms):                          520113.99
Median TTFT (ms):                        526194.47
P99 TTFT (ms):                           1067230.41
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          363.05
Median TPOT (ms):                        356.14
P99 TPOT (ms):                           736.93
---------------Inter-token Latency----------------
Mean ITL (ms):                           360.61
Median ITL (ms):                         273.54
P99 ITL (ms):                            1525.31
==================================================

双卡的并发的情况下，吞吐可以线性增长，但是相较于1P1D来说，prefill的时间没有改善。

多机器配置

DeepSeek R1 8xH20 x2 台机器，每台机器RDMA配置16个 MT2910 Family [ConnectX-7] 做8个bond。

8TP x 2PP 的部署方案，如果后面EP支持的话可能会有更好的效果。

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
============ Serving Benchmark Result ============
Backend:                                 vllm
Traffic request rate:                    inf
Max reqeuest concurrency:                not set
Successful requests:                     1000
Benchmark duration (s):                  234.47
Total input tokens:                      303481
Total generated tokens:                  187870
Total generated tokens (retokenized):    186116
Request throughput (req/s):              4.26
Input token throughput (tok/s):          1294.33
Output token throughput (tok/s):         801.26
Total token throughput (tok/s):          2095.59
Concurrency:                             363.04
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   85122.29
Median E2E Latency (ms):                 82826.18
---------------Time to First Token----------------
Mean TTFT (ms):                          31789.26
Median TTFT (ms):                        17669.77
P99 TTFT (ms):                           100110.92
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          770.73
Median TPOT (ms):                        341.77
P99 TPOT (ms):                           9445.55
---------------Inter-token Latency----------------
Mean ITL (ms):                           284.74
Median ITL (ms):                         214.68
P99 ITL (ms):                            745.14
==================================================

sglang tp 16的配置，sglang不支持pp，sglang明显要快一些，主要原因应该是sglang支持了MTP，vLLM目前还没有。

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
============ Serving Benchmark Result ============
Backend:                                 sglang
Traffic request rate:                    inf
Max reqeuest concurrency:                not set
Successful requests:                     1000
Benchmark duration (s):                  190.92
Total input tokens:                      306113
Total generated tokens:                  197108
Total generated tokens (retokenized):    195033
Request throughput (req/s):              5.24
Input token throughput (tok/s):          1603.38
Output token throughput (tok/s):         1032.43
Total token throughput (tok/s):          2635.81
Concurrency:                             488.50
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   93263.23
Median E2E Latency (ms):                 86230.17
---------------Time to First Token----------------
Mean TTFT (ms):                          39722.57
Median TTFT (ms):                        43590.80
P99 TTFT (ms):                           60010.86
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          1529.43
Median TPOT (ms):                        270.69
P99 TPOT (ms):                           37619.47
---------------Inter-token Latency----------------
Mean ITL (ms):                           276.88
Median ITL (ms):                         158.45
P99 ITL (ms):                            945.60
==================================================

发表于 更新于 Disqus：

CacheBlend的主要目标是在一些RAG场景下，多个文档Chunk之间不能像多轮对话那样构成Prefix Cache。

对于位置编码来说，RoPE得到的注意力是绝对位置无关的，所以两个下三角放到对应的位置就可以。但是，两个Chunk之间的交叉注意力机制实际上是空的，如果单纯这样使用会丢失交叉注意力的信息。

论文中提到了一个例子。

在比较两个球员进球数的场景中，就损失了球员之间的交叉信息。

为了弥补下面那个空的矩形的注意力，同时尽量节省计算量，根据一些insights提出了一种选择性计算的方法。

注意力矩阵是稀疏的，其中差异较大（颜色较深）的部分对最终结果贡献较大。另外，对于多层的Transformer，前几层的注意力差异往往会一直保持下去。因此，CacheBlend会完全重算第一层的注意力，然后标记差值较大的token，用它们的注意力来代表两个Chunk之间的完全注意力矩阵。当然，这还是有损的。

其中的权衡在于选择多少比例的token来代表两个Chunk之间的注意力矩阵，这个比例是可以调整的。

一个平衡点在于从异构存储中加载KVCache和重计算KVCache的时间。如果选择性计算token的时间大于加载的时间，则可以将这个过程流水线化。在计算当前层时，加载下一层的KVCache。

比率r%刚好满足计算时间和加载时间相等。

尽管第一层高差异化的token的注意力比较重要，但只看第一层似乎不太合理。CacheBlend会用一个比r%更大的范围选择token，然后每一层逐渐递减r%，以容许更多的可能性。只从第一层选择r%可能会丢失一些重要信息。

总体来说，利用交叉注意力的稀疏性选择性重计算Chunk之间的交叉注意力，并平衡加载和计算，使得计算过程和加载是并行的，时间损耗无影响。在精度上，选择一个较大的r%再每层缩小到理想r%来满足准确性。最终实验效果显示，与完全重算的交叉注意力相比，结果还是很接近的。

发表于 更新于 Disqus：

vLLM的PD分离

vLLM的PD分离是指vLLM的Prefill和Decode分离到不同的实例中执行。

新增配置

新增 KVTransferConfig 配置，决定了实例的类型。如果是 prefill 则 role 为 producer，如果是 decode 则 role 为 consumer，并且要设定传输的方法。

  • is_kv_transfer_instance 判断是否是 PD 分离的实例。

代码实现

在 ./vllm/worker/model_runner.py 中：

  1. 在计算之前执行 need_rev_kv，检查是否是 consumer，且当前 run 是不是 prefill。然后调用 get_kv_transfer_group().recv_kv_caches_and_hidden_states。
  2. 在计算之后执行 need_send_kv，检查是否开启配置 producer，并且当前 run 是不是 prefill（对于以前未分离的结构来说，decode实例要经历prefill阶段，
    但是prefill已经被prefill实例做掉了，所以要等着接受prefill的KVCache，不需要重复计算prefill了）。然后调用 get_kv_transfer_group().send_kv_caches_and_hidden_states。

KVTransfer 实例

get_kv_transfer_group 会返回一个 KVTransfer 的实例，是一个全局实例，初始化方式如下。其中的 rank 0 代表 prefill，rank 1 代表 decode。

1
2
3
4
_KV_TRANSFER = kv_transfer.KVTransferAgent(
    rank=get_world_group().rank,
    local_rank=get_world_group().local_rank,
    config=vllm_config)

Transfer 实现

Transfer 的实现在 vllm/distributed/kv_transfer，这种解耦的设计是为了对接多种实现，比如 Mooncake 的开源 TransferEngine。Transfer 内部会调用 connector 的 send 和 recv 方法，这个方法是一个抽象方法，需要子类实现。目前有两种实现：Mooncake 的 transfer 和 PyNccl 的 transfer。

1
2
3
4
5
6
7
8
9
10
11
12
# Register various connectors here.
# The registration should not be done in each individual file, as we want to
# only load the files corresponding to the current connector.
KVConnectorFactory.register_connector(
    "PyNcclConnector",
    "vllm.distributed.kv_transfer.kv_connector.simple_connector",
    "SimpleConnector")

KVConnectorFactory.register_connector(
    "MooncakeConnector",
    "vllm.distributed.kv_transfer.kv_connector.simple_connector",
    "SimpleConnector")

Connector 依赖

Connector 依赖 kv_pipe 的实现。

  • from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator 用来实现 PyNccl 的 kv_pipe。其中的 Send 和 Recv 会依赖 NCCL 的集合通信实现。
  • 如果是 Mooncake pipe，import mooncake_vllm_adaptor as mva 这个模块，基于 ZeroMQ 的通信，通过 pickle 去序列化 tensor。
1
2
3
4
def _send_impl(self, tensor: torch.Tensor) -> None:
    """Implement the tensor sending logic."""
    value_bytes = pickle.dumps(tensor)
    self.transfer_engine.send_bytes(value_bytes)

KV Lookup Buffer 实现

另外还有一种 kv_lookup_buffer 的实现，抽象的接口是非阻塞 insert 和阻塞的 drop_select。

  • Producer 调用 insert，consumer 调用 drop_select。目前 SimpleBuffer 也是基于 Pipe 去实现的，insert 变 send，drop_select 变 recv。
  • 如果有一些中心化的 KVCacheBuffer 的话可能可以不用基于 Pipe 的实现。比如可以基于分布式的 LMCache?

Prefill 启动

vLLM 目前的实现是基于 connector 的。Prefill 的启动时通过设置 max_token 为 1 来执行，当生成了 bonus token 以后转而去调用 decode 的实例。

发表于 更新于 Disqus：

DeepSeek V3分析

During the pre-training stage, training DeepSeek-V3 on each trillion tokens requires only 180K
H800 GPU hours, i.e., 3.7 days on our cluster with 2048 H800 GPUs.

DeepSeek实现了非常便宜的训练成本，是一个700B的MoE模型。

基础设施

  • 计算集群：在配备 2048 个 NVIDIA H800 GPU 的集群上训练，节点内通过 NVLink 和 NVSwitch 连接，节点间使用 InfiniBand 互连。
  • 训练框架：基于 HAI - LLM 框架，采用 16 路管道并行（PP）、64 路专家并行（EP）和 ZeRO - 1 数据并行（DP）。设计 DualPipe 算法减少管道气泡并重叠计算与通信，开发高效的跨节点全对全通信内核，优化内存占用，无需使用昂贵的张量并行（TP）。
  • FP8 训练：提出 FP8 混合精度训练框架，对多数计算密集型操作采用 FP8 精度，对部分关键操作保留原始精度，引入细粒度量化策略、提高累积精度、采用 E4M3 格式及在线量化，还降低了内存和通信开销。
  • 推理与部署：部署在 H800 集群上，通过分离预填充和解码阶段确保服务水平目标（SLO）和高吞吐量。预填充阶段最小部署单元为 4 节点 32 个 GPU，采用特定并行策略和冗余专家策略确保负载均衡；解码阶段最小部署单元为 40 节点 320 个 GPU，采用相应并行策略和冗余专家策略，并探索动态冗余策略。
  • 硬件设计建议：针对通信硬件，期望未来硬件能卸载通信任务，统一网络接口；针对计算硬件，建议提高 FP8 GEMM 累积精度、支持细粒度量化、在线量化和转置 GEMM 操作。

并行度配置

在prefill阶段，attention模块采用4路张量并行+8路数据并行，moe模块采用32路专家并行。这样并行的目的是在满足首token时延的要求下，最大化系统吞吐（和训练任务类似）。

在decode阶段，DeepSeek-V3采取320路专家并行（256个小专家+64个热点专家），有效降低解码时延，并缓解负载不均衡的问题。

DeepSeek-V3 采用了多种并行策略，包括 16 路流水线并行（PP），这一策略有助于提高训练效率，加快模型的处理速度。同时，还应用了 64 路专家并行（EP），且在 8 个节点上进行，能够充分发挥多节点的计算优势。此外，ZeRO-1 数据并行（DP）也被运用到训练中，进一步提升了模型的训练效果。

ZeRO-1 优化器被切分到不同的GPU上。 《大模型动力引擎——PyTorch性能与显存优化手册》有提到这个优化，总结的很好。

假设我们有N=64块GPU进行数据并行训练，在ZeRO-1阶段，优化器的状态量首先被分散存储到所有GPU中，此时单张GPU上的内存使用量骤降到(4+4+8/64)*7.5=60.9GB。ZeRO-2阶段进一步地将模型的梯度也分散存储，此时单张GPU上的内存使用量便是(4+(4+8)/64)7.5=31.4GB。而ZeRO-3阶段将模型的参数也分散存储到N个节点，此时每张GPU的内存消耗只有(4+4+8)/647.5=1.875GB。从单卡需要120GB到仅需不到2GB内存，这个优化效果是不是有点惊艳？不过需要再次强调的是，这样巨大的显存优化是有代价的，显存切分的程度越高，相应的通信开销也会增加。因此，根据实际需求合理地进行显存切分是非常重要的。

MLA

采用类似 LoRA 的架构，借助一个低秩矩阵 “compressed laten vector”，kvcache 仅需对低秩的 key-value 对以及附带旋转位置编码（RoPE）的 key 进行缓存。

MoE

除了针对 Top k、routed experts 运用添加了激活函数的加权求和方式外，还额外引入了 shared experts。在 gate 的激活函数里增添一个 bias，以此来化解 balance 失衡的难题，在训练阶段，通过调节这个 bias 对 balance 状况予以奖惩，这一调节过程被称作 bias update speed。
就一个 batch、一个序列而言，每个 token 倘若倾向于特定的一些 expert，那么未被选中的 expert 实际上仅相当于训练了极小的 batch size，或者极短的序列，正因如此，才有了这样一种策略，用以平衡 expert 的 batch size 以及序列当中的 token 数量，毕竟序列通常都很长。
DeepSeek-V3 着重凭借辅助损失策略达成负载均衡，与此同时，引入互为补充的序列平衡损失，以防单个序列内部出现极度不平衡的现象。

MTP

类似于 speculative decoding，它同样会计算多个 token，不过具体方式存在一定差异。其 embedding 与 output head 是共用的，这一点和 sd 里的 Medusa 有所不同，Medusa 是由多个头来推测不同位置，而 MTP 则是依靠多个相同的头（只是 attention 有别）去推断不同位置。

MTP 的核心目的在于提升主模型的性能表现，在推理阶段能够直接将 MTP 模块舍去，主模型依旧可以独自正常运作。不仅如此，MTP 模块还能够应用于推测解码环节，以此进一步优化生成延迟问题，让整个流程更加高效流畅。

DualPipe

双流水线pipeline的优化。它实现了前向和后向过程中计算与通信阶段的重叠，有效解决了跨节点专家并行带来的通信负载问题。

FP8

能够不依赖硬件能力做FP8精度的训练，这个点是非常厉害的。

首先，为提高模型训练速度，大部分核心计算操作（尤其是 GEMM 运算），均采用 FP8 精度实现。这些 GEMM 运算接收 FP8 格式的张量输入，输出 BF16 或 FP32 格式的结果。如图6所示，线性运算相关的三个 GEMM 操作，包括 Fprop（前向传播）、Dgrad（激活值反向传播）和 Wgrad（权重反向传播），均采用 FP8 执行。这种设计策略理论上将计算速度提升至原有 BF16 方法的两倍。同时，FP8 格式的 Wgrad GEMM 使得激活值能够以 FP8 格式存储用于反向传播，显著降低了内存使用量。

发表于 更新于 Disqus：

LoRA一般的设定是认为微调任务应该只需要在一个更小的子空间去训练即可不需要复用基座模型的大空间，从而实现低成本的微调。
LoRA的前提是问题是不是在子空间能得到最优解。在线性回归 y=W x 中，如果最优 W * 是高秩的，那么对 W 施加低秩假设永远不会导致最优解，无论使用什么优化器。

Gradient Low-Rank Projection (GaLore) 允许全参数学习，但比 LoRA 等常见的低秩自适应方法更具内存效率。

使用单个批处理大小从头开始预训练 LLaMA7B 模型至少需要 58 GB 内存（14GB 用于可训练参数，42GB 用于 Adam 优化器状态和权重梯度，2GB 用于激活函数）。这使得训练在消费级 GPU 上不可行，例如具有 24GB 内存的 NVIDIA RTX 4090。

他证明梯度可能具有低秩结构，如果我们能够在优化器状态中保留梯度的一个小 “核心” 的梯度统计信息，而不是完整的梯度本身，那么内存消耗就可以大幅降低。这就引出了 GaLore 策略。
他的关键思想是利用权重矩阵 W 的梯度 G 上做LoRA，而不是试图将权重矩阵本身近似为低秩。他的核心逻辑用Torch写出来如下：

1
2
3
4
5
6
7
8
9
for weight in model.parameters():
    grad = weight.grad
    # original space -> compact space
    lor_grad = project(grad)
    # update by Adam, Adafactor, etc.
    lor_update = update(lor_grad)
    # compact space -> original space 
    update = project_back(lor_update) 
    weight.data += update

发表于 更新于 Disqus：

Sequence Parallelism

假设有4个chunk，切四份。

初始化状态，每个GPU都有自己的 Qn Kn，可以计算出对应的注意力矩阵，然后类似AllReduce的方式传递切分的K。

第一步环形传递K，然后再算一次注意力矩阵。

第二步环形传递K，然后再算一次注意力矩阵。

第三步全部传完，得到完整的Sn。

然后 Sn 和 Vn 的计算也是类似的，经过三次环形传递Vn，然后每一份可以单独和小s的那一份做乘法。

所以K和V的传播都要经历 3 次(N-1)的集合通信。