DeepSeek 的算法杠杆：如何用推理优化撬动硬件生态

import torch
import torch.nn as nn
import math

class LatentKVCompressor(nn.Module):
    """
    简化版 MLA 思路：将 KV 联合压缩到低维 latent 空间，
    推理时只缓存 latent，计算时解压回 K 和 V。
    """

    def __init__(self, n_heads: int, d_model: int, d_latent: int):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.d_latent = d_latent  # 压缩后的 latent 维度，远小于 2 * d_head

        # 压缩投影：KV -> latent（推理时只缓存这一步的输出）
        self.kv_compress = nn.Linear(2 * self.d_head, d_latent, bias=False)

        # 解压投影：latent -> K 和 latent -> V（推理时动态计算，不缓存）
        self.k_decompress = nn.Linear(d_latent, self.d_head, bias=False)
        self.v_decompress = nn.Linear(d_latent, self.d_head, bias=False)

    def forward(self, x: torch.Tensor, cache: torch.Tensor | None = None):
        """
        x: [batch, seq_len, d_model] — 当前 token 的 hidden state（已按 head 拆分）
        cache: [batch, past_len, d_latent] — 之前 token 的 latent KV cache
        返回: attention 输出, 更新后的 latent cache
        """
        batch, seq_len, _ = x.shape

        # 1. 拆分出当前 token 的 K 和 V 原始表示
        #    实际模型中 x 已经是 per-head 的，这里简化处理
        kv_raw = torch.cat([
            x.view(batch, seq_len, self.n_heads, self.d_head).transpose(1, 2),  # K
            x.view(batch, seq_len, self.n_heads, self.d_head).transpose(1, 2),  # V
        ], dim=-1)  # [batch, n_heads, seq_len, 2*d_head]

        # 2. 压缩到 latent 空间 — 这就是我们要缓存的东西
        latent = self.kv_compress(kv_raw)  # [batch, n_heads, seq_len, d_latent]

        # 3. 拼接历史 cache
        if cache is not None:
            latent_full = torch.cat([cache, latent], dim=2)  # [batch, n_heads, past+cur, d_latent]
        else:
            latent_full = latent

        # 4. 动态解压回 K 和 V（不缓存，用完即弃）
        k = self.k_decompress(latent_full)  # [batch, n_heads, past+cur, d_head]
        v = self.v_decompress(latent_full)

        # 5. 标准 attention 计算
        q = x.view(batch, seq_len, self.n_heads, self.d_head).transpose(1, 2)
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_head)
        attn = torch.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)  # [batch, n_heads, seq_len, d_head]

        out = out.transpose(1, 2).contiguous().view(batch, seq_len, -1)
        return out, latent.detach()  # 返回输出 + 新增的 latent cache

    def cache_bytes_per_token(self):
        """每个 token 的 KV cache 字节数（FP16）"""
        return self.n_heads * self.d_latent * 2

    def original_cache_bytes_per_token(self):
        """原始 MHA 每个 token 的 KV cache 字节数（FP16）"""
        return self.n_heads * 2 * self.d_head * 2


# --- 使用示例 ---
compressor = LatentKVCompressor(
    n_heads=32,
    d_model=4096,
    d_latent=128,  # 压缩到 128 维，原始 KV 是 256 维
)

print(f"原始 KV Cache: {compressor.original_cache_bytes_per_token()} bytes/token")
print(f"压缩后 Cache:  {compressor.cache_bytes_per_token()} bytes/token")
print(f"压缩比: {compressor.original_cache_bytes_per_token() / compressor.cache_bytes_per_token():.1f}x")

# 模拟推理：逐 token 生成，只缓存 latent
x = torch.randn(1, 1, 4096)  # 第一个 token
out, cache = compressor(x)
for step in range(10):  # 生成 10 个后续 token
    x = torch.randn(1, 1, 4096)
    out, new_latent = compressor(x, cache=cache)
    cache = torch.cat([cache, new_latent], dim=2)  # 只追加 latent，不追加完整 KV

原始 KV Cache: 4096 bytes/token
压缩后 Cache:  2048 bytes/token
压缩比: 2.0x

# 用 vLLM 部署 DeepSeek-V2（自动启用 MLA KV cache 压缩）
python -m vllm.entrypoints.openai.api_server \
  --model deepseek-ai/DeepSeek-V2 \
  --max-model-len 65536 \
  --gpu-memory-utilization 0.90 \
  --trust-remote-code

# 验证 MLA 是否生效：观察显存占用对比
# 同样 64K 上下文，MLA 下显存占用约为标准 MHA 的 1/25
nvidia-smi --query-gpu=memory.used,memory.total --format=csv -l 5