执行摘要

• 一句话：reland QSA 压缩 gather 越界修复，extend 行号夹紧
• 推荐动作：值得精读但不必逐行推敲——它只有一行代码，重点是理解取舍而非实现。建议关注三点：第一，阅读 update_key_state_and_compress 的 extend 与 paged eager 两个分支，确认 clamp 只作用于“有界性缺失”的那一条路径；第二，向作者确认未填满 group 的池化结果是否会被后续完整 group 覆盖，若不会，则应改为 mask/重算而非重复最后一行；第三，建议补一个构造“chunk 行数非 compress_ratio 整数倍”的单测，并说明 Extra、AMD ROCm 10 两条 CI 的失败是否为已知问题。若团队正在安排 #33426、#39061、#39219 的 reland，本 PR 可作为“回退后如何最小化恢复”的参考样例。

功能与动机

PR 标题直接点明动机：Reland fix(qsa): clamp the compress gather to the rows (#38346)，即重新落地 #38346 的 QSA 压缩 gather 越界修复；PR body 只保留了模板，未补充额外说明。结合 #39405（Revert #38346, #33426, #39061 和 #39219）的历史上下文可知，原修复在一次“批量回退误合并 PR”的操作中被连带撤销，本 PR 是选择性恢复。代码侧动机来自 update_key_state_and_compress 的 extend 分支：成员行号由 compress_member_rows[:, None] + arange(compress_ratio) 构造，当本轮 forward 的 token_k 行数不是 compress_ratio 的整数倍时，最后一个 group 的尾部索引会超出 source_keys.shape[0]，直接对 source_keys 做 gather 会触发越界读（非法访存或崩溃）。文件中已有的 RuntimeError（“QSA metadata is missing the precomputed write plan”）表明该路径此前缺乏防御，本次 clamp 是补上同一类边界保护。

实现拆解

1. 变更入口：唯一改动位于 python/sglang/srt/layers/attention/qsa/qsa_indexer.py 的 QsaIndexer.update_key_state_and_compress。该方法在每轮 forward 后把已写入的 key 按 compress_ratio 分组、池化并落到压缩槽位 metadata.write_locs，是 QSA 索引器从「原始 key 状态」到「压缩 key」的关键一跳。

2. 问题定位在 extend 分支：方法内先做 metadata.is_cuda_graph 早退，再校验 write_locs 是否为空。真正构造行号的只有 extend 分支：member_rows = metadata.compress_member_rows.long()，随后 group_locs = member_rows[:, None] + torch.arange(self.compress_ratio, ...)，source_keys = token_k。这里的隐含假设是“本轮 chunk 内每个 group 都完整”，而 chunked prefill 的尾部 chunk 往往只承载部分 group 成员，于是 group_locs 的最大值可达 rows + compress_ratio - 1，越界。

3. 修复动作：在 source_keys = token_k 之后立即插入 group_locs = group_locs.clamp_max(source_keys.shape[0] - 1)。选择 clamp 而不是重写索引或加 mask，是因为 clamp 让越界成员退化为指向最后一行：索引合法、显存访问合法，配合后续 average_pool_qsa_keys 的池化不会引入非法值；代价是未填满的 group 会重复统计最后一行。另一分支（paged eager / speculative fallback）通过 metadata.compress_group_ring_locs 或 _group_ring_slots 取 pending ring 窗口，窗口天然有界，故未做同样处理。

4. 下游影响：group_locs 在 self._use_fused_compress(pool) 为真时传给 self._fused_compress_store(pool, group_locs, compressed_locs, source_keys=..., source_rope=...)，为假时用于 key_groups = source_keys[group_locs] 与 source_rope[group_locs[:, 0]]。由于 clamp 被放在 _use_fused_compress 判断之前，融合内核路径与 eager pool 路径同时受保护，避免了“只修一条路径”的常见漏改。

5. 配套情况：本次 changed_files_count = 1，+1/-0，没有新增或修改任何测试文件、配置项或文档；唯一的非代码动作是作者 Qiaolin-Yu 在评论区执行 /tag-and-rerun-ci。CI 状态显示 Latest PR Test (Base) 仍在运行，Extra 与 AMD ROCm 10 两个流水线为失败，PR 在只有一个 commit（作者 ehuaa）的情况下被合并。

关键文件：

• python/sglang/srt/layers/attention/qsa/qsa_indexer.py（模块 注意力索引器；类别 source；类型 core-logic；符号 update_key_state_and_compress, average_pool_qsa_keys, _fused_compress_store）: 唯一变更文件，也是 QSA 索引器压缩写入路径的唯一落点；+1 行 clamp 同时保护融合内核与 eager 两条下游路径。

关键符号：update_key_state_and_compress, average_pool_qsa_keys

关键源码片段

python/sglang/srt/layers/attention/qsa/qsa_indexer.py

唯一变更文件，也是 QSA 索引器压缩写入路径的唯一落点；+1 行 clamp 同时保护融合内核与 eager 两条下游路径。

# python/sglang/srt/layers/attention/qsa/qsa_indexer.py
# QsaIndexer.update_key_state_and_compress：把本轮 forward 已写入的 key 按 compress_ratio 分组池化，
# 再写入压缩槽位（metadata.write_locs）。成员行号有两个来源：extend chunk 的 packed token_k，
# 或 paged eager（speculative fallback）的 pending ring 窗口。
​
if metadata.is_cuda_graph:
    self._compress_decode_cuda_graph(metadata)
    return
​
if metadata.write_locs is None:
    raise RuntimeError(
        "QSA metadata is missing the precomputed write plan; the "
        "sparse-attention backend derives it from the batch lengths"
    )
if metadata.write_locs.numel() == 0:
    return
​
group_end_positions = metadata.compress_group_positions.long()
compressed_locs = metadata.write_locs
​
if is_extend:
    # Extend chunk 按 group 对齐：每个 group 的成员行 = 组首行 + [0, compress_ratio)。
    member_rows = metadata.compress_member_rows.long()
    group_locs = member_rows[:, None] + torch.arange(
        self.compress_ratio, device=member_rows.device, dtype=torch.long
    )
    source_keys = token_k
    # 本次 reland 的关键一行：本轮 chunk 未必刚好装下 compress_ratio 的整数倍个成员，
    # 最后一个 group 的尾部行号会超过 token_k 的行数，直接 gather 即越界访存；
    # 夹到 [0, rows - 1] 后索引恒合法，未填满的成员退化为重复取最后一行参与池化。
    group_locs = group_locs.clamp_max(source_keys.shape[0] - 1)
    source_rope = metadata.extend_rope_matrix
    if source_rope is None:
        source_rope = build_rope_position_matrix(rope_positions, token_k.shape[0])
else:
    # Paged eager 行（speculative fallback）每次最多完成一个 group，
    # 成员正好是 pending ring 窗口，窗口天然有界，因此不需要 clamp。
    group_locs = metadata.compress_group_ring_locs
    if group_locs is None:
        group_locs = self._group_ring_slots(
            metadata,
            group_end_positions,
            metadata.compress_sequence_ids.long(),
        )
    source_keys = pool.get_qsa_key_state_buffer(self.layer_id)
    source_rope = pool.qsa_rope_position_buffer
​
# clamp 放在 fused 判断之前，因此 _fused_compress_store 与 eager 取数路径同时受保护。
if self._use_fused_compress(pool):
    self._fused_compress_store(
        pool,
        group_locs,
        compressed_locs,
        source_keys=source_keys,
        source_rope=source_rope,
    )
    return
​
key_groups = source_keys[group_locs]
pooled = average_pool_qsa_keys(key_groups)
compressed_rope_positions = self._rope_from_matrix(source_rope[group_locs[:, 0]])

评论区精华

本 PR 的 review 材料基本为空：review_comments_count = 0、Review 审核 = []，唯一的评论是作者自己的 /tag-and-rerun-ci CI 触发指令。因此不存在就 clamp 语义、边界 group 正确性展开的交锋记录，也没有 reviewer 对“reland 是否需要重新走审批”提出公开质疑（PR body 仍保留 checklist 未勾选状态）。需要读者自行判断的争议点集中在两处：其一，clamp 让未填满 group 的最后一行被重复参与平均池化，属于“以数值偏差换稳定性”的取舍，材料中没有验证其偏差是否被后续完整 group 覆盖写回；其二，原 PR 是在批量回退中被撤销的，本次单独 reland 未附带回退原因分析或新增回归测试，讨论链条被动依赖 #38346 / #39405 的历史上下文。

• reland 是否需要重新审查与测试补齐 (testing): 改动最终合入 main；但材料中没有解释两条失败流水线是否为已知失败，也未补充针对越界场景的回归测试。
• clamp 与掩码重算的语义取舍 (correctness): 作者选择 clamp 作为最小改动方案，隐含前提是该边界 group 的写入值可被后续完整 group 覆盖或可容忍近似；该前提未在 PR 中被显式验证。

风险与影响

• 风险：
  1. 数值正确性风险（主要）：clamp_max 会把越界成员指向 source_keys.shape[0] - 1，即同一行被重复计入 average_pool_qsa_keys。若某个 group 的池化结果会被当作最终压缩 KV 使用，那么该 group 的键值不再等于其真实成员的均值；只有在“该 group 后续会被完整覆盖写回”这一前提下才无害，而材料中没有任何测试或注释证明这一点。

1. 缺少测试覆盖：改动涉及 QSA 压缩索引主路径，但 PR 未新增/修改测试文件（对照同类边界修复 PR#38941 会附带 test/registered/unit/mem_cache/test_free_kv_row_coalesce.py 这类单测）。这意味着越界场景既无法在 CI 中复现，也无法防止将来重构时回归。

2. reland 过程风险：#38346 曾因“误合并”被 #39405 回退，本次 reland 依旧没有 review 记录、body checklist 全部未勾选，且 Extra 与 AMD ROCm 10 两条 CI 流水线显示失败后即合入。如果失败来自与本改动无关的环境问题，风险可控；但没有任何材料证明这一结论。

3. 兼容性/性能：clamp_max 是逐元素算子，张量规模为 [group 数, compress_ratio]，开销可忽略；eager 路径不引入额外 kernel。兼容性上唯一的行为差异是：以前越界即崩（或触发 CUDA 断言），现在静默产出边界 group 的近似池化值，故障模式从“可见崩溃”变为“潜在静默偏差”，这对可观测性不利。
   - 影响：影响范围限定在启用 QSA 索引器压缩路径的模型与运行模式：只有 update_key_state_and_compress 的 extend 分支（chunked prefill）在“尾部 group 未填满”时才会命中该边界。用户侧收益是消除一次潜在的越界访存/崩溃，特别是在长序列 chunked prefill 与 CP（context parallel）组合下更易触发；代价是边界 group 的压缩 key 由“重复最后一行”的池化结果构成。系统层面不改变接口、元数据结构或配置项，对调度、内存池、HiCache、PD 分离等其他模块零影响。团队层面，本 PR 是 #39405 批量回退后的第一笔“选择性恢复”，会影响后续对 #33426、#39061、#39219 的 reland 节奏与审查标准。
   - 风险标记：核心注意力路径, 缺少测试覆盖, reland 缺审查记录, 边界 group 池化语义变化, CI 未全绿即合入, 故障从崩溃转为静默偏差

关联脉络

• PR #38346 fix(qsa): clamp the compress gather to the rows this forward has: 本 PR 的原始修复，被 #39405 批量回退后由本次 reland 原样恢复，是理解改动意图的直接来源。
• PR #39405 [misc] Revert #38346, #33426, #39061 and #39219: 批量回退四个误合并 PR 的操作，其 key_files 中包含 qsa_indexer.py；本 PR 是这批回退后的选择性恢复。
• PR #39061 Fix MUSA detection in compiled prefill path: 与 #38346 同批被回退的兄弟 PR，同样落在 prefill/注意力路径，可作为判断后续 reland 节奏与审查标准的参照。