本次计划开源GMM系列4个算子，lighting_indexer，lighting_indexer_kl_loss算子
场景：GroupMatmul 是 MoE（Mixture of Experts）网络中的核心算子族，将已按专家分组的 token 打包后一次完成所有分组的矩阵乘。MoE 训练链路中需求多样：前向要按专家对 token 分组计算，反向要按 K 维分组求权重梯度并累加，且大模型张量并行下每卡只负责一部分"本地专家"。
概览：面向 MoE 训练链路自研的 4 个精简变体——gmm_local_exp、gmm_local_exp_with_zero（前向）、gmm_add、gmm_k_dim（反向权重梯度），支持 Atlas A2/A3，自带实测性能数据。
gmm_local_exp
gmm_local_exp_with_zero
gmm_add
gmm_k_dim
算子粒度对应关系总览：
grouped_matmul_add
grouped_matmul
两者均实现"分组矩阵乘 + 原地累加"：C[g] += X_g @ W_g。MT 版专为 MoE 反向权重梯度设计（X_g^T 转置求导），社区版为通用累加。
C[g] += X_g @ W_g
X_g^T
group_list
groupList
c
[G, M, N]
yRef
trans_a
transposeX
transposeWeight
aicore_num
两者均支持"按 K 维分组、各组独立输出"：Y[g] = X_g @ W_g（MT 支持 trans_a 走反向 X_g^T @ W_g）。
Y[g] = X_g @ W_g
X_g^T @ W_g
均按 M 维（token 轴）分组做矩阵乘 Y[cuM:cuM+M_g] = X[cuM:cuM+M_g] @ W[g]^T。核心差异在"本地专家子集"语义：
Y[cuM:cuM+M_g] = X[cuM:cuM+M_g] @ W[g]^T
exp_start_idx
exp_end_idx
[LocalExpNum, N, K]
is_b_nz
type_promotion
社区 8 个算子中没有"先清零输出再分组矩阵乘"的变体。MT 版通过 AIV 向量核先将整个输出 Y 清零、再由 AIC 矩阵核执行计算，避免未激活专家位置残留脏数据，省去额外 alloc + zero 的内存与 Kernel 开销。这是社区能力空白，也是本次开源在 MoE 前向场景下的独有补充。
Y
alloc + zero
1. 直击 MoE 训练链路，补齐社区推理/量化导向的空白 社区 gmm 算子族以推理 + 量化为重心；MT 的 4 变体完整覆盖 MoE 训练两端（前向 local_exp 系列 + 反向 gmm_add/gmm_k_dim），且均有实测性能数据。
local_exp
2. 本地专家子集计算能力（社区缺失） exp_start_idx / exp_end_idx 支持张量并行下每卡只计算本卡专家对应 token 段，减少冗余计算与通信——社区 grouped_matmul 无此语义。
3. 输出清零变体（社区无对应） gmm_local_exp_with_zero 以 AIV 预清零 + AIC 计算的双核流水保证数值正确性，是社区能力空白，可直接作为社区新特性回灌。
4. 反向梯度专用 API（单接口切换） torch_npu.gmm 一个接口通过 c 是否传入自动切到 gmm_add（原地累加梯度）或 gmm_k_dim（独立输出），配合 trans_a 原生支持 X_g^T 反向转置、aicore_num 核数控制，开发者零额外配置。
torch_npu.gmm
5. 轻量精简、易移植 无量化/激活融合复杂度，每变体仅几个输入/属性，代码量小、与社区同构（均基于 Ascend C Matmul），可作为社区 grouped_matmul 各 groupType 场景的参考实现低成本回灌。
6. 性能数据透明 README 内置 Atlas A2 实测 MFU：前向最高 79.83%（local_exp）、反向 11%~44%，社区 README 均未提供性能数据。
社区 gmm 族中，gmm_add 与 gmm_k_dim 分别有 grouped_matmul_add 与 grouped_matmul 可对应（社区为通用/量化超集，MT 为训练专用精简实现），gmm_local_exp 有部分对应（缺本地专家子集），gmm_local_exp_with_zero 无对应。本次开源以一套双平台、带实测性能的 MoE 训练专用 gmm 变体，补齐社区在反向梯度、本地专家子集、输出清零三方面的能力空白，可低成本回灌增强社区官方算子。
场景：LightningIndexer（LKI）是稀疏注意力（MoE、NSA 等模型）的路由算子，通过计算 Query 与 Key 的注意力得分，经逐头权重加权归约后做 TopK 排序，输出每个 Query token 最重要的 K 个 Key block 索引及分值，供后续 Sparse FlashAttention 使用。其反向/损失算子 LightningIndexerGradKLLoss 融合 KL 散度损失与对 query_index / key_index / weight 的梯度计算。
LightningIndexerGradKLLoss
query_index / key_index / weight
lightning_indexer
lightning_indexer_v2
lightning_indexer_grad_kl_loss
sparse_lightning_indexer_grad_kl_loss
dense_lightning_indexer_grad_kl_loss
dense_lightning_indexer_kl_loss_grad
MT 版 lightning_indexer_grad_kl_loss 的 def/kernel 文件均以 sparse_lightning_indexer_grad_kl_loss_* 命名，与社区算子同名——两者为同源实现。
sparse_lightning_indexer_grad_kl_loss_*
两者功能高度一致：Indices = Top-k{ [1]₁ₓg @ [(W @ [1]₁ₓSₖ) ⊙ ReLU(Q_index @ K_indexᵀ)] }，输入 query/key/weights/actual_seq_lengths/block_table，属性 layout_query/layout_key/sparse_count/sparse_mode/pre_tokens/next_tokens/return_values。核心差异在 MT 的稀疏路由增强属性：
Indices = Top-k{ [1]₁ₓg @ [(W @ [1]₁ₓSₖ) ⊙ ReLU(Q_index @ K_indexᵀ)] }
query/key/weights/actual_seq_lengths/block_table
layout_query/layout_key/sparse_count/sparse_mode/pre_tokens/next_tokens/return_values
block_len
init_num
local_num
cur_seq_lengths_*
actual_seq_lengths_*
社区 V2 为另一形态的独立实现：输入 q/k/w/cu_seqlens_q/k/seqused_q/k/cmp_residual_k/block_table/output_idx_offset/metadata，属性 topk/max_seqlen_q/layout_q/layout_k/mask_mode/cmp_ratio/return_value，引入 cmp_ratio 压缩采样与 metadata 机制。MT 版与其非同源、功能侧重不同：MT 以 block 粒度聚合 + sink/局部窗口的路由增强为核心，V2 以压缩比/元数据驱动为特色，两者互补。
q/k/w/cu_seqlens_q/k/seqused_q/k/cmp_residual_k/block_table/output_idx_offset/metadata
topk/max_seqlen_q/layout_q/layout_k/mask_mode/cmp_ratio/return_value
cmp_ratio
两者功能一致：给定 query/key（重建 FlashAttention softmax 分布 P）、query_index/key_index（低秩压缩）、weight、sparse_indices、softmax_max/sum，融合计算 KL 散度损失与 dQ_index / dK_index / dW 梯度。MT 的 def/kernel 文件命名与社区完全同名，属同源实现。差异：
dQ_index / dK_index / dW
pre_tokens
next_tokens
block_size
社区另有 dense_lightning_indexer_grad_kl_loss / dense_lightning_indexer_kl_loss_grad：稠密场景下 query/key/query_index/key_index 无需稀疏化处理，是不同于 MT 稀疏版（sparse_indices 驱动）的另一条实现路线，场景互补而非竞争。
1. 社区前向 V1 缺失的稀疏路由增强（MT 独有） block_len（block 粒度聚合）、init_num（sink token 固定保留）、local_num（对角线局部窗口）是稀疏注意力路由的关键能力，社区 lightning_indexer（V1）与 lightning_indexer_v2 均未覆盖。本次开源可直接作为社区前向算子的增强特性回灌，补齐 V1 → 增强版的能力缺口。
2. 反向损失算子与社区同源，可双向验证 MT 的 lightning_indexer_grad_kl_loss 与社区 sparse_lightning_indexer_grad_kl_loss 同源同名，开源后可为社区提供原始参考实现与对照基准，便于社区追溯算法来源、回灌 pre_tokens/next_tokens/block_size 等 MT 侧属性语义。
pre_tokens/next_tokens/block_size
3. 成对交付、训练链路完整 前向（输出 sparse_indices/sparse_values，供 SFA 使用）+ 反向损失（KL + 梯度 + loss）成对提供，天然衔接 Sparse FlashAttention 前反向，构成 NSA/MoE 稀疏注意力训练闭环。
sparse_indices/sparse_values
4. int64 序列长度 cur_seq_lengths_* / actual_seq_lengths_* 采用 int64，支持更长序列与更大 Batch，优于社区前向 V1 的 int32。
/assign @weili_zju
This issue can not be assigned to weili_zju. Please try to assign to the repository members.
背景
本次计划开源GMM系列4个算子，lighting_indexer，lighting_indexer_kl_loss算子
第一部分：GroupMatmul 族
一、背景
场景：GroupMatmul 是 MoE（Mixture of Experts）网络中的核心算子族，将已按专家分组的 token 打包后一次完成所有分组的矩阵乘。MoE 训练链路中需求多样：前向要按专家对 token 分组计算，反向要按 K 维分组求权重梯度并累加，且大模型张量并行下每卡只负责一部分"本地专家"。
概览：面向 MoE 训练链路自研的 4 个精简变体——
gmm_local_exp、gmm_local_exp_with_zero（前向）、gmm_add、gmm_k_dim（反向权重梯度），支持 Atlas A2/A3，自带实测性能数据。算子粒度对应关系总览：
gmm_addgrouped_matmul_addgmm_k_dimgrouped_matmul（groupType=2，K 轴分组）gmm_local_expgrouped_matmul（groupType=0，M 轴分组）gmm_local_exp_with_zero二、按算子粒度的差异性对比
1. gmm_add ↔ 社区
grouped_matmul_add两者均实现"分组矩阵乘 + 原地累加"：
C[g] += X_g @ W_g。MT 版专为 MoE 反向权重梯度设计（X_g^T转置求导），社区版为通用累加。gmm_addgrouped_matmul_addgroup_listint32，直接给每段 K 行数groupListINT64，cumsum 累积和c3D[G, M, N]，每专家独立梯度槽，in-placeyRef单 2D 张量原地累加trans_a（反向X_g^T）transposeX/transposeWeightaicore_num可指定核数2. gmm_k_dim ↔ 社区
grouped_matmul（K 轴分组）两者均支持"按 K 维分组、各组独立输出"：
Y[g] = X_g @ W_g（MT 支持trans_a走反向X_g^T @ W_g）。gmm_k_dimgrouped_matmul[G, M, N]3. gmm_local_exp ↔ 社区
grouped_matmul（M 轴分组）均按 M 维（token 轴）分组做矩阵乘
Y[cuM:cuM+M_g] = X[cuM:cuM+M_g] @ W[g]^T。核心差异在"本地专家子集"语义：gmm_local_expgrouped_matmulexp_start_idx/exp_end_idx只计算本地专家子集（张量并行下每卡仅算本卡专家）[LocalExpNum, N, K]（trans_b=True）+is_b_nz（NZ 支持）type_promotion（输出提 fp32）4. gmm_local_exp_with_zero ↔ 社区无对应
社区 8 个算子中没有"先清零输出再分组矩阵乘"的变体。MT 版通过 AIV 向量核先将整个输出
Y清零、再由 AIC 矩阵核执行计算，避免未激活专家位置残留脏数据，省去额外alloc + zero的内存与 Kernel 开销。这是社区能力空白，也是本次开源在 MoE 前向场景下的独有补充。三、自研算子的价值总结
1. 直击 MoE 训练链路，补齐社区推理/量化导向的空白
社区 gmm 算子族以推理 + 量化为重心；MT 的 4 变体完整覆盖 MoE 训练两端（前向
local_exp系列 + 反向gmm_add/gmm_k_dim），且均有实测性能数据。2. 本地专家子集计算能力（社区缺失）
exp_start_idx/exp_end_idx支持张量并行下每卡只计算本卡专家对应 token 段，减少冗余计算与通信——社区grouped_matmul无此语义。3. 输出清零变体（社区无对应）
gmm_local_exp_with_zero以 AIV 预清零 + AIC 计算的双核流水保证数值正确性，是社区能力空白，可直接作为社区新特性回灌。4. 反向梯度专用 API（单接口切换）
torch_npu.gmm一个接口通过c是否传入自动切到gmm_add（原地累加梯度）或gmm_k_dim（独立输出），配合trans_a原生支持X_g^T反向转置、aicore_num核数控制，开发者零额外配置。5. 轻量精简、易移植
无量化/激活融合复杂度，每变体仅几个输入/属性，代码量小、与社区同构（均基于 Ascend C Matmul），可作为社区
grouped_matmul各 groupType 场景的参考实现低成本回灌。6. 性能数据透明
README 内置 Atlas A2 实测 MFU：前向最高 79.83%（local_exp）、反向 11%~44%，社区 README 均未提供性能数据。
四、结论
社区 gmm 族中，
gmm_add与gmm_k_dim分别有grouped_matmul_add与grouped_matmul可对应（社区为通用/量化超集，MT 为训练专用精简实现），gmm_local_exp有部分对应（缺本地专家子集），gmm_local_exp_with_zero无对应。本次开源以一套双平台、带实测性能的 MoE 训练专用 gmm 变体，补齐社区在反向梯度、本地专家子集、输出清零三方面的能力空白，可低成本回灌增强社区官方算子。第二部分：LightningIndexer 族
一、背景
场景：LightningIndexer（LKI）是稀疏注意力（MoE、NSA 等模型）的路由算子，通过计算 Query 与 Key 的注意力得分，经逐头权重加权归约后做 TopK 排序，输出每个 Query token 最重要的 K 个 Key block 索引及分值，供后续 Sparse FlashAttention 使用。其反向/损失算子
LightningIndexerGradKLLoss融合 KL 散度损失与对query_index / key_index / weight的梯度计算。算子粒度对应关系总览：
lightning_indexerlightning_indexer（V1，同源）lightning_indexerlightning_indexer_v2lightning_indexer_grad_kl_losssparse_lightning_indexer_grad_kl_loss（同源同名）lightning_indexer_grad_kl_lossdense_lightning_indexer_grad_kl_loss/dense_lightning_indexer_kl_loss_grad二、按算子粒度的差异性对比
1. lightning_indexer ↔ 社区
lightning_indexer（V1，同源）两者功能高度一致：
Indices = Top-k{ [1]₁ₓg @ [(W @ [1]₁ₓSₖ) ⊙ ReLU(Q_index @ K_indexᵀ)] }，输入query/key/weights/actual_seq_lengths/block_table，属性layout_query/layout_key/sparse_count/sparse_mode/pre_tokens/next_tokens/return_values。核心差异在 MT 的稀疏路由增强属性：lightning_indexerlightning_indexerblock_len（KV block 粒度 1/2/4/8/16）、init_num（sink token 固定选取）、local_num（对角线局部窗口）cur_seq_lengths_*int64actual_seq_lengths_*int322. lightning_indexer ↔ 社区
lightning_indexer_v2（独立实现）社区 V2 为另一形态的独立实现：输入
q/k/w/cu_seqlens_q/k/seqused_q/k/cmp_residual_k/block_table/output_idx_offset/metadata，属性topk/max_seqlen_q/layout_q/layout_k/mask_mode/cmp_ratio/return_value，引入cmp_ratio压缩采样与 metadata 机制。MT 版与其非同源、功能侧重不同：MT 以 block 粒度聚合 + sink/局部窗口的路由增强为核心，V2 以压缩比/元数据驱动为特色，两者互补。3. lightning_indexer_grad_kl_loss ↔ 社区
sparse_lightning_indexer_grad_kl_loss（同源）两者功能一致：给定 query/key（重建 FlashAttention softmax 分布 P）、query_index/key_index（低秩压缩）、weight、sparse_indices、softmax_max/sum，融合计算 KL 散度损失与
dQ_index / dK_index / dW梯度。MT 的 def/kernel 文件命名与社区完全同名，属同源实现。差异：lightning_indexer_grad_kl_losssparse_lightning_indexer_grad_kl_losspre_tokens/next_tokens/block_size4. lightning_indexer_grad_kl_loss ↔ 社区 dense 变体（部分对应）
社区另有
dense_lightning_indexer_grad_kl_loss/dense_lightning_indexer_kl_loss_grad：稠密场景下 query/key/query_index/key_index 无需稀疏化处理，是不同于 MT 稀疏版（sparse_indices 驱动）的另一条实现路线，场景互补而非竞争。三、自研算子的价值总结
1. 社区前向 V1 缺失的稀疏路由增强（MT 独有）
block_len（block 粒度聚合）、init_num（sink token 固定保留）、local_num（对角线局部窗口）是稀疏注意力路由的关键能力，社区lightning_indexer（V1）与lightning_indexer_v2均未覆盖。本次开源可直接作为社区前向算子的增强特性回灌，补齐 V1 → 增强版的能力缺口。2. 反向损失算子与社区同源，可双向验证
MT 的
lightning_indexer_grad_kl_loss与社区sparse_lightning_indexer_grad_kl_loss同源同名，开源后可为社区提供原始参考实现与对照基准，便于社区追溯算法来源、回灌pre_tokens/next_tokens/block_size等 MT 侧属性语义。3. 成对交付、训练链路完整
前向（输出
sparse_indices/sparse_values，供 SFA 使用）+ 反向损失（KL + 梯度 + loss）成对提供，天然衔接 Sparse FlashAttention 前反向，构成 NSA/MoE 稀疏注意力训练闭环。4. int64 序列长度
cur_seq_lengths_*/actual_seq_lengths_*采用 int64，支持更长序列与更大 Batch，优于社区前向 V1 的 int32。