X3 讲解：Metal GEMM，从 threadgroup tiling 到 simdgroup_matrix

① 先看全貌：两个版本

② matmul_tiled：C4 v1 的直译

drills_TILED（13–30 行）

uint2 t = thread_position_in_threadgroup.xy;
uint2 g = threadgroup_position_in_grid.xy;
uint row = g.y * TS + t.y, col = g.x * TS + t.x;
threadgroup float As[16][16], Bs[16][16];
for (uint k0 = 0; k0 < K; k0 += TS) {
    As[t.y][t.x] = (...) ? float(A[row * K + k0 + t.x]) : 0.0f;    // 越界填 0
    Bs[t.y][t.x] = (...) ? float(B[(k0 + t.y) * N + col]) : 0.0f;
    threadgroup_barrier(mem_flags::mem_threadgroup);                // ① 都写完了才能读
    for (uint k = 0; k < TS; ++k) acc += As[t.y][k] * Bs[k][t.x];
    threadgroup_barrier(mem_flags::mem_threadgroup);                // ② 都读完了才能覆盖
}

和 C4 的 matmul_smem_kernel 逐行对应：__shared__ → threadgroup，__syncthreads() → threadgroup_barrier，threadIdx → thread_position_in_threadgroup。两次 barrier 的理由也完全一样。

grid 是线程数（drills95 行）：grid = (⌈N/16⌉·16, ⌈M/16⌉·16, 1)，不是 C4 里的 (⌈N/32⌉, ⌈M/32⌉) 个 block。这里特意把线程数补齐到 16 的倍数，让每个 threadgroup 都是完整的 16×16：barrier 要求 threadgroup 里的所有线程都执行到，加载 tile 时也需要每个位置都有线程。

③ matmul_simd：4 个 simdgroup，各管四分之一

drills_SIMD（32–81 行）

32 × 32 的 C tile 被分成 4 个 16 × 16 的象限，每个 simdgroup 负责一个；每个象限又是 2 × 2 个 8 × 8 的块，对应 4 个累加器 c00、c01、c10、c11：

sg_r = sg / 2，sg_c = sg % 2（37 行）：simdgroup 0 在左上，1 在右上，2 在左下，3 在右下。每个格子是一个 8×8 的 simdgroup_float8x8 矩阵。

simdgroup_float8x8 c00 = make_filled_simdgroup_matrix<float, 8, 8>(0.0f);   // c01、c10、c11 同理
simdgroup_float8x8 a0, a1, b0, b1;

simdgroup_float8x8 是什么：一个 8 × 8 的 float 矩阵，它的 64 个元素分散存放在 simdgroup 的 32 个线程的寄存器里（每个线程 2 个）。具体哪个线程拿哪个元素由硬件决定，你不需要知道：只能通过 simdgroup_load、simdgroup_store 和 simdgroup_multiply_accumulate 整体操作它。这和 G1 讲的 layout 是同一个概念，只是这里的 layout 是硬件规定的。

④ 主循环：加载、四次 MMA

协作加载

drills51–58 行

for (uint e = tid; e < BM * BK; e += 128) {        // 32 × 16 = 512 个数，每线程 4 个
    uint r = e / BK, k = e % BK;
    As[e] = (...) ? float(A[(row0 + r) * K + k0 + k]) : 0.0f;
    uint kb = e / BN, c = e % BN;
    Bs[e] = (...) ? float(B[(k0 + kb) * N + col0 + c]) : 0.0f;
}

和 C4 v2 的 e 循环一样：128 个线程一起把 A 的 32×16 块和 B 的 16×32 块搬进 threadgroup 内存，快维度连续，越界填 0。这里 As 和 Bs 都是自然布局（As 是 [BM][BK]，Bs 是 [BK][BN]），不像 C4 v2 那样把 A 转置存放，因为 simdgroup_load 可以直接按 stride 读 8×8 块。

四次 multiply_accumulate

drills59–69 行

for (uint kk = 0; kk < BK; kk += 8) {                         // BK = 16，走两步
    simdgroup_load(a0, As + (sg_r * 16 + 0) * BK + kk, BK);   // 我这个象限的上 8 行
    simdgroup_load(a1, As + (sg_r * 16 + 8) * BK + kk, BK);   // 下 8 行
    simdgroup_load(b0, Bs + kk * BN + sg_c * 16 + 0, BN);     // 左 8 列
    simdgroup_load(b1, Bs + kk * BN + sg_c * 16 + 8, BN);     // 右 8 列
    simdgroup_multiply_accumulate(c00, a0, b0, c00);          // c00 += a0 @ b0
    simdgroup_multiply_accumulate(c01, a0, b1, c01);
    simdgroup_multiply_accumulate(c10, a1, b0, c10);
    simdgroup_multiply_accumulate(c11, a1, b1, c11);
}

simdgroup_load(矩阵, 指针, stride)：指针指向 8×8 块的左上角元素，stride 是“下一行比这一行往后多少个元素”。比如 a1 的左上角在 As 的第 sg_r·16 + 8 行第 kk 列，As 每行 BK = 16 个元素，所以指针是 As + (sg_r·16 + 8)·16 + kk，stride 是 16。

这是 C4 v2 外积的块版本：

题目页提到 MLX 的 steel GEMM 快 3 倍多，主要原因之一就是每个 simdgroup 持有更多累加器（更大的外积，复用度更高），这和 C4 里 TM、TN 从 4 加到 8 是同一个方向。题目页的实验建议正是把 2×2 改成 4×2。

⑤ 写回：为什么要经过 Cs

drills72–80 行

simdgroup_store(c00, Cs + (sg_r * 16 + 0) * BN + sg_c * 16 + 0, BN);   // 4 个累加器先存进 threadgroup 内存
...
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint e = tid; e < BM * BN; e += 128) {
    uint r = e / BN, c = e % BN;
    if (row0 + r < M && col0 + c < N) C[(row0 + r) * N + col0 + c] = T(Cs[e]);   // 逐元素、带 mask 写回
}

simdgroup_store 一次写一整个 8×8 块，没有 mask。在矩阵边缘，一个 8×8 块可能只有一部分落在 C 里面，直接存到 C 就会越界写。所以先把整个 32×32 的结果存进 threadgroup 内存 Cs（这里永远不会越界），再由 128 个线程逐个元素地检查边界、写回 C，顺便做 T(...) 的类型转换。

这多用了 32 × 32 × 4 = 4 KB 的 threadgroup 内存和一次 barrier。M、N 都是 32 的倍数时其实可以直接 store 到 C；MLX 的 steel GEMM 就区分了这两种情况：mlxstore_result 和带边界检查的 mlxstore_result_safe。

⑥ grid 和 threadgroup

drillsmatmul_simd host 端（100–105 行）

grid = (((N + 31) // 32) * 128, (M + 31) // 32, 1)       # 线程数
threadgroup = (128, 1, 1)

x 方向：每个 32 列的 tile 需要 128 个线程，所以 x 方向的线程数是 tile 数 × 128，组成 ⌈N/32⌉ 个 threadgroup。y 方向：threadgroup 的 y 大小是 1，所以 ⌈M/32⌉ 个线程就是 ⌈M/32⌉ 个 threadgroup。kernel 里用 threadgroup_position_in_grid.xy 拿到 tile 坐标。

n = 2048 时：x 方向 64 × 128 = 8192 个线程（64 个 threadgroup），y 方向 64 个，一共 4096 个 threadgroup。

⑦ 对照 MLX 的 steel GEMM

• mlxBlockMMA：每个 simdgroup 持有 mlxTM × TN 个 8×8 片段，相当于这里的 2×2 推广成可配置的大小；底层调用的同样是 mlxsimdgroup_multiply_accumulate，片段类型是 mlxmetal::simdgroup_matrix。
• mlxBlockLoader：协作加载，每个线程一次读 mlxn_reads 个连续元素（向量化），对应这里的 e 循环。
• mlxGEMMKernel：把 loader 和 mma 组合起来的主循环。
• Mac 上跑 LLM 的主力其实是量化 matmul：decode 时 mlxqmv_fast_impl（矩阵乘向量，memory-bound），prefill 时 mlxqmm_t_impl；反量化在 mlxdequantize 里和乘加融合。

⑧ 常见坑与自测

• grid 用 threadgroup 数：X1 讲过的语义问题，GEMM 里尤其容易错，因为 x 和 y 两个方向的换算还不一样（x 要乘 128，y 不用）。
• simdgroup_load 的 stride 写错：As 的行宽是 BK，Bs 的行宽是 BN，写反了读到的是错位的块。
• 直接 simdgroup_store 到 C：边缘 tile 越界写。
• 忘了 #include <metal_simdgroup_matrix>（drillsheader，85 行）：编译找不到 simdgroup_float8x8。

源码与资料索引

本仓库：drillssolutions/x3_metal_matmul.py · drillsmine/x3_metal_matmul.py · drillstests/test_x3_metal_matmul.py · drillscontent/x3.md

MLX steel：mlxmma.h · mlxloader.h · mlxgemm.h · mlxquantized.h · mlxsdpa_vector.h

文档：Metal Shading Language Specification（SIMD-group matrix functions 一节）· MLX · Custom Metal Kernels