端侧推理 KV Cache 内存优化：基于 PagedAttention 的 WebGPU 显存碎片治理

/* MD / 富文本中的 .toc（含博客园搬家等嵌套结构）；.toc-box 在侧栏，不受影响 */
#content_views .toc,
/* 编辑器常在目录前后插入空 p（:empty 仍占 20px），一并去掉避免顶空隙 */
#content_views.markdown_views > p:empty:has(+ .toc),
#content_views.markdown_views > .toc + p:empty,
/* 富文本旧版目录标记 */
#content_views.htmledit_views #main-toc,
#content_views.htmledit_views #hr-toc,
#content_views.htmledit_views p[id*=\”-toc\”] {
display: none !important;
}
/* 目录去掉后，紧跟的首个标题不再多出一块上边距 */
#content_views.markdown_views > .toc + h1,
#content_views.markdown_views > .toc + h2,
#content_views.markdown_views > .toc + h3,
#content_views.markdown_views > .toc + h4,
#content_views.markdown_views > .toc + p:empty + h1,
#content_views.markdown_views > .toc + p:empty + h2,
#content_views.markdown_views > .toc + p:empty + h3,
#content_views.markdown_views > .toc + p:empty + h4 {
margin-top: 0 !important;
}

在浏览器端跑 WebLLM 或本地大模型推理时，很多开发者一开始都能顺利跑通“单轮 Hello World”。然而只要一旦进入多轮对话、或者输入一段超过 1,000 字的业务上下文，控制台就会突然抛出致命的 GPUOutOfMemoryError，页面瞬时崩溃。

造成这一崩溃的元凶，通常不是模型本身的静态权重参数，而是 Transformer 自回归生成中疯狂膨胀的显存黑洞——KV Cache（Key-Value 键值对缓存）。

在自回归解码（Decode）阶段，模型每生成一个新的 Token，为了避免重复计算历史 Context 的注意力权重，必须把历史所有序列的 Key 和 Value 向量保存在显存中。随着对话轮次的增加，KV Cache 的体积呈线性暴涨。

在云端高性能服务器上，我们可以挥霍数十 G 的 H100 显存；但在用户的浏览器端，WebGPU 分配给单个 Tab 标签页的显存配额通常被严格限制在 1GB ~ 2GB 以内。

如果依然采用传统的“静态连续显存预分配（Contiguous Pre-allocation）”，整个系统将陷入灾难性的显存碎片化泥潭。

为了在寸土寸金的端侧环境中让模型承载长达数千字的上下文，必须将云端推理利器 vLLM 的核心思想——分页注意力（PagedAttention），移植并改良至浏览器的 WebGPU WGSL 算子管线中。本文系统解析我们在端侧实现显存碎片治理的工程实践。

为什么传统的 KV Cache 分配在浏览器端注定破产？

要看懂 PagedAttention 的革命性，我们先剖析传统静态连续分配的三大“原罪”：

[ 传统连续分配: 预留 2048 上下文空间 (假设占用 800MB 显存) ]
┌────────────────────────────────────────────────────────┐
│ [已用 200 Token] │ 闲置浪费 1848 Token 空间 │
└────────────────────────────────────────────────────────┘
– 内部碎片 (Internal Fragmentation): 真实对话才刚开始，80% 以上显存被白白冻结！
– 外部碎片 (External Fragmentation): 无法找到一块连续的 800MB 物理显存，直接 OOM！
– 动态扩容拷贝开销: 一旦超过预留长度，必须重开更大 Buffer 并执行耗时的全量显存拷贝

PagedAttention 核心架构：操作系统虚拟分页的降维打击

PagedAttention 的灵感直接源自操作系统的虚拟内存分页机制（Virtual Memory Paging）：它彻底消灭了“必须在连续物理显存中存放连续上下文”的教条，把 KV Cache 打散为一个个固定尺寸的微小显存块（Block / Page）。

[ 逻辑 Token 序列 (Token 0 ~ Token 47) ]
┌───────────────┬───────────────┬───────────────┐
│ Block 0 (0~15)│ Block 1(16~31)│ Block 2(32~47)│
└───────────────┴───────────────┴───────────────┘
│ │ │ (通过 Block Table 虚拟映射)
▼ ▼ ▼
[ WebGPU 物理分页显存池 (由离散的 16-Token 物理块拼装而成) ]
┌───────────────┬───────────────┬───────────────┬───────────────┐
│ 物理块 #12 │ 物理块 #3 │ 物理块 #88 │ …其他空闲块 │
└───────────────┴───────────────┴───────────────┴───────────────┘

• 块尺寸（Block Size）：通常设定为 16 个 Tokens。每个块仅占用几兆字节显存；
• 按需分配（Allocate on Demand）：模型每新生成 16 个 Token，才从空闲块池中申请一个物理块，内部碎片率被压制到 3% 以下；
• 块映射表（Block Table）：在 CPU 侧维护一张轻量级的数字索引数组，记录逻辑块号与物理显存块号的对应关系，在调用着色器时传递给 GPU。

WGSL 计算着色器实现：带非连续寻址的 Paged Attention

在 WebGPU 的 WGSL 着色器中，传统注意力的索引计算 index = token_idx * head_dim 变成了基于映射表的二次寻址：

// WGSL 片段: paged_attention.wgsl

struct BlockTable {
entries: array<u32>, // 记录每个逻辑 block 对应的物理 block_id
};

@group(0) @binding(0) var<storage, read> q: array<f32>;
@group(0) @binding(1) var<storage, read> k_cache_pool: array<f32>; // 物理分页显存池
@group(0) @binding(2) var<storage, read> block_table: BlockTable;

const BLOCK_SIZE: u32 = 16u;
const HEAD_DIM: u32 = 64u;

// 核心：将逻辑上下文位置转换为物理显存池索引
fn get_k_cache_offset(token_idx: u32, head_idx: u32) -> u32 {
let logical_block_idx = token_idx / BLOCK_SIZE;
let block_offset = token_idx % BLOCK_SIZE;

// 通过 Block Table 捞出真实的物理块号
let physical_block_id = block_table.entries[logical_block_idx];

// 定位物理地址偏移量
return (physical_block_id * BLOCK_SIZE + block_offset) * HEAD_DIM;
}

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
let token_idx = global_id.x;
// 基于 get_k_cache_offset 在非连续物理块中完成 Q * K 点积注意力计算…
}

注意这套设计的绝妙之处：在 GPU 计算核心看来，它虽然读取的是散落在不同物理地址上的数据，但由于现代 GPU 强大的 L2 缓存预取机制，只要同一个 Workgroup 内部的数据保持在 16 字节对齐，其内存访问吞吐（Memory Bandwidth）几乎没有发生任何退化！

前端 JavaScript 显存池管理器 (BlockSpaceManager)

在主线程侧，我们用一套极简的环形空闲链表管理所有物理块：

// src/ai/blockSpaceManager.ts

export class BlockSpaceManager {
private totalBlocks: number;
private freeBlocks: number[] = [];
private blockSizeTokens: number;

constructor(totalMemoryBytes: number, bytesPerToken: number, blockSizeTokens = 16) {
this.blockSizeTokens = blockSizeTokens;
const bytesPerBlock = bytesPerToken * blockSizeTokens;
this.totalBlocks = Math.floor(totalMemoryBytes / bytesPerBlock);

// 初始化所有可用物理块索引
for (let i = 0; i < this.totalBlocks; i++) {
this.freeBlocks.push(i);
}
}

/**
* 为当前会话追加分配一个新物理块
*/
public allocateBlock(): number {
if (this.freeBlocks.length === 0) {
throw new Error('[PagedKV] 显存物理块池耗尽，触发紧急会话驱逐或云端回退！');
}
return this.freeBlocks.pop()!;
}

/**
* 会话结束或滑动窗口过期时，回收物理块
*/
public freeBlock(physicalBlockId: number): void {
this.freeBlocks.push(physicalBlockId);
}

public getUtilization(): number {
return (1 – this.freeBlocks.length / this.totalBlocks) * 100;
}
}

实测收益：显存利用率暴增 4 倍

我们在 MacBook Pro (M3 Pro) 与主流 Windows 独显 PC 上，针对 1.5B 参数规模的端侧模型进行了极限长文本多轮对话测试：

实测数据表明：

• 在同样的 1GB WebGPU 显存上限下，PagedAttention 使得模型能够支撑的有效上下文长度直接拉升了整整 4 倍；
• 显存利用效率从可怜的 22% 飙升至近 90%；
• 多轮会话由于不再需要预留大块连续内存，浏览器再也没有发生过一次由显存碎片引起的 OOM 崩溃。

结语

在端侧算力这个充满物理约束的舞台上，唯有把每一兆显存、每一个字节的利用率压榨到极限，才能把原本属于服务器集群的大模型能力真正变成用户浏览器里的随身利器。将系统级的分页思想搬进 WebGPU 的着色器管线，不仅展现了现代前端架构在底层系统工程上的无限可能，更为端侧 AI 的大规模商业化落地踩平了最深的一道泥坑。