Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

一、背景信息 （必填）

Ascend 950 上的 Sort 算子在非末轴小轴场景中可通过 no-transpose 路径避免前后转置，但现有 two-stage 路径主要按相邻 inner 位置分批。当 innerSize 较小、outerSize 较大或 batch 分布不均时，单批规模和活跃核数不足，无法充分利用多核；非连续 GM 数据的逐段搬运及 int64 索引 UB 占用也会限制 batch 大小。部分极端 shape 下，no-transpose 路径的 gather/scatter 固定开销还可能高于成熟的 transpose 路径。

本需求面向有收益的非末轴小轴 shape，新增与末轴 Sort、KthValue 隔离的专项 two-stage 调度，并完善 host 侧批处理、UB 规划和 aclnn 路由收益判断。

二、价值/作用 （必填）

1. 对适合 no-transpose 的非末轴小轴场景，直接处理原始 [outer, axis, inner] 布局，减少前后转置带来的额外 GM 搬运。
2. 支持按完整 outer slice 分组批处理，优先铺满可用核，并在同等核数下减少每核批次数和空闲槽位。
3. schId 11 使用 uint32 轴内索引作为 UB 中间表示，仅在写回 GM 时转换为输出索引类型，降低 UB 占用并扩大可选 batch。
4. 对 FP32 Sort32 场景搜索更宽的 inner tile，并为 merge 路径增加 bank rotation padding，改善 UB 利用与访存冲突。
5. 对收益不足的超小 inner row、大整数轴和超大输出继续选择 transpose 路径；KthValue、末轴 Sort 及已有 fallback 路径保持原有行为。

三、设计方案 （必填）

3.1 使能方式（涉及哪些框架：如Aclnn直调、Pytorch训练等）

• aclnnSort 直调：op_api 根据 shape、dtype 和输出规模判断是否启用 no-transpose。
• Sort GE/kernel 通路：host tiling 为有收益的非末轴 two-stage 场景选择新增 schId 11。

3.2 总体设计

3.2.1 算子支持的数据类型

本次优化不改变 Sort 的接口与既有 dtype 范围。Ascend 950 kernel 路径沿用现有支持类型：int8、int16、int32、int64、uint8、uint16、uint32、uint64、float16、float32、bfloat16；输出索引沿用 int32/int64 配置。aclnnSort 的 bool 输入继续使用既有 uint8 转换路径。

3.2.2 host侧设计

1. 新增 Sort 专用非末轴 two-stage 调度 schId 11。末轴和物理连续的 innerSize == 1 场景继续使用 schId 6；KthValue 继续使用原有批处理策略，避免新策略改变既有二进制行为。
2. 当一个目标 batch 能容纳至少两个完整 outer slice 时，令 batchSize = outerSlicesPerBatch * innerSize，通过 keyParams4 下发每批 outer slice 数，并将 innerLoopNum 设为 1。
3. 非分组场景采用“优先更多活跃核，其次更少每核批次，最后更少空闲槽位”的 batch 搜索策略；分组场景保留完整 outer slice 边界。
4. schId 11 的最终索引 UB 元素宽度按 uint32 估算；FP32 且 axisLen 不超过 Sort32 阈值时，优先尝试 32、16、8 等更宽 inner chunk，并由 UB 估算逐项校验。
5. merge 型非末轴路径对 4 个 UB block 整数倍的 value/index 行步长增加一个 block padding，轮转相邻行起始 bank。
6. aclnn no-transpose 收益判断增加以下回退条件：
   • inner row 小于一个搬运 block 且 outerSize 不小于 1024；
   • 整数 dtype 且 axisLen 不小于 512；
   • values 与 int64 indices 的估算输出总量不小于 1 GiB。

3.2.3 kernel侧设计

1. 新增 SortGroupedOuterSmallAxisTwoStage，复用既有 two-stage 排序基类，只隔离非末轴 GM 映射和批处理。
2. grouped batch 将 GM [outer, axis, inner] 聚集为 UB [segment, axis]：inner 行满足 block 对齐时使用三维 NDDMA，否则使用 SIMT gather。
3. 非分组的 strided batch 使用二维 NDDMA 完成 [axis, inner] 到密集 UB 的搬运。
4. 排序完成后通过 SIMT scatter 恢复原始 GM 布局；grouped 路径在 UB 中保存 uint32 轴内索引，并在写回时转换为实际输出索引类型。
5. 为 MTE2、Vector、SIMT 和 MTE3 生产消费关系补充事件同步，确保 UB 生命周期和写回顺序正确。

3.3 支持硬件

Ascend 950PR / Ascend 950DT。

3.4 算子约束限制

• 仅在 RegBase/Ascend 950 路径启用。
• aclnn no-transpose 候选要求非末轴、维数不超过 8，且 axisLen 位于 [2, 2048]；收益判断不满足时回退前后 transpose。
• grouped outer 仅用于可按完整 outer slice 分批的 two-stage 场景；无法满足 UB、batch 或并行度约束时回退非分组 two-stage 或既有非末轴 merge/radix 路径。
• 不新增接口属性，不改变升序/降序及 values/indices 输出语义。

💡 备注（选填）

补充 Sort host tiling UT 覆盖 strided two-stage、innerSize == 1 连续布局、完整 outer slice 分组、紧凑索引和 FP32 wide-inner-tile，并补充 KthValue 回归用例验证其原有 two-stage 路由不受影响。