Pull Request已成功合入, 合并人@ascend-robot
(感谢 QianZH97 的贡献)变更摘要
本 PR 为源自 FBGEMM 的 triton_dense_to_jagged 算子新增 GPU 端 ATK 测试迁移:该算子将 dense tensor 转为 jagged tensor,并支持 add/mul 元素级融合。改动新增 Triton kernel 实现与基于 BaseApi 的 ATK 测试 API 封装(TritonDenseToJaggedApi),覆盖 NPU/CUDA/CPU 设备选择、jagged offsets 生成与 kernel 调用参数配置。
主要改动
- 新增 Triton kernel 实现:新增
triton_dense_to_jagged内核及融合函数tensor_elementwise_add/tensor_elementwise_mul,按 offsets 区间从 dense 读取数据,通过jagged_mask/dense_mask掩码写回 jagged 输出,并在operation_function为 "add"/"mul" 时与operation_jagged_value_ptr做元素级融合计算。 - 新增 dense index 推导辅助函数:新增
_jagged_offsets_to_dense_indice,对多维 jagged tensor 通过二分查找将 offsets 递归换算为对应的 dense 索引(越界情况返回 -1),作为 kernel 的dense_indices输入。 - 新增 ATK API 封装类:新增
TritonDenseToJaggedApi(@register("triton_dense_to_jagged")),实现init_by_input_data与__call__,解析dense、jagged_offsets、operation_function等输入,生成输出 jagged tensor,并计算 grid、stride、thread block 尺寸等 kernel 启动参数。 - 新增 offsets 约束生成逻辑:新增
_generate_jagged_offsets,针对 3D/4D dense 确定性生成满足约束的 1 层/2 层 offsets(首元素为 0、末值落在约束区间、相邻差值不超过上限),保证不同机器上结果一致。


| 阶段 | 任务名 | 状态 | 详情 |
|---|---|---|---|
| 编译构建 | PR-build-SAST-check | ✅ COMPLETED | >>> |
| Build-2.7.1_arm | ✅ COMPLETED | >>> | |
| Build-2.10.0_arm | ✅ COMPLETED | >>> | |
| Build-2.7.1_X86 | ✅ COMPLETED | >>> | |
| Build-2.10.0_X86 | ✅ COMPLETED | >>> | |
| 恶意代码检查 | Antipoison | ✅ COMPLETED | >>> |
| 编码安全与规范检查 | pre-commit | ✅ COMPLETED | >>> |
| clang-tidy | ✅ COMPLETED | >>> | |
| 开源片段检查 | SCA | ✅ COMPLETED | >>> |
| 开发者测试 | presmoke | ✅ COMPLETED | >>> |
| 流水线 | PR-pipeline_ops-rec | ✅ COMPLETED | >>> |
- compile : 运行流水线
- retry : 重试流水线所有失败子任务
- retry <任务名> : 仅重试指定失败子任务
- stop : 停止流水线


Thanks for your pull-request.
The full list of commands accepted by me can be found at here.
You can get sig-info at here.
You can self-configure the PR merge rules for this repository. For more details, please refer to Here.
PR Approval Progress
✅ Congratulations! All modules have met the lgtm and approve requirements.
Module Approval Details
| module | lgtm status | approve status |
|---|---|---|
| repo-Ascend/ops-rec | ✅ 文一晴, li-da-xiong, m0_57292843 (3/2) | ✅ li-da-xiong (1/1) |
💡 Tip:
- Committer can comment
/approveor/lgtm- Commenting
/approveimplies both code review (lgtm) and intent to merge (approve)


ascend docs pipeline is running...


✅ 跳过 docs ci 检查,没有需要检查的文档文件


/lgtm


/lgtm


[P2] offset 输入值被丢弃,仅按 dense shape 重新生成
问题
_generate_jagged_offsets(dense, jagged_offsets, device) 函数中,用户传入的 jagged_offsets 仅用于取 dtype/device,其数值完全被按 dense shape 重建的确定性 offset 覆盖。
影响
用户在 YAML 中配置的特定 offset 分布不生效,测试无法覆盖非均匀序列长度场景。
建议
应直接使用传入的 offset tensor(仅做 device 转换和约束校验),仅在 offset 为空时才自动生成。


[P2] block_offset 使用 dense stride 索引 jagged 指针,stride 不一致时写入位置错误
问题
block_offset = offset_row * dense_row_stride + offset_col * dense_col_stride 同时用于 dense load(正确)和 jagged store(tl.store(jagged_value_ptr + block_offset, ...))。
但 jagged tensor 的行 stride 是 jagged_value_row_stride(= N),与 dense_row_stride 不一定相等。当 dense 非连续或 shape 不匹配时,jagged store 的写入位置会错位。
影响
在 contiguous 且 D 相同的场景下两者一致,不影响结果。但 kernel 缺乏对 stride 一致性的校验,存在潜在错误风险。
建议
jagged store 应使用 jagged_value_row_stride 计算行偏移,或添加断言 assert jagged_value_row_stride == dense_row_stride。


/approve


Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.


Pull Request 已合并或已关闭。
If you want to solve this problem, you can click here to do it in the FAQs.


合入来源
问题/功能描述
[feat] 新增triton_dense_to_jagged算子atk测试迁移(gpu部分)
修改方案描述
triton_dense_to_jagged 算子源自 FBGEMM,将 dense tensor 转为 jagged tensor。含 GPU 端 ATK 精度与性能测试自定义 API 封装,支持 add/mul 元素级融合。
是否涉及UT/ST
开发自检