KV Cache 与 PagedAttention（vllm-musa 实战）

一、为什么 KV Cache 是 LLM 推理的"心脏"

回忆 m8-l1 的推理流水线：Prefill 阶段一次性处理整段 prompt 并建立 KV Cache，Decode 阶段每生成一个 token 都需要读取并扩展 KV Cache。没有 KV Cache，每生成一个 token 都要重算前面所有 token 的注意力——复杂度从 O(N) 退化到 O(N²)。

KV Cache 的代价是显存。以 Qwen3-32B（80 层、num_kv_heads=8、head_dim=128、FP16）为例：

每 token KV Cache 显存 = 2(K+V) × 80(layers) × 8(kv_heads) × 128(head_dim) × 2(bytes) = 327 KB
4K 上下文 1 个请求 = 327 KB × 4096 ≈ 1.3 GB
batch=32, 4K 上下文 = 1.3 GB × 32 ≈ 42 GB

一张 MTT S5000（80GB HBM）的 PH100，模型权重已经吃掉 ~64GB（FP16），留给 KV Cache 的只有 ~16GB——也就是约 12 个 4K 请求的容量。显存就是吞吐天花板。

二、PagedAttention 的核心思想

vLLM 团队借鉴操作系统虚拟内存的设计：把连续的 KV Cache 切成固定大小的 Block（默认 16 token/block），通过 Block Table 间接寻址。

• Block：显存中的物理块，固定 16 token × head_dim × num_kv_heads × 2(K+V) × dtype_bytes
• Block Table：每个请求维护一份逻辑 block → 物理 block 映射表
• 动态分配：需要时才申请新 block，请求结束立即归还到 free pool
• Copy-on-Write：多个请求共享前缀（system prompt）时共享同一份物理 block

结果：内部碎片从平均 60-80% 降到 <4%（最后一个未满 block 的浪费）；外部碎片为零（block 等大）。同等显存吞吐提升 2~4 倍。

三、在 MUSA 上跑 PagedAttention（vllm-musa 实战）

3.1 环境准备

本节示例基于以下软硬件配置（截至 2026-06-01）：

• GPU：MTT S5000（PH100 芯片，80GB HBM）单卡
• MUSA SDK：5.1.0（2026-05-18 发布）
• vllm-musa：master @ PR #46（PagedAttention 优化 commit）
• torch_musa：v2.9.0
• 测试模型：Qwen3-32B（FP16）

3.2 Step-by-Step：启用 PagedAttention

1. 拉取最新镜像：

   docker pull registry.mthreads.com/mcconline/musa-pytorch-release-public:rc4.3.0-v2.9.0-S5000-py310

2. 启动容器：

   docker run -d --name vllm-paged --gpus all --shm-size 32g \
     -v /data/models:/models registry.mthreads.com/mcconline/musa-pytorch-release-public:rc4.3.0-v2.9.0-S5000-py310 sleep infinity

3. 安装含 PR #46 的 vllm-musa：

   docker exec -it vllm-paged bash
   pip install -U git+https://github.com/MooreThreads/vllm-musa.git@main

4. 启动服务（关键参数：--block-size 16、--gpu-memory-utilization 0.92）：

   python -m vllm.entrypoints.openai.api_server \
     --model /models/Qwen3-32B \
     --block-size 16 \
     --gpu-memory-utilization 0.92 \
     --max-num-seqs 64 \
     --enable-prefix-caching

5. 验证 PagedAttention 已启用——日志中应包含：

   INFO: PagedAttention enabled, block_size=16, num_gpu_blocks=12480
   INFO: Maximum concurrency for 4096 tokens per request: 48.75x

四、PH100 + vllm-musa 实测对比表

测试场景：Qwen3-32B / ShareGPT 1000 条对话 / 平均 input=512 tok / output=256 tok / concurrency=32。

从 baseline 到全开 PagedAttention + Prefix Caching，吞吐提升 4.8 倍，TTFT P99 降低 74%。这就是 vllm-musa PR #46 + #52 带来的实际收益。

五、Block 内部寻址：从代码看 PagedAttention

vllm-musa 的 PagedAttention kernel 入口（简化版）：

// vllm/csrc/musa/paged_attention_kernel.mu
__global__ void paged_attention_v2_kernel(
    float* __restrict__ out,
    const float* __restrict__ q,           // [num_seqs, num_heads, head_size]
    const float* __restrict__ k_cache,     // [num_blocks, num_kv_heads, head_size/x, block_size, x]
    const float* __restrict__ v_cache,     // [num_blocks, num_kv_heads, head_size, block_size]
    const int* __restrict__ block_tables,  // [num_seqs, max_num_blocks_per_seq]
    const int* __restrict__ seq_lens,      // [num_seqs]
    int block_size,
    int num_kv_heads,
    float scale
) {
    const int seq_idx = blockIdx.y;
    const int head_idx = blockIdx.x;
    const int seq_len = seq_lens[seq_idx];
    const int num_blocks = (seq_len + block_size - 1) / block_size;

    // 1. 通过 block_table 间接访问物理 block
    for (int block_idx = 0; block_idx < num_blocks; block_idx++) {
        const int physical_block_number = block_tables[seq_idx * max_blocks + block_idx];
        const float* k_block_ptr = k_cache + physical_block_number * block_stride_k;
        // 2. 计算这一个 block 内的 attention scores
        ...
    }
}

关键点：block_tables[seq_idx][block_idx] 把逻辑地址翻译成物理 block 编号，逻辑连续 → 物理可以离散。这就是 vLLM 论文标题"虚拟内存"的灵魂。

六、本课检查清单（动手前必读）

• ✅ MUSA SDK ≥ 5.1.0（5.0.x 不支持新的 block_size=16 内核）
• ✅ vllm-musa 包含 PR #46（2026-05-20 之后 main 分支）
• ✅ --block-size 显式设为 16（默认值可能因版本不同）
• ✅ --gpu-memory-utilization 不超过 0.92（PH100 80GB）
• ✅ 用 GenAI-Perf 测 throughput + TTFT，对比是否达到 ~3200 tok/s 基线
• ✅ 验证显存利用率 >90%（mthreads-gmi 看显存占用）

七、本课小结与下节预告

本课讲清了三件事：

1. KV Cache 是 LLM 推理的吞吐瓶颈——显存效率决定并发数
2. PagedAttention 用 OS 虚拟内存思路解决碎片问题——block 化 + 间接寻址 + 动态分配
3. vllm-musa 在 PH100 上的实测收益：吞吐 4.8×、TTFT 降 74%

下节 m8-l3《投机解码与 EAGLE-3（vllm-musa）》将进入 Decode 阶段的另一类优化：让一个 step 一次返回 N 个 token，把 TPOT 再压低 2~3 倍。