已合并
[feat] 新增triton_dense_to_jagged算子atk测试迁移(gpu部分) #128
[feat] 新增triton_dense_to_jagged算子atk测试迁移(gpu部分) #128
已合并
QianZH97创建于 9月11日
QianZH97
QianZH97成员
9月11日

合入来源

问题/功能描述

[feat] 新增triton_dense_to_jagged算子atk测试迁移(gpu部分)

修改方案描述

triton_dense_to_jagged 算子源自 FBGEMM,将 dense tensor 转为 jagged tensor。含 GPU 端 ATK 精度与性能测试自定义 API 封装,支持 add/mul 元素级融合。

是否涉及UT/ST

开发自检

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 QianZH97 的贡献)
QianZH97QianZH97成员
9月11日 关联了里程碑:ops-rec 26.2.0
atomgit-bot
atomgit-bot
9月11日 评论:

变更摘要

本 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、末值落在约束区间、相邻差值不超过上限),保证不同机器上结果一致。
likedislike
不准确?
atomgit-bot
atomgit-bot
9月11日 评论:

代码审查

✅ 未发现问题

likedislike
不准确?
ascend-robot
ascend-robot成员
9月11日 评论:
流水线 PR-pipeline_ops-rec#613 [ commitID:7f79d484 ] 已完成
>>>代码风格自动修复执行成功(无修复内容)
阶段 任务名 状态 详情
编译构建 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 >>>
此流水线已支持下列评论快捷指令,仅PR创建者和白名单成员评论有效
  • compile : 运行流水线
  • retry : 重试流水线所有失败子任务
  • retry <任务名> : 仅重试指定失败子任务
  • stop : 停止流水线
likedislike
ascend-robotascend-robot成员
9月11日 添加了label:ascend-cla/yes
ascend-robot
ascend-robot成员
9月11日 评论:

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 /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)
likedislike
ascend-robotascend-robot成员
9月11日 添加了label:ci-pipeline-running
ascend-robot
ascend-robot成员
9月11日 评论:

CLA Signature Pass

QianZH97, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
ascend-robot
ascend-robot成员
9月11日 评论:

ascend docs pipeline is running...

likedislike
ascend-robotascend-robot成员
9月11日 添加了label:docs-ci-pipeline-running
ascend-robot
ascend-robot成员
9月11日 评论:

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

likedislike
ascend-robotascend-robot成员
9月11日 删除了label:docs-ci-pipeline-running
ascend-robotascend-robot成员
9月11日 添加了label:docs-ci-pipeline-success
ascend-robotascend-robot成员
9月11日 删除了label:ci-pipeline-running
ascend-robotascend-robot成员
9月11日 添加了label:ci-pipeline-passed
iinoiino
iinoiino成员
9月11日 评论:

/lgtm

likedislike
wenyiqing
wenyiqing成员
9月11日 评论:

/lgtm

likedislike
ascend-robotascend-robot成员
9月11日 添加了label:lgtm
zengxiong成员9月11日进行代码检视2
experimental/triton/atk_test/triton_dense_to_jagged/gpu_atk_test/triton_dense_to_jagged_api.py
@@ -0,0 +231,4 @@
231+ jagged_offsets.append(torch.tensor(offset_tensor, device=device))
232+ 
233+ # ==================== 生成满足约束的 jagged_offsets ====================
234+ jagged_offsets = _generate_jagged_offsets(dense, jagged_offsets, device)
zengxiong9月11日评论:

[P2] offset 输入值被丢弃,仅按 dense shape 重新生成

问题
_generate_jagged_offsets(dense, jagged_offsets, device) 函数中,用户传入的 jagged_offsets 仅用于取 dtype/device,其数值完全被按 dense shape 重建的确定性 offset 覆盖。

影响
用户在 YAML 中配置的特定 offset 分布不生效,测试无法覆盖非均匀序列长度场景。

建议
应直接使用传入的 offset tensor(仅做 device 转换和约束校验),仅在 offset 为空时才自动生成。

likedislike
QianZH97
QianZH97成员
9月12日 评论:

这是因为atk工具基于yaml和generate限制文件生成的jagged_offsets只会确定shape和数据类型,值只有在实际测试的时候才会随机生成,不可用(jagged_offsets必须是累加和)。所以只有在api文件中生成jagged_offsets的值

zengxiong成员9月11日进行代码检视2
experimental/triton/atk_test/triton_dense_to_jagged/gpu_atk_test/triton_dense_to_jagged.py
@@ -0,0 +107,4 @@
107+ 
108+ for _i in range(begin, end, thread_block_row_size):
109+ offset_col = tl.arange(0, thread_block_col_size)
110+ block_offset = offset_row[:, None] * dense_row_stride + offset_col[None, :] * dense_col_stride
zengxiong9月11日评论:

[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。

likedislike
QianZH97
QianZH97成员
9月12日 评论:

这是FBGEMM源码,本次只迁移源码,不对内容改动

li-da-xiong成员
9月12日 评论:

/approve

likedislike
ascend-robotascend-robot成员
9月12日 添加了label:approved
ascend-robotascend-robot成员
9月12日 合入了pull request
ascend-robot
ascend-robot成员
9月12日 评论:

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

likedislike
ascend-robot
ascend-robot成员
9月12日 评论:

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

likedislike