02 - 推理优化技术¶

class KVCache:
    """KV Cache实现"""

    def __init__(self, num_layers, batch_size, num_heads, max_seq_len, head_dim):
        self.num_layers = num_layers
        self.batch_size = batch_size
        self.num_heads = num_heads
        self.head_dim = head_dim

        # 预分配内存
        self.k_cache = torch.zeros(
            num_layers, batch_size, num_heads, max_seq_len, head_dim
        )
        self.v_cache = torch.zeros(
            num_layers, batch_size, num_heads, max_seq_len, head_dim
        )

        # 当前缓存长度
        self.seq_len = 0

    def update(self, layer_idx, new_k, new_v):
        """
        更新KV Cache

        Args:
            layer_idx: 层索引
            new_k: [batch, num_heads, new_seq_len, head_dim]
            new_v: [batch, num_heads, new_seq_len, head_dim]
        """
        new_seq_len = new_k.size(2)

        # 写入新值
        self.k_cache[layer_idx, :, :, self.seq_len:self.seq_len+new_seq_len] = new_k
        self.v_cache[layer_idx, :, :, self.seq_len:self.seq_len+new_seq_len] = new_v

        # 更新长度
        if layer_idx == 0:  # 只在第一层更新
            self.seq_len += new_seq_len

    def get(self, layer_idx):
        """获取当前层的KV Cache"""
        return (
            self.k_cache[layer_idx, :, :, :self.seq_len],
            self.v_cache[layer_idx, :, :, :self.seq_len]
        )

import torch
import torch.nn as nn

def quantize_tensor(x, num_bits=8):
    """
    简单的对称量化

    Args:
        x: 输入张量
        num_bits: 量化位数

    Returns:
        x_quant: 量化后的张量
        scale: 缩放因子
    """
    qmax = 2 ** (num_bits - 1) - 1
    qmin = -(2 ** (num_bits - 1))

    # 计算缩放因子
    max_val = x.abs().max()
    scale = max_val / qmax

    # 量化
    x_quant = torch.clamp(torch.round(x / scale), qmin, qmax)

    return x_quant, scale

def dequantize_tensor(x_quant, scale):
    """反量化"""
    return x_quant * scale

class QuantizedLinear(nn.Module):
    """量化线性层"""

    def __init__(self, in_features, out_features, num_bits=8):
        super().__init__()  # super()调用父类方法
        self.in_features = in_features
        self.out_features = out_features
        self.num_bits = num_bits

        # 存储量化后的权重
        self.register_buffer('weight_quant', torch.randint(-128, 127, (out_features, in_features)))
        self.register_buffer('weight_scale', torch.tensor(1.0))

        self.bias = nn.Parameter(torch.zeros(out_features))

    def forward(self, x):
        # 反量化权重
        weight = self.weight_quant.float() * self.weight_scale

        # 正常计算
        return F.linear(x, weight, self.bias)

    @torch.no_grad()  # 禁用梯度计算，节省内存（推理时使用）
    def quantize(self, weight_fp32):
        """量化权重"""
        weight_quant, scale = quantize_tensor(weight_fp32, self.num_bits)
        self.weight_quant.copy_(weight_quant)
        self.weight_scale.copy_(scale)

class ContinuousBatchingScheduler:
    """连续批处理调度器"""

    def __init__(self, max_batch_size, max_seq_len):
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len
        self.active_requests = []
        self.waiting_queue = []

    def add_request(self, request):
        """添加新请求到等待队列"""
        self.waiting_queue.append(request)

    def schedule(self):
        """
        调度请求

        返回当前批次的请求列表
        """
        # 1. 移除已完成的请求
        self.active_requests = [r for r in self.active_requests if not r.is_done()]

        # 2. 从等待队列补充新请求
        while (len(self.active_requests) < self.max_batch_size and
               self.waiting_queue):
            request = self.waiting_queue.pop(0)
            self.active_requests.append(request)

        return self.active_requests

    def step(self):
        """执行一步生成"""
        batch = self.schedule()

        # 构建batch输入
        input_ids = [r.get_next_input() for r in batch]

        # 批量推理
        outputs = model.generate_batch(input_ids)

        # 分发结果
        for request, output in zip(batch, outputs):  # zip按位置配对多个可迭代对象
            request.append_token(output)

def speculative_decoding(draft_model, target_model, input_ids, max_new_tokens):
    """
    投机采样

    Args:
        draft_model: 小模型（如7B）
        target_model: 大模型（如70B）
        input_ids: 输入token IDs
        max_new_tokens: 最大生成token数
    """
    accepted_tokens = []

    while len(accepted_tokens) < max_new_tokens:
        # 1. 小模型生成k个候选token
        k = 5
        draft_tokens = draft_model.generate(input_ids, max_new_tokens=k)

        # 2. 大模型验证（一次前向传播）
        target_logits = target_model(input_ids + draft_tokens)

        # 3. 逐个验证
        for i, draft_token in enumerate(draft_tokens):  # enumerate同时获取索引和元素
            # 计算接受概率
            q = draft_model.get_prob(input_ids + accepted_tokens, draft_token)
            p = target_model.get_prob(input_ids + accepted_tokens, draft_token)

            # 接受条件
            if torch.rand(1).item() < min(1, p / q):
                accepted_tokens.append(draft_token)
            else:
                # 拒绝，从调整后的分布采样
                new_token = sample_from_distribution(p - q)
                accepted_tokens.append(new_token)
                break

        # 更新输入
        input_ids = input_ids + accepted_tokens[-k:]

    return accepted_tokens

def benchmark(model, tokenizer, prompts, max_tokens=100):
    """
    基准测试
    """
    import time

    total_tokens = 0
    total_time = 0

    for prompt in prompts:
        inputs = tokenizer(prompt, return_tensors="pt")

        start = time.time()
        outputs = model.generate(**inputs, max_new_tokens=max_tokens)
        end = time.time()

        num_tokens = outputs.shape[1] - inputs.input_ids.shape[1]
        total_tokens += num_tokens
        total_time += end - start

    throughput = total_tokens / total_time
    latency = total_time / len(prompts)

    print(f"吞吐量: {throughput:.2f} tokens/s")
    print(f"平均延迟: {latency:.2f} s")
    print(f"每token延迟: {total_time/total_tokens*1000:.2f} ms")