ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

Day 9·2 flash_attn_single_q_q8_neon——单 q 单遍扫全 KV

Day 9·2 flash_attn_single_q_q8_neon——单 q 单遍扫全 KV 源码精读篇本文为源码/方法论精读无独立实测文中数字均引述仓库 docs 的板端实测记录一句话导读flash_attn_single_q_q8_neon单 query 对全量 KV 恒定单遍的 NEON 注意力内核逐段导读 Q 量化一次、逐 token 打分与在线 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 的分数当场决定分数 ≤ Mw 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] khper-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。段 1Q 量化一次第 6609–6625 行/* 1. Q → INT8每 32 元素 1 max-abs scaleQ8_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*32i] (int8_t)clamp(qb[i] * iq, -128, 127); }q8 × q8 的点积要求两边都是 int8——K 已经是 int8所以把 Q 也量一次两者就能直接走 vdot6-1 的i8x16_dot_s32。Q 只有 128 维量化成本微不足道且全函数只量这一次。段 2在线 softmax 运行态第 6627–6666 行8-2 已拆float M -1e9f, S 0.0f; /* 单遍 KV 扫描逐 tokenonline 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 */ ... }三个要点块 scale 的归位K 的 int8 值与 Q 的 int8 值点积得到“无 scale”的 int32 部分和真正的浮点值要乘qsc[b]Q 每 32 块 scale×ksK 每 token 头 scale——整数点积只算骨架浮点标定在累加时补齐Day 6–7 的 dotprod 纪律在 KV 上的再现kv_bs分块t / kv_bs定位块、t % kv_bs定位块内 token8-3 的布局在线 rescale遇新 max 只缩放已累计量KV 扫完即得未归一结果。段 3归一输出函数尾部attn_out[i] / S连同最终 scale 处理一次除法收尾。3. 改动后果把在线 rescale 去掉改回“先存全部分数”假设改成朴素两遍先把 seq_len 个分数都算好存进数组float scores[8192]之类再找 max、再统一 exp——功能等价代价三样多一遍 KV 读取打分一遍、加权 V 又一遍 → KV 扫描 ×2q8 KV 的带宽红利白省一半多一份分数数组8K 上下文要 32KB 栈/堆暂存还破坏 cache 局部性失去“流式”能力在线版可以在 KV 逐块到达跨进程磁盘 KVDay 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.cflash_attn_single_q_q8_neon第 6594 行起、技术文档.md 第 293 行。开源仓库Kestrel-LLM (Gitee)源码可得双许可学习 / 学术研究免费下篇预告单遍扫全 KV 很优雅——可它“算得准吗”下一篇 9-3 做 q8 KV 与 fp32 KV 的精度对照看量化在注意力输出上到底差多少。关键词q8 KV、flash attention、NEON、单遍扫描、在线 softmax上一篇Day 9·1 KV 也量化——q8 KV 把缓存与 decode 带宽压到一半下一篇Day 9·3 q8 KV 精度对照——量化进注意力输出差多少
返回列表