RadixTopK 算子在 int64 大 shape(sortLen 超 int32 上限,INT32_SORT_LEN_MAX)与 WS 变体路径上存在三处缺陷,且 host 侧 UT 的 tiling key 期望值只编码 3 位,导致 int64 主路径恒测不到:
INT32_SORT_LEN_MAX
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。
math/radix_top_k/op_kernel/radix_top_k_ws.h
HandleLastRoundBoundary
Add
tileTopK
CopyTileTopKUb2Ws
SToMTE3Sync()
PipeBarrier<PIPE_V>
SToMTE3Sync
AddTileHist2TileTopKInWs
VToMTE3Sync
i64 索引序列越界写 i64 变体的多段 regime 会命中 tileNum_ > MAX_TILE_NUM_IN_UB。此时 CreateVecIndex4TopK 中 MAX_TILE_NUM_IN_UB - tileNumAlign_ 的无符号相减下溢成巨大值,再被 tileLen_ 兜住,最终 GetWithOffset(maxIndexLen, tileNumAlign_ * 4) 的偏移会写出 6144 元素 buffer。
tileNum_ > MAX_TILE_NUM_IN_UB
CreateVecIndex4TopK
MAX_TILE_NUM_IN_UB - tileNumAlign_
tileLen_
GetWithOffset(maxIndexLen, tileNumAlign_ * 4)
跨核前缀和累加溢出 前驱 core 边界 bin 累计元素数(CopyOutBoundaryBinCumSum → ComputeCumSumPrev)总量可达 sortLen > 2^31(int64 形状),而 UB 变体的 cumSumValuePrev / remainCoreBoundaryNum 为 int32;WS 变体则另有一份只差累加器类型的 ComputeCumSumPrevI64 实现。
CopyOutBoundaryBinCumSum
ComputeCumSumPrev
sortLen > 2^31
cumSumValuePrev
remainCoreBoundaryNum
ComputeCumSumPrevI64
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)恒测不到。
math/radix_top_k/tests/ut/op_host/test_radix_top_k_tiling.cpp
ComputeExpectedTilingKey
sorted/largest/isLargeShape
isInt64Shape
isLargeShape_ || isInt64Shape_
return 0
💡 备注(选填) 修复方式见关联 PR:包含三处缺陷修复(VToMTE3Sync、CreateVecIndex4TopK 前置防御、ComputeCumSumPrev 统一 int64)、static_assert(LARGE_SORT_LEN < INT32_SORT_LEN_MAX) 不变量固化、int64 主路径 host UT 补充,以及 README 约束说明(输出为选中元素在排序轴上的索引序,不保证有序)与注释更新。
static_assert(LARGE_SORT_LEN < INT32_SORT_LEN_MAX)
关联 PR:https://gitcode.com/cann/ops-math/pull/6396
/assign @Nice_try
一、问题描述 (必填)
RadixTopK 算子在 int64 大 shape(sortLen 超 int32 上限,
INT32_SORT_LEN_MAX)与 WS 变体路径上存在三处缺陷,且 host 侧 UT 的 tiling key 期望值只编码 3 位,导致 int64 主路径恒测不到: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。i64 索引序列越界写
i64 变体的多段 regime 会命中
tileNum_ > MAX_TILE_NUM_IN_UB。此时CreateVecIndex4TopK中MAX_TILE_NUM_IN_UB - tileNumAlign_的无符号相减下溢成巨大值,再被tileLen_兜住,最终GetWithOffset(maxIndexLen, tileNumAlign_ * 4)的偏移会写出 6144 元素 buffer。跨核前缀和累加溢出
前驱 core 边界 bin 累计元素数(
CopyOutBoundaryBinCumSum→ComputeCumSumPrev)总量可达sortLen > 2^31(int64 形状),而 UB 变体的cumSumValuePrev/remainCoreBoundaryNum为 int32;WS 变体则另有一份只差累加器类型的ComputeCumSumPrevI64实现。UT 覆盖缺口
math/radix_top_k/tests/ut/op_host/test_radix_top_k_tiling.cpp的ComputeExpectedTilingKey只编码sorted/largest/isLargeShape3 位,未编码isInt64Shape(bit3),int64 主路径(sortLen > 2^31)恒测不到。二、环境信息 (可选)
三、重现步骤 (可选)
INT32_SORT_LEN_MAX),indices_dtype = DT_INT64;isLargeShape_ || isInt64Shape_⇒ needWorkspace=true),K 配额需在末轮边界 bin 处补齐;HandleLastRoundBoundary中 V→MTE3 之间的同步与CreateVecIndex4TopK的 buffer 偏移(tileNumAlign_ 达 MAX_TILE_NUM_IN_UB 时)。四、预期结果 (可选)
tileTopK与 MTE3 读取之间同步正确(VToMTE3Sync),k 配额不再错乱;CreateVecIndex4TopK在不满足不变量时return 0,下游走「现算索引」分支,不再发生越界写;💡 备注(选填)
修复方式见关联 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