DeepSeek V4.1 Flash 工程实践(二):从权重装配到八卡文本链路

34 阅读17分钟

上一篇讲了 Engram:把训练好的记忆放进主机内存,再用查表、AVX-512 BF16 和 CPU/GPU 重叠,让它参与推理。这一篇回到 GPU 主干。

把模型分成八段,只完成了设备分工;要得到正确文本,还得把模型的状态接完整。 DeepSeek V4.1 Flash 的压缩 KV 会被多层共享,稀疏选择会被后续层复用,mHC 的混合系数还要跨子层传递。设备边界切在这些依赖中间,不能让它们随局部变量一起消失。

我们在八张 AMD Radeon Pro W7900D 上接入官方 safetensors 权重,从权重解释、ROCm 算子、会话状态一路修到输出 head。下面按这条真实路径展开。本文记录的是接入与后续修复,不把今天仍在推进的其他功能混进当时的验收结论。

先确定一轮推理到底要经过什么

这次语言主干有 40 层、5,120 维 hidden,mHC 使用四路残差表示。八卡入口按连续完整层分段,layer_ends 是每段的排他上界:

layer_ends: [5, 10, 15, 20, 25, 30, 35, 40]
GPU负责的主干层这一段遇到的特殊状态
0L0–L4L1 注入 Engram;L2 发布压缩 KV 与 selection
1L5–L9先继续使用 L2 的状态,L8 发布新的一组
2L10–L14继续使用 L8 的状态;L14 注入 Engram 并发布新状态
3L15–L19使用 L14 的共享压缩 KV 与 selection
4L20–L24L20 发布后半段共享 KV;L24 更新 selection
5L25–L29使用 L20 的 KV;L28 更新 selection
6L30–L34使用 L20 的 KV;L32 更新 selection
7L35–L39使用 L20 的 KV;L36 更新 selection,之后完成主干末层

层号从零开始。这里的八卡是连续层流水线:每一层的 attention、router 和专家计算由负责它的设备完成。表中每卡五层,描述的是层的放置,不表示八张卡在同一时刻处理同一个 token 的八个独立片段。

关闭投机解码时,一轮完整路径可以概括为:

文本 → tokenizer / 对话编码 → embedding
     → GPU0:L0–L4 → GPU1:L5–L9 → … → GPU7:L35–L39
     → mHC 最终折叠 → final norm → LM head → 采样
     → 下一个 token 回到入口

prefill 可以把多个输入块排进流水线;普通单路 decode 的下一 token 则依赖这一轮末尾的采样结果。因此,八卡提供了容量与分工,但吞吐还取决于依赖链、块大小和各段耗时。

第一步:把文件里的张量装成模型需要的权重

权重装配首先回答三个问题:张量叫什么、逻辑形状是什么、字节该怎样解释。

Safetensors 与 GGUF 解决的不是同一层问题

Safetensors 更接近一个安全、简单的张量容器。它记录张量名、dtype、shape 和对应的字节区间,支持单文件或由 index JSON 管理的多分片权重。它不执行反序列化代码,适合保存官方训练权重,也方便按张量或按行读取。但“一个张量叫 layers.2.attn.wq_a.weight 时应当怎样参与 DeepSeek V4.1 推理”,仍由模型配置和实现解释。

GGUF 除了张量目录,还通常携带模型架构、上下文参数、tokenizer 等推理元数据,并定义 GGML 生态中的量化类型。它更像面向分发和本地推理的模型包:一个文件可以同时告诉引擎“这是什么模型”和“这些 packed bytes 使用哪种量化格式”。代价是转换过程已经决定了张量命名、量化方案和元数据约定;第三方 GGUF 是否完整表达新架构,仍需逐项校验。

二者的区别不能简化成“Safetensors 是高精度,GGUF 是低精度”。Safetensors 可以保存 FP8、MXFP4 和 BF16 等不同 dtype;GGUF 也可以包含 F32、F16、BF16 或多种量化张量。真正的区别在于容器契约、元数据约定和量化表示方式。

维度SafetensorsGGUF
主要角色通用、安全的张量存储面向推理的模型与量化容器
模型结构通常结合外部 config 与模型代码解释通常由文件 metadata 描述
tokenizer通常是外部文件通常可随 GGUF metadata 携带
常见组织官方 checkpoint,多文件分片单文件或按体积分片的发布包
zLLM 读取方式按名字、shape、dtype 读取张量或指定行解析 metadata、张量目录和 GGML 量化类型,按需取得矩阵

zLLM 的“格式无关”发生在模型编排层

zLLM 没有要求 Safetensors 和 GGUF 在字节层长成同一种样子。实现分成四层:

文件容器
  SafetensorsStore / GgufReader
        ↓
模型权重装配
  张量命名、shape 校验、逻辑权重角色
        ↓
统一的准备后权重与 Backend capability
  LinearWeight / PreparedLayer / ExpertSource
        ↓
DeepSeek V4.1 runtime
  attention → shared KV / selection → MoE → mHC

容器层只负责 metadata、张量索引和原始字节访问。Safetensors 路径将 BlockFP8、MXFP8、MXFP4 和 dense 张量装配成带有明确逻辑形状的模型权重;GGUF 路径解析自身量化类型,准备成同一后端可以消费的线性权重和专家来源。

两条路径在“准备权重”之前可以不同。例如当前 DeepSeek 实现分别有 Safetensors 与 GGUF 的 layer preparation,因为它们的张量名、融合方式和量化切片并不相同;准备之后则进入同一份 DeepSeekV4PreparedLayer 和同一套 layer runtime。attention、MoE、KV 生命周期与 mHC 顺序不会因为文件扩展名再写一遍。

因此,这里的格式无关不是“任何 GGUF 都能自动替换官方 checkpoint”,而是:模型算法不绑定文件容器,容器差异在读取与权重装配边界收敛。 新格式仍要提供完整的架构 metadata、张量映射和后端支持;缺少 Engram、mHC、共享压缩 KV 或对应量化 kernel 时,引擎应明确拒绝,而不是带着错误假设继续生成。

zLLM 的权重层会探测顶层命名空间。例如当前 loader 在发现 language_model.embed.weight 时采用 language_model. 前缀,否则使用无前缀名称。这样能兼容实际遇到的不同 checkpoint 组织方式;前缀来自张量目录,不能凭下载目录名猜测。

同样,文件里“都是量化权重”也不足以决定执行方式:

权重形态装配时必须知道什么
BlockFP8数据矩阵与二维 scale 网格,以及行、列两个方向的分块大小
MXFP8每行沿输入维度分组的 codes 与 scale;不能当成二维分块
MXFP4 专家packed 数据、分组 scale、专家编号,以及 gate/up/down 的逻辑形状
未量化张量dtype、shape,以及它在 norm、混合或投影中的用途

专家权重还存在独立存储与合并存储两种组织。当前实现既能读取逐 expert 命名的矩阵,也能从合并的三维专家张量里按 expert 对应的行区间取出数据。对 FFN 而言,gate/up 是 [intermediate, hidden],down 是 [hidden, intermediate];装配时要同时保持 packed 列与逻辑列之间的关系。

这些差异在权重层处理。模型编排拿到的是可解释的矩阵与专家来源,随后由后端选择相应执行路径。

第二步:FP8 的错误,藏在 scale 的二维布局里

接入中一次很具体的修复,是官方 V4.1 主干线性权重的 32×32 BlockFP8。官方推理实现明确使用这套分块。官方模型实现

对一个 [M, N] 矩阵,二维 scale 网格的形状为:

[ceil(M / 32), ceil(N / 32)]

元素 (r, c) 使用的 scale 来自 (r / 32, c / 32),这里取整数商。行、列两个方向都决定选哪一个 scale。行独立的 MXFP8 则沿每行的输入维度分组,即使同样写着“32”,也不是同一种寻址。

我们为这次修复加入了一个很小的测试:64×64 的 codes 全部编码同一个值,四个 32×32 块分别使用对应 1、2、4、8 倍的 scale。解码后,四个象限就应分别得到 1、2、4、8。这个测试直接检查行方向和列方向是否都跨到了正确的 scale。

还有一个细节:小矩阵的 scale shape 可能无法区分 32×32 与旧的 128×128 布局。因此 loader 要结合模型版本解释格式,不能只看 scale 张量有几个元素。

修复也分成两层。80b8d3e5 校正权重装配,07cf53f1 补齐 ROCm 的 32×32 BlockFP8 路径。CPU 能正确解码,并不自动意味着 GPU kernel 使用了相同的 scale 索引。

第三步:分清本层 KV、共享压缩 KV 和 selection

普通的“每层一个 KV cache”直觉,在这里需要展开成几份不同的数据。

滑窗 KV 属于每一层,用于该层的局部历史。当前配置的窗口为 128。

压缩 KV 由指定源层产生,后续一组层共享使用。zLLM 的 V4.1 映射是:

发布源层使用范围ratio
L2L2–L72
L8L8–L132
L14L14–L192
L20L20–L391

L0、L1 不走这条压缩分支。L20 之后的 ratio=1 仍表示使用由 L20 产生的 1:1 表示,不能解释成“关闭压缩分支,让各层自己生成一份”。压缩比描述粒度,源层映射描述数据来自哪里,两者都要保留。

这可以看成 DSA(DeepSeek Sparse Attention)的进一步工程化。DSA 首先解决的是历史维度的稀疏:面对不断增长的上下文,indexer 不再让当前 query 对全部历史 KV 做完整 attention,而是先为历史建立较轻的索引表示,计算相关性并选出 Top-K 位置,随后只读取这些位置的 KV 完成 attention。上下文越长,被跳过的无关历史越多,稀疏带来的收益越明显。

V4.1 又把稀疏从时间轴推进到了网络深度:相邻层不必各自生成一份压缩历史,也不必每层都重新做一次完整选择。一个 KV source 层负责发布这一组层共用的压缩 KV 与 index key;index source 层按设定的节奏重新计算 selection;其他层直接复用最近的结果。它形成两种正交的稀疏:

维度解决的问题运行时行为
基于历史的稀疏当前 query 没必要访问所有历史 tokenindexer 从历史中选 Top-K,只对入选位置做 attention
跨层的稀疏相邻层没必要重复构造近似的历史表示与选择结果源层发布压缩 KV、index key 或 selection,组内后续层共享

这种共享仍然保留每一层自己的 attention 计算:各层用自己的 hidden state 产生 query,也维护自己的滑窗 KV;跨层复用的是对远端历史的压缩表示,以及在有效层范围内沿用的历史位置。计算保留层间差异,同时把重复度高的历史整理和检索工作从每层一次,改为按组产生、按需刷新。

从 L20–L39 最容易看出这种解耦:压缩 KV 始终来自 L20,但 selection 会在 L20、L24、L28、L32、L36 更新。换句话说,历史数据由谁提供何时重新判断哪些历史值得读取是两条独立时间线。较稳定、体量较大的压缩历史可以跨更多层复用,较轻的选择结果则可以更频繁刷新。

这与我们在 GLM 5.3 路径中处理 DSA 的思路异曲同工。GLM 5.3 的完整 indexer 层先建立历史索引并产生 selection,后续 IndexShare 层复用最近一次有效选择;连续的 MTP 步骤也会在状态完备时复用 selection。两个模型的具体权重和层映射不同,但工程目标一致:先在长历史中找出值得计算的位置,再让后续计算复用这次检索,而不是在每一层重复扫描整段历史。

selection 是 indexer 为各 query 选出的历史位置。当前配置在 L2、L8、L14、L20、L24、L28、L32、L36 更新选择,其余层复用最近发布的结果。它与压缩 KV 也不是同一份状态:L24 可以更新选择,同时继续读取 L20 的压缩历史。

在 runtime 中,把这三份状态分开,才能准确表达“本层写什么、读谁的历史、是否重新选择”。

第四步:跨卡时,状态不能在 stage 入口清零

第一个典型例子是 GPU0 到 GPU1。

GPU0 执行到 L2 时已经产生共享历史与 selection。GPU1 从 L5 开始,L5–L7 还需要它们,直到 L8 才发布新的一组。如果每个 stage 入口都创建一份全新的局部共享状态,L5 接收到 hidden,却找不到应当沿用的选择结果。

cc0be412 的修复让 selection 成为流水 work item 的一部分:进入 stage 时接回上一段发布的状态,完成本段后继续传给下一段。压缩 KV 的源层则通过全模型的层映射查找,不再依赖“本 stage 内最近见过哪个 source”。

这两种状态采用不同的管理方式:

数据在链路中的管理方式
hidden随当前输入块流过各个 stage
selection随输入块延续,遇到新的 indexer 层时更新
压缩 KV 历史保存在会话 cache 表中,消费者按源层映射访问
层本地滑窗 KV由负责该层的会话状态维护

这里的“随输入块延续”是逻辑依赖,并不意味着每次都要把数据下载到 CPU。后续优化又进一步改变了它的物理位置。

第五步:同一会话跨卡共享,不同会话彼此隔离

跨 stage 的来源修正后,还出现了另一个生命周期问题:每个 stage 各自从模板 fork 一张全层 cache 表。

这样得到的每一张表都可能合法,但它们不是同一张表。GPU0 把 L2 的历史写进自己的实例,GPU1 去另一份实例里读取 L2,就会读到空状态。

307c5e85 将建立会话改为统一入口:先为这个会话 fork 一份全层 cache 表,再把同一份共享引用交给八个 stage。 Engram 的序列状态也在会话层统一创建,由同一会话内的 stage 共用。

这同时界定了共享范围。模型和只读权重可以跨会话复用;会被追加和修改的序列历史必须属于具体会话。重置会话时,不仅要处理各层局部 cache,也要清理共享历史与 Engram 的 token 状态。

“共享”能成立的前提,是明确共享的是哪份数据、由哪个会话拥有,以及写入何时可见。

第六步:mHC 的 pre-mix 也要跨层、跨卡延续

权重和 KV 都接上之后,还不能只检查 hidden 的形状。

V4.1 的 mHC 维护四路残差表示,子层输入需要按系数折叠。我们早期沿用了较独立的子层处理方式,后续根据实际语义改成连续传递 pre-mix:

上一层 FFN 发布的 pre
  → 本层 attention 的输入折叠

本层 attention 发布的 pre
  → 本层 FFN 的输入折叠

本层 FFN 发布的 pre
  → 下一层 attention,或最终输出 head 的折叠

首层没有上游输入时,使用实现规定的初始处理。到最后一层,则使用最后 FFN 发布的 pre 折叠展开态,再进入 final norm 与 LM head。

8919f797 因此让层执行返回 hidden 和待传递的 pre,并将 pre 一起带过 stage 边界。比如 GPU0 结束 L4 后,GPU1 的 L5 不只需要 L4 的 hidden,也需要对应的 pre。

这种错误很难只靠 shape 检查发现:来自不同子层的系数可能形状相同,却代表不同执行时刻的状态。数值链路需要核对“谁产生、谁消费”,不能只核对“能否相乘”。

正确之后:把共享状态留在需要它的设备上

最初跑通时,共享状态存在并能到达消费者,是第一目标。随后我们才优化它的物理访问。

一项是压缩 KV 的本地镜像。最初后续卡可以读取源卡上的历史,但长 prefill 会反复扫描同一份远端数据。后续代码在消费者侧维护本地副本,在布局和已有前缀不变时只复制新增行;扩容、布局变化或前缀不一致时重新建立需要的内容。这样增加了一部分本地显存占用,换取计算时的本卡访问。

另一项是selection 留在 GPU。早期路径先把索引下载成 host 向量,再上传给后续层。优化后的 ROCm selection 保存设备 buffer 引用,同卡复用,跨卡使用有序 P2P 传递,避免每个消费层反复经过 CPU。

两项都不能只把指针换过去。异步任务在消费完成前必须持有源数据,生产、复制和消费之间必须有正确的顺序;临时分配还要转换为能够安全跨 stage 使用的生命周期。逻辑共享与物理放置分开,才有空间从“能运行”走到“运行得快”。

如何确认这条链真的成立

这次验证按不同层次进行,没有用某一个“能出字”的结果替代全部检查:

  • 格式级:用独立 scale 象限测试检查 BlockFP8 的二维解释。
  • 模型编排级:小尺寸全前向测试覆盖压缩共享、候选选择与 Engram,执行 prefill 和 decode,并检查 Engram 开关对结果的影响。
  • 真实权重级:装配官方 checkpoint,执行完整 40 层、最终折叠、head 与采样。
  • 优化回归级:固定输入与生成预算,比较完整输出文本哈希,并分别记录首字延迟与生成速度。

还有一条需要公开说明的边界:官方模型包含候选块粗筛后的两级选择。在本文核对的 ROCm 接入配置里,candidate_source_layer=Nonecandidate_topk_blocks=0,使用单级 Top-K;两级候选逻辑已有 CPU reference,但不能据此把 ROCm 路径写成完整复现了官方所有选择步骤。两种选择策略也不能未经验证就宣称对所有输入等价。

第一篇给出的 50K 首字延迟约 26.06 秒、单路约 20.45 token/s,来自后续优化后的固定文本测试,关闭 DSpark。它们不是某个 FP8 修复或状态修复单独带来的性能收益,也不是这次写作重新跑出的数据。

从文件到八卡文本生成,最重要的工作是保持同一个模型在不同位置上的连续性:字节解释相同、历史来源明确、会话边界清楚、混合状态顺序正确。性能优化随后改变数据在哪里、什么时候搬运,而不能丢掉这些关系。

下一篇继续讲图像:如何从预处理、视觉塔和 aligner 的中间张量找到差异,再把图文输入接进这条语言主干。

工程依据:zLLM 的 V4.1 接入记录,以及 80b8d3e507cf53f1cc0be412307c5e858919f797e9b539e4 的实际改动。模型配置参考官方推理配置

原文链接:zhuai.tech/blog/deepse…