> For AI agents: the complete documentation index is available at https://shijin.space/llms.txt, the full documentation bundle is available at https://shijin.space/llms-full.txt.

# CUDA 与 Tensor Core 编程

# MOSAIC-Sieve 的 CUDA 与 Tensor Core 编程

## 共同的问题：短归约点积、融合筛选与稀疏输出

分桶与桶内搜索看起来是两种操作，计算核心却很相近：对两组短向量进行大量点积，再只保留满足条件的关系。

将两组向量按行排列为 X 和 Y，可以写成：

$$
X\in\mathbb{Z}^{M\times K},\qquad Y\in\mathbb{Z}^{N\times K},\qquad D=XY^{\mathsf T}.
$$

X、Y 的元素采用 INT8 表示，K 是点积的归约维度。当前实现按 32 的倍数补齐并选择计算规格，支持的 K 为 64、96、128、160、192。对于大批量任务，通常有 K 远小于 M、N：两份输入是高而窄的矩阵，其乘积却有 M×N 个元素。

实际需要的是通过筛选的记录，而不是完整的 D：

$$
\mathcal{E}=\{(i,j,\text{payload})\mid P(D_{ij},\text{metadata})=\mathrm{true}\}.
$$

`payload` 表示需要保留的输出信息，`metadata` 表示行列元数据。P 表示由阈值、范数等信息决定的判定。典型情况下，命中数量远少于全部比较数量，但稀疏程度取决于数据与筛选条件，不能作为固定的容量保证。

- **分桶**：X 是待分配向量，Y 是中心；点积满足相似性条件时，记录向量所属的中心。输出主要是成员关系。
- **桶内搜索**：同一桶中的向量与自身比较，计算相当于 Gram 矩阵 XXᵀ。利用对称性只处理不重复的向量对，再筛出可能产生更短向量的候选，保留索引及所需点积等信息。

因此，这类任务可以描述为**短归约维度的稠密点积，融合阈值筛选，产生稀疏记录输出**。“稀疏”描述的是筛选后的输出，输入与乘加计算仍是稠密的，当前路径使用普通 INT8 Tensor Core MMA。

## 计算一小块，筛选一小块

如果先生成完整点积矩阵，再启动另一个 kernel 筛选，就需要写出并重新读取大量最终会被丢弃的结果。这里将筛选融合到点积 kernel 中：一小块结果在寄存器里完成 K 维累加后，立即判断并输出命中项。

可以将主要计算过程概括为：

```text
加载一组可复用的输入向量
for 每个中心面板或另一组向量面板:
    预取下一面板
    用 Tensor Core 完成当前小块的全部 K 维累加
    在寄存器中比较阈值
    汇总并输出命中记录
    确认预取完成后切换面板
```

## 从线程块到 Tensor Core 指令

线程块负责输出区域中的一块或一条带状区域，多个 warp 分担其中的小块。当前实现同时保留分桶面板、普通搜索 tile、条带搜索及小桶路径；它们共享点积基础，但不强求相同的线程块形状。

典型分桶路径固定一组 64 行输入，逐步扫描中心面板；条带搜索固定一组向量，与多个相邻面板比较。短 K 使一侧输入片段有机会完整缓存在寄存器中，供多个面板复用。另一侧使用共享内存双缓冲，在计算当前面板时预取下一面板。

关键指令承担不同层次的工作：

- **`cp.async`**：将数据从设备全局内存异步复制到共享内存，为面板预取提供基础。当前路径使用 16 字节搬运；读取前仍要等待复制完成，并在需要时同步线程块。它与前文讨论的主机到 GPU 传输是两层不同的数据移动。参见 [Ampere 异步数据复制](https://docs.nvidia.com/cuda/ampere-tuning-guide/index.html#asynchronous-data-copy-from-global-memory-to-shared-memory)。
- **`ldmatrix`**：将共享内存中的数据按 warp 协作方式装入寄存器片段，匹配 MMA 的操作数分布。代码中的 `.b16` 是加载的数据单元格式，这条路径中的数值仍按 INT8 参与计算。
- **`mma.sync`**：一个 warp 协作完成矩阵乘加。当前使用 `m16n8k32` 的 INT8×INT8→INT32 形式，每次沿 K 累加 32 个分量；例如 K=192 需要六步累加才能得到对应小块的完整点积。片段布局与指令约束见 [NVIDIA Warp 级 MMA 说明](https://docs.nvidia.com/cutlass/4.8.0/media/docs/pythonDSL/guides/mma/wmma_programming.html)。

当前这组 INT8 MMA 与异步复制路径面向 SM80 及以上兼容架构。MMA 和矩阵加载是 warp 协作操作，边界处理要保持参与线程满足指令要求，而不是让部分线程独自执行矩阵指令。参见 [PTX MMA 指令说明](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-mma)。

### MMA 与 WMMA：封装程度与布局控制

MMA 是矩阵乘加（Matrix Multiply-Accumulate），基本操作为 D = AB + C。使用 Tensor Core 时，可以通过较高层的 WMMA 接口组织计算，也可以直接使用 PTX 的 `mma.sync` 指令控制寄存器片段。

CUDA C++ 的 `nvcuda::wmma` 提供 `fragment`、`load_matrix_sync`、`mma_sync` 和 `store_matrix_sync`，分别表示矩阵片段及其加载、乘加和存储。开发者指定矩阵块形状、数据类型与内存布局，由接口封装片段在 warp 各线程中的分布。PTX 中也有 `wmma.load`、`wmma.mma`、`wmma.store` 指令族；因此，讨论 WMMA 时需要区分 C++ 接口与 PTX 指令。

**WMMA 便于表达矩阵运算，fragment 内部的元素映射则由实现决定。** 对所有元素施加相同操作，例如统一缩放，可以直接处理累加器片段；若操作需要知道元素对应的具体行、列，通常先将一个结果块存入共享内存等具有明确布局的位置，再按坐标处理。参见 [CUDA WMMA 编程说明](https://docs.nvidia.com/cuda/archive/13.1.0/cuda-programming-guide/05-appendices/cpp-language-extensions.html)。

**PTX `mma.sync` 为具体的指令形状和数据类型规定了寄存器片段与矩阵坐标的对应关系。** 开发者自行组织操作数，并按映射解释累加器。以当前使用的 `m16n8k32` 为例，一个 warp 协作计算 16×8 的输出块，每条指令沿归约维度推进 32 个分量。配合 `ldmatrix` 等加载指令，可以衔接共享内存布局、寄存器片段和乘加计算。参见 [PTX MMA 指令与片段布局](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-mma)。

本项目的筛选需要将每个点积与对应向量的范数、阈值和索引关联起来。直接使用 `mma.sync`，可以根据累加器位置找到这些信息，在寄存器中完成判断，并只输出命中记录。这使矩阵计算、筛选和输出压缩能够紧密衔接。

这种控制也带来维护成本：加载布局、线程协作和寄存器占用都需要自行安排。WMMA 更适合从常规分块矩阵乘出发实现计算；直接使用 MMA 则适合需要精细控制片段布局与后处理的场景。具体选择应以完整 kernel 的性能和实现复杂度为依据。

## 把稀疏输出也纳入 kernel 设计

筛选后的输出方式需要与计算布局一起设计。

分桶时，同一输入行的命中中心可以先形成位掩码，通过 warp 内交换汇总，再写出中心编号。这里主要使用 shuffle、位计数和位扫描，避免为每一个点积保存一个结果值。

搜索时，多数比较不产生候选，可以先用 `__any_sync` 判断一个 warp 是否存在命中；确认有命中后，再计算各线程需要写出的数量。通过 ballot 或 shuffle 前缀和分配局部位置，由一个线程用 `atomicAdd` 申请连续输出区间，其他线程各写自己的位置。这样可以减少逐候选争抢全局计数器的开销。相关方法见 [Warp 聚合原子操作](https://developer.nvidia.com/blog/cuda-pro-tip-optimized-filtering-warp-aggregated-atomics/)。

两条路径仍有不同的输出容量约束：分桶需要关注每个向量可保留的成员数，搜索需要关注候选缓冲容量。评估性能时应同时检查溢出与丢弃，不能把少输出了结果当成优化收益。

## 关键优化及其取舍

- **沿长维度复用，控制寄存器压力。** K 较短时，缓存一侧输入、沿另一侧持续扫描，能减少重复加载。但输入片段与累加器都占用寄存器，扩大 tile 或条带未必更快；当前分桶路径的部分规格会每完成 16 列结果就筛选一次，缩短累加器的存活时间。
- **让布局同时适合搬运与 MMA。** 全局数据沿向量坐标连续组织，共享内存布局则配合加载片段与对齐需求。当前采用 K+16 字节的共享内存行跨度；这是针对具体访问方式的布局选择，效果需要结合 bank 冲突、共享内存占用和实测判断。
- **减少重复比较和通用边界判断。** 桶内搜索跳过对称位置与自身配对，将完整的非对角区域与边界区域分开处理。小桶使用较轻的任务组织，避免少量有效比较承担过大的分块开销。
- **保持阈值与数值语义一致。** INT8 输入采用 INT32 累加，在 K≤192 的范围内，整数点积可容纳于 INT32；这不意味着对原始向量的量化没有误差。筛选还要与范数、阈值的量化和取整规则一致，尤其要核对临界命中，不能只检查平均误差。

这类 kernel 的优化目标是提高“完成点积、筛选并输出有效记录”的整体速度。仅测 MMA 吞吐，会漏掉加载、比较、压缩和输出成本。应覆盖不同 K、桶大小和命中率，并将结果放回[《实践经验》](https://shijin.space/projects/mosaic-sieve-practical-experience/index.md)中的完整流水线验证，确认局部优化确实改善了整体推进速度。

资料依据：MOSAIC-Sieve v3 的 CUDA 分桶与桶内搜索实现，以及文中引用的 NVIDIA 官方资料。核对版本：`1b22e06cceed`；整理日期：2026-10-20。本文讨论 kernel 内部的数据组织，与[《架构设计》](https://shijin.space/projects/mosaic-sieve-architecture/index.md)中的模块协作相衔接。

## 引用此笔记

- [项目概览](https://shijin.space/projects/mosaic-sieve/index.md)
- [实践经验](https://shijin.space/projects/mosaic-sieve-practical-experience/index.md)