背景痛点：自回归生成的双重浪费
C模型每生成一个新 token，都要把之前已经算过的 key、value 重新计算一遍。以 7 B 参数、40 层、hidden_size=4096 的模型为例，序列长度从 128 增长到 2048 时，单条请求的 FLOPs 呈平方级放大，GPU 内存占用也从 1.3 GB 膨胀到 21 GB。线上服务若采用静态批处理，还会出现“短序列等长序列”的 bubble 时间，进一步拉低吞吐量。总结起来，冗余主要来自两点：
- 计算冗余：历史 token 的 K/V 被反复重算。
- 内存冗余：激活值与缓存同时驻留显存，无法共享。
下面三条优化策略围绕“减少重复计算、压缩内存、提高批利用率”展开，全部在 PyTorch 2.2 + CUDA 12.1 环境验证通过，硬件为 A100-40 GB。
KV Cache：用空间换时间的最优实践
实现方式很简单：在每一层 attention 模块中维护两个张量cache_k与cache_v，形状[batch, num_heads, max_seq_len, head_dim]。当新 token 到来时，仅对新增位置执行一次 QKV 投影，再把结果拼接到缓存区。
import torch import torch.nn as nn class CachedAttention(nn.Module): def __init__(self, hidden_size: int, num_heads: int, max_seq: int = 2048): super().__init__() self.nh = num_heads self.hd = hidden_size // num_heads self.max_seq = max_seq self.qkv = nn.Linear(hidden_size, 3 * hidden_size, bias=False) self.o = nn.Linear(hidden_size, hidden_size, bias=False) @torch.inference_mode() def forward(self, x: torch.Tensor, layer_idx: int, kv_cache: tuple[torch.Tensor, torch.Tensor], pos: int) -> torch.Tensor: b, t, _ = x.shape qkv = self.qkv(x).chunk(3, dim=-1) q, k, v = [y.view(b, t, self.nh, self.hd).transpose(1, 2) # [b, nh, t, hd] for y in qkv] # 写入缓存 cache_k, cache_v = kv_cache cache_k[:, :, pos:pos+t, :] = k cache_v[:, :, pos:pos+t, :] = v # 读取完整历史 k_full = cache_k[:, :, :pos+t, :] v_full = cache_v[:, :, :pos+t, :] att = (q @ k_full.transpose(-2, -1)) * (self.hd ** -0.5) att = att.softmax(dim=-1) out = (att @ v_full).transpose(1, 2).contiguous().view(b, t, -1) return self.o(out)注意点：
- 使用
@torch.inference_mode()关闭自动求图，显存占用下降 8% 左右。 - 预先分配最大长度缓存，避免
torch.cat带来的碎片化；若长度超过阈值，再一次性torch.empty重新分配并拷贝。
动态批处理：把 bubble 压到最低
静态批要求同一批请求长度对齐，导致 30% 以上计算被填充 token 浪费。Dynamic Batching 采用“连续调度 + 异步填充”策略：
- 维护一个待调度队列，按“先到先服务”排序。
- 每次调度前计算可合并的“长度和”是否小于
max_tokens，若满足则拼接成一个新批。 - 推理完成后，立即把已完成请求弹出，再尝试把新请求插入空缺位置。
核心代码（简化版）：
class ContinuousBatch: def __init__(self, max_tokens: int = 8192): self.max_t = max_tokens self.queue: list[Request] = [] def try_schedule(self) -> torch.Tensor | None: tot = 0 idx = 0 for r in self.queue: if tot + r.len > self.max_t: break tot += r.len idx += 1 if idx == 0: return None batch_seq = [r.tokens for r in self.queue[:idx]] return pad_and_stack(batch_seq) # [bs, max_len] def finish_one(self, rid: int): self.queue = [r for r in self.queue if r.id != rid]线上实测：在平均长度 512、标准差 300 的 Poisson 到达流下，动态批把 GPU 利用率从 58% 提升到 87%，P99 延迟仅增加 6%。
8-bit 量化：精度与速度的再平衡
采用bitsandbytes的Linear8bitLt实现权重 8-bit 存储、16-bit 计算，可在几乎不掉点（perplexity ↑0.02）的情况下，把模型体积从 13 GB 压缩到 3.4 GB，单卡即可部署 2 倍副本。关键是对nn.Linear做 Monkey Patch：
import bitsandbytes as bnb def replace_8bit(model: nn.Module, threshold=6.0): for name, m in model.named_children(): if isinstance(m, nn.Linear): bias = m.bias is not None new_m = bnb.nn.Linear8bitLt( m.in_features, m.out_features, bias=bias, has_fp16_weights=False, threshold=threshold) new_m.weight = bnb.nn.Int8Params( m.weight.data.to(torch.int8), requires_grad=False) if bias: new_m.bias = m.bias setattr(model, name, new_m) else: replace_8bit(m, threshold)精度补偿：对lm_head及第一层embed_tokens保持原始精度，可再降 perplexity 0.01。量化后 KV Cache 同样压缩为 8-bit，显存再省 40%。
性能验证：数据说话
| Batch Size | 基线吞吐 (tok/s) | 优化后吞吐 (tok/s) | 加速比 | 平均延迟 (ms) | GPU 内存 (GB) |
|---|---|---|---|---|---|
| 1 | 31.2 | 97.5 | 3.12× | 256 | 6.8 |
| 8 | 228.0 | 742.3 | 3.26× | 278 | 14.2 |
| 16 | 412.1 | 1320.5 | 3.21× | 295 | 22.5 |
| 32 | OOM | 1980.7 | — | 323 | 38.9 |
测试条件：7 B 模型，序列长度 1024，生成 256 个新 token，A100-40 GB。可见在 32 批规模下，基线已 OOM，而优化方案仍能维持 1980 tok/s 的吞吐。
避坑指南：生产环境的三条经验
KV Cache 碎片
预分配时按“最大长度 × 2”留余量，并定期调用torch.cuda.empty_cache()回收空闲块；若仍出现 OOM，可用cudaMallocAsync版本 PyTorch 2.1+，打开PYTORCH_CUDA_ALLOC_CONF=backend:native。量化后 prompt 编码异常
8-bit 权重仅影响Linear计算，embedding 查表阶段仍用原精度；若发现首 token 延迟飙高，检查是否把nn.Embedding也误替换。CUDA kernel 配置
对 A100，设置setattr(torch.backends.cuda, 'sdp_kernel', 'flash')启用 FlashAttention，再把max_split_size_mb=128可缓解长序列核函数 launch 过多导致的 CPU 调度开销。
延伸思考：Attention 稀疏化还能走多远？
在 8 K 以上长文本场景，KV Cache 再次成为瓶颈。可尝试：
- 局部-全局稀疏：每层仅保留最近 1 K 的密集交互，其余 token 用 1/8 步长稀疏。
- 低秩投影：对历史 key 做 SVD 压缩，把
head_dim映射到 64 维再计算 attention。
实验方法：保持生成长度 4096，对比不同稀疏度下的 perplexity 与吞吐。初步结果显示，稀疏率 50% 时，吞吐再提 1.8×，perplexity 仅增 0.03，值得继续深挖。
动手把方案跑起来
如果你希望一站式体验上述三条优化策略的完整流程，不妨尝试「从0打造个人豆包实时通话AI」动手实验。实验里把 KV Cache、动态批、量化全部封装成可插拔组件，并给出逐行中文注释，本地单卡即可复现。我亲测按照文档走完，半小时就能把 7 B 模型的推理速度提升 3 倍，显存占用降到原来的 1/3，对中级开发者来说非常友好。把代码拉回自己的业务线，再按文内调参表格微调，就能快速拿到生产级收益。
