执行摘要

• 一句话：NIXL push 连接器支持 PP+HMA 预填充按层名路由
• 推荐动作：值得精读。重点看 base_worker.py 中 _set_region_layers、_requires_layer_name_routing、_align_remote_regions_by_layer 三个方法以及 _compute_desc_ids 的 layer-name 分支："按层名身份而非 region 索引路由" + "缺失即响亮失败"的设计对任何跨进程资源映射问题都有借鉴价值；完成语义与 #56104 的兼容处理也展示了如何在推进新特性时避免与并行重构冲突。若你维护 kv-connector 或分布式 KV 传输，建议同时跟进 #50499。

功能与动机

PR body 明确说明动机："Enable attention HMA with pipeline-parallel prefill in NixlPushConnector, following #45880. Producer PP stages and a PP1 decoder may pool the same layers into different physical KV-cache regions. This PR matches layers by name and builds descriptors using each side's region geometry." 此前 push 连接器在 PP>1 且 HMA 时直接抛 NotImplementedError，pull 模式虽有 PP+HMA 支持但实现路径不同；该 PR 补上 push 模式缺口，使 PP 分片 prefiller 可以把 KV 写入 PP1 decoder 的 HMA 池化 region。

实现拆解

实现按以下 6 步拆解：

1. 元数据与 wire 协议升级（metadata.py）：NixlAgentMetadata 新增 region_members: list[list[str]] 字段，记录每个广告 region 共享的 layer 名；NIXL_CONNECTOR_VERSION 从 11 升至 12。原因：PP 下 P/D 的 region 几何可能不同，必须把 layer 身份带上握手线；版本号升级保证不兼容的旧 peer 在兼容哈希阶段被优雅拒绝（两阶段解码）。

2. 本地 layer 布局推导（base_worker.py）：在注册 KV cache 时调用 _set_region_layers(region_layers)，先断言同一 layer 不跨多个 NIXL region，再通过 _requires_layer_name_routing()（条件为 _supports_pp_hma and pp_size > 1 and _is_hma_required and not _has_mamba）判断是否启用以名路由；启用时从 kv_cache_config.transfer_group_index_by_layer 派生 _transfer_layer_names、_transfer_layer_region_indices、_transfer_layer_group_ids 三元组。默认（pull、PP1、非 HMA、含 Mamba）保持空元组，完全走旧路径，回归风险被压缩。

3. 远端 region 按层对齐（_align_remote_regions_by_layer）：在 add_remote_agent 握手阶段调用。先断言远端通告了 region_members 且各元数据数组长度一致，建立 layer → region 反向映射并拒绝重复通告；再按本地 _transfer_layer_names 顺序重排 kv_caches_base_addr、block_lens、block_strides、region_num_blocks、region_group_ids、region_mem_types、region_names，并把 region_members 归一化为每 region 一层，使第二次对齐成为 no-op。任何本地 layer 在远端缺失都会断言失败——宁可握手失败也不静默写陈旧 KV。

4. descriptor 计算与完成语义：_compute_desc_ids 新增 layer-name 路由分支：当 _transfer_layer_group_ids 非空时，按 layer 顺序用各 region 累计容量作为偏移（region_offsets = np.cumsum([0, *region_num_blocks[:-1]])），物理分配只注册一次。push_worker.py 的 _xfer_blocks_for_req 对 groups_differ 的拒绝增加豁免：仅当未启用 layer-name 路由时才拒绝 P/D 不同 cache-group 布局。完成语义方面，移除 _pop_done_transfers 的 is_send 透传：所有 transfer handle 结束后即上报完成，receive 侧的去重与 pending notif 处理移到 get_finished()，push WRITE 的完成由 consumer notif / lease expiry 决定。

5. 启动门禁：__init__ 中把全部 PP 相关检查收进 if self.pp_size > 1 分组：pull + HMA 拒绝（提示改用 push）、Mamba/SSM 混合布局拒绝、consumer 侧 PP 拒绝。NixlPushConnectorWorker 置 _supports_pp_hma = True。不等 block size 与非 MLA 的 decode TP > prefill TP 在 topology/NIXL 注册之前拒绝。

6. 测试与文档配套：test_nixl_push_connector.py 新增 _agent_metadata、_layer_routing_worker 辅助与 5 个 layer 路由测试（descriptor 路由、元数据往返、非 HMA 路径保持、缺层响亮失败、pooled region 扩张），并把完成语义测试参数化为 PP1/PP2+HMA 两组；test_nixl_desc_geometry.py 参数化 push_pp 覆盖 PP 下 region_members 通告与 descriptor 几何；test_nixl_connector_hma.py 补充 _transfer_layer_group_ids 桩。两份文档分别阐述按名路由的设计不变量与兼容性限制。

关键文件：

• vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py（模块 KV连接器；类别 source；类型 core-logic；符号 _set_region_layers, _requires_layer_name_routing, _align_remote_regions_by_layer, _compute_desc_ids）: 实现核心：新增 _set_region_layers、_requires_layer_name_routing、_align_remote_regions_by_layer 三个方法，改造 _compute_desc_ids 支持 layer-name 路由，并在 __init__ 重组 PP 门禁。
• vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py（模块 元数据；类别 source；类型 data-contract；符号 NixlAgentMetadata, NIXL_CONNECTOR_VERSION）: wire 协议变更点：NixlAgentMetadata 新增 region_members 字段，NIXL_CONNECTOR_VERSION 11 → 12，是 P/D 互操作边界。
• vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py（模块 推送连接器；类别 source；类型 core-logic；符号 _supports_pp_hma, _xfer_blocks_for_req）: push 连接器声明 _supports_pp_hma = True，并在 _xfer_blocks_for_req 中对 P/D 不同 cache-group 布局增加 layer-name 路由豁免。
• tests/v1/kv_connector/unit/test_nixl_push_connector.py（模块 单元测试；类别 test；类型 test-coverage；符号 test_completed_push_send_waits_for_consumer_notif, test_different_region_groups_require_layer_routing, _agent_metadata, _layer_routing_worker）: 最大测试文件（+398/-2）：覆盖 layer 路由 descriptor 计算、元数据往返、非 HMA 路径保持、缺层响亮失败、pooled region 扩张，并把完成语义测试参数化为 PP1/PP2+HMA。
• tests/v1/kv_connector/unit/test_nixl_desc_geometry.py（模块 单元测试；类别 test；类型 test-coverage；符号 test_overlaid_transfer_groups_share_region_geometry）: 把 test_overlaid_transfer_groups_share_region_geometry 参数化为 push/pull 两条路径，验证 PP 下 region_members 通告与 descriptor 几何（含 num_descs、dst_region_num_blocks）。
• tests/v1/kv_connector/unit/test_nixl_connector_hma.py（模块 单元测试；类别 test；类型 test-coverage）: 为 _compute_desc_ids 的 mock 补充 _transfer_layer_group_ids 空元组默认值，保持既有 HMA 测试在新增分支下的兼容。
• docs/design/nixl_kv_push_connector.md（模块 设计文档；类别 docs；类型 documentation）: 新增 "Pipeline parallelism and hybrid KV caches" 一节，系统阐述按名路由的设计不变量（成员唯一、缺层失败、仅本地布局决定 member 顺序）。
• docs/features/nixl_connector_compatibility.md（模块 特性文档；类别 docs；类型 documentation）: 补充 push PP+HMA 的能力边界与限制清单（decode 侧 PP 不支持、Mamba/SSM 不支持、P/D block size 必须一致、decode TP 不得大于 prefill TP）。

关键符号：_set_region_layers, _requires_layer_name_routing, _align_remote_regions_by_layer, _compute_desc_ids, _xfer_blocks_for_req, add_remote_agent, _pop_done_transfers, get_finished

关键源码片段

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

wire 协议变更点：NixlAgentMetadata 新增 region_members 字段，NIXL_CONNECTOR_VERSION 11 → 12，是 P/D 互操作边界。

# Version History:
# 1: Initial version with compatibility checking
# 2: Add remote_request_id to kv_transfer_params
# ...
# 11: Add per-region transfer geometry and memory types to NixlAgentMetadata
# 12: Add per-region member names for PP push
#
# 版本升级意味着旧 peer 解码新消息会失败；配合两阶段握手
#（先比较 compatibility_hash，再解码 agent_metadata_bytes）
# 保证不兼容的 P/D 组合在握手期被优雅拒绝。
NIXL_CONNECTOR_VERSION: int = 12
​
​
@dataclass
class NixlAgentMetadata:
    ...
    region_num_blocks: list[int] | None = None
    region_group_ids: list[int] | None = None
    region_names: list[str] | None = None
    region_mem_types: list[str] | None = None
    dcp_size: int = 1
    pcp_size: int = 1
    # Layer names sharing each advertised region, in region order.
    # PP 分片 prefill 靠它把本 stage 的 layer 与远端 region 对齐。
    region_members: list[list[str]] = field(default_factory=list)

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

push 连接器声明 _supports_pp_hma = True，并在 _xfer_blocks_for_req 中对 P/D 不同 cache-group 布局增加 layer-name 路由豁免。

class NixlPushConnectorWorker(NixlBaseConnectorWorker):
    """Push-specific (WRITE) worker logic."""
​
    # 区分 push 与 pull 的 NIXL 兼容哈希因子。
    _TRANSFER_MODE: str = "push"
​
    # push 模式声明支持 PP + HMA；pull 连接器保持 _supports_pp_hma = False，
    # 因此 PP>1 + HMA 组合会在 __init__ 时被明确拒绝并提示改用 push。
    _supports_pp_hma = True
​
    def _xfer_blocks_for_req(self, req_id: str, meta: ReqMeta):
        ...
        local_region_groups = self.region_group_ids
        remote_region_groups = self.dst_region_group_ids[engine_id]
        groups_differ = local_region_groups != remote_region_groups
        # 开启 layer-name 路由后，P/D 两侧的 region group 布局允许不同
        #（HMA 的 pooled 区域由 layer 名对齐）；未开启时保持原有拒绝逻辑。
        if groups_differ and not self._transfer_layer_group_ids:
            raise NotImplementedError(
                "NixlPushConnector does not support different producer and "
                "consumer cache-group layouts"
            )

评论区精华

核心讨论围绕命名、完成语义与校验时机展开：

• NickLucche 总评认可方向："I think this is much better than the original MemberTransferPlan I initially reviewed! This looks like the right direction at first pass."
• 关于 _pop_done_transfers 的 is_send 透传，NickLucche 指出 #56104 已改动该函数："we might not need to thread it through _pop_done_transfers"。作者随后移除 is_send 管道，把 receive 侧工作移入 get_finished()，让完成的 WRITE 直接上报 finished_sending 并清理 producer lease。
• 命名与属性归属：NickLucche 质疑类属性写法与 "member" 命名过于泛化；作者改为实例属性，并把 _supports_member_identity 重命名为 _supports_pp_hma，"member" 统一改为 "layer name"。
• CodeRabbit 提出两个正确性问题：① 应拒绝 block_size_ratio != 1 的 member routing（"This fails the handshake before a transfer starts"）；② 校验若发生在 add_remote_agent() 之后会遗留未跟踪的 NIXL agent（"Repeated invalid handshakes can retain untracked NIXL agents"）。作者以提交 "Validate PP HMA peers before registration" 将不等 block size 与非 MLA TP 扩张检查前移到注册之前。

• njhill 的 nit 建议把 PP 检查统一收进 if self.pp_size > 1，作者照做并确认。

• is_send 透传 vs 完成语义重构 (design): 采用与 #56104 兼容的完成语义：WRITE 完成由 consumer notif / lease expiry 判定，而非本地 handle；_pop_done_transfers 只负责全部 handle 结束后上报。

• 不等 block size 应在握手期拒绝 (correctness): 最终在 add_remote_agent 中通过 self.block_size != nixl_agent_meta.block_size 拒绝并在文档声明 Equal P/D block sizes required。
• 校验时机与远端 agent 泄漏 (correctness): 提交 "Validate PP HMA peers before registration" 把不等 block size 与非 MLA TP 扩张检查前移到 topology/NIXL 注册之前。
• 成员命名与类属性归属 (design): 作者改为实例属性、"member" 精确化为 "layer name"、_supports_member_identity 重命名为 _supports_pp_hma。
• 缺层时的失败策略 (correctness): 以断言方式在握手期响亮失败，作为既定的安全策略。
• PP 启动检查分组 (style): 作者照做，提交 "Group pipeline-parallel startup checks" 统一分组并保持校验顺序与错误信息。

风险与影响

• 风险：
  1. wire 协议兼容性：NIXL_CONNECTOR_VERSION 11 → 12 属于不兼容变更，P/D 必须同步升级；混合版本部署会在兼容哈希阶段失败。这是预期行为，但升级窗口需要统筹。
  2. 正确性风险集中在按名路由：_align_remote_regions_by_layer 依赖 transfer_group_index_by_layer 与 region_members 的一致性，任何 layer 命名错位都会在握手期断言失败而非静默转移——故障模式安全，但需要两侧 checkpoint/模型命名严格一致。
  3. 完成语义变更影响共用路径：_pop_done_transfers 与 get_finished 是 push/pull 共用代码，改动虽被回归测试覆盖（含失败 handle 后存活 handle 完成的场景），但 pull 模式 receive 完成路径仍存在回归面。
  4. 异常路径 agent 泄漏：CodeRabbit 指出校验晚于 add_remote_agent 会遗留未跟踪 agent；已通过前移校验缓解，但异常路径下 remove_remote_agent 的兜底仍建议后续关注。
  5. 验证缺口：PR body 明确承认 E2E 未在 rebase 后重跑，模型评测结果属于更早 commit；TRITON_FLASHINFER 复合后端在 cache 注册阶段失败，未纳入支持范围。
     - 影响：对用户：解锁 push 模式下 PP 分片 prefill + HMA 的部署组合（如 Gemma-4-26B-A4B-it 的 PP2/TP1 → PP1/TP1），GSM8K strict 精度 96.0% vs 直连 baseline 95.0%，并完成 217,085 次外部命中验证 decoder 零本地 prefill。pull 模式 PP+HMA 依旧在启动时被拒绝，行为不变。对系统：新增 wire 字段与版本号，P/D 部署需同步升级；descriptor 计算在握手期增加常数级重排开销，运行时无额外成本。对团队：kv-connector 模块的核心演进，为 #50499（packed MLA）提供可复用的 _set_region_layers 机制，且完成语义与 #56104 对齐，降低后续合并冲突概率。
     - 风险标记：wire 协议版本升级, 完成语义变更, 异常路径 agent 泄漏风险, P/D 必须同版本部署, TRITON_FLASHINFER 未覆盖

关联脉络

• PR #55055 [BugFix][Model Runner V2][Spec Decode] Fix decode instance's multi-layer MTP kv caches during P/D: 与本 PR 共用 tests/v1/kv_connector/unit/test_nixl_push_connector.py，同属 push 连接器 P/D KV 生命周期与完成语义路径的连续迭代。
• PR #50499 [KVConnector][NIXL] Packed MLA（标题按 PR body 推断）: PR body 声明为 stacked follow-up（"Packed MLA is the stacked follow-up #50499, which should land after this PR"）；作者在 review 中说明 _set_region_layers 会被其复用。
• PR #56104 [NIXL] 完成处理重构（review 中引用，标题按上下文推断）: NickLucche 在该 PR 中改动了 _pop_done_transfers；本 PR 据此移除 is_send 透传，保持完成语义兼容，降低后续合并冲突。
• PR #45880 [KVConnector][NIXL] attention-HMA 基础支持（PR body 引用，标题按上下文推断）: 本 PR 明确 "following #45880"，在其 HMA 支持之上扩展 PP 场景；pull 模式 PP/HMA 由 #43368/#48263 走不同实现，原型 #47981 已被本方案取代。