已关闭
[Bug-Report|缺陷反馈]: RadixTopK WS 变体 V→MTE3 同步缺失、i64 索引序列越界及跨核前缀和累加溢出 #3604
Nice try创建于  22 天前关闭于  21 天前
Nice try
Nice try成员
22 天前 创建

一、问题描述 (必填)

RadixTopK 算子在 int64 大 shape(sortLen 超 int32 上限,INT32_SORT_LEN_MAX)与 WS 变体路径上存在三处缺陷,且 host 侧 UT 的 tiling key 期望值只编码 3 位,导致 int64 主路径恒测不到:

  1. WS 变体同步缺失(k 配额错乱)
    math/radix_top_k/op_kernel/radix_top_k_ws.h 的 HandleLastRoundBoundary 中,Add 在 V 流水线写 tileTopK,随后 MTE3 通过 CopyTileTopKUb2Ws 读同一块 UB。原代码用 SToMTE3Sync():PipeBarrier<PIPE_V> 只保证 V 流水线内部顺序,SToMTE3Sync 只覆盖 S→MTE3,两者都覆盖不到 V→MTE3,MTE3 可能读到 Add 之前的旧值,造成该段 k 配额错乱。同文件 AddTileHist2TileTopKInWs 对同样的 Add→CopyOut 序列使用的是 VToMTE3Sync。

  2. i64 索引序列越界写
    i64 变体的多段 regime 会命中 tileNum_ > MAX_TILE_NUM_IN_UB。此时 CreateVecIndex4TopK 中 MAX_TILE_NUM_IN_UB - tileNumAlign_ 的无符号相减下溢成巨大值,再被 tileLen_ 兜住,最终 GetWithOffset(maxIndexLen, tileNumAlign_ * 4) 的偏移会写出 6144 元素 buffer。

  3. 跨核前缀和累加溢出
    前驱 core 边界 bin 累计元素数(CopyOutBoundaryBinCumSum → ComputeCumSumPrev)总量可达 sortLen > 2^31(int64 形状),而 UB 变体的 cumSumValuePrev / remainCoreBoundaryNum 为 int32;WS 变体则另有一份只差累加器类型的 ComputeCumSumPrevI64 实现。

  4. UT 覆盖缺口
    math/radix_top_k/tests/ut/op_host/test_radix_top_k_tiling.cpp 的 ComputeExpectedTilingKey 只编码 sorted/largest/isLargeShape 3 位,未编码 isInt64Shape(bit3),int64 主路径(sortLen > 2^31)恒测不到。

二、环境信息 (可选)

  • 昇腾硬件型号:ASCEND910B、ASCEND910_93(radix_top_k 配置目录仅含 ascend910b/ascend910_93)
  • CANN版本:(可补充)
  • 操作系统:(可补充)
  • 其他相关软件版本:(可补充)

三、重现步骤 (可选)

  1. 平台为 ASCEND910B 或 ASCEND910_93;
  2. 调用 RadixTopK:输入 x 为 FLOAT16/BFLOAT16,最后一维 sortLen > 2^31-1(INT32_SORT_LEN_MAX),indices_dtype = DT_INT64;
  3. 命中 WS + int64 变体(isLargeShape_ || isInt64Shape_ ⇒ needWorkspace=true),K 配额需在末轮边界 bin 处补齐;
  4. 观察 HandleLastRoundBoundary 中 V→MTE3 之间的同步与 CreateVecIndex4TopK 的 buffer 偏移(tileNumAlign_ 达 MAX_TILE_NUM_IN_UB 时)。

四、预期结果 (可选)

  • V 写 tileTopK 与 MTE3 读取之间同步正确(VToMTE3Sync),k 配额不再错乱;
  • CreateVecIndex4TopK 在不满足不变量时 return 0,下游走「现算索引」分支,不再发生越界写;
  • 跨核前缀和统一按 int64 累加(UB/WS 共用基类实现),int64 大 shape 下不再溢出;
  • host UT 覆盖 int64 主路径(tilingKey bit3)与 indices 非 int64 时 tiling 失败。

💡 备注(选填)
修复方式见关联 PR:包含三处缺陷修复(VToMTE3Sync、CreateVecIndex4TopK 前置防御、ComputeCumSumPrev 统一 int64)、static_assert(LARGE_SORT_LEN < INT32_SORT_LEN_MAX) 不变量固化、int64 主路径 host UT 补充,以及 README 约束说明(输出为选中元素在排序轴上的索引序,不保证有序)与注释更新。

关联 PR:https://gitcode.com/cann/ops-math/pull/6396

likedislike
Nice tryNice try成员
22 天前 添加了label:bug-report
Nice try
Nice try成员
22 天前 评论:

/assign @Nice_try

likedislike
CANN-robotCANN-robot成员
22 天前 将 Nice_try 设为负责人
CANN-robotCANN-robot成员
21 天前 关闭了 issue
CANN-robotCANN-robot成员
21 天前 添加了label:resolved