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

Backgroud（背景信息）

MiniMax M3模型中，稀疏注意力MSA（Minimax Sparse Attention）模块的 Indexer 支路用轻量级的 index Q/K 为每个 KV sparse block 打分，再由 topk 选块喂给正式的 sparse attention。其中的“打分”计算流程当前无高性能算子实现，只能通过小算子拼接实现，性能较差且在模型中占比高，存在使用AscendC开发融合算子实现上述计算过程的性能优化空间。

Benefit / Necessity （价值/作用）

价值：开发融合算子实现，可优化性能，提升推理速度。

应用场景：Minimax M3相关模型部署场景。

Design（设计方案）

算子名：MsaIndexScore

aclnn接口名：aclnnMsaIndexScore

torch_extension：msa_index_score

具体实现路径为：基于当前ops-transformer仓已引入的catlass三方库，进行模板化组装，最终以对应功能的torch接口形式交付调用。

index_q ──┐
          ├─► [本算子 MsaIndexScore] ─► score[H, T, NB] ─► torch.topk ─► topk_idx ─► sparse attention
index_k ──┘  (QK^T + mask + 分段 max)                             ↑
  (已写入 paged cache)                                        本算子不覆盖

算子参数说明

$T$ = total_q（varlen 扁平后的 query token 总数），$B$ = batch，
$H_q$ = num_q_heads（index heads），$H_{kv}$ = num_kv_heads，$D$ = head_dim，
$P$ = block_size，$NP$ = KV cache 的物理 page 数，$MB$ = 每 batch 最大 block 数。

计算过程

设：

• 请求（request）索引 $b \in [0, B)$；
• q_len 以 cumsum 形式给出，$\mathrm{cu}_q[b]$ 为第 $b$ 个请求首个 query token 的全局扁平下标，$L_q^{(b)} = \mathrm{cu}_q[b{+}1] - \mathrm{cu}_q[b]$；
• start_loc[b] 为该请求首个 query token 在自身 KV 序列中的位置（等价于 Triton 的 prefix_lens / context_lens）；
• kv_len[b] 为该请求当前可见的 KV 总长度；
• $H_q$ 为 query/index head 数，$H_{kv}$ 为 KV head 数，GQA 组大小 $g = H_q / H_{kv}$，$h_{kv} = \lfloor h / g \rfloor$；
• block_size $= P$（默认 128），block 数 $NB^{(b)} = \lceil kv\_len[b] / P \rceil$。

对全局 token $t$（属于请求 $b$）、head $h$、逻辑 block $\beta$：

$$q_{\text{pos}}(t) = \texttt{start\_loc}[b] + \bigl(t - \mathrm{cu}_q[b]\bigr)$$

$$\mathrm{score}[h, t, \beta] \;=\; \max_{\substack{j \in [\beta P,\ \beta P + P) \\ j \le q_{\text{pos}}(t) \\ j < kv\_len[b]}} \Bigl\langle \mathbf{q}_{t,h},\ \tilde{\mathbf{k}}_{j, h_{kv}} \Bigr\rangle$$

其中非量化：$\tilde{\mathbf{k}} = \texttt{key\_cache}[\texttt{block\_table}[\ldots]]$（与 query 同 dtype）。

int8 前融合量化（MSA 预量化 K cache）：

$$\tilde{\mathbf{k}}_{j,h_{kv}} = \mathrm{cast}_{\mathrm{fp}}\bigl(\texttt{key\_cache}^{\mathrm{int8}}[\mathrm{page}, j\bmod P, h_{kv}, :]\bigr) \cdot \texttt{scale}[\mathrm{page}, h_{kv}, j\bmod P]$$

scale shape = $[NP, N_{kv}, P]$（per-token + per-head 反量化）。实现上：
AIV 将 int8 page DataCopy+Cast 到 GM scratch → FFTS K_READY → AIC Mmad(Q, cast(K)) →
S_READY → AIV 在 mask 前对 S 做列乘 scale（因 $Q\cdot(K s)=(Q\cdot\mathrm{cast}(K))\cdot s$）。

若窗口内无任何可见 $j$，则 $\mathrm{score} = \texttt{fill\_value}$（默认 $-\infty$）。

核心运算 = 矩阵乘（Mmad）+ 段长为 $P$ 的行方向分段最大值（ReduceMax）。顺序：DOT →（可选 DEQUANT 列乘）→ MASK → MAX。