LLM 新人的压榨本地 DeepSeek V4 Flash 记
0. 背景
最近 DeepSeek-V4-Flash-0731 模型能力的提升令人瞩目,刚好自己有一台 AMD Strix Halo, 于是就想看看在这台统一内存的 APU 上能跑多快。因此,我跟着 Fable-5 学习了一些 LLM 的基础知识,让 AI 帮我进行了一些优化实验,最终把 batch-1 decode 从 14.7 提升到 16.7 tok/s,开启 DSpark 再推到 23.8 tok/s。
这篇文章主要写给和我一样有体系结构/系统背景、想学推理框架和 GPU 算子优化的读者。模型本身的数学 (attention 在算什么、RoPE 为什么有效) 只是被优化的对象: 文中会给出必要的直觉, 但公式看不懂完全不影响主线。在优化器眼里, 每个算子都只是"读多少字节、做多少乘加、发几次 kernel launch"– 建立这个视角才是本文的真正目的, profiler 和 roofline 会带你走完全程。
1. 实验环境
- 模型: DeepSeek-V4-Flash-0731, unsloth UD-IQ3_XXS GGUF (98 GB)
- 硬件: AMD Strix Halo (gfx1151, RDNA3.5), 40 CU, 256-bit LPDDR5X-8000 (理论 256 GB/s), 128 GB 统一内存, 32 MB Infinity Cache
- 软件栈: llama.cpp, ROCm/HIP 后端, 开启 flash attention, batch-1 解码
- 代码: 本文全部优化已按独立 commit 整理在 cyyself/llama.cpp 的 dsv4-flash-optimizations 分支
- 硬约束: 零精度损失。每一项优化都必须在 temperature 0 下产生逐字节一致的输出。不允许重新量化。
顺带说明: 虽然是 AMD GPU + ROCm, 但本文的改动都落在 ggml-cuda 目录 — llama.cpp 的 HIP 后端没有单独的 AMD kernel 实现, 就是用 AMD 工具链直接编译 NVIDIA 后端的 ggml/src/ggml-cuda/ 源码树 (CUDA 到 HIP 的 API 映射由 vendors/hip.h 里的一层宏完成)。
2. DSV4-Flash 究竟是怎么推理的
先修知识: 写给 0 基础读者的 LLM 推理速成
(以下内容我也是在 Fable-5 的帮助下整理的, 仅供速成, 不保证完全准确。若你熟悉 LLM, 可以直接跳过。阅读时也不必逐条弄懂: 每个词条真正需要带走的只是张量的形状、字节量和计算/访存模式; 数学细节看不懂就放心跳过, 完全不影响后面的优化章节 — 事实上整场优化战役没有用到任何一条公式的推导。)
-
Token 与自回归解码。 文本被切成 token (大约一个词或词片, 来自 129k 大小的词表)。推理分两个阶段: prefill 一次性并行处理整段 prompt; decode 逐个生成 token — 把上一个 token 送进模型, 得到"下一个 token 是词表中每个词的概率", 选一个, 再循环。对体系结构读者, 最关键的事实是: decode 每生成一个 token, 都要把本步被激活的全部权重从内存完整读一遍 — 对稠密模型这就是所有权重; 对本文的 MoE 模型 (见下文), 是全部非 expert 权重加上被路由选中的 6/256 个 expert, 合计约 8 GiB/token — 其中非 expert 部分约 6.4 GiB (占 80%), expert 部分约 2.1 GiB (98 GB 里, 每步只有这些真正被触碰)。这个比例值得记住: batched 推理时, 所有序列共享同一份非 expert 权重, 即每 token 读取里约 80% 是可跨请求复用的; 只有 expert 那约 20% 因各 token 路由不同而难以摊销 (后文有详细账目)。这 8 GiB 不会命中任何缓存, 这就是 decode 是带宽受限负载的根本原因。
-
Hidden state 与 token embedding。 token 进入模型后由一个定长向量表示, 称为 hidden state — 可以想成这个 token 的"工作寄存器组", 模型的每一层都读取并更新它, 层间传递的就是它。这个向量的初值来自 token embedding: 模型带一张训练出来的 [129k x 4096] 查找表, 词表中每个 token 对应一行向量 (训练的副产品是语义相近的词, 行向量也相近); token 以整数编号进场, 查表取行即得初始 hidden state — 纯 gather 操作, 零算术, 每 token 只读一行 16 KB, 访存可忽略。它的镜像操作在整条流水线的末端且并不便宜: output head 把最终 hidden state 映射回 129k 个分数 (logits), 是一次完整读取 [4096 x 129k] 矩阵 (本模型 0.4 GiB) 的 GEMV, 每个 token 都要付一遍。它的维度是模型设计时定下的超参数 (即 n_embd / hidden size), 各模型不同, 大体随模型规模增长; 本文的 DSV4-Flash 取 4096, 下文所有 "4096" 都源于此。数据类型上它是 F32 激活, 只有 KB 级大小, 全程缓存常驻 — 与动辄 GiB 级的权重形成对照。
-
GEMV 与 GEMM。 BLAS 术语: GEMV = 矩阵 x 向量乘 (General Matrix-Vector multiply), GEMM = 矩阵 x 矩阵乘。这个区别是全文的轴心: GEMV 里矩阵的每个元素只被用一次 — 读进来, 乘对应的向量分量, 累加, 丢弃, 注定带宽受限; GEMM ([M x N] x [N x B]) 里矩阵的每个元素要与 B 列各乘一次 — 同样的访存换来 B 倍的计算, 算术强度随 B 线性上升, B 足够大时负载翻转为计算受限, 分块复用和 tensor core 才有用武之地。对应到 LLM: batch-1 (也就是传统的单流推理)decode 全是 GEMV; prefill、batched serving 和 Speculative decoding 都是小 B 的 GEMM。后文出现的 mmvq/mmvf, 就是 llama.cpp 里量化/浮点权重的 GEMV kernel 家族。
-
Batch 从哪里来, batch-1 意味着什么。 上一条说算术强度随 B 增长, 那这个 B 在现实里是什么? 对 decode 而言, batch = 同一步里并行推进的序列数, 直接等于服务端同时在生成的请求数。数据中心的 serving 系统靠 batching 把几十上百个并发请求拼在一起: 权重只读一遍,但将计算中的 B 个 GEMV 变成了 GEMM,每次权重访存同时推理 B 个用户。照这个账算, B 到 50 左右算术强度就跨过本机的 ridge point (后文 roofline 一节实测约 95 FLOP/byte), 负载从带宽受限翻转为计算受限 — 云端 GPU 常年工作在 B=64-256 的计算受限区, 这正是它们吞吐惊人的原因。而个人设备上只有一个用户在等一个回答: batch-1 既是常态, 也是最坏情况 — 每读 8 GiB 权重只换 1 个 token, 零摊销, 本文的一切优化都发生在这个场景里。单用户想"凭空造出 batch"只有一条路: 不能跨请求摊销, 就跨未来的 token 摊销 — 即下文的投机解码。
-
每层在算什么。 每层两大块: attention (让当前 token 与历史 token 交换信息) 和 FFN (对每个 token 独立做非线性变换)。FFN 的具体形态是三个权重矩阵 (gate/up/down) 和四步计算:
g = W_gate @ x # GEMV, [n_ff x 4096] 矩阵乘输入向量 u = W_up @ x # GEMV, 同上, 与 g 并行无依赖 h = silu(g) * u # 逐元素: 非线性门控, 数据量仅 KB 级 y = W_down @ h # GEMV, 投回 4096 维(记号:
@是矩阵乘, 借自 NumPy;*是逐元素乘。) 注意这套公式对 batch 是透明的: 单 token 时 x 是 [4096] 向量, 三个@都是 GEMV; 一旦一次喂 B 个 token — 不管来自 B 个并发请求、prefill 的一段 prompt, 还是投机验证的 B 个候选 — x 就变成 [4096 x B] 的矩阵, 同样三行自动变成 GEMM: 权重仍只读一遍, 但每个元素与 B 列各乘一次, 算术强度乘 B。这就是上一条 GEMV/GEMM 词条所说"负载翻转"的微观出处。x 的来源: 它就是那个 4096 维 hidden state — 本层 attention 的输出加回 residual stream 后, 过一个 rms_norm (数值归一化, 防止逐层累加导致数值漂移) 即成为 FFN 的输入; FFN 的输出 y 再加回 residual stream, 交给下一层。第一层的源头是 token embedding 查表。(在本文的模型里, 这条链上还套着 hyper-connections: attention 结果先写回四条流、再次混合归一化后才进 FFN, 见后文的数据流图。)
它访问权重的方式极其朴素: 三个矩阵各按行连续地从头到尾流式读一遍, 每个权重元素读进来参与一次乘加就永远丢弃 — decode 下就是三次背靠背的 GEMV, 零复用、零局部性, 这正是"算术强度低到不能再低"的含义。attention 的四个投影矩阵 (wq/wkv/wo 系列) 访问方式完全相同, 只有中间的打分-加权一步例外 (它读的是 KV cache 而不是权重)。MoE 只是在这之上加了一层间接寻址: router 选出的 expert 编号作为索引, 先从 expert 池里 gather 出对应的三元组矩阵, 再做同样的流式读取 (llama.cpp 里对应带索引的矩阵乘 mul_mat_id)。
-
Attention 的具体计算。 一句话: 为当前 token 算一个 query, 与历史每个 token 的 key 比相似度, 按相似度加权读取它们的 value — 一次按内容寻址 (content-addressable) 的读取, 代价随上下文长度线性增长。以 DSV4-Flash 的实际尺寸画出来:
hidden state x [4096] (F32) | | wq_a / wq_b 投影 (Q8_0 权重, 读入时现场解压) v q: 64 个 head x 576 维 KV cache (F16): 每 token 每层 | 1 个 576 维 latent = 1.15 KB | (43 层合计 ~50 KB/token) | | v v 打分 score[h,i] = q[h] . k[i] <- 576 维点积 -> 1 个 F32 标量 | 530 上下文: 64 x 530 个分数 v softmax (逐 head 把 n_ctx 个分数归一化成权重) | v 加权求和 out[h] = sum_i w[h,i] * v[i] <- v 取 latent 的 512 维切片 | v 64 个 head x 512 维, 拼接 -> wo_a/wo_b 投影 -> 回到 [4096]数据类型的分工全文通用: 权重量化存储、kernel 内现场解压; 激活全程 F32; KV cache 是 F16 — 每 token 约 50 KB 的 cache 足迹, 对比 8 GiB 的权重读取, 这就是为什么短上下文时 KV 访存可忽略、32k 时才涨到每 token 约 1.6 GiB 的量级。
-
KV cache。 历史 token 的 k/v 每一步都要用, 而且不变 — 所以缓存下来 (标准的 memoization), 每步只为新 token 追加一行。本模型的 MLA (Multi-head Latent Attention) 更进一步: 不为 64 个 head 各存一份 k/v, 而是只存一个 576 维的压缩 latent 向量, 用时再投影展开 — cache 体积小一个量级, 代价是解码时多做几次小矩阵乘。
-
RoPE (旋转位置编码)。 让点积能感知"相隔多少个 token"。(这是全文数学最密的一条词条, 跳过它不影响任何后续内容 — 对优化而言只需记住最后一句话: 它是一个很便宜、但出现次数非常多的 elementwise kernel。) 做法: 把向量按相邻两个分量分成数对, 每一对看成二维平面上的一个点, 按该 token 的位置 p 转一个角度:
角度 a = p x f # p: token 位置 (0, 1, 2, ...), f: 这一对固定的频率 x0' = x0*cos(a) - x1*sin(a) x1' = x0*sin(a) + x1*cos(a) # 标准二维旋转, 就这么多"固定的频率"是指: 一个 64 维的向量被分成 32 对, 每一对被分配一个常数频率 f — 它只由"这是第几对"决定 (第 0 对最快, 往后逐对等比递减, 到最后一对慢几个量级), 是模型定义的一部分, 与 token 的内容无关、也不随位置变化。会随位置变的只有角度 a = p x f: 想象一块有 32 根指针的钟, 每根指针有自己恒定的转速 (频率), 位置 p 是流逝的时间, 第 i 根指针此刻的角度就是 p x f_i。快指针几个 token 就转一圈, 用来分辨近距离; 慢指针上万个 token 才转一圈, 用来分辨远距离 — 合起来覆盖整个上下文范围。
它起作用靠旋转的一个基本性质: 两个被旋转过的向量做点积, 结果只取决于角度差 — q 在位置 m (转了 m x f), k 在位置 n (转了 n x f), 点积里留下的就是 (m-n) x f, 即相对距离, 与绝对位置无关。计算上就是逐元素的 sin/cos 乘加, 很便宜, 但在这个模型里出现的次数非常多。
-
MoE (Mixture of Experts)。 FFN 的稀疏化: 不是一个大 FFN, 而是 256 个小 FFN (expert), 由一个小 router 网络给每个 token 挑 6 个执行。看本模型的权重构成就明白这个设计的意义: 我们的模型 98 GiB 里, routed experts 占约 90 GiB (93%), 非 expert 权重 (attention、shared expert、output head、embedding 等) 只有约 7 GiB。每 token 实际读取的是: 非 expert 部分几乎全读 (~6.4 GiB) + expert 部分只读 6/256 (~2.1 GiB)。用体系结构的类比, expert 池是对 256 个 bank 的稀疏 gather — 总容量大而单次访问量可控。
-
量化。 权重以 2-8 bit 的压缩格式存放 (Q8_0 = 8 bit 线性量化; IQ3_XXS 约 3.06 bit/权重, 用 codebook 编码), GEMV kernel 边读边解压 (dequant) 边做点积。压缩节省带宽, 解压消耗 ALU 指令 — 后文有一个关键测量正是这两者的交锋。
-
Sampling 与 temperature。 模型每步的原始输出是词表上 129k 个分数 (logits), 经 softmax 变成概率分布; 从这个分布里"挑出最终 token"的策略就是 sampling, 由一条参数链控制。先明确选择机制本身: sampling 是按概率掷骰子, 不是"选概率最大的" — 若分布是 "the" 60%、"a" 30%、"an" 10%, 那么每次生成有 40% 的概率选中的不是最大候选。这就是随机性的来源。temperature 改变的则是骰子各面的大小: 在 softmax 之前把所有 logits 除以 T — T 越大, 各候选的概率被拉得越平, 掷出冷门的机会越大, 输出越"有创造性"; T=1 按模型给出的原始分布掷; T 趋近 0 时分布被推向极端, 最大候选的概率逼近 100%, 掷骰子退化为必中 argmax (即 greedy, 完全确定, 本文全部实验的设定)。直觉上 T 控制的是"冷门 token 有多大机会被掷中": 代码和数学任务常用低温求稳, 创意写作用高温求变化。其余常见参数都是截断器, 在掷骰子前把长尾剪掉以免采出乱码: top-k 只保留概率前 k 名再归一化, top-p (nucleus) 保留累计概率达 p 的最小集合, min-p 丢弃概率不足最高项 p 倍的候选。它们组成一条 sampling 链, 依次作用在 logits 上:
logits [129k] (F32, 模型对词表每个词的打分) | | 除以 T (temperature) v softmax -> 概率分布 | | 截断器链: top-k -> top-p -> min-p (剪长尾, 剩余重新归一化) v 掷骰子 -> 下一个 token temperature 对分布形状的影响 (概率随候选 token 的分布, 示意): token T -> 0 (greedy) T = 1 (原始) T 大 (拉平) "the" ########## ###### ### "a" . ### ### "an" . # ## "one" . . # ... . . # 唯一候选, 头部集中, 长尾抬高, 完全确定 偶有意外 更随机/更发散 -
Speculative decoding (投机解码 / speculative draft)。 decode 的串行性来自一条具体的数据依赖: 要算第 n+1 个 token, 输入必须是第 n 步掷骰子选出的那个 token — 它得先查 embedding、走完 43 层、把自己的 k/v 写进 KV cache, 第 n+1 步的 attention 才有东西可读。用体系结构语言说, 每一步的输出 (选中的 token 和新写入的 cache 行) 恰是下一步的输入, 一条跨步的 RAW 依赖链, 乱序和流水都无从下手 — 但依赖是可以"猜测执行"的: 用一个小而快的 draft 模型先顺序猜出 k 个候选 token (它小, 猜得快也便宜), 再让大模型把这 k 个候选一次性并行验证 — 相当于把 k 步 GEMV 合并成一次小批量 GEMM, 8 GiB 的权重读一遍摊销给 k 个 token。用一段 token 流看它如何工作 (k=3):
已生成: ... 今天 天气 draft 猜: [真] [好] [,] <- 小模型接力猜 3 个, 便宜 大模型验证: [真] [好] [。] <- 一遍 pass, 逐位给出自己的选择 == == != 本轮采纳: 真 好 。 <- 一致则接受; 不一致处用大模型的, 其后作废同样拿到这 3 个 token, batch-1 下的普通 decode 要读 3 遍权重 (一步一遍、无人分摊), 投机只读约 1 遍。猜错也不亏什么: 最坏情况 (第一个就错) 本轮仍然产出 1 个大模型自己算的 token, 等价于普通 decode 的一步。
两个模型的体量差距值得有个直观感受。目标模型 DSV4-Flash: 总参数约 285B (98 GB, 3-bit 级量化), 每 token 激活约 13B、读 8 GiB; dspark draft 模型: 约 5.5B 参数 (11 GB, BF16), 每猜一个 token 只读约 0.7 GiB。也就是参数量差约 50 倍, 单 token 访存差约 11 倍 -- 一轮猜 3 个候选的全部开销, 大约只相当于大模型四分之一个 pass。这么小的模型为什么能猜中六七成? 因为语言生成里大部分 token 是"容易的": 语法接续、固定搭配、代码模板, 小模型足以胜任; 真正需要"思考"的少数决策点, 会以验证失败的形式自动交还给大模型把关 -- 投机解码本质上是按难度分工。此外 dspark 这个 draft 并非独立瞎猜: 它以大模型的中间 hidden state 作为输入起草, 相当于站在大模型的肩膀上做短程外推, 这也是它命中率高的原因之一。
是否接受, 由 sampling 参数说了算: 每个位置, 大模型按用户设定的 sampling 链选出自己的 token, 与 draft 逐位比对, 相同才接受。 temperature 0 下"接受"就是两个模型的 argmax 一致, 接受率最高 (本文实测的 0.61-0.74 皆为此设定); 温度升高后大模型一侧也在掷骰子, draft 即使猜中概率最高的 token 也会被随机 miss, 接受率随之下滑, 低到一定程度, 加速就被 draft 自身的开销吃掉; top-k/min-p 这类截断器收窄候选集, 能挽回一部分损失。temp > 0 时输出本就是随机的, 投机保证每个 token 仍严格出自大模型的分布 -- 只影响速度, 不改分布。体系结构类比再贴切不过: 这就是分支预测加推测执行 -- 猜对了流水线满载, 猜错了 squash 重来, 只是这里的"预测器"本身是个小神经网络。第 6 节有它与本文 kernel 优化叠加后的实测。
-
为什么能要求逐字节一致。 temperature 0 意味着每步都取概率最高的 token (argmax), 完全确定。因此任何数学上等价的优化都应产出一模一样的文本 -- 这让"逐字节 diff"成为可行且极其严格的正确性标准。
逐层在算什么, 以及 prefill 与 decode 的关系
层的宏观图景。 前面说过, 每个 token 由一个 4096 维的 hidden state 表示; 43 层做的事情是对它的逐层精炼: 每层先用 attention 从历史 token 的 hidden state 里按相关性取回信息、混入当前向量, 再用 FFN 对它做一次非线性变换; residual stream 则像一个累加器, 每层的结果都是"加"在原向量上而不是替换它 (这也是层可以堆得很深的原因)。43 层是一条固定深度的流水线, 每个 token 都要完整走一遍, 最后的向量经 output head 映射成对下一个 token 的预测分布。
同一条流水线, 两种驱动方式。 prefill 和 decode 跑的是完全相同的每层计算, 差别只在一次喂多少 token -- 但这个差别彻底改变了负载性质:
- Prefill (处理 prompt): 宽而并行。 N 个 token 同时进入流水线, 每个矩阵乘从 GEMV 变成 GEMM ([权重] x [N 个激活向量]): 权重读一遍服务 N 个 token, 算术强度乘以 N。attention 部分是因果的 (每个 token 只看得到它之前的), N 个 token 的 KV 一次性批量算好、灌入 KV cache, compressor states 也在这里批量建立。prefill 决定的是"首 token 延迟" (time-to-first-token)。
- Decode (逐 token 生成): 窄而串行。 自回归意味着硬数据依赖 -- 第 n+1 个 token 的输入是第 n 个的输出, 未来 token 之间无法并行。于是每步都是 batch-1 的 GEMV 洪流: 权重读一遍只服务 1 个 token, 算术强度 2 FLOP/byte, 纯带宽受限。每步在 KV cache 和各 state 后面追加一行。decode 决定的是生成吞吐 -- 也就是本文优化的全部对象。
权重构成与可复用性。 速成一节说过 MoE 会削弱 batching 收益, 用本模型的权重构成把这句话变成数字:
| 权重类别 | 总大小 | 每 token 读取 | batching 下的可摊销性 |
|---|---|---|---|
| 非 expert (attention, shared expert, output head 等) | ~7 GiB | ~6.4 GiB (几乎全读) | 完美 -- 每个 token 读的是同一份, B 个序列共享一遍读取 |
| routed experts | ~90 GiB | ~2.1 GiB (6/256) | 打折 -- 各 token 选各的 expert, B 增大时被触碰的 expert 集合一路膨胀到近乎全量 |
也就是说, 占 93% 的 expert 权重, 在单 token 视角下只贡献 25% 的访存; 而 batching 能完美摊销的, 恰恰是那 7 GiB 人人都要读的非 expert 部分。这就是 MoE 的"读取经济学": 用巨大的容量换取每 token 的低访存, 代价是牺牲一部分 batching 的规模效应。
DSV4-Flash 每层的具体结构
有了以上概念, 再看这个模型为什么不是普通 transformer。它每层 (共 43 层, n_embd 4096) 都在做什么:
MLA attention (576/512 head)。 它对标准 attention 做了两件"省"的事, 都不难懂:
一是投影走"窄门" (low-rank bottleneck)。从 4096 维算出 64 个 head 的 q, 本来一个大矩阵一步到位; MLA 改成两步接力 -- 先把 4096 维压进一个 1024 维的窄通道 (wq_a), 再从窄通道展开成 64 个 head (wq_b)。两个"小"矩阵的总体积远小于一个"大"矩阵, 权重更少, 每 token 要读的字节也更少。代价是多一次矩阵乘和一点表达能力, 换来的是访存 -- 对带宽受限的 decode 来说是划算的买卖。
二是 KV cache 只存"浓缩摘要"。标准做法要为 64 个 head 各存一份 k/v; MLA 只为每个 token 存一份 576 维的压缩摘要 (latent), 用的时候再展开给各个 head -- cache 体积直接小一个量级 (前文 50 KB/token 就是这么来的)。
工程细节有两个, 记住结论即可: 每个 head 的向量分成两段, 大段 (512 维) 只管"内容"、小段 (64 维) 只管"位置", RoPE 只旋转位置段 -- 于是计算图里 q 和 kv 都有一串 view -> rope(尾部) -> concat 的拆接操作, 输出侧甚至要做一次反向旋转 (de-rope, 把摘要里带的位置旋转"拧"回去)。另外 K/V head 宽达 576/512 维, 超出了 flash-attention 快速 vector kernel 的支持范围, 解码只能走较慢的 tile kernel -- 这两个细节在后文的优化里都会再次登场。
Lightning indexer (DSA)。 一个"给大 attention 当索引的小 attention": 每层先用一套低成本的投影对全部历史 KV 条目打相关性分, 做 top-k, 正式的 attention 只对被选中的条目计算。作用类似为按内容寻址的读取建了一个粗筛索引 -- 长上下文时省下的是主 attention 的算力和带宽, 代价是这套索引自身的投影、rope 和排序开销。
Hyper-connections (hc_mult=4)。 普通 transformer 的层间信息通过一条 4096 维的 residual stream 传递 (每层的输出加回输入, 类似流水线里的旁路)。这个模型携带的是四条并行的流: 每层先用一个小矩阵乘 (hc_fn, F32, [24 x 16384]) 加 sigmoid gating 计算"读哪几条、各占多少权重", 把四条流混合成一条工作流交给 attention/FFN, 算完再按另一组门控权重写回四条流。可以类比成四组可切换的工作寄存器堆, 读写端口的选择是模型学出来的。每层这个包裹做两次 (attention 一次, FFN 一次)。
Compressor state machines (CSA / HCA / LID)。 完整 KV cache 只保留近处的历史; 更远的历史被压缩成固定大小的摘要 -- 类似存储层次: 近处全细节, 远处只留概要。每层有三套这样的滑动状态 (attention 用的 CSA/HCA, indexer 用的 LID), 每个 token 都要: 从状态缓冲区 gather 出一个窗口, 逐通道做 softmax 加权平均 (一种学习出来的 pooling, 权重来自一个并行维护的打分状态), 归一化、rope, 再把新条目写回。计算本身很轻, 但算子数量惊人 -- 优化前仅这部分就占每 token 约 1,000 次 kernel launch。
MoE FFN。 router 是一个 [4096 x 256] 的小矩阵乘: 输出 256 个分数, 过 sigmoid 加 bias 后取 top-6, 这 6 个 expert 各自执行一个标准 FFN -- 三个矩阵乘 (gate/up/down) 加 SwiGLU 非线性 (逐元素的 silu(gate) x up), 结果按 router 权重加权求和; 另有一个所有 token 都走的 shared expert。这个模型的特别之处是 SwiGLU 带 clamp (在 gate/up 矩阵乘和 GLU 之间截断数值, 为了数值稳定) -- 这个小小的 clamp 后来成了第一项优化的主角。最后, 43 层走完后, 一个 [4096 x 129k] 的 output head 矩阵乘把最终向量映射成整个词表的分数 (logits), argmax 得到下一个 token。
把一层的数据流画出来是这样 (hc 包裹每层出现两次: attention 一次, FFN 一次):
x0 x1 x2 x3 <- 4 residual streams (hyper-connections)
| | | |
+---+---+---+--> hc mix: hc_fn matmul + sigmoid gates -> one working stream x
|
rms_norm
|
+--- MLA attention ---+------------------------------------------------+
| | |
| q path: x -> wq_a -> q_lora -> wq_b -> per-head [ nope | pe ] |
| rope only the pe tail |
| kv path: x -> wkv -> kv latent [ nope | pe ] -> rope the pe tail |
| -> KV cache (single latent head) |
| |
| lightning indexer: score history keys, keep only top-k |
| CSA / HCA / LID: softmax-pooled compressed history states |
| | |
| FA tile (K=576 / V=512): q x [ selected cache + compressed ] |
| -> out -> de-rope a 64-dim slice -> wo_a -> wo_b |
+----------------------------------------------------------------------+
|
4 streams <-- hc gates write the attention result back
| | | |
+---+---+---+--> hc mix again -> x
|
rms_norm
|
+--- MoE FFN ---------+------------------------------------------------+
| router (BF16): pick 6 of 256 experts |
| 6 routed experts (IQ3_XXS / IQ2_XS) + shared expert (Q6_K) |
| all clamped SwiGLU; weighted-sum the expert outputs |
+----------------------------------------------------------------------+
|
4 streams <-- hc gates write the FFN result back
| | | |
v v v v
next layer (x43 total) ... -> final hc head mix -> output head -> logits
Attention 与 FFN: 计算和访存各自在做什么
把一层里两大块的每一步摊开看 (以下均为 decode、单 token、本模型的实际尺寸):
Attention 这一侧:
- Q 投影 (两级 GEMV)。 输入 x [4096] 先乘 wq_a [4096 x 1024] 得到 q_lora, 再乘 wq_b [1024 x 32768] 展开成 64 个 head x 512 维。访存: 两个矩阵合计约 40 MB (Q8_0), 每个字节读进来只参与 2 FLOP 的乘加, 用完即弃 -- 教科书级的纯流式读取, 无任何复用。计算: 约 70 MFLOP, 对 20 TFLOPS 的机器不到 4 us 的活, 时间全花在等内存上。
- KV 投影与 cache 追加。 x 乘 wkv 得到 576 维 latent, 对 64 维 pe tail 做 rope, 追加进 KV cache。写流量每层约 1 KB -- 可忽略。这里体现 MLA 的设计意图: cache 里存的是压缩 latent, 不是 64 份展开的 k/v。
- 打分与加权求和 (FA tile kernel)。 q 与历史每个位置的 latent 做点积 -> softmax -> 加权求和。访存: 读整个 KV cache, 每层 n_ctx x 576 x 2 B -- 上下文 530 时约 0.6 MB, 32k 时约 37 MB, 这是全层唯一随上下文线性增长的访存。计算: 64 head x n_ctx x (576+512) x 2 FLOP, 530 上下文时约 74 MFLOP -- 这是每层里算术强度最高的一块 (数据在 cache 内被 64 个 head 复用), 但绝对量依然轻。softmax 本身访存和计算都极小, 属于延迟型算子。
- 输出投影 (又是两级 GEMV)。 attention 结果先 de-rope 一个切片, 再过 wo_a/wo_b 两个矩阵回到 4096 维 -- 又是几十 MB 的权重流。
FFN (MoE) 这一侧:
- Router。 一个 [4096 x 256] 的 BF16 GEMV (2 MB) 输出 256 个分数, sigmoid 加 bias 后 argsort 取 top-6。排序和 gather 是典型的延迟型小 kernel: 数据量以 KB 计, 时间全是 launch 和访存延迟。
- 6 个 routed expert。 每个 expert 是三个 GEMV: gate [4096 x 2048]、up [4096 x 2048]、down [2048 x 4096], IQ3_XXS/IQ2_XS 编码下约 10 MB/expert, 6 个合计约 50 MB/层。地址模式值得注意: 6 个 expert 在 90 GiB 的 expert 池里位置离散 (稀疏 gather), 但每个 expert 内部是连续大块 -- 对 DRAM 仍然友好; 真正的代价在解压: codebook 反量化每字节要花约 0.9 条 VALU 指令, 是这些流跑不满带宽的原因 (见第 3 节)。
- 逐元素 SwiGLU。 silu(gate) x up 加 clamp, 作用在 [2048 x 6] 上 -- 数据量以 KB 计, 却要占三四个 kernel launch, 是"launch 税"的典型来源。
- Shared expert 与合并。 所有 token 都走的一份 FFN (Q6_K, 约 19 MB/层), 与 6 个 expert 的输出按 router 权重加权求和, 经 hc gates 写回 residual streams。
每层每 token 的访存台账 (43 层合计即全模型的约 8 GiB):
| 访存对象 | 每层读取量 | 特征 |
|---|---|---|
| attention 权重 (wq_a/wq_b/wkv/wo_a/wo_b) | ~114 MB | Q8_0 流式, 已顶满带宽Roofline |
| routed experts (6/256) | ~50 MB | IQ3_XXS/IQ2_XS, 解压 ALU 受限 |
| shared expert | ~19 MB | Q6_K 流式 |
| hc_fn / router / compressor 投影等 | ~8 MB | F32/BF16 小矩阵 |
| KV cache 读取 | 0.6 MB @ 530 ctx, ~37 MB @ 32k | 唯一随上下文增长的项 |
两点总结: 其一, attention 和 FFN 的访存主体都是权重, 不是激活 -- 激活向量只有 KB 级, 全程缓存常驻; 其二, 双方的算术强度同为 2 FLOP/byte, 在 roofline 上落在同一个位置, 所以下一节可以把整层当成一个统一的带宽问题来记账。
结果是: 一步解码就是一张 约 8,000 个节点 (约 5,000 次 kernel launch) 的计算图, 要读取 约 8 GiB 权重 -- 其中大头是 4.8 GiB 的 Q8_0 attention 权重 (UD 量化有意把 MLA 保持在 8 bit; 只有 routed experts 是 2-3 bit), 以及约 2.1 GiB 的 expert 切片。
3. 硬件 Roofline
从这一章起, 模型数学退场, 主线回到体系结构的地盘 -- 也是本文真正想教的东西。前两章哪怕只记住了一件事, 记住这个就够: 每生成一个 token, GPU 要流式读取约 8 GiB 权重、执行约 5,000 次 kernel launch。 后面的全部分析和优化, 都只围绕这两个数字展开。
| Roofline | 理论值 | 实测值 |
|---|---|---|
| 内存带宽 | 256 GB/s | 218 GB/s (大尺寸 Q8_0 GEMV 流式读取, 85%) |
| 计算 (int8 dot / MMQ) | 约 29.7 TFLOPS fp32 双发射 | 20.7 TFLOPS (Q4_0 MMQ) |
| rocBLAS F16 / F32 GEMM | - | 3.7 / 0.47 TFLOPS (针对 gfx1151 几乎没有调优) |
Ridge point(最佳的计算/访存比): 约 95 FLOP/byte。 Batch-1 解码对每个权重元素只做 2 FLOP, 每元素 0.38-1.06 字节 -- 算术强度 (arithmetic intensity) 只有 2-5 FLOP/byte, 比 ridge 低 20-50 倍。解码是最极端的带宽受限负载; CU 的浮点利用率只有个位数百分比。
由此得到一个硬性的 token 速率天花板: 每 token 8.05 GiB, 除以可达的流式读取时间 (混合量化类型下约 44.5 ms) = 如果所有非流式读取的微秒都免费, 上限是 22.5 tok/s。机器起点是 14.7 tok/s = 68 ms/token, 意味着在 44.5 ms 的流式读取地板之上, 还压着约 23.5 ms/token 的开销。
为了把这件事说得具体, 下图是 GPU 为每一个 token 必须执行的工作。每步解码要顺序走完全部 43 层; 图中每个方框都是一串 kernel launch, 其中权重密集的方框被钉死在内存Roofline上:
one token = one pass through this, x43 layers, then the output head
token embedding
|
v
+--- layer l (x43) ---------------------------------------------------------+
| |
| hyper-connection mix hc_fn matmul [24 x 16384] F32 + sigmoid gates |
| | |
| MLA attention q/kv low-rank projections <- Q8_0 GEMVs, the |
| | rope on the 64-dim pe tail biggest stream |
| | FA tile kernel (576/512) |
| | lightning indexer (DSA) scoring + top-k |
| | |
| compressor states CSA / HCA / LID: gather sliding state, |
| | softmax-pool per channel, norm, rope, rewrite |
| | |
| MoE FFN router (BF16) -> pick 6 of 256 experts |
| | expert GEMVs <- IQ3_XXS / IQ2_XS streams |
| | shared expert <- Q6_K stream, clamped SwiGLU |
| | |
| hyper-connection combine (4 residual streams) |
+---------------------------------------------------------------------------+
|
v
output head (Q6_K GEMV over the 129k vocab) -> logits -> sample 1 token
下图则是未优化版本 68 ms 的实测去向 (kernel 级 trace 测量, 1 个字符 = 1 ms)。最上面的块被钉在内存 Roofline 上, 不改变字节数就无法缩小; 它下面的所有行, 才是本文要"退税"的开销:
weight streaming |============================================| 44.5 ms at the 210 GB/s roof
- Q8_0 attention |========================| 24.5 ms
- IQ3/IQ2 experts|=============| 12.9 ms (dequant ALU-bound)
- Q6_K shexp+head|=======| 7.1 ms
tiny kernels |==========| ~10 ms 1-4 us each, x1000s
dispatch + gaps |======| ~ 6 ms ~2 us/launch floor
host + sampling |====| ~ 4.5 ms
attention + indexer|==| ~ 2.5 ms FA tile, top-k
BF16 matvecs (MLA) |==| ~ 2.3 ms latency-bound
total ~68 ms -> 14.7 tok/s
streaming floor alone: 44.5 ms -> 22.5 tok/s ceiling
这张图让策略一目了然: GPU 在每个 token 上花三分之二的时间做那件它无法再加速的事 (以 Roofline 速度流式读取权重), 剩下三分之一在做成千上万件小事 -- 它们之所以存在, 只因为这个架构足够复杂。整场优化战役, 打的就是图里下面那几行。
另一个关键结论: 流式读取本身已经贴着 Roofline 了 -- 硬件计数器显示 Q8_0 GEMV 的内存单元占用率达 98.3%。这一点对不同量化格式并不一视同仁。用 test-backend-ops 对同一个 [16384 x 8192] 的 DRAM-resident GEMV (权重体积远超 32 MB Infinity Cache) 扫一遍各量化类型, 用"权重字节数 / 实测耗时"折算有效带宽:
| 量化类型 | bit/权重 | 权重体积 | 实测耗时 | 有效带宽 |
|---|---|---|---|---|
| Q8_0 | 8.5 | 142.6 MB | 623.1 us | 228.9 GB/s |
| Q6_K | 6.5625 | 110.1 MB | 484.1 us | 227.4 GB/s |
| Q5_0 | 5.5 | 92.3 MB | 409.5 us | 225.3 GB/s |
| Q5_K | 5.5 | 92.3 MB | 413.9 us | 222.9 GB/s |
| IQ4_XS | 4.25 | 71.3 MB | 315.9 us | 225.7 GB/s |
| Q4_K | 4.5 | 75.5 MB | 353.3 us | 213.7 GB/s |
| IQ3_XXS | 3.0625 | 51.4 MB | 284.6 us | 180.6 GB/s |
| IQ2_XS | 2.3125 | 38.8 MB | 280.9 us | 138.2 GB/s |
规律一目了然: 线性量化 (Q8_0/Q5_0) 和 k-quants (Q4_K/Q5_K/Q6_K/IQ4_XS) 的反量化足够便宜, 全部贴着 213-229 GB/s 的硬件Roofline流动; 而 codebook 类的 IQ 量化掉到 138-181 GB/s。(这一批扫描的 Q8_0 比正文他处引用的 218 GB/s 略高: 不同批次的测量随时钟状态有约 ±5% 的正常波动, 不影响任何结论。)最能说明问题的是最后两行: IQ2_XS 比 IQ3_XXS 少读 25% 的字节, 耗时却几乎一样 (280.9 vs 284.6 us) -- wall time 被钉在指令发射上而不是字节上。PMC 计数器给出定量解释: 它们的 codebook 反量化每流式读取一个字节要消耗约 0.9 条 VALU 指令 (是 Q8_0 的 15 倍), ALU 发射占用约 74%、内存占用 84% -- 双重受限, 任何缓存技巧都救不了 (我们试过; 见下文"阵亡名单")。这也解释了在线看到的 expert 流带宽 (141-173 GB/s, 形状更小的 expert 切片): 模型里 2-3 bit 的 routed experts 恰好全押在了最吃 ALU 的格式上。
所以整场优化战役可以归结为一句话: 在不改变任何一个字节数学运算的前提下, 削减那约 23 ms 的非流式时间。
4. 如何找到优化点
方法论和修复本身同样重要:
- rocprofv3 提供 kernel 级真相。 完整 kernel trace, 按 kernel 间隙聚类重建出每个 token 的 pass, 得到逐 kernel 的时间账本: 约 5,000 次 launch 各花多少。结论: 约 770 个权重 GEMV 已在硬件速度运行, 另有数千个 1-4 us 的 elementwise kernel, 单个微不足道, 合计约 10 ms -- 纯粹的 launch/延迟税。
- 一次性图 dump 提供计算图级真相。 一个环境变量控制的遍历器, 打印每个节点的形状和来源。它找出了状态机 (占 launch 的 20.7%)、占用率不足的矩阵乘、以及被 VIEW 节点打断而让现有 fusion 匹配器失效的算子链。
- 硬件计数器裁决假设。
MemUnitBusy和SQ_INSTS_VALU区分出"内存已饱和" (Q8_0: 别动它) 和"指令受限" (IQ 量化: 需要重写 kernel, 而不是加缓存)。 - 只认在线 A/B。 整个项目最重要的教训: 孤立的 kernel benchmark 在这台机器上会撒谎。 一个快 42% 的 rms_norm 让服务端整体退步 2.5%; 把 codebook 放进 LDS 毫无收益还倒亏 1.2%。缓存常驻的微基准无法复现 DRAM 轮转和依赖链条件。每一项改动都通过重启真实服务、对比稳态 tok/s 来评判 -- 无论孤立测试看起来多漂亮, 在线退步一律回滚。
- 以逐字节一致为正确性标准。 相同 prompt, temperature 0, 每次改动后与未打补丁的版本逐字节对比。融合 kernel 里显式强制算子边界舍入 (
__fmul_rn/__fadd_rn), 防止编译器悄悄把乘加链收缩成会改变结果的 FMA。
下面用三个真实例子展示这些工具具体怎么用、输出长什么样、以及从输出走到结论的推理。
例 1: rocprofv3 kernel trace -- 给 5,000 次 launch 记账
对正在服务的进程整段抓 kernel trace, 解码约 130 个 token 后停掉:
$ rocprofv3 --kernel-trace --stats -d rocprof -o llsrv -- \
./build/bin/llama-server -m DeepSeek-V4-Flash-0731-UD-IQ3_XXS-00001-of-00004.gguf \
-ngl 99 -fa on -c 32768 -ub 32 --no-mmap --port 8041
$ # 另一个终端: 发一个正常的 chat 请求, 等生成结束后停掉 server
--stats 直接产出按 kernel 名聚合的账本 llsrv_kernel_stats.csv。按总耗时排序 (节选; 这段 trace GPU 总忙时 9.0 s, 共 681,388 次 dispatch ≈ 每 token 5,200 次):
kernel calls total avg 占比
mul_mat_vec_q<Q8_0> 40640 4174 ms 102.7 us 46.3% <- attention 权重 GEMV
mul_mat_vec_q<IQ3_XXS> 9525 1010 ms 106.0 us 11.2% <- expert GEMV
mul_mat_vec_q<Q6_K> 16132 719 ms 44.6 us 8.0%
mul_mat_vec_q<IQ2_XS> 6350 624 ms 98.2 us 6.9% <- expert GEMV
mul_mat_vec_f<float> 13719 194 ms 14.1 us 2.1% <- [24 x 16384] 瘦长 f32
flash_attn_tile<512,...> 5461 122 ms 22.3 us 1.4%
quantize_q8_1 80428 115 ms 1.4 us 1.3% <- 每 token 约 590 次!
rms_norm_f32<...> 28425 158 ms 4.5-7.2 us 1.7%
...
(avg < 5 us 的小 kernel 共 44 种, 554,083 次 call, 合计 1,019 ms)
同一张表读出三件事: (a) 前四行量化 GEMV 合计 72%, 平均耗时与"张量字节数 / 218 GB/s"吻合 -- 大头已经贴着硬件速度跑, kernel 本身没得改, 只能减少字节或提高占空比; (b) 55 万次 avg < 5 us 的小 kernel 几乎不搬数据, 却吃掉约 1 秒 GPU 时间, 还没算它们之间的 launch 间隙 -- 这就是 launch 税, 直接指向 state-pool 融合算子和 HC gate 融合; (c) mul_mat_vec_f 在 [24 x 16384] 这类瘦长矩阵上平均 14.1 us -- 24 行只发 24 个 workgroup, 40 个 CU 大半闲着, 指向 split-K。逐条 dispatch 的时间轴在 llsrv_kernel_trace.csv 里, 按 dispatch 间隙聚类即可重建每个 token 的 pass, "流式读取只占 pass 时间 65%" 的占空比结论就是这么算出来的。
例 2: rocprofv3 PMC 计数器 -- 判定 memory-bound 还是 ALU-bound
kernel trace 只能告诉你谁慢, 说不出为什么慢。硬件性能计数器 (PMC) 可以。对 DRAM-resident 的 GEMV perf case 采 MemUnitBusy:
$ rocprofv3 --pmc MemUnitBusy -d pmc -o q8 -- \
./build/bin/test-backend-ops perf -b ROCm0 -o MUL_MAT
$ python3 - <<'PY'
import csv
rows = [r for r in csv.DictReader(open('pmc/q8_counter_collection.csv'))
if 'mul_mat_vec_q' in r['Kernel_Name']]
v = [float(r['Counter_Value']) for r in rows]
print(f"n={len(v)} mean={sum(v)/len(v):.1f}")
PY
n=3409 mean=98.3
Q8_0 GEMV 的内存单元占用率 98.3% -- 内存管线已经打满, 这类 kernel 进入"别碰"名单。换成 IQ 量化的 expert 形状, 再加上指令计数器:
$ rocprofv3 --pmc SQ_INSTS_VALU SQ_BUSY_CYCLES GL2C_HIT_sum -d pmc2 -o iq3_xxs -- \
./build/bin/test-backend-ops perf -b ROCm0 -o MUL_MAT_ID
把 SQ_INSTS_VALU 除以每次 dispatch 流式读取的权重字节数: IQ3_XXS 每字节要发 0.93 条 VALU 指令, IQ2_XS 0.87 条, 约是 Q8_0 的 15 倍; 同时 ALU 发射占用约 74%、内存占用 84% -- 双重受限。这一条测量直接宣判了"把 codebook 搬进 LDS 加速解码"想法的死刑 (瓶颈不在查表的访存, 在解码指令流本身的串行发射), 后来的在线 A/B (-1.2%) 只是补了一枪。
例 3: 图 dump -- 找到 fusion 缺口 (以及它会骗你的地方)
kernel trace 看到的是执行结果, 计算图 dump 看到的是结构。llama-context 里加了一个环境变量控制的一次性遍历器, 打印每个节点的算子、形状和来源:
$ LLAMA_SPEC_STATS=1 LLAMA_GRAPH_DUMP=1 ./build/bin/llama-server ... 2> graph.log
$ grep -nE "CLAMP|GLU" graph.log | head
马上看到 clamped-SwiGLU 的 fusion 缺口 (下面是 shared-expert dense 链, MoE 的 MUL_MAT_ID 链同形):
gdump 116 MUL_MAT ffn_gate-0 [2048,1,1,1] <- blk.0.ffn_gate_shexp.weight <- ffn_norm-0
gdump 117 CLAMP ffn_gate_clamped-0 [2048,1,1,1] <- ffn_gate-0
gdump 118 MUL_MAT ffn_up-0 [2048,1,1,1] <- blk.0.ffn_up_shexp.weight <- ffn_norm-0
gdump 119 CLAMP ffn_up_clamped-0 [2048,1,1,1] <- ffn_up-0
gdump 120 GLU ffn_swiglu_limited-0 [2048,1,1,1] <- ffn_gate_clamped-0 <- ffn_up_clamped-0
ggml-cuda 本来就有 MUL_MAT, MUL_MAT, GLU 的融合匹配器, 但 DSV4 的两个 CLAMP 插在中间, 模式匹配不上 -- 修复就是给匹配器扩展 MM, CLAMP, MM, CLAMP, GLU 模式, kernel 里加 clamp 收尾。再看 HC gate 链, 就明白为什么第一版 gate 融合一个也没匹配上:
gdump 5 MUL_MAT hc_mixes-0 [24,1,1,1] <- blk.0.hc_attn_fn.weight[16384,24]
gdump 6 VIEW hc_mixes-0 (view) [4,1,1,1] <- hc_mixes-0[24,1,1,1]
gdump 7 VIEW blk.0.hc_attn_scale.weight (view) [1,1,1,1]
gdump 8 MUL node_8 [4,1,1,1] <- 上面两个 view
gdump 9 VIEW blk.0.hc_attn_base.weight (view) [4,1,1,1]
gdump 10 ADD node_10 [4,1,1,1]
gdump 11 UNARY node_11 (sigmoid) [4,1,1,1]
gdump 12 SCALE hc_pre-0 [4,1,1,1]
MUL / ADD / UNARY 之间穿插着 VIEW 节点 -- 只看相邻节点的匹配器必然落空, 必须跳过 no-op 节点再比对。对 dump 做统计还能给融合立项排优先级: compressor 状态机相关节点每 pass 约 1,021 次 launch (占 20.7%), 这就是 GGML_OP_DSV4_STATE_POOL 立项的依据。
一个重要的反面教训: 图 dump 反映的是构图, 不是执行。 dump 里 RMS_NORM 和后面的 MUL 看起来是两个节点, 我们据此追了一条"补上 norm+mul 融合"的线索 -- 结果 CUDA backend 在执行期早就把它们融合了 (执行期融合不改图, dump 里看不见), 白忙一场。图 dump 负责提出假设, 判决永远交给 trace 和在线 A/B。
5. 落地的优化
下表按落地顺序排列, 每一行对应 dsv4-flash-optimizations 分支上的一个独立 commit, 可以逐个对照阅读。后续我打算把这些改动整理后逐步提交到 llama.cpp 上游。
| 改动 | 做了什么 | tok/s |
|---|---|---|
| (基线) | 14.70 | |
| Clamped-SwiGLU 融合 | ggml-cuda 本来就有融合的 gate+up+GLU matvec, 但 DSV4 在矩阵乘和 GLU 之间的 clamp 算子破坏了模式匹配。给匹配器扩展了 MM, CLAMP, MM, CLAMP, GLU 模式, kernel 里加了 clamp 收尾。43 层 MoE 中融合了 41 层。 |
14.97 |
| 瘦长 F32 matvec 的 Split-K | Hyper-connection 混合矩阵是 [24 行 x 16384] -- 原 kernel 每行一个 workgroup: 40 个 CU 上只有 24 个 block。Split-K 路径 (按列切片算部分和 + 微型归约) 把 kernel 从 15.9 us 降到 6.6 us。 | 15.40 |
| BF16 rows-split | 对 MoE router ([256 x 4096] BF16, 延迟受限的微小 block) 做同样的处理。 | 15.44 |
GGML_OP_DSV4_STATE_POOL |
新的融合算子, 把每条 compressor state 链 -- 补零、gather、窗口拆分、per-channel softmax、加权池化: 约 13-15 次 launch -- 折叠成一个 kernel (ratio<=8 用寄存器数组变体, 128 条目的 HCA 用循环变体)。后来又扩展为直接消费状态追加的 concat。每 token 减少约 1,500 次 launch。 | 16.08 |
| HC gate 融合 | y = s*sigmoid(x*a+b)+c 链 (4 个 kernel x 172 个 gate) 融合成一个 kernel。第一版一个也没匹配上 -- VIEW 节点穿插在链中间; 匹配器必须跳过 no-op 节点。 |
16.53 |
| Rope-region | 新原语 ggml_rope_ext_region: 只对 [start, start+n) 维做 rope, 其余原样传递。由于 DSV4 的 nope/pe 是同一父张量的两个 view, 所有 view -> rope -> concat 三元组 (q、kv、两处 compressor 尾部; 约 170 处) 都塌缩成单个 rope 节点。在 ggml core + CPU + CUDA 实现, start=0 时可证明与原行为等价。 |
16.68 |
| FA tile 配置 | RDNA 的 576-head tile 配置原为 (256 线程, occupancy 4, 64 宽 tile); 扫描发现 (512, 1, 128, 128) 孤立测试快 22%。但 FA 只占整个 pass 的 1.5%, 在线收益: +0.04。 | 16.71 |
最终: 14.70 -> 16.71 tok/s (+13.7%), 每 token 减少约 2,000 次 kernel launch, 全程输出逐字节一致。 整 token 的有效内存带宽从约 127 升到约 139 GiB/s -- 不是因为哪条流变快了, 而是流式读取的占空比从 65% 提高到 70% 以上。
6. 完整结果, 包括投机解码
以上全部针对 batch-1 非 draft 解码, 但收益并不局限于此。DeepSeek-V4-Flash 配有 dspark draft 模型 (一个 11 GB 的 BF16 sidecar 模型, 每轮以 block 形式一次 draft 3 个 token), 而 temperature 0 下的投机解码是数值精确的: 目标模型会校验每一个 draft token, 所以输出与非 draft 解码逐字节相同 -- 我们真的逐字节对比验证过。
在未打补丁和最终版本上用完全相同的协议重测 (512 token 生成, temp 0, draft n_max=3):
| 配置 | 优化前 | 优化后 | 提升 |
|---|---|---|---|
| 非 draft 解码 | 14.70 tok/s | 16.71 | +13.7% |
| dspark, 散文 prompt | 20.26 tok/s | 20.93 | +3.3% |
| dspark, 代码 prompt | 20.81 tok/s | 23.76 | +14.2% |
Kernel 优化与投机解码是复利关系: 目标模型的 batch-4 校验 pass 同样运行融合状态池、rope-region 和 elementwise 融合 (只有限定 batch-1 的 matvec 融合缺席)。提醒一句: 表中的接受率是 temperature 0 下的数字 -- 按速成一节的分析这是接受率的上限情形, 更高温度的采样会让接受率与加速比一同缩水。代码类 prompt 受益最大, 因为其更高的 draft 接受率 (约 0.74 vs 约 0.61)。
优化后这台机器的精确解码峰值吞吐: 23.8 tok/s -- 比整个旅程的起点高 62%。
7. 阵亡名单 (同样重要)
有十一个假设死于在线实测。值得记录的几个:
- IQ3_XXS codebook 进 LDS: 针对"gather 污染流式读取"的"显然"修复。退步 1.2% -- gather 本来就便宜地命中缓存; 真正的开销是反量化的指令发射。
- 给 mmvq 加大 warp 数: gfx1151 落到一个未调优的默认值 (1 warp, 每 block 1 行), 看起来像白捡的钱; 但 Q8_0 用 nwarps=8 (W7900 的调优值) 退步 10.6%。未调优的默认值恰好是最优的。
- 小 block rms_norm: 独立测试 -42%, 在线 wall clock +2.5% 损失。
8. Roofline 对结局的裁定
在 16.71 tok/s 时, 剩余 budget 是: 44.5 ms 被 DRAM 饱和的流式读取 (不改变字节数不可缩减)、真实数学运算 (attention、indexer、compressor 投影)、以及约 3,500 次剩余 launch 各自约 1.5-2 us 的硬件派发地板。唯一超过 0.3 ms 的已识别项, 是为 RDNA3.5 重写 IQ2_XS/IQ3_XXS 的反量化指令序列。
9. 总结
数字上的结果: batch-1 decode 从 14.70 到 16.71 tok/s (+13.7%), 全程输出逐字节一致; 每 token 的 kernel launch 从约 5,200 降到约 3,500; 叠加 DSpark 投机解码后 20.9-23.8 tok/s。
但比数字更想传递的是方法 -- 如果你也想在自己的机器上做这样一场优化, 带走(告诉你的AI)这六条就够了:
- 先算 roofline, 再动手。 8 GiB/token ÷ 218 GB/s = 44.5 ms 的流式地板、22.5 tok/s 的硬上限。有了这个数, 每项优化值多少毫秒、总共还剩多少空间, 全部可以预算; 没有它, 优化就是盲人摸象。
- profiler 三件套各司其职。 kernel trace 记账 (谁慢), PMC 计数器定罪 (为什么慢), 图 dump 找结构 (在哪融合)。三者交叉验证 -- 单独用任何一个都会被骗 (第 4 章的三个例子各演示了一次)。
- 只认在线 A/B。 孤立 benchmark 会撒谎: 快 42% 的 kernel 让整机退步 2.5%。稳态 tok/s 是唯一判据, 无论孤立数字多漂亮。
- 零精度损失不是束缚, 是资产。 逐字节一致让每次 A/B 自带正确性证明, 任何数值回归无处藏身; 守住它只需要
__fmul_rn/__fadd_rn级别的细心。 - 阵亡名单与战果同等重要。 十来条被证伪的假设划定了这台机器的能力边界, 也替后来者省下同样的弯路 -- 写下来才算数。
- 知道何时收手。 roofline 判定剩余空间只剩 kernel 重写; 想再快, 出路是改变问题本身 -- 投机解码把 8 GiB 权重读取摊销到多个 token 上, 或者干脆少读字节 (重量化)。