Inference - 控制输出质量

"推理通过策略性的 token 选择，将概率分布转化为连贯的文本。"

训练构建了模型，但推理决定了它的输出。理解解码策略、采样参数和优化技术对于在生产环境中控制模型行为至关重要。本文档涵盖自回归生成、解码算法、采样参数以及大规模部署 LLM 的性能优化技术。

自回归生成

生成循环

LLM 以自回归方式生成文本——一次一个 token，每个新 token 以所有先前 token 为条件。

基本实现

import torch
import torch.nn.functional as F

def generate_autoregressive(
    model,
    input_ids: torch.Tensor,
    max_new_tokens: int = 100,
    temperature: float = 1.0,
    top_k: int = None,
    top_p: float = 1.0,
    eos_token_id: int = None
) -> torch.Tensor:
    """
    自回归生成 token。

    Args:
        model: 语言模型
        input_ids: 输入 token ID (batch, seq_len)
        max_new_tokens: 最大生成 token 数
        temperature: 采样温度
        top_k: 保留前 k 个 token
        top_p: 核采样阈值
        eos_token_id: 序列结束 token ID
    """
    batch_size, seq_len = input_ids.shape
    current_ids = input_ids.clone()

    for step in range(max_new_tokens):
        # 前向传播
        with torch.no_grad():
            outputs = model(current_ids)
            logits = outputs.logits[:, -1, :]  # (batch, vocab_size)

        # 应用温度
        logits = logits / temperature

        # 应用 top-k 过滤
        if top_k is not None:
            top_k_logits, top_k_indices = torch.topk(logits, top_k)
            logits = torch.full_like(logits, float('-inf'))
            logits.scatter_(1, top_k_indices, top_k_logits)

        # 应用 top-p（核采样）过滤
        if top_p < 1.0:
            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

            # 移除累积概率超过阈值的 token
            sorted_indices_to_remove = cumulative_probs > top_p
            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
            sorted_indices_to_remove[..., 0] = False

            indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
            logits[indices_to_remove] = float('-inf')

        # 采样 token
        probs = F.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)  # (batch, 1)

        # 追加到序列
        current_ids = torch.cat([current_ids, next_token], dim=1)

        # 检查 EOS
        if eos_token_id is not None and (next_token == eos_token_id).all():
            break

    return current_ids

# 使用示例
input_text = "The future of AI is"
input_ids = tokenizer.encode(input_text, return_tensors="pt")
output_ids = generate_autoregressive(model, input_ids, max_new_tokens=50, temperature=0.8)
output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
print(output_text)

解码策略

策略对比

策略	描述	速度	质量	多样性	适用场景
贪心搜索	始终选择最高概率	最快	适合事实性任务	无	事实问答、代码
束搜索	保留前 k 个假设	慢	高质量	低	翻译、摘要
采样	从概率分布中采样	快	可变	高	创意写作
核采样 (Top-p)	从最小顶部概率质量中采样	快	好	高	通用聊天、助手
Top-k	从前 k 个 token 中采样	快	好	中高	平衡生成
对比搜索	平衡概率 + 退化惩罚	中	很高	中	长文本内容
MCTS	带前瞻的树搜索	很慢	最佳	中	复杂推理

贪心搜索

def greedy_search(model, input_ids: torch.Tensor, max_new_tokens: int) -> torch.Tensor:
    """
    贪心解码：始终选择最可能的 token。
    """
    current_ids = input_ids.clone()

    for _ in range(max_new_tokens):
        outputs = model(current_ids)
        next_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True)
        current_ids = torch.cat([current_ids, next_token], dim=1)

    return current_ids

# 示例：贪心搜索可能重复 "The the the the..." 如果陷入循环

束搜索

def beam_search(
    model,
    input_ids: torch.Tensor,
    max_new_tokens: int,
    num_beams: int = 4,
    length_penalty: float = 1.0
) -> torch.Tensor:
    """
    束搜索：在每一步保留前 k 个假设。

    Args:
        model: 语言模型
        input_ids: 输入 token ID
        max_new_tokens: 最大生成 token 数
        num_beams: 跟踪的束数量
        length_penalty: 惩罚短序列（1.0 = 无惩罚）
    """
    batch_size = input_ids.shape[0]
    input_ids = input_ids.repeat_interleave(num_beams, dim=0)

    # 初始化束
    beam_scores = torch.zeros(batch_size * num_beams, device=input_ids.device)
    beam_scores[1::num_beams] = float('-inf')  # 只有第一个束有效

    for step in range(max_new_tokens):
        outputs = model(input_ids)
        next_token_logits = outputs.logits[:, -1, :]
        next_token_scores = F.log_softmax(next_token_logits, dim=-1)

        # 添加束分数
        vocab_size = next_token_scores.shape[-1]
        next_scores = beam_scores.unsqueeze(-1) + next_token_scores

        # 重塑以进行 top-k 选择
        next_scores = next_scores.view(batch_size, num_beams * vocab_size)
        next_scores, next_tokens = torch.topk(next_scores, k=num_beams, dim=-1)

        # 将扁平索引转换为 (束, token) 对
        beam_indices = next_tokens // vocab_size
        token_indices = next_tokens % vocab_size

        # 更新束
        input_ids = input_ids.view(batch_size, num_beams, -1)
        input_ids = input_ids[torch.arange(batch_size).unsqueeze(-1), beam_indices]
        input_ids = input_ids.reshape(batch_size * num_beams, -1)

        next_token_ids = token_indices.view(-1, 1)
        input_ids = torch.cat([input_ids, next_token_ids], dim=-1)

        beam_scores = next_scores.view(-1)

    # 选择最佳束
    input_ids = input_ids.view(batch_size, num_beams, -1)
    best_beam_indices = beam_scores.view(batch_size, num_beams).argmax(dim=-1)
    output_ids = input_ids[torch.arange(batch_size), best_beam_indices]

    return output_ids

# 使用示例
output = beam_search(model, input_ids, max_new_tokens=50, num_beams=4)

对比搜索

对比搜索将概率建模与退化惩罚相结合：

def contrastive_search(
    model,
    input_ids: torch.Tensor,
    max_new_tokens: int,
    top_k: int = 5,
    alpha: float = 0.6,
    momentum: float = 0.5
) -> torch.Tensor:
    """
    对比搜索：平衡概率和与上下文的相似度。

    Args:
        model: 语言模型
        input_ids: 输入 token
        max_new_tokens: 最大生成 token 数
        top_k: 候选 token 数量
        alpha: 重复惩罚（0 = 纯采样，1 = 纯退化惩罚）
        momentum: 更新累积概率的权重

    参考:
        Su, J., et al. (2022). "A Contrastive Framework for Neural Text Generation"
        https://arxiv.org/abs/2202.01855
    """
    current_ids = input_ids.clone()
    cumulative_probs = None

    for _ in range(max_new_tokens):
        # 前向传播
        with torch.no_grad():
            outputs = model(current_ids)
            logits = outputs.logits[:, -1, :]
            hidden_states = outputs.hidden_states[-1][:, -1, :]  # 最后一层隐藏状态

        # 获取 top-k 候选
        top_k_logits, top_k_indices = torch.topk(logits, top_k)
        top_k_probs = F.softmax(top_k_logits, dim=-1)

        # 计算模型置信度（概率质量）
        model_confidence = top_k_probs.max(dim=-1)[0]

        # 计算退化惩罚（与之前 token 的相似度）
        # 使用当前隐藏状态和之前状态的余弦相似度
        prev_hidden = outputs.hidden_states[-1][:, :-1, :]  # 所有之前的隐藏状态

        # 与最近的 token 计算相似度
        if prev_hidden.size(1) > 0:
            # 与最后几个 token 的相似度（局部上下文）
            recent_hidden = prev_hidden[:, -min(5, prev_hidden.size(1)):, :]
            similarities = F.cosine_similarity(
                hidden_states.unsqueeze(1),
                recent_hidden,
                dim=-1
            ).max(dim=1)[0]
            degeneration_penalty = similarities.max()
        else:
            degeneration_penalty = torch.tensor(0.0)

        # 更新累积概率
        if cumulative_probs is None:
            cumulative_probs = model_confidence
        else:
            cumulative_probs = momentum * cumulative_probs + (1 - momentum) * model_confidence

        # 选择最大化以下目标的 token：(1-alpha) * 概率 - alpha * 退化
        scores = (1 - alpha) * top_k_probs - alpha * degeneration_penalty

        # 选择最佳 token
        best_idx = scores.argmax(dim=-1)
        next_token = top_k_indices[range(len(best_idx)), best_idx].unsqueeze(-1)

        current_ids = torch.cat([current_ids, next_token], dim=1)

    return current_ids

MCTS（蒙特卡洛树搜索）解码

适用于需要更深层前瞻的任务：

class MCTSNode:
    """蒙特卡洛树搜索解码的节点。"""
    def __init__(self, token_id: int, parent=None):
        self.token_id = token_id
        self.parent = parent
        self.children = []
        self.visits = 0
        self.total_value = 0.0
        self.prior_prob = 0.0

    def ucb_score(self, c_puct: float = 1.0) -> float:
        """选择的上置信界。"""
        if self.visits == 0:
            return float('inf')

        exploitation = self.total_value / self.visits
        exploration = c_puct * math.sqrt(math.log(self.parent.visits) / self.visits) if self.parent else 0

        return exploitation + exploration

def mcts_decode(
    model,
    input_ids: torch.Tensor,
    max_new_tokens: int,
    num_simulations: int = 50,
    c_puct: float = 1.0,
    temperature: float = 1.0
) -> torch.Tensor:
    """
    用于前瞻解码的蒙特卡洛树搜索。

    代价更高但可以为复杂任务找到更好的序列。
    """
    current_ids = input_ids.clone()

    for step in range(max_new_tokens):
        # 构建搜索树
        root = MCTSNode(token_id=None)

        # 运行模拟
        for _ in range(num_simulations):
            # 选择：使用 UCB 遍历树
            node = root
            path = []

            while node.children:
                # 选择 UCB 分数最高的子节点
                node = max(node.children, key=lambda n: n.ucb_score(c_puct))
                path.append(node)

            # 扩展：如果不是终止节点则添加新子节点
            if len(path) < 5:  # 限制深度
                # 获取模型预测
                with torch.no_grad():
                    outputs = model(current_ids)
                    logits = outputs.logits[:, -1, :] / temperature
                    probs = F.softmax(logits, dim=-1)

                # 采样候选 token
                token = torch.multinomial(probs, num_samples=1).item()
                child = MCTSNode(token, parent=node)
                node.children.append(child)
                node.prior_prob = probs[0, token].item()
                path.append(child)

            # 评估：使用模型估计价值
            with torch.no_grad():
                # 构建完整序列进行评估
                eval_ids = current_ids.clone()
                for node in path[1:]:
                    eval_ids = torch.cat([eval_ids, torch.tensor([[node.token_id]])], dim=1)

                outputs = model(eval_ids)
                # 价值 = 负损失（越高越好）
                value = -outputs.loss.item() if hasattr(outputs, 'loss') else 0.0

            # 反向传播：更新价值
            for node in path:
                node.visits += 1
                node.total_value += value

        # 模拟后选择最佳子节点
        if root.children:
            best_child = max(root.children, key=lambda n: n.visits)
            next_token = torch.tensor([[best_child.token_id]])
            current_ids = torch.cat([current_ids, next_token], dim=1)
        else:
            break

    return current_ids

带温度的采样

def sample_with_temperature(
    model,
    input_ids: torch.Tensor,
    max_new_tokens: int,
    temperature: float = 1.0
) -> torch.Tensor:
    """
    带温度缩放的概率分布采样。
    """
    current_ids = input_ids.clone()

    for _ in range(max_new_tokens):
        outputs = model(current_ids)
        logits = outputs.logits[:, -1, :] / temperature

        # 从 softmax 分布采样
        probs = F.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)

        current_ids = torch.cat([current_ids, next_token], dim=1)

    return current_ids

# 温度效果
# T=0.1: 非常确定性，几乎等于贪心
# T=0.5: 集中，主要是高概率 token
# T=1.0: 标准采样
# T=1.5: 更有创意，包含低概率 token
# T=2.0+: 非常随机，通常不连贯

采样参数

温度

控制采样中的随机性：

def apply_temperature(logits: torch.Tensor, temperature: float) -> torch.Tensor:
    """
    对 logits 应用温度缩放。

    较低温度 -> 更尖锐的分布
    较高温度 -> 更平坦的分布
    """
    return logits / temperature

# 温度对概率分布的影响
logits = torch.tensor([2.0, 1.0, 0.0, -1.0, -2.0])

print("温度效果:")
for temp in [0.1, 0.5, 1.0, 2.0]:
    scaled = apply_temperature(logits, temp)
    probs = F.softmax(scaled, dim=0)
    print(f"  T={temp:.1f}: {probs.tolist()}")

输出：

温度效果:
  T=0.1: [0.97, 0.03, 0.00, 0.00, 0.00]  # 非常尖锐
  T=0.5: [0.67, 0.24, 0.07, 0.01, 0.00]  # 集中
  T=1.0: [0.50, 0.27, 0.12, 0.07, 0.04]  # 平衡
  T=2.0: [0.39, 0.32, 0.16, 0.09, 0.05]  # 平坦

Top-K vs Top-P（核采样）

def apply_top_k(logits: torch.Tensor, top_k: int) -> torch.Tensor:
    """
    过滤保留前 k 个 token。
    """
    top_k_logits, top_k_indices = torch.topk(logits, top_k)
    filtered = torch.full_like(logits, float('-inf'))
    filtered.scatter_(0, top_k_indices, top_k_logits)
    return filtered

def apply_top_p(logits: torch.Tensor, top_p: float) -> torch.Tensor:
    """
    核采样：保留累积概率质量 >= top_p 的最小顶部集合。
    """
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

    # 找到要移除的索引
    sorted_indices_to_remove = cumulative_probs > top_p
    sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()
    sorted_indices_to_remove[0] = False

    # 散射回原始顺序
    indices_to_remove = sorted_indices_to_remove.scatter(0, sorted_indices, sorted_indices_to_remove)
    logits[indices_to_remove] = float('-inf')

    return logits

# 示例：Top-k vs Top-p
logits = torch.randn(50000)  # 大词汇表

# Top-k：始终保留恰好 k 个 token
top_k_filtered = apply_top_k(logits, top_k=50)

# Top-p：保留可变数量的 token
top_p_filtered = apply_top_p(logits, top_p=0.9)
num_kept = (top_p_filtered != float('-inf')).sum()
print(f"Top-p=0.9 保留了 {num_kept} 个 token")

频率和存在惩罚

def apply_frequency_penalty(
    logits: torch.Tensor,
    token_ids: torch.Tensor,
    frequency_penalty: float = 0.0,
    presence_penalty: float = 0.0
) -> torch.Tensor:
    """
    应用频率和存在惩罚。

    Args:
        logits: 模型 logits (vocab_size,)
        token_ids: 之前生成的 token ID
        frequency_penalty: 基于频率的惩罚（越高 = 重复越少）
        presence_penalty: 基于存在的惩罚（二值）
    """
    # 统计 token 频率
    unique_tokens, counts = torch.unique(token_ids, return_counts=True)

    # 应用惩罚
    for token, count in zip(unique_tokens, counts):
        # 频率惩罚：与计数成正比
        logits[token] -= frequency_penalty * count

        # 存在惩罚：二值（存在或不存在）
        logits[token] -= presence_penalty

    return logits

# 示例
generated_tokens = torch.tensor([10, 20, 10, 30, 10, 20])  # 10: 3次, 20: 2次, 30: 1次
logits = torch.randn(50000)

# 带惩罚时，token 10 和 20 被更重地惩罚
logits_penalized = apply_frequency_penalty(
    logits.clone(),
    generated_tokens,
    frequency_penalty=0.5,
    presence_penalty=0.1
)

参数参考表

参数	范围	效果	使用场景
temperature	0.0 - 2.0	随机性	0.2: 编码、事实 / 0.8: 聊天 / 1.2: 创意
top_k	1 - 100	多样性	1: 贪心 / 40-50: 平衡 / 100: 非常多样
top_p	0.1 - 1.0	质量过滤	0.5: 集中 / 0.9: 标准 / 1.0: 无过滤
frequency_penalty	0.0 - 2.0	减少重复	0.0: 无 / 0.5: 适度 / 1.0: 强
presence_penalty	0.0 - 2.0	鼓励多样性	0.0: 无 / 0.5: 适度 / 1.0: 强

性能优化

KV 缓存

Key-Value 缓存存储之前 token 的注意力键和值，避免重复计算。

复杂度对比：

无缓存：每步 O(N²)
有缓存：每步 O(N)

class KVCache:
    """
    用于高效自回归生成的 Key-Value 缓存。
    """
    def __init__(self, batch_size: int, num_layers: int, num_heads: int, head_dim: int, max_len: int = 2048):
        self.batch_size = batch_size
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.max_len = max_len

        # 预分配缓存
        # 形状：(num_layers, batch_size, num_heads, max_len, head_dim)
        self.keys = torch.zeros(num_layers, batch_size, num_heads, max_len, head_dim)
        self.values = torch.zeros(num_layers, batch_size, num_heads, max_len, head_dim)
        self.current_len = 0

    def update(self, layer_idx: int, keys: torch.Tensor, values: torch.Tensor):
        """
        更新特定层的缓存。

        Args:
            layer_idx: 层索引
            keys: 新的键 (batch, num_heads, seq_len, head_dim)
            values: 新的值 (batch, num_heads, seq_len, head_dim)
        """
        seq_len = keys.size(2)
        if self.current_len + seq_len > self.max_len:
            raise ValueError(f"缓存溢出: {self.current_len + seq_len} > {self.max_len}")

        # 存储键和值
        self.keys[layer_idx, :, :, self.current_len:self.current_len+seq_len, :] = keys
        self.values[layer_idx, :, :, self.current_len:self.current_len+seq_len, :] = values

    def get(self, layer_idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        """
        获取某层的缓存键和值。

        Returns:
            keys: (batch, num_heads, current_len, head_dim)
            values: (batch, num_heads, current_len, head_dim)
        """
        return (
            self.keys[layer_idx, :, :, :self.current_len, :],
            self.values[layer_idx, :, :, :self.current_len, :]
        )

    def increment(self, n: int):
        """增加当前长度。"""
        self.current_len += n

# 带 KV 缓存的修改版注意力
def attention_with_cache(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    cache: KVCache,
    layer_idx: int,
    use_cache: bool = True
) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
    """
    带 KV 缓存的注意力计算。

    Args:
        q: Query (batch, num_heads, seq_len, head_dim)
        k: Key (batch, num_heads, seq_len, head_dim)
        v: Value (batch, num_heads, seq_len, head_dim)
        cache: KV 缓存
        layer_idx: 当前层索引
        use_cache: 是否使用/更新缓存

    Returns:
        output: 注意力输出 (batch, num_heads, seq_len, head_dim)
        (k, v): 用于缓存的键和值
    """
    if use_cache:
        # 用新的键和值更新缓存
        cache.update(layer_idx, k, v)

        # 获取所有键和值（缓存 + 新的）
        k_all, v_all = cache.get(layer_idx)
    else:
        k_all, v_all = k, v

    # 计算注意力分数
    scores = torch.matmul(q, k_all.transpose(-2, -1)) / math.sqrt(q.size(-1))
    attn_weights = F.softmax(scores, dim=-1)
    output = torch.matmul(attn_weights, v_all)

    return output, (k, v)

投机解码

使用较小的草稿模型预测 token，然后用较大模型验证。

def speculative_decode(
    draft_model,
    target_model,
    input_ids: torch.Tensor,
    max_new_tokens: int,
    spec_len: int = 4
) -> torch.Tensor:
    """
    使用草稿模型的投机解码。

    Args:
        draft_model: 更小、更快的模型
        target_model: 更大、更好的模型
        input_ids: 输入 token
        max_new_tokens: 最大生成 token 数
        spec_len: 投机前瞻的 token 数
    """
    current_ids = input_ids.clone()

    while current_ids.size(1) < input_ids.size(1) + max_new_tokens:
        # 步骤 1：草稿模型预测 spec_len 个 token
        draft_tokens = []
        draft_probs = []

        draft_input = current_ids
        for _ in range(spec_len):
            with torch.no_grad():
                draft_outputs = draft_model(draft_input)
                draft_logits = draft_outputs.logits[:, -1, :]
                draft_prob = F.softmax(draft_logits, dim=-1)

            # 采样 token
            draft_token = torch.multinomial(draft_prob, num_samples=1)
            draft_tokens.append(draft_token)
            draft_probs.append(draft_prob)

            draft_input = torch.cat([draft_input, draft_token], dim=1)

        draft_tokens = torch.cat(draft_tokens, dim=1)

        # 步骤 2：目标模型验证草稿 token
        with torch.no_grad():
            target_outputs = target_model(current_ids)
            target_logits = target_outputs.logits[:, -1:, :]  # 单个 token
            target_prob = F.softmax(target_logits, dim=-1)

        # 步骤 3：接受或拒绝每个草稿 token
        accepted_tokens = []

        for i, (token, draft_p) in enumerate(zip(draft_tokens.T, draft_probs)):
            # 获取该 token 的目标概率
            target_p = target_prob[0, :, token]

            # 以概率 target_p / draft_p 接受
            accept_prob = target_p / draft_p[0, token]

            if torch.rand(1) < accept_prob:
                accepted_tokens.append(token)
                current_ids = torch.cat([current_ids, token.unsqueeze(0)], dim=1)

                # 更新目标模型用于下一次验证
                with torch.no_grad():
                    target_outputs = target_model(current_ids)
                    target_logits = target_outputs.logits[:, -1:, :]
                    target_prob = F.softmax(target_logits, dim=-1)
            else:
                # 拒绝：从目标分布重新采样
                resampled = torch.multinomial(target_prob[0], num_samples=1)
                accepted_tokens.append(resampled)
                current_ids = torch.cat([current_ids, resampled.unsqueeze(0)], dim=1)
                break  # 停止投机，重新同步

    return current_ids

# 加速：10 倍模型大小比例时为 2-3 倍

2025：Medusa - 多头投机解码

Medusa（2024）通过向基础模型添加多个解码头来消除对单独草稿模型的需求。

工作原理：

训练：向基础模型添加轻量级预测头
推理：所有头并行预测 token
验证：基础模型同时验证所有预测
加速：2.2-3.6 倍，质量损失极小

相比草稿模型方法的优势：

无需单独的草稿模型
更好的质量（头在同一模型上训练）
更容易部署（单个模型）
随模型大小扩展

性能对比：

方法	加速	草稿模型	质量	设置
传统投机	2-3x	需要	~99%	复杂（2 个模型）
Medusa-1	2.2x	不需要	~99%	简单（1 个模型 + 头）
Medusa-2	3.6x	不需要	~98.5%	简单（1 个模型 + 头）

2025：QuantSpec - 量化 KV 投机解码

QuantSpec（2024）将投机解码与量化 KV 缓存相结合，实现自我投机。

关键创新：

分层 KV 量化：不同 token 使用不同精度级别
自我投机：无需单独的草稿模型
4 位 KV 缓存：减少 75% 内存
2.5 倍加速：质量下降极小

2025：EAGLE - Ngram 投机解码

EAGLE（2024）使用基于 ngram 的草稿，无需训练单独的草稿模型。

关键方法：

Ngram 模式：利用频繁的 token 序列
无需训练：适用于任何基础模型
2-3 倍加速：适用于重复性文本

量化

降低精度以加速计算和减少内存。

def quantize_int8(weight: torch.Tensor) -> tuple[torch.Tensor, float, float]:
    """
    将权重量化为 8 位整数。

    Args:
        weight: FP32 权重张量

    Returns:
        int8_weight: 量化后的权重
        scale: 缩放因子
        zero_point: 零点
    """
    # 计算缩放和零点
    qmin, qmax = -128, 127
    fp_min = weight.min().item()
    fp_max = weight.max().item()

    scale = (fp_max - fp_min) / (qmax - qmin)
    zero_point = qmin - fp_min / scale

    # 量化
    int8_weight = torch.round(weight / scale + zero_point).clamp(qmin, qmax).to(torch.int8)

    return int8_weight, scale, zero_point

def dequantize_int8(int8_weight: torch.Tensor, scale: float, zero_point: float) -> torch.Tensor:
    """将 INT8 权重反量化回 FP32。"""
    return scale * (int8_weight.float() - zero_point)

# 量化格式
quantization_formats = {
    "FP32": {"bits": 32, "range": "全精度", "speed": "1x"},
    "FP16": {"bits": 16, "range": "半精度", "speed": "2x"},
    "BF16": {"bits": 16, "range": "Brain float", "speed": "2x"},
    "INT8": {"bits": 8, "range": "整数量化", "speed": "4x"},
    "INT4": {"bits": 4, "range": "激进量化", "speed": "6x"},
}

# 权衡：INT8/INT4 损失质量但获得速度
# 典型退化：INT8 (~1-2%), INT4 (~5-10%)

2025：KV 缓存量化

KV 缓存是长上下文推理的内存瓶颈。量化缓存可大幅减少内存。

最先进技术：

技术	策略	内存减少	质量影响
CommVQ	交换向量量化	8-16x	~1-2%
ShadowKV	带选择性量化的影子表示	4-8x	`<1%`
KV Quant	按通道量化	4x	~2%
QuantSpec	分层量化	4-16x	~2%

2025：基于树的投机

STree、SpecInfer 和 TreeAttention 使用树结构进行并行投机。

关键概念： 不再顺序投机，而是并行投机多个分支并一起验证。

性能对比：

方法	并行度	加速	使用场景
顺序投机	1 个分支	2-3x	通用文本
基于树	4-8 个分支	3-5x	多样化输出
STree	自适应	4-6x	混合 SSM-Transformer

交互式参数演示

def compare_decoding_strategies(
    model,
    tokenizer,
    prompt: str,
    strategies: list[dict]
) -> dict[str, str]:
    """
    在同一提示上比较不同的解码策略。

    Args:
        model: 语言模型
        tokenizer: 分词器
        prompt: 输入提示
        strategies: 策略配置列表

    Returns:
        策略名称到输出的映射字典
    """
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    results = {}

    for config in strategies:
        name = config.pop('name', 'unknown')

        output_ids = generate_autoregressive(
            model,
            input_ids,
            max_new_tokens=100,
            **config
        )

        output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
        results[name] = output_text

        # 重置配置用于下一次迭代
        config['name'] = name

    return results

# 示例比较
prompt = "Once upon a time"
strategies = [
    {"name": "贪心", "temperature": 0.0},
    {"name": "集中", "temperature": 0.3, "top_p": 0.9},
    {"name": "平衡", "temperature": 0.8, "top_p": 0.9},
    {"name": "创意", "temperature": 1.2, "top_k": 50},
    {"name": "狂放", "temperature": 1.5, "top_k": 100},
]

results = compare_decoding_strategies(model, tokenizer, prompt, strategies)

for name, output in results.items():
    print(f"\n{name.upper()}:")
    print(output)

关键要点

解码策略决定输出质量和多样性
温度控制随机性与确定性
Top-k 和 Top-p 平衡多样性和连贯性
KV 缓存对高效生成至关重要
量化以质量换取速度和内存
2025 推理优化：
- Medusa：多头投机（2.2-3.6 倍加速，无需草稿模型）
- QuantSpec：量化 KV + 自我投机（2.5 倍加速，75% 内存减少）
- KV 缓存量化：CommVQ/ShadowKV（4-16 倍内存减少）
- 基于树的投机：并行分支验证（3-5 倍加速）

2025 推理优化常见问题

Q：何时应该使用 Medusa 而非传统投机解码？

A: 使用 Medusa 当：

你想要单模型部署（无草稿模型管理）
有训练预算（需要训练 Medusa 头）
想要跨任务一致的 2.2-3.6 倍加速
优先考虑更简单的部署

使用传统投机解码 当：

你已有可用的草稿模型
无训练预算（使用预训练模型）
需要与现有基础设施的最大兼容性
能管理双模型部署

2025 年结论：Medusa 正在成为新部署的默认选择，因为其简单性和竞争力性能。

Q：什么导致 KV 缓存成为瓶颈？

A: KV 缓存内存随序列长度线性增长：

内存 = 2 × num_layers × num_heads × seq_len × head_dim × bytes_per_value
对于 Llama 3 70B（80 层, 64 头, 128 维, FP16）：
- 8k token：每批次约 2 GB
- 128k token：每批次约 32 GB
- 1M token：每批次约 256 GB（单 GPU 不可能）

2025 年解决方案：

GQA（分组查询注意力）：4-8 倍减少（Llama 3 已使用）
KV 量化：额外 4-16 倍减少（CommVQ、ShadowKV）
Ring Attention：跨多 GPU 分布

经验法则：对于 32k+ 上下文，始终使用 KV 量化。

Q：如何选择正确的量化策略？

A: 决策框架：

使用 FP16/BF16 当：

质量至关重要（基准测试、研究）
内存不受限（A100/H100）
短上下文（< 8k token）

使用 INT8 当：

生产部署，可接受 1-2% 的质量损失
消费级 GPU（RTX 4090 等）
标准上下文（8k-32k token）

使用 INT4 当：

需要最大吞吐量
边缘部署或移动端
可接受 5-10% 的质量下降
长上下文（32k+ token）

使用分层量化（QuantSpec）当：

超长上下文（100k+ token）
需要平衡内存和质量
最近的 token 需要高精度，旧的 token 可以压缩

2025 年最佳实践：从 INT8 开始，仅在内存受限时使用 INT4。100k+ 上下文考虑 QuantSpec。

Q：采样参数对推理速度有什么影响？

A: 采样参数通过以下方式影响速度：

温度：

低（0.1-0.3）：更快（类似贪心），较少多样性
高（1.0-1.5）：更慢（更多计算），更多多样性
影响：约 10-20% 的速度差异

Top-k：

小（1-10）：计算 softmax 更快
大（50-100）：需要考虑更多 token
影响：约 5-10% 的速度差异

投机兼容性：

贪心/temp=0：完美适合投机（80-90% 接受率）
高温：较低接受率（50-60%）

2025 年优化提示： 使用投机实现最大速度时，使用 temperature=0.3-0.5, top_k=40。这平衡了速度和质量。

Spring AI 流式实现

Spring AI 使用 Server-Sent Events (SSE) 和响应式编程提供实时 LLM 响应的流式支持。

SSE 流式传输

// Spring AI SSE 流式传输
@RestController
@RequestMapping("/api/chat")
public class ChatController {
    private final ChatClient chatClient;

    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> streamChat(@RequestParam String message) {
        return chatClient.prompt()
            .user(message)
            .stream()
            .content();  // 返回 Flux<String> 用于流式传输
    }
}

响应式流式服务

// 流式 LLM 响应服务
@Service
public class StreamingChatService {
    private final ChatClient chatClient;

    // 流式响应到前端
    public Flux<String> streamResponse(String userMessage) {
        return chatClient.prompt()
            .user(userMessage)
            .stream()
            .content()
            .doOnNext(token -> log.debug("收到 token: {}", token))
            .doOnComplete(() -> log.info("流完成"))
            .doOnError(error -> log.error("流错误", error));
    }

    // 带进度跟踪的流式传输
    public Flux<StreamingResponse> streamWithProgress(String userMessage) {
        AtomicInteger tokenCount = new AtomicInteger(0);

        return chatClient.prompt()
            .user(userMessage)
            .stream()
            .content()
            .map(token -> {
                int count = tokenCount.incrementAndGet();
                return new StreamingResponse(token, count, false);
            })
            .concatWithValues(StreamingResponse.done());
    }

    // 流式响应记录
    public record StreamingResponse(
        String token,
        int tokenCount,
        boolean isDone
    ) {
        public static StreamingResponse done() {
            return new StreamingResponse("", 0, true);
        }
    }
}

参数对比示例

// 采样参数对输出的影响
@Service
public class ParameterComparisonService {
    private final ChatClient chatClient;

    // 代码生成：低温度以保持一致性
    public String generateCode(String description) {
        return chatClient.prompt()
            .user("Write code to: " + description)
            .options(OpenAiChatOptions.builder()
                .temperature(0.2)  // 低 = 更确定性
                .topP(0.95)
                .maxTokens(1500)
                .build())
            .call()
            .content();
    }

    // 创意写作：更高温度
    public String generateStory(String prompt) {
        return chatClient.prompt()
            .user("Write a story about: " + prompt)
            .options(OpenAiChatOptions.builder()
                .temperature(0.9)  // 高 = 更有创意
                .topP(0.9)
                .maxTokens(2000)
                .presencePenalty(0.5)  // 鼓励新想法
                .build())
            .call()
            .content();
    }

    // 摘要：低温度，集中
    public String summarize(String text) {
        return chatClient.prompt()
            .user("Summarize: " + text)
            .options(OpenAiChatOptions.builder()
                .temperature(0.3)  // 集中
                .topP(1.0)  // 考虑所有 token
                .maxTokens(500)
                .build())
            .call()
            .content();
    }
}

带背压的流式传输

// 为慢速客户端处理背压
@Service
public class BackpressureStreamingService {
    private final ChatClient chatClient;

    public Flux<String> streamWithBackpressure(String message, int bufferSize) {
        return chatClient.prompt()
            .user(message)
            .stream()
            .content()
            .onBackpressureBuffer(BufferOverflowStrategy.DROP_OLDEST)
            .doOnNext(token -> log.debug("发送 token: {}", token));
    }

    // 限速流式传输
    public Flux<String> streamRateLimited(String message, int tokensPerSecond) {
        Duration delay = Duration.ofMillis(1000 / tokensPerSecond);

        return chatClient.prompt()
            .user(message)
            .stream()
            .content()
            .delayElements(delay);
    }
}

WebSocket 流式传输

// 用于双向通信的 WebSocket 流式传输
@Controller
public class WebSocketChatController {
    private final StreamingChatService chatService;

    @MessageMapping("/chat")
    public Flux<String> handleChatMessage(String message) {
        return chatService.streamResponse(message)
            .doOnNext(token -> log.debug("发送 token: {}", token));
    }
}

参考文献

Radford, A., Wu, J., Child, R., et al. (2019). "Language Models are Unsupervised Multitask Learners." OpenAI technical report.

链接： https://d4mucfpksywv.cloudfront.net/better-language-models/language-models.pdf

GPT-2 论文，引入核采样和解码策略。

Holtzman, A., Buys, J., Du, L., et al. (2020). "The Curious Case of Neural Text Degeneration." ACL 2020.

链接： arXiv:1904.09751

分析神经文本生成中的重复和退化问题。

Levi, D., Khashabi, D., & Roth, D. (2023). "Speculative Decoding: Accelerating LLM Inference." arXiv preprint.

链接： arXiv:2211.17192

引入投机解码以加速推理。

Dettmers, T., Pagnoni, A., Holtzman, A., & Zettlemoyer, L. (2024). "QLoRA: Efficient Finetuning of Quantized LLMs." ICLR 2024.

链接： arXiv:2305.14314

量化感知训练用于大型语言模型。

Inference - 控制输出质量

自回归生成

生成循环

基本实现

解码策略

策略对比

贪心搜索

束搜索

对比搜索

MCTS（蒙特卡洛树搜索）解码

带温度的采样

采样参数

温度

Top-K vs Top-P（核采样）

频率和存在惩罚

参数参考表

性能优化

KV 缓存

投机解码

2025：Medusa - 多头投机解码

2025：QuantSpec - 量化 KV 投机解码

2025：EAGLE - Ngram 投机解码

量化

2025：KV 缓存量化

2025：基于树的投机

交互式参数演示

关键要点

2025 推理优化常见问题

Spring AI 流式实现

SSE 流式传输

响应式流式服务

参数对比示例

带背压的流式传输

WebSocket 流式传输

参考文献

AiDIY 助手

AiDIY 助手

自回归生成​

生成循环​

基本实现​

解码策略​

策略对比​

贪心搜索​

束搜索​

对比搜索​

MCTS（蒙特卡洛树搜索）解码​

带温度的采样​

采样参数​

温度​

Top-K vs Top-P（核采样）​

频率和存在惩罚​

参数参考表​

性能优化​

KV 缓存​

投机解码​

2025：Medusa - 多头投机解码​

2025：QuantSpec - 量化 KV 投机解码​

2025：EAGLE - Ngram 投机解码​

量化​

2025：KV 缓存量化​

2025：基于树的投机​

交互式参数演示​

关键要点​

2025 推理优化常见问题​

Spring AI 流式实现​

SSE 流式传输​

响应式流式服务​

参数对比示例​

带背压的流式传输​

WebSocket 流式传输​

参考文献​

自回归生成

生成循环

基本实现

解码策略

策略对比

贪心搜索

束搜索

对比搜索

MCTS（蒙特卡洛树搜索）解码

带温度的采样

采样参数

温度

Top-K vs Top-P（核采样）

频率和存在惩罚

参数参考表

性能优化

KV 缓存

投机解码

2025：Medusa - 多头投机解码

2025：QuantSpec - 量化 KV 投机解码

2025：EAGLE - Ngram 投机解码

量化

2025：KV 缓存量化

2025：基于树的投机

交互式参数演示

关键要点

2025 推理优化常见问题

Spring AI 流式实现

SSE 流式传输

响应式流式服务

参数对比示例

带背压的流式传输

WebSocket 流式传输

参考文献