状态（Status）： Draft
作者（Authors）： @husichao
创建日期（Created）： 2026-09-04
更新日期（Updated）： 2026-09-04
相关 Issue/PR： 待关联

1. 概述

1.1 简介与目标

继续使用现有 dsa，在 token index key 之后增加可选 KPool 分支：

• 4 个 token index key 压缩成 1 个 pool key；
• Lightning Indexer 从 pool key 中选择 512 个 pool；
• 将 pool 展开为 2048 个 token，并追加当前 query 所在未满分组中的 0～3 个可见 token；
• 将有效 token id 整理到行前部，再用 -1 补到融合算子支持的宽度；
• 将补齐后的索引一次传给 Sparse Flash Attention（SFA）；
• Indexer 仍按 DSA 两阶段方案训练，KPool 新增参数由 pool 级 Indexer Loss 更新。

对 GLM-5.3-Flash，原始索引宽度最大为 2048 + 3 = 2051。Ascend 训练路径补齐到
CANN 前向和反向均支持的 K=3072：

512 pool × 4 token + 最多 3 个尾部 token = 最多 2051 个有效槽位
2051 个槽位 + 1021 个 -1 padding = SFA K=3072

本文将 query 所在未满 4-token 分组中的已可见 token 称为“尾部 token（tail）”。Padding 只扩展
sparse_indices 的最后一维，原始序列及 K/V Tensor 保持不变。

2. 用例分析

本方案覆盖以下生产场景：

• GLM-5.3-Flash 的 4-token KPool 索引和 512-pool TopK；
• BSND 与 TND（packed sequence）布局，pool 不跨样本或 document；
• Dense warm-up 与 Sparse training 两阶段 Indexer Loss；
• TP、CP、FSDP、recompute 和 checkpoint；
• 长序列下的算法正确性、精度、吞吐、峰值显存和断点续训验收。

约束是保持现有 token DSA 行为兼容，并在 Ascend 训练路径使用 CANN 前向、反向均支持的 SFA 规格。

3. 方案设计

3.1 总体流程

名称	职责
DSAIndexer	生成 index query、token index key 和 weight
KPool	将每 4 个 token index key 压成一个 pool key
Lightning Indexer	对 pool 打分并选择 512 个完整 pool
尾部 token	当前未满 4-token 分组中，对该 query 已因果可见的 0～3 个 token
索引构造	展开 pool、追加尾部 token、整理有效索引并补齐到 3072
SFA	根据最终 token id 一次完成稀疏 Attention
Indexer Loss	两阶段训练 Indexer；TopK 本身不参与反向

灰色是原 DSA，绿色是 KPool 新增流程，橙色是需要适配的现有模块。原 token DSA 沿用
2048 个 token id 的现有路径。

3.2 配置

增加最小配置：

dsa_indexer_candidate_mode: str = "token"   # token | kpool
dsa_indexer_pool_size: int = 1               # GLM-5.3-Flash 为 4
dsa_indexer_pool_always_select_tail: bool = True

校验规则：

• candidate_mode=kpool 时，第一版只接受 pool_size=4；
• index_topk % pool_size == 0，GLM-5.3-Flash 为 2048 / 4 = 512 个 pool；
• KPool 使用 pool 级 Indexer Loss；
• token 模式保持原配置、原算子和原 checkpoint 行为。

SFA 的 K=3072 是由 index_topk 和 pool_size 推导出的算子规格，不对用户暴露为算法配置。

3.3 KPool 与索引构造

DSAIndexer 原有 query、key、weight 投影和打分公式不变。KPool 只压缩 token_k，新增参数
与 GLM-5.3-Flash checkpoint 对齐：

index_kpool_compress_gate: [index_dim, hidden_size]
index_kpool_compress_ape:  [pool_size, index_dim]

BSND 直接按 batch 中每条样本的有效长度分 pool；TND 先按 actual_seq_klen 的累计边界
拆出 document，再使用 document-local 位置分 pool。两种布局仅在序列边界和 token id 映射上
不同，后续的 pool 压缩、TopK、展开、tail 追加和补齐流程相同。

简化伪代码：

def kpool_indices(q_index, token_k, weight, hidden_states, seq_meta):
    POOL_SIZE = 4
    POOL_TOPK = 2048 // POOL_SIZE       # 512
    RAW_WIDTH = 2048 + POOL_SIZE - 1    # 2051
    SFA_WIDTH = 3072                    # CANN 训练支持的下一档 K

    pool_keys = []
    pool_token_ids = []

    # BSND 按 batch 划定边界；TND 按 actual_seq_klen 的累计边界划定 document。
    # pool 不能跨样本或 packed document。
    for doc in split_by_sequence(hidden_states, token_k, seq_meta):
        complete_pools = doc.length // POOL_SIZE
        width = complete_pools * POOL_SIZE

        keys4 = doc.token_k[:width].reshape(
            complete_pools, POOL_SIZE, index_dim
        )
        gates4 = linear(
            doc.hidden[:width], index_kpool_compress_gate
        ).reshape(complete_pools, POOL_SIZE, index_dim)

        # pool 内位置维做 FP32 softmax。
        alpha = softmax(
            gates4.float() + index_kpool_compress_ape[None].float(),
            dim=1,
        )
        pool_keys.append(sum(alpha.to(keys4.dtype) * keys4, dim=1))
        pool_token_ids.append(
            arange(width).reshape(complete_pools, POOL_SIZE)
        )

    pool_k, pool_meta = pack_pool_keys(pool_keys, seq_meta)

    # cmp_ratio/cmp_residual_k 使算子按每个 query 的局部位置计算可见完整 pool。
    # 可见 pool 不足 512 时，未使用槽位返回 -1。
    selected_pool_ids, selected_pool_scores = lightning_indexer_v2(
        q=q_index,
        k=pool_k,
        w=weight,
        topk=POOL_TOPK,
        q_seq_meta=seq_meta,
        k_seq_meta=pool_meta,
        mask_mode=CAUSAL,
        cmp_ratio=POOL_SIZE,
        cmp_residual_k=sequence_lengths(seq_meta) % POOL_SIZE,
    )

    # 每个有效 pool id 展开成 4 个原始 token id。
    selected_token_ids = expand_pool_ids(
        selected_pool_ids, pool_token_ids, invalid_value=-1
    )                                           # 固定 2048 槽位

    # 每个 query 单独追加当前未满分组中的可见 token，最多 3 个。
    q_pos = document_local_query_positions(seq_meta)
    tail_count = (q_pos + 1) % POOL_SIZE
    tail_start = q_pos + 1 - tail_count
    offsets = arange(POOL_SIZE - 1)
    tail_ids = where(
        offsets < tail_count[..., None],
        tail_start[..., None] + offsets,
        -1,
    )                                           # 固定 3 槽位

    raw_indices = concat([selected_token_ids, tail_ids], dim=-1)
    assert raw_indices.shape[-1] == RAW_WIDTH

    # CANN 要求每行有效值在前、无效值在后。
    # 排序不是 Attention 语义要求，但便于确定性和算子校验。
    compact_indices = compact_valid_prefix(raw_indices, invalid_value=-1)
    compact_indices = sort_valid_prefix(compact_indices, invalid_value=-1)
    sfa_indices = pad(compact_indices, width=SFA_WIDTH, value=-1)

    # pool id/score 留给 Indexer Loss；3072 token id 只交给 SFA。
    return sfa_indices, selected_pool_ids, selected_pool_scores, pool_meta

以第 2050 个 token（零基位置 2049）为例：该 query 可见 512 个完整 pool，当前未满分组包含
2 个尾部 token。该行最多包含 2048 个完整-pool token、2 个尾部 token 和 1022 个 -1，
最终宽度为 3072。

3.4 一次调用融合 SFA

NVIDIA NeMo AutoModel 的 KPool 路径先产生 index_topk + index_kpool - 1 个索引槽位，
再将有效索引整理为连续前缀，并补齐到 FlashMLA 要求的 512 对齐宽度。本方案采用相同的
索引组织方式，并按 CANN 反向支持的 K 规格将 2051 补齐到 3072。

attention_out, softmax_max, softmax_sum = sparse_flash_attention(
    q, k, v,
    sparse_indices=sfa_indices,     # [..., 3072]，有效前缀 + -1 后缀
    sparse_block_size=1,
    actual_seq_lengths_query=actual_seq_qlen,
    actual_seq_lengths_kv=actual_seq_kvlen,
)

融合 SFA 的输入约定为：

• 完整 pool token 和尾部 token 位于同一条索引序列；
• 每行有效索引形成连续前缀，无效槽位形成连续的 -1 后缀；
• SFA 在全部有效 token 上执行一次 softmax，并使用其原生前向和反向；
• SFA 输入宽度固定为 3072，其中至少 1021 个槽位为 -1 padding。

CANN SFA 前向允许 sparse_size > 0，并要求有效索引在前、无效索引在后；当前反向明确支持
K=1024、2048、3072...。当前 MindSpore 公共 ops.sparse_flash_attention 仍把 K 固定为
2048，因此 P2 需要将该接口扩展到 CANN 已支持的 K=3072，或接入能透传 3072 的底层封装。
这属于 MindSpore 接口接入，不要求新增 CANN SFA 算法。

下表说明选择 K=3072 的依据：

方案	前向	反向	结论
2051 → 3072，block=1	CANN 支持	CANN 支持 K=3072	当前方案
直接K=2051，block=1	CANN 支持	CANN 不支持 K=2051	仅作前向对照
K=513，block=4	A2/A3 前向支持	K=513、block=4 均不支持反向	不用于训练

3.5 Indexer Loss

KPool 沿用普通 DSA 的 KL 监督和两阶段训练，变化是把监督和打分的候选粒度从 token 提升为
完整 pool。KPool Indexer 及其新增的 compress gate、APE 参数加入优化器。

索引展开、Tail 追加和补齐只服务 SFA，不进入 Indexer Loss。两阶段的 Attention 和参数更新为：

阶段	Attention	参数更新
Dense warm-up	Dense Attention	冻结主模型；pool 级 dense Indexer Loss 更新 Indexer
Sparse training	K=3072 融合 SFA	主模型由 LM Loss 更新；Indexer 由 pool 级 sparse Indexer Loss 更新

主 Attention 产生的注意力概率作为停止梯度的监督目标。每 4 个 token 的目标概率相加得到
pool 目标概率，再与 Indexer 产生的 pool score 分布计算 KL Loss：

target_pool_prob[j] = sum(target_token_prob[4*j : 4*j+4])

TopK 是离散的，两阶段都由 Indexer Loss 更新 Indexer 参数；LM Loss 不经过 TopK 回传。

规则：

• 只监督完整且因果可见的 pool；尾部 token 尚未形成完整 pool，因此没有对应的 pool score；
• pool target 重新归一化；没有完整可见 pool 的 query 跳过 loss；
• 梯度从 pool score 回到 index query、pool key 和 weight；
• pool key 的梯度继续回到 token index key、compress gate 和 APE；
• SFA 补出的 -1 槽位只服务 Attention 算子对齐，不进入 Indexer Loss。

因此阶段二虽然不直接对 TopK id 求导，但 Indexer 分数更新后，后续 step 的 pool 排序和选择会变化。

3.6 并行与 checkpoint

• TP：沿用现有 DSA 的 index head 布局；KPool gate/APE 使用对应布局。
• CP：沿用现有 DSA 的 key 侧收集；在得到 document 全局连续的 token key/gate 后再分 pool，
  避免 pool 跨 CP 边界断开。新增的是 pool 边界处理，不引入新的并行语义。
• FSDP：新增 gate/APE 参与分片、保存和恢复。
• Recompute：索引构造必须确定，重计算得到同一组 pool/尾部 token id。
• Checkpoint：KPool 模式必须加载或初始化 gate/APE；GLM-5.3-Flash 权重名需建立明确映射。

BSND 按样本、TND 按 cu_seqlens 中的 document 独立分 pool；pool 和尾部 token 都以
document 为边界生成。

3.7 代码改动

文件	改动
mindformers/parallel_core/transformer_config.py	增加 KPool 配置和校验
mindformers/parallel_core/transformer_config_utils.py	配置字段透传
.../experimental_attention_variant/dsa_kpool.py	新增 pool 压缩、展开、尾部索引、有效值整理和补齐
.../experimental_attention_variant/dsa_indexer.py	KPool 分支接入 Lightning Indexer V2；原打分逻辑不改
.../experimental_attention_variant/dsa.py	SFA 接收 K=3072 padded indices
.../experimental_attention_variant/dsa_attention.py	装配 token/KPool 两条路径和两阶段 loss
.../experimental_attention_variant/dsa_indexer_loss.py	增加 dense/sparse pool 级 loss 和梯度
mindformers/pynative/base_models/gpt/parallelize.py	新参数的 TP/CP/FSDP/recompute 策略
tests/st/test_ut/test_pynative/test_transformers/test_dsa/	数学、前反向、边界、并行和回归测试

不新增顶层 dsa_kpool 类型；experimental_attention_variant="dsa" 保持不变。

3.8 分阶段交付

P1：单卡算法正确性

• 实现 4-token pool、512-pool TopK、尾部 token 追加和 2051→3072 索引补齐；
• 实现 BSND/TND document 边界；
• 实现 pool 级 dense/sparse Indexer Loss 的参考版本；
• 与 dense masked Attention、未 padding 的 2051-token 实现对齐；
• 回归原 token DSA。

P2：融合算子与训练闭环

• 接入 Lightning Indexer V2 的 cmp_ratio=4 和 cmp_residual_k；
• 让 MindSpore SFA 前向、反向透传 K=3072；
• 接入 pool 级 Indexer Loss 的训练算子；
• 完成 dense warm-up 与 sparse training 单卡闭环；
• 验证 BF16/FP16、长序列和异常输入。

P3：生产并行与工程化

• 适配 TP、CP、FSDP、recompute；
• 完成 GLM-5.3-Flash checkpoint 映射、保存和恢复；
• 验证 packed sequence、多卡确定性和断点续训；
• 完成长稳、吞吐和显存验收。

4. 测试设计

类别	必测内容
Pool	gate softmax 轴、FP32 计算、APE、pool key 梯度
边界	S<4、S%4=1/2/3、可见 pool 少于 512、空 pool
索引	pool 展开、逐 query 尾部 token、有效前缀、-1 padding、K=3072
Attention	3072 padded SFA 与 2051/dense 对照的前向和反向
Loss	两阶段 pool target、无完整 pool、无候选跳过、gate/APE 梯度
并行	单卡对 TP/CP/FSDP，TND document 与 CP 边界一致
回归	token DSA 的输出、loss 和 checkpoint 行为不变

关键验收关系：

K=3072 padded SFA ≈ K=2051 对照 ≈ dense masked Attention

5. 缺点与风险

生产前需要重点收敛以下风险：

• 当前缺少带反向的 KPool Compress 和索引构造训练融合算子，小算子组合可能引起断流或 Host 下发瓶颈；
• 2051 → 3072 至少引入 1021 个 -1 槽位，索引宽度相对 K=2048 增加 50%，需实测 SFA 耗时和峰值显存；
• 当前未找到 GLM-5.3 KPool Indexer Loss 的开源训练实现，只能验证本方案参考公式与融合实现一致，不能证明复现官方训练行为。

6. 现有技术

本方案参考 NVIDIA NeMo AutoModel 的 KPool 索引组织与 padding 路径，并结合 MindSpore/CANN SFA 的前向、反向约束适配 Ascend 训练。

参考资料

• GLM-5.3-Flash config：index_kpool=4、index_topk=2048
• NVIDIA NeMo AutoModel：KPool 索引构造
• NVIDIA NeMo AutoModel：稀疏索引 compact/sort/padding
• MindSpore Sparse Flash Attention
• CANN SFA 前向
• CANN SFA 反向

7. 未解决问题

• Pool 级 Indexer Loss 是否与 GLM-5.3-Flash 官方训练定义一致；
• 两个训练融合算子的交付形态、反向支持和性能收益；
• K=3072 padding 在目标长序列、并行配置上的吞吐与峰值显存是否满足生产门槛。