第 8 课：二维 Tile 与转置

本课目标与完成标准

学完后你应能：构造二维 pointer tensor，正确处理非方阵与双向 mask。

通过必须同时满足：

• 独立补齐本课唯一的代码填空题，并通过给定检查；
• 三个问答题均说明因果链，而不是只报术语；
• 能指出至少一个正确性边界和一个性能取舍；
• 总分不低于 8/10，且没有一票否决级概念错误。

前置关系

• 课程前置：Python、PyTorch 张量、CUDA 基本线程/内存概念
• 本课在路线中的作用：二维 tile 用广播后的行/列 offset 构造 pointer matrix；tl.trans 交换编译期维度。

核心心智模型

1. 它是什么，解决什么问题

矩阵转置（Matrix Transpose: $Y = X^T$）是最基础但对硬件内存层次极具杀伤力的内存重排算子。 在行优先（Row-Major）物理存储下，如果直接按标量或朴素方式转置，面临非此即彼的内存性能断崖：

• 若以行连续方式合并读取（Coalesced Read），则写回目标矩阵时将跨行跨列产生巨大的跨步写入（Strided Write），导致 L2 缓存行污染与极低的写带宽；
• 反之，若写回保持合并，则读取必定离散跨步。 在传统 CUDA 编程中，解决此问题的标准范式是使用带 +1 Padding 的共享内存（Shared Memory）充当片上数据交换中转站，极度繁琐且容易引发 Bank Conflict。 本课解决的核心工程问题是：在 Triton 中如何使用二维广播语法构建高维 Pointer Tensor 与二维谓词掩码，并借助原生编译原语 tl.trans 将复杂的片上寄存器/共享内存转置重排交给编译器自动优化，实现非方阵与任意步长下的安全合并读写。

2. 它如何工作

2D 分块转置与内存步长重塑架构如图 8-1 与图 8-2 所示：

2D Tile 坐标与指针轴交换

输入分块 (BLOCK_M, BLOCK_N)

转置操作 tl.trans(x) ➔ (BLOCK_N, BLOCK_M)

按列转置写回目标地址

图 8-1：二维分块转置与内存步长重塑架构

2D Program 划分 (pid_m, pid_n)

合并加载源矩阵分块

Triton 编译器自动在片上执行无冲突转置

合并写回目标转置矩阵

图 8-2：2D 分块转置全局合并读写流程

1. 二维网格分块映射： 发射网格为 grid = (cdiv(M, BM), cdiv(N, BN))。每个 Program 处理一个形状为 (BM, BN) 的二维 Tile。

2. 二维广播指针与双向掩码构造：

   • 行偏移向量 rm = pid_m * BM + tl.arange(0, BM)，形状为 (BM,)；
   • 列偏移向量 rn = pid_n * BN + tl.arange(0, BN)，形状为 (BN,)；
   • 通过广播语法 rm[:, None]（形状 (BM, 1)）与 rn[None, :]（形状 (1, BN)）合成形状为 (BM, BN) 的二维指针矩阵： $$\text{ptrs} = x + rm[:, \text{None}] \times s_{xm} + rn[\text{None}, :] \times s_{xn}$$
   • 二维输入掩码：mask = (rm[:, None] < M) & (rn[None, :] < N)，形状同样为 (BM, BN)。

3. 合并加载与片上向量转置（tl.trans）：

   • 执行 tile = tl.load(ptrs, mask=mask)。此时 tile 处于片上寄存器，形状为 (BM, BN)。
   • 调用 tl.trans(tile)，在编译期将寄存器块的逻辑维度从 (BM, BN) 交换为 (BN, BM)。底层 Triton 编译器（Triton-GPU Dialect）会自动在片上寄存器或利用 Bank 优化的共享内存中调度高效的数据重排，对开发者完全透明。

4. 交换指针与输出掩码写回：

   • 输出矩阵的目标形状为 $(N, M)$，目标地址指针通过将坐标轴反转构造： $$\text{out\_ptrs} = out + rn[:, \text{None}] \times s_{om} + rm[\text{None}, :] \times s_{on}$$ 其形状为 (BN, BM)。
   • 输出掩码维度同步反转：out_mask = (rn[:, None] < N) & (rm[None, :] < M)，形状为 (BN, BM)。
   • 执行 tl.store(out_ptrs, tl.trans(tile), mask=out_mask)，实现读写双方均处于行连续的最佳合并访存状态。

3. 正确性条件与常见误区

• 输入掩码与输出掩码的形状失配（Shape Mismatch）： 输入 Tile 的形状是 (BM, BN)，其掩码 mask 形状也是 (BM, BN)。 输出 Tile 由于被转置，形状变成了 (BN, BM)。如果直接将输入的 mask 传给 tl.store，当 $BM \neq BN$ 时，Triton 编译器会直接抛出张量形状不匹配编译错误；即使 $BM == BN$，若处理非方阵边界块（如 $M=31, N=65$），输入掩码中 $M$ 轴与 $N$ 轴的布尔值被错误复用，会导致边界处数据被严重切碎截断或写入越界垃圾。

• 只测方阵（Square Matrix）的虚假正确性： 若仅测试 $M=N$ 的方阵，输入步长与输出步长、行块与列块完全对称，会掩盖大量关于 $s_{xm}$ 与 $s_{xn}$、som 与 son 混淆的严重逻辑 Bug。必须使用长宽极度不对称的矩形矩阵（如 $31 \times 65$）进行边界压力验证。

• 非连续输入与输出连续契约： 虽然代码中输入地址使用了通用步长 x.stride(0) 和 x.stride(1)，但输出 out 必须分配连续内存，或者显式指定符合转置期望的目标步长。

4. 性能与工程取舍

• Tile 尺寸选择（32x32 vs 64x64）： $32 \times 32$（1024 元素）是一个通用且均衡的选择。对于 FP32，单 Tile 占用 4KB 显存数据； 若选择更大 Tile（如 $64 \times 64 = 4096$ 元素），能提升内存突发传输长度，但在 num_warps 不变时单个 Program 占用的寄存器变为 4 倍，可能导致 SM 上并发活动 Warp 数量骤降（Occupancy 降低）。

• 非方阵 Tile（$64 \times 16$）与对齐优化： 若已知某一维度的物理步长极大，选择不对称的 Tile（如 $BM=16, BN=64$）有时能更好地匹配内存控制器的 128-byte 突发访问，但增加了代码复杂性。

具体演示

以输入形状 $M = 31, N = 65$、分块大小 $BM = 32, BN = 32$ 为例：

• 网格大小：$\lceil 31 / 32 \rceil = 1$ 个 M 块，$\lceil 65 / 32 \rceil = 3$ 个 N 块，总共发射 $1 \times 3 = 3$ 个 2D Program。
• Program (0, 2) 处理输入尾块 $[0:32, 64:96]$（左闭右开）：
  • 行范围 $rm = [0..31]$，前 31 行有效，最后 1 行越界（$31 < 31$ 为 False）；
  • 列范围 $rn = [64..95]$，仅第 64 列有效，后 31 列越界（$rn < 65$ 仅 1 个 True）；
  • 输入掩码 mask 形状为 $(32, 32)$，仅左上角 $31 \times 1$ 的区域为 True；
  • 转置后 tl.trans(tile) 形状变为 $(32, 32)$；
  • 输出掩码 out_mask 形状为 $(32, 32)$，仅左上角 $1 \times 31$ 的区域为 True；
  • 写回到输出矩阵的 $[64, 0:31]$ 区域（第 64 行的前 31 列）。

请在阅读后先合上这一节，用自己的语言复述“输入状态 → 中间状态 → 输出状态”，再做练习。

实践任务：唯一代码填空题

补齐转置值。

规则：只能修改 TODO/______ 所在位置；不要删除断言或放宽误差。代码注释说明了每个边界条件。

import torch
import triton
import triton.language as tl

@triton.jit
def transpose_kernel(x, out, M: tl.constexpr, N: tl.constexpr,
                     sxm: tl.constexpr, sxn: tl.constexpr,
                     som: tl.constexpr, son: tl.constexpr,
                     BM: tl.constexpr, BN: tl.constexpr):
    rm = tl.program_id(0) * BM + tl.arange(0, BM)
    rn = tl.program_id(1) * BN + tl.arange(0, BN)
    ptrs = x + rm[:, None] * sxm + rn[None, :] * sxn
    mask = (rm[:, None] < M) & (rn[None, :] < N)
    tile = tl.load(ptrs, mask=mask)
    out_ptrs = out + rn[:, None] * som + rm[None, :] * son
    out_mask = (rn[:, None] < N) & (rm[None, :] < M)
    tl.store(out_ptrs, ______, mask=out_mask)  # TODO: 交换 tile 两维

def transpose(x):
    assert x.ndim == 2
    M, N = x.shape
    out = torch.empty((N, M), device=x.device, dtype=x.dtype)
    grid = (triton.cdiv(M, 32), triton.cdiv(N, 32))
    transpose_kernel[grid](x, out, M, N, x.stride(0), x.stride(1),
                           out.stride(0), out.stride(1), BM=32, BN=32)
    return out

for shape in ((2, 3), (31, 65), (128, 96)):
    x = torch.randn(shape, device="cuda")
    torch.testing.assert_close(transpose(x), x.T)

检查方法

在 CUDA/Triton 环境运行本单元格；断言覆盖规则尺寸和非规则尾块。首次 JIT 不计入性能。

提交时请给出：补齐后的代码、实际运行输出（环境不可用时注明“仅静态审查”）以及对失败用例的解释。

Q1

不要背定义：请从输入、状态变化和输出三个阶段解释“二维 Tile 与转置”的工作机制。

你的答案：

Q2

输出 mask 为何不能直接复用输入 mask？

你的答案：

Q3

对非连续输入，这个 stride-aware kernel 还需哪些 wrapper 检查？

你的答案：

评分与通过规则

• 代码 4 分：正常输入 2 分，边界输入 1 分，解释实现 1 分；
• Q1～Q3 各 2 分；
• 一票否决：结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。

需要提示时按四级机制请求：概念区域 → 具体方向 → 关键局部 → 完整答案。

Tip💡 点击展开参考答案与详细代码实现

参考答案（仅 answer 分支）

先完成题目再核对。即使代码一致，也要能解释关键步骤，并尝试更换一个输入规模。

import torch
import triton
import triton.language as tl

@triton.jit
def transpose_kernel(x, out, M: tl.constexpr, N: tl.constexpr,
                     sxm: tl.constexpr, sxn: tl.constexpr,
                     som: tl.constexpr, son: tl.constexpr,
                     BM: tl.constexpr, BN: tl.constexpr):
    rm = tl.program_id(0) * BM + tl.arange(0, BM)
    rn = tl.program_id(1) * BN + tl.arange(0, BN)
    ptrs = x + rm[:, None] * sxm + rn[None, :] * sxn
    mask = (rm[:, None] < M) & (rn[None, :] < N)
    tile = tl.load(ptrs, mask=mask)
    out_ptrs = out + rn[:, None] * som + rm[None, :] * son
    out_mask = (rn[:, None] < N) & (rm[None, :] < M)
    tl.store(out_ptrs, tl.trans(tile), mask=out_mask)

def transpose(x):
    assert x.ndim == 2
    M, N = x.shape
    out = torch.empty((N, M), device=x.device, dtype=x.dtype)
    grid = (triton.cdiv(M, 32), triton.cdiv(N, 32))
    transpose_kernel[grid](x, out, M, N, x.stride(0), x.stride(1),
                           out.stride(0), out.stride(1), BM=32, BN=32)
    return out

for shape in ((2, 3), (31, 65), (128, 96)):
    x = torch.randn(shape, device="cuda")
    torch.testing.assert_close(transpose(x), x.T)

Q1 参考答案

从三个阶段拆解其运行机制：

1. 输入阶段（2D 空间解构与广播指针生成）：Python 调度层以二维网格 (cdiv(M, BM), cdiv(N, BN)) 发射 Program。内核中通过广播切片构造出二维索引网格 rm[:, None] 与 rn[None, :]，结合源张量物理步长生成形状为 (BM, BN) 的二维物理指针矩阵 ptrs，并在输入掩码 (rm < M) & (rn < N) 保护下合并读取至寄存器块 tile。

2. 状态变化阶段（片上多维张量转置）：调用 tl.trans(tile)，在编译期将寄存器块的逻辑形状由 (BM, BN) 交换为 (BN, BM)。底层 Triton 编译器自动生成无 Bank Conflict 的高效数据交换指令序列，在片上寄存器文件或共享内存中极速完成几何维度的转置重排，整个过程无需任何全局内存访问。

3. 输出阶段（目标步长重构与反向掩码写回）：以反转后的坐标网格生成目标矩阵物理指针 rn[:, None] * som + rm[None, :] * son，其形状为 (BN, BM)。结合同步转置后的二维输出掩码 out_mask = (rn < N) & (rm < M)，通过 tl.store 以行合并连续写回全局显存，完成高带宽、无冲突的矩阵转置。

Q2 参考答案

输出掩码不能直接复用输入掩码，原因在于张量维度与几何边界的根本变化：

1. 张量形状（Shape）不兼容： 输入 Tile 的形状是 (BM, BN)，因此输入掩码 mask 的形状严格为 (BM, BN)； 而转置后写回的输出 Tile 形状为 (BN, BM)，tl.store 要求写回数据、输出指针与掩码的形状必须位级对齐。若 $BM \neq BN$，直接传入输入掩码会在编译期引发张量维度不匹配错误（Shape Mismatch）。

2. 边界保护逻辑翻转： 即使在方阵 Tile（$BM = BN$）下，矩阵本身的维度通常是非方阵（$M \neq N$）。输入掩码判定的是第 0 维是否小于 $M$、第 1 维是否小于 $N$；而写回目标矩阵时，行数变成了 $N$、列数变成了 $M$。若直接复用输入掩码，会导致原本针对行维度的边界约束被错误施加到列维度，导致尾部边界数据被错误裁剪或写入未分配的非法显存。

Q3 参考答案

虽然内核内部使用了通用步长计算（Stride-aware），但一个健壮的生产级 Wrapper 仍需施加以下严密检查：

1. 张量维度与设备检查：断言输入张量严格为二维（x.ndim == 2），且驻留在正确的 CUDA 设备上。

2. 步长与索引范围检查： PyTorch 不支持负步长（x[::-1] 会直接报错），所以不必处理负步长；输入是 expand 产生的 0 步长或其他重叠视图时，本 kernel 只读输入，结果仍然正确。需要检查的是：输出 out 不能与输入共享存储（原地转置非方阵会读到已被覆盖的数据）；当 rm * sxm 等偏移可能超过 $2^{31}-1$ 时，int32 索引会溢出，需要把偏移转为 int64（如 rm.to(tl.int64)）或在 wrapper 中拒绝过大的张量。

3. 输出张量属性对齐：必须保证输出张量 out 的 device、dtype 与输入严格一致，且输出张量在显存中必须具备可预测的连续性或合法步长（out.is_contiguous()）。

参考资料

• Triton Tutorials
• Triton language API

资料用于建立事实基线；面试回答仍需用自己的语言组织。

← 上一课
第 7 课：稳定 Softmax 与融合 下一课 →
第 9 课：Blocked Matmul 与…