Closed
[Requirement|需求建议]: aclnnTopk 支持 sortLen 超 int32 上限的 RadixTopK int64 索引 #6667
Nice tryAdd on  29 days agoClosed at  28 days ago
Nice try
Nice tryCollaborators
29 days ago Create

Backgroud(背景信息)
aclnnTopk 接口当前在排序轴长度 sortLen 超过 int32 上限(2147483647)时,RadixTopK 判定直接不支持,超大 shape 只能回落到 TopKV2 排序方案;而 TopKV2 在该量级会产生约 59.5GB workspace 导致 OOM。RadixTopK 已提供 int64 大 shape 能力,但 aclnn 层尚未接入。

Origin(信息来源)
(待补充:提出该需求的部门/团队)

Benefit / Necessity (价值/作用)
使 aclnnTopk 在 sortLen 超过 int32 上限时走 RadixTopK int64 模板,避免 TopKV2 在该量级约 59.5GB 的 workspace OOM,补齐超大 shape 下 int64 索引的支持能力。

Design(设计方案)

  • 使能方式:Aclnn 直调(index/gather_elements/op_api/aclnn_topk.cpp)。
  • IsRadixTopKSupported 增加 indicesDType 入参,sortLen > int32 上限时要求 indices 为 int64 且 k 比例不设限,即判为支持。
  • positiveDim == lastDim 分支对 int64 索引直接透传 RadixTopK 输出,绕过 Sort/GatherElementsV2。
  • 支持硬件:跟随 RadixTopK(ascend910b / ascend910_93)。

关联 PR:https://gitcode.com/cann/ops-nn/pull/10291

likedislike
Nice tryNice tryCollaborators
29 days ago Added label: requirement
Nice try
Nice tryCollaborators
29 days ago Comments:

/assign @Nice_try

likedislike
CANN-robotCANN-robotCollaborators
29 days ago Assign Nice_try as responsible person
CANN-robotCANN-robotCollaborators
28 days ago Close issue
CANN-robotCANN-robotCollaborators
28 days ago Added label: resolved