LLM 新人的压榨本地 DeepSeek V4 Flash 记

Table of Contents

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. 实验环境

顺带说明: 虽然是 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 这一侧:

  1. 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 的活, 时间全花在等内存上。
  2. KV 投影与 cache 追加。 x 乘 wkv 得到 576 维 latent, 对 64 维 pe tail 做 rope, 追加进 KV cache。写流量每层约 1 KB -- 可忽略。这里体现 MLA 的设计意图: cache 里存的是压缩 latent, 不是 64 份展开的 k/v。
  3. 打分与加权求和 (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 本身访存和计算都极小, 属于延迟型算子。
  4. 输出投影 (又是两级 GEMV)。 attention 结果先 de-rope 一个切片, 再过 wo_a/wo_b 两个矩阵回到 4096 维 -- 又是几十 MB 的权重流。

FFN (MoE) 这一侧:

  1. Router。 一个 [4096 x 256] 的 BF16 GEMV (2 MB) 输出 256 个分数, sigmoid 加 bias 后 argsort 取 top-6。排序和 gather 是典型的延迟型小 kernel: 数据量以 KB 计, 时间全是 launch 和访存延迟。
  2. 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 节)。
  3. 逐元素 SwiGLU。 silu(gate) x up 加 clamp, 作用在 [2048 x 6] 上 -- 数据量以 KB 计, 却要占三四个 kernel launch, 是"launch 税"的典型来源。
  4. 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. 如何找到优化点

方法论和修复本身同样重要:

  1. rocprofv3 提供 kernel 级真相。 完整 kernel trace, 按 kernel 间隙聚类重建出每个 token 的 pass, 得到逐 kernel 的时间账本: 约 5,000 次 launch 各花多少。结论: 约 770 个权重 GEMV 已在硬件速度运行, 另有数千个 1-4 us 的 elementwise kernel, 单个微不足道, 合计约 10 ms -- 纯粹的 launch/延迟税
  2. 一次性图 dump 提供计算图级真相。 一个环境变量控制的遍历器, 打印每个节点的形状和来源。它找出了状态机 (占 launch 的 20.7%)、占用率不足的矩阵乘、以及被 VIEW 节点打断而让现有 fusion 匹配器失效的算子链。
  3. 硬件计数器裁决假设。 MemUnitBusySQ_INSTS_VALU 区分出"内存已饱和" (Q8_0: 别动它) 和"指令受限" (IQ 量化: 需要重写 kernel, 而不是加缓存)。
  4. 只认在线 A/B。 整个项目最重要的教训: 孤立的 kernel benchmark 在这台机器上会撒谎。 一个快 42% 的 rms_norm 让服务端整体退步 2.5%; 把 codebook 放进 LDS 毫无收益还倒亏 1.2%。缓存常驻的微基准无法复现 DRAM 轮转和依赖链条件。每一项改动都通过重启真实服务、对比稳态 tok/s 来评判 -- 无论孤立测试看起来多漂亮, 在线退步一律回滚。
  5. 以逐字节一致为正确性标准。 相同 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)这六条就够了:

  1. 先算 roofline, 再动手。 8 GiB/token ÷ 218 GB/s = 44.5 ms 的流式地板、22.5 tok/s 的硬上限。有了这个数, 每项优化值多少毫秒、总共还剩多少空间, 全部可以预算; 没有它, 优化就是盲人摸象。
  2. profiler 三件套各司其职。 kernel trace 记账 (谁慢), PMC 计数器定罪 (为什么慢), 图 dump 找结构 (在哪融合)。三者交叉验证 -- 单独用任何一个都会被骗 (第 4 章的三个例子各演示了一次)。
  3. 只认在线 A/B。 孤立 benchmark 会撒谎: 快 42% 的 kernel 让整机退步 2.5%。稳态 tok/s 是唯一判据, 无论孤立数字多漂亮。
  4. 零精度损失不是束缚, 是资产。 逐字节一致让每次 A/B 自带正确性证明, 任何数值回归无处藏身; 守住它只需要 __fmul_rn/__fadd_rn 级别的细心。
  5. 阵亡名单与战果同等重要。 十来条被证伪的假设划定了这台机器的能力边界, 也替后来者省下同样的弯路 -- 写下来才算数。
  6. 知道何时收手。 roofline 判定剩余空间只剩 kernel 重写; 想再快, 出路是改变问题本身 -- 投机解码把 8 GiB 权重读取摊销到多个 token 上, 或者干脆少读字节 (重量化)。

Leave a Reply

Your email address will not be published. Required fields are marked *

Back to Top