【大模型】- 推理优化

宋 算法
  2026-08-19 09:12 2026-08-19 09:12 创建   2026-08-21 16:20:12 2026-08-21 16:20:12 更新      2.1k 字  9 分钟

推理优化

加速语言模型推理

类型： 学习　｜　语言： Python　｜　🏷 前置：《推理优化》（本系列第 11 篇）

学习目标

  • 理解推理优化的关键技术
  • 实现 KV 缓存加速自回归生成
  • 使用连续批处理提高吞吐量
  • 优化注意力计算
  • 部署高性能推理服务

推理瓶颈

LLM 推理的主要瓶颈：

  1. 内存带宽：加载模型权重
  2. 计算：注意力和 FFN 计算
  3. 延迟：逐 token 生成
  4. 批处理效率：GPU 利用率

KV 缓存

KV 缓存是最重要的优化之一：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class KVCache:
    def __init__(self, max_batch_size, max_seq_len, n_heads, head_dim, device):
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len
        
        # 预分配缓存
        self.key_cache = torch.zeros(
            max_batch_size, n_heads, max_seq_len, head_dim,
            device=device
        )
        self.value_cache = torch.zeros(
            max_batch_size, n_heads, max_seq_len, head_dim,
            device=device
        )
        
        self.cur_len = 0
    
    def update(self, input_pos, k_val, v_val):
        """更新缓存"""
        batch_size = k_val.shape[0]
        
        self.key_cache[:batch_size, :, input_pos:input_pos + 1, :] = k_val
        self.value_cache[:batch_size, :, input_pos:input_pos + 1, :] = v_val
        
        self.cur_len = input_pos + 1
    
    def get(self, seq_len):
        """获取缓存"""
        return self.key_cache[:, :, :seq_len, :], self.value_cache[:, :, :seq_len, :]

使用 KV 缓存的注意力

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
def attention_with_cache(query, key, value, kv_cache=None, input_pos=None):
    """使用 KV 缓存的注意力计算"""
    if kv_cache is not None and input_pos is not None:
        # 更新缓存
        kv_cache.update(input_pos, key, value)
        
        # 获取完整缓存
        key, value = kv_cache.get(input_pos + 1)
    
    # 标准注意力
    attn_weights = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(key.shape[-1])
    attn_weights = torch.softmax(attn_weights, dim=-1)
    attn_output = torch.matmul(attn_weights, value)
    
    return attn_output

连续批处理

连续批处理允许动态添加和移除请求：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
class ContinuousBatchScheduler:
    def __init__(self, max_batch_size, max_seq_len):
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len
        
        self.pending_requests = []
        self.active_requests = []
        self.completed_requests = []
    
    def add_request(self, request):
        """添加新请求"""
        self.pending_requests.append(request)
    
    def schedule(self):
        """调度请求"""
        # 尝试填充批次
        while (len(self.active_requests) < self.max_batch_size and 
               self.pending_requests):
            request = self.pending_requests.pop(0)
            
            # 检查是否有足够的内存
            if self.can_schedule(request):
                self.active_requests.append(request)
            else:
                self.pending_requests.insert(0, request)
                break
    
    def can_schedule(self, request):
        """检查是否可以调度请求"""
        # 检查序列长度
        if request.seq_len + 1 > self.max_seq_len:
            return False
        
        # 检查内存（简化）
        return True
    
    def step(self, model):
        """执行一步推理"""
        if not self.active_requests:
            return []
        
        # 准备批次
        batch = self.prepare_batch()
        
        # 前向传播
        with torch.no_grad():
            logits = model(**batch)
        
        # 处理输出
        completed = self.process_outputs(logits)
        
        return completed
    
    def prepare_batch(self):
        """准备批次输入"""
        input_ids = []
        attention_mask = []
        
        for request in self.active_requests:
            input_ids.append(request.input_ids)
            attention_mask.append(request.attention_mask)
        
        # 填充到最大长度
        max_len = max(len(ids) for ids in input_ids)
        
        padded_input_ids = []
        padded_attention_mask = []
        
        for ids, mask in zip(input_ids, attention_mask):
            padding_length = max_len - len(ids)
            padded_input_ids.append(
                torch.cat([ids, torch.zeros(padding_length, dtype=torch.long)])
            )
            padded_attention_mask.append(
                torch.cat([mask, torch.zeros(padding_length, dtype=torch.long)])
            )
        
        return {
            "input_ids": torch.stack(padded_input_ids),
            "attention_mask": torch.stack(padded_attention_mask)
        }

Flash Attention

Flash Attention 优化注意力计算：

1
2
3
4
5
6
7
8
9
10
11
from torch.nn.functional import scaled_dot_product_attention

def flash_attention(query, key, value, is_causal=True):
    """使用 Flash Attention"""
    return scaled_dot_product_attention(
        query, key, value,
        is_causal=is_causal,
        enable_flash=True,
        enable_math=False,
        enable_mem_efficient=False
    )

Flash Attention 2

1
2
3
4
5
6
7
8
9
# Flash Attention 2 需要安装
# pip install flash-attn

from flash_attn import flash_attn_func

def flash_attention_2(query, key, value):
    """Flash Attention 2"""
    # query, key, value: (batch, seqlen, nheads, headdim)
    return flash_attn_func(query, key, value, causal=True)

投机解码

使用小模型猜测，大模型验证：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
class SpeculativeDecoder:
    def __init__(self, draft_model, target_model, tokenizer, k=5):
        self.draft_model = draft_model
        self.target_model = target_model
        self.tokenizer = tokenizer
        self.k = k  # 每次猜测的 token 数
    
    def generate(self, prompt, max_new_tokens):
        """投机解码生成"""
        input_ids = self.tokenizer.encode(prompt, return_tensors="pt")
        
        generated = input_ids.clone()
        
        while generated.shape[1] < max_new_tokens:
            # 小模型猜测 k 个 token
            draft_tokens = self.draft_generate(generated, self.k)
            
            # 大模型验证
            accepted = self.target_verify(generated, draft_tokens)
            
            # 更新生成序列
            generated = torch.cat([generated, accepted], dim=1)
            
            if self.tokenizer.eos_token_id in accepted:
                break
        
        return generated
    
    def draft_generate(self, input_ids, k):
        """小模型生成 k 个 token"""
        tokens = []
        
        for _ in range(k):
            with torch.no_grad():
                outputs = self.draft_model(input_ids)
                logits = outputs.logits[:, -1, :]
                probs = torch.softmax(logits, dim=-1)
                token = torch.multinomial(probs, 1)
                tokens.append(token)
                input_ids = torch.cat([input_ids, token], dim=1)
        
        return torch.cat(tokens, dim=1)
    
    def target_verify(self, input_ids, draft_tokens):
        """大模型验证"""
        # 拼接输入
        full_input = torch.cat([input_ids, draft_tokens], dim=1)
        
        with torch.no_grad():
            outputs = self.target_model(full_input)
            logits = outputs.logits[:, input_ids.shape[1]-1:-1, :]
        
        # 逐个验证
        accepted = []
        for i in range(draft_tokens.shape[1]):
            probs = torch.softmax(logits[:, i, :], dim=-1)
            draft_token = draft_tokens[:, i]
            
            # 计算接受概率
            draft_prob = probs.gather(1, draft_token.unsqueeze(1))
            target_prob = torch.softmax(
                self.target_model(full_input[:, :input_ids.shape[1]+i]).logits[:, -1, :],
                dim=-1
            ).gather(1, draft_token.unsqueeze(1))
            
            accept_prob = min(1.0, target_prob.item() / draft_prob.item())
            
            if torch.rand(1).item() < accept_prob:
                accepted.append(draft_token.unsqueeze(1))
            else:
                # 从目标模型采样
                new_token = torch.multinomial(target_prob, 1)
                accepted.append(new_token)
                break
        
        return torch.cat(accepted, dim=1) if accepted else draft_tokens[:, :1]

推理服务

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn

app = FastAPI()

class GenerationRequest(BaseModel):
    prompt: str
    max_new_tokens: int = 512
    temperature: float = 0.7
    top_p: float = 0.9

class GenerationResponse(BaseModel):
    generated_text: str
    tokens_generated: int
    time_taken: float

@app.post("/generate", response_model=GenerationResponse)
async def generate(request: GenerationRequest):
    """生成文本"""
    try:
        # 分词
        inputs = tokenizer(request.prompt, return_tensors="pt").to(model.device)
        
        # 生成
        import time
        start_time = time.time()
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=request.max_new_tokens,
                temperature=request.temperature,
                top_p=request.top_p,
                do_sample=True
            )
        
        time_taken = time.time() - start_time
        
        # 解码
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        tokens_generated = outputs.shape[1] - inputs["input_ids"].shape[1]
        
        return GenerationResponse(
            generated_text=generated_text,
            tokens_generated=tokens_generated,
            time_taken=time_taken
        )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# 启动服务
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

性能监控

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import time
from dataclasses import dataclass

@dataclass
class InferenceMetrics:
    total_requests: int = 0
    total_tokens: int = 0
    total_time: float = 0.0
    avg_latency: float = 0.0
    throughput: float = 0.0

class InferenceMonitor:
    def __init__(self):
        self.metrics = InferenceMetrics()
        self.start_time = time.time()
    
    def record_request(self, tokens_generated, time_taken):
        """记录请求"""
        self.metrics.total_requests += 1
        self.metrics.total_tokens += tokens_generated
        self.metrics.total_time += time_taken
        
        # 更新平均值
        self.metrics.avg_latency = (
            self.metrics.total_time / self.metrics.total_requests
        )
        
        elapsed = time.time() - self.start_time
        self.metrics.throughput = self.metrics.total_tokens / elapsed
    
    def print_stats(self):
        """打印统计信息"""
        print(f"总请求数: {self.metrics.total_requests}")
        print(f"总 token 数: {self.metrics.total_tokens}")
        print(f"平均延迟: {self.metrics.avg_latency:.4f}s")
        print(f"吞吐量: {self.metrics.throughput:.2f} tokens/s")

优化技术总结

技术 效果 复杂度 适用场景
KV 缓存 高 低 所有场景
Flash Attention 高 中 长序列
连续批处理 中 高 服务部署
投机解码 中 高 低延迟
量化 高 中 内存受限

总结

推理优化通过多种技术加速 LLM 推理。KV 缓存、Flash Attention 和连续批处理是关键优化。选择合适的优化技术需要平衡延迟、吞吐量和资源约束。

下一步

下一课将构建一个完整的 LLM 管道，整合所有优化技术。

📚 本文改编自 AI Engineering from Scratch（MIT License · 作者 Rohit Ghumare），中文内容来自官方中文镜像。原课程共 503 课 · 20 阶段 · 免费开源，教程网站见 aiengineeringfromscratch.com。

📝 自我检查（课程配套测验）

Q1（学前） LLM 推理的主要瓶颈是什么？

A. 磁盘 I/O
B. 内存带宽——自回归生成需逐 token 读取全部模型权重
C. 网络延迟
D. 分词速度

答案： B　解析： LLM 推理是 memory-bound：每生成一个 token 需读取全部模型权重。7B FP16 模型每 token 约 14GB 内存读取。优化重点是减少内存传输而非 raw 计算。

Q2（学前） KV cache 在 LLM 推理中的作用是什么？

A. 存储模型权重
B. 缓存先前 token 的 key 和 value 向量，避免每步重复计算
C. 缓存最终输出文本
D. 存储训练数据

答案： B　解析： 自回归生成中，先前 token 的 K/V 向量不变。KV cache 存储它们，每步只计算新 token 的 K/V。这将对已见 token 的 O(n) 注意力计算变为 O(1)，但随序列长度增加内存。

Q3（学后） 什么是 speculative decoding（推测解码）？

A. 用更大模型生成所有 token
B. 小 draft 模型快速生成候选 token，大 target 模型并行验证，接受匹配 token 以加速
C. 随机猜测 token
D. 跳过某些层

答案： B　解析： 推测解码用小/fast draft 模型提出多个候选 token，大 target 模型一次前向并行验证。接受的 token 使有效吞吐量超过逐 token 生成，draft 与 target 分布接近时尤其有效。

Q4（学后） 连续 batching 为什么比静态 batching 更适合 LLM 服务？

A. 它使用更大 batch
B. 请求完成即从 batch 移除、新请求加入，提高 GPU 利用率——静态 batching 需等最长请求完成
C. 它消除 KV cache
D. 它只适用于训练

答案： B　解析： 静态 batching 中，batch 内一个长请求会阻塞其他已完成的请求。连续 batching（vLLM、TGI）动态管理 batch 组成，完成的序列立即被新请求替换，显著提高吞吐量。

Q5（学后） PagedAttention（vLLM）解决什么问题？

A. 模型训练速度
B. KV cache 内存碎片和浪费——像 OS 虚拟内存一样分页管理 KV cache
C. 网络带宽
D. 分词效率

答案： B　解析： KV cache 为最大序列长度预分配，短请求浪费内存。PagedAttention 将 KV cache 分成固定大小块，按需分配，类似虚拟内存分页，在相同 GPU 内存下显著提高 batch 大小和吞吐量。

  • 标题: 【大模型】- 推理优化
  • 作者: 宋
  • 创建于 : 2026-08-19 09:12:00
  • 更新于 : 2026-08-21 16:20:12
  • 链接: https://sxl-space.tk/2026/08/19/010_LLM/010_LLM-12-InferenceOptimization/
  • 版权声明: 版权所有 © 宋，禁止转载。
目录
【大模型】- 推理优化
  1. 1. 推理优化
  2. 2. 学习目标
  3. 3. 推理瓶颈
  4. 4. KV 缓存
    1. 4.1. 使用 KV 缓存的注意力
  5. 5. 连续批处理
  6. 6. Flash Attention
    1. 6.1. Flash Attention 2
  7. 7. 投机解码
  8. 8. 推理服务
  9. 9. 性能监控
  10. 10. 优化技术总结
  11. 11. 总结
  12. 12. 下一步
  13. 13. 下一课将构建一个完整的 LLM 管道，整合所有优化技术。
  14. 14. 📝 自我检查（课程配套测验）