Skip to content

第十章 推理

image

Inference工作示意图

  Inference是大型语言模型真正发挥价值的核心环节,也是训练阶段所学习知识转化为实际能力的唯一途径。无论是智能问答、文本生成、代码补全,还是搜索、翻译等应用场景,模型都需要通过推理过程完成对输入信息的理解、计算与输出。

  随着 Agent 的快速发展,以 Claude Code、Codex、WorkBuddy 等 AI 开发工具的新一代应用,其核心能力也几乎都建立在 LLM 推理之上。虽然这些工具可以接入不同的大模型 API,但实际使用体验却存在明显差异:性能更强的模型通常具有更高的任务完成率、更稳定的推理过程以及更好的代码生成质量。然而仅有高质量的输出还不够,在 Agent 场景下,一个复杂任务往往需要经历多轮规划、调用工具、执行代码以及反思修正等推理过程,推理延迟会被不断累积。因此,模型的响应速度、吞吐能力以及推理成本,同样成为影响整体开发效率和用户体验的重要因素。可以说,现代 LLM 的竞争已经不仅仅体现在"回答得是否正确",更体现在"回答得是否足够快、足够经济"

  那么,影响 LLM 推理的瓶颈会来自哪里?业界又采用了哪些技术来提升推理速度、降低推理成本?在分析这些优化方法之前,我们首先需要建立评价推理性能的基本指标,理解什么样的模型才算"推理得快",以及影响推理效率的关键因素有哪些。接下来,我们将围绕这些问题,对 LLM 推理优化技术进行系统分析。

学习目标

读完本章后,应能够围绕下面几个问题给出完整回答:

 1.如何评价一个 LLM 的推理性能,其中TTFT、首 token 之后的 token 间延迟和 Throughput 分别衡量什么,它们之间为什么不能简单地相互替代?

 2.训练阶段与推理阶段在输入来源、并行方式、显存占用和性能瓶颈上有什么差异?为什么训练可以并行处理整段序列,而自回归推理通常必须逐 token 生成?

 3.为什么推理要区分 Prefill 和 Decode?KV Cache 解决了什么重复计算问题,又为什么会成为长上下文 Decode 的显存和带宽负担?

 4.如何用 FLOPs、HBM 数据传输量和算术强度判断一个推理阶段更可能是 compute-bound 还是 memory-bound?为什么增大 batch size、采用 GQA/MQA、量化能够改善推理效率?

 5.当前推理优化技术分别改变了哪一类瓶颈?SSM、扩散模型、推测解码以及 MTP 在并行性、上下文一致性、KV Cache、延迟和吞吐量之间作出了什么取舍?

带着这几个疑问,我们继续接下来的分析!

  本章的论述路线也与这些问题一一对应:先比较训练和推理的基本机制,再介绍 Transformer 自回归推理、Prefill/Decode 与 KV Cache;随后用算术强度定量分析 Attention 和 MLP 的瓶颈,并讨论延迟、吞吐量、批处理和提示词压缩;最后转向 SSM、扩散模型、推测解码和 MTP 等结构或系统层面的优化方法。这样可以把“模型如何生成 token”与“为什么生成得快或慢”以及“如何进一步加速”串成一条完整主线。

10.1 推理与训练

  训练与推理虽各自受不同瓶颈约束,却并非独立阶段;训练塑造能力,推理呈现能力,而推理的需求又反过来影响训练设计。二者共同构成了构建大规模语言模型的完整系统循环。

10.1.1 推理与训练的差异

  在LLM中,训练与推理都涉及模型的前向计算,但它们的目标、计算模式和资源瓶颈存在本质差异。对于自回归推理并行化训练,假设需要生成第 i 个token,则模型在该位置的条件概率可以统一写为:

P(tokenitoken1,token2,,tokeni1)

训练与推理在单步预测的数学形式上相同 —— 都是对词表做概率分布估计并从中选择,但差别在于 tokeni 前文的来源(真实vs生成)以及是否把预测结果作为后续输入和是否能并行化

训练与推理的预测下一个词的依据对比

阶段下一步预测 tokeni 的输入来源
训练来自真实标签(ground-truth)及其之前的输入序列,模型可利用完整序列并行计算,同时通过因果掩码保证每个位置只访问前面token。
推理来自模型先前预测生成的token及其之前的文本,每步预测依赖前一步输出必须逐步生成(auto-regressive)。

1.训练阶段

  训练阶段的主要目标是 优化模型参数,训练需要执行完整的前向与反向传播,保留中间激活用于梯度计算,并在多 GPU 间同步梯度。因此训练过程高度依赖算力,瓶颈主要来自:

  • 大规模矩阵乘法的FLOPs
  • 多张显卡并行运行涉及的跨设备、统一设备内通信量
  • 激活、梯度和优化器状态带来的额外内存占用

训练阶段的显存占用主要由四部分组成:激活、梯度、优化器状态以及模型参数。

对于训练阶段的自回归 Transformer,我们输入完整目标序列并使用 因果掩码(causal mask):每个位置只能访问前面token(真实标签),无法看到未来token。由于训练输入是已知的完整序列,模型可以将整个序列视作一个大批次,通过一次或少量大矩阵乘法并行计算所有位置的前向表示,同时因果掩码保证自注意力的顺序依赖。中间激活会被保留用于反向传播与参数更新。正因如此,尽管自注意力涉及成对位置的注意力计算,训练仍能在时间维度高度并行,这也是Transformer在大规模训练中高效利用GPU的关键原因。

2.推理阶段

  推理阶段的目标使用固定参数生成输出序列。推理通常采用自回归生成方式,生成每个token时,模型需要参考之前预测生成的所有token的信息,也就是每一层的Key、Value(KV cache)。

  可以把LLM的推理过程想象成在写作文:每写一个字都要回头看前面已经写好的内容。随着上下文长度增加,这些“参考资料”(KV cache)也越来越多,占用的显存越来越大,而且每次访问都需要时间,就像翻阅厚厚笔记本一样慢。因此,推理的速度并不主要受算力限制,而是受显存容量和内存带宽的制约

总的来说就是:训练的瓶颈在算力,而推理的瓶颈在显存与带宽

为什么在自回归LLM的训练阶段,模型预测下一个token时通常以真实标签的前序tokens作为输入即teacher-forcing,而不是以模型先前预测的token作为输入?

  在训练阶段,使用teacher-forcing有两个主要原因:一方面,它提供稳定且明确的监督信号,便于模型快速收敛;另一方面,它允许对整段序列进行并行计算,从而大幅提高训练效率。如果在训练中将模型先前预测的token作为下一步输入,则可能导致隐藏状态中累积预测偏差。此外,由于采样(预测生成的token)操作是非微分的,会破坏标准的基于最大似然的逐位置损失计算和反向传播,从而增加训练难度。

举例:训练 vs 推理并行对比图

训练阶段(全序列并行化,使用ground-tru原有的前文)
----------------------------------------------------
输入序列:     x1            x2           x3            x4
模型前向:    ┌────┐      ┌────┐       ┌────┐         ┌────┐
             │f(x1)│     │f(x1,x2)│   │f(x1,x2,x3)│  │f(x1,x2,x3,x4)│
             └────┘      └────┘       └────┘         └────┘
输出序列:     y1           y2           y3             y4
说明:
1. 所有token的预测y_i可并行计算(Transformer一次前向即可得到全部y_i)。
2. 位置i的预测条件为真实前文(x1...x{i-1}),因causal mask保证无法看到未来token。
3. 预测y_i仅用于计算训练损失例如cross-entropy,不会作为下一步输入。

----------------------------------------------------
推理阶段(自回归生成,逐token依赖模型生成的前文)
初始输入:   x1  
Step 1:     y1 = f(x1)
Step 2:     y2 = f(x1, y1)
Step 3:     y3 = f(x1, y1, y2)
Step 4:     y4 = f(x1, y1, y2, y3)
说明:
1每一步都必须等待上一步完成,无法并行。
2. 位置 i 的前文是模型自己已生成的token(不是ground-truth)。
3. y_i用作下一步输入并不断扩展序列。

----------------------------------------------------

  在LLM的推理阶段,每一次生成一个token都依赖于之前已经生成的所有token,因此生成token的流程在时间维度上具有严格的顺序性,无法并行一次性输出多个token。但需要强调的是,这并不意味着推理内部的计算无法并行,每一步在计算新的KV、注意力权重以及前馈网络时,这些操作本质上仍是大型矩阵运算,完全可以在GPU上进行高度并行化处理。因此,推理阶段的特点可以概括为在时间维度上是顺序的(自回归约束),在计算维度上是并行的(矩阵运算加速),训练阶段可以并行处理整个序列而推理阶段只能顺序处理生成token,但每个token的计算内部依然高效并行化。

  到这里,我们已经回答了第二个基础问题的一半:训练和推理都执行前向计算,但训练利用已知的完整序列进行时间维度上的并行计算,并为反向传播保存大量中间状态;推理则使用模型自己生成的前文,受到严格的自回归依赖。接下来还需要进一步说明:这种顺序生成在 Transformer 中具体如何实现,以及 KV Cache 为什么能够缓解重复计算、却又带来新的内存压力。

10.1.2 训练与推理的联系

  尽管训练与推理在操作机制、性能瓶颈和优化方向上存在明显差异,但二者在整个模型生命周期中紧密关联:

  • 训练的根本目的在于优化推理行为。 无论采用监督学习、预训练、监督微调还是对齐方法,其结果最终体现在推理阶段的生成质量、稳定性和一致性上。推理是模型面向用户呈现所有能力的方式,因此训练过程的每一项设计——数据选择、损失函数、正则化方法等方法最终都影响推理表现。

  • 推理本身是训练过程不可或缺的一部分。 模型验证、能力评估、指令遵循测试等都依赖推理过程。在训练中,模型必须通过推理生成响应,再利用这些响应计算奖励或者损失并更新参数。因此,从系统角度来看推理并非仅出现在部署环节,而是贯穿训练的整个迭代周期。

  • 训练中对模型结构的设计决定推理的可优化性。 推理阶段受限于显存、KV缓存大小、延迟要求、带宽等多方面瓶颈,因此模型在训练阶段的结构设计必须为推理的可执行性和性能提供保障。

    • 结构设计影响推理的资源占用。 例如Transformer的结构、注意力复杂度、参数规模、MoE的专家布局、是否有KV压缩等,都会在推理时决定显存占用与延迟表现。设计不当的结构可能在训练时表现良好,但在推理时根本无法高效运行。

    • 结构设计影响推理的可优化空间。 是否能够使用FlashAttention、长上下文机制、推理加速策略等,都依赖模型结构在训练阶段是否为这些优化留下空间,一个推理友好的结构设计能在部署阶段获得更高吞吐、更低延迟和更强的扩展性。这些设计往往必须在训练前便已确定,而无法在推理阶段临时补救。因此,推动模型推理效率的提升,需要在训练阶段就纳入系统性的结构设计考虑。

FlashAttention依然是全局注意力机制,但将原本需要一次性计算和存储的大型注意力矩阵拆分为多个可在高速片上缓存中处理的小块,通过分块计算显著减少显存读写量,从而降低内存访问开销并显著提升推理与训练效率。

10.2 推理原理分析

  因此,训练阶段决定模型学到了什么,推理阶段决定这些能力能否以可接受的速度和成本交付给用户。下面进入推理原理分析:先明确性能指标,再沿着一次请求的 Prefill 和 Decode 流程,定位计算、显存和带宽分别在哪里成为瓶颈。

  训练是一次性的成本,而推理却会被不断重复执行,随着大模型真正落地到各种应用场景,推理的规模正在急速扩大。Sam Altman曾公开提到,OpenAI现在每天要生成超过1000亿个词;像Cursor这种成熟的AI开发工具,每天也会产出数十亿行实际被用户采纳的代码。这些数字说明了一个事实:推理已经成为大模型成本中的主力,而非训练。为了更清晰地理解推理效率,我们通常从首 token 等待时间、生成阶段的 token 间隔、单请求总耗时以及系统总吞吐量几个维度来衡量。

  • TTFT(Time To First Token):用户发出请求到看到第一个 token 的时间,主要受请求排队和 Prefill 计算影响。
  • ITL(Inter-Token Latency):首 token 之后相邻输出 token 的平均间隔,更直接反映 Decode 阶段的流式输出速度。
  • E2E Latency:从请求到达到最终 token 输出完毕的总时间,可近似理解为排队时间、TTFT 和后续 token 生成时间的总和。
  • Throughput:系统单位时间内处理或生成的 token 数量,通常用于衡量多请求服务的整体效率。

需要注意的是,高吞吐量并不等于低延迟。在大规模推理系统中,系统总体吞吐量可以很高即每秒产出大量token,但个别对话请求的响应时间仍可能很长——原因包括上下文特别长、为该请求分配的计算资源较少、或请求在调度队列中等待等。吞吐量衡量的是系统层面的总体效率,而延迟关注的是单个用户或单次对话请求的体验,因为用户直接感受到的是其带来的等待。

  接下来先从 Transformer 的一次自回归推理开始,说明请求如何经历 Prefill 和 Decode、KV Cache 在哪里产生并被复用;然后再把这些过程转换为 FLOPs、HBM 数据传输量和算术强度,解释不同阶段为什么会呈现不同的性能瓶颈。最后,我们再讨论如何通过调度、输入的提示词压缩和模型结构改造,针对这些瓶颈进行优化。

10.2.1 Transformer

48108f8a684becfe656b925ff929365f

图10.1 Transformer推理过程

自回归Transformer的LLM推理中,一个核心优化点在于对已经存在的上下文,不需要每生成一个token就重新做完整前向计算。为此,推理通常分为两个阶段:预填充逐步生成

2d8603a3413b3e848e256268422914ca

图10.2 自回归Transformer推理

  • 预填充(Prefill)阶段,模型一次性处理用户输入的所有prompt token,并在每一层、每个注意力头中计算出这些token的Key与Value,记为 K,V 。这些向量随后会被保存进KV Cache的结构中。保存KV的目的很简单——后续每生成一个新token时,我们只需要对这个新token做一次线性投影得到Query(Q),再让Q去与已经缓存好的所有K做点积,并用已有的V做加权求和即可。这样就不需要反复对旧token重新做线性映射与矩阵乘法,从而避免了大量重复计算。

  • 生成(Decode)阶段,模型用上一步得到的新token embedding生成它的Query,执行

Attention(Q,K,V)=softmax(QKdk)V

其中 (K,V) 全部来自KV Cache,得到的输出logits是模型对词汇表中每个token的未归一化分数,其经过softmax归一化后就是词汇表的概率分布,用于采样生成下一个token。然后把这个新token在各层的新 (K,V) 继续追加到缓存中,循环直到终止符或达到最大长度。

KV Cache会占用显存,其规模大致与token数 × 层数 × 注意力头数 × 隐藏层维度成正比。

需要注意——KV Cache适用于自回归的推理模式,因为在自回归设置中过去的token一旦计算完成就不会被未来token改变,故其对应的 (K,V) 可以复用。

 ①对于双向注意力如BERT、encoder-only等或者频繁修改输入上下文的场景如掩码语言模型、文本编辑器等,任一token的变动都会影响其他token的表示与Query–Key的相对关系,从而改变注意力权重即 QK 和应加权的 V 。因此在这些非自回归推理场景中,已缓存的 (K,V) 不能保持有效——一旦序列有改动就必须重新计算全部 (K,V) ,其不适合使用KV Cache

 ②在扩散模型中,注意力模块同样会计算Key、Value,但其含义与自回归模型中的 KV cache并不相同。扩散模型中的KV是针对所有空间位置在一次前向计算中并行生成的,并且通常在每一个扩散步骤(diffusion step)中重新计算,因此不会像自回归生成那样在时间维度上随生成过程逐步累积。当模型引入外部条件如文本条件时,由条件编码器产生的KV可以在多个扩散步骤之间复用,以减少重复计算,但这种缓存仅针对不随扩散步骤变化的条件表示,并不等同于自回归模型中随序列长度不断增长,并在整个生成过程中持续占用显存的KV Cache。

为什么当前大多数生成型模型如GPT系列、LLaMA系列、DeepSeek系列等采用自回归注意力架构,而不是像BERT那样的双向注意力?

  1. 高效逐步生成:自回归模型按顺序预测下一个token,可用KV Cache前面计算的Key、Value,避免重复计算,实现流式生成和在线交互的高效推理。
  2. 全局一致性 vs 计算开销:双向attention能访问前后上下文,提高理解和连贯性,但生成时需要多次全序列attention、迭代更新,计算和内存开销大,交互延迟高。
  3. 工程折中与扩展性:自回归架构训练易扩展,虽然牺牲了即时全局修改能力,但在大规模部署中是最实用的折中方案。

LLM自回归方式推理简易代码实现

python
import numpy as np
# 词表
vocab = ["I", "love", "deep", "learning", "<EOS>"]
token2id = {w: i for i, w in enumerate(vocab)}
id2token = {i: w for i, w in enumerate(vocab)}
vocab_size = len(vocab)

# embedding维度
d_model = 8
np.random.seed(42)
E = np.random.randn(vocab_size, d_model)

# 每个token对下一个token的概率分布(行:当前token,列:下一个token)
manual_logits = np.array([
    [0.0, 3.0, 0.0, 0.0, 0.0],   # I -> love
    [0.0, 0.0, 0.0, 3.0, 0.0],   # love -> learning
    [0.0, 3.0, 0.0, 0.0, 0.0],   # deep -> love
    [0.0, 0.0, 0.0, 0.0, 3.0],   # learning -> <EOS>
    [0.0, 0.0, 0.0, 0.0, 0.0],   # <EOS> -> nothing
])

# softmax
def softmax(x):
    e = np.exp(x - np.max(x))
    return e / e.sum()

# KV Cache + 可控生成
def generate_kv_readable(prefix, max_len=10, min_len=3):
    print(f"输入前缀: {prefix}")
    ids = [token2id[w] for w in prefix]

    K_cache = np.zeros((0, d_model))
    V_cache = np.zeros((0, d_model))

    for step in range(max_len):
        last_id = ids[-1]
        x = E[last_id]

        # Q, K, V
        q = x @ np.eye(d_model)  # 简化: Q=K=V=embedding
        k = x
        v = x

        # 更新KV Cache
        K_cache = np.vstack([K_cache, k])
        V_cache = np.vstack([V_cache, v])

        # 注意力
        att_scores = K_cache @ q
        att_scores /= np.sqrt(d_model)
        att_weights = softmax(att_scores)
        context = att_weights @ V_cache

        # 使用手动logits + softmax
        logits = manual_logits[last_id]
        probs = softmax(logits)

        # 用argmax确保可控生成
        next_id = int(np.argmax(probs))
        next_token = id2token[next_id]

        print(f"第 {step+1} 步预测 → {next_token} (p={probs[next_id]:.2f})")
        print(f"  Attention 权重: {np.round(att_weights, 2)}")

        # 最小生成长度控制 <EOS>
        if next_token == "<EOS>" and step+1 < min_len:
            # 强制选择概率第二高的token
            sorted_ids = np.argsort(probs)[::-1]
            for sid in sorted_ids:
                if sid != token2id["<EOS>"]:
                    next_id = sid
                    next_token = id2token[next_id]
                    break

        if next_token == "<EOS>":
            ids.append(next_id)
            break
        ids.append(next_id)
    print("\n最终生成序列:", " ".join(id2token[i] for i in ids))

# 测试 
generate_kv_readable(["I"])

输入 ["I", "love", "deep", "learning", "<EOS>"],generate_kv_readable(["I", "deep"])

输出 最终生成序列: I deep love learning <EOS>

输出生成序列有问题不是KV Cache或代码逻辑本身的问题,而是随机采样 + 未训练参数矩阵造成的。

  目前主流的大语言模型推理体系采用的是decoder-only、基于自回归的Transformer。这种结构在推理过程中可以使用KV Cache来复用历史token的Key、Value,从而避免重复计算。但随着上下文变长,KV Cache会带来显著的显存与带宽开销,因此近年来出现了许多为改进长上下文推理效率或降低KV负担等问题的Transformer变体比如GQA、MHA、MLA、稀疏注意力机制、DSA等。以及MINMAX研究的混合注意力机制对推理明显加速,其架构以一种结构化模式融合了线性注意力机制和Softmax注意力机制即在每7个采用线性注意力的Transnormer之后,插入一个使用Softmax注意力的Transformer模块,总共构成80层。

MINMAX的线性注意力机制+局部注意力机制

ee1e4333575a9664a1f8c985866d1660

图10.3 MINMAX混合注意力机制

  线性注意力机制通过将传统的点积Softmax注意力线性化,实现了在序列处理中的高效计算。在训练阶段,其复杂度为 O(nd2) ,并且在推理阶段可以通过递归更新累积项 KV 来高效生成序列,这相比传统Softmax注意力在推理阶段仍然需要 O(n2d) 的计算,显得尤为高效。

a8ae5949f18d0ab1adbe8c40d1c0f840

图10.4 推理时间对比

  在自回归生成场景中,我们希望模型在生成每个词时只能看到前面已经生成的词,不能“偷看”后面的词。如果直接使用 ϕ(Q)(KV) 的方法,相当于把整句话的所有词的关键信息一次性混在一起再去计算每个词的输出,这会导致模型在计算当前词时意外看到未来的词信息,违反了自回归规则。为了避免这种情况,需要在计算时一步一步累加:每生成一个词,就把当前词的key、value加到累加器里,下一步生成下一个词时只使用已经累加过的前缀信息。这样,每个词的输出都只依赖前面已经生成的词,而看不到未来。

逐步累加的方法不能像训练或全句并行那样一次性快速计算所有词,必须按顺序一点一点生成,因此推理速度比全序列并行慢。但相比局部注意力机制,线性注意力仍然在推理阶段具有明显的计算优势,因为它在每一步的计算量是线性增长的,而非平方增长。

  因此,尽管线性注意力早在九年前就被提出由于其设计比较复杂,目前主流开源大语言模型——包括LLaMA3、Qwen2.5、DeepSeekV3和Mistra等仍未将线性注意力作为自回归生成的默认方案。

fc8be482273494340f651bcc6a6541f7

图10.5 全局注意力机制 vs 线性注意力机制

线性注意力机制的生成任务影响

python
import numpy as np

# ---- 配置 ----
L = 5          # 序列长度 
d = 2          # 特征维度
np.random.seed(42) # 设置随机种子以保证结果可复现

# 随机生成 Q, K, V 矩阵,形状都是(L, d)
Q = np.random.rand(L, d)
K = np.random.rand(L, d)
V = np.random.rand(L, d)

# 核映射φ(x),这里用ReLU激活函数作为核函数的近似。
# 在线性注意力中,使用核函数来近似softmax
# 目的是将QK^T形式的二次复杂性O(L^2)降低为O(L * d)
def phi(x):
    # ReLU
    return np.maximum(x, 0)

# ---- 错误方式:右乘 K^T V (包含未来信息、非因果计算) ----
# phi(K)是形状(L, d)的矩阵,phi(K).T是形状(d, L)的矩阵,V是形状(L, d)的矩阵
KV_all = (phi(K).T @ V) 

# 计算输出Y
# 分子:phi(Q) (L, d)乘以KV_all (d, d)得到(L, d),这是未经归一化的Attention结果。
Numerator = phi(Q) @ KV_all

# 计算全局归一化因子Z
# phi(K).sum(axis=0, keepdims=True) 对K沿着序列长度L求和,得到 (1, d)
# .T 转换为 (d, 1)
# 这里的归一化因子Z也是全局计算的,没有考虑因果性。
Denominator = phi(Q) @ phi(K).sum(axis=0, keepdims=True).T

# Y_wrong:最终结果(L, d),每一行Y_wrong[i]都受到了整个序列K和V的影响
Y_wrong = Numerator / Denominator
print("错误方式 (包含未来信息):")
print(np.round(Y_wrong, 3))

# ---- 正确方式:前缀累加器 (只用前缀、因果计算) ----
# 这种方式是自回归模型在推理时应采用的线性注意力计算方法。
# S: 累加器,用于存储 K^T V 的前缀和,S的形状是(d, d),用于累积K^T V的乘积
S = np.zeros((d, d)) 
# Z: 归一化因子累加器,Z的形状是(d, 1),用于累积K向量之和
Z = np.zeros((d, 1)) 
Y_correct = []

# 遍历序列中的每一个位置i
for i in range(L):
    # 取出当前位置i的K, V, Q向量,并应用核映射
    # ki, vi, qi形状均为(d, 1)
    ki, vi, qi = phi(K[i:i+1]).T, V[i:i+1].T, phi(Q[i:i+1]).T
    
    # 增量更新S:S += K[i]^T @ V[i]
    # 累积所有历史token的(Kᵢᵀ Vᵢ),构建前缀信息池用来在当前步骤从所有过去信息中“读”出需要的内容
    S += ki @ vi.T            
    
    # 增量更新Z:Z += K[i]^T
    # 前缀累加所有K_i,作为注意力归一化因子;只包含所有过去的信息,确保当前token不会看到未来内容
    Z += ki                   
    
    # 计算当前位置i的输出y_i,分子:qi^T @ S_i;分母:qi^T @ Z_i
    # y_i形状是(1, d),只依赖于i之前(含 i)的信息
    y_i = (qi.T @ S) / (qi.T @ Z)
    
    Y_correct.append(y_i.flatten())

Y_correct = np.array(Y_correct)
print("\n正确方式 (只用前缀):")
print(np.round(Y_correct, 3))

# ---- 对比差异 ----
print("\n差异 (错误 - 正确):")

# 理论上,Y_wrong[i]应该包含Y_correct[i]加上未来信息的影响。
print(np.round(Y_wrong - Y_correct, 3))

输出(线性注意力机制的影响)

差异 (错误未添加累加 - 正确添加累加):

I -> [ 0.05 -0.196]

like -> [-0.097 -0.037]

deep -> [-0.058 0.028]

😄-> [-0. -0.]

这段代码展示了自回归模型LLM推理中,“偷看未来”会产生编码偏差,以及前缀累加如何保证因果性,差异向量(错误方式未使用累加 - 正确前缀累加)量化了未来信息泄露对每个token的影响,错误方式使用一次性全局计算 KTV ,在每步输出中引入了未来token的信息。结果显示:

  • 第一个token “I” 受到所有4个未来token的影响,偏差最大。
  • 随着序列推进,中间token受到未来个数token的影响逐渐减少,偏差依次减小。
  • 最后一个token “😄” 没有未来token,正确和错误方式输出完全一致,差异为零。

  前缀累加是线性注意力在自回归任务中的核心技术,它通过增量更新在保持 O(L) 计算复杂度的同时,严格保证因果性,防止未来信息泄露,从而确保每个token的表示仅依赖已生成的前文。因果掩码阻止模型看到未来,预测任务迫使模型根据前文推断下一个token,而注意力机制则帮助模型从前文抓取关键线索,三者协同作用使模型能够自动学习语言规律,实现高效且准确的自回归预测。

目前线性注意力机制的使用,还有最近 Kimi K3 的 KDA 等,在前面的章节提到过了。


  前面的 Transformer 机制说明了“数据为什么会被重复计算或重复搬运”,而算术强度进一步回答“这些计算相对于 HBM 数据搬运是否足够多”。因此,下面不再只描述 KV Cache 和矩阵乘法的形状,而是统一用 FLOPs/Byte 估计不同 kernel 的计算密集程度。

10.2.2 分析算术强度

符号定义

  • B:批大小batch size,一次处理的文本序列数量(推理通常 B=1);
  • D:token embedding 维度,即隐藏层维度;
  • T:下一次输出的 next token 数量。其中 Prefill 阶段为输入处理序列长度,而 Decode 阶段为 1;
  • S:已经生成的 token 数量(用于生成 next tokens 输入的上下文);
  • F:前馈网络中的中间维度(通常约为 D 的四倍);
  • H:注意力头数(Q 的头数);
  • Hkv :KV 的头数,MHA 时 Hkv = H;GQA/MQA 时 Hkv < H.

算术强度(Arithmetic Intensity)

  算术强度(AI)是衡量 Kernel级别 的计算任务计算密集程度的核心指标,用于预测任务的性能瓶颈是 compute-bound 还是 memory-bound

AI=FLOPsBytes Transferred
  • FLOPs:浮点运算次数,即矩阵乘法中加法、乘法等操作的总量,一次加法、乘法操作算是1 FLOPs。而一次矩阵运算中同时包含乘加操作,算作为2 FLOPs
  • Bytes Transferred:数据传输量,即Kernel执行该任务时需要从HBM读取和写入的数据总量(字节)。

  后续算术强度分析均假设数据类型为 bf16,即每个元素占 2 字节。Transformer 单层主要由 Attention 与 MLP 两部分组成,分别分析其 Arithmetic Intensity。无论是 MLP 还是 Attention,其 AI 分析本质上都需要考虑计算过程中的数据移动,包括从 HBM 加载输入数据与模型权重,经过片上存储进行数据复用,完成计算后将输出结果写回 HBM。

对于 FlashAttention 等优化方法,由于其通过 tiling、片上缓存以及算子融合减少了中间张量的 HBM 读写次数,从而降低了数据传输量,因为 FLOPs 几乎没有变化,因此提高了Arithmetic Intensity。

  为了统一接下来的分析,后续默认每个完整 kernel 均包含最终输出写回 HBM 的过程。Transformer 每个组成部分的数据流动过程可以抽象为:

text
HBM → L2 Cache → L1 Cache / Shared Memory → Registers → Tensor Core
      ↑                                                    ↓
      └────── L2/L1/Shared Memory/Registers 中的数据复用 ────┘
输出:Registers → L1/L2 → HBM

上面的箭头是对单个矩阵乘法 kernel 的数据路径的抽象:Tensor Core 通常从 Registers 中取数据,累加结果也先保存在 Registers ,最后才通过缓存层级写回 HBM。L2、Shared Memory 和Registers中的复用依赖具体的 tiling、算子融合和缓存命中情况,不能简单理解为所有权重都会完整驻留在 L2。后续公式统一把 HBM 读写作为主要数据传输量,而实际算子融合程度不同 HBM 数据流量也会不同。

 ① MLP 部分的核心操作是矩阵乘法。输入 x 的形状为 (B,T,D) ,通常先将前两维展平为 N=BT ,再计算形如 XW 的 GEMM。计算时需要从 HBM 读取输入和权重矩阵,完成计算后将结果写回 HBM:

a. SwiGLU 激活 3 个权重矩阵

SwiGLU(x)=Swish(xWgate)(xWup),MLP(x)=SwiGLU(x)Wdown

其中,权重矩阵形状 WgateRD×FWupRD×FWdownRF×D 。计算过程中共 3 次矩阵乘法,每次形状为 (B,T,D)×(D,F)(B,T,F)×(F,D)

FLOPs=3×2BTDF=6BDFT

  数据传输——读取输入 x(形状 BTD )+ 读取 3 个权重矩阵(形状均为 DF ) + 中间结果和最终结果写回(两个中间结果 BTF 逐元素相乘后产生 BTF ,最终输出 BTD )。只计算与 HBM 相关的矩阵乘法部分的输入输出:

Bytes Transferred=2(BTD+3DF+2BTF+BTD)=4BDT+6DF+4BFT

✍️注:每项乘以 2 是因为 bf16 占 2 byte。其中 BTD 出现两次分别对应读取输入和写回最终输出,2BTF 对应 gate 和 up 两个中间结果写回,逐元素相乘使用

AISwiGLU=6BDFT4BDT+6DF+4BFT

b.其他 GLU 变体之类的激活函数(也有门控)

  所有 GLU 类激活函数的结构都是 GLU(x)=σ(xWgate)(xWup) ,再接 Wdown 下投影。因此同样有 3 个权重矩阵、3 次矩阵乘法,FLOPs 与 SwiGLU 相同,Bytes Transferred 也相同。

c.标准 FFN 的激活函数(无门控,2 个权重矩阵)

  计算原理为 FFN(x)=ReLU(xW1)W2 ,这个过程中只有 2 次矩阵乘法,2个矩阵:

FLOPs=2×2BDFT=4BDFTBytes Transferred=2(BTD+2DF+BTF+BTD)=4BDT+4DF+2BFTAIFFN=4BDFT4BDT+4DF+2BFT

📖MLP 算术强度的简化

  以 SwiGLU 为例,将分母各项除以 2DF

AISwiGLU=3BDFT2BDT+3DF+2BFT

  当 DBFB 时,分母中 3DF 项远大于含 BFTBDT 的项,近似有:

AIMLP3BDFT3DF=BT

  标准 FFN 类似:AI4BDFT/4DF=BT

💡结论——当模型维度远大于 batch size 时,MLP 部分的算术强度近似为 BT


  MLP 的主要复用来自输入 token 在多个输出通道上的复用,以及权重 tile 在多个 token 上的复用,因此, Prefill 中 BT 较大时,矩阵乘法更容易获得较高算术强度,而 Decode 中每次通常只有一个新 token,矩阵形状变窄,权重读写的固定成本难以被充分摊薄。明确这一点后,下面分析 Attention,它除了 Q/K/V 和输出投影外,还要处理随上下文长度 S 增长的核心注意力。

 ② Attention 部分包含两类计算:线性投影,即计算 Q/K/V 、输出投影 以及 核心注意力,即QKT 分数计算 + Attn×V 加权求和( Attn=QKT )。线性投影部分、核心注意力部分有两种不同的情况 MHA 、GQA/MQA 。

🌟线性投影部分:

a.MHA 标准多头注意力 , Hkv = H

  投影矩阵 WQ,WK,WV,Wout 均为 D×D ,共 4 个:

FLOPsproj=4×2BTD2=8BTD2Bytes Transferredproj=2(BTD+4D2+BTD)=4BDT+8D2

  读取输入 xBTD )+ 4 个权重矩阵( 4D2 )+ 写回输出( BTD ):

AIproj, MHA=8BTD24BDT+8D2=2BTDBT+2D

DBT 时, AIBT ;当 BTD 时, AI2D

b.GQA/MQA 组查询, Hkv < H

  K、V的投影矩阵变为 D×Dkv ,其中 Dkv=(Hkv/H)D ,Q 和 O 的投影仍为 D×D

FLOPsproj,GQA=2×2BTD2+2×2BTDDkvH=8BTD2Bytes Transferredproj,GQA=2(BDT+2D2+2DDkv+BDT)=4BDT+4D(D+Dkv)AIproj,GQA=8BTD24BDT+4D(D+Dkv)

  权重参数量减少,从而减少了从 HBM 读取权重的数据量。GQA 的优势主要体现在减少 KV cache 的显存占用和带宽消耗,而非减少线性投影部分的 FLOPs

🌟核心注意力部分:

输入—— Q 形状 (B,H,T,dh)K 形状 (B,Hkv,S,dh)V 形状 (B,Hkv,S,dh) ,其中 dh=D/H

a.MHA 对核心注意力的影响

  • QKT(B,H,T,dh)×(B,H,dh,S)(B,H,T,S)FLOPs=2BHTSdh=2BTSD
  • Attn×V(B,H,T,S)×(B,H,S,dh)(B,H,T,dh)FLOPs=2BHTSdh=2BTSD .
FLOPscore,MHA=4BTSD

 数据传输(从 HBM 读 Q/K/V,写回输出):

Bytes Transferredcore=2(BTD+BSD+BSD+BTD)=4BDT+4BDS

  读取 Q( BTD )+ 读取 K cache( BSD )+ 读取 V cache( BSD )+ 写回输出( BTD )。

AIcore,MHA=4BTSD4BDT+4BDS=TST+S

✍️注:核心注意力部分的算术强度与 D 无关,仅取决于序列长度

b.GQA/MQA 对核心注意力的影响

  GQA 下 KV cache 大小缩小为 (B,Hkv,S,dh)FLOPs 不变因为 Q 仍有 H 个头,但 Bytes Transferred 中读取 KV cache 的部分缩小:

Bytes Transferredcore, GQA=2(BTD+2BHkvSdh+BTD)=4BDT+4BSDkvAIcore, GQA=4BTSD4BDT+4BSDkv=TST+S(Dkv/D)

GQA 通过减少 KV cache 的传输量,提高了核心注意力部分的算术强度,使 Decode 阶段更容易接近 compute-bound。

📖Attention 算术强度的简化

  将线性投影和核心注意力合并(以 MHA 为例):

FLOPs=8BTD2+4BTSDBytes Transferred=(4BDT+8D2)+(4BDT+4BDS)=8BDT+8D2+4BDSAIAttention=8BTD2+4BTSD8BDT+8D2+4BDS

  前面的公式给出了单个 kernel 的近似值,但实际推理还取决于请求处在哪个阶段。Prefill 一次处理一段已知输入,Decode 则每次只处理新生成的 token,即二者的 TS、矩阵形状和数据复用方式不同,因此必须分开讨论——

 对于 Prefill、Decode 阶段,Attention 算术强度简化结果不一样,而Prefill 与 Decode 阶段的关键区别在于 TS 的取值:

阶段T 的含义S 的含义
Prefill输入Prompt的长度等于T(self-attention,Q/K/V来自于统一输入)
Decode每次生成的 token 数量已生成的全部上下文长度

Prefill 阶段(T = S)

DBFB 的情况中有:

 MLP 算术强度表示 AIMLPBT=BS

 两种情况中的 Attention 核心部分算术强度,代入 T=S 可得 AIcoreSSS+S=S2

 同理可得,两种情况中的 Attention 线性投影部分算术强度 AIprojBT=BS .

 对于长序列S,算术强度很高,Attention也算是 compute-bound。

💡Prefill 阶段结论:只要 batch size 足够大或序列足够长,并且模型计算维度远大于 batch size,主要的性能瓶颈限制是 compute-bound,这就是 Prefill 阶段的吞吐量主要受 GPU 算力限制。

Decode 阶段(T = 1,只考虑逐 token 生成 )

DBFB 的情况中有:

 MLP 算术强度表示 AIMLPB×1=B

B = 1 时算术强度仅为 1,远低于 GPU 的 ,算是 memory-bound。增大 batch size 可以提高算术强度,这是 continuous batching 的核心动机之一。

 两种情况中的 Attention 核心部分算术强度代入 T=1 ,可得 AIcore=1×S1+S=SS+11(当 S1) ,算术强度趋近于 1,意味着每读取一个字节只做约一次浮点运算,是 memory-bound。

 两种情况中的 Attention 线性投影部分算术强度:AIprojB ,与 MLP 类似, BB 较小时是 memory-bound。

💡Decode 阶段结论:在典型推理场景(B=1 或较小)下,大部分性能瓶颈限制是 memory-bound。因此,在 Decode 阶段性能瓶颈主要在于 HBM 带宽——模型需要从显存读取全部权重参数和 KV cache,每读取一次仅做极少量计算。

这也是为什么推理优化的核心方向包括:

  • 增大 batch size(continuous batching):提高 MLP 和投影的算术强度;
  • GQA/MQA:减少 KV cache 读取量,提高核心注意力的算术强度;
  • 量化(INT8/INT4):减少权重和 KV cache 的字节数;
  • 模型并行:利用更多 HBM 带宽总量.

一个值得深入思考的问题是:为什么增大 batch size 往往能够提升算术强度,从而提高 LLM 的计算吞吐量,但这种收益是否存在上限,其边界又由什么因素决定?可以从以下 3 个角度进行分析:

算力饱和与“拼车效应”

  增大批次大小 B 是让GPU从“空闲”转向“饱和”的关键手段。可以把GPU推理想象成一辆大客车:模型权重是沉重的车身,输入数据是乘客。如果一次只拉一个乘客( B=1 ),每趟行程都要拖动几十GB的车身,非常浪费;而一次拉满 B 个乘客,固定成本就被分摊。在底层这意味着矩阵运算的并行性得到充分利用,算术强度显著提高——每从显存搬运一次数据,就能进行更多计算,从而充分发挥GPU的并行能力,提高整体吞吐量。

LLM推理阶段(自回归生成 + KV Cache)

          ┌─────────────┐
          │ 输入请求队列    ← 多个请求
          └─────┬───────┘

        ┌───────┴─────────┐
        │ 拼合多条请求一起   ← 提高总体吞吐
        └───────┬─────────┘

        ┌───────┴─────────┐
        │  GPU 计算单元    
        │  ┌─────────┐    
        │  │token_1  │<── K[K_1] V[V_1]
        │  └─────────┘  
        │  ┌─────────┐   
        │  │token_2  │<── 必须等Token_1完成
        │  └─────────┘
        |       |
        |       ↓   KV cache更新为K[K1,K2] V[V1,V2]
        │  ┌─────────┐
        │  │token_3  │<── 必须等Token_2完成
        │  └─────────┘
        |       |
        |       ↓   KV cache更新为K[K1,K2,K3] V[V1,V2,V3]
        |    ......
        └───────┬─────────┘

        ┌───────┴─────────┐
        │   KV Cache      │  ← 存储已生成token的键值对
        │   token_1...N   │  ← 每步生成都要访问全部历史
        └───────┬─────────┘

          单条请求延迟 ↑
          (串行生成 + 记忆负担)

说明:
1. 动态合批提升系统吞吐,但单条请求延迟仍受KV Cache访问限制。
2. 随序列长度增长,KV Cache越大,延迟越高。
3. 每生成一个新的token直接在原有KV Cache上追加这个的KV即键值对,不会生成新的cache。

串行诅咒与记忆负担

  然而,LLM的自回归生成特性决定了它无法像传统神经网络那样完全并行,也不同于训练阶段可以并行计算token。无论批次多大,生成过程都像“成语接龙”:必须先写完上一个字,才能决定下一个字。更麻烦的是为了保持上下文连贯,GPU每生成一个token,都需要访问此前所有token的状态。KV Cache存储了每个已生成token的键值对,用于下一步注意力计算。随着序列长度增加,这个“记忆库”会越来越大,使原本轻量的计算变成频繁查阅历史信息的重负,严重影响单条请求的延迟。

内存墙与速度瓶颈

  这就引出了推理的终极瓶颈——内存墙。在解码阶段,GPU的计算单元往往空闲等待显存将庞大的KV Cache数据搬运过来,而实际计算时间非常短。虽然动态合批可以通过同时处理多条请求提升整体吞吐量,但这只是让“大客车一次拉更多乘客”,并不能让车速本身变快。对于单个用户,其生成延迟仍然被显存带宽限制:GPU算力再强,如果数据读取速度供不上,生成速度就无法提升。

解码阶段会依据模型输出的词汇表概率分布,选取一个token作为生成结果如取最大概率、按分布采样等,并将其作为下一步的输入继续生成。

总的来说就是,批量 B 可以显著提升系统总体吞吐,但自回归生成和不断增长的KV Cache决定了单条请求的延迟最终受显存带宽限制。随着生成推进,即使系统仍能对多条请求进行合批,单条请求自身每一步解码的计算仍然只能处理1个token,并且随着KV Cache变大,其算术强度不断下降,性能表现会逐渐逼近单独处理即B≈1时的有效计算强度。


  算术强度分析解释了“为什么不同阶段受不同硬件资源限制”,但用户感受到的模型性能还要通过 TTFT、ITL/E2E Latency 和 Throughput 来描述;同时,调度策略会把单请求的 Decode 组织成系统级批处理。因此,下面把内核级瓶颈和服务级指标连接起来。

10.2.3 延迟 vs 吞吐量

  在分析算术强度之后,还需要理解其在实际 LLM 自回归生成推理场景中的性能影响。其中,Latency与 Throughput 是衡量 LLM 推理性能的两个核心指标。

核心性能指标

指标优化目标
延迟降低,以提升用户交互流畅度
吞吐量提高,以提升系统整体处理能力资源利用率

  值得注意的是,追求极致延迟通常针对单请求优化;而追求极致吞吐量通常倾向于增加批次,但批次大小并非吞吐量或延迟的唯一决定因素——大批次不一定带来高吞吐量,小批次也不必然意味着低延迟,实际效果还取决于GPU利用率、请求长度分布KV Cache效率以及调度策略等多重因素。

批量处理

  KV Cache减轻了计算和带宽压力后,GPU核心利用率提升,可以通过批量处理进一步优化吞吐量。

批量处理方式

  1. 静态Batching:将固定数量的请求或token打包。

  2. 动态、持续Batching

    • 在等待生成新token时立即打包下一个请求。
    • 多请求拼Batch,不同长度和进度的请求动态组合,提高GPU并行利用率。
text
请求1: T1,T2,T3 | 请求2: T1,T2 -> 动态拼成Batch -> GPU并行计算

增益:提高整体吞吐量,特别适合在线服务、请求长短不一的场景。

权衡

目标权衡点说明
延迟 vs 吞吐量批次大小增加batch size → 吞吐量提高,但单请求延迟可能增加
效率 vs 精度优化技术选择KV Cache或量化等优化必须保持自回归因果性语义准确性
显存 vs 性能KV Cache占用随序列长度线性增长,限制最大序列长度或批次大小

总结

  1. KV Cache 能以额外显存和带宽开销换取更低的重复计算成本,但是当其规模过大时内存压力会抵消其收益,反而降低生成速度。
  2. 批量处理提升整体吞吐量。
  3. 结合两者可在高吞吐量和低延迟之间实现工程平衡。
  4. 工程优化必须考虑显存、序列长度、请求长度分布等实际因素。

  其中,批处理主要改善生成阶段的资源利用率,而输入 prompt 的长度决定了 Prefill 工作量、初始 KV Cache 大小以及后续每一步需要读取的历史长度。由此,提示词压缩成为另一条直接作用于输入规模和 KV Cache 的优化路径。

10.2.4 提示词压缩

  在自回归LLM中,用户输入的提示词(prompt)并不是“免费”的——模型首先将提示分词并映射为嵌入,然后在每一层计算对应的Key、Value并把这些中间KV存入缓存,以便后续逐步生成时复用。因此,过长或过于详尽的prompt不仅会占用大量显存尤其是在上下文窗口变长时,还会增加每个token的注意力计算量,从而拉长延迟并降低整体推理效率。此外,过多、冗余或无关的信息还可能“稀释”模型的注意力分配,导致输出质量并不随上下文长度线性提高,反而变差。为在尽量缩短prompt的同时保留关键信息,常见的提示压缩策略可分为两类:

1. 硬提示压缩

  硬提示压缩旨在通过缩短长度降低复杂度来简化自然语言提示,同时仍能有效引导模型生成期望的响应。常用方法包括两类:

7fc0df8fcb8724d813992dc8fdc290bd

图10.6 硬提示词压缩

  • 选择prompt中的关键内容:

    • 基于句子、token相关性筛选,利用句子编码器或文本嵌入,将prompt的句子或token与查询内容进行匹配,按相关性排序,剔除低相关性部分;对于大型LLM,可根据模型容量限制适当减少token数量,保留语义最相关的token。
    • 动态控制token数量,可以通过强化学习训练一个token控制器,根据查询难度动态调整prompt长度,以兼顾效率与输出质量。
  • 改写与压缩提示:

    • 简洁化表达,将冗长提示压缩为更短、清晰的表述,同时遵守生成长度限制,保持原句子语义不变;
    • 保留核心上下文,通过动态注意力或文档摘要策略,从大量文档中提取最重要的内容,融合到简化后的prompt中。

2. 软提示压缩

  软提示词摆脱了对离散token的依赖,通过学习一组连续向量来替代或补充自然语言提示。在训练阶段,这些向量作为可学习参数,通过反向传播引导模型产生期望输出;在推理阶段,这些向量可直接作为Transformer输入的前缀或以特定方式注入注意力机制,无需映射回离散token,直接以embedding的形式参与计算,常见的软提示词策略可分为两类:

80033b09-1b76-4fbd-8fd9-0693bc063cea

图10.7 软提示词压缩

  • 固定LLM参数的方法:通过优化输入表示,在保持原始LLM参数冻结的同时,提供了多样化的提示压缩策略,这些方法显著提升了计算效率与内存利用率,使得LLM能更高效地应用于广泛的任务场景。
  • 基于要点词元的方法:通过将上下文转化为紧凑且可复用的词元提供了一种强大的输入长度缩减框架,这类方法带来了显著的效率提升,但通常需要更新LLM的参数,以实现有效的压缩与集成。

其中软提示的优点是参数量小、可复用,并且在不微调模型主体权重的情况下即可获得良好性能。然而,需要注意的是:

  • 推理开销:软提示并不自动降低推理计算,因为它们仍需参与整个模型的前向传播和注意力计算;
  • 可解释性:软提示是连续向量,在冻结LLM主体的情况下被优化为引导模型生成特定输出,但向量本身不对应任何可读的词语或语义概念。每个向量存在于高维空间中,不同维度的组合共同作用于模型行为,因此可解释性较低且不直观;
  • 泛化能力:软提示通常是针对特定任务或训练数据学习的,引导效果对训练数据之外的情况较敏感,泛化能力通常不如显式的自然语言提示或对模型主体的微调,对prompt风格变化尤其敏感。

从推理机制上看,硬提示词与软提示词都会扩展输入序列长度,并参与自注意力计算。其不同之处在于,硬提示词由离散token经embedding映射得到,而软提示词则以通过学习好的额外可学习参数得到的连续向量(不改变LLM原有结构)直接注入模型输入表示中。

3. 在视觉层面进行prompt压缩

47a0cbac2c0e680d7ecc62303f44340d

图10.8 DeepSeek-OCR结构

在视觉层面进行prompt压缩可以分为三个步骤:

Step1 文本 → 图像:将原始文本prompt转化为一张图片,可以是文字渲染图或经过特殊视觉编码的图像表示;

Step2 模型读取:通过OCR或视觉编码模块如CLIP的文本编码器或专门的图像文本理解模块提取图像中的信息;

Step3 压缩效果:视觉编码将长文本压缩为固定维度的图像表示,从而减少序列长度token数量对自回归模型的计算压力。

10.3 相关研究

  前面的讨论主要围绕标准自回归 Transformer 展开:通过 KV Cache 减少重复计算,通过 GQA/MQA、算子融合、批处理和提示词压缩降低数据搬运或调度成本。但这些方法大多仍保留“逐 token 生成”的基本范式。接下来转向模型结构和生成范式本身,观察能否从记忆形式、并行方式或大小模型协作等层面改变瓶颈。

  本节按三条思路展开:状态空间模型用连续状态替代显式的长 KV Cache;扩散/块扩散模型放松 token 级自回归约束,以换取块内并行;草稿模型与推测解码则保留目标模型的自回归语义,通过“小模型起草、大模型验证”摊薄大模型的串行解码成本。它们都在追求更快推理,但所承担的质量、训练成本和系统复杂度并不相同。

10.3.1 状态空间模型

76abd2a89506f7f433850cc9ec07ae58

图10.9 状态空间模型

  状态空间模型(SSM)可以被视为对传统RNN的一种扩展,其在信号处理和长序列建模方面表现出色。近年来,有研究尝试将SSM应用于自然语言上下文建模,尤其在捕捉远程依赖方面具有优势。与自回归Transformer依赖显式KV缓存不同,SSM通过隐状态的连续更新以及训练阶段学习得到的权重矩阵实现对长期依赖信息的建模。在长序列处理上,SSM的计算复杂度低于标准自回归Transformer,因此速度通常更快。

  然而,部分SSM实现仍可能在处理极长序列时面临长期依赖保持的挑战。具体而言,在训练过程中,隐藏状态的反向传播可能出现梯度消失或爆炸,导致长期信息衰减从而在某些文本生成或理解任务中影响模型的整体性能。为了改善SSM在长期依赖建模中的能力,出现了一些新的线性注意力机制,例如:

  • Mamba:在标准SSM的基础上优化了状态矩阵和数值稳定性,提升长序列建模能力。
  • 其他基于RNN的方法(如LSTM):不是SSM的直接变体,其通过门控机制解决长期依赖问题。
  • S5:在S4的基础上,S5对连续时间状态空间进行离散化,使卷积核能够高效作用于整个序列。这种离散化方法在自然语言处理等任务中非常有效,同时通过调整输出结构和状态更新方式,在一定程度上缓解了长期依赖信息衰减的问题。

  SSM 的核心变化是“如何保存历史”:它把历史信息压缩到递归更新的状态中,从而没有了显式 KV Cache 的依赖,但它仍然需要处理序列状态(通过存储的隐藏中间状态)。下面进一步改变“如何生成”:扩散模型尝试在位置维度或块内并行恢复多个 token,以减少严格逐 token 解码带来的串行等待。

10.3.2 扩散模型

ffe1238b-f064-4372-a17e-d782ee244351

图10.10 扩散模型

  在10.2.1 Transformer中提到,虽然扩散模型也可以存储一些不随扩散更新的KV状态信息,但由于扩散模型可以在空间位置上并行生成多个token,因此其推理速度在很大程度上不受显存限制影响,特别是在处理大批量生成任务时优势明显。

  然而,扩散模型的设计擅长处理空间维度的信息,而自然语言本质上是时间序列,存在严格的因果依赖关系。这就导致扩散模型在生成长文本或有逻辑依赖的内容时容易出现问题,例如生成一段代码时可能速度很快,但容易出现明显的逻辑错误或语法错误。

  换句话说,扩散模型在完成推理的速度上有优势,但在保证上下文一致性和因果逻辑方面不如自回归模型。因此,想要把扩散模型运用于LLM中需要面临的挑战: 1.上下文一致性不足

  • 扩散模型本质上是非自回归(NAR)并行生成,难以严格保证长序列的因果依赖与全局语义连贯性。
  • 对于LLM任务,尤其是长文本生成或复杂推理,单纯并行生成的token容易出现逻辑或语义不一致。

2.离散token对齐问题

  • 传统扩散模型依赖预测噪声或score,再映射回离散token,这会在token分布空间引入额外偏差。
  • 若直接将扩散生成序列作为LLM输入,需要处理token对齐和概率分布匹配问题。

3.推理效率与 KV Cache 调度 尽管扩散模型并行生成可降低逐token前向计算,但在与自回归模型结合时,需要设计合理的proposal‑verify流程,保证KV Cache调用次数显著减少,否则仍可能产生高延迟或显存压力。

LLaDA 2.0率先将扩散语言模型做到千亿体量发展过程

1.掩码去噪语言模型(MDLMs)背景

  掩码去噪语言模型提供了一种新的文本生成思路。与AR模型不同,MDLM将文本生成视为一个迭代去噪的过程:在每次前向传播中,部分词元会被随机掩码,模型的任务是根据未被掩码的上下文恢复原始词元。

  这种范式的转变让研究者开始尝试从零训练MDLM,以探索其潜力。例如,LLaDA表明——一个完全从零训练的8B规模MDLM在表现能力上可以与同等规模的自回归模型相媲美。进一步地,LLaDA-MoE首次将混合专家架构 引入MDLM,结果显示基于MoE的MDLM在效率和能力上均优于密集模型,证明了MDLM与先进MoE架构的兼容性和可扩展性。

  需要注意的是,由于MDLMAR模型在训练动态上存在根本差异,AR模型常用的训练策略和超参数往往不适用于MDLM。为此,一些研究工作如 Quakka和OpenMoE2专门研究了MDLM的缩放特性和训练策略,为新兴的去噪生成范式提供了系统性指导。

2.利用自回归模型初始化去噪模型

  考虑到AR模型具备强大的知识容量和优秀性能,近期研究尝试用预训练AR模型初始化MDLM或DLM,以降低训练成本并缩小性能差距,例如:

  • DiffusionLLaMADream-7B

    • 在训练过程中采用掩码退火策略,逐步将因果注意力过渡为双向注意力。
    • 同时结合基于CART的损失重加权方案,平衡词元级学习动态。
  • RND1

    • 在初始化阶段就将AR模型的因果注意力直接转换为双向注意力。
    • 为保留原AR模型的知识密集能力,训练时限制了密集层参数更新,防止灾难性遗忘。

3.块扩散语言模型(BDLMs)

  块扩散语言模型提出了一种混合生成范式,将扩散建模与自回归建模结合起来:

  • 块内部:通过扩散过程恢复被掩码的词元。

  • 块之间:以自回归方式顺序生成。

  • 优点:

    • 支持可变长度文本生成。
    • 可以复用KV-cache,提高推理效率。

基于AR模型初始化的BDLM,例如SDAR在不同块大小和优化策略下,能够达到与AR基础模型相当的性能。

块(Block)的定义说明

LLM中引入的扩散模型,所谓的“块”指的是在文本序列的位置维度上严格连续的一段token,其划分目的是在保持语言因果结构的前提下引入并行去噪。这一概念不同于Transformer中基于语义相似性、结构信息或计算效率所进行的分块。

扩散块的划分示意

token索引值:   1    2    3    4    5    6    7    8    9    10   11
               │    │    │    │    │    │    │    │    │    │   │
片段:          x1   x2   x3   x4   x5   x6   x7   x8   x9  x10  x11
              └───────条件输入不参与扩散───────┘   └──── block ────┘

同一个block中的token会被并行处理即同时生成、修改。

4.现有局限

  虽然这些方法在中小规模(7B~30B)上表现良好,但仍存在几个挑战:

  • 可扩展性有限,基于AR初始化的扩散模型在超大规模上训练的可行性尚未充分验证。
  • 训练效率低,块扩散训练步骤多,限制了在大规模语料上训练超大模型的效率。

5.LLaDA2.0训练范式

  研究不是从零开始训练扩散模型,而是采用在AR模型训练过程采用扩散风格输入数据,也就是AR迁移学习转换为扩散模型,这一范式转换有3个阶段:(1) 从AR到全序列掩码扩散模型的持续预训练;(2) 块扩散预训练,实现从标记级到块级扩散建模的过渡;(3) 后续进行有监督微调、直接偏好优化,以实现对齐并进行任务专项优化。

9d77db32117af2ebdc105b2e4e41e841

图10.11 LLDaA2.0训练过程图示

Step 1:渐进式块大小预热

  渐进式块大小预热阶段的核心思想,是通过逐步增大生成块的大小,引导模型逐渐适应更大范围的并行去噪与上下文建模能力

  在训练初期,模型以接近自回归的方式工作为每个block的大小较小在极端情况下等价于单个token,模型主要学习在严格因果约束也就是遵循时间序列情况下进行稳定预测。随着训练推进block大小逐步增大,模型开始在块内并行预测多token,同时仍保持块间的自回归顺序,这一过程可以理解为:

  • 块内感受域逐步扩大 随着block大小的增加,模型需要在同一个block内联合建模多个连续位置的token。这些token在训练时通常处于被噪声扰动、被mask的状态,模型通过一次前向传播block内的多个位置进行并行预测与一致性建模,而不再局限于单一token的逐步生成;
  • 注意力结构逐步转变 训练早期,模型主要遵循严格的token级自回归因果注意力,随着 block扩大模型在block内允许使用非因果(双向)注意力或去噪式建模,以充分利用块内上下文信息进行联合恢复。同时,block与block之间仍然保持自回归的因果顺序,从而在保证长程依赖建模稳定性的以及训练模型有边思考边修正的能力。

  通过这种渐进式训练策略,模型无需在早期就直接面对大规模并行去噪的困难任务,而是平滑地从自回归建模过渡到扩散、去噪风格的表征学习。这有助于:

  • 稳定训练过程,避免大block带来的优化不稳定;
  • 逐步提升模型对长上下文、复杂掩码模式以及块内联合一致性的建模能力;
  • 使AR初始化的模型权重更自然地适配后续的扩散模型生成范式。

由于不同块之间的token可能在语义上不相关,但注意力机制仍会尝试分配权重,这使得扩散模型在跨块建模时存在形成虚假依赖关系的风险,从而可能引发语义混淆以及双向注意力机制训练不稳定。为避免这种跨块干扰,我们引入文档级注意力掩码(mask),将自注意力限制在单个块内部,从而保证上下文建模的一致性和训练稳定性。

Step2:对第一阶段得到的模型进行稳定训练

  在第一阶段的大规模训练完成后,当块大小固定为4096并且模型切换至MDLM模式时,每个块都会占用一定的KV缓存。此时,跨块注意力计算中与当前块关联度较低的信息可以被屏蔽,从而显著降低计算开销,使数据在MDLM范式下得以更高效处理。

  随着模型适应该模式,稳定训练阶段的重点转向在经过筛选的大规模语料库上进行广泛训练,以进一步提升第一阶段模型的扩散表征能力。在此阶段,每个输入序列作为单个块进行处理,块大小固定为4096,使块内注意力覆盖整个序列,从而保证块内上下文完整建模,并等效于经典的MDLM设置。

Step 3:全局注意力转化为局部注意力

  在前两个阶段的大规模MDLM训练完成后,模型已经学会了全局上下文信息。接下来,通过逐步将块大小从4096缩小到更小的值,模型逐步从全局块条件化过渡到局部块条件化BDLM模式。块大小减小意味着注意力计算只在单个块内进行,这显著降低了显存消耗并加快推理速度,同时保留局部块内的上下文信息以理解语义。为了不完全丢失长距离依赖,模型在衰减过程中仍有部分全局注意力,使其能够一定程度上捕捉跨块关系,从而在提高效率的同时保留对远距离文本的建模能力。

BDLM和MDLM都是扩散语言模型中的注意力策略。MDLM强调捕获全局依赖,支持块内与跨块的全局注意力以建模长程关系;而BDLM则侧重局部上下文,主要将注意力限制在单个块内,从而在满足必要因果约束的前提下显著提高计算效率并降低显存开销。

LLaDA2.0成果   在更大模型LLaDA2.0-flash中,这一潜力尤为明显。该模型平均得分达73.18,与Qwen3-30B-A3B-Instruct-2507(73.60)等顶尖AR模型不相上下。更重要的是,LLaDA2.0-flash在复杂生成任务编码能力智能体能力以及高级数学等方面均优于同类AR模型,显示出扩散架构的潜在优势

  值得注意的是,LLaDA2.0系列展示了基于扩散的语言模型的巨大潜力。它们不仅可以作为AR的可扩展替代方案,而且还能与AR兼容。虽然在一些通用测试上,它们与传统模型的差距正在快速缩小,但在代码生成和工具使用等复杂、结构化、智能体任务中,扩散模型已经表现出更强的能力。这意味着,未来在语言生成领域中扩散模型有望成为一个非常有前景的研究方向。

  扩散和块扩散说明了一种更激进的方向,通过放松 token 级因果顺序换取并行效率,但这也带来离散 token 对齐、上下文一致性和 KV Cache 调度等问题。如果希望保留目标模型的自回归分布,又尽量减少目标模型逐 token 的调用次数,就自然想到推测解码这个办法。

10.3.3 推测解码

9a1fd54c917fb77830954a2bcfea8384

图10.12 大模型与小模型协作

  大型语言模型在完成一些文本任务时以牺牲推理延迟为代价取得显著效果,因此有大小模型协作完成任务的思路解决在尽可能保持能力的前提下,减少延迟问题。关于大小模型协作有两个策略,它们虽然都在解决同一个问题,但是有着本质的不同:

  • 推测解码(Speculative Decodeing):小模型以自回归方式一次生成一段候选token序列,大模型对该序列进行并行前向计算并验证其与目标分布的一致性;若在某一位置出现不一致,则接受此前一致的token前缀,并从该位置回滚然后大模型接管生成,这一过程不断重复直到完成相应问题为止。
  • 级联解码:小模型对当前生成步骤进行一次前向计算,并基于其预测token分布(生成token之前)的不确定性 (如最大概率或熵) 决定该步token的生成是由小模型直接生成还是交由大模型生成一个置信度,整个过程就是小模型预测相关概率是否达到阈值,达到则继续按照预测让小模型生成token,否则由大模型重新响应计算生成token,重复直到完成相应的问题为止。

一句话总结:推测解码通过“并行验证生成token的一致性”来减少大模型的串行解码步数;级联解码通过“基于token分布的不确定性进行风险感知路由”来减少大模型的调用频率。

10

动画演示:对GSM8K数据集中一道算术题(玛丽一共有30只羊。她每天从其中一半那里取1公斤牛奶,另一半从另一半那里取2公斤牛奶,那么她每天收集多少牛奶?)的推测级联和推测解码方法进行比较。草稿token以黄色显示,已验证token以红色显示。推测级联方法得出了正确答案,且速度比推测解码更快。

Google团队在分析推测解码与级联解码两种策略时发现:推测解码通过允许小模型生成与目标分布可能存在偏差的草稿序列,并由大模型并行进行可接受性验证,从而显著减少大模型的串行解码开销;而级联解码则通过路由机制在生成前判断当前步骤是否需要调用大模型,以保证最终生成质量不低于大模型本身,但该策略由于引入了额外的决策与模型切换过程,在延迟上的优化幅度通常不如推测解码显著。

c87d2598268d9fb609e2dbb2806c7abc

图10.13 不同数据集中的Gemma 2B与Gemma 27B使用的质量与拒绝率曲线图( \gamma=1 )

  图 10.13表明,大模型一致性检验小模型生成token拒绝率高于大模型置信度检验小模型预测token分布,因此通过大模型概率判断小模型生成block的可接受性,比直接比较模型分布更能平衡质量与延迟。于是Google团队在分析推测解码与级联解码的质量–成本权衡后提出了推测级联(Speculative Cascades)——该方法构建由多个规模递增的模型组成的级联结构,小模型以自回归方式生成草稿block,并由更大的模型对该block进行并行可接受性验证;若验证通过,则直接接受该block并由小模型继续生成后续草稿,否则丢弃该block并升级由更大模型重新生成。通过这种方式,在保证最终输出质量由最大模型兜底的前提下,最大化利用小模型完成可接受生成从而实现推理阶段质量、成本与延迟的折中。

推测级联中block指的是:由某一个模型在一次自回归草稿阶段连续生成的、长度为k的token序列。

推理级联和级联解码有相似的思路,前者处理对象是block,而后者处理的是token

举例说明推理级联的使用 假设问题为:“Datawhale 成立的时间是什么时候?”

  1. 小模型草拟起步 小模型快速生成候选block,例如:[Datawhale, 成立, 时间, 于, …],相当于一次性block作为草稿。

  2. 大模型验证 大模型对该block做前向计算判断其可接受性,而不是逐token强制生成。

  3. 递延规则判定

    • 接受:小模型生成的block被认为可靠 → 继续生成下一个block。
    • 拒绝:小模型block被认为不可靠 → 由大模型重生成该block,并由小模型继续生成后续block同时大模型继续验证。
  4. 迭代生成 该过程不断循环,直到生成完整答案,保证最终输出由大模型兜底,而小模型尽量完成可接受生成。

优势说明:这种“小模型起草 + 大模型验证 + 回退”的方法非常灵活,递延规则可针对不同场景定制。例如可以优先保证重要信息由大模型审核,而非关键位置由小模型直接生成,从而在提升推理速度的同时保持答案质量

15cd6d92fd28f280aed55007e61c39b3

图10.14 大模型与小模型协作效果

图10.14中推测级联变体(蓝色和橙色)在数学推理和总结任务中,相较于标准推测解码(绿色星)实现了更好的质量与延迟权衡

  根据上图可以分析到,在相同token产生的速度情况下,大小模型协作的方法在token预测准确度、内容顺序一致性评估明显更好,证明这种方法的有效性,也为加速LLM推理提供了一种新的思路。

Speculative Decoding 是当前许多 LLM 推理框架采用的重要优化技术,例如 vLLM、DSpark、JetSpec 等。一个值得思考的问题是:为什么在相同的推理场景下,小模型生成一个 token 的速度通常比大模型更快,计算成本也更低?这个问题可以从生成一个 token 所需的计算量与内存访问开销进行分析。

模型在生成 token 的过程中,需要依次经过所有 Transformer Layer 完成一次前向计算。因此,模型规模越大,每生成一个 token 所需的计算量和内存访问成本通常也越高。为了便于理解,下面以一个简化模型为例进行估算:假设模型共有 25个Transformer Layer,隐藏维度为 128,采用 16个Head(即每个 Head 的维度为8),数据类型为 FP8(1 Byte)。那么,仅考虑每层隐藏状态的读写,其数据量计算的内存开销约为:

Activation=25×128×1=3200 Bytes(3.2KB)

实际推理过程中,每层还需要读取 Attention、MLP 等模块的大量模型权重并访问 KV Cache,因此真实内存访问量会远高于上述估算的 Activation,并且会随着 LLM 的模型扩大而增加

从计算量来看,每个 Transformer Layer 主要由 Attention 和 MLP 两部分组成,其计算复杂度可近似表示为:

Per Layer FLOPs16d2+2nd

其中, d 为隐藏维度, n 为当前上下文长度。因此,该模型生成一个 token 的计算量可近似表示为:

Total FLOPs25×(16×1282+2×n×128)

从上述的简单估算(实际计算量会因采用标准FFN、SwiGLU、MLA 等不同模型结构而有所差异)可以看出,即使仅生成一个 token,也需要依次完成所有 Transformer Layer 的前向计算,并不断读取模型权重、访问 KV Cache 等数据。随着模型参数规模不断增大,每个 token 所需的计算量和内存访问量都会相应增加,因此推理延迟也会随之上升。这也是在相同硬件条件下,大模型生成 token 通常比小模型更慢的重要原因之一。

10.4 MTP + 推测解码

  前面介绍的推测解码能否获得明显加速,取决于草稿模型能否连续提出较长且容易被目标模型接受的 token block。由此可以进一步追问:能否在训练阶段就让模型学习未来多个 token 的连续条件依赖,为推理阶段的多 token drafting 提供更好的表示基础?MTP(Multi-Token Prediction)正是围绕这一问题展开的。

  需要先区分两个概念:MTP 是训练阶段的预测目标或结构设计,推测解码是推理阶段的 proposal–verify 算法;MTP 可以为推测解码提供帮助,但并不等于推测解码本身,也不会自动把一次前向变成多个已经被目标模型验证的 token。

10.4.1 DeepSeek MTP

MTP vs no MTP效果对比:

  DeepSeek-V3的MTP并不是使用多个独立输出头基于同一个backbone hidden representation并行预测未来token,而是采用按时间顺序预测深度递进的级联式结构——后续预测深度会基于前一阶段演化得到的预测表示(stage representation)继续进行条件化建模。

📖Tips:这种递进过程发生在连续隐表示空间中,而不是显式地将前一步生成的离散token再输入模型进行自回归的rollout。因此,MTP保留的是预测深度之间在多时间的隐表示层面条件依赖关系,而非传统自回归sampling意义下的token-level rollout。

DeepSeek MTP vs Traditional Parallel MTP:

text
DeepSeek Sequential MTP              | Traditional Parallel MTP
h_t                                  | same hidden representation h_t
 ↓                                   |   ├── predict x_{t+1} (token1)
stage1 representation                |   ├── predict x_{t+2} (token1, token2)
 ↓ predict x_{t+1}                   |   ├── predict x_{t+3} (token1, token2, token3)
 ↓                                   |   └── ...
stage2 representation                |
 ↓ predict x_{t+2}                   |
 ↓                                   |
stage3 representation                |
 ↓ predict x_{t+3}                   |                           
...                                  |
0abeb3c5-ccef-4def-a37c-c7034a8cb7f9

图10.15 DeepSeek MTP

DeepSeek V3 报告原文:Our principle of maintaining the causal chain of predictions is similar to that of EAGLE, but its primary objective is speculative decoding , whereas we utilize MTP to improve training.

  DeepSeek V3 MTP本身属于训练目标设计,而推测解码属于推理阶段的加速算法。前者并不直接加速模型训练,而是能够为后续多token推测解码提供更适合的representation学习基础。

10.4.2 Gemma4 MTP

image

图10.16 Gemma 4推理加速的设计

2026 年 5 月,谷歌在Gemma 4的推理系统中引入了基于draft model与推测解码的多token drafting结构,用于提升LLM推理速度。虽然Gemma 4与DeepSeek V3都涉及对未来多个token的预测,但两者的目标与实现方式并不完全相同

  • DeepSeek V3 MTP属于训练阶段的多步预测目标,其核心目的是增强模型对未来语义演化与多步生成一致性的建模能力
  • Gemma 4则是在推理阶段多token预测推测解码结合,用于减少target model逐token自回归解码(decoding)带来的推理开销。

两者的共同点在于——都不再局限于传统单步next-token prediction的单步条件预测目标,而是尝试基于已有上下文的中间状态对未来多个token进行连续条件化预测

10 17

图10.17 草案模型协作示意图

  在Gemma 4中,draft model会在target model提供的上下文状态(如共享的KV Cache或部分中间激活状态)的条件基础上自主预测未来多个token,然后再由target model对这些“草稿”token序列进行验证。target model会接受其中与自身预测一致的前缀;若出现不一致,则从第一个不匹配的位置恢复target model的正常自回归解码。

  考虑KV Cache情况,Transformer(层数为 Nlayers )生成单token的主要计算开销通常仍近似与模型层数与hidden dimension( d )成正相关,其复杂度可近似写为:

O(Nlayersd2)

  因此,不直接让target model自身完成多个token rollout,是因为大参数模型逐token生成的计算成本较高。对于具有更多层数与更大隐藏层维度的target model,其单token推理成本通常远高于参数较小的draft model。结合推测解码 + MTP ,Gemma 4能够显著提升LLM的整体推理速度。

  除了刚才提到的改进以外,还有当词表权重规模较大时完整矩阵乘法运算与softmax计算会带来较高的计算与内存带宽开销。为了进一步降低这一部分成本,Gemma 4的推理系统结合了基于词表分块与候选筛选的优化思路,其核心思想包括:

10 18

图10.18 draft model优化处理

  • 将大规模词表通过聚类的方法划分为多个候选词表分块,提取相应的embedding representation;
  • 根据当前的logits,对不同词表分块计算候选得分;
  • 仅在得分较高的候选词表分块上执行后续矩阵乘法运算与softmax计算,从而减少完整词表计算带来的开销。

  因此,Gemma 4的整体推理加速并不仅仅依赖speculative decoding、draft model与KV Cache共享,还结合了高效词表分块与候选筛选优化。

10.5 总结

  本章围绕一个核心问题展开:在模型能力基本确定之后,为什么同一个 LLM 在不同硬件、不同请求规模和不同生成阶段会表现出不同的速度与成本,又应当如何针对瓶颈进行优化?答案可以概括为一条链路:

训练得到能力自回归推理呈现能力Prefill/Decode 产生不同数据流算术强度决定硬件瓶颈调度、缓存、压缩和结构改造进行优化

开头问题回顾

开头问题本章结论
如何评价推理性能?TTFT 主要反映请求进入模型后到首 token 的等待;ITL/TPOT 反映 Decode 阶段相邻 token 的输出间隔;E2E Latency 反映单请求完整耗时;Throughput 反映系统单位时间处理的 token 数。低延迟和高吞吐是不同目标,需要结合场景评价。
训练和推理为什么不同?训练使用已知的真实前文和因果掩码,可以在时间维度并行计算整段序列,并保存激活用于反向传播;推理使用模型生成的前文,必须在 token 之间串行,但每一步内部的矩阵运算仍可并行。
Prefill、Decode 和 KV Cache 解决什么问题?Prefill 一次处理 prompt 并建立历史 KV;Decode 每次处理一个新 token,并读取已有 KV。KV Cache 避免重复计算旧 token 的 K/V,但会随上下文长度增长,带来显存容量和 HBM 带宽压力。
如何判断 compute-bound 或 memory-bound?用 FLOPs 除以 HBM 数据传输量得到算术强度,并与硬件的计算峰值/带宽比比较。Prefill 通常因 token 数较多而更容易提高复用、接近 compute-bound;小 batch Decode 通常因权重和 KV Cache 搬运占主导而接近 memory-bound。
优化方法分别解决什么瓶颈?FlashAttention、算子融合和 tiling 减少中间张量读写;GQA/MQA 减少 KV Cache;量化减少权重和缓存字节数;continuous batching 摊薄权重读取并提高吞吐;提示词压缩减少 Prefill 和初始缓存规模;SSM、扩散改变记忆或生成范式;推测解码减少目标模型的串行解码工作,MTP 则可为多 token drafting 提供更好的训练表示基础。

  因此,本章最重要的不是记住某一种“万能加速方法”,而是建立按阶段定位瓶颈的习惯:先区分 Prefill 与 Decode,再判断问题主要来自计算量、HBM 流量、KV Cache、请求调度还是自回归串行依赖,最后选择相应的优化。所有方法都存在代价——更大的 batch 可能增加单请求等待,缓存压缩可能损失信息,非自回归生成可能削弱因果一致性,草稿模型需要额外模型和验证流程;真正有效的推理系统必须在质量、延迟、吞吐、显存和成本之间做整体权衡。

思考

1)扩散模型(DMs)应用于LLM推理时,展现出非自回归(NAR)的并行生成优势,有效降低了推理延迟和KV Cache显存占用。然而,NAR特性牺牲了序列的上下文强依赖性,有什么办法能将扩散模型的并行效率与自回归模型的上下文连贯性相结合?

通过离散扩散模型并行提出token候选,并由自回归模型进行上下文一致性验证,在保证生成质量的同时显著降低推理延迟与KV Cache开销,同时跳过传统噪声预测步骤,使proposal与AR的目标空间对齐。

扩散模型生成的token作为预选序列(proposal)

2)在长上下文推理中,Transformer的自注意力具有 O(n2) 复杂度,并且KV Cache随上下文长度线性膨胀。为了降低推理显存占用并缓解长序列“遗忘”问题,是否存在更高效的方法在减少KV储存成本的同时维持甚至提升模型的长程依赖建模能力?

为了解决长上下文推理中Transformer自注意力的 O(n2) 计算复杂度以及KV Cache随上下文长度线性膨胀的问题,可以探索构建一种选择性、分层化、预测驱动的KV Cache管理框架。该框架在 Transformer中引入可学习的多层参数模块类似MoE,用于在编码阶段提取Prompt中的关键信息,从而减少冗余token进入KV Cache并降低提示词带来的缓存开销。同时,通过门控机制控制信息保留,采用多时间尺度的层级记忆区分短期与长期信息,并结合预测需求对KV Cache进行动态增删。在显著降低推理显存占用与计算成本的同时,该方法有望维持甚至增强模型对长距离文本依赖的建模能力。

3)自回归生成是一种严格的逐步依赖机制,一旦模型在某个步骤生成错误token,就会把错误传递到后续所有步骤,当前模型只能“继续生成”,无法“回溯修正”。是否存在结构性改进,使模型能够在发现局部错误后,只修改必要的前文token并重新继续生成,而无需从头开始?

是否可以通过改进双向注意力机制或构建双向+自回归的混合架构,让推理过程具备可回溯、可纠错能力......

参考文献