PyTorch CUDA 索引宽度实战：int 与 int64_t 的取舍、大张量越界修复与 index_t 分派机制

当 PyTorch 的 CUDA kernel 在接近 2^31 个元素的大张量上出现结果错乱或内存越界时，根因往往不是算法逻辑，而是 32 位 int 索引算术溢出。本文基于 PyTorch 仓库中 .claude/skills/cuda-index-width/SKILL.md 所沉淀的索引宽度决策技能展开，结合 aten 源码中 canUse32BitIndexMath、AT_DISPATCH_INDEX_TYPES 与 CUDA_KERNEL_LOOP_TYPE 的真实实现，讲清楚如何定位溢出点、选择最经济的修复方案，以及如何用回归测试和二进制体积评估来验证修复的正确性与代价。读完本文，你将能够在修改任何 PyTorch CUDA kernel 时给出有依据的索引类型决策，而不是盲目把所有索引变量改成 int64_t。

一、问题起点：32 位索引在哪里溢出

CUDA kernel 最常见的错误是“在 2^31 个元素附近突然出错”：元素索引、blockIdx * blockDim + threadIdx 的乘积、或者 size * stride 的偏移超过了 int32_t 的最大值。修复的第一步不是动手改类型，而是找到具体那条可以超过 32 位的表达式，并判断它执行在什么位置。技能文档给出的分类是：

• 每个 CTA 或每组输出元素的一次性 setup 计算：优先在第一次乘法处做一次局部 64 位转换，其余保持 32 位；
• 热点逐元素线性索引（循环变量、取模/除法分解、data[index] 访问）：将 kernel 模板化到 index_t，按 int 与 int64_t 分派；
• 算法本身只有 32 位假设、无法支持 64 位：与其静默溢出，不如加一条明确的 TORCH_CHECK(canUse32BitIndexMath(...)) 快速失败；
• Grid 维度溢出：光改算术类型不够，必须对该维度增加分块（tiling）/跨步（striding）策略，因为 CUDA 启动维度受硬件上限约束（gridDim.x <= 2^31 - 1，gridDim.y/z <= 65535）。

PyTorch 测试里已经存在对这一边界的显式防护。例如 test/test_cuda.py 中就有 “Make sure input.numel() > INT_MAX is handled” 的注释，说明仓库自身也把 numel() > INT_MAX 当作一类需要专门覆盖的回归场景。

二、标准工具链：canUse32BitIndexMath 与索引分派宏

技能文档要求：不要写临时检查，直接使用仓库中已有的标准工具：

#include <ATen/native/CanUse32BitIndexMath.h>
#include <ATen/cuda/detail/KernelUtils.h>

canUse32BitIndexMath 到底检查了什么

接口声明在 CanUse32BitIndexMath.h：

TORCH_API bool canUse32BitIndexMath(const at::TensorBase &t, int64_t max_elem=std::numeric_limits<int32_t>::max());

其实现位于 IndexingUtils.cpp，可以确认它检查的不只是 numel()：

bool canUse32BitIndexMath(const TensorBase& t, int64_t max_elem) {
  auto elements = t.sym_numel();
  if (elements >= max_elem) {
    return false;
  }
  ...
  c10::SymInt offset = 0;
  auto linearId = elements - 1;
  // NOTE: Assumes all strides are positive, which is true for now
  for (auto i = t.dim() - 1; i >= 0; --i) {
    auto curDimIndex = linearId % t.sym_size(i);
    auto curDimOffset = curDimIndex * t.sym_stride(i);
    offset += curDimOffset;
    linearId /= t.sym_size(i);
  }
  if (offset >= max_elem) {
    return false;
  }
  return true;
}

也就是说它同时校验两点：元素总数是否达到 max_elem（默认 INT_MAX），以及最后一个元素在 stride 展开后所需的最大 storage 偏移是否达到 max_elem。这正是技能文档强调的“numel() 很小的张量如果是一个大 strided view，仍然可能需要 64 位偏移”的源码依据——实现里是逐个维度做 curDimIndex * sym_stride(i) 累加，大 stride 会直接把偏移推过 32 位边界。

在已经包含 CUDA detail 头文件的 .cu 里，at::cuda::detail::canUse32BitIndexMath 也可作为别名使用；at::native:: 版本则是通用入口。调用时记得检查所有偏移会被所选索引类型计算的张量，而不仅是输出张量——多张量算子里任何一个参与者的偏移越界都会导致错误。

AT_DISPATCH_INDEX_TYPES：按 Int / Long 双路分派

分派宏定义在 Dispatch.h：

#define AT_DISPATCH_INDEX_TYPES(TYPE, NAME, ...)     \
  AT_DISPATCH_SWITCH(                                \
      TYPE,                                          \
      NAME,                                          \
      AT_PRIVATE_CASE_TYPE_USING_HINT(               \
          at::ScalarType::Int, index_t, __VA_ARGS__) \
          AT_PRIVATE_CASE_TYPE_USING_HINT(           \
              at::ScalarType::Long, index_t, __VA_ARGS__))

它把 ScalarType::Int / ScalarType::Long 映射到一个统一的别名 index_t（分别对应 int32_t / int64_t），lambda 捕获后 kernel 即可模板化。真实用例如 SoftMax.cu 的 host 侧 softmax 启动器：

AT_DISPATCH_INDEX_TYPES(
    at::native::canUse32BitIndexMath(input, INT_MAX) ? ScalarType::Int : ScalarType::Long,
"host_softmax_launcher", [&] {
    ...
    cunn_SpatialSoftMaxForward<scalar_t, accscalar_t, scalar_t, index_t, Epilogue>
      <<<grid, block, smem_size, stream>>>(...);

这种“canUse32BitIndexMath(...) ? Int : Long”的三元表达式在 aten/src/ATen/native/cuda/ 下被大量复用，例如 Indexing.cu、EmbeddingBag.cu、ScatterGatherKernel.cu、AveragePool2d.cu 等文件，是该仓库事实上的标准模式。

CUDA_KERNEL_LOOP_TYPE：让 grid-stride 循环对任意宽度都正确

grid-stride 循环宏定义在 KernelUtils.h：

// int64_t _i_n_d_e_x specifically prevents overflow in the loop increment.
#define CUDA_KERNEL_LOOP_TYPE(i, n, index_type)                         \
  int64_t _i_n_d_e_x = ((int64_t) blockIdx.x) * blockDim.x + threadIdx.x;           \
  for (index_type i=_i_n_d_e_x; _i_n_d_e_x < (n); _i_n_d_e_x+=blockDim.x * gridDim.x, i=_i_n_d_e_x)

#define CUDA_KERNEL_LOOP(i, n) CUDA_KERNEL_LOOP_TYPE(i, n, int)

从源码结构看，这个宏的设计很讲究：循环计数器 _i_n_d_e_x 始终是 int64_t，保证 blockDim.x * gridDim.x 的累加不会在 32 位下回绕，而暴露给用户代码的下标 i 的宽度由模板参数 index_type 决定。头文件注释也解释了边界安全：即使最后一轮迭代后 _i_n_d_e_x += ... 溢出 32 位，此时 _i_n_d_e_x >= n，循环不会再进入，溢出值不会被使用。同文件中的 GET_BLOCKS 则负责把 int64_t 的 N 安全地折算成 int 的 block 数，并在超出硬件调度能力时给出 TORCH_INTERNAL_ASSERT——这呼应了前面“grid 维度溢出不能只靠改算术类型”的分类。

三、四类修复模式（Fix Patterns）

以下四种模式覆盖了技能文档给出的全部修复路径，前两类附完整代码，后两类给出决策要点。

模式 1：局部基础偏移溢出——只在 setup 处做一次 64 位转换

如果只有 base pointer 偏移可能溢出、且它是在热点循环之外计算的，那么 kernel 其余部分保持不动：

int64_t plane = blockIdx.x;
input = input + plane * strideD;
output = output + plane * osizeH * osizeW;

这样避免了 kernel 模板实例化的翻倍，内层循环算术保持 32 位。适用条件是 tile 内部的各维度仍然放得下 int。这是代价最小的方案，应优先考虑。

模式 2：grid-stride 线性 kernel——模板化 index_t 并分派

如果循环变量、取模/除法分解或最终的 data[index] 访问可能超过 32 位，按标准模式模板化：

template <typename scalar_t, typename index_t>
__global__ void kernel(index_t n, const scalar_t* in, scalar_t* out) {
  CUDA_KERNEL_LOOP_TYPE(index, n, index_t) {
    out[index] = in[index];
  }
}

AT_DISPATCH_INDEX_TYPES(
    canUse32BitIndexMath(out, INT_MAX) && canUse32BitIndexMath(in, INT_MAX)
        ? ScalarType::Int
        : ScalarType::Long,
    "kernel_index_type",
    [&] {
      kernel<scalar_t, index_t><<<blocks, threads, 0, stream>>>(n, in, out);
      C10_CUDA_KERNEL_LAUNCH_CHECK();
    });

注意两个细节：判断条件里输入和输出都要过一遍 canUse32BitIndexMath；启动后紧跟 C10_CUDA_KERNEL_LAUNCH_CHECK() 检查启动失败。技能文档特别提醒：优先采用这种条件分派而不是无条件把循环下标改成 int64_t，因为热循环里的 64 位除法/取模开销在 GPU 上可能是可测量的（GPU 没有原生的 64 位整数除法指令，需要多次软件模拟运算）。

模式 3：TensorInfo/accessor 与带 stride 的 kernel

当偏移是由 size/stride 计算出来的（例如使用 TensorInfo 或 TensorAccessor 的 kernel），只在所有参与张量都通过对应索引类型的 canUse32BitIndexMath 时才分派 32 位。再次强调：numel() 小的张量若为大 strided view 仍需要 64 位偏移——这一点已被 IndexingUtils.cpp 中逐维 stride 偏移累加的实现所背书。

模式 4：只能 32 位的 kernel——快速失败而不是静默溢出

如果支持 64 位索引需要更大的算法重写，或会超出 CUDA 启动维度限制，就明确失败：

TORCH_CHECK(
    canUse32BitIndexMath(input) && canUse32BitIndexMath(output),
    "op_name: tensors must fit into 32-bit index math");

使用前提是该算子已有文档化或可接受的大小限制；不要因为图省事，把一个被报告的“正确性 bug”（其实可以通过模式 1/2/3 修复）变成新的尺寸限制。

四、二进制体积与性能权衡

对 index_t 做模板化会让每个受影响的 kernel 在每个 scalar dtype 与 memory-format 特化上都再复制一份实例。在给多个 kernel 加索引分派之前，先问自己：溢出是在热路径上，还是只在一处 setup 表达式里？

当两种方案难分高下时，技能文档给出了一组可操作的 A/B 对比流程：

1. 从干净 diff 出发，在同一构建环境分别构建每个候选方案；

2. 记录受影响的 CUDA 目标文件与库的体积变化：

   stat -c '%s %n' build/aten/src/ATen/CMakeFiles/torch_cuda.dir/native/cuda/<file>.cu.o torch/lib/libtorch_cuda.so
   

3. 检查 kernel 符号是否因模板实例化而倍增：

   nm -S --size-sort -C torch/lib/libtorch_cuda.so | rg '<kernel_name>|index_t|long|int'
   

4. 如果热循环算术发生了变化，用小张量与张大张量做代表性基准测试；不要拿 sanitizer 运行时的数据当作性能结论。

由此得到默认决策表：

场景	推荐做法
一两个 64 位 setup 乘法	局部 64 位转换（模式 1）
逐元素索引可能超过 32 位	index_t 分派（模式 2）
现有 kernel 家族已在分派 index_t	沿用以有的模式扩展
需要改动很多按 scalar 特化的 kernel	先评估二进制体积，再决定是否全部模板化

五、大索引修复的回归测试

修复索引溢出时，回归测试要恰好跨过当初失败的那条边界：

• 选择仍能触发溢出的最小 dtype 与最小输出尺寸，控制测试成本；

• 断言 tensor.numel() > torch.iinfo(torch.int32).max，或者断言具体的偏移边界值；

• 采样点在边界两侧各取一些值，避免对超大张量做整体 CPU 比对；

• 用 compute-sanitizer 跑一遍原始复现脚本以暴露内存错误：

  CUDA_LAUNCH_BLOCKING=1 PYTORCH_NO_CUDA_MEMORY_CACHING=1 compute-sanitizer --tool memcheck --error-exitcode=99 <python> repro.py
  

在 PyTorch 的测试组织上，技能文档建议把这类回归加在相关的 pooling/indexing 测试附近，并对昂贵用例使用 @largeTensorTest 与相应的 device 装饰器做门控——这与 test/test_cuda.py 中针对 numel() > INT_MAX 的既有测试风格一致。

六、Review 检查清单

合入任何索引宽度修改前，按以下清单逐项核对：

• [ ] 已定位到确切溢出的表达式，而不是全文件无差别改类型；
• [ ] 64 位数学只用于确实需要的表达式；热索引需要时则模板化 kernel；
• [ ] canUse32BitIndexMath 覆盖了每一个偏移使用所选索引类型的张量（含 strided view 输入）；
• [ ] CUDA 启动维度仍满足硬件上限；grid 维度溢出已通过分块/跨步解决；
• [ ] 回归测试在修复前失败、修复后通过，或者报告中说明了为何没有重跑修复前的失败用例；
• [ ] 新增 index_t 分派导致 kernel 实例翻倍时，二进制体积与性能影响已被提及并量化。

小结

PyTorch 对 CUDA 索引宽度的处理形成了一套完整闭环：canUse32BitIndexMath（IndexingUtils.cpp）负责判定“这个张量能不能安全用 32 位索引”，AT_DISPATCH_INDEX_TYPES（Dispatch.h）负责按 Int/Long 双路实例化 kernel，CUDA_KERNEL_LOOP_TYPE（KernelUtils.h）保证 grid-stride 循环在两种宽度下都正确。修复策略上，能局部转换就不模板化，能分派就不无条件 64 位，无法支持就快速失败；验证策略上，用恰好跨过 INT_MAX 边界的回归测试与 compute-sanitizer 双保险。掌握这条从溢出定位、工具选型、代码修改到体积/性能评估的完整链路，是维护 PyTorch CUDA 算子时处理大张量问题的标准工作方式。