SGLang之KV Cache内存池机制

鱿鱼圈 Lv4
  2026-05-28 23:50 2026-05-28 23:50 创建   2026-06-05 23:13:38 2026-06-05 23:13:38 更新      

本文档详细解析 SGLang 的 KV Cache 内存池机制，包括核心数据结构、分配流程，以及 KDA（Kimi Delta Attention）线性注意力的缓存机制。

一、整体架构概览

SGLang 的 KV Cache 内存池分为三层结构（见 python/sglang/srt/mem_cache/memory_pool.py）：

1
2
3
4
5
6
7
8
9
10
11
┌─────────────────────────────────────────────────────┐
│  ReqToTokenPool                                      │
│  请求 → token 位置映射（间接索引表）                  │
│  [req_pool_idx][token_position] → kv_index           │
├─────────────────────────────────────────────────────┤
│  TokenToKVPoolAllocator                              │
│  管理 kv_index 的分配/回收（空闲 token 槽位管理）     │
├─────────────────────────────────────────────────────┤
│  KVCache (实际物理存储)                              │
│  k_buffer / v_buffer 等张量                          │
└─────────────────────────────────────────────────────┘

核心设计思想：通过间接索引实现非连续的 KV 存储，从而支持 prefix cache 共享、灵活的内存回收和 paged attention。

二、核心数据结构

2.1 两张"表"的概念模型

理解内存池的关键是把它看成两张表的行号分配：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
┌─ 表1: req_to_token ──────────────────────────────────────────────┐
│  形状: (max_requests+1, max_context_len)                          │
│  dtype: int32                                                     │
│                                                                   │
│  行 0: padding (dummy, CUDA graph 填充用)                          │
│  行 1: [5, 8, 12, 3, 7, ...]  ← 请求A的各token对应的kv位置        │
│  行 2: [1, 9, 14, 6, ...]     ← 请求B的各token对应的kv位置        │
│  行 3: (空闲)                                                     │
│  行 4: [2, 11, 4, ...]        ← 请求C的各token对应的kv位置        │
│  ...                                                              │
└───────────────────────────────────────────────────────────────────┘
         ↓ 每个单元格存储的值就是下面表的行号 ↓

┌─ 表2: k_buffer / v_buffer ───────────────────────────────────────┐
│  形状: (max_total_tokens+1, head_num, head_dim) × num_layers      │
│                                                                   │
│  行 0: padding (dummy)                                            │
│  行 1: 某个 token 的 K/V 向量                                     │
│  行 2: 另一个 token 的 K/V 向量                                   │
│  行 3: ...                                                        │
│  ...                                                              │
│  行 N: 又一个 token 的 K/V 向量                                   │
└───────────────────────────────────────────────────────────────────┘

三、ReqToTokenPool 详解

3.1 数据结构

1
2
3
4
5
6
7
class ReqToTokenPool:
    def __init__(self, size, max_context_len, device, enable_memory_saver):
        self._alloc_size = size + 1  # +1 for padding row at index 0
        self.req_to_token = torch.zeros(
            (self._alloc_size, max_context_len), dtype=torch.int32, device=device
        )
        self.free_slots = list(range(1, self._alloc_size))  # Python list
  • req_to_token: 一张二维表，行是请求，列是 token 位置，值是 KV buffer 行号
  • free_slots: Python list，存储当前空闲的行号
  • 行 0 是 padding 行：CUDA graph 中 padded batch 默认 req_pool_indices=0，dummy 读写落到此行

3.2 alloc：分配请求行号

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def alloc(self, reqs: list[Req]) -> Optional[List[int]]:
    # 1. 找出已有行号的请求（chunked prefill 继续使用旧行号）
    reusing = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None]

    # 2. 计算需要新分配的数量
    need_size = len(reqs) - len(reusing)
    if need_size > len(self.free_slots):
        return None  # 内存不足

    # 3. 从空闲列表头部取行号
    select_index = self.free_slots[:need_size]
    self.free_slots = self.free_slots[need_size:]

    # 4. 赋值给请求
    offset = 0
    for r in reqs:
        if r.req_pool_idx is None:
            r.req_pool_idx = select_index[offset]
            offset += 1

    return [r.req_pool_idx for r in reqs]

语义： req_pool_idx = 2 意味着这个请求的 token 位置信息存储在 req_to_token[2, :] 这一行中。

比喻： 就像酒店的房间号分配。每个请求是一个住客，分配了房间号后，这个房间里可以存放该请求的所有 token 位置信息（最多 max_context_len 个）。

3.3 free：释放请求行号

1
2
3
def free(self, req: Req):
    self.free_slots.append(req.req_pool_idx)  # 归还行号
    req.req_pool_idx = None

3.4 write：写入映射关系

1
2
3
4
5
def write(self, indices, values):
    self.req_to_token[indices] = values

# 使用示例：
# req_to_token[req_pool_idx, prefix_len:seq_len] = out_cache_loc

四、TokenToKVPoolAllocator 详解

4.1 数据结构

1
2
3
4
5
6
7
8
9
class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
    def __init__(self, size, dtype, device, kvcache, need_sort):
        super().__init__(size, 1, dtype, device, kvcache, need_sort)
        self.clear()

    def clear(self):
        # 行号 0 是 padding, 从 1 开始
        self.free_pages = torch.arange(1, self.size + 1, dtype=torch.int64, device=device)
        self.release_pages = torch.empty((0,), dtype=torch.int64, device=device)
  • free_pages: GPU tensor，存储当前可用的 KV buffer 行号
  • release_pages: 延迟释放的行号（需要排序合并后才能复用）
  • need_sort: 是否需要排序（支持 radix cache 时需要，以保证相邻 token 物理连续）

4.2 alloc：分配 KV 存储行号

1
2
3
4
5
6
7
8
9
10
11
12
def alloc(self, need_size: int):
    # 如果需要排序且空闲不足，先合并 release_pages
    if self.need_sort and need_size > len(self.free_pages):
        self.merge_and_sort_free()

    if need_size > len(self.free_pages):
        return None  # 内存不足

    # 从 free_pages 头部切片
    select_index = self.free_pages[:need_size]
    self.free_pages = self.free_pages[need_size:]
    return select_index  # 返回 GPU tensor

语义： 返回值 [5, 8, 12] 意味着三个 token 的 K/V 向量分别存储在 k_buffer[layer][5]、k_buffer[layer][8]、k_buffer[layer][12]。

比喻： 就像仓库的储物格号码。每个 token 需要一个格子来存放它的 K/V 数据。

4.3 free：释放 KV 行号

1
2
3
4
5
6
7
8
def free(self, free_index: torch.Tensor):
    if self.is_not_in_free_group:
        if self.need_sort:
            self.release_pages = torch.cat((self.release_pages, free_index))
        else:
            self.free_pages = torch.cat((self.free_pages, free_index))
    else:
        self.free_group.append(free_index)  # 批量释放模式
  • need_sort=True 时，释放的行号先进入 release_pages（延迟合并排序，保证后续分配连续性）
  • need_sort=False 时，直接追加到 free_pages 尾部

4.4 PagedTokenToKVPoolAllocator

当 page_size > 1 时使用分页版本（allocator.py:362），管理逻辑类似但以页为单位分配，一页包含多个连续 token slot。

五、两个 alloc 如何配合：完整数据流

5.1 Extend（Prefill）阶段

以一个请求有 5 个 token（prefix 2 个已缓存 + 新增 3 个）为例：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
步骤1: alloc_req_slots()  [common.py:398]
     → ReqToTokenPool.alloc(reqs)
     → 得到 req_pool_idx = 2  (这个请求使用 req_to_token 表的第2行)

步骤2: alloc_token_slots()  [common.py:302]
     → TokenToKVPoolAllocator.alloc(3)  (新增3个token需要3个KV位置)
     → 得到 out_cache_loc = [17, 42, 89]  (KV buffer的3个空闲行号)

步骤3: write_cache_indices()  [common.py:104]
     → 把映射关系写入 req_to_token 表:

     req_to_token[2, 0] = 5    ← prefix token 0 (之前已缓存在KV行5)
     req_to_token[2, 1] = 11   ← prefix token 1 (之前已缓存在KV行11)
     req_to_token[2, 2] = 17   ← 新 token 2 (刚分配的)
     req_to_token[2, 3] = 42   ← 新 token 3 (刚分配的)
     req_to_token[2, 4] = 89   ← 新 token 4 (刚分配的)

关键代码路径（common.py:134-150）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
for i in range(len(reqs)):
    req_idx = req_pool_indices_cpu[i]
    prefix_len = prefix_lens_cpu[i]
    seq_len = seq_lens_cpu[i]
    extend_len = extend_lens_cpu[i]

    # 写入 prefix 部分（来自 radix cache 共享）
    req_to_token_pool.write(
        (req_idx, slice(0, prefix_len)),
        prefix_tensors[i],
    )
    # 写入新分配的 extend 部分
    req_to_token_pool.write(
        (req_idx, slice(prefix_len, seq_len)),
        out_cache_loc[pt : pt + extend_len],
    )

5.2 Decode 阶段

每步生成 1 个新 token（common.py:524）：

1
2
3
4
5
6
7
8
9
步骤1: (req_pool_idx 已有，不需要重新分配)

步骤2: TokenToKVPoolAllocator.alloc(batch_size)
     → 得到 out_cache_loc = [103, 104, ...]  (每个请求1个新位置)

步骤3: write:
     req_to_token[2, 5] = 103  ← 请求A新生成的 token 5
     req_to_token[4, 8] = 104  ← 请求C新生成的 token 8
     ...

关键代码（common.py:559）：

1
2
3
batch.req_to_token_pool.write(
    (batch.req_pool_indices, batch.seq_lens), out_cache_loc.to(torch.int32)
)

5.3 Attention 计算时的查找路径

已知 req_pool_idx=2, seq_len=5, layer=3，查找全部 K/V：

1
2
3
4
5
6
1. indices = req_to_token[2, 0:5] → [5, 11, 17, 42, 89]    (5个KV行号)

2. k = k_buffer[3][[5, 11, 17, 42, 89]]   (gather 操作, shape: [5, head_num, head_dim])
   v = v_buffer[3][[5, 11, 17, 42, 89]]   (gather 操作, shape: [5, head_num, head_dim])

3. output = attention(q, k, v)             (标准注意力计算)

六、为什么需要两层间接索引？

不连续存储的核心好处

  1. Prefix cache 共享： 不同请求如果有相同的 prefix，它们的 req_to_token 行中前几列可以存相同的 KV 行号 – 物理上共享同一份 KV 数据
  2. 灵活回收： 可以单独释放某些 token 的 KV（如 SWA 窗口外的 token），无需移动其他数据
  3. 避免碎片： 无需做 compaction（数据搬移），释放后的行号可直接被新 token 复用
  4. 支持 chunked prefill： 一个请求可以分多次填充 KV，每次分配新的不连续位置

如果没有间接层会怎样？

每个请求的 token 必须连续存放在 KV buffer 中：

  • Prefix cache 无法共享（相同 prefix 必须复制一份）
  • 请求完成后留下"空洞"，需要 compaction 或浪费内存
  • Chunked prefill 需要预分配连续空间

七、两个 alloc 的对比总结

维度 ReqToTokenPool.alloc TokenToKVPoolAllocator.alloc
分配什么 req_to_token 表的行号 k_buffer/v_buffer 的行号
粒度 每个请求 1 个 每个新 token 1 个
存储介质 Python list (CPU) GPU tensor
生命周期 请求存活期间一直持有 prefix cache 可跨请求共享复用
调用时机 请求首次进入 running batch 每次 extend/decode 生成新 token
释放时机 请求完成时 token 被从 radix tree 驱逐时
数量级 ~max_running_requests (少) ~max_total_tokens (多，通常比请求数大10-100倍)
OOM 影响 限制并发请求数 限制总缓存 token 数

八、KDA（Kimi Delta Attention）的缓存机制

8.1 KDA 概述

KDA 是 Kimi 系列模型使用的线性注意力机制（kda_backend.py:130）。Kimi Linear 是一个混合架构：

  • 部分层使用 KDA（线性注意力）-- 无 per-token KV cache，维护固定大小递归状态
  • 部分层使用标准 full attention（MLA）-- 使用上述 per-token KV cache

8.2 KDA 层的状态结构

KDA 不使用传统的 per-token K/V cache，而是维护两种递归状态（类似 Mamba/SSM）：

从 KimiLinearStateShape.create（configs/mamba_utils.py:197）:

1
2
3
4
5
6
7
# causal conv1d 的滑窗历史
conv_state_shape = (conv_kernel - 1, proj_size / tp)
# 例: (3, 4096) -- 卷积核大小4, 隐藏维度4096

# Delta Rule 的递归隐藏状态
temporal_state_shape = (num_heads / tp, head_dim, head_dim)
# 例: (32, 128, 128) -- 32个头, 每个头一个 128x128 矩阵

物理存储布局：

  • conv_state: [num_kda_layers, pool_size+1, conv_kernel-1, dim]
  • temporal_state (ssm_state): [num_kda_layers, pool_size+1, num_heads, head_dim, head_dim]

关键特性： 每个请求对应一个 ssm_state slot，大小是固定的 head_dim x head_dim 矩阵，不随序列长度增长。

8.3 HybridReqToTokenPool – 混合内存池

HybridReqToTokenPool（memory_pool.py:487）专为 KDA + full attention 混合架构设计：

1
2
3
4
5
6
7
8
9
HybridReqToTokenPool
├── req_to_token (继承自 ReqToTokenPool)
│   → 用于 full attention 层的 token → kv 映射
├── MambaPool (mamba_pool)
│   → 用于 KDA 层的 conv_state + ssm_state 存储
├── req_index_to_mamba_index_mapping: torch.Tensor
│   → 请求行号 → mamba state slot 的映射
└── mamba_ping_pong_track_buffer
    → overlap schedule 使用的双缓冲

分配流程（memory_pool.py:572）：

1
2
3
4
5
6
7
8
9
10
11
def alloc(self, reqs):
    # 1. 为 full attention 层分配 req_to_token 行
    select_index = super().alloc(reqs)  # ReqToTokenPool.alloc

    # 2. 为每个请求从 mamba_pool 分配一个 state slot
    for req in reqs:
        mid = self.mamba_pool.alloc(1)  # 分配 mamba state 槽位
        req.mamba_pool_idx = mid[0]

    # 3. 记录映射关系
    self.req_index_to_mamba_index_mapping[select_index] = mamba_indices

8.4 HybridLinearKVPool – 混合 KV Cache

HybridLinearKVPool（memory_pool.py:1389）组合了两种存储：

1
2
3
4
5
6
7
8
9
10
class HybridLinearKVPool(KVCache):
    def __init__(self, ...):
        # Full attention 层：正常的 per-token KV cache
        self.full_kv_pool = MHATokenToKVPool(...) 或 MLATokenToKVPool(...)

        # KDA 层：引用同一个 MambaPool（不使用 per-token KV）
        self.mamba_pool = mamba_pool

        # Full attention 层的 layer_id 映射
        self.full_attention_layer_id_mapping = {id: i for i, id in enumerate(full_attention_layer_ids)}
  • Full attention 层：通过 get_kv_buffer(layer_id) 访问 k_buffer[layer]、v_buffer[layer]
  • KDA 层：通过 mamba2_layer_cache(layer_id) 获取 conv_states 和 ssm_states

8.5 KDA 运行时的缓存使用

Decode 阶段（kda_backend.py:138）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def forward_decode(self, layer, mixed_qkv, a, b):
    # 获取当前层的 mamba state
    layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
    conv_states = layer_cache.conv[0]         # [pool_size+1, conv_kernel-1, dim]
    ssm_states = layer_cache.temporal         # [pool_size+1, num_heads, head_dim, head_dim]
    cache_indices = forward_metadata.mamba_cache_indices  # 当前 batch 的 slot 索引

    # 1. Causal Conv1d 更新（原地更新 conv_states）
    qkv = causal_conv1d_update(mixed_qkv, conv_states, ..., conv_state_indices=cache_indices)

    # 2. Delta Rule 递归更新（Triton kernel）
    #    内部逻辑:
    #    h = ssm_states[cache_indices]   # 加载状态
    #    h *= exp(g)                      # 门控衰减
    #    v -= sum(h * k)                  # delta 修正
    #    v *= beta                        # beta 门控
    #    h += k ⊗ v                       # 状态更新
    #    o = sum(h * q)                   # 输出
    #    ssm_states[cache_indices] = h    # 写回
    output = self.kernel_dispatcher.decode(q, k, v, a, b,
        ssm_states=ssm_states, cache_indices=cache_indices, ...)

Extend/Prefill 阶段（kda_backend.py:178）

1
2
3
4
5
6
7
def forward_extend(self, layer, forward_batch, mixed_qkv, a, b):
    # 使用 chunk-based 并行算法处理长序列
    # chunk_size = 64, 支持 variable-length (cu_seqlens)
    output = self.kernel_dispatcher.extend(q, k, v, g, beta,
        ssm_states=ssm_states,       # initial_state 从 mamba pool 加载
        cache_indices=cache_indices,  # final_state 原地写回
        ...)

Chunk-based 算法（fla/kda.py:1025）将长序列分块，组合 inter-chunk 递归 + intra-chunk 并行计算。

8.6 KDA vs Full Attention 内存对比

特性 Full Attention (per-token KV) KDA (递归状态)
内存增长 O(seq_len x heads x head_dim) 线性增长 O(heads x head_dim^2) 固定
存储位置 k_buffer[layer][token_index] temporal[layer][slot_index]
索引方式 req_to_token 间接索引 mamba_cache_indices 直接索引
Prefix cache radix tree 共享 KV 行号 fork（复制 state）
Decode 复杂度 Attention O(seq_len) 递归更新 O(head_dim^2)
示例大小 (per req) 4096 tokens x 32 heads x 128 dim x 2 x 2B = 64MB 32 heads x 128 x 128 x 4B = 2MB

九、完整架构图

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
37
38
39
40
┌─ 请求调度层 ────────────────────────────────────────────────────────┐
│   Scheduler → alloc_for_extend(batch) / alloc_for_decode(batch)      │
│                                                                      │
│   alloc_req_slots():                                                 │
│     → HybridReqToTokenPool.alloc(reqs)                               │
│        ├→ ReqToTokenPool.alloc(): 分配 req_pool_idx (表1行号)         │
│        └→ MambaPool.alloc(): 分配 mamba_pool_idx (state slot)         │
│                                                                      │
│   alloc_token_slots():                                               │
│     → TokenToKVPoolAllocator.alloc(num_tokens): 分配 kv 行号 (表2行号)│
│                                                                      │
│   write_cache_indices():                                             │
│     → req_to_token[req_pool_idx, prefix:seq_len] = kv行号            │
└──────────────────────────────────────────────────────────────────────┘
                              ↓
┌─ Forward 执行层 ─────────────────────────────────────────────────────┐
│                                                                       │
│   Full Attention Layer (layer 0, 4, 8, ...):                          │
│     kv_indices = req_to_token[req_pool_idx, 0:seq_len]                │
│     k = k_buffer[layer][kv_indices]   (gather)                        │
│     v = v_buffer[layer][kv_indices]   (gather)                        │
│     output = attention(q, k, v)                                       │
│                                                                       │
│   KDA Layer (layer 1, 2, 3, 5, 6, 7, ...):                           │
│     layer_cache = mamba2_layer_cache(layer_id)                        │
│     conv_states = layer_cache.conv[0]                                 │
│     ssm_states = layer_cache.temporal                                 │
│     causal_conv1d_update(qkv, conv_states[cache_indices])             │
│     delta_rule_update(q, k, v, ssm_states[cache_indices])             │
│     → 无需 per-token KV indexing, 只需 per-request state slot         │
│                                                                       │
└──────────────────────────────────────────────────────────────────────┘
                              ↓
┌─ 释放层 ─────────────────────────────────────────────────────────────┐
│   release_kv_cache(req):                                              │
│     → tree_cache.cache_finished_req(req)  (可能保留给 prefix cache)    │
│     → token_to_kv_pool_allocator.free(kv_indices)  (释放KV行号)       │
│     → req_to_token_pool.free(req)  (释放请求行号)                     │
│     → mamba_pool.free(mamba_pool_idx)  (释放state slot)               │
└──────────────────────────────────────────────────────────────────────┘

十、关键源码文件索引

文件 内容
python/sglang/srt/mem_cache/memory_pool.py ReqToTokenPool, MambaPool, HybridReqToTokenPool, KVCache 及各实现类
python/sglang/srt/mem_cache/allocator.py TokenToKVPoolAllocator, PagedTokenToKVPoolAllocator
python/sglang/srt/mem_cache/common.py alloc_for_extend, alloc_for_decode, write_cache_indices 等调度逻辑
python/sglang/srt/layers/attention/linear/kda_backend.py KDAAttnBackend (KDA 的 forward_decode/extend)
python/sglang/srt/layers/attention/fla/kda.py chunk_kda (chunk-based KDA prefill 算法)
python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py fused_sigmoid_gating_delta_rule_update (KDA decode kernel)
python/sglang/srt/configs/kimi_linear.py KimiLinearConfig (模型配置, kda_layers 定义)
python/sglang/srt/configs/mamba_utils.py KimiLinearStateShape, KimiLinearCacheParams (state 形状定义)
python/sglang/srt/layers/radix_linear_attention.py RadixLinearAttention (KDA 层的 nn.Module)
  • 标题: SGLang之KV Cache内存池机制
  • 作者: 鱿鱼圈
  • 创建于 : 2026-05-28 23:50:00
  • 更新于 : 2026-06-05 23:13:38
  • 链接: https://yuyanqi.com/2026/05/28/SGLang之KV Cache内存池机制详解/
  • 版权声明: 本文章采用 CC BY-NC-SA 4.0 进行许可。
评论
目录
SGLang之KV Cache内存池机制
  1. 一、整体架构概览
  2. 二、核心数据结构
    1. 2.1 两张"表"的概念模型
  3. 三、ReqToTokenPool 详解
    1. 3.1 数据结构
    2. 3.2 alloc：分配请求行号
    3. 3.3 free：释放请求行号
    4. 3.4 write：写入映射关系
  4. 四、TokenToKVPoolAllocator 详解
    1. 4.1 数据结构
    2. 4.2 alloc：分配 KV 存储行号
    3. 4.3 free：释放 KV 行号
    4. 4.4 PagedTokenToKVPoolAllocator
  5. 五、两个 alloc 如何配合：完整数据流
    1. 5.1 Extend（Prefill）阶段
    2. 5.2 Decode 阶段
    3. 5.3 Attention 计算时的查找路径
  6. 六、为什么需要两层间接索引？
    1. 不连续存储的核心好处
    2. 如果没有间接层会怎样？
  7. 七、两个 alloc 的对比总结
  8. 八、KDA（Kimi Delta Attention）的缓存机制
    1. 8.1 KDA 概述
    2. 8.2 KDA 层的状态结构
    3. 8.3 HybridReqToTokenPool – 混合内存池
    4. 8.4 HybridLinearKVPool – 混合 KV Cache
    5. 8.5 KDA 运行时的缓存使用
    6. 8.6 KDA vs Full Attention 内存对比
  9. 九、完整架构图
  10. 十、关键源码文件索引