深入 CPU 矩阵乘法微内核：AVX-512 指令级并行与缓存层次优化实战

2025 年，随着端侧 AI 推理需求爆发和 CPU-GPU 异构计算架构的成熟，优化 CPU 路径上的矩阵乘法（GEMM）微内核已成为推理引擎性能工程的关键一环。本文从缓存层次、指令调度、寄存器分块三个维度，深入剖析 AVX-512 微内核的设计空间，给出可直接落地的工程实践。

一、为什么 2025 年还要关心 CPU GEMM

GPU 在训练和大 batch 推理中占绝对优势，但以下场景使得 CPU GEMM 微内核优化依然至关重要：

• 端侧推理：Copilot+ PC、Apple Silicon Mac、ARM 手机上的本地 LLM 推理，CPU 路径承担 prefill 和小 batch decode
• 首 token 延迟：当 batch=1 时，GPU kernel 启动开销（~10μs）占比不可忽视，CPU 已预热缓存
• 稀疏/量化推理：INT4/INT8 反量化 + GEMM 混合操作在 CPU 上天然比 GPU 灵活
• 混合部署：大型推理集群中 CPU 承担 prefill 或 speculative decoding 的验证阶段

以 llama.cpp 为例，其 AVX-512 GEMM 路径在 Intel Sapphire Rapids 上可实现超过 200 TOPS（INT8），对于 7B 模型首 token 延迟贡献显著。

二、算法基础：为什么朴素的 O(n³) 实现如此缓慢

标准矩阵乘法 C += A × B 的内存访问复杂度为 O(n³)，而算术操作同样为 O(n³)，这意味着每字节的内存访问对应 O(1) 次运算——即 运算强度（arithmetic intensity）为常数。对于现代 CPU（skylake-X 以上），每个核心每周期可执行 64 字节加载 + 32 FLOP，运算强度需求至少为 32/64 = 0.5 FLOP/byte，朴素的逐元素访问远达不到此阈值。

核心优化思路：分块（tiling），将矩阵分解为适合各级缓存的小块，将运算强度提升至 O(block_size)。

2.1 缓存层次分块策略

L1 缓存（32KB）  → 微内核（micro-kernel）：MR × NR 的小面板
L2 缓存（1MB）   → 宏内核（macro-kernel）：KC × NC 的中面板
L3 缓存（共享）  → 打包（packing）：重新排列 A/B 矩阵以连续访问

关键参数（针对 AVX-512，乘法与加法指令延迟 4 周期，吞吐量 2/cycle）： - MR = 6 行（6 个 ZMM 寄存器用于累加） - NR = 32 列（2 个 ZMM 加载 B 行，每行 16 个 float） - KC = 256~512（A 的 K 维度分块，适配 L2）

2.2 寄存器压力分析

x86-64 AVX-512 提供 32 个 ZMM 寄存器（每个 512 bit = 16 个 float）。

在 6×16 的微内核中： - 6 个 ZMM 用于 C 的 6×16 累加器 - 1 个 ZMM 用于广播加载 A 的单个元素 - 2 个 ZMM 用于加载 B 的 16+16 个元素 - 1 个 ZMM 作为临时/地址计算 - 剩余 22 个 ZMM 可用于预取或展开

三、AVX-512 微内核：从零实现

下面展示一个针对 Intel Ice Lake / Sapphire Rapids 的 6×32 FP32 微内核。关键优化：使用 VEXTRACTF64x4 和 VFMADD231PS 最大化 FMA 端口利用率。

// avx512_gemm_microkernel.c
// 6 rows × 32 cols FP32 GEMM micro-kernel
// C[6×32] += A[6×KC] * B[KC×32]
// A: row-major, leading dimension KC
// B: row-major, leading dimension NR (=32)

void sgemm_microkernel_6x32(
    int KC,
    const float* __restrict__ A,  // 6 x KC
    const float* __restrict__ B,  // KC x 32
    float* __restrict__ C,        // 6 x 32
    int ldc)
{
    // 累加器：6 行 × 32 列，每行 2 个 ZMM（16 floats each）
    __m512 c00, c01;  // row 0, cols 0-15 and 16-31
    __m512 c10, c11;
    __m512 c20, c21;
    __m512 c30, c31;
    __m512 c40, c41;
    __m512 c50, c51;

    // 加载 C 初始值
    c00 = _mm512_loadu_ps(&C[0*ldc + 0]);
    c01 = _mm512_loadu_ps(&C[0*ldc + 16]);
    c10 = _mm512_loadu_ps(&C[1*ldc + 0]);
    c11 = _mm512_loadu_ps(&C[1*ldc + 16]);
    c20 = __mm512_loadu_ps(&C[2*ldc + 0]);
    c21 = _mm512_loadu_ps(&C[2*ldc + 16]);
    c30 = _mm512_loadu_ps(&C[3*ldc + 0]);
    c31 = _mm512_loadu_ps(&C[3*ldc + 16]);
    c40 = _mm512_loadu_ps(&C[4*ldc + 0]);
    c41 = _mm512_loadu_ps(&C[4*ldc + 16]);
    c50 = _mm512_loadu_ps(&C[5*ldc + 0]);
    c51 = _mm512_loadu_ps(&C[5*ldc + 16]);

    const float* a_ptr = A;
    const float* b_ptr = B;

    for (int k = 0; k < KC; k++) {
        // 加载 B 的一行（32 个 float）
        __m512 b0 = _mm512_loadu_ps(b_ptr);      // B[k][0:15]
        __m512 b1 = _mm512_loadu_ps(b_ptr + 16);  // B[k][16:31]

        // A 的当前列元素广播到 ZMM，然后 FMA
        __m512 a0 = _mm512_set1_ps(a_ptr[0]);  // A[k][0]
        c00 = _mm512_fmadd_ps(a0, b0, c00);
        c01 = _mm512_fmadd_ps(a0, b1, c01);

        __m512 a1 = _mm512_set1_ps(a_ptr[1]);  // A[k][1]
        c10 = _mm512_fmadd_ps(a1, b0, c10);
        c11 = _mm512_fmadd_ps(a1, b1, c11);

        __m512 a2 = _mm512_set1_ps(a_ptr[2]);
        c20 = _mm512_fmadd_ps(a2, b0, c20);
        c21 = _mm512_fmadd_ps(a2, b1, c21);

        __m512 a3 = _mm512_set1_ps(a_ptr[3]);
        c30 = _mm512_fmadd_ps(a3, b0, c30);
        c31 = _mm512_fmadd_ps(a3, b1, c31);

        __m512 a4 = _mm512_set1_ps(a_ptr[4]);
        c40 = _mm512_fmadd_ps(a4, b0, c40);
        c41 = _mm512_fmadd_ps(a4, b1, c41);

        __m512 a5 = _mm512_set1_ps(a_ptr[5]);
        c50 = _mm512_fmadd_ps(a5, b0, c50);
        c51 = _mm512_fmadd_ps(a5, b1, c51);

        a_ptr += 6;   // 下一列（A 按列存放的微面板）
        b_ptr += 32;  // B 的下一行
    }

    // 写回 C
    _mm512_storeu_ps(&C[0*ldc + 0], c00);
    _mm512_storeu_ps(&C[0*ldc + 16], c01);
    _mm512_storeu_ps(&C[1*ldc + 0], c10);
    _mm512_storeu_ps(&C[1*ldc + 16], c11);
    _mm512_storeu_ps(&C[2*ldc + 0], c20);
    _mm512_storeu_ps(&C[2*ldc + 16], c21);
    _mm512_storeu_ps(&C[3*ldc + 0], c30);
    _mm512_storeu_ps(&C[3*ldc + 16], c31);
    _mm512_storeu_ps(&C[4*ldc + 0], c40);
    _mm512_storeu_ps(&C[4*ldc + 16], c41);
    _mm512_storeu_ps(&C[5*ldc + 0], c50);
    _mm512_storeu_ps(&C[5*ldc + 16], c51);
}

四、性能分析：端口压力与瓶颈定位

4.1 FMA 端口利用

在 Ice Lake / SPR 上，每周期可执行 2 条 512-bit FMA。6×32 微内核每个 k 迭代执行 12 条 FMA（6 rows × 2 ZMM），理论峰值需要 6 周期完成。

指令序列分析： - 6 条 _mm512_fmadd_ps（行 0-5，组 1）+ 6 条（组 2） - FMA 延迟 4 周期，但吞吐量 2/cycle → 需要足够的独立 FMA 来填充延迟

每行 2 条独立 FMA → 可重叠 2 行（2×4 周期），6 行完全展开后可以隐藏延迟。实际瓶颈通常在 数据加载 而非计算。

4.2 内存带宽需求

每 k 迭代加载： - A: 6 × 4 = 24 bytes（广播展开后） - B: 32 × 4 × 2 = 256 bytes（两个 ZMM 加载） - 总计：280 bytes/iter

SPR 内存带宽 ~150 GB/s，每周期 ~48 bytes @ 3GHz。12 条 FMA 占 6 周期 = 可加载 288 bytes → 刚好匹配。

4.3 实测性能数据（Intel Xeon w9-3495X，56 核，DDR5-4800）

峰值：56×2×2×32×3.0 = 11520 GFLOPS，内存带宽限制约 1200 GFLOPS

关键观察：小尺寸矩阵受 L1/L2 延迟影响，超大尺寸接近内存带宽上限。

五、进阶优化：打包（Packing）与预取

5.1 为什么需要 Pack

朴素实现中 A 和 B 的访问模式存在跨步（stride），导致 TLB 失效和 Cache line 浪费。Pack 操作在 L3 级别将子矩阵复制为连续的微面板。

// 打包 A 的微面板：6 × KC → 连续内存
void pack_a_6xKC(const float* A, int lda, float* packed, int KC) {
    for (int k = 0; k < KC; k++) {
        for (int i = 0; i < 6; i++) {
            packed[i] = A[i * lda + k];
        }
        packed += 6;
    }
}

// 打包 B 的微面板：KC × 32 → 连续内存（逐行复制）
void pack_b_KCx32(const float* B, int ldb, float* packed, int KC) {
    for (int k = 0; k < KC; k++) {
        memcpy(packed, &B[k * ldb], 32 * sizeof(float));
        packed += 32;
    }
}

5.2 软件预取策略

在微内核中手动插入 _mm_prefetch 隐藏 DRAM 延迟：

for (int k = 0; k < KC; k++) {
    // 预取 B 后面第 8 行的数据（~128 L1 cache lines ahead）
    _mm_prefetch((const char*)(b_ptr + 8 * 32), _MM_HINT_T1);

    // 微内核计算...
}

5.3 KC 尺寸选择

KC 影响 L2 缓存驻留的 B 面板大小（KC × NR × 4 bytes）。

最优 KC 约 256-384，平衡了 L2 利用率和 TLB 覆盖。

六、INT8 量化推理实战路径

实际 AI 推理中 INT8 GEMM 远比 FP32 常用：

// INT8 GEMM with AVX-512 VNNI (Ice Lake+)
// C_int32[6×32] += A_int8[6×KC] * B_int8[KC×32]
void igemm_microkernel_6x32_vnni(
    int KC,
    const int8_t* __restrict__ A,
    const int8_t* __restrict__ B,
    int32_t* __restrict__ C,
    int ldc)
{
    __m512i c00, c01, c10, c11, c20, c21;
    __m512i c30, c31, c40, c41, c50, c51;

    // 加载 INT32 累加器...

    for (int k = 0; k < KC; k += 4) {
        // VPSRLVD + VPBROADCASTD + VPMADDWD + VPADDD
        // 每 4 个 INT8 元素对产生 1 个 INT32 累加
        __m512i b0 = _mm512_loadu_si512(b_ptr);
        __m512i b1 = _mm512_loadu_si512(b_ptr + 64);

        // 对每行 A 执行 DPBD（Dot Product of Byte）
        __m512i a0 = _mm512_set1_epi32(*(int32_t*)(a_ptr));
        c00 = _mm512_dpbusd_epi32(c00, a0, b0);
        c01 = _mm512_dpbusd_epi32(c01, a0, b1);
        // ... 重复 6 行

        a_ptr += 6 * 4;  // 6 行 × 4 字节
        b_ptr += 128;    // 2 × 64 字节
    }
}

INT8 VNNI 指令通过 VPDPBUSD 实现单周期 4 对 INT8 乘加，理论峰值是 FP32 的 4 倍——这也是 llama.cpp 中量化内核的性能基础。

七、与 BLIS/OpenBLAS 的关系

实际工程中不需要从头手写。BLIS（BLAS-like Library Instantiation System）和 OpenBLAS 提供了参数自动调优的 GEMM：

# BLIS 架构自动检测 + 微内核选择
export BLIS_ARCH_TYPE=haswell  # 自动检测 x86-64 微架构
gcc -O3 -march=native dgemm.c -lblis -o dgemm_blbench

BLIS 的设计哲学与本文完全一致： 1. 定义微内核接口（≤ 20 参数） 2. 针对每种架构手工优化微内核 3. 外层分块循环选择最优参数（KC, MC, NC）

对于自定义深度学习框架，可以直接调用 BLIS 微内核接口，上层自己实现张量内存管理。

八、工程落地的几点建议

1. 不要在 product 代码中硬编码微内核：选择 oneDNN、BLIS、或苹果 Accelerate 框架作为底层，换取多架构覆盖。

2. profile 先行：使用 perf stat -e cycles,instructions,cache-misses 定位瓶颈是计算受限还是内存受限。

3. 分块参数自适应：不同型号 CPU 的最优 KC 差异可达 2×，应建立参数数据库或运行时选择。

4. 警惕 AVX-512 频降（clock throttling）：在混合 256-bit/512-bit 负载下可能触发降频，需要实测确定净性能增益。

5. 与 AMX（Advanced Matrix Extensions）协同：Sapphire Rapids 引入的 AMX 指令可在 TILE 寄存器上执行 2D 矩阵乘，INT8 峰值再翻倍，值得在推理路径优先采用。

AVX-512 微内核是连接算法理论与硅片性能的关键一层。理解其设计空间，不仅有助于高性能计算开发，也能帮助我们在推理系统选型时做出正确的技术决策：知道何时 CPU 路径何时值得优化、何时该卸载到 GPU 或 NPU，本身就是架构师的核心能力。