我们已经优化了 AIGC 的方方面面，但还有一个根本性的物理瓶颈悬在头顶：自回归（Autoregressive）。
无论你的 NPU 有多快，LLM 生成 Token 必须是一个接一个的：生成了"今"，才能生成"天"，生成了"天"，才能生成"气"。这种串行依赖性，使得昂贵的 GPU/NPU 在生成每一个 Token 之间不得不进行大量的内存读写，而计算单元往往处于"吃不饱"的状态。
投机采样（Speculative Decoding） 是打破这一诅咒的"黑魔法"。它利用一个小模型（Draft Model）快速"猜"出未来几个 Token，然后让大模型（Target Model）并行验证。
在 AIGC 的推理赛道上，如果你还在一个字一个字地生成，那你已经输在了起跑线。
DeepSeek、Llama 等大模型的推理延迟主要来自于显存带宽（Memory Bandwidth）。读取 70B 模型的权重只需要毫秒级，但如果只是为了生成一个 token，这巨大的带宽开销就显得极不划算。
Speculative Decoding 的核心思想是：既然读一次权重这么贵，不如读一次多算几个 Token。
- Draft ：让一个小模型（比如 1B）快速猜出未来 5 个词：
["今天", "天气", "非常", "不错", "。"] - Verify：让大模型（70B）一次性并行计算这 5 个位置的概率。
- Accept/Reject：对比大模型的判断，接受前 N 个正确的词。
这个过程中，最关键的一步是验证（Verification） 。它不仅仅是简单的 if a == b，而是涉及复杂的概率分布比对（Rejection Sampling）。在 ops-nn 仓库中，CANN 工程师利用 Vector 单元的并行掩码（Mask）能力，将这一逻辑硬化为算子。
极速推理核心
- CANN 组织主页： https://atomgit.com/cann
- ops-nn 源码仓库： https://atomgit.com/cann/ops-nn
一、 验证的难题：并行中的分支
在 CPU 上写验证逻辑很简单：
python
for i in range(k):
if target_token[i] == draft_token[i]:
accept_count += 1
else:
break但在 NPU 上，这是一个典型的 Serial Dependency（串行依赖） 问题：第 2 个词是否被接受，取决于第 1 个词是否被接受。如果我们在 Kernel 中使用标量循环，性能会非常差。
ops-nn 的解法是：全并行计算 + 前缀扫描（Prefix Scan） 。先不管依懒性，并行判断所有位置是否符合要求，生成一个 Mask 向量（如 [1, 1, 0, 1, 1]），然后通过一条指令找到第一个 0 的位置，截断即可。
二、 代码实战：构建 Speculation Verify Kernel
我们来编写一个基于 Greedy Strategy（贪婪策略，即直接比对 Token ID）的验证算子核心。这常用于代码生成等确定性场景。
Ascend C 核心代码逻辑
cpp
#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t MAX_DRAFT_LEN = 16; // 假设最大猜测长度
class KernelSpecVerify {
public:
__aicore__ inline KernelSpecVerify() {}
__aicore__ inline void Init(GM_ADDR draft_tokens, GM_ADDR target_tokens, GM_ADDR out_len, int32_t num_seqs) {
// draft_tokens: 小模型猜的 [Batch, Max_Draft]
// target_tokens: 大模型算的 Top1 [Batch, Max_Draft]
draftGm.SetGlobalBuffer((__gm__ int32_t *)draft_tokens);
targetGm.SetGlobalBuffer((__gm__ int32_t *)target_tokens);
outLenGm.SetGlobalBuffer((__gm__ int32_t *)out_len); // 输出每句话接受了多少个
m_num_seqs = num_seqs;
pipe.InitBuffer(inDraft, 1, MAX_DRAFT_LEN * sizeof(int32_t));
pipe.InitBuffer(inTarget, 1, MAX_DRAFT_LEN * sizeof(int32_t));
pipe.InitBuffer(outQueue, 1, sizeof(int32_t)); // 标量输出
}
__aicore__ inline void Process() {
// 循环处理 Batch 中的每一个 Sequence
for (int32_t i = 0; i < m_num_seqs; i++) {
VerifySeq(i);
}
}
private:
__aicore__ inline void VerifySeq(int32_t batch_idx) {
LocalTensor<int32_t> drafts = inDraft.AllocTensor<int32_t>();
LocalTensor<int32_t> targets = inTarget.AllocTensor<int32_t>();
// 1. 搬运数据
// 每次搬运一行 (Max_Draft_Len)
DataCopy(drafts, draftGm[batch_idx * MAX_DRAFT_LEN], MAX_DRAFT_LEN);
DataCopy(targets, targetGm[batch_idx * MAX_DRAFT_LEN], MAX_DRAFT_LEN);
// 2. 并行比对 (Vector Compare)
// 申请一个 Mask Tensor，用于存储比对结果
// cmp_res[i] = (drafts[i] == targets[i]) ? 1 : 0
LocalTensor<uint8_t> cmp_res = inDraft.AllocTensor<uint8_t>(); // 复用或新申请
// Ascend C Compare 指令
// Compare(cmp_res, drafts, targets, CMP_EQ, MAX_DRAFT_LEN);
// --- 核心逻辑: 寻找第一个不匹配的位置 (Find First Zero) ---
// 在 SIMD 编程中，这通常通过 Convert To BitMask 然后 Count Leading Ones 来实现
// 或者使用 Ascend C 的特定归约指令
// 假设我们得到了一个 cmp_res = [1, 1, 1, 0, 1] (注意第4个0之后即使是1也是无效的)
// 方法 A: 累积与运算 (Cumulative And) - 模拟前缀扫描
// [1, 1, 1, 0, 1] -> [1, 1, 1, 0, 0]
// 但这通常需要 Log(N) 步
// 方法 B: 移动到标量单元处理 (Scalar Unit)
// 对于 Draft 长度较短 (如 5-10)，直接用 Scalar 循环可能更快且不阻塞 Vector
// 但 ops-nn 追求极致，通常会用 Vector Reduce
// 这里演示一个向量化的思路：
// 1. Cast Mask to FP16
// 2. 找到第一个 0 的索引
// 简化实现：将数据倒出到 UB 的 Scalar 区域进行快速扫描
// 因为 MAX_DRAFT_LEN 很小，Scalar 循环开销极低
int32_t accepted_len = 0;
for (int32_t k = 0; k < MAX_DRAFT_LEN; k++) {
if (drafts.GetValue(k) == targets.GetValue(k)) {
accepted_len++;
} else {
break; // 遇到第一个错误，立即停止
}
}
// 3. 输出结果
LocalTensor<int32_t> outT = outQueue.AllocTensor<int32_t>();
outT.SetValue(0, accepted_len);
DataCopy(outLenGm[batch_idx], outT, 1);
inDraft.FreeTensor(drafts);
inTarget.FreeTensor(targets);
// ... (释放其他)
outQueue.FreeTensor(outT);
}
private:
TPipe pipe;
TQue<QuePosition::VECIN, 1> inDraft, inTarget;
TQue<QuePosition::VECOUT, 1> outQueue;
GlobalTensor<int32_t> draftGm, targetGm, outLenGm;
int32_t m_num_seqs;
};3. 代码进阶：从 Greedy 到 Rejection Sampling
上面的代码仅演示了最简单的相等比对。在 Rejection Sampling（拒绝采样） 场景中，逻辑会复杂得多：
- 概率回读：输入不再是 Token ID，而是 Draft Model 和 Target Model 的概率分布 和 。
- 随机判定：如果 ，接受；否则，以概率 拒绝。
- 随机数生成：这需要在 Kernel 内生成一组随机数 。
在 ops-nn 仓库的完整实现中，你会看到大量使用 Vector Div（向量除法） 和 Vector Compare（向量比较） 来并行处理这组概率公式，最后才收敛到 Scalar 进行截断。这种"大量并行计算 + 少量标量控制"的模式，是异构编程的精髓。
三、 Tree Attention：更高级的投机
普通的 Speculative Decoding 是线性的（猜一个序列）。但更激进的做法是 Tree Speculation （猜一棵树）。
Draft Model 可以输出多个分支：
- 分支 A: "今天" -> "天气"
- 分支 B: "今天" -> "是"
Target Model 需要验证这棵树上的所有节点。这意味着：
- Attention Mask 变了：不再是简单的三角矩阵，而是基于树结构的 Mask。
- Gather Index 变了：KV Cache 的读取不再连续，而是沿着树的分支跳跃。
在 AtomGit 的 ops-nn 仓库中，针对 Tree Attention 的实现，展示了如何通过构建特殊的 Topology Mask（拓扑掩码） 喂给 FlashAttention 算子，使得一次 Attention 计算就能覆盖树上的所有假设路径。这是 NPU 算力利用率的巅峰展示。
四、 为什么说 Speculative Decoding 是未来？
- 打破摩尔定律的限制 ：
单卡的显存带宽增长速度已经跟不上模型参数的增长速度。Speculative Decoding 是用"多余的 FLOPs（算力）"来换取"宝贵的 Bandwidth（带宽）"。因为 NPU 的算力往往是过剩的，而带宽是瓶颈。 - 端侧大模型的救星 ：
在手机或 PC 上跑大模型，内存带宽很小。通过一个小模型（Draft）跑在 CPU 或 NPU 小核上，大模型跑在 NPU 大核上，可以显著提升用户体验。ops-nn的轻量化算子实现对此至关重要。
五、 结语：算力的赌局
投机采样本质上是一场赌局：赌小模型猜得准。赌赢了，推理速度翻倍；赌输了，浪费点电费（回退）。
而 CANN 的 ops-nn 仓库，通过极致优化的验证算子，将"赌输"的代价（Verification Overhead）降到了最低。它确保了即使小模型猜得不准，大模型的验证过程也快如闪电，几乎不占用额外时间。
如果你想让你的 AIGC 应用快人一步，不要仅仅关注模型本身，去 AtomGit 上看看 ops-nn 是如何处理这些精妙的"投机"逻辑的。
加入极速阵营：
- CANN 开发者社区： https://atomgit.com/cann
- 探索投机采样源码： https://atomgit.com/cann/ops-nn