fix(cuda): guard CPU index tensors in index/_unsafe_index boxing

[ethanZZZZZZZZZZZZ]

问题

Fixes #86

boxing 路线上,aten::index.Tensor / aten::_unsafe_index.Tensor 的索引列表中一旦含有 CPU 张量(PyTorch 高级索引的合法用法,PyTorch 会自行把 CPU 索引搬到设备侧),生成的 CUDA boxing kernel 会无条件对它做零拷贝元数据改写(BoxToCuda),把一个存储在 host 内存上的张量伪装成 cuda:0,厂商 kernel 设备检查直接拒绝:

RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cuda:0)

影响真实模型路径:transformers 推测解码(generation/utils.py 中 q[:, torch.arange(candidate_length), new_candidate_input_ids],arange 在 CPU、ids 在设备上),导致 HF 官方 UT 回归:

tests/models/whisper/test_modeling_whisper.py::WhisperStandaloneDecoderModelTest::test_assisted_decoding_sample

同一代码点还有第二个缺陷:at::index 抛异常(越界、索引 dtype 不合法等)时,已 box 的张量不会被还原,伪造的 CUDA 元数据永久残留——同进程后续所有涉及该张量的 op 连环错乱,且在 pytest 渲染失败 repr 时直接段错误(torch/_tensor_str.py → torch/_tensor.py __iter__),失败信息丢失、批量测试结果不可信。

根因

scripts/codegen_ops.py 的 gen_optlist 模板(仅 index.Tensor/_unsafe_index.Tensor 两个 op 走 special_optlist 类别)生成的 box 循环条件为:

if (opt.has_value() && opt->defined()) { BoxToCuda(*opt); ... }

缺少 is_privateuseone() 检查。BoxToCuda(csrc/aten/device_boxing.h:50-52)是无条件改写 TensorImpl 设备字段的裸原语;仓库其余调用点(DeviceBoxingGuard::maybe_box、BoxTensorListToCuda)都带该守卫,唯此遗漏。同时 api(...) 之后的 UnboxToFlagos 序列写在正常返回路径上,无异常安全。

修复

• box 循环仅作用于 is_privateuseone() 的张量;CPU 索引原样透传,由 PyTorch 按原生语义处理(与 DeviceBoxingGuard::maybe_box 的既有惯例对齐);
• 调用包进 try/catch(...),失败路径同样还原全部已 box 张量的元数据;
• csrc/aten/generated/cuda_kernels.cc 为重新生成产物(diff 仅 IndexTensorKernelCuda / PrivUnsafeIndexTensorKernelCuda 两个函数),生成器两次运行幂等。

验证(MetaX C550 ×8,driver 3.8.1,MACA 3.8.0.23,torch 2.10.0+cpu boxing 环境,transformers v5.12.1)

验证项	修复前	修复后
最小复现(CPU/flagos/混合索引、_unsafe_index、异常后元数据卫生共 7 项)	前 5 项 FAIL + 同进程连环污染	7/7 PASS,数值与 CPU 对拍逐位一致
WhisperStandaloneDecoderModelTest::test_assisted_decoding_sample	FAIL	PASS
whisper 全部 assisted decoding UT(-k assisted,6 个)	—	6 passed
环境冒烟(factory/linear/add/mean/mm/Module.to)	—	PASS

备注

• 本 PR 仅含 scripts/codegen_ops.py 模板改动 + 重新生成的 csrc/aten/generated/cuda_kernels.cc;其余生成产物(flaggems_python_kernels.cc、*_flaggems.conf)未变。
• 其余生成器类别走 DeviceBoxingGuard(RAII)或 BoxTensorListToCuda,天然具备守卫与异常安全,本次修复使 gen_optlist 这一特例对齐。

[lvyufeng]

绑定issue