0.8MB 跑通 Qwen|第 9-2 篇:推理引擎的 flash_attn_single_q_q8_neon——单 q 单遍扫全 KV

简介: 本系列《0.8MB跑通Qwen》聚焦ARM零依赖纯C推理引擎,精读`flash_attn_single_q_q8_neon`内核:通过Q单次量化、q8 KV单遍扫描、在线softmax与尾部归一,显著降低带宽与内存开销,实现在RK3588等端侧设备高效运行Qwen3-VL系列模型。(239字)

系列:《0.8MB 跑通 Qwen:从零实现 ARM 零依赖纯 C 推理引擎》(30 天 × 90 篇) | 适配模型:Qwen3-VL-8B-Instruct(千问3_VL_8B_Instruct)· Qwen3-VL-2B-Instruct · Qwen3-30B-A3B | 测试设备:RK3588(4×Cortex-A76 + 4×Cortex-A55,aarch64)

系列总纲:《0.8MB 跑通 Qwen》30 天 90 篇 · 总纲(阿里云社区)

上一篇:9-1《KV 也量化:q8 KV 把缓存与 decode 带宽压到一半》 | 下一篇:9-3《q8 KV 精度对照实验》

源码精读篇:本文为源码/方法论精读,无独立实测;文中数字均引述仓库 docs 的板端实测记录

一句话导读:推理引擎的 flash_attn_single_q_q8_neon:单 query 对全量 KV 恒定单遍的 NEON 注意力内核,逐段导读 Q 量化一次、逐 token 打分与在线 rescale、尾部归一,讲清单遍比两遍省在哪三处。

关键词:手搓 Qwen 推理引擎、千问大模型推理、Qwen3-VL、零依赖纯 C、q8 KV、flash attention、NEON、单遍扫描、在线 softmax

导语:decode 每个新词只有一个 query,却要扫遍全部历史 KV。手搓 Qwen 推理引擎用"恒定单遍"的 NEON 内核 flash_attn_single_q_q8_neon 把在线 softmax 与 q8 KV 合体:量化 Q 一次、扫 KV 一遍、边扫边 rescale、读完即出结果。这篇逐段导读它省在哪。

9-1 把 KV 压成了 q8,可 decode 要"每个 token 扫一遍全部历史 KV"——怎么扫才不浪费?答案是一个"恒定单遍"的 NEON 内核:flash_attn_single_q_q8_neon。它把 8-2 的在线 softmax 与 q8 KV 合体:量化 Q 一次、扫 KV 一遍、边扫边 rescale、读完即出结果。

1. 知识点:decode 的注意力为什么能"单遍"

decode 时 query 只有一个(当前生成词),要对它和全部历史 KV 算注意力。教科书写法(8-1)要三遍:先算完全部分数、再 softmax、再加权 V——这意味着把全部分数存下来(或扫两遍 KV)。

单遍版的数学基础就是 8-2 的在线 softmax:维护 M(已见分数最大值)与 S(已见 exp 和),每来一个 token 的分数,当场决定:

  • 分数 ≤ M:w = exp(score - M),直接累进 V 加权和;
  • 分数 > M:先把已累加的和整体乘 exp(M_old - M_new)(rescaling),再继续。

于是 KV 只扫一遍、分数不必物化、输出是"边扫边累积"的——这就是"恒定单遍":无论上下文多长,KV 只被读一次。

2. 对应代码:逐段导读(第 6594–6670 行)

函数签名先交代了 q8 KV 的"三件套"(第 6594–6602 行):

static void flash_attn_single_q_q8_neon(
    float *restrict attn_out,           /* [hd] 归一化输出 */
    const float *restrict q,            /* [hd] query 激活(fp32) */
    int8_t *const *k_cache_q8,          /* 每层 INT8 K 块数组 */
    int8_t *const *v_cache_q8,
    const float *restrict kscale,       /* k_scale[l] + kh:per-token per-head */
    const float *restrict vscale,
    int seq_len, int kv_dim, int kv_bs, int nkv,
    int kh_off, int hd, float scale)

注意 kscale 注释"per-token per-head"——正是 9-1 那个每 (token, 头) 的 scale。

段 1:Q 量化一次(第 6609–6625 行)

/* 1. Q → INT8(每 32 元素 1 max-abs scale,Q8_0 同构) */
for (int b = 0; b < nblk; b++) {
   
    float qm = 0.0f;
    for (int i = 0; i < 32; i++) qm = fmaxf(qm, fabsf(qb[i]));
    if (qm < 1e-6f) qm = 1.0f;
    qsc[b] = qm * KVQ_INV;            /* KVQ_INV = 1/127 */
    float iq = 127.0f / qm;
    for (int i = 0; i < 32; i++) qi[b*32+i] = (int8_t)clamp(qb[i] * iq, -128, 127);
}

q8 × q8 的点积要求两边都是 int8——K 已经是 int8,所以把 Q 也量一次,两者就能直接走 vdot(6-1 的 i8x16_dot_s32)。Q 只有 128 维,量化成本微不足道,且全函数只量这一次。

段 2:在线 softmax 运行态(第 6627–6666 行,8-2 已拆)

float M = -1e9f, S = 0.0f;
/* 单遍 KV 扫描(逐 token,online softmax 分块语义) */
for (int t = 0; t < seq_len; t++) {
   
    const int8_t *kt = k_cache_q8[t / kv_bs] + (size_t)(t % kv_bs) * kv_dim + kh_off;
    float ks = kscale[(size_t)t * nkv] * KVQ_INV;
    /* Q·K^T:每 32 块一个 int32 累加,块 scale 分别乘后求和 */
    float s = 0.0f;
    for (int b = 0; b < nblk; b++) {
    ...acc = i8x16_dot_s32(...); s += dot32 * qsc[b]; }
    s = s * ks * scale;
    /* online softmax:新 max → rescale 已累积 VKQ */
    if (s > M) {
    S *= expf(M - s); M = s; /* acc 同步缩放 */ ... }
    else {
    vsf = expf(s - M); }
    S += vsf;
    /* VKQ += INT8 V × (vsf × vs),softmax 权重保持 F32 */
    ...
}

三个要点:

  1. 块 scale 的归位:K 的 int8 值与 Q 的 int8 值点积得到"无 scale"的 int32 部分和,真正的浮点值要乘 qsc[b](Q 每 32 块 scale)× ks(K 每 token 头 scale)——整数点积只算骨架,浮点标定在累加时补齐(Day 6–7 的 dotprod 纪律在 KV 上的再现);
  2. kv_bs 分块:t / kv_bs 定位块、t % kv_bs 定位块内 token(8-3 的布局);
  3. 在线 rescale:遇新 max 只缩放已累计量,KV 扫完即得未归一结果。

段 3:归一输出(函数尾部):attn_out[i] /= S(连同最终 scale 处理),一次除法收尾。

3. 改动后果:把在线 rescale 去掉,改回"先存全部分数"

假设改成朴素两遍:先把 seq_len 个分数都算好存进数组(float scores[8192] 之类),再找 max、再统一 exp——功能等价,代价三样:

  1. 多一遍 KV 读取:打分一遍、加权 V 又一遍 → KV 扫描 ×2(q8 KV 的带宽红利白省一半);
  2. 多一份分数数组:8K 上下文要 32KB 栈/堆暂存,还破坏 cache 局部性;
  3. 失去"流式"能力:在线版可以在 KV 逐块到达(跨进程磁盘 KV,Day 12)时边读边算;两遍版必须整段就绪。

这正是技术文档里那句"q8-KV decode 恒定单遍"的含义(技术文档.md 第 293 行)——单遍不是风格,是带宽与流式能力的共同要求。

诚实标注:文档还注明"x86 实验开关 VLLM_ATTN_ONLINE 未进入本 aarch64 树"——即 ARM 发布树固定走本内核的在线路径,没有可切换的"非在线"后端。本节"去掉 rescale"是概念对照,板端无对应开关。

4. 学员调试任务

  • A 档(板端动手):在 vllm_safetensors.c 第 6594 行函数内打断点(用 1-3 的 Debug 构建姿势),跑一次生成,观察:进入函数时 seq_len 是多少、单次调用内 for (t) 是否恰好扫 seq_len 遍、M 是否单调不降。
  • B 档(纯读源码):逐行标出"Q 量化 / 打分 / rescale / 归一"四段的起止行号,回答:为什么 Q 只量化一次而 KV 的 scale 每 token 都要乘?

预期输出:你能画出"单 q × 全 KV 单遍"的数据流(Q→int8、逐 token 打分+在线 softmax、尾部归一),并说清它比两遍版省在哪三处。

收尾

  • 本篇源码点名:vllm_safetensors.c(flash_attn_single_q_q8_neon 第 6594 行起)、技术文档.md 第 293 行。
  • 开源仓库:Kestrel-LLM (Gitee)(AGPL-3.0-or-later 或商业许可,二选一)
  • 下篇预告:单遍扫全 KV 很优雅——可它"算得准吗"?下一篇 9-3 做 q8 KV 与 fp32 KV 的精度对照,看量化在注意力输出上到底差多少。
相关文章
|
1月前
|
调度 C++
0.8MB 跑通 Qwen|第 13-2 篇:草稿 + 验证——推理引擎里 spec decode 的实现
本系列《0.8MB跑通Qwen》聚焦ARM端零依赖纯C推理引擎,实测RK3588上运行Qwen3-VL多模态模型;详解speculative decode实现——基于n-gram草稿、批量验证与无损回放,确保输出位级一致,兼顾 correctness 与工程可控性。(239字)
123 1
|
1月前
|
缓存 C语言 C++
0.8MB 跑通 Qwen|第 9-3 篇:推理引擎的 q8 KV 精度对照——量化进注意力,输出差多少
本系列《0.8MB跑通Qwen》手搓零依赖纯C推理引擎,适配Qwen3-VL多尺寸模型,在RK3588上实测q8 KV量化:输出误差仅≈0.01,端到端PPL损失<1%,带宽降为52%,精度与效率达成教科书级平衡。(239字)
160 1
|
1月前
|
缓存
0.8MB 跑通 Qwen|第 8-3 篇:推理引擎的 KV 缓存结构——按 token 还是按头存
本系列《0.8MB跑通Qwen》聚焦ARM端零依赖纯C推理引擎,适配Qwen3-VL多模态模型,在RK3588上实测。本文精析KV缓存的token-major布局设计原理,揭示其如何兼顾prefill与decode访存效率,直击带宽瓶颈。(239字)
109 1
|
1月前
|
缓存 人工智能 索引
0.8MB 跑通 Qwen|第 12-1 篇:LLM 推理的跨进程恢复——服务重启了,会话不能断
本系列《0.8MB跑通Qwen》聚焦ARM端零依赖纯C推理引擎,实现在RK3588上部署Qwen3-VL等多尺寸模型。本文详解`--disk-kv`机制:将KV Cache落盘持久化,支持进程重启后按前缀精准恢复,实现“会话不断连”,2K上下文prefill加速达23.6×,真正打通边缘AI服务可用性最后一环。(239字)
131 1
|
1月前
|
NoSQL 调度 C++
0.8MB 跑通 Qwen|第 14-3 篇:推理引擎的批处理正确性边界——margin 0.71 vs 0.032
本篇精读Qwen3-VL系列推理引擎的“位级一致性”本质:批量与串行输出大多逐字节一致,但因浮点累加顺序差异,在near-tie(近平局)场景下存在翻盘风险——margin(如0.71 vs 0.032)决定是否一致。这是工程结果,非数学承诺。
104 0
|
1月前
|
缓存 NoSQL 区块链
0.8MB 跑通 Qwen|第 15-2 篇:推理引擎的 prefill 与 decode——两条路径为何分开
本篇详解Qwen推理引擎中prefill与decode双路径设计:prefill一次性处理整段prompt(如18 token),批量写入KV缓存;decode逐token循环生成,追加KV。通过RK3588真机gdb断点实证,明确二者独立入口、状态流转与性能动因,手搓零依赖纯C引擎的核心逻辑。(239字)
170 0
|
1月前
|
缓存 算法 API
0.8MB 跑通 Qwen|第 11-3 篇:推理引擎里 KV 缓存的一生——内存、容量与淘汰
本系列《0.8MB跑通Qwen》聚焦ARM端零依赖纯C推理引擎开发,实测RK3588平台流畅运行Qwen3-VL-2B/8B等多版本。本文详解KV缓存“出生—续用/覆盖—死亡”三段式生命周期,揭示预分配、逻辑覆盖、进程级销毁本质,破除“命中即加速、淘汰即释放”认知误区。(239字)
139 0
|
1月前
|
缓存 安全 API
0.8MB 跑通 Qwen|第 11-2 篇:大模型推理的前缀缓存键——怎么知道"两轮一样"?
本系列《0.8MB跑通Qwen》聚焦ARM端零依赖纯C推理引擎,实测RK3588上高效运行Qwen3-VL多模态模型。本文详解前缀缓存核心机制:以token序列最长公共前缀(LCP)为键,实现KV复用;并揭示“完全相同反不复用”的安全设计——确保prefill刷新logits,杜绝空响应。
144 0
|
1月前
|
调度
0.8MB 跑通 Qwen|第 13-3 篇:推理引擎的回退与收益边界——哪些场景 spec 才划算
本文精析Speculative Decoding在ARM端的收益边界:基于RK3588实测,揭示“62%命中率仍变慢”的根源——单步Decode仅44ms时,验证开销与KV回滚反致负增益;明确开启条件:长上下文、高重复性、大模型、贪婪采样。零依赖纯C,适配Qwen3-VL系列。
112 0
|
1月前
|
缓存 自然语言处理 安全
0.8MB 跑通 Qwen|第 11-1 篇:大模型推理中多轮对话的浪费——上一轮 prefill 算完就被扔了
本系列《0.8MB跑通Qwen》聚焦ARM端零依赖纯C推理引擎,实测RK3588平台高效运行Qwen3-VL多模态模型;核心突破是默认启用前缀KV复用,使多轮对话prefill耗时骤降27倍,显著消除历史重复计算,大幅提升端侧长会话效率。(239字)
134 0