[KVConnector][1/N] PP-aware handshake aggregation and intermediate-PP output plumbing

[zixi-qi]

Purpose

The goal of this PR is to make kv connector / engine / model runner / gpu worker PP-aware, to layout the foundations to support PP(pipeline parallelism) + PD disaggregation.

• This PR is connector agnostic, the changes here are needed for both NIXL connector and Mooncake connector to support PP + PD.
• This PR is a pure PP-aware refactor and does not introduce any behavior changes.

Changes

• Widen `TransferTopology._engines` to `dict[(engine_id, pp_rank), …]`; extend `EngineTransferInfo` with `remote_pp_rank` / `start_layer` / `end_layer` (default `0`). Existing public methods keep their names and call shapes; PP-aware helpers gain an optional `remote_pp_rank: int = 0`.
• New `SupportsPP` ABC next to `SupportsHMA` (abstract `set_xfer_handshake_metadata_pp_aware` + `supports_pp(connector)` helper). `MultiConnector` inherits `SupportsPP` and dispatches per child.
• Worker emits `{(pp_rank, tp_rank): metadata}`. Engine core converts at the connector boundary: PP-aware connectors get the tuple-keyed dict; non-PP connectors get the unwrapped `{tp_rank: metadata}` via the existing `set_xfer_handshake_metadata`.
• Non-last-PP-rank returns `EMPTY_MODEL_RUNNER_OUTPUT` (optionally carrying `kv_connector_output`) instead of `None`, so the KV output aggregator's non-None invariant holds.

Test Plan

• New unit tests: `test_transfer_topology_sharded.py`, `test_handshake_pp_aggregation.py`, `test_pp_intermediate_output.py` — 9 passed.
• Existing `test_nixl_connector.py` (59 passed, 2 skipped — same skips on main) and `test_mooncake_connector.py` pass unchanged.
• `pre-commit run --all-files` clean.
• Ran PD disagg and non-PD + PP cases e2e and verified everything still work as expected

AI-assisted PR. The submitter has reviewed every line.

[njhill]

@zixi-qi fyi I've also included the "intermediate-PP output plumbing" part of this in #43732 which cleans up some of the PP kv connector handling more generally.

[zixi-qi]

@zixi-qi fyi I've also included the "intermediate-PP output plumbing" part of this in #43732 which cleans up some of the PP kv connector handling more generally.

Sounds good! I will rebase once your PR merges

[zixi-qi]

• Rebase on main
• [TODO] Some changes overlap with [Core] Cleanup KVConnector handling with PP + fix MRV2  #43732 will rebase again after it merges first

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zixi-qi.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[njhill]

Thanks @zixi-qi, this looks mostly good to me.

However I think the SupportsPP interface is a bit heavy, it would be cleaner to manage this inside the KVConnectorBase_V1 class, and just check for existence of the pp_aware method on self rather than using a marker class.

Would be good for @NickLucche to bless the TransferTopology updates too when he has time!

[zixi-qi]

Thanks @zixi-qi, this looks mostly good to me.

However I think the SupportsPP interface is a bit heavy, it would be cleaner to manage this inside the KVConnectorBase_V1 class, and just check for existence of the pp_aware method on self rather than using a marker class.

Would be good for @NickLucche to bless the TransferTopology updates too when he has time!

Updated to remove SupportsPP interface, thanks for the suggestion!

[njhill]

Thanks @zixi-qi that looks much cleaner. I think we can clean up some comments further though.

One question - what happens if a non-PP-aware connector is used in a PP context. Can that still be valid? If not, we should think about where to validate this.

[zixi-qi]

Thanks @zixi-qi that looks much cleaner. I think we can clean up some comments further though.

One question - what happens if a non-PP-aware connector is used in a PP context. Can that still be valid? If not, we should think about where to validate this.

Thanks for the suggestion! I have cleaned up the excessive comments, will check the stacked PRs to do the same as well.

non-PP-aware connector is used in a PP context.
Good catch on this! I added a validation in the default set_xfer_handshake_metadata_pp_aware implementation that throws if pp_rank > 0 in this scenario

[zixi-qi]

cc @NickLucche would be great if you could take a look as well :)