执行摘要

• 一句话：kpool 索引器迁入 glm5next 并拆分后端
• 推荐动作：值得精读，但重点不在 729 行的搬运量，而在三处设计决策：一是 vllm/models/glm5next/sparse_indexer.py 这个 22 行 dispatcher——它决定了所有平台行为，CodeRabbit 指出的「非 ROCm 即 NVIDIA」回落值得你判断是否要改成显式平台判定；二是 CustomOp 的 forward_hip / forward_native dispatch 契约，任何把实现按硬件拆文件的改动都要先确认被覆写的方法名与 dispatch_forward() 的绑定规则一致，否则就是 Codex 那条 P1；三是 helper 抽取边界——只抽纯索引/形状逻辑、把 kernel 调用留在后端文件，这个尺子本 PR 把握得比较准，可以照搬到其他模型。若你在 ROCm 上使用 custom_ops=["all"] 或 enforce-eager，建议先确认 forward_hip 是否已补上再升级；若你的代码直接 import 了 vllm.model_executor.layers.sparse_attn_indexer_kpool，需要同步迁移到 vllm.models.glm5next.sparse_indexer。

功能与动机

PR body 明确表达了三个诉求：一是「Move sparse_attn_indexer_kpool out of vllm/model_executor/layers/ and into the glm5next model folder」，即让只服务 GLM-5.3-Flash 的模型私有 kernel 回归 vllm/models/glm5next/，与已有的 glm5next/nvidia/ops/kpool_compress.py、glm5next/amd/ops/kpool_compress.py 目录约定对齐；二是「split by hardware backend」，原文件用 TYPE_CHECKING + elif current_platform.is_rocm() 选择 kpool_ops，函数体内还有 if current_platform.is_rocm(): ...triton... else: ops... 的运行时分支，AMD 与 NVIDIA 代码互相纠缠；三是把与硬件无关的 decode 布局 helper 抽到 common/ 供两端共用，避免双实现漂移。PR body 同时声明该变更由 AI 辅助（Kimi Code CLI）完成并先开 draft，便于人工逐行复核；验证依据是 gsm8k 0.9371 与 aime25 pass@4=0.9333。

实现拆解

1. 目录分层：一份实现拆成四份文件。vllm/model_executor/layers/sparse_attn_indexer_kpool.py 重命名为 vllm/models/glm5next/nvidia/sparse_indexer.py（+76/-391）承载 CUDA 路径；新增 vllm/models/glm5next/amd/sparse_indexer.py（+729）承载 ROCm 路径；新增 vllm/models/glm5next/common/sparse_indexer.py（+170）放平台无关逻辑；新增 vllm/models/glm5next/sparse_indexer.py（+22）作为 dispatcher，成为调用方唯一可见入口。

2. 公共 helper 抽取（本 PR 最实质的契约变化）。把 _build_decode_scatter_indices、_scatter_decode_tokens_by_request、_decode_topk_seq_lens、_fill_causal_indices、_fill_short_decode_causal_indices、_gather_workspace_shapes、kv_cache_as_quant_view 以及常量 RADIX_TOPK_WORKSPACE_SIZE、MXFP4_BLOCK_SIZE 移入 common/sparse_indexer.py。这些函数只做索引构造、padded [batch_size, next_n] 网格与扁平布局的换算、workspace 形状与 dtype 推导，不调用平台 kernel，因此 AMD/NVIDIA 可共用；_decode_topk_seq_lens 的 docstring 记录了非均匀 decode batch 下按扁平 positions[:n] 取值会导致行错位乃至越界 block-table 读的原因。

3. NVIDIA 侧去分支。nvidia/sparse_indexer.py 删除三类平台耦合：if TYPE_CHECKING / elif is_rocm() / else 的 kpool_ops 选择（改为直接 import nvidia 实现）、elif is_xpu(): from vllm._xpu_ops import xpu_ops 及 forward_native 的 XPU 路由（commit 1e2089b 理由：glm5next 在 package import 阶段就对 XPU 抛 NotImplementedError，分支不可达）、标准写入路径上的 if current_platform.is_rocm(): indexer_k_quant_and_cache_triton(...) else: ops.indexer_k_quant_and_cache(...)，只保留 _custom_ops 版本（含 cp_gather_indexer_k_quant_cache）。

4. AMD 侧独立成文件。amd/sparse_indexer.py 模块级 from vllm.models.glm5next.amd.ops import kpool_compress as kpool_ops 取代运行时分派，函数内按需 import vllm.v1.attention.ops.rocm_aiter_mla_sparse 的 indexer_k_quant_and_cache_triton 与 cp_gather_indexer_k_quant_cache_triton，并重新定义 _kpool_compress_insert、sparse_attn_indexer_kpool、SparseAttnIndexerKpool（__init__、forward_native）。

5. 调用方、测试与 CI 配套。vllm/models/glm5next/common/attention.py 的 import 由 vllm.model_executor.layers.sparse_attn_indexer_kpool 改为 vllm.models.glm5next.sparse_indexer，这是仓库内唯一正式调用点；tests/v1/attention/test_sparse_indexer_decode_seq_lens.py 改为从 glm5next.common.sparse_indexer 导入 helper，并按 current_platform.is_rocm() 选择 amd/nvidia 模块做 dispatch 断言，测试注释同步更新了 package bootstrap 顺序。没有新增测试用例、没有改 Buildkite 配置，验证依赖 4 次 /ci run 与一次 /amd-ci run nightly。

关键文件：

• vllm/models/glm5next/amd/sparse_indexer.py（模块 AMD 后端；类别 source；类型 data-contract；符号 _kpool_compress_insert, sparse_attn_indexer_kpool, SparseAttnIndexerKpool, init）: 新增 729 行 AMD/ROCm 实现，把原先内联在同一文件里的 ROCm 分支（aiter + Triton 的 indexer_k_quant_and_cache_triton、cp_gather_indexer_k_quant_cache_triton）独立出来；同时 Codex 的 P1 缺陷（在 ROCm 上只实现 forward_native 而未覆写 forward_hip）正落在这个文件上。
• vllm/models/glm5next/nvidia/sparse_indexer.py（模块 NVIDIA 后端；类别 source；类型 rename-or-move；符号 sparse_attn_indexer_kpool, SparseAttnIndexerKpool, _kpool_compress_insert, forward_native）: 由 vllm/model_executor/layers/sparse_attn_indexer_kpool.py 重命名而来，是 NVIDIA 路径的唯一实现；删除了 TYPE_CHECKING + is_rocm() 的 kpool_ops 选择、xpu_ops 分支与 ROCm 运行时分支，改用 common 的共享 helper。
• vllm/models/glm5next/common/sparse_indexer.py（模块 共享辅助；类别 source；类型 data-contract；符号 _build_decode_scatter_indices, _scatter_decode_tokens_by_request, _decode_topk_seq_lens, _fill_causal_indices）: 新增的共享 helper 层，承载 AMD/NVIDIA 共用的 decode 布局与 workspace 形状逻辑（含 _decode_topk_seq_lens 的 padded 网格换算）；这是本 PR 最实质的契约变化，也是两端一致性的单点依赖。
• vllm/models/glm5next/sparse_indexer.py（模块 平台分发；类别 source；类型 entrypoint；符号 SparseAttnIndexerKpool）: 新增 22 行平台 dispatcher，是调用方唯一可见的入口；CodeRabbit 在此指出「非 ROCm 即 NVIDIA」的静默回落问题（未采纳修改）。
• vllm/models/glm5next/common/attention.py（模块 注意力层；类别 source；类型 import-adjustment）: 仓库内唯一正式调用点，import 从 vllm.model_executor.layers.sparse_attn_indexer_kpool 改为 vllm.models.glm5next.sparse_indexer，是本次搬迁对调用链的全部影响面。
• tests/v1/attention/test_sparse_indexer_decode_seq_lens.py（模块 单元测试；类别 test；类型 test-coverage；符号 test_kpool_ops_dispatch_matches_platform, test_short_decode_fills_exact_causal_rows, test_short_decode_leaves_buffer_unchanged_for_sparse_context）: CPU 端 decode 布局测试同步迁移：helper 改从 glm5next.common.sparse_indexer 导入，并按平台选择 amd/nvidia 模块做 dispatch 断言；同时更新了 package bootstrap 顺序的注释。

关键符号：sparse_attn_indexer_kpool, _kpool_compress_insert, _decode_topk_seq_lens, _build_decode_scatter_indices, _scatter_decode_tokens_by_request, _fill_short_decode_causal_indices, _gather_workspace_shapes, kv_cache_as_quant_view, SparseAttnIndexerKpool.forward_native

关键源码片段

vllm/models/glm5next/amd/sparse_indexer.py

新增 729 行 AMD/ROCm 实现，把原先内联在同一文件里的 ROCm 分支（aiter + Triton 的 indexer_k_quant_and_cache_triton、cp_gather_indexer_k_quant_cache_triton）独立出来；同时 Codex 的 P1 缺陷（在 ROCm 上只实现 forward_native 而未覆写 forward_hip）正落在这个文件上。

# vllm/models/glm5next/amd/sparse_indexer.py（新增，+729）
# ROCm 实现独立成文件后，模块级 import 取代了原先的运行时平台分派：
# 搬迁前 kpool_ops 由 if TYPE_CHECKING / elif is_rocm() 在同一个文件里选择。
import torch
​
import vllm.envs as envs
from vllm import _custom_ops # noqa: F401 # 注册 torch.ops._C 中的 kernel
from vllm._aiter_ops import rocm_aiter_ops
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
from vllm.model_executor.custom_op import CustomOp
from vllm.models.glm5next.amd.ops import kpool_compress as kpool_ops
from vllm.models.glm5next.common.sparse_indexer import (
    RADIX_TOPK_WORKSPACE_SIZE, # 常量与平台无关 helper 都来自 common/
    _build_decode_scatter_indices,
    _decode_topk_seq_lens,
    _fill_causal_indices,
    _fill_short_decode_causal_indices,
    _gather_workspace_shapes,
    _scatter_decode_tokens_by_request,
    kv_cache_as_quant_view,
)
​
​
def _kpool_compress_insert(
    k: torch.Tensor,
    gate_score: torch.Tensor,
    ape: torch.Tensor,
    kv_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    kpool: int,
    head_dim: int,
    round_scale: bool,
) -> None:
    """把 ``kpool`` 个连续 token 池化成一个 fp8 的 K，写到 pool slot 上。
​
    ``slot_mapping`` 是 pool 粒度的（spec 里 compress_ratio == kpool）：
    只有每个完整 pool 的最后一个 token 携带有效（>= 0）slot，pool 内其余
    token 都是 -1。这里把每个位置都当作 pool 完成候选，未完成的在 kernel
    内部被 mask 掉；先压缩有效行反而要在 eager prefill 路径上多两次 device
    sync，数值上却没有任何收益。假定 chunk 起点与 pool 对齐。
    """
    n = slot_mapping.shape[0]
    # batch 小于一个 pool 时不可能有 pool 完成，同时也保证下面的
    # clamp 后 gather 下标不越界。
    if n < kpool:
        return
    pos = torch.arange(n, device=k.device)
    valid = slot_mapping >= 0
    # 丢掉起点落在 batch 之前的 pool（前导 padding），它们的 gate / k 数据
    # 本来就没有定义。
    write_mask = valid & (pos >= kpool - 1)
    offs = torch.arange(kpool, device=k.device)
    idx = (pos - (kpool - 1)).clamp_min(0)[:, None] + offs[None, :]
    kpool_ops.kpool_compress_and_write_cache(
        kv_cache,
        k[idx], # [n, kpool, head_dim]
        gate_score[idx],
        ape,
        slot_mapping.to(torch.int64),
        pool_size=kpool,
        head_dim=head_dim,
        write_mask=write_mask,
        round_scale=round_scale,
        write_cache=True,
        return_compressed=False,
    )
​
​
# 类实现（SparseAttnIndexerKpool.__init__ / forward_native）同样搬到这里。
# 注意：Codex 的 P1 指出 ROCm 上 CustomOp.dispatch_forward() 绑定的是
# forward_hip，若只覆写 forward_native，继承链会走到 forward_cuda 并抛
# NotImplementedError；启用 custom_ops=["all"] 或 enforce-eager 时会命中。

vllm/models/glm5next/nvidia/sparse_indexer.py

由 vllm/model_executor/layers/sparse_attn_indexer_kpool.py 重命名而来，是 NVIDIA 路径的唯一实现；删除了 TYPE_CHECKING + is_rocm() 的 kpool_ops 选择、xpu_ops 分支与 ROCm 运行时分支，改用 common 的共享 helper。

# vllm/models/glm5next/nvidia/sparse_indexer.py
# 由 vllm/model_executor/layers/sparse_attn_indexer_kpool.py 整体搬迁而来，
# 现在只保留 NVIDIA 路径：原先 TYPE_CHECKING + is_rocm() 的 kpool_ops 二选一、
# is_xpu() 的 xpu_ops 导入，以及标准写入路径上的 ROCm 运行时分支全部删除。
from vllm.models.glm5next.common.sparse_indexer import (
    RADIX_TOPK_WORKSPACE_SIZE, # 常量与布局 helper 统一来自 common/
    _build_decode_scatter_indices,
    _decode_topk_seq_lens,
    _fill_causal_indices,
    _fill_short_decode_causal_indices,
    _gather_workspace_shapes,
    _scatter_decode_tokens_by_request,
    kv_cache_as_quant_view,
)
from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops
​
if current_platform.is_cuda_alike():
    from vllm import _custom_ops as ops
​
​
# 下面是 sparse_attn_indexer_kpool 的核心写入分支（函数体内部片段）。
if not skip_k_cache_insert:
    assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet"
    if index_kpool > 1 and gate_score is not None and compress_ape is not None:
        # kpool prefill 写入：把连续 index_kpool 个 prefill token 按
        # softmax(gate + ape) 加权求和，再 Hadamard 变换量化成 fp8 写到 pool slot。
        # decode token 排在 batch 最前面，它们所属 pool 的早期 token 不在本 batch，
        # 因此这里只处理 prefill，其余交给 has_decode 分支的 tail kernel。
        # compress_ratio == index_kpool 使得 slot_mapping 是 pool 粒度。
        n_prefill = num_tokens - num_decode_tokens
        if n_prefill > 0:
            # decode token 排在前面，prefill token 紧随其后。
            prefill_slice = slice(num_decode_tokens, num_tokens)
            _kpool_compress_insert(
                k[prefill_slice],
                gate_score[prefill_slice],
                compress_ape,
                kv_cache,
                slot_mapping[prefill_slice],
                index_kpool,
                head_dim,
                round_scale=(scale_fmt is not None),
            )
            # 把每个请求尚未凑满的 prefill pool 落到 tail cache，让 decode
            # （含 PD 传输之后）能接着补完；tail slot 用请求内的 pos % kpool 定位。
            # 只处理 batch 尾部 token 会漏掉多请求 prefill 里除最后一个之外的请求。
            if tail_kv_cache is not None and tail_prefix is not None:
                tail_meta = attn_metadata.get(_resolve_layer_name(tail_prefix))
                if tail_meta is not None:
                    assert isinstance(tail_meta, DeepseekV32IndexerMetadata)
                    kpool_ops.kpool_seed_tail_cache(
                        tail_kv_cache,
                        k[prefill_slice],
                        gate_score[prefill_slice],
                        tail_meta.slot_mapping[prefill_slice],
                        index_kpool,
                        head_dim,
                    )
    else:
        # 标准路径：逐 token fp8 量化 + scatter，覆盖全部 token。
        # 搬迁前这里是 if current_platform.is_rocm(): ...triton... else: 的运行时
        # 分支；NVIDIA 侧现在无条件走 _custom_ops，对应的 ROCm 实现改为
        # amd/sparse_indexer.py 中独立的 indexer_k_quant_and_cache_triton 调用。
        assert scale_fmt is not None
        ops.indexer_k_quant_and_cache(
            k,
            kv_cache,
            slot_mapping,
            quant_block_size,
            scale_fmt,
        )

vllm/models/glm5next/sparse_indexer.py

新增 22 行平台 dispatcher，是调用方唯一可见的入口；CodeRabbit 在此指出「非 ROCm 即 NVIDIA」的静默回落问题（未采纳修改）。

# vllm/models/glm5next/sparse_indexer.py（新增，+22）
# 硬件隔离的入口模块：调用方只 import 这一个符号，实现按后端拆到
# nvidia/ 与 amd/ 两个目录，公共 helper 放在 common/。
from vllm.platforms import current_platform
​
if current_platform.is_rocm():
    # ROCm 走 aiter + Triton 实现（indexer_k_quant_and_cache_triton 等）。
    from .amd.sparse_indexer import (
        SparseAttnIndexerKpool,
    )
else:
    # 注意：这里只做二分判断，CPU / XPU / 未知平台同样会落到 NVIDIA 分支，
    # CodeRabbit 建议改成显式 CUDA 判定并对其他平台直接报错，
    # 该建议在合并前未被采纳。
    from .nvidia.sparse_indexer import (
        SparseAttnIndexerKpool,
    )
​
__all__ = [
    "SparseAttnIndexerKpool",
]

评论区精华

review 由机器人意见加一个人类 approve 构成，最有价值的两条都指向拆分本身引入的契约问题。

1）Codex（chatgpt-codex-connector[bot]）在 vllm/models/glm5next/amd/sparse_indexer.py:668 给出 P1：ROCm 上 CustomOp.dispatch_forward() 绑定的是 forward_hip 而非 forward_native，AMD 子类只在 forward_native 放置实现，继承链会落到 forward_cuda 抛 NotImplementedError，命中场景是 custom_ops=["all"] 或 enforce-eager；建议让 forward_hip 承载实现、forward_native 委托。这是本次 review 唯一的 P1，合并前未见对应修复 commit。

2）CodeRabbit（coderabbitai[bot]）在 vllm/models/glm5next/sparse_indexer.py:17 提 Minor：dispatcher 的 if is_rocm(): amd else: nvidia 会让 CPU/未知平台静默回落到 NVIDIA 实现，应改为显式 CUDA 判定并对其他平台报错；该建议未被采纳。

3）人类侧只有 Isotr0py 的 APPROVED（无正文）。CI 侧 AndreasKaratzas 指出 AMD CI 超时「I think it timed out due to the queue congestion」，判定为排队问题而非代码问题；mergify[bot] 提示过 merge conflict 需要 rebase，作者用多次 Merge branch 'main' into ... 解决，7 个 commit 中有 4 个是合并 main。

• ROCm 上 CustomOp 绑定 forward_hip，而 AMD 子类只实现 forward_native (correctness): 属于拆分时引入的真实功能缺陷，但合并前未见对应修复 commit，当前仅以 ROCm 默认路径（forward_native）不触发来自证；需要在后续 PR 或小补丁中确认 forward_hip 是否已覆写。
• dispatcher 对非 ROCm 平台一律回落到 NVIDIA 实现 (design): 未被采纳，最终以 if is_rocm(): amd else: nvidia 合入；风险由 glm5next 在 XPU 上 package import 即抛 NotImplementedError 部分抵消。
• AI 辅助生成的 729 行 AMD 实现如何验证 (testing): 以端到端 eval 代替单元测试作为正确性证据；新增测试缺失，现有 CPU 测试只覆盖共享 helper 的 decode 布局与模块名断言，无法覆盖 AMD 类的方法分派问题。
• AMD CI 超时与 merge conflict 的推进节奏 (other): AMD CI 超时被判定为排队拥堵而非代码问题；作者通过 4 次合并 main 解决冲突，7 个 commit 中 4 个是 Merge branch 'main'。

风险与影响

• 风险：
  1. ROCm 方法名契约（P1，未修复）：vllm/models/glm5next/amd/sparse_indexer.py 只在 forward_native 实现，而 ROCm 上 CustomOp.dispatch_forward() 绑定 forward_hip，继承下去会调用 forward_cuda 抛 NotImplementedError；custom_ops=["all"] 或 enforce-eager 的 ROCm 部署会直接失败。
  2. dispatcher 静默回落：vllm/models/glm5next/sparse_indexer.py 对 CPU/XPU/未知平台一律导入 NVIDIA 实现，把报错点推到更深的算子调用处，可读性差，且掩盖了平台不支持的语义。
  3. import 路径 breaking：vllm.model_executor.layers.sparse_attn_indexer_kpool 模块路径消失，仓库内调用点与测试已同步，但仓库外的脚本、插件或复现用例会 ImportError。
  4. 共享 helper 的隐式约定：_decode_topk_seq_lens、kv_cache_as_quant_view、_gather_workspace_shapes 被 AMD 复用，其正确性依赖 NVIDIA kpool_compress 侧「pool_topk 按 padded 网格排布」「tail 用 pos % kpool 定位」等语义；两端语义若不同，共用 helper 会引入跨平台 bug。
  5. 测试覆盖不足：test_kpool_ops_dispatch_matches_platform 只断言 kpool_ops.__name__ 含 .amd. / .nvidia.，无法捕获 dispatch_forward 层面的方法名错误；CPU 单测只覆盖共享 helper 的 decode 布局，AMD 分支靠端到端 eval 兜底。
  6. AI 生成代码复核强度：729 行新 AMD 文件只换来 2 条 inline comment，review 深度与代码量不匹配（PR body 已主动说明需人工逐行复核）。
     - 影响：用户侧：NVIDIA 与 ROCm 的 GLM-5.3-Flash 推理行为预期不变，作者以 gsm8k 0.9371、aime25 pass@4=0.9333 佐证；唯一可能被感知的变化是 ROCm 在开启 custom_ops 或 enforce-eager 时直接抛 NotImplementedError。系统侧：稀疏索引器从通用层目录下沉到模型目录，「模型私有 kernel 不污染 model_executor/layers」的分层在实践中进一步落地，common/nvidia/amd + dispatcher 的三分法可被其他模型专属 kernel 复用。团队侧：AMD/NVIDIA 两条线可在各自文件独立演进，减少 merge 与 review 的相互干扰，代价是三处文件需要同步维护 helper 契约。影响范围集中在 glm5next 模型与前缀索引器路径，不涉及 dtype、KV cache 存储格式或对外 API schema，也没有改 buildkite/部署配置。
     - 风险标记：ROCm dispatch 路径可能抛 NotImplementedError, import 路径 breaking 变更, 缺少 AMD 端单元测试, dispatcher 静默回落, AI 辅助生成大块新代码

关联脉络

• PR #56176 [ROCm] [Bugfix] Enable Load and Inference of GLM-5.3-Flash Quark MXFP4 Checkpoint: 同属 GLM-5.3-Flash 支持主线，且改动了本 PR 同步更新的调用点 vllm/models/glm5next/common/attention.py，两者共同构成 glm5next 目录的修复与重构序列。
• PR #56845 [Refactor] Remove dead kernel code: 该清理 PR 删除了 vllm/models/glm5next/amd/ops/kpool_compress.py 等死代码，说明 glm5next 的 amd/ops 目录正在被持续整理，与本 PR 把 AMD 索引器落到 glm5next/amd/ 的方向一致。
• PR #56254 [DSA] Wire DeepGEMM sparse MQA logits into the DeepSeek V4.1 indexer: 同属 MLA 稀疏索引器族，涉及 kv_cache 量化视图与 DeepGEMM mqa logits 的布局约定；本 PR 抽出的 _gather_workspace_shapes / kv_cache_as_quant_view 与那条线共用同一套 workspace 契约。