第 4 课：二维 Grid、Stride 与广播

本课目标与完成标准

学完后你应能：实现矩阵按列广播 bias，区分逻辑边界与物理 stride。

通过必须同时满足：

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

前置关系

课程前置：Python、PyTorch 张量、CUDA 基本线程/内存概念
本课在路线中的作用：广播不复制 bias；每行 program 使用相同列下标读取 b，并按 x stride 找输入。

核心心智模型

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

在深度学习的前向与反向计算中，线性层（Linear Layer）的输出通常需要加上偏置向量： $Y = X + B$ ，其中 $X \in \mathbb{R}^{M \times N}$ ， $B \in \mathbb{R}^{N}$ 。根据 NumPy/PyTorch 广播语义，长度为 $N$ 的偏置向量 $B$ 需要沿行维度（ $M$ 维度）复制并与每一行相加。如果显式将 $B$ 复制成 $M \times N$ 的稠密矩阵，将浪费大量全局显存并引入无谓的内存搬运。本课解决的核心工程问题是：如何在 Triton 中通过将 Program 映射到二维逻辑网格（或一维 Program 处理一整行），结合张量的物理跨步（Stride）代数，实现零内存开销的隐式广播（Implicit In-Register Broadcasting），并深刻辨析“逻辑维度边界（Logical Shape）”与“物理内存步长（Physical Stride）”的本质解耦。

2. 它如何工作

二维 Stride 寻址与广播的数据通路如图 4-1 与图 4-2 所示：

二维张量行优先 Stride 映射

二维坐标 (row, col)

一维物理下标 = row * stride_row + col * stride_col

连续跨步读取二维矩阵块

图 4-1：Triton 二维 Block 与行优先 Stride 计算代数

1D 偏置 (BLOCK_N,)

切片扩展维度 [None, :] 变成 (1, BLOCK_N)

沿行方向自动广播与 (BLOCK_M, BLOCK_N) 矩阵相加

图 4-2：Triton 广播机制（Broadcasting）执行流

网格与行映射：主机端发射 $M$ 个 Program 块：grid = (M,)。每个 Program 的 row = tl.program_id(0) 精准负责矩阵的一整行。
列向量与步长寻址：通过 cols = tl.arange(0, BLOCK) 生成列索引向量，计算输入矩阵的物理偏移量： $\text{xoffs} = \text{row} \times s_{xm} + \text{cols} \times s_{xn}$ 其中 $s_{xm}$ 为行步长（跨行跳跃的元素个数）， $s_{xn}$ 为列步长（通常为 1）。这使得内核天然兼容任意内存布局（如连续、转置或带 Padding 的对齐内存）。
偏置向量的零拷贝广播：偏置向量 $B$ 只依赖于列索引：b + cols。所有 $M$ 个并发执行的 Program 块在相同的逻辑列位置上读取相同的偏置数据 $B$ 。底层 L2 Cache 会对偏置张量进行高命中率广播缓存，省去显式张量扩充。
受保护的计算与写回：使用掩码 mask = cols < N 抑制超出逻辑列数 $N$ 的无效访存，直接在寄存器中执行 xv + bv，随后写回连续输出矩阵：out + row * N + cols。

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

逻辑列宽 vs 物理跨步（混淆 $N$ 与 $s_{xm}$ ）：常见致命错误是将谓词掩码写为 mask = cols < sxm。在张量存在行对齐填充（Row Padding）或属于大矩阵切片时，物理行步长 $s_{xm}$ 往往严格大于逻辑有效列数 $N$ （例如 $N=17$ ，但为了 128-bit 对齐，底层分配了 $s_{xm}=32$ ）。若使用 cols < sxm，内核将误把 Padding 区域的未初始化垃圾数据读入并写回输出，破坏张量边界。
块大小为 2 的幂次约束（Power of Two）： Triton 的 tl.arange(0, BLOCK) 要求 BLOCK 必须是 2 的幂次。若列数 $N = 17$ ，必须将 BLOCK 向上对齐到 triton.next_power_of_2(17) = 32，并通过 mask = cols < N 屏蔽后 15 个 lane。
输出张量的连续性假设：若写回公式写为 out + row * N + cols，则隐式假设了输出矩阵 out 必须是行优先且完全紧凑连续的（out.stride(0) == N, out.stride(1) == 1）。若输出也是非连续切片，必须传入输出自身的步长。

4. 性能与工程取舍

一维 Program 一整行 vs 二维分块（Tiling）：当列宽 $N$ 较小（如 $N \le 2048$ 或 4096）时，一维 Program 处理一整行具有极其简洁的代码结构，无跨 Block 协调开销；但当 $N$ 极大（如 LLM 词表维度 $N = 32000, 128000$ ）时，单个 Program 分配如此巨大的 BLOCK 会瞬间耗尽 SM 寄存器配额，导致编译失败或严重的寄存器溢出（Register Spill）。此时必须将网格重构为二维分块 grid = (M, cdiv(N, BLOCK_N))。
Cache 广播效应：由于 $M$ 个 Program 几乎并发发射并读取同一段偏置地址 $B[0:N]$ ，偏置向量在首次读取后基本命中 L2（以及同一 SM 上的 L1），对 HBM 带宽的额外消耗很小。

具体演示

以 $M = 3, N = 17$ 为例：

BLOCK = triton.next_power_of_2(17) = 32。
网格发射 3 个 Program：pid=0, 1, 2。
每个 Program 中：
- cols = [0, 1, ..., 16, 17, ..., 31]
- mask = [True x 17, False x 15]
- row 0 加载 x[0, 0:17] 与 b[0:17]
- row 1 加载 x[1, 0:17] 与 b[0:17]
- row 2 加载 x[2, 0:17] 与 b[0:17]
每个 Program 各自读取一次 $b$ 的 17 个浮点数，后两次读取通常命中缓存；写回时后 15 个无效位被 mask 屏蔽。

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

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

补齐 bias 地址。

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

Code

import torch
import triton
import triton.language as tl

@triton.jit
def bias_kernel(x, b, out, M: tl.constexpr, N: tl.constexpr,
                sxm: tl.constexpr, sxn: tl.constexpr, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < N
    xoffs = row * sxm + cols * sxn
    xv = tl.load(x + xoffs, mask=mask)
    bv = tl.load(b + ______, mask=mask)  # TODO: bias 只依赖列
    tl.store(out + row * N + cols, xv + bv, mask=mask)

def add_bias(x, b):
    assert x.ndim == 2 and b.ndim == 1 and x.shape[1] == b.numel()
    M, N = x.shape
    out = torch.empty((M, N), device=x.device, dtype=x.dtype)
    block = triton.next_power_of_2(N)
    bias_kernel[(M,)](x, b, out, M, N, x.stride(0), x.stride(1), BLOCK=block)
    return out

for shape in ((3, 17), (4, 128)):
    x = torch.randn(shape, device="cuda"); b = torch.randn(shape[1], device="cuda")
    torch.testing.assert_close(add_bias(x, b), x + b)

检查方法

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

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

Q1

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

你的答案：

Q2

把 mask 写成 cols < sxm 为什么对 padding stride 错？

你的答案：

Q3

N 很大时怎样扩展为二维 grid 并保持输出地址正确？

你的答案：

评分与通过规则

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

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

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

参考答案（仅 answer 分支）

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

Code

import torch
import triton
import triton.language as tl

@triton.jit
def bias_kernel(x, b, out, M: tl.constexpr, N: tl.constexpr,
                sxm: tl.constexpr, sxn: tl.constexpr, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < N
    xoffs = row * sxm + cols * sxn
    xv = tl.load(x + xoffs, mask=mask)
    bv = tl.load(b + cols, mask=mask)
    tl.store(out + row * N + cols, xv + bv, mask=mask)

def add_bias(x, b):
    assert x.ndim == 2 and b.ndim == 1 and x.shape[1] == b.numel()
    M, N = x.shape
    out = torch.empty((M, N), device=x.device, dtype=x.dtype)
    block = triton.next_power_of_2(N)
    bias_kernel[(M,)](x, b, out, M, N, x.stride(0), x.stride(1), BLOCK=block)
    return out

for shape in ((3, 17), (4, 128)):
    x = torch.randn(shape, device="cuda"); b = torch.randn(shape[1], device="cuda")
    torch.testing.assert_close(add_bias(x, b), x + b)

Q1 参考答案

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

输入阶段（逻辑解耦与步长物理寻址）：Python 包装层获取输入矩阵的维度 $(M, N)$ 与物理步长 $(s_{xm}, s_{xn})$ ，以 (M,) 为网格发射内核。每个 Program 块通过 tl.program_id(0) 锁定当前行索引 row，并通过 tl.arange(0, BLOCK) 构建列索引向量 cols。结合物理步长精确计算二维内存地址 x + row * sxm + cols * sxn，偏置则仅寻址 b + cols。
状态变化阶段（带掩码合并加载与片上广播计算）：在掩码 mask = cols < N 保护下，SM 执行合并内存事务将该行的有效元素载入寄存器块 xv。同时，各行 Program 并发从显存加载相同的偏置向量至寄存器块 bv（借助硬件 L1/L2 Cache 实现高速单播/广播）。随后执行寄存器逐元素相加 xv + bv。
输出阶段（紧凑布局带掩码写回）：利用相同的 mask 谓词，将相加结果通过 tl.store 写入目标地址 out + row * N + cols。越界 lane 的内存写操作被硬件完全屏蔽，保证紧凑目标张量的内存纯净度。

Q2 参考答案

在张量体系中，“逻辑有效列数 $N$ ”与“物理行步长 $s_{xm}$ ”具有根本区别：

逻辑列数 $N$ 代表当前张量在语义上真正包含的有效特征数。
物理行步长 $s_{xm}$ 代表在显存物理排布中从第 $i$ 行跳转到第 $i+1$ 行所需的内存偏移量。当底层显存为提升内存对齐效率而引入了 Padding（例如为了 128-bit 向量对齐，逻辑列 $N=17$ 被填充为物理步长 $s_{xm}=32$ ），或者输入张量是截取自更大张量的一个切片时，物理跨步 $s_{xm}$ 严格大于逻辑维度 $N$ （ $s_{xm} > N$ ）。如果将掩码误写为 cols < sxm，则第 $N$ 到 $s_{xm}-1$ 之间的填充内存（包含未初始化的脏数据或邻近切片数据）也会被判定为合法有效并参与计算，随后被强行写回输出张量，导致越界污染和严重的数据静默损坏。

Q3 参考答案

当 $N$ 极大（例如 $N = 65536$ ）导致单 Program 无法在寄存器中容纳一整行时，必须将网格重构为二维分块架构：

二维网格划分：发射网格设置为 grid = (M, triton.cdiv(N, BLOCK_N))，固定较小的分块尺寸（如 BLOCK_N = 1024）。
二维坐标与列偏移解构：内核中通过 row = tl.program_id(0) 获取行号，通过 block_col_id = tl.program_id(1) 获取当前处理的列块号。当前 Program 的列偏移向量计算为： cols = block_col_id * BLOCK_N + tl.arange(0, BLOCK_N) 谓词掩码保持为：mask = cols < N。
输出地址保持正确：
- 输入寻址：x + row * sxm + cols * sxn
- 偏置寻址：b + cols
- 输出寻址：若输出为连续矩阵，地址公式为 out + row * N + cols；若输出同样包含特定步长，地址公式为 out + row * out_stride_m + cols * out_stride_n。通过在输出指针基地址上精准加上全局二维逻辑偏移 row * N + cols，并配合 mask 保护尾块，即可在二维并行网格下保证全域访存地址的绝对正确性。

参考资料

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

← 上一课
第 3 课：逐元素激活与融合边界 下一课 →
第 5 课：行归约与 Padding 单位元

--- title: "第 4 课：二维 Grid、Stride 与广播" jupyter: python3 --- ## 本课目标与完成标准学完后你应能：实现矩阵按列广播 bias，区分逻辑边界与物理 stride。通过必须同时满足： - 独立补齐本课唯一的代码填空题，并通过给定检查； - 三个问答题均说明因果链，而不是只报术语； - 能指出至少一个正确性边界和一个性能取舍； - 总分不低于 8/10，且没有一票否决级概念错误。 ## 前置关系 - 课程前置：Python、PyTorch 张量、CUDA 基本线程/内存概念 - 本课在路线中的作用：广播不复制 bias；每行 program 使用相同列下标读取 b，并按 x stride 找输入。 ## 核心心智模型 ### 1. 它是什么，解决什么问题在深度学习的前向与反向计算中，线性层（Linear Layer）的输出通常需要加上偏置向量：$Y = X + B$，其中 $X \in \mathbb{R}^{M \times N}$，$B \in \mathbb{R}^{N}$。根据 NumPy/PyTorch 广播语义，长度为 $N$ 的偏置向量 $B$ 需要沿行维度（$M$ 维度）复制并与每一行相加。如果显式将 $B$ 复制成 $M \times N$ 的稠密矩阵，将浪费大量全局显存并引入无谓的内存搬运。本课解决的核心工程问题是：如何在 Triton 中通过将 Program 映射到二维逻辑网格（或一维 Program 处理一整行），结合张量的物理跨步（Stride）代数，实现**零内存开销的隐式广播（Implicit In-Register Broadcasting）**，并深刻辨析“逻辑维度边界（Logical Shape）”与“物理内存步长（Physical Stride）”的本质解耦。 ### 2. 它如何工作二维 Stride 寻址与广播的数据通路如图 4-1 与图 4-2 所示： ```{=html} <div class="fd fd-flow" role="img"><div class="fd-stack"><div class="fd-comp fd-td"><div class="fd-sg" data-nid="Stride2D" data-depth="0"><div class="fd-sg-title">二维张量行优先 Stride 映射</div><div class="fd-stack"><div class="fd-comp fd-td"><div class="fd-rank"><div class="fd-node fd-box" data-nid="Coord">二维坐标 (row, col)</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="Formula">一维物理下标 = row * stride_row + col * stride_col</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="Mem">连续跨步读取二维矩阵块</div></div></div></div></div></div></div><i class="fd-e" data-s="Coord" data-t="Formula" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Formula" data-t="Mem" data-k="solid" data-h="1" data-b="0"></i></div> ``` <p class="caption" align="center"><em>图 4-1：Triton 二维 Block 与行优先 Stride 计算代数</em></p> ```{=html} <div class="fd fd-flow" role="img"><div class="fd-stack"><div class="fd-comp fd-lr" data-minw="640"><div class="fd-rank"><div class="fd-node fd-box" data-nid="Bias">1D 偏置 (BLOCK_N,)</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="Expand">切片扩展维度 [None, :] 变成 (1, BLOCK_N)</div></div><div class="fd-gap" style="--gl:44px;--gt:38px"></div><div class="fd-rank"><div class="fd-node fd-box" data-nid="Broad">沿行方向自动广播与 (BLOCK_M, BLOCK_N) 矩阵相加</div></div></div></div><i class="fd-e" data-s="Bias" data-t="Expand" data-k="solid" data-h="1" data-b="0"></i><i class="fd-e" data-s="Expand" data-t="Broad" data-k="solid" data-h="1" data-b="0"></i></div> ``` <p class="caption" align="center"><em>图 4-2：Triton 广播机制（Broadcasting）执行流</em></p> 1. **网格与行映射**：主机端发射 $M$ 个 Program 块：`grid = (M,)`。每个 Program 的 `row = tl.program_id(0)` 精准负责矩阵的一整行。 2. **列向量与步长寻址**：通过 `cols = tl.arange(0, BLOCK)` 生成列索引向量，计算输入矩阵的物理偏移量： $$ \text{xoffs} = \text{row} \times s_{xm} + \text{cols} \times s_{xn} $$ 其中 $s_{xm}$ 为行步长（跨行跳跃的元素个数），$s_{xn}$ 为列步长（通常为 1）。这使得内核天然兼容任意内存布局（如连续、转置或带 Padding 的对齐内存）。 3. **偏置向量的零拷贝广播**：偏置向量 $B$ 只依赖于列索引：`b + cols`。所有 $M$ 个并发执行的 Program 块在相同的逻辑列位置上读取相同的偏置数据 $B$。底层 L2 Cache 会对偏置张量进行高命中率广播缓存，省去显式张量扩充。 4. **受保护的计算与写回**：使用掩码 `mask = cols < N` 抑制超出逻辑列数 $N$ 的无效访存，直接在寄存器中执行 `xv + bv`，随后写回连续输出矩阵：`out + row * N + cols`。 ### 3. 正确性条件与常见误区 - **逻辑列宽 vs 物理跨步（混淆 $N$ 与 $s_{xm}$）**：常见致命错误是将谓词掩码写为 `mask = cols < sxm`。在张量存在行对齐填充（Row Padding）或属于大矩阵切片时，物理行步长 $s_{xm}$ 往往严格大于逻辑有效列数 $N$（例如 $N=17$，但为了 128-bit 对齐，底层分配了 $s_{xm}=32$）。若使用 `cols < sxm`，内核将误把 Padding 区域的未初始化垃圾数据读入并写回输出，破坏张量边界。 - **块大小为 2 的幂次约束（Power of Two）**： Triton 的 `tl.arange(0, BLOCK)` 要求 `BLOCK` 必须是 2 的幂次。若列数 $N = 17$，必须将 `BLOCK` 向上对齐到 `triton.next_power_of_2(17) = 32`，并通过 `mask = cols < N` 屏蔽后 15 个 lane。 - **输出张量的连续性假设**：若写回公式写为 `out + row * N + cols`，则隐式假设了输出矩阵 `out` 必须是行优先且完全紧凑连续的（`out.stride(0) == N, out.stride(1) == 1`）。若输出也是非连续切片，必须传入输出自身的步长。 ### 4. 性能与工程取舍 - **一维 Program 一整行 vs 二维分块（Tiling）**：当列宽 $N$ 较小（如 $N \le 2048$ 或 4096）时，一维 Program 处理一整行具有极其简洁的代码结构，无跨 Block 协调开销；但当 $N$ 极大（如 LLM 词表维度 $N = 32000, 128000$）时，单个 Program 分配如此巨大的 `BLOCK` 会瞬间耗尽 SM 寄存器配额，导致编译失败或严重的寄存器溢出（Register Spill）。此时必须将网格重构为二维分块 `grid = (M, cdiv(N, BLOCK_N))`。 - **Cache 广播效应**：由于 $M$ 个 Program 几乎并发发射并读取同一段偏置地址 $B[0:N]$，偏置向量在首次读取后基本命中 L2（以及同一 SM 上的 L1），对 HBM 带宽的额外消耗很小。 ## 具体演示以 $M = 3, N = 17$ 为例： - `BLOCK = triton.next_power_of_2(17) = 32`。 - 网格发射 3 个 Program：`pid=0, 1, 2`。 - 每个 Program 中： - `cols = [0, 1, ..., 16, 17, ..., 31]` - `mask = [True x 17, False x 15]` - `row 0` 加载 `x[0, 0:17]` 与 `b[0:17]` - `row 1` 加载 `x[1, 0:17]` 与 `b[0:17]` - `row 2` 加载 `x[2, 0:17]` 与 `b[0:17]` - 每个 Program 各自读取一次 $b$ 的 17 个浮点数，后两次读取通常命中缓存；写回时后 15 个无效位被 mask 屏蔽。请在阅读后先合上这一节，用自己的语言复述“输入状态 → 中间状态 → 输出状态”，再做练习。 ## 实践任务：唯一代码填空题补齐 bias 地址。规则：只能修改 `TODO`/`______` 所在位置；不要删除断言或放宽误差。代码注释说明了每个边界条件。 ```{python} #| course_role: exercise import torch import triton import triton.language as tl @triton.jit def bias_kernel(x, b, out, M: tl.constexpr, N: tl.constexpr, sxm: tl.constexpr, sxn: tl.constexpr, BLOCK: tl.constexpr): row = tl.program_id(0) cols = tl.arange(0, BLOCK) mask = cols < N xoffs = row * sxm + cols * sxn xv = tl.load(x + xoffs, mask=mask) bv = tl.load(b + ______, mask=mask) # TODO: bias 只依赖列 tl.store(out + row * N + cols, xv + bv, mask=mask) def add_bias(x, b): assert x.ndim == 2 and b.ndim == 1 and x.shape[1] == b.numel() M, N = x.shape out = torch.empty((M, N), device=x.device, dtype=x.dtype) block = triton.next_power_of_2(N) bias_kernel[(M,)](x, b, out, M, N, x.stride(0), x.stride(1), BLOCK=block) return out for shape in ((3, 17), (4, 128)): x = torch.randn(shape, device="cuda"); b = torch.randn(shape[1], device="cuda") torch.testing.assert_close(add_bias(x, b), x + b) ``` ### 检查方法在 CUDA/Triton 环境运行本单元格；断言覆盖规则尺寸和非规则尾块。首次 JIT 不计入性能。提交时请给出：补齐后的代码、实际运行输出（环境不可用时注明“仅静态审查”）以及对失败用例的解释。 ### Q1 不要背定义：请从输入、状态变化和输出三个阶段解释“二维 Grid、Stride 与广播”的工作机制。 **你的答案：** ### Q2 把 mask 写成 `cols < sxm` 为什么对 padding stride 错？ **你的答案：** ### Q3 N 很大时怎样扩展为二维 grid 并保持输出地址正确？ **你的答案：** ## 评分与通过规则 - 代码 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 bias_kernel(x, b, out, M: tl.constexpr, N: tl.constexpr, sxm: tl.constexpr, sxn: tl.constexpr, BLOCK: tl.constexpr): row = tl.program_id(0) cols = tl.arange(0, BLOCK) mask = cols < N xoffs = row * sxm + cols * sxn xv = tl.load(x + xoffs, mask=mask) bv = tl.load(b + cols, mask=mask) tl.store(out + row * N + cols, xv + bv, mask=mask) def add_bias(x, b): assert x.ndim == 2 and b.ndim == 1 and x.shape[1] == b.numel() M, N = x.shape out = torch.empty((M, N), device=x.device, dtype=x.dtype) block = triton.next_power_of_2(N) bias_kernel[(M,)](x, b, out, M, N, x.stride(0), x.stride(1), BLOCK=block) return out for shape in ((3, 17), (4, 128)): x = torch.randn(shape, device="cuda"); b = torch.randn(shape[1], device="cuda") torch.testing.assert_close(add_bias(x, b), x + b) ``` ### Q1 参考答案从三个阶段拆解其运行机制： 1. **输入阶段（逻辑解耦与步长物理寻址）**：Python 包装层获取输入矩阵的维度 $(M, N)$ 与物理步长 $(s_{xm}, s_{xn})$，以 `(M,)` 为网格发射内核。每个 Program 块通过 `tl.program_id(0)` 锁定当前行索引 `row`，并通过 `tl.arange(0, BLOCK)` 构建列索引向量 `cols`。结合物理步长精确计算二维内存地址 `x + row * sxm + cols * sxn`，偏置则仅寻址 `b + cols`。 2. **状态变化阶段（带掩码合并加载与片上广播计算）**：在掩码 `mask = cols < N` 保护下，SM 执行合并内存事务将该行的有效元素载入寄存器块 `xv`。同时，各行 Program 并发从显存加载相同的偏置向量至寄存器块 `bv`（借助硬件 L1/L2 Cache 实现高速单播/广播）。随后执行寄存器逐元素相加 `xv + bv`。 3. **输出阶段（紧凑布局带掩码写回）**：利用相同的 `mask` 谓词，将相加结果通过 `tl.store` 写入目标地址 `out + row * N + cols`。越界 lane 的内存写操作被硬件完全屏蔽，保证紧凑目标张量的内存纯净度。 ### Q2 参考答案在张量体系中，“逻辑有效列数 $N$”与“物理行步长 $s_{xm}$”具有根本区别： - 逻辑列数 $N$ 代表当前张量在语义上真正包含的有效特征数。 - 物理行步长 $s_{xm}$ 代表在显存物理排布中从第 $i$ 行跳转到第 $i+1$ 行所需的内存偏移量。当底层显存为提升内存对齐效率而引入了 Padding（例如为了 128-bit 向量对齐，逻辑列 $N=17$ 被填充为物理步长 $s_{xm}=32$），或者输入张量是截取自更大张量的一个切片时，物理跨步 $s_{xm}$ 严格大于逻辑维度 $N$（$s_{xm} > N$）。如果将掩码误写为 `cols < sxm`，则第 $N$ 到 $s_{xm}-1$ 之间的填充内存（包含未初始化的脏数据或邻近切片数据）也会被判定为合法有效并参与计算，随后被强行写回输出张量，导致越界污染和严重的数据静默损坏。 ### Q3 参考答案当 $N$ 极大（例如 $N = 65536$）导致单 Program 无法在寄存器中容纳一整行时，必须将网格重构为二维分块架构： 1. **二维网格划分**：发射网格设置为 `grid = (M, triton.cdiv(N, BLOCK_N))`，固定较小的分块尺寸（如 `BLOCK_N = 1024`）。 2. **二维坐标与列偏移解构**：内核中通过 `row = tl.program_id(0)` 获取行号，通过 `block_col_id = tl.program_id(1)` 获取当前处理的列块号。当前 Program 的列偏移向量计算为： `cols = block_col_id * BLOCK_N + tl.arange(0, BLOCK_N)` 谓词掩码保持为：`mask = cols < N`。 3. **输出地址保持正确**： - 输入寻址：`x + row * sxm + cols * sxn` - 偏置寻址：`b + cols` - 输出寻址：若输出为连续矩阵，地址公式为 `out + row * N + cols`；若输出同样包含特定步长，地址公式为 `out + row * out_stride_m + cols * out_stride_n`。通过在输出指针基地址上精准加上全局二维逻辑偏移 `row * N + cols`，并配合 `mask` 保护尾块，即可在二维并行网格下保证全域访存地址的绝对正确性。 ## 参考资料 - [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} [**← 上一课**<br>第 3 课：逐元素激活与融合边界](lesson03_relu.qmd){.pagination-link .pagination-prev} [**下一课 →**<br>第 5 课：行归约与 Padding 单位元](lesson05_row_sum.qmd){.pagination-link .pagination-next} :::