问题描述

在 cann/ops-math 的固定源码版本中，TopKV2 的 Host InferShape 只校验输入 rank
和 dim，没有校验动态输入 k 是否为非负值且不超过排序轴长度。仓库的
top_k_v2_proto.h 要求输入最后一维至少为 k，同仓 TopK/TopKV2D 文档也规定
0 <= k <= x[dim]；AICPU 运行时还会对这两个非法方向返回参数错误。

然而，Host InferShape 对 x=[2,32] 接受 k=33 和 k=-1，并分别生成
[2,33]、[2,-1] 的两个输出 shape。这样非法输入在图形状阶段被标记为成功，后续
阶段才可能失败或携带非法 shape 继续处理，Host 与运行时的参数契约不一致。

环境信息

• 仓库/commit：cann/ops-math / 0a2ce5b57caec6068d9e5658b740c2d41482aa15
• OS/架构：Linux 5.10.0/aarch64
• Ascend 硬件：Ascend 910B3（本次仅验证 Host InferShape）
• Driver/Toolkit/OPP：CANN 9.1.0；本结论限定为固定源码 Host 行为
• 编译器和 Python：GCC/G++ 11.4.0；Python 3.12.13；CMake 3.22.1

重现步骤

1. 检出 0a2ce5b57caec6068d9e5658b740c2d41482aa15。
2. 按仓库 Host UT 配置构建 top_k_v2：启用 ENABLE_TEST=TRUE、OP_HOST_UT=TRUE、
   ASCEND_COMPUTE_UNIT=ascend950 和 ENABLE_UT_EXEC=TRUE。
3. 在 TopKV2 注册 InferShape 回调上构造输入 x shape [2,32]、标量 k，
   dim=-1，输出预留两个 Tensor。
4. 分别运行 k=32、k=33 和 k=-1，执行过滤器
   TopKV2InferShape.*。

预期结果

k=32 应返回 GRAPH_SUCCESS 并生成 [2,32]。对于 k=33 和 k=-1，应返回
GRAPH_FAILED 或等价参数错误，不应生成超过输入轴长度或为负的输出维度。

实际结果

Host InferShape 返回 GRAPH_SUCCESS：

• k=32：输出 [2,32]；
• k=33：输出 [2,33]；
• k=-1：输出 [2,-1]。

固定源码 Host UT 构建运行 12/12 通过；重复过滤运行 4/4 通过，退出码均为 0。非法
用例的测试日志明确记录为 actual_contract_probe=GRAPH_SUCCESS。

日志/最小复现

源码定位：math/top_k_v2/op_host/top_k_v2_infershape.cpp:25-67，其中
InferShapeForTopKCommon 在 :58-62 直接将 k 写入输出 shape。约束定位：
math/top_k_v2/op_graph/top_k_v2_proto.h:29；运行时参数检查可见
math/top_k_v2_d/op_kernel_aicpu/top_k_v2_d_aicpu.cpp:551-557。

关键结果：

validation_case k=32 expected=GRAPH_SUCCESS
validation_case k=33 expected=GRAPH_FAILED actual_contract_probe=GRAPH_SUCCESS output_shape=[2,33]
validation_case k=-1 expected=GRAPH_FAILED actual_contract_probe=GRAPH_SUCCESS output_shape=[2,-1]
test_exit_code=0

影响与规避方式

影响是非法参数未在形状推导阶段被拒绝，可能使后续图校验、内存规划或运行期才暴露
错误，增加定位成本并造成 Host/运行时行为差异。临时规避是调用方在提交 TopKV2 前
自行保证 0 <= k <= x[dim]；不应依赖 Host InferShape 来替换该校验。

建议的回归测试

补充 k=0、k=dim_size、k=dim_size+1、k=-1 和 dim 为负数的 Host InferShape
测试，要求合法边界成功、两个非法方向明确返回失败，并同时覆盖 INT32/INT64 的 k
输入路径（适用于支持该类型的 TopKV2 变体）。

关联 PR

• https://gitcode.com/cann/ops-math/merge_requests/4987