Backgroud（背景信息）

ArgMaxGrad 在 Ascend950/arch35 上按不带 D 的原型交付后，轴下标（A2 由融合 pass 传入的 assist 张量）改为内核自生成。当前实现存在两个可优化点：

1. 轴下标每段都要物化到 UB（每元素 4 字节 int32）再读回参与比较，既占 UB 又多走一趟写；UB 被占满后段长（colsPerChunk）被压小。
2. BUFFER_NUM 固定为 1，双缓冲未开启，MTE2 搬入 → V 计算 → MTE3 搬出完全串行；实测大形状每核也只有 3~8 段，没有任何流水重叠。

另外 UB 记账采用「先按每元素字节数估一个段长，再回头断言没超」的写法，一旦记账口径与内核 InitBuffer 不一致就会静默踩 UB（本算子此前的 UB 越界缺陷即出于此）。

Origin（信息来源）

昇腾算子 arch35（Ascend950）迁移交付线自验发现，非外部用户反馈。

Benefit / Necessity （价值/作用）

Ascend950PR 实测（ttk -d false -b release，--run 12）：

• 大形状 6 例：中位 1.13x，最好 1.35x；
• 小形状 8 例：中位 1.05x，最好 1.42x；
• 收益集中在 inner>1 的三档（逐行直算 / 紧排铺 tile / 按行补齐），inner==1 基本持平。

同时 UB 记账改为按容量反解后，段长由「Total() <= ubSize 的最大整宽解」直接求出，消除了记账口径与 InitBuffer 不一致导致静默踩 UB 的一类风险。

Design（设计方案）

1. VF 入口引入轴下标来源档 AssistSrc：
   • ARANGE：inner==1 时本段沿被选轴连续，Reg::Arange(reg, kStart + 车道偏移) 一条指令生成；
   • SCALAR：inner>1 的单行段整段同一个 k，Reg::Duplicate 进寄存器；
   • UB：仅保留「一个寄存器块跨多行」的紧排/补齐两档（k 在寄存器内算不出闭式，仍需物化）。
2. tiling 侧 i32BufBytes 按档记账：inner==1 或逐行直算档记 0。
3. 用腾出的 UB 把 BUFFER_NUM 从 1 提到 2，使 MTE2 与 V 重叠；fp16 每元素占用 ~8.1B → ~4.1B → ~8.2B，UB 总量基本持平。
4. UB 记账改为 UbLayout（字段与 InitBuffer 调用一一对应）+ SolveColsPerChunk 二分反解。

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