llama.cpp 高级特性：推测解码与语法约束生成

1. 背景与目标

大语言模型（LLM）的自回归生成是典型的“内存墙”和工作负载：每个 token 的生成都需要完整读取所有参数并计算注意力，顺序依赖使得 GPU 利用率极低。推测解码（Speculative Decoding）通过用一个轻量级的“草稿模型”（draft model）快速预测未来多个 token，再让目标模型并行验证整段草稿，从而实现加速且不改变输出分布。这一思想来自 2023 年 Leviathan 等人的 ICML 论文《Fast Inference from Transformers via Speculative Decoding》，llama.cpp 中将其落地为多种策略（n-gram、EAGLE3、Medusa 等），并将验证流程与采样管道深度整合。

同时，在结构化生成（如 JSON 输出、函数调用）场景中，我们希望模型输出的每个 token 都符合给定的形式语法。GBNF（GGML BNF）允许用户定义类似 BNF 的规则，引擎在采样时动态施加约束，保证输出完全合法。这涉及语法解析、自动机构建、logit 掩码以及拒绝采样机制。

本文将深入剖析 llama.cpp 的推测解码验证流程和基于 GBNF 的约束采样实现，聚焦于 draft-target 的交互细节以及语法采样器如何在 logit 层面工作。所有分析均基于源码（common/speculative.cpp、common/sampling.cpp），并附有精确的文件路径和行号。

2. 核心设计与实现

2.1 推测解码的整体控制流

推测解码的核心是“草稿-验证”循环，位于 common/speculative.cpp 中。每次迭代执行：

1. 草稿模型（draft）自回归地生成一个候选 token 序列。
2. 目标模型（target）一次性解码整条草稿序列，同时计算出每个位置的目标分布 logits。
3. 验证过程逐个位置比较草稿 token 在目标分布下的概率，并根据一个拒绝采样规则决定接受或重新采样。

// common/speculative.cpp:700-730 (草稿生成循环，仅保留置信度高的 token)
common_batch_clear(batch);

for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
    if (!drafting[seq_id]) {
        continue;
    }

    auto * smpl = smpls[seq_id].get();
    common_sampler_sample(smpl, ctx_dft, i_batch, true);
    const float * prenorm = llama_get_embeddings_nextn_ith(ctx_dft, i_batch);
    ++i_batch;

    const auto * cur_p = common_sampler_get_candidates(smpl, true);
    const llama_token id = cur_p->data[0].id;

    // 仅当草稿 token 概率高于阈值 p_min 时才继续，否则终止本序列的草稿
    if (cur_p->data[0].p < params.p_min) {
        drafting[seq_id] = false;
        n_drafting--;
        continue;
    }
    // ... 将 id 加入草稿序列
}

这段代码揭示了一个重要细节：草稿模型不是无条件生成 max_draft_len 个 token，而是只要任意一步的 top-1 概率低于 params.p_min 就会立刻停止该序列的草稿。这样做有两个好处：第一，低置信度的草稿很可能在验证阶段被拒绝，浪费目标模型的计算；第二，提前终止避免了草稿模型在不确定区域“胡编乱造”，减少了无用的显存带宽占用。

2.2 拒绝采样验证

2.3 EAGLE3 草稿生成的特征提取

EAGLE3（来自论文《EAGLE-3: Scaling Up Inference Acceleration of Large Language Models》）的草稿模型并非独立的小网络，而是共享目标模型的隐藏状态。其核心思想是：目标模型某些层的隐藏状态已经包含了丰富的上下文信息，可以输入一个小型解码器直接预测下一个 token，而无需完整计算注意力。

在 llama.cpp 中，EAGLE3 的草稿生成分为特征提取（编码器）和逐 token 生成（解码器）两步。common/speculative.cpp 中 eagle3_draft() 函数实现了这两步：

1. 从目标模型的指定层提取 输入嵌入，拼接成特征矩阵 features_buf，维度为 [n_tokens, n_layers * n_embd_tgt]。
2. 使用草稿模型的编码器将特征投影到草稿维度 [n_embd_dec]，得到每组 token 的表示 g_embd。
3. 在草稿解码器中，位置 $P+1$ 的 token 与位置 $P$ 的 g_embd 配对，共同解码出位置 $P+1$ 的 token。

但交互式生成中，位置 $P+1$ 的 token 可能来自之前的批次，尚未确定。llama.cpp 使用了一种 延迟边界转发（deferred boundary forwarding）技术：保存上一批次最后一个 g_embd 以及其位置，在下一批次开头完成配对。相关代码：

// common/speculative.cpp:610-620 (保存 pending_g_last 和位置)
const int32_t n_rows = end - beg + 1;
verify_pos_first[seq_id] = batch_in.pos[beg];
pending_pos_last[seq_id] = batch_in.pos[end];
verify_g_rows[seq_id]    = n_rows;
verify_g[seq_id].resize((size_t) n_rows * n_embd_dec, 0.0f);
std::memcpy(verify_g[seq_id].data(), g_embd + (size_t) beg * n_embd_dec, row_bytes * n_rows);
std::memcpy(pending_g_last[seq_id].data(), g_embd + (size_t) end * n_embd_dec, row_bytes);

这一操作避免了对整个序列重新运行编码器，仅仅是复制一个隐藏向量和记录位置。当验证阶段确定了实际接受的 token 数量后，通过之前保存的 verify_g 快照可以恢复正确的 pending_g_last，确保了跨批次调度的一致性，且没有引入额外的草稿模型前向传播。

2.4 GBNF 约束采样的集成

GBNF（可以理解为一种简化的巴科斯-诺尔形式）允许用户以类似正则表达式的规则定义生成语法。llama.cpp 的语法采样器接受一个 GBNF 规则（或 JSON schema 转换而来的规则），内部编译成确定性有限自动机（DFA）。在每个生成步骤，采样器接收候选 token 数组，并将那些无法在当前语法状态中出现的 token 的 logit 设置为 -INFINITY（即概率 0），从而“硬屏蔽”非法 token。

然而，采样过程并不直接“只看语法”，而是与温度、top-k、top-p 等采样器在管道中协作。在 common_sampler_sample 中，语法采样器的调用时机存在两种模式：

• grammar-first 模式：先应用语法屏蔽，再执行其他采样器。这是默认行为，确保采样出的 token 一定合法。
• 后验拒绝模式：先进行普通采样，然后检查采样结果是否合法；若不合法，重新应用语法屏蔽后重采样。

后验模式的实现代码如下：

// common/sampling.cpp:400-418 (拒绝采样路径)
if (grammar_first || !grammar_should_apply(gsmpl)) {
    return id;
}

// 检查采样的 token 是否符合语法
{
    llama_token_data       single_token_data       = { id, 1.0f, 0.0f };
    llama_token_data_array single_token_data_array = { &single_token_data, 1, -1, false };
    llama_sampler_apply(grmr, &single_token_data_array);

    const bool is_valid = single_token_data_array.data[0].logit != -INFINITY;
    if (is_valid) {
        return id;
    }
}

// 重新采样，但这次先应用语法采样器再应用采样链
gsmpl->set_logits(ctx, idx);
llama_sampler_apply(rbudget,  &cur_p);
if (grammar_should_apply(gsmpl)) {
    llama_sampler_apply(grmr,  &cur_p);
}

这里将采出的单个 token 重新包装成一个 llama_token_data_array 送入语法采样器，通过检查 logit 是否被置为 -INFINITY 来判断合法性。若非法，则丢弃该 token，并重新设置 logits（即恢复原始未修改的 logits），然后依次应用推理预算采样器、语法采样器和整个采样链。这种“先试一个，不行再走安全路径”的策略在高概率合法时避免了全量屏蔽的开销（因为全量屏蔽需要为词表中的每一个 token 测试语法合法性，复杂度 $O(V)$）。但这种策略仅在生成 token 大概率符合语法时有效；如果语法严格导致大量拒绝，反而会增加开销。工程上的权衡非常清晰。

值得注意的是，语法采样器还可以与推理预算（reasoning budget） 协作：在模型输出“思考内容”时，可以暂时关闭输出格式约束（如 JSON 语法），让模型自由生成，待推理阶段结束后再重新启用语法，强制最终输出合法。

// common/sampling.cpp:323-333 (语法懒惰应用条件)
static bool grammar_should_apply(struct common_sampler * gsmpl) {
    if (!gsmpl->grmr) return false;
    if (!gsmpl->rbudget) return true;
    if (gsmpl->params.grammar_lazy) {
        const auto state = common_reasoning_budget_get_state(gsmpl->rbudget);
        return state == REASONING_BUDGET_IDLE || state == REASONING_BUDGET_DONE;
    }
    return true;
}

3. 优化技术细节

3.1 草稿生成精度与性能的平衡

草稿模型必须足够快，否则其带来的并行验证收益会被自身开销抵消。EAGLE3 的草稿模型本身很小（通常只包含若干层），并且在生成过程中使用了 分块编码器 来适应 ubatch 限制，避免一次调度过大的批次导致失败：

// common/speculative.cpp:590-600 (分块编码)
for (int32_t i = 0; i < n_tokens; i += n_ubatch_dft) {
    const int32_t n_chunk = std::min(n_ubatch_dft, n_tokens - i);
    llama_batch enc_batch = {
        /*.n_tokens =*/ n_chunk,
        /*.embd     =*/ features_buf.data() + (size_t) i * n_embd_enc,
        // 其余字段为 nullptr
    };
    llama_encode(ctx_dft, enc_batch);
    // ...
}

将较长的提示分成多个 chunk 依次编码，最后拼接结果。牺牲了部分 GPU 利用率，但保证了在任意长度输入下都能稳定运行。

另外，在草稿 token 采样后，立刻通过 llama_get_embeddings_nextn_ith 获取该位置的预归一化隐藏状态（而不是重新计算），作为下一位置解码的 g_embd 输入，避免了额外的前向计算。这种“复用中间结果”的 trick 在序列生成中极大降低了草稿模型的成本。

3.2 后端采样的预计算候选

对于那些不支持语法约束的硬件后端（如某些 GPU），采样模块支持使用后端提前采样候选列表 (llama_get_sampled_probs_ith)，直接提供排序好的概率值和 token ID，CPU 端无需遍历整个词表。这在无语法约束时能将采样开销从 $O(V \log V)$ 降至 $O(n_\text{candidates})$，对推测解码这种高频采样场景尤为重要。

3.3 语法自动机的动态编译

虽然未在提供的代码片段中详细展开，但 GBNF 规则在首次使用时会被编译成一个 DFA，并存储在 common_sampler 的 grmr 字段中。该 DFA 维护当前状态，以及对每个 token 的转移逻辑。对于每个 token，只需查表即可知道是否合法，然后直接设置 logit。这种做法的空间开销与词表大小和语法状态数成正比，但对于多数应用场景（如 JSON 输出）而言是可接受的。

4. 性能与权衡

推测解码的加速比理论上由草稿模型的接受率决定。论文指出，当草稿模型与目标模型相近、接受率较高时，可接近 2-3 倍加速。llama.cpp 通过设置 p_min 参数（如 0.8）仅保留高置信草稿，进一步提高了接受率，同时减少了无效验证。但引入 EAGLE3 等基于特征的草稿模型后会增加额外的特征提取开销，且对延迟敏感的流式应用，验证步骤的批次大小受草稿长度限制，不能充分利用 GPU 并行性。

语法约束采样带来了确定性的输出格式，但其硬屏蔽机制会显著减小有效候选空间，极端情况下可能使得采样退化到贪心搜索（因为只剩一个合法 token），降低多样性。同时，后验拒绝模式在语法复杂时可能导致频繁重采样，增加 CPU 负担。

5. 相关论文

• Fast Inference from Transformers via Speculative Decoding (Leviathan, Kalman, Matias, 2023, ICML)[2302.01318]：奠定了推测解码的理论基础，llama.cpp 采用了其拒绝采样验证算法。
• EAGLE-3: Scaling Up Inference Acceleration of Large Language Models (Zhang et al., 2024)：提出通过目标模型隐藏状态预测后续 token，llama.cpp 中的 EAGLE3 草稿策略即基于此。
• FlashAttention (Dao et al., 2022) 等不直接相关，但在注意力加速层面推动了整个推理系统。

6. 总结

llama.cpp 将推测解码和约束采样深度融入自回归管道，展现出工程上的精致权衡。从草稿的置信度剪枝到延迟边界转发，从语法拒绝采样到与推理预算的协同，每个细节都服务于实际部署效率。理解这些实现有助于进一步优化本地推理系统，无论是为特定硬件定制草稿策略，还是设计更为灵活的结构化生成接口。