DeepSeek MLA — Prefix Cache 与 Chunked Prefill 深度解析

1. 核心矛盾：Cache 格式 vs Prefill 格式

DeepSeek MLA 的独特之处在于：存储在 KV Cache 中的数据格式和Prefill 阶段注意力计算需要的数据格式是不一样的。这是 MLA 区别于 MHA/GQA/MQA 的根本工程挑战。

2. MLA KV Cache 的存储结构

2.1 每 token 的缓存内容

MLA 的 KV Cache 不像传统 MHA 那样存储每个 head 的 K 和 V，而是只存储联合压缩的 latent 向量：

2.2 FP8 量化存储 (vLLM)

vLLM 中对 MLA KV Cache 进一步应用了 FP8 量化：

3. 两条注意力路径：Non-Absorbed (Prefill) vs Absorbed (Decode)

4. vLLM 的实现方案 vLLM

4.1 核心代码架构

vLLM 的 MLA 注意力后端位于 vllm/v1/attention/backends/mla/common.py，核心类为 MLACommonBackend。

4.2 _forward_prefill 的两阶段设计

当一个请求同时包含 prefix-cached tokens (context) 和 new tokens (extend) 时，vLLM 把 prefill 拆分为两个阶段：

4.3 merge_attn_states: 合并两阶段结果

两个阶段各自产出 (output, lse)，最后通过 merge_attn_states 进行数值稳定的加权合并：

# 伪代码 — merge_attn_states 的核心逻辑
def merge_attn_states(output_a, lse_a, output_b, lse_b):
    # lse = log-sum-exp of attention scores
    # 利用 log-sum-exp trick 做数值稳定的加权平均
    max_lse = torch.max(lse_a, lse_b)
    weight_a = torch.exp(lse_a - max_lse)
    weight_b = torch.exp(lse_b - max_lse)

    merged_output = (weight_a * output_a + weight_b * output_b) / (weight_a + weight_b)
    return merged_output

4.4 _compute_prefill_context 的分块处理

对于 prefix cached 的 context 部分，vLLM 使用 context chunk 机制来控制内存：

# 伪代码 — context chunk 处理流程
for chunk_idx in range(num_context_chunks):
    # 1. 从 Paged KV Cache 读取一个 chunk 的 c_t 和 k_rope
    cache_kv_c, cache_k_pe = _gather_cache(
        kv_cache, block_table, chunk_start, chunk_len)

    # 2. 解压缩: 通过 kv_b_proj 将 c_t 投影回全量 K/V
    #    kv_b_proj = [W_UK; W_UV] 拼接, shape: [d_c, n_heads*(d_nope+d_v)]
    kv_full = cache_kv_c @ kv_b_proj   # [chunk_len, 512] × [512, n_heads*(128+128)]
    k_nope = kv_full[:, :n_heads*d_nope].reshape(chunk_len, n_heads, d_nope)
    v = kv_full[:, n_heads*d_nope:].reshape(chunk_len, n_heads, d_v)

    # 3. 拼接 K = [K_nope; K_rope] 
    k = torch.cat([k_nope, cache_k_pe.expand(..., n_heads, ...)], dim=-1)

    # 4. 跑 FlashAttention, 得到 chunk_output 和 chunk_lse
    chunk_out, chunk_lse = flash_attention(q_for_context, k, v)

    # 5. 将 chunk 结果累积合并
    ctx_output, ctx_lse = merge_attn_states(
        ctx_output, ctx_lse, chunk_out, chunk_lse)

4.5 vLLM Paged KV Cache 对 MLA 的适配

5. SGLang 的实现方案 SGLang

5.1 RadixAttention 与 MLA KV Cache

SGLang 使用 RadixTree 数据结构来管理 prefix cache（称为 RadixAttention），比 vLLM 的 block hash 方案更高效地支持前缀共享：

5.2 MLA KV Buffer 格式

SGLang 中 MLA 的 KV 缓存以单一 buffer 存储：

# SGLang MLA KV Buffer 结构
kv_buffer: List[Tensor]   # [num_layers]
# 每层 shape: [max_total_tokens, kv_lora_rank + qk_rope_head_dim]
#           = [max_total_tokens, 512 + 64]
#           = [max_total_tokens, 576]
# 单一 buffer, interleaved: [c_t_0..c_t_511 | k_rope_0..k_rope_63]

5.3 Chunked Prefix Cache 优化 (FlashAttention3)

SGLang 的核心创新是 Chunked Prefix Cache 优化，目前仅在 FlashAttention3 后端可用：

5.4 SGLang 支持的 MLA Attention 后端

6. 核心数据流全景图：Prefix Cache → Chunked Prefill → Attention Output

7. 性能对比与工程取舍

7.1 Prefix Cache 收益分析

7.2 解压缩的额外开销

与标准 MHA 的 prefix cache "零开销命中" 不同，MLA 的 prefix cache 命中后仍有解压缩开销：

7.3 vLLM vs SGLang 实现对比

8. 总结

关键技术要点