执行摘要

• 一句话：修复 NIXL Mamba 状态块在多槽 block table 下的配对错误
• 推荐动作：立即合入以修复 K3 等 Mamba 模型的 P/D 正确性问题。建议后续跟进 ZhanqiuHu 的建议，将 Attention 组和 SSM 组的裁剪逻辑统一到公共基类，减少分散的 clipping 点，降低维护成本。同时可考虑在 Mamba 管理器中直接避免分配无用的 scratch 槽。

功能与动机

PR body 指出 NIXL P/D 传输假设请求的 Mamba 组只有一个本地状态块，但在 speculative config 下 Mamba 管理器会为每个请求分配 trailing scratch slots，而 Mamba prefix caching 使块列表成为多槽表。原本的假设导致：spec 在 P 和 D 上时，P 的 stale scratch 字节被 RDMA 覆盖 D 的零初始化 scratch 槽；spec 仅在 P 上时，单块规则读取远程尾部，导致静默状态损坏；align 模式下头裁剪可能将 D 的状态块与 P 的占位块配对。Issue #46694 修复了 Attention 组的 clip，但 Mamba 组未覆盖。

实现拆解

1. 调度器端：扩展 block ID 裁剪方法
   在 base_scheduler.py 中将 get_sw_clipped_blocks 重写为 get_exchange_clipped_blocks，新增 SSM 感知的块列表裁剪逻辑：根据 _ssm_spec_blocks（每组的 speculative scratch slot 数）和 _ssm_state_slots_are_positional（是否依赖位置索引）裁剪掉无用槽位。裁剪后的列表保证 P 与 D 之间 1:1 配对或位置对齐。pull_scheduler.py 和 push_scheduler.py 中的调用点同步改为新方法。

2. Worker 端：重构前缀缓存阶段的 SSM 对齐
   在 base_worker.py 的 _apply_prefix_caching 中，替换原有的单本地块断言（assert num_local_blocks == 1）为基于 num_local_blocks - num_remote_blocks <= 1 的差值上限检查。对于“all”模式的多槽列表，根据列表长度大小关系裁剪头部或尾部，确保位置对应；对于不可配对的状态（差值 > 1）则断言失败，避免静默数据错误。

3. 测试覆盖新增
   在 test_nixl_connector_hma.py 中新增三个测试：test_exchange_clipped_blocks_ssm_single_state 验证单状态模式下裁剪结果；test_exchange_clipped_blocks_ssm_positional_states 验证“all”模式下位置对齐的正确性；test_apply_prefix_caching_ssm_unpairable_slots_rejected 验证不可配对场景正确触发断言。test_nixl_push_connector.py 中适配了辅助函数 stub。

关键文件：

• vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py（模块 KV 传输调度；类别 source；类型 core-logic；符号 get_exchange_clipped_blocks, get_sw_clipped_blocks）: 核心变更：新增 SSM 状态的裁剪逻辑，重命名并扩展了 block ID 裁剪方法，决定 P/D 间传输的列表内容。
• vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py（模块 KV 传输工作；类别 source；类型 core-logic；符号 _apply_prefix_caching）: Worker 端前缀缓存阶段的 SSM 对齐逻辑重写，将单本地块断言替换为差值约束，保证多槽列表位置对齐或失败。
• tests/v1/kv_connector/unit/test_nixl_connector_hma.py（模块 测试用例；类别 test；类型 test-coverage；符号 test_apply_prefix_caching_ssm_unpairable_slots_rejected, test_exchange_clipped_blocks_ssm_single_state, test_exchange_clipped_blocks_ssm_positional_states）: 新增三个测试用例覆盖 SSM 状态配对的三种关键场景：单状态裁剪、多槽位置对齐、不可配对拒绝。
• vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py（模块 KV 传输调度；类别 source；类型 core-logic）: 调用点方法替换，将 get_sw_clipped_blocks 改为 get_exchange_clipped_blocks。
• vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_scheduler.py（模块 KV 传输调度；类别 source；类型 core-logic）: 同上，调用点方法替换。
• tests/v1/kv_connector/unit/test_nixl_push_connector.py（模块 测试用例；类别 test；类型 test-coverage）: 适配辅助函数 stub 以匹配新方法名。

关键符号：get_exchange_clipped_blocks, _apply_prefix_caching, test_apply_prefix_caching_ssm_unpairable_slots_rejected, test_exchange_clipped_blocks_ssm_single_state, test_exchange_clipped_blocks_ssm_positional_states

关键源码片段

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py

核心变更：新增 SSM 状态的裁剪逻辑，重命名并扩展了 block ID 裁剪方法，决定 P/D 间传输的列表内容。

# vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
​
# In __init__:
# Trailing scratch slots that mamba managers co-allocate per request
# for speculative decoding; None for non - SSM groups.
self._ssm_spec_blocks = [
    g.kv_cache_spec.num_speculative_blocks
    if isinstance(g.kv_cache_spec, MambaSpec)
    else None
    for g in kv_cache_config.kv_cache_groups
]
# Only "all" mode keeps a state per block position; the other modes
# keep a single running state in the last non - speculative slot.
self._ssm_state_slots_are_positional = (
    vllm_config.cache_config.mamba_cache_mode == "all"
)
​
def get_exchange_clipped_blocks(
    self, block_ids: BlockIds, clip_ssm: bool = True
) -> BlockIds:
    """Clip a request's block lists down to the transferable blocks.
​
    Sliding - window groups keep only the in - window tail.
    SSM groups keep only their state - bearing slots: trailing speculative
    scratch slots always go, and in single - state cache modes so does
    everything before the running state. "all" mode keeps its remaining
    slots, which the worker pairs position - wise.
​
    Use at every block - id exchange point. Pass clip_ssm = False for
    per - step partial lists (host - buffer save), where SSM strip does not
    apply.
    """
    if len(block_ids) == 0 or not self._is_hma_required:
        return block_ids
    assert len(block_ids) == len(self.blocks_per_sw), (
        "Number of KV cache groups must match"
    )
    clipped = []
    for i, blocks in enumerate(block_ids):
        if n_sw := self.blocks_per_sw[i]:
            # Sliding window clip: keep tail.
            blocks = blocks[-n_sw:]
        elif (
            clip_ssm
            and blocks
            and (n_spec_blocks := self._ssm_spec_blocks[i]) is not None
        ):
            # SSM groups: strip trailing speculative scratch slots.
            n_spec = min(n_spec_blocks, len(blocks))
            blocks = blocks[:-n_spec] if n_spec else blocks
            if not self._ssm_state_slots_are_positional:
                # Single - state modes: keep only the last slot (running state).
                blocks = blocks[-1:]
        clipped.append(blocks)
    return tuple(clipped)

评论区精华

ZhanqiuHu 提议将所有缓存组的 clipping 集中到一处（类似 #46694 对 FA 的处理），但承认不阻塞本 PR。njhill 同意复杂度问题，但希望能尽快修复 K3 的 P/D 正确性，后续再优化结构。NickLucche 认为不应该分配这些 scratch slots 而是根本避免，但 LGTM。整体 reviewer 一致认为修复正确且必要。

• 集中裁剪逻辑的提议 (design): njhill 认同，但希望先尽快修复 K3，后续再重构。
• 不应分配 scratch slots 的讨论 (design): 未进一步讨论，当前修复可接受。
• 合并冲突 (other): njhill 最终 merge 了 main 并解决冲突。

风险与影响

• 风险：风险主要在兼容性：新逻辑依赖 _ssm_spec_blocks 和 _ssm_state_slots_are_positional 的正确配置，若未来引入新的 Mamba 缓存模式但未同步更新此处，可能导致配对错误。另外，get_exchange_clipped_blocks 替换了 get_sw_clipped_blocks，所有调用点（pull_scheduler.py、push_scheduler.py、以及可能的外部集成）必须确保正确调用，否则滑动窗口裁剪会失效或遗漏 SSM 裁剪。Worker 端的断言变更将不可配对状态从静默错误转为主动失败，可能暴露新场景下未预期的列表长度差异，但这是安全改进。
• 影响：直接影响：启用 NIXL KV Connector 且使用 Mamba 模型（如 Kimi K3）并在 speculative decoding 或 prefix caching 场景下的用户可避免静默状态损坏或崩溃。间接影响：所有使用 NIXL 的 P/D 传输流程（pull/push）均受益于更严格的配对契约。影响范围限定于 KV Connector 模块，不触及前端、调度核心或模型推理。
• 风险标记：核心路径变更（调度器&worker）, 新增配置属性依赖, Mamba 缓存模式兼容性

关联脉络

• PR #46694 [P/D][Bugfix] Fix PD async KV load lookahead handling for MTP spec decode: 之前修复了 Attention 组在 speculative decoding 下的 lookahead 处理，本 PR 是同一方向但对 Mamba 组的扩展。