Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

点积内核与深K孪生

大多数内核家族每次只消费收缩维的一个深度步。每一步加载打包 A 面板的一列,广播打包 B 的一个元素,给每个累加器寄存器发出一条 FMA(或加宽乘加)。

两个 AVX-512 扩展打破了这个节奏。它们各自把好几个深度步折叠进一条指令。VNNI 的 vpdpbusd 把 4 个连续的 i8 深度步乘进 16 个 i32 通道中的每一个。AVX-512 BF16 的 vdpbf16ps 把 2 个连续的 bf16 深度步折叠进每个 f32 通道。

一条一次消费多个深度步的指令,要求这些深度步在内存中相邻。对浮点来说,这还会改变累加的舍入方式。所以 gemmkit 给点积内核单独建了内核家族和 ISA 令牌,而不是把它们藏在共享微内核的一个分支里。

本页依次讲三件事:承载点积内核的两道接缝、两个具体的点积内核本身,以及两个窄类型家族共享的深收缩路线。

为什么点积指令需要独立家族

点积内核与它的加宽(widen)兄弟恰好在两处不同。这两处差异分别落在引擎的两条不同扩展轴上。家族与令牌的划分见标量与内核家族

打包布局是家族的事。KernelFamily::pack_lhspack_rhs 不接受 ISA 参数,所以不同的交织方式只能由家族本身决定。这正是 Bf16DotGemmMixedGemm<bf16> 的兄弟家族、而不是它内部一个分支的原因。

内层循环是令牌的事。只有真正具备 vpdpbusdvdpbf16ps 的 CPU 才能执行它。所以这条指令藏在一个 KernelSimd 方法后面,只有具备点积能力的令牌才覆写这个方法。

具体的点积家族是:每指令折叠 4 个深度步的 IntGemmVnni(及其重量化变体 IntGemmVnniQ),以及每指令折叠 2 个深度步的 Bf16DotGemm。再加上本页末尾要讲的两个 f32 输出孪生家族,凑齐这一整套。

DEPTH_MULTIPLE 与 k 组打包

每指令折叠 Q 个深度步的家族声明 const DEPTH_MULTIPLE: usize = Q,默认值是 1gemmkit/src/kernel.rs 里写明的契约是这样运作的:家族的打包例程要写出 width * kc.next_multiple_of(Q) 个元素的面板,并把深度尾部补齐。驱动层按同样补齐后的深度为打包面板定步长,让两边保持锁步。对所有普通家族来说,DEPTH_MULTIPLE = 1,每一次 next_multiple_of 调用都会退化为恒等。

布局本身来自一个共享例程,gemmkit/src/pack.rs 里的 pack_kgroup_panels。它是交织下标算术的唯一权威来源。

普通的 pack_panels 按深度为主序存储面板。每个深度步存 width 个连续的先导元素:LHS 是 mr 行,RHS 是 nr 列。

pack_kgroup_panels 则把深度轴按 Q 个一组分组,让同一通道连续的 Q 个深度值变得连续。面板内第 g 组、第 i 个通道、组内第 t 个位置,落在偏移 g*width*Q + i*Q + t 处。

这正是一条点积指令要读的形状。一个 64 字节的 A 寄存器覆盖 LANES 行乘 Q 个连续深度元素。一个 B 组把某一列 Q 个连续深度值当作一次 32 位加载来广播。

这个共享打包器还有两个要点值得一提。它接受一个逐元素变换 xform:对 bf16 是恒等变换,对 VNNI 的 A 操作数是 +128 偏置。打包器把每个填充位置都填成 xform(0)。填充位置指超出块的先导位置,以及超出 kc 的深度位置。这让填充始终与活跃元素保持一致。

交织布局无法原地读取。所以每个点积家族都设置 FORCE_PACK_LHSFORCE_PACK_RHS,压过驱动层基于代价的打包决策。点积内核永远要付打包成本。这恰好是下文各个门槛所要对冲的开销。

pack.rs 测试里有一个字节级 oracle,用朴素写法重新实现该布局。它验证真实例程在宽度尾部、深度尾部、带步长源等各种情形下都逐位复现该布局。

在消费侧,KernelSimd::dot_accumulate 是各家族用来替代加宽 FMA 循环的接缝。它的默认实现是 unreachable!。只有具备点积能力的令牌会覆写它,也只有点积家族会调用它。

Avx512VnniAvx512Bf16 之所以是独立于 Avx512F 的令牌,是因为 #[target_feature] 按令牌生效。_mm512_dpbusd_epi32 需要 avx512vnni 的代码生成上下文。Avx512F::vectorize 只建立 avx512f 上下文,提供不了这个。

该方法接收真实、未补齐的 kc。它从补齐深度的面板中读取 ceil(kc / Q) 个指令组。一切符号或偏置修正都在内部完成,所以返回时累加器已经持有真实的 sum_k(A*B)

折叠之所以放在这个专用接缝上,而不是通用的 accumulate_tile 上,是有明文记录的原因的:折叠深度步会重塑累加的舍入方式,而 accumulate_tile 的契约不允许这样做。

i8 经 vpdpbusd

vpdpbusd 计算的是无符号乘有符号的点积。它的第一操作数取 u8,第二操作数取 i8。而 GEMM 要的是有符号乘有符号。

修正是代数式的,不是逐元素的。LHS 打包把每个字节偏移 +128,送进无符号域。这个变换就是 gemmkit/src/kernel/int.rs 里的 vnni_a_xform。它使用常量 VNNI_A_BIAS = 128,这个常量只在 gemmkit/src/simd.rs 里定义一次,让打包与修正永远不会失配。

sum_k((A+128)*B) = sum_k(A*B) + 128*sum_k(B)。所以内核通过减去逐列修正项 128 * sum_k(B[k][j]) 来还原真实乘积。Avx512Vnni::dot_accumulate 在向量循环之前,先用一小段标量代码对有符号的打包 B 面板求出这些列和,然后在最后把广播后的修正从每个累加器里减去。

填充在这套方案里也是配合好的。A 的填充是 xform(0) = 128,修正项恰好抵消它的贡献。B 的填充是 0,对乘积和列和都没有贡献。

i32 累加按环绕方式进行,而环绕加法在模 2^32 下满足结合律。所以把求和重新按四个一组分组,再加上偏置修正,与升序 k 的加宽求和逐位相等。IntGemmVnni 与加宽版的 IntGemm 在任何输入上都产生完全相同的输出。

ISA 的选择因此永远不会改变 i8 的结果。这比可复现性契约所要求的还要强。那份契约只承诺在固定机器和固定配置下结果可复现,并不承诺跨内核选择时逐位一致。

正是这种可以随时换内核的自由,被小规模并行回退门槛拿来利用。VNNI 的打包在两个操作数上都是强制的。在一个小型多线程问题上,这道打包壁垒有可能压过它本该节省下来的计算量。

触发这个回退需要同时满足三个条件。ISA 选择必须是自动的。并行度必须是 Rayon(n)n != 1。而且 m*n*k 必须低于 GEMMKIT_I8_VNNI_MIN_PAR_MNK,它的默认值是 768^3。三个条件同时成立时,dispatch/int.rs 才会把调用交给可原地读取的加宽内核。

串行运行和大规模并行运行都继续用 VNNI。强制设定 GEMMKIT_REQUIRE_ISA=avx512vnni 会完全关闭这道门槛,因为强制锁定就必须运行指定的那个内核。

预打包 RHS 路径同样绕过这道门槛,原因有二。其一,k 四元组交织的缓冲区只有 VNNI 家族能消费。其二,这道门槛本来要对冲的打包壁垒,已经在预打包阶段被摊销过一次了。VNNI 的 RHS 打包本来在每次调用中都是强制的,所以预打包在这里带来的收益,比任何一个能原地读取操作数的内核都大。

bf16 经 vdpbf16ps

Bf16DotGemm 是浮点这一侧的对应者。它的 DEPTH_MULTIPLE2。两个操作数都按 k 对交织打包,每一对存成一个 32 位的 __m512bh 元素。dot_accumulate 每个累加器、每个对步发出一条 vdpbf16ps

累加之后的一切都与 MixedGemm<bf16> 逐字共用,包括 alpha 折叠,以及加宽读、收窄写的 epilogue,都经由共享的 mixed_epilogue 辅助函数完成。这个家族保持 OUT_IS_ACC = false,让整个收缩在 f32 里累加、只向 bf16 舍入一次。

数值故事与 VNNI 有一处本质区别。vdpbf16ps 融合的二项点积,舍入方式与两条独立的加宽 FMA 不同。所以点积内核的结果与加宽路径只是容差相等,不是逐位相等。

这恰好就是引擎一致性标准所允许的:结果必须在固定的输入、环境与配置下可复现,不需要跨内核选择逐位一致。点积内核本身完全确定。串行、并行、预打包运行共用同一个内核和同一套打包布局,因此彼此逐位复现。

这条路径上没有任何规模门槛。只要 CPU 报告支持 avx512bf16,自动选择就会优先选 Bf16DotGemm,因为它相对普通加宽路径是结构性的胜出。和 VNNI 不同,这里没有小规模并行回退。

只有三条特殊路径改道是例外。gemv、small_mn、以及小 k 形状,都刻意留在 MixedGemm<bf16> 的加宽接缝上。微小或退化的输出折叠不了什么,点积打包的深度补齐在那里纯属亏损。i8 分发出于同样的理由,把自己的微小形状也改道到加宽内核。

深 K 问题

OUT_IS_ACC = false 用一种结构性代价换来单次舍入。驱动层运行 kc = k,用一个覆盖整个收缩的深度面板。这取代了每个同质家族原本都有的缓存模型 kc 切片(见分块与缓存模型)。

一次微块调用要读的 RHS 微面板于是有 nr * k * sizeof(N) 字节。这个微面板一旦超出 L2 缓存,一个列条带里 m/mr 次微块调用中的每一次,都要转而从 L3 或 DRAM 里把它流过来。更大的 mr * k LHS 微面板也一样要从那里流过来。这样形成的悬崖很陡峭。微面板还在 L2 里时吞吐接近峰值,一旦装不下就会下降。

gemmkit/src/dispatch/mixed.rs 里的接入门把这个微面板大小拿去和一个字节阈值比较:

#![allow(unused)]
fn main() {
let engage_deep_k = NR
    .checked_mul(t.k)
    .and_then(|x| x.checked_mul(core::mem::size_of::<N>()))
    .is_some_and(|bytes| bytes > crate::cache::deep_k_engage_bytes());
if engage_deep_k {
    run_deep_k_twin::<N, Fam::Twin, S, MR_REG, NR>(simd, &t, par, ws);
    return;
}
}

这个阈值就是 GEMMKIT_DEEP_KC_BYTES 旋钮,非零时原样采用。默认值 0 把阈值推导为每工作线程有效 L2 容量的一半。

用 L2 的一半,而不是整个 L2,是刻意的选择。如果把门槛设成整个 L2 大小,会接入得太晚,晚到微面板早已装不下了,也就错过了悬崖。用一半留出空间给工作集里的其他部分,能在还来得及避开悬崖的时候就接入孪生路线。

checked_mul 这条链条是失败即关闭的。一个广播操作数可能带着逻辑上荒谬的 k 通过校验。溢出的尺寸必须转而落到单面板上。单面板自己的打包尺寸检查会拒绝这个问题,而不是接入一个会对着这个 k 永远多切片下去的孪生。

f32 输出孪生

在门槛之上,分发根本不运行窄类型家族。一个小小的 DeepKTwin trait 把每个窄家族映射到它的 f32 输出孪生。MixedGemm<N> 映射到 MixedGemmF32<N>Bf16DotGemm 映射到 Bf16DotGemmF32。每个孪生唯一的类型变化是 Out = f32 = Acc

这一处变化让 OUT_IS_ACC 回到默认值 true。驱动层普通的多切片 K 分块随即原样适用,每个切片的面板重新变得 L2 常驻。这正是孪生存在的全部意义。

打包布局和累加循环都逐字沿用窄家族的。MixedGemmF32 复用 pack_panels 和共享的加宽 FMA 辅助函数。Bf16DotGemmF32 复用 pack_kgroup_panelsdot_accumulate。这些累加辅助函数只触及 KernelSimd 接缝的输入侧。所以它们产出的累加器,与窄家族本会算出的结果逐字节相同。

孪生以 alpha = 1beta = 0 运行,写入一块 m x n 列主序的 f32 暂存缓冲区。这块缓冲区取自一个专用的 Workspace。深 K 按定义就是大 k 场景,所以一块 m*n 的 f32 分配可以忽略不计。把它独立出来,还能让池化的打包工作区留给孪生驱动使用。

之后,一次向量化扫描计算 narrow(alpha*scratch + beta*widen(C))。这逐个操作复刻了 mixed_epilogue 的算术,包括同一个 store_out 收窄步骤。

让这条路线不止于近似的关键,在于切片是怎么衔接起来的。孪生的微内核通过第三个 KernelSimd<N, N, f32, f32> 接缝(gemmkit/src/kernel/mixed.rs 里的 twin_seed),从暂存缓冲区为累加器寄存器播种。在一个累加切片上,它把运行中的部分和加载进寄存器,继续升序 k 的累加链。它从不把某个切片单独从零求和、再事后相加。

f32 的存取往返是精确的。所以多切片之和恰好就是单面板之和,只是在切片边界处被切开而已。对 beta 属于 {0, 1} 的情形,深 K 结果与单面板结果逐字节相同。

对一般的 beta,结果就只能做到容差一致了。原因很平常。单面板在完整微块上用一条 FMA 融合 beta*C + AB,但在边缘微块上是不融合地组合这两项。没有哪一种扫描方式能同时匹配这两种情形。

串行运行与并行运行在任何情况下都保持逐位一致。孪生驱动的分块不依赖线程数,最后那次扫描又是逐元素的。

点积孪生还多一条对齐规则。驱动层把分块 kc 向上取整到 DEPTH_MULTIPLE 的倍数,这样内部的切片边界就永远不会切开一个 k 对。被切开的一对会在收缩中途补零,从而错误地重组融合点积。有了这条规则,只有最后的短尾会被补齐,与单面板情形完全一致。

有三条路线刻意保留单面板。它们是门槛之下的浅 k、融合 epilogue 的入口,以及预打包 RHS 路径。在预打包路径上,DEPTH_MULTIPLE > 1 的缓冲区要求整个收缩必须是同一个深度切片。驱动层用一个硬性 assert 强制这一点,因为违反它会静默地错位读取微面板。

这些等价性声明都有直接的测试。gemmkit/tests/deep_k_narrow.rs 拨动 GEMMKIT_DEEP_KC_BYTES:取值 1 会在任何 k 下强制走孪生路线,取值 usize::MAX 会强制走单面板。这个测试在宿主机选中的任意 ISA 上,检验 beta 属于 {0, 1} 时逐字节相等,一般 beta 时容差一致。调优旋钮与其余旋钮一起记录了这个旋钮。