| # 面试：推理优化深度 |
|
|
| > 本仓库 `src/models/lm/model.py` 的推理实现，覆盖 KV Cache、采样策略、流式生成 |
|
|
| ## 0. 推理流程概览 |
|
|
| ``` |
| 输入 prompt |
| │ |
| ▼ |
| Token Embedding |
| │ |
| ▼ |
| ┌──────────────────────────────┐ |
| │ Prefill 阶段 │ 处理整个 prompt |
| │ (一次性计算所有 token) │ |
| └──────────────┬───────────────┘ |
| │ |
| ▼ |
| ┌──────────────────────────────┐ |
| │ Decode 阶段 │ 逐 token 生成 |
| │ (KV Cache 加速) │ |
| └──────────────┬───────────────┘ |
| │ |
| ▼ |
| 输出序列 |
| ``` |
|
|
| --- |
|
|
| ## Q1. KV Cache 是什么？为什么快？ |
|
|
| ### 核心思想 |
|
|
| 缓存已算的 K/V，每步只算新 token 的 Q 与已有 K/V 做注意力，避免重算。 |
|
|
| ### 本仓库实现（`src/core/attention.py:39-42`） |
|
|
| ```python |
| def forward(self, x, start_pos, freqs_cos, freqs_sin, mask=None): |
| # 计算 Q/K/V |
| xq = self.q_norm(self.q_proj(x)) |
| xk = self.k_norm(self.k_proj(x)) |
| xv = self.v_proj(x) |
| |
| # 应用 RoPE |
| xq, xk = apply_rotary_pos_emb(xq, xk, freqs_cos, freqs_sin) |
| |
| # KV Cache 简单拼接实现 |
| if past_key_value is not None: |
| xk = torch.cat([past_key_value[0], xk], dim=2) |
| xv = torch.cat([past_key_value[1], xv], dim=2) |
| past_key_value = (xk, xv) |
| |
| # 计算注意力 |
| attn_output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=mask) |
| return attn_output, past_key_value |
| ``` |
|
|
| ### 显存节省 |
|
|
| 假设： |
| - batch_size=B, seq_len=S, num_layers=L |
| - num_kv_heads=H, head_dim=D |
| - 精度=fp16（2 bytes） |
|
|
| **无 KV Cache**： |
| - 每步计算量 = B × S² × L × H × D |
| - 显存 = B × S × L × H × D × 2 bytes |
|
|
| **有 KV Cache**： |
| - 每步计算量 = B × S × L × H × D |
| - 显存 = B × S × L × H × D × 2 bytes（但只需算一次） |
|
|
| > 面试点：KV Cache 为什么能加速？→ 避免重复计算 K/V，每步只需算新 token 的 Q |
|
|
| --- |
|
|
| ## Q2. Prefill vs Decode 阶段 |
|
|
| ### Prefill 阶段 |
|
|
| - 处理整个 prompt |
| - 一次性计算所有 token 的 K/V |
| - 计算量大，但只需做一次 |
|
|
| ### Decode 阶段 |
|
|
| - 逐 token 生成 |
| - 每步只算新 token 的 Q |
| - 计算量小，但需要很多步 |
|
|
| ### 本仓库实现（`src/models/lm/model.py:73-113`） |
|
|
| ```python |
| def generate(self, input_ids, max_new_tokens=200, temperature=0.6, top_k=5, top_p=0.8): |
| for _ in range(max_new_tokens): |
| # Prefill 阶段：处理整个 prompt |
| if idx_cond.shape[1] > 1: |
| logits, _ = self(idx_cond) |
| # Decode 阶段：只处理最后一个 token |
| else: |
| logits, _ = self(idx_cond, start_pos=start_pos) |
| |
| # 采样 |
| logits = logits[:, -1, :] / temperature |
| idx_next = self.sample(logits, top_k=top_k, top_p=top_p) |
| |
| # 更新序列 |
| idx_cond = torch.cat([idx_cond, idx_next], dim=1) |
| start_pos += 1 |
| ``` |
|
|
| > 面试点：为什么 Prefill 和 Decode 要分开处理？→ Prefill 可以并行处理所有 token，Decode 只能逐 token 处理 |
|
|
| --- |
|
|
| ## Q3. Logits to Keep 优化（`src/models/lm/model.py:65-66`） |
|
|
| ### 问题 |
|
|
| 在生成时，我们只需要最后一个 token 的 logits，但标准实现会计算所有 token 的 logits。 |
|
|
| ### 本仓库优化 |
|
|
| ```python |
| def forward(self, input_ids, ..., logits_to_keep=1): |
| # 仅计算最后 N 个 token 的 logits |
| hidden_states = hidden_states[:, -logits_to_keep:] |
| logits = self.lm_head(hidden_states) |
| ``` |
|
|
| ### 显存节省 |
|
|
| 假设 vocab_size=64000，seq_len=32768： |
| - 不优化：`32768 × 64000 × 2 bytes ≈ 4GB` |
| - 优化后：`1 × 64000 × 2 bytes ≈ 128KB` |
|
|
| > 面试点：为什么可以只算最后 N 个？→ 生成时只需要最后一个 token 的 logits 来采样下一个 token |
|
|
| --- |
|
|
| ## Q4. 采样策略 |
|
|
| ### Top-k 采样 |
|
|
| ```python |
| def top_k_logits(logits, k): |
| # 只保留概率最高的 k 个 token |
| values, indices = torch.topk(logits, k) |
| # 其他 token 设为 -inf |
| logits[logits < values[:, -1:]] = float('-inf') |
| return logits |
| ``` |
|
|
| ### Top-p 采样（Nucleus Sampling） |
|
|
| ```python |
| def top_p_logits(logits, p): |
| # 按概率排序 |
| sorted_logits, sorted_indices = torch.sort(logits, descending=True) |
| cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) |
| |
| # 移除累积概率超过 p 的 token |
| sorted_indices_to_remove = cumulative_probs > p |
| sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() |
| sorted_indices_to_remove[..., 0] = 0 |
| |
| # 恢复原始顺序 |
| indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) |
| logits[indices_to_remove] = float('-inf') |
| return logits |
| ``` |
|
|
| ### Temperature 采样 |
|
|
| ```python |
| def temperature_scale(logits, temperature): |
| # temperature < 1: 分布更尖锐，更确定 |
| # temperature > 1: 分布更平滑，更随机 |
| return logits / temperature |
| ``` |
|
|
| > 面试点：Top-k 和 Top-p 的区别？→ Top-k 保留固定数量的 token，Top-p 保留累积概率达到 p 的 token |
|
|
| --- |
|
|
| ## Q5. Repetition Penalty |
|
|
| ### 问题 |
|
|
| 生成时可能重复之前的 token，导致输出质量下降。 |
|
|
| ### 本仓库实现（`src/models/lm/model.py:73-113`） |
|
|
| ```python |
| def generate(self, input_ids, ..., repetition_penalty=1.2): |
| # 对已生成的 token 施加惩罚 |
| if input_ids.shape[1] > 1: |
| # 计算每个 token 的出现次数 |
| for i in range(input_ids.shape[1]): |
| token_id = input_ids[0, i].item() |
| logits[0, token_id] /= repetition_penalty |
| ``` |
|
|
| ### 为什么用除法？ |
|
|
| - 除以惩罚系数，降低已生成 token 的概率 |
| - 保持概率分布的相对顺序 |
|
|
| > 面试点：repetition_penalty=1.0 表示什么？→ 无惩罚，等于没有使用 repetition penalty |
| |
| --- |
| |
| ## Q6. 流式生成（Streaming） |
| |
| ### 本仓库实现（`src/models/lm/model.py:73-113`） |
| |
| ```python |
| def generate(self, input_ids, ..., stream_callback=None): |
| for _ in range(max_new_tokens): |
| # ... 前向传播 ... |
| |
| # 流式回调 |
| if stream_callback: |
| stream_callback(idx_next) |
| |
| # 更新序列 |
| idx_cond = torch.cat([idx_cond, idx_next], dim=1) |
| ``` |
| |
| ### 为什么需要流式生成？ |
|
|
| 1. **用户体验**：实时看到生成结果 |
| 2. **早停**：用户可以在生成完成前停止 |
| 3. **调试**：实时观察生成过程 |
|
|
| > 面试点：流式生成如何实现？→ 通过回调函数，在每步生成后返回当前 token |
|
|
| --- |
|
|
| ## Q7. GQA 对推理的影响 |
|
|
| ### KV Cache 节省 |
|
|
| 假设： |
| - num_attention_heads = 8 |
| - num_key_value_heads = 4 |
| - head_dim = 64 |
|
|
| **MHA（Multi-Head Attention）**： |
| - KV Cache = 2 × B × S × L × 8 × 64 × 2 bytes |
|
|
| **GQA（Grouped-Query Attention）**： |
| - KV Cache = 2 × B × S × L × 4 × 64 × 2 bytes |
|
|
| **节省**：50% |
|
|
| ### 本仓库实现（`src/core/attention.py:13-16`） |
|
|
| ```python |
| class Attention(nn.Module): |
| def __init__(self, config): |
| self.n_local_heads = config.num_attention_heads # 8 |
| self.n_local_kv_heads = config.num_key_value_heads # 4 |
| self.n_rep = self.n_local_heads // self.n_local_kv_heads # 2 倍复制 |
| ``` |
|
|
| > 面试点：GQA 如何减少 KV Cache？→ KV 头数从 n_heads 减到 n_kv_heads，KV Cache 减少 n_heads/n_kv_heads 倍 |
|
|
| --- |
|
|
| ## Q8. Flash-Attention 对推理的影响 |
|
|
| ### 核心思想 |
|
|
| IO 感知的注意力计算，避免物化完整的 N×N 注意力矩阵。 |
|
|
| ### 本仓库的 Flash-Attention 条件（`src/core/attention.py:28, 44`） |
|
|
| ```python |
| if (seq_len > 1 and |
| (not self.causal or past_key_value is None) and |
| attention_mask is None): |
| # 使用 Flash Attention |
| ``` |
|
|
| ### 为什么有条件限制？ |
|
|
| 1. `seq_len > 1`：单 token 无需注意力 |
| 2. `not self.causal or past_key_value is None`：Flash Attention 对 causal mask 支持有限 |
| 3. `attention_mask is None`：Flash Attention 不支持自定义 mask |
|
|
| > 面试点：Flash-Attention 对推理有什么好处？→ 减少 HBM 读写，降低延迟 |
|
|
| --- |
|
|
| ## Q9. 批量推理优化 |
|
|
| ### 问题 |
|
|
| 逐条推理效率低，需要批量处理。 |
|
|
| ### 解决方案 |
|
|
| 1. **Padding**：将不同长度的序列 padding 到统一长度 |
| 2. **Dynamic Batching**：根据序列长度动态调整 batch size |
| 3. **Continuous Batching**：不等待整个 batch 完成，动态添加新请求 |
|
|
| ### 本仓库的批量推理（`src/models/lm/model.py:73-113`） |
|
|
| ```python |
| def generate(self, input_ids, ...): |
| # 支持 batch_size > 1 |
| for _ in range(max_new_tokens): |
| logits, _ = self(idx_cond) |
| # ... 采样 ... |
| ``` |
|
|
| > 面试点：Padding 的缺点是什么？→ 浪费计算资源，短序列需要 padding 到长序列长度 |
|
|
| --- |
|
|
| ## Q10. 量化推理 |
|
|
| ### 问题 |
|
|
| FP16 精度显存占用高，推理速度慢。 |
|
|
| ### 解决方案 |
|
|
| 1. **INT8 量化**：将权重从 FP16 量化到 INT8 |
| 2. **INT4 量化**：将权重从 FP16 量化到 INT4 |
| 3. **GPTQ**：基于二阶信息的量化方法 |
| 4. **AWQ**：激活感知的量化方法 |
|
|
| ### 本仓库的量化支持 |
|
|
| ```python |
| # 通过 config.dtype 控制精度 |
| config.dtype = 'float16' # FP16 |
| config.dtype = 'bfloat16' # BF16 |
| ``` |
|
|
| > 面试点：量化会损失多少精度？→ 取决于量化方法和位数，INT8 通常损失很小，INT4 可能有明显损失 |
|
|
| --- |
|
|
| ## Q11. 推理显存估算 |
|
|
| ### KV Cache 显存 |
|
|
| 假设： |
| - batch_size=B, seq_len=S, num_layers=L |
| - num_kv_heads=H, head_dim=D |
| - 精度=fp16（2 bytes） |
|
|
| KV Cache 显存 = `2 × B × S × L × H × D × 2 bytes` |
|
|
| ### 模型参数显存 |
|
|
| 假设： |
| - hidden_size=d, num_layers=L |
| - 精度=fp16（2 bytes） |
|
|
| 模型参数显存 = `12 × L × d × d × 2 bytes`（Q/K/V/O 四个投影矩阵） |
|
|
| ### 总显存 |
|
|
| 总显存 ≈ 模型参数显存 + KV Cache 显存 |
|
|
| > 面试点：如何估算推理显存？→ 模型参数显存 + KV Cache 显存 |
|
|
| --- |
|
|
| ## Q12. 推理延迟估算 |
|
|
| ### Prefill 延迟 |
|
|
| 假设： |
| - batch_size=B, seq_len=S, hidden_size=d |
| - FLOPS = 2 × B × S² × d |
| |
| ### Decode 延迟 |
| |
| 假设： |
| - batch_size=B, hidden_size=d |
| - FLOPS = 2 × B × S × d |
| |
| ### 总延迟 |
| |
| 总延迟 ≈ Prefill 延迟 + Decode 延迟 × 生成长度 |
| |
| > 面试点：如何优化推理延迟？→ 使用 Flash-Attention、KV Cache、量化等方法 |
| |
| --- |
| |
| ## Q13. 推理服务设计 |
| |
| ### 问题 |
| |
| 如何设计高并发的推理服务？ |
| |
| ### 解决方案 |
| |
| 1. **动态批处理**：根据请求到达时间动态组 batch |
| 2. **请求排队**：使用消息队列管理请求 |
| 3. **负载均衡**：将请求分发到多个 GPU |
| 4. **模型并行**：将模型拆分到多个 GPU |
| |
| ### 本仓库的推理服务（`src/serve/`） |
| |
| ```python |
| class RealtimeSession: |
| def __init__(self, model): |
| self.model = model |
| self.vad = SileroVAD() # 语音活动检测 |
| |
| def process(self, audio): |
| # 1. VAD 检测 |
| if not self.vad.detect(audio): |
| return None |
| |
| # 2. 推理 |
| output = self.model.generate(audio) |
| |
| return output |
| ``` |
| |
| > 面试点：如何提高推理吞吐量？→ 使用动态批处理、模型并行、量化等方法 |
| |
| --- |
| |
| ## Q14. 推理与训练的区别 |
| |
| ### 训练 |
| |
| - 需要反向传播 |
| - 需要梯度存储 |
| - 需要优化器状态 |
| - 显存占用高 |
| |
| ### 推理 |
| |
| - 只需要前向传播 |
| - 不需要梯度存储 |
| - 不需要优化器状态 |
| - 显存占用低 |
| |
| ### 本仓库的切换 |
| |
| ```python |
| # 训练时 |
| model.train() |
| loss = model(input_ids, labels=labels) |
|
|
| # 推理时 |
| model.eval() |
| with torch.no_grad(): |
| logits = model(input_ids) |
| ``` |
| |
| > 面试点：为什么推理时要 `torch.no_grad()`？→ 节省显存，避免存储梯度 |
| |
| --- |
| |
| ## Q15. 推理优化的未来方向 |
| |
| ### 当前瓶颈 |
| |
| 1. **内存墙**：显存带宽限制推理速度 |
| 2. **计算墙**：GPU 计算能力限制吞吐量 |
| 3. **延迟墙**：逐 token 生成限制响应速度 |
| |
| ### 未来方向 |
| |
| 1. **投机采样**：用小模型预测，大模型验证 |
| 2. **模型并行**：将模型拆分到多个 GPU |
| 3. **硬件优化**：使用专用推理芯片 |
| 4. **算法优化**：设计更高效的注意力机制 |
| |
| > 面试点：投机采样是什么？→ 用小模型快速生成候选，大模型验证并选择，提高生成速度 |
| |
