第 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 所示：

flowchart LR
    subgraph TileSwap["2D Tile 坐标与指针轴交换"]
        InTile["输入分块 (BLOCK_M, BLOCK_N)"] --> Trans["转置操作 tl.trans(x) ➔ (BLOCK_N, BLOCK_M)"]
        Trans --> OutTile["按列转置写回目标地址"]
    end

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

flowchart TD
    Program2D["2D Program 划分 (pid_m, pid_n)"] --> Load["合并加载源矩阵分块"]
    Load --> SMemTrans["Triton 编译器自动在片上执行无冲突转置"]
    SMemTrans --> Store["合并写回目标转置矩阵"]

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

二维网格分块映射：发射网格为 grid = (cdiv(M, BM), cdiv(N, BN))。每个 Program 处理一个形状为 (BM, BN) 的二维 Tile。
二维广播指针与双向掩码构造：
- 行偏移向量 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)。
合并加载与片上向量转置（tl.trans）：
- 执行 tile = tl.load(ptrs, mask=mask)。此时 tile 处于片上寄存器，形状为 (BM, BN)。
- 调用 tl.trans(tile)，在编译期将寄存器块的逻辑维度从 (BM, BN) 交换为 (BN, BM)。底层 Triton 编译器（Triton-GPU Dialect）会自动在片上寄存器或利用 Bank 优化的共享内存中调度高效的数据重排，对开发者完全透明。
交换指针与输出掩码写回：
- 输出矩阵的目标形状为 $(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$ 元素），能提升内存突发传输长度，但单个 Program 占用的寄存器翻倍，可能导致 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:31,64:95][0:31, 64:95]：
- 行范围 $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:30]$ 行列区域，位级无缝拼合。

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

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

补齐转置值。

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

Code

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 分；
一票否决：结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。

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

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

参考答案（仅 answer 分支）

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

Code

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 参考答案

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

输入阶段（2D 空间解构与广播指针生成）：Python 调度层以二维网格 (cdiv(M, BM), cdiv(N, BN)) 发射 Program。内核中通过广播切片构造出二维索引网格 rm[:, None] 与 rn[None, :]，结合源张量物理步长生成形状为 (BM, BN) 的二维物理指针矩阵 ptrs，并在输入掩码 (rm < M) & (rn < N) 保护下合并读取至寄存器块 tile。
状态变化阶段（片上多维张量转置）：调用 tl.trans(tile)，在编译期将寄存器块的逻辑形状由 (BM, BN) 交换为 (BN, BM)。底层 Triton 编译器自动生成无 Bank Conflict 的高效数据交换指令序列，在片上寄存器文件或共享内存中极速完成几何维度的转置重排，整个过程无需任何全局内存访问。
输出阶段（目标步长重构与反向掩码写回）：以反转后的坐标网格生成目标矩阵物理指针 rn[:, None] * som + rm[None, :] * son，其形状为 (BN, BM)。结合同步转置后的二维输出掩码 out_mask = (rn < N) & (rm < M)，通过 tl.store 以行合并连续写回全局显存，完成高带宽、无冲突的矩阵转置。

Q2 参考答案

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

张量形状（Shape）不兼容：输入 Tile 的形状是 (BM, BN)，因此输入掩码 mask 的形状严格为 (BM, BN)；而转置后写回的输出 Tile 形状为 (BN, BM)，tl.store 要求写回数据、输出指针与掩码的形状必须位级对齐。若 $BM \neq BN$ ，直接传入输入掩码会在编译期引发张量维度不匹配错误（Shape Mismatch）。
边界保护逻辑翻转：即使在方阵 Tile（ $BM = BN$ ）下，矩阵本身的维度通常是非方阵（ $M \neq N$ ）。输入掩码判定的是第 0 维是否小于 $M$ 、第 1 维是否小于 $N$ ；而写回目标矩阵时，行数变成了 $N$ 、列数变成了 $M$ 。若直接复用输入掩码，会导致原本针对行维度的边界约束被错误施加到列维度，导致尾部边界数据被错误裁剪或写入未分配的非法显存。

Q3 参考答案

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

张量维度与设备检查：断言输入张量严格为二维（x.ndim == 2），且驻留在正确的 CUDA 设备上。
张量底层存储重叠与非负步长检查（Non-overlapping Strides）：若输入张量是通过非法的跨步切片或负步长反向切片构造的（例如 x[::-1, :] 带有负步长），内核中以基地址做线性偏移的公式可能产生地址回卷或越界；若底层存储存在重叠（Overlapping Memory），必须在调用前通过 x.contiguous() 规范化。
输出张量属性对齐：必须保证输出张量 out 的 device、dtype 与输入严格一致，且输出张量在显存中必须具备可预测的连续性或合法步长（out.is_contiguous()）。

参考资料

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

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

--- title: "第 8 课：二维 Tile 与转置" jupyter: python3 --- ::: {.callout-note appearance="simple"} **参考答案版**：包含完整代码实现与问答题深度参考解析。建议先独立完成练习题后再展开核对。 ::: ## 本课目标与完成标准学完后你应能：构造二维 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 所示： ```{mermaid} flowchart LR subgraph TileSwap["2D Tile 坐标与指针轴交换"] InTile["输入分块 (BLOCK_M, BLOCK_N)"] --> Trans["转置操作 tl.trans(x) ➔ (BLOCK_N, BLOCK_M)"] Trans --> OutTile["按列转置写回目标地址"] end ``` 图 8-1：二维分块转置与内存步长重塑架构 ```{mermaid} flowchart TD Program2D["2D Program 划分 (pid_m, pid_n)"] --> Load["合并加载源矩阵分块"] Load --> SMemTrans["Triton 编译器自动在片上执行无冲突转置"] SMemTrans --> Store["合并写回目标转置矩阵"] ``` 图 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$ 元素），能提升内存突发传输长度，但单个 Program 占用的寄存器翻倍，可能导致 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:31, 64:95]$： - 行范围 $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:30]$ 行列区域，位级无缝拼合。请在阅读后先合上这一节，用自己的语言复述“输入状态 → 中间状态 → 输出状态”，再做练习。 ## 实践任务：唯一代码填空题补齐转置值。规则：只能修改 `TODO`/`______` 所在位置；不要删除断言或放宽误差。代码注释说明了每个边界条件。 ```{python} #| course_role: exercise 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 分； - 一票否决：结果碰巧正确但核心因果链错误、删除边界检查、把未运行结果说成实测。需要提示时按四级机制请求：概念区域 → 具体方向 → 关键局部 → 完整答案。 ::: {.callout-tip collapse="true" title="💡 点击展开参考答案与详细代码实现" icon="false"} ## 参考答案（仅 answer 分支）先完成题目再核对。即使代码一致，也要能解释关键步骤，并尝试更换一个输入规模。 ```{python} #| course_role: solution 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. **张量底层存储重叠与非负步长检查（Non-overlapping Strides）**：若输入张量是通过非法的跨步切片或负步长反向切片构造的（例如 `x[::-1, :]` 带有负步长），内核中以基地址做线性偏移的公式可能产生地址回卷或越界；若底层存储存在重叠（Overlapping Memory），必须在调用前通过 `x.contiguous()` 规范化。 3. **输出张量属性对齐**：必须保证输出张量 `out` 的 `device`、`dtype` 与输入严格一致，且输出张量在显存中必须具备可预测的连续性或合法步长（`out.is_contiguous()`）。 ## 参考资料 - [Triton Tutorials](https://triton-lang.org/main/getting-started/tutorials/) - [Triton language API](https://triton-lang.org/main/python-api/triton.language.html) 资料用于建立事实基线；面试回答仍需用自己的语言组织。 ::: ::: {.course-pagination} [**← 上一课** 第 7 课：稳定 Softmax 与融合](lesson07_softmax.qmd){.pagination-link .pagination-prev} [**下一课 →** 第 9 课：Blocked Matmul 与...](lesson09_matmul.qmd){.pagination-link .pagination-next} :::