Closed
[Bug-Report|缺陷反馈]: RadixTopK WS 变体 V→MTE3 同步缺失、i64 索引序列越界及跨核前缀和累加溢出 #3604
Nice tryAdd on 23 days agoClosed at 22 days ago
23 days ago Added label: bug-report
23 days ago Assign Nice_try as responsible person
22 days ago Close issue
22 days ago Added label: resolved


一、问题描述 (必填)
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