rtos_llm_opt/02-kv-cache-memory.md

方向2: KV Cache 与内存管理

1. 问题陈述

KV Cache是LLM推理中最大的内存消费者，也是RTOS内存管理的核心痛点：

KV Cache Size ≈ 2 × N_layers × batch_size × seq_len × hidden_dim × 2 bytes

Example (Qwen2.5-7B, batch=32, seq_len=4096):
  = 2 × 32 × 32 × 4096 × 4096 × 2 bytes
  ≈ 6.8 GB

Example (Qwen2.5-1.5B, batch=8, seq_len=2048):
  = 2 × 32 × 8 × 2048 × 2048 × 2 bytes
  ≈ 2.2 GB

RTOS场景下KV Cache管理的关键问题：

1. 内存碎片 — 频繁分配/释放导致碎片化
2. 带宽竞争 — KV读写与计算同时竞争DDR
3. 内存带宽瓶颈 — Attention是memory-bound而非compute-bound
4. 多请求KV Cache隔离 — 多流场景下的内存隔离与复用

2. KV Cache结构分析

2.1 KV Cache布局

┌─────────────────────────────────────────────────────┐
│                   KV Cache Layout                    │
│                                                      │
│  Layer 0: [K_0 | V_0]  → Shape: [batch, head, seq, hd] |
│  Layer 1: [K_1 | V_1]  → Shape: [batch, head, seq, hd] |
│  ...                                                 │
│  Layer N: [K_N | V_N]  → Shape: [batch, head, seq, hd] |
│                                                      │
│  Total = 2 × N_layers × batch × seq × hidden × dtype |
└─────────────────────────────────────────────────────┘

Attention计算公式:
  Output = Softmax(QK^T / √d_k) × V

Q, K, V均为实时维护的Tensor。推理过程中：
  - Prefill阶段: K, V一次性填充
  - Decode阶段: K, V逐tokenappend

2.2 KV Cache的内存访问特征

Access Pattern          | Read  | Write | Latency | Bandwidth
------------------------|-------|-------|---------|----------
Prefill K/V init        |  Low  | High  | Low     | Very High
Decode K append         |  High | Medium| Medium  | Medium
Attention QK^T          |  High |  Low  | Medium  | High
Attention SV            |  High | High  | Medium  | High

关键洞察：Attention是 memory-bound 的，内存带宽比计算算力更容易成为瓶颈。

3. 内存管理策略

3.1 Memory Pool 设计

核心思路：为KV Cache预分配固定大小的内存池，避免运行时malloc/free。

// KV Cache Memory Pool
typedef struct {
    uint8_t *base;              // 内存池基地址
    size_t total_size;          // 总大小
    size_t used_size;           // 已使用
    size_t max_block;           // 最大block大小
    uint32_t n_blocks;          // block数量
    uint32_t *free_map;         // 空闲块位图
    uint32_t alloc_count;       // 当前分配次数
} kv_cache_pool_t;

// Block结构 — 每个layer一个block
typedef struct {
    struct kv_cache_pool_t *pool;
    uint8_t *data;              // 数据指针
    size_t size;                // 大小
    uint32_t layer_id;          // 所属layer
    uint32_t batch_id;          // 所属batch
    bool locked;                // 是否锁定(防释放)
} kv_block_t;

Pool设计决策：

策略	描述	优点	缺点
Fixed-size pool	预分配固定大小，不动态扩容	零碎片，O(1)分配	可能浪费或不够用
Buddy system	powers-of-2分块	低碎片	内碎片最多50%
Slab allocator	预分配固定类型cache	高效同类分配	不同size需多个slab
Region pool	按请求分配region	便于多请求隔离	region间可能碎片

推荐: 在嵌入式场景用Fixed-size pool，在边缘场景用Slab + Region组合。

3.2 分层KV Cache管理

┌──────────────────────────────────────────────┐
│  L1: SRAM Cache (芯片内, ~100KB)             │
│  最近访问的KV block, 超低延迟(<10ns)          │
│  策略: LRU, hardware-managed                  │
├──────────────────────────────────────────────┤
│  L2: DDR Memory Pool (RAM, 2-8GB)            │
│  所有KV Cache, 中延迟(~100ns)                 │
│  策略: Slab allocator + region per request    │
├──────────────────────────────────────────────┤
│  L3: Flash/SSD (持久化, 10-100GB)           │
│  冷KV Cache, 高延迟(~100μs)                   │
│  策略: 按需swap, 预读                         │
└──────────────────────────────────────────────┘

3.3 多请求KV Cache隔离

Request A: [Layer 0 Block | Layer 1 Block | ...]  → Pool Region A
Request B: [Layer 0 Block | Layer 1 Block | ...]  → Pool Region B
Request C: [Layer 0 Block | Layer 1 Block | ...]  → Pool Region C

Region管理:
┌──────────┬──────────┬──────────┬──────────┐
│ Region A │ Region B │ Region C │   Free   │
│ 256MB    │ 128MB    │ 256MB    │ 512MB    │
└──────────┴──────────┴──────────┴──────────┘

每个Region包含:
  - Header: 大小、引用计数、是否锁定
  - Data: KV Cache数据
  - Metadata: 序列长度、是否有效、引用task

4. 内存带宽优化

4.1 带宽竞争模型

总带宽 = DDR Bandwidth (e.g., LPDDR5: 6400 Mbps × 4 = 25.6 GB/s)

带宽分配:
  KV Cache Read (Attention):    █████████████████████  40%
  KV Cache Write (Append):      ████████████           25%
  Weight Read (FFN/Attn):       ████████████████████  30%
  Input/Output:                 ████████               5%

问题: KV Cache读+写 = 65% 带宽
      加上Weight Read = 95% 带宽
      仅剩5% 给其他任务

4.2 带宽调度策略

策略1: 错峰访问

Time →
KV Cache Write: [████████][            ][████████]
Weight Read:    [        ][████████████][        ]
Input/Output:   [███████][            ][        ]
              t0        t1           t2

策略2: DMA Batching

Before (no batching):
  KV Write 1 → DMA → DDR (small chunk)
  KV Write 2 → DMA → DDR (small chunk)
  KV Write 3 → DMA → DDR (small chunk)
  Total DMA overhead: 3 × overhead

After (batching):
  KV Write 1,2,3 → DMA burst → DDR (large chunk)
  Total DMA overhead: 1 × overhead

RTOS实现:
  1. KV Write task将多个小DMA请求放入队列
  2. DMA Controller task批量处理
  3. 使用RTOS queue传递DMA描述符

策略3: Bandwidth-aware Scheduling

typedef struct {
    uint64_t bandwidth_allocated;  // 已分配带宽
    uint64_t bandwidth_used;        // 已使用带宽
    uint64_t peak_bandwidth;        // 峰值带宽
    uint64_t window_ms;             // 滑动窗口大小
} bandwidth_tracker_t;

// 任务请求带宽时的检查
bool can_allocate_bandwidth(task_t *task, uint64_t required_bw) {
    if (bandwidth_tracker.used + required_bw <= BANDWIDTH_LIMIT) {
        bandwidth_tracker.used += required_bw;
        return true;
    }
    // 等待或排队
    vTaskSuspend(task);
    return false;
}

4.3 In-Place KV Update

减少KV Cache的write bandwidth：

Before (old):
  for each new token:
    allocate new KV block      // 分配+写
    write new KV values        // 写
    update pointer             // 写

After (in-place):
  for each new token:
    KV_buffer += step_size     // 指针移动(O(1))
    write KV values in-place   // 只写一次

实现要点：

• 预分配连续的KV内存(避免碎片)
• 使用环形buffer管理seq_len(自然复用)
• 维护valid长度指针，不实际移动数据

5. KV Cache替换策略

5.1 缓存淘汰算法

当KV Cache满时，需要选择替换策略：

策略                  | 复杂度 | 命中率 | 适用场景
----------------------|--------|--------|---------
LRU                  | O(1)   | Good   | 通用
LFU                  | O(log n)| Better | 长对话
Sliding Window       | O(1)   | Good   | 固定上下文
Hybrid (Window+LFU)  | O(1)   | Best   | 生产环境

RTOS友好实现 (LRU with doubly-linked list):

typedef struct kv_cache_node {
    uint32_t layer;
    uint32_t token_pos;
    struct kv_cache_node *prev;
    struct kv_cache_node *next;
    uint8_t *data;
    uint64_t last_access;
} kv_cache_node_t;

// LRU: 访问时移动到链表头部
void kv_cache_access(uint32_t layer, uint32_t pos) {
    kv_cache_node_t *node = find_node(layer, pos);
    if (node) {
        node->last_access = get_tick_ms();
        move_to_head(&lru_list, node);  // O(1)
    }
}

// Evict: 从链表尾部淘汰
void kv_cache_evict(void) {
    kv_cache_node_t *victim = lru_list.tail;
    remove_from_list(victim);
    free_block(victim);
}

5.2 压缩策略

策略          | 压缩比 | 精度损失 | 带宽节省 | 计算开销
--------------|--------|----------|----------|---------
Float16       | 1.0x   | 0%       | 基准     | 无
Int8          | 2.0x   | ~2%      | 50%      | 低
Int4          | 4.0x   | ~5%      | 75%      | 中
Sparsification| 2-8x   | 1-10%   | 50-87%   | 中-高
Paged KV      | variable| 0%     | variable | 低

6. 各硬件平台的内存管理差异

6.1 MCU级 (256KB~2MB)

约束:
- 内存极小，无法容纳完整KV Cache
- 通常只支持一个请求
- 模型参数也需压缩

策略:
- KV Cache固定大小(如256 token)
- 使用连续内存池，零碎片
- KV Cache满后触发OOM处理(丢弃 oldest)
- 可能需要在Flash上swap

6.2 SoC级 (2-8GB)

约束:
- 内存有限但可容纳小模型KV Cache
- 多请求时内存竞争严重
- NPU有独立的memory pool

策略:
- 分区管理: CPU pool + NPU pool
- 使用Paged Memory (类似vLLP)
- NPU-Shared Memory通过DMA同步

6.3 Edge盒子 (8-16GB)

约束:
- 内存充足但带宽可能受限
- 多请求+大模型场景

策略:
- 分层缓存: SRAM(热点) + DDR(主流) + SSD(冷数据)
- 预读策略: 预测下一个token的KV位置
- 多流隔离: CGroup级别内存限制

6.4 Server级

约束:
- 内存充足(NVMe/DDR)
- 主要问题是带宽和NUMA

策略:
- NUMA-aware KV placement
- NVMe作为KV Cache扩展
- 多GPU间KV Cache同步

7. 关键设计决策

决策点	选项	推荐	理由
内存分配器	malloc/free / Pool / Slab	Slab + Region	零碎片+多请求隔离
KV Cache布局	连续 / Paged / Sparse	Paged	灵活+低碎片
替换策略	LRU / LFU / Sliding	Sliding Window	计算复杂度O(1)
压缩策略	FP16 / Int8 / Int4	Int8	精度-带宽权衡最优
带宽管理	固定分配 / 动态	Bandwidth-aware	避免带宽饱和
多流隔离	共享池 / Region	Region	公平+可分析

8. 开放研究问题

1. Dynamic KV Cache Sizing: 运行时根据负载自动调整KV Cache大小？
2. Cross-device KV: 在CPU-NPU间共享/分发KV Cache？
3. KV Cache Compression: 有损压缩的RTOS友好实现？
4. Predictive Allocation: 预测最大seq_len，预分配最优大小？
5. KV Cache Migration: 请求迁移时的KV Cache热迁移？

---

最后更新: 2026-09-17