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
import torch.nn.functional as F

def quantize_tensor(x, num_bits=8):
    """
    简单的对称量化（per-tensor）。

    说明:
    - 这是教学版实现，重点展示缩放和裁剪逻辑
    - 4-bit 这里仍用 int8 容器存储，方便演示；真实 4-bit 部署通常会做打包存储
    - 生产环境常见的是 per-channel / per-group 量化，而不是单一全局 scale

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

    Returns:
        x_quant: 量化后的张量
        scale: 缩放因子
    """
    if not 2 <= num_bits <= 8:
        raise ValueError("num_bits must be between 2 and 8")

    x = x.float()
    qmax = 2 ** (num_bits - 1) - 1
    qmin = -(2 ** (num_bits - 1))

    # 计算缩放因子
    max_val = x.abs().max()
    scale = (max_val / qmax).clamp_min(torch.finfo(x.dtype).eps)

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

    return x_quant.to(torch.int8), scale

def dequantize_tensor(x_quant, scale):
    """反量化"""
    return x_quant.float() * 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

        # 用int8作为教学示例的通用容器
        self.register_buffer(
            'weight_quant',
            torch.zeros(out_features, in_features, dtype=torch.int8)
        )
        self.register_buffer('weight_scale', torch.tensor(1.0, dtype=torch.float32))

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

    def forward(self, x):
        # 反量化权重
        weight = dequantize_tensor(self.weight_quant, 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)

import torch

def normalize_positive_part(p: torch.Tensor, q: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
    """
    将 [p - q]_+ 重新归一化为合法分布。
    p, q 均为 shape [vocab_size] 的概率向量。
    """
    diff = torch.clamp(p - q, min=0.0)
    z = diff.sum()
    if z.item() <= eps:
        return p / p.sum().clamp_min(eps)
    return diff / z

@torch.no_grad()
def speculative_accept_reject(draft_tokens, draft_probs, target_probs):
    """
    对一段 draft proposals 执行一次接受/拒绝步骤。

    Args:
        draft_tokens: 长度为 k 的候选 token 序列
        draft_probs: 长度为 k 的列表，第 t 项是 q_t(.)
        target_probs: 长度为 k+1 的列表，第 t 项是 p_t(.)，
                      最后一项用于"全部接受后再补采一个 token"
    """
    accepted = []

    for step, draft_token in enumerate(draft_tokens):
        token_id = int(draft_token)
        q_t = draft_probs[step]
        p_t = target_probs[step]

        accept_prob = min(
            1.0,
            (p_t[token_id] / q_t[token_id].clamp_min(1e-12)).item(),
        )

        if torch.rand(()).item() < accept_prob:
            accepted.append(token_id)
            continue

        corrected_dist = normalize_positive_part(p_t, q_t)
        replacement = torch.multinomial(corrected_dist, num_samples=1).item()
        return accepted + [replacement], False

    # 如果 k 个候选都被接受，再从 target 的下一步分布补采 1 个 token
    extra_token = torch.multinomial(target_probs[len(draft_tokens)], num_samples=1).item()
    return accepted + [extra_token], True

def speculative_decoding_step(prefix_ids, draft_model, target_model, k=4):
    """
    教学版接口约定:
    - draft_model.propose(prefix_ids, k) -> (draft_tokens, draft_probs)
    - target_model.score(prefix_ids, draft_tokens) -> target_probs
      其中 target_probs 长度为 k+1
    """
    draft_tokens, draft_probs = draft_model.propose(prefix_ids, k)
    target_probs = target_model.score(prefix_ids, draft_tokens)
    return speculative_accept_reject(draft_tokens, draft_probs, target_probs)

from vllm import LLM, SamplingParams

# 使用 n-gram 投机解码（无需额外 draft model）
llm = LLM(
    model="Qwen/Qwen2.5-7B-Instruct",
    tensor_parallel_size=1,
    speculative_model="[ngram]",       # 使用 n-gram 模型
    num_speculative_tokens=5,          # 预测 5 个 token
    ngram_prompt_lookup_max=4,         # 用前 4 个 token 预测下一个
)

sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
outputs = llm.generate(["The future of AI is"], sampling_params)

# 也可以使用小型 draft model
# speculative_model="Qwen/Qwen2.5-0.5B"

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")