一句话BLOCK_M / BLOCK_N 每个 BlockProgram负责计算的那一小块子矩阵的尺寸。它把一个大矩阵切成很多小 tile每个 tile 分给一个 Block 去算。以矩阵乘法为例C A × B A: [M, K] B: [K, N] C: [M, N] 比如 M1024, N1024, K512不可能一个 Block 算整个 1024×1024 的 C所以切块设 BLOCK_M 128, BLOCK_N 128 C [1024 × 1024] 被切成 BLOCK_N128 ├────┤ ┌────┬────┬────┬────┬────┬────┬────┬────┐ │ │ │ │ │ │ │ │ │ ↑ │(0,0)│(0,1)│(0,2)│(0,3)│(0,4)│(0,5)│(0,6)│(0,7)│ │ │ │ │ │ │ │ │ │ │ │ ├────┼────┼────┼────┼────┼────┼────┼────┤ │ │ │ │ │ │ │ │ │ │ │ │(1,0)│(1,1)│ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ ├────┼────┼────┼────┼────┼────┼────┼────┤ │ BLOCK_M128 │ │ │ │ │ │ │ │ │ │ │(2,0)│ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ │ ├────┼────┼────┼────┼────┼────┼────┼────┤ │ │ │ │ │ │ │ │ │ │ │ │... │ │ │ │ │ │ │ │ │ └────┴────┴────┴────┴────┴────┴────┴────┘ ↓ 共 (1024/128) × (1024/128) 8 × 8 64 个 tile → 启动 64 个 Block每个 Block 做什么Block (pid_m2, pid_n3) 的任务 计算 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): # K512, BLOCK_K32 → 迭代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_K32 ──→ ↑ │ │ 每次加载一小条 每次加载一小条 │ │ 迭代 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_start128, k : k32] 到 smem // 加载 B[k : k32, col_start : col_start128] 到 smem // acc A_smem × B_smem } // 写回 C[row_start:row_start128, col_start:col_start128] } // 启动 dim3 block(256); dim3 grid(N / BLOCK_N, M / BLOCK_M); // (8, 8) 64 个 block matmulgrid, block(A, B, C, M, N, K);Tritontriton.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*128128, pid_n*128 : pid_n*128128] acc tl.zeros((BLOCK_M, BLOCK_N), dtypetl.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_M128, BLOCK_N128, BLOCK_K32)在 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 大BLOCK 小Shared Memory 用量大可能超限小计算访存比高好低差Grid 大小并行度小可能填不满 SM大好Register 压力大每线程累加器多小典型值128 / 25616 / 32经验法则 BLOCK_M × BLOCK_N × sizeof(float) ≈ 每个线程的累加器大小 比如 BLOCK_M128, BLOCK_N128, 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 只解决一小块最后拼起来就是完整结果。