BLOCK_M, BLOCK_N 是干什么的

一句话

BLOCK_M / BLOCK_N = 每个 Block（Program）负责计算的那一小块子矩阵的尺寸。

它把一个大矩阵切成很多小 tile，每个 tile 分给一个 Block 去算。

以矩阵乘法为例

C = A × B A: [M, K] B: [K, N] C: [M, N] 比如 M=1024, N=1024, K=512

不可能一个 Block 算整个 1024×1024 的 C，所以切块：

设 BLOCK_M = 128, BLOCK_N = 128 C [1024 × 1024] 被切成： BLOCK_N=128 ├────┤ ┌────┬────┬────┬────┬────┬────┬────┬────┐ │ │ │ │ │ │ │ │ │ ↑ │(0,0)│(0,1)│(0,2)│(0,3)│(0,4)│(0,5)│(0,6)│(0,7)│ │ │ │ │ │ │ │ │ │ │ │ ├────┼────┼────┼────┼────┼────┼────┼────┤ │ │ │ │ │ │ │ │ │ │ │ │(1,0)│(1,1)│ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ ├────┼────┼────┼────┼────┼────┼────┼────┤ │ BLOCK_M=128 │ │ │ │ │ │ │ │ │ │ │(2,0)│ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ ├────┼────┼────┼────┼────┼────┼────┼────┤ │ │ │ │ │ │ │ │ │ │ │ │... │ │ │ │ │ │ │ │ │ └────┴────┴────┴────┴────┴────┴────┴────┘ ↓ 共 (1024/128) × (1024/128) = 8 × 8 = 64 个 tile → 启动 64 个 Block

每个 Block 做什么？

Block (pid_m=2, pid_n=3) 的任务： 计算 C 的第 2 行 tile、第 3 列 tile： C[256:384, 384:512] = A[256:384, :] × B[:, 384:512] ↑ 128行 ↑ 128列 但 K 方向也要分块（BLOCK_K），迭代累加： for k in range(0, K, BLOCK_K): # K=512, BLOCK_K=32 → 迭代16次 加载 A_tile [128 × 32] 到 smem 加载 B_tile [32 × 128] 到 smem C_tile += A_tile × B_tile # [128×32] × [32×128] = [128×128]

┌─────────────────────────────────────────────────┐ │ Block (2,3) 的工作： │ │ │ │ A [128 × 512] B [512 × 128] │ │ ┌──┬──┬──┬──┬──┐ ┌──┐ │ │ │32│32│32│32│..│ │ │ │ │ │ │ │ │ │ │ × │ │ = C_tile [128×128]│ │ │ │ │ │ │ │ │ │ │ │ └──┴──┴──┴──┴──┘ └──┘ │ │ ←── BLOCK_K=32 ──→ ↑ │ │ 每次加载一小条 每次加载一小条 │ │ 迭代 16 次累加 │ └─────────────────────────────────────────────────┘

对应到代码

CUDA

#define BLOCK_M 128 #define BLOCK_N 128 #define BLOCK_K 32 __global__ void matmul(float* A, float* B, float* C, int M, int N, int K) { // 我是哪个 tile？ int pid_m = blockIdx.y; int pid_n = blockIdx.x; // 我负责 C 的哪一块？ int row_start = pid_m * BLOCK_M; // 比如 256 int col_start = pid_n * BLOCK_N; // 比如 384 // 在 K 方向迭代 float acc[BLOCK_M][BLOCK_N] = {0}; // 每个线程负责 acc 的一小部分 for (int k = 0; k < K; k += BLOCK_K) { // 加载 A[row_start : row_start+128, k : k+32] 到 smem // 加载 B[k : k+32, col_start : col_start+128] 到 smem // acc += A_smem × B_smem } // 写回 C[row_start:row_start+128, col_start:col_start+128] } // 启动 dim3 block(256); dim3 grid(N / BLOCK_N, M / BLOCK_M); // (8, 8) = 64 个 block matmul<<<grid, block>>>(A, B, C, M, N, K);

Triton

@triton.jit def matmul_kernel( A_ptr, B_ptr, C_ptr, M, N, K, BLOCK_M: tl.constexpr, # 128 BLOCK_N: tl.constexpr, # 128 BLOCK_K: tl.constexpr, # 32 ): pid_m = tl.program_id(0) pid_n = tl.program_id(1) # 我负责 C 的 [pid_m*128 : pid_m*128+128, pid_n*128 : pid_n*128+128] acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for k in range(0, K, BLOCK_K): a = tl.load(A_ptr + ...) # [BLOCK_M, BLOCK_K] = [128, 32] b = tl.load(B_ptr + ...) # [BLOCK_K, BLOCK_N] = [32, 128] acc += tl.dot(a, b) # [128, 128] tl.store(C_ptr + ..., acc) # 启动 grid = (M // BLOCK_M, N // BLOCK_N) # (8, 8) matmul_kernel[grid](A, B, C, M, N, K, BLOCK_M=128, BLOCK_N=128, BLOCK_K=32)

在 Attention 中的含义

Attention: O = softmax(Q × K^T) × V Q: [seq_len, head_dim] 比如 [4096, 576] K: [seq_len, head_dim] 比如 [4096, 576] V: [seq_len, head_dim_v] 比如 [4096, 512] BLOCK_M = 64 → 每个 Block 处理 64 个 Query token BLOCK_N = 64 → 每次加载 64 个 KV token（= TOPK_BLOCK_SIZE） ┌──────────────────────────────────────────────────────────┐ │ Q [4096 × 576] │ │ ┌──┐ │ │ │64│ ← BLOCK_M: 这个 Block 负责的 64 个 query │ │ └──┘ │ │ │ │ × K^T [576 × 4096] │ │ ┌──┬──┬──┬──┬──┬──┐ │ │ │64│64│64│64│..│64│ ← BLOCK_N: 每次加载64个KV │ │ └──┴──┴──┴──┴──┴──┘ │ │ │ │ = Score [64 × 4096] → softmax → × V → O [64 × 512] │ │ │ │ 迭代 4096/64 = 64 次（或 TopK 后只迭代 32 次） │ └──────────────────────────────────────────────────────────┘

BLOCK_M / BLOCK_N 怎么选？

经验法则： BLOCK_M × BLOCK_N × sizeof(float) ≈ 每个线程的累加器大小 比如 BLOCK_M=128, BLOCK_N=128, 256 个线程： 每线程累加器 = 128×128 / 256 = 64 个 float = 64 个 register 加上其他变量，总共 ~128 registers/thread → 合理

总结

BLOCK_M = 每个 Block 在 M 方向（行/query）上处理多大 BLOCK_N = 每个 Block 在 N 方向（列/KV）上处理多大 BLOCK_K = 在 K 方向（reduction/内积）上每次加载多大 它们决定了： 1. Grid 大小 = (M/BLOCK_M) × (N/BLOCK_N) → 启动多少个 Block 2. Shared Memory 大小 = BLOCK_M×BLOCK_K + BLOCK_K×BLOCK_N 3. 每个线程的计算量 = (BLOCK_M × BLOCK_N) / num_threads 4. K 方向迭代次数 = K / BLOCK_K

本质上就是分治：大问题切小，每个 Block 只解决一小块，最后拼起来就是完整结果。