@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7802c20c4e

[KVConnector][NIXL] Support pipeline-parallel prefill in push mode (#45880)

Signed-off-by: zixi-qi <zixi@inferact.ai> Co-authored-by: Nick Hill <nickhill123@gmail.com>

qizixi committed Jul 8, 2026 at 16:49 UTC 7802c20c4e02d1c403fc4140d9b8e7b971e99d20

15 files changed +506 -104

.buildkite/test_areas/disaggregated.yaml

+13

@@ -39,6 +39,19 @@ steps:

39 - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh

40 - FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

42 +- label: Push NixlConnector PP prefill PD accuracy (4 GPUs)

43 + key: push-nixlconnector-pp-prefill-pd-accuracy-4-gpus

44 + timeout_in_minutes: 30

45 + working_dir: "/vllm-workspace/tests"

46 + num_devices: 4

47 + source_file_dependencies:

48 + - vllm/distributed/kv_transfer/kv_connector/v1/nixl/

49 + - tests/v1/kv_connector/nixl_integration/

50 + - tests/v1/kv_connector/nixl_push_integration/

51 + commands:

52 + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh

53 + - bash v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh

54 +

55 - label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)

56 key: dp-ep-distributed-nixlconnector-pd-accuracy-tests-4-gpus

57 timeout_in_minutes: 30

examples/disaggregated/disaggregated_serving/disagg_proxy_pushconnector_demo.py

+15 -4

@@ -9,9 +9,9 @@ the same; the difference is in how P and D coordinate the KV transfer:

9 * Pull mode: proxy forwards P's ``kv_transfer_params`` (including

10 ``remote_block_ids``) to D, and D pulls KV from P via NIXL READ.

11 * Push mode: proxy hands D **only** P's coordinates

12 - (``remote_engine_id``, ``remote_host``, ``remote_port``, ``tp_size``)

13 - and the shared ``remote_request_id``. D registers its locally allocated

14 - blocks with P over a NIXL notification; P then pushes the KV to D via

12 + (``remote_engine_id``, ``remote_host``, ``remote_port``, ``tp_size``,

13 + ``pp_size``) and the shared ``remote_request_id``. D registers its locally

14 + allocated blocks with P over a NIXL notification; P then pushes the KV to D via

15 NIXL WRITE.

17 Launch multiple vLLM instances configured with ``NixlPushConnector`` and

@@ -26,6 +26,7 @@ disagg_proxy_pushconnector_demo.py \

26 --prefill-kv-host 10.0.0.1 \

27 --prefill-side-channel-port 5600 \

28 --prefill-tp-size 1 \

29 + --prefill-pp-size 1 \

30 --port 8000

31 """

@@ -89,6 +90,7 @@ class PushProxy:

90 prefill_kv_host: str,

91 prefill_side_channel_port: int,

92 prefill_tp_size: int,

93 + prefill_pp_size: int,

94 custom_create_completion: Callable[[Request], StreamingResponse] | None = None,

95 custom_create_chat_completion: Callable[[Request], StreamingResponse]

96 | None = None,

@@ -110,6 +112,7 @@ class PushProxy:

112 "remote_host": prefill_kv_host,

113 "remote_port": prefill_side_channel_port,

114 "tp_size": prefill_tp_size,

115 + "pp_size": prefill_pp_size,

116 }

117

118 self.custom_create_completion = custom_create_completion

@@ -198,6 +201,7 @@ class PushProxy:

201 "prefill_kv_host": self.push_metadata["remote_host"],

202 "prefill_side_channel_port": self.push_metadata["remote_port"],

203 "prefill_tp_size": self.push_metadata["tp_size"],

204 + "prefill_pp_size": self.push_metadata["pp_size"],

205 }

206

207 # ── push-mode request handling ──────────────────────────────────── #

@@ -265,7 +269,7 @@ class PushProxy:

269 generator = self.forward_request(

270 f"http://{decode_instance}{path}", decode_request, headers

271 )

268 - return StreamingResponse(generator)

272 + return StreamingResponse(generator, media_type="application/json")

273

274 async def create_completion(self, raw_request: Request):

275 try:

@@ -316,6 +320,7 @@ class PushProxyServer:

320 prefill_kv_host=args.prefill_kv_host,

321 prefill_side_channel_port=args.prefill_side_channel_port,

322 prefill_tp_size=args.prefill_tp_size,

323 + prefill_pp_size=args.prefill_pp_size,

324 custom_create_completion=create_completion,

325 custom_create_chat_completion=create_chat_completion,

326 )

@@ -420,6 +425,12 @@ def parse_args():

425 default=1,

426 help="Tensor parallel size of the prefill vLLM instance",

427 )

428 + parser.add_argument(

429 + "--prefill-pp-size",

430 + type=int,

431 + default=1,

432 + help="Pipeline parallel size of the prefill vLLM instance",

433 + )

434 return parser.parse_args()

435

436

tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

-1

@@ -40,7 +40,6 @@ sw_attn_configs=(

40 # Gemma4: SW + cross-layer KV sharing

41 "GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-4-E2B-it VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"

42 )

43 -

43 # Select config array based on DP_EP env var

44 if [[ -n "${DP_EP:-}" ]]; then

45 configs=("${dp_ep_configs[@]}")

tests/v1/kv_connector/nixl_integration/nixl_side_channel_probe.py

+2 -1

@@ -15,6 +15,7 @@ def parse_args() -> argparse.Namespace:

15 parser = argparse.ArgumentParser()

16 parser.add_argument("--host", required=True)

17 parser.add_argument("--port", required=True, type=int)

18 + parser.add_argument("--pp-rank", default=0, type=int)

19 parser.add_argument("--rank", default=0, type=int)

20 parser.add_argument("--timeout-ms", default=1000, type=int)

21 return parser.parse_args()

@@ -37,7 +38,7 @@ def main() -> None:

38 sock.setsockopt(zmq.RCVTIMEO, args.timeout_ms)

39 try:

40 sock.connect(make_zmq_path(args.host, args.port))

40 - sock.send(msgspec.msgpack.encode((GET_META_MSG, args.rank)))

41 + sock.send(msgspec.msgpack.encode((GET_META_MSG, args.pp_rank, args.rank)))

42 sock.recv()

43 finally:

44 sock.close()

tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

+12 -6

@@ -49,13 +49,16 @@ else

49 KV_EXTRA_CONFIG=''

50 fi

52 +# Connector: default pull NixlConnector; NixlPushConnector enables PP prefill.

53 +KV_CONNECTOR=${KV_CONNECTOR:-NixlConnector}

54 +

55 # Build the kv-transfer-config for P and D

56 if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then

54 - KV_CONFIG_P='{"kv_connector":"NixlConnector","kv_role":"kv_producer"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'

55 - KV_CONFIG_D='{"kv_connector":"NixlConnector","kv_role":"kv_consumer"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'

57 + KV_CONFIG_P='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_producer"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'

58 + KV_CONFIG_D='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_consumer"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'

59 else

57 - KV_CONFIG_P="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_producer\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"

58 - KV_CONFIG_D="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_consumer\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"

60 + KV_CONFIG_P="{\"kv_connector\":\"$KV_CONNECTOR\",\"kv_role\":\"kv_producer\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"

61 + KV_CONFIG_D="{\"kv_connector\":\"$KV_CONNECTOR\",\"kv_role\":\"kv_consumer\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"

62 fi

64 # Models to run

@@ -72,6 +75,7 @@ fi

75 NUM_PREFILL_INSTANCES=${NUM_PREFILL_INSTANCES:-1} # Default to 1

76 NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-1} # Default to 1

77 PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1}

78 +PREFILLER_PP_SIZE=${PREFILLER_PP_SIZE:-1} # >1 requires NixlPushConnector

79 DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}

80 GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.2}

81 PREFILL_BLOCK_SIZE=${PREFILL_BLOCK_SIZE:-128}

@@ -138,8 +142,9 @@ run_tests_for_model() {

142 # Calculate GPU ID - we'll distribute across available GPUs

143 GPU_ID=$((i % $(get_num_gpus)))

144 NEXT_GPU=${GPU_ID}

141 - # If PREFILLER_TP_SIZE is more than 1

142 - for (( j=1; j < PREFILLER_TP_SIZE; j++ )); do

145 + # Reserve TP*PP GPUs for the prefiller (TP shards across PP stages).

146 + PREFILLER_WORLD_SIZE=$((PREFILLER_TP_SIZE * PREFILLER_PP_SIZE))

147 + for (( j=1; j < PREFILLER_WORLD_SIZE; j++ )); do

148 NEXT_GPU=$(((GPU_ID + j) % $(get_num_gpus)))

149 GPU_ID="${GPU_ID},${NEXT_GPU}"

150 done

@@ -161,6 +166,7 @@ run_tests_for_model() {

166 --block-size ${PREFILL_BLOCK_SIZE} \

167 --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \

168 --tensor-parallel-size $PREFILLER_TP_SIZE \

169 + --pipeline-parallel-size $PREFILLER_PP_SIZE \

170 --kv-transfer-config '$KV_CONFIG_P'"

171 if [[ "$ENFORCE_EAGER" == "1" ]]; then

172 BASE_CMD="${BASE_CMD} --enforce-eager"

tests/v1/kv_connector/nixl_push_integration/config_sweep_accuracy_test.sh new

+41

@@ -0,0 +1,41 @@

1 +#!/usr/bin/env bash

2 +set -euo pipefail

3 +

4 +# NixlPushConnector (push-mode PD) integration tests. Kept separate from the

5 +# pull-mode sweep in ../nixl_integration so push coverage can grow on its own;

6 +# the per-config runner is shared.

7 +SCRIPT="v1/kv_connector/nixl_integration/run_accuracy_test.sh"

8 +IMPORT_CANARY="v1/kv_connector/nixl_integration/test_nixl_imports.py"

9 +

10 +echo "=== Running NIXL import canary ==="

11 +python3 -m pytest -s -x "${IMPORT_CANARY}"

12 +

13 +# PP=2 prefill -> TP=2 decode: each PP stage writes its layer slice into the

14 +# decoder's matching region range (non-MLA and MLA contiguous-slice paths).

15 +configs=(

16 + "KV_CONNECTOR=NixlPushConnector GPU_MEMORY_UTILIZATION=0.6 PREFILLER_PP_SIZE=2 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2" # non-MLA, Qwen3-0.6B

17 + "KV_CONNECTOR=NixlPushConnector GPU_MEMORY_UTILIZATION=0.8 PREFILLER_PP_SIZE=2 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA

18 +)

19 +

20 +run_tests() {

21 + local label=$1

22 + local extra_args=$2

23 +

24 + echo "=== Running push tests (${label}) ==="

25 + for cfg in "${configs[@]}"; do

26 + local -a cfg_parts extra_args_parts

27 + read -r -a cfg_parts <<< "$cfg"

28 + read -r -a extra_args_parts <<< "$extra_args"

29 +

30 + echo "-> Running with ${cfg} ${extra_args:+and ${extra_args}}"

31 + # Use 'env' to safely set variables without eval; keep argv splitting

32 + # safe and SC2086-clean via arrays.

33 + if ! env "${cfg_parts[@]}" bash "${SCRIPT}" "${extra_args_parts[@]}"; then

34 + echo "❌ Test failed for config: ${cfg} ${extra_args:+(${extra_args})}"

35 + exit 1

36 + fi

37 + done

38 + echo "✅ All ${label} push tests passed!"

39 +}

40 +

41 +run_tests "default backend" ""

tests/v1/kv_connector/unit/test_nixl_connector.py

+18 -13

@@ -408,11 +408,10 @@ def test_kv_transfer_handshake(dist_init):

408 decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)

409 expected_agent_metadata = decoder.decode(metadata.agent_metadata_bytes)

410

411 - # The scheduler connector expects metadata to be in

412 - # dict[int, KVConnectorHandshakeMetadata], where the first key is

413 - # the dp_rank, the second key is the tp_rank.

411 + # The scheduler connector expects metadata keyed by

412 + # (pp_rank, tp_rank).

413 scheduler_connector = scheduler.get_kv_connector()

415 - scheduler_connector.set_xfer_handshake_metadata({0: metadata})

414 + scheduler_connector.set_xfer_handshake_metadata_pp_aware({(0, 0): metadata})

415

416 # Simulate a request that finishes prefill, which returns

417 # corresponding NixlConnectorMetadata for decode instance.

@@ -502,8 +501,14 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

501 )

502

503 def _nixl_handshake(

505 - self, host: str, port: int, remote_tp_size: int, expected_engine_id: str

506 - ) -> dict[int, str]:

504 + self,

505 + host: str,

506 + port: int,

507 + remote_tp_size: int,

508 + expected_engine_id: str,

509 + remote_pp_size: int = 1,

510 + notif_agents_only: bool = False,

511 + ) -> dict[tuple[int, int], str]:

512 # Mimic slow _nixl_handshake, as well as bypass zmq communication.

513 time.sleep(self._hand_shake_latency)

514 # These should've been done in register_kv_caches(), called by

@@ -533,7 +538,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

538 # When remote tp_size > local tp_size, handshake with multiple

539 # remote ranks.

540 num_handshakes = 1 if tp_ratio > 0 else -tp_ratio

536 - remote_agents: dict[int, str] = {}

541 + remote_agents: dict[tuple[int, int], str] = {}

542 for remote_tp_rank in range(num_handshakes):

543 remote_agent_name = self.add_remote_agent(

544 NixlAgentMetadata(

@@ -554,7 +559,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

559 remote_tp_rank=remote_tp_rank,

560 remote_tp_size=remote_tp_size,

561 )

557 - remote_agents[remote_tp_rank] = remote_agent_name

562 + remote_agents[(0, remote_tp_rank)] = remote_agent_name

563 return remote_agents

564

565

@@ -749,7 +754,7 @@ class TestNixlHandshake:

754

755 def check_handshake(remote_tp_size: int):

756 tp_ratio = remote_tp_size // local_tp_size

752 - assert set(remote_agents.keys()) == set(range(tp_ratio))

757 + assert set(remote_agents.keys()) == {(0, r) for r in range(tp_ratio)}

758

759 remote_engine_id = worker.REMOTE_ENGINE_ID

760 remote_info = worker.transfer_topo.get_engine_info(remote_engine_id)

@@ -2044,7 +2049,7 @@ def test_shutdown_cleans_up_resources(default_vllm_config, dist_init):

2049 # P TP = 2 * D TP case, we should register 2 local handles

2050 worker.src_xfer_handles_by_tp_ratio = {-2: [456, 457]}

2051 worker.dst_xfer_side_handles = {"engine1": {0: 789}}

2047 - worker._remote_agents = {"engine1": {0: "agent1"}}

2052 + worker._remote_agents = {"engine1": {(0, 0): "agent1"}}

2053 # _cleanup_remote_engine (called by shutdown) also clears these:

2054 worker.kv_caches_base_addr["engine1"] = {0: [0xABC]}

2055 worker.dst_num_blocks["engine1"] = 50

@@ -2097,7 +2102,7 @@ def _setup_worker_with_remote_engine(

2102 )

2103

2104 engine_id = "remote-engine-1"

2100 - worker._remote_agents[engine_id] = {0: "agent_0", 1: "agent_1"}

2105 + worker._remote_agents[engine_id] = {(0, 0): "agent_0", (0, 1): "agent_1"}

2106 worker.dst_xfer_side_handles[engine_id] = {0: 100, 1: 200}

2107 worker.kv_caches_base_addr[engine_id] = {0: [0xABC]}

2108 worker.dst_num_blocks[engine_id] = 50

@@ -2993,7 +2998,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)

2998 local_block_len=worker.block_size * 4096,

2999 )

3000 worker._remote_agents[remote_engine_id] = {

2996 - rank: f"agent_p{rank}" for rank in range(prefill_tp_size)

3001 + (0, rank): f"agent_p{rank}" for rank in range(prefill_tp_size)

3002 }

3003 worker.dst_xfer_side_handles = {

3004 remote_engine_id: {rank: 100 + rank for rank in range(prefill_tp_size)}

@@ -3044,7 +3049,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)

3049

3050 # Broadcast goes to ranks {1, 2, 3} only, never to the read target.

3051 expected_recipients = {

3047 - worker._remote_agents[remote_engine_id][r]

3052 + worker._remote_agents[remote_engine_id][(0, r)]

3053 for r in range(1, prefill_tp_size)

3054 }

3055 assert {agent for agent, _ in send_notif_calls} == expected_recipients

tests/v1/kv_connector/unit/test_nixl_push_connector.py

+177 -1

@@ -29,9 +29,11 @@ from typing import Any

29 from unittest.mock import MagicMock, patch

31 import msgspec

32 +import pytest

34 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (

35 PUSH_REG_NOTIF_PREFIX,

36 + NixlAgentMetadata,

37 NixlConnectorMetadata,

38 )

39 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (

@@ -512,7 +514,7 @@ class TestPushWriterNotifs:

514 # Pretend the writer thread already forwarded a completion notif

515 # for a request whose KV is being received.

516 request_id = "req-recv-1"

515 - w._recving_metadata[request_id] = MagicMock()

517 + w._recving_metadata[request_id] = MagicMock(pp_size=1)

518 # Compose the standard completion notif: req_id:tp_size.

519 notif_msg = f"{request_id}:1".encode()

520 w._pending_completion_notifs.put(notif_msg)

@@ -813,3 +815,177 @@ class TestPushWriterNegative:

815 assert w._reqs_to_send[rid] >= now

816 # Unknown request must not be inserted by the heartbeat path.

817 assert "req-unknown" not in w._reqs_to_send

818 +

819 +

820 +# ----------------------------------------------------------------- #

821 +# Pipeline-parallel producer (push-mode PP-disagg) #

822 +# ----------------------------------------------------------------- #

823 +

824 +

825 +class TestPushPipelineParallel:

826 + """PP-sharded producer: per-stage completion counting, pp_size plumbing,

827 + and remote-region slicing."""

828 +

829 + def test_completion_waits_for_one_notif_per_pp_stage(self):

830 + """Each PP stage WRITEs its own layers and sends one notif; D must

831 + collect pp_size notifs before reporting the recv done."""

832 + w = _StubWriterWorker.fresh()

833 + w.transfer_topo = MagicMock()

834 + request_id = "req-pp-2"

835 + w._recving_metadata[request_id] = MagicMock(pp_size=2)

836 + notif = f"{request_id}:1".encode()

837 +

838 + # First stage: counted, not yet done.

839 + w._pending_completion_notifs.put(notif)

840 + assert w._get_new_notifs() == set()

841 + assert request_id not in w._recving_transfers

842 + assert w.consumer_notification_counts_by_req[request_id] == 1

843 +

844 + # Second (final) stage: now reported done.

845 + w._pending_completion_notifs.put(notif)

846 + assert w._get_new_notifs() == set()

847 + assert request_id in w._recving_transfers

848 + assert request_id not in w.consumer_notification_counts_by_req

849 +

850 + def test_req_meta_reads_pp_size_from_kv_transfer_params(self):

851 + """D learns the producer's pp_size from kv_transfer_params (forwarded

852 + by the proxy) and defaults to 1 when absent."""

853 + metadata = NixlConnectorMetadata()

854 + params = {

855 + "remote_block_ids": ([0],),

856 + "remote_engine_id": "p-engine",

857 + "remote_request_id": "p-req",

858 + "remote_host": "localhost",

859 + "remote_port": 1234,

860 + "tp_size": 1,

861 + "pp_size": 2,

862 + }

863 + metadata.add_new_req_to_recv("req", ([0],), params)

864 + assert metadata.reqs_to_recv["req"].pp_size == 2

865 +

866 + params.pop("pp_size")

867 + metadata.add_new_req_to_recv("req-default", ([0],), params)

868 + assert metadata.reqs_to_recv["req-default"].pp_size == 1

869 +

870 + def test_add_remote_agent_slices_remote_regions_to_local_pp_window(self):

871 + """With PP>1 the producer registered regions for the full model, but

872 + this worker holds only a contiguous layer slice; add_remote_agent

873 + trims the remote region list to [offset : offset + num_local_regions]

874 + before building descriptors. We stop right after the slice via a

875 + sentinel on the next collaborator call."""

876 + block_len = 4096 * 16

877 + w = _StubWriterWorker.fresh() # seeds writer-thread state for teardown

878 + w.pp_size = 2

879 + w._remote_region_offset = 2 # this worker owns layers [2, 4)

880 + w.block_len_per_layer = [block_len, block_len] # 2 local layers

881 + w.nixl_wrapper = MagicMock()

882 +

883 + class _StopAfterSlice(RuntimeError):

884 + pass

885 +

886 + w.transfer_topo = MagicMock()

887 + w.transfer_topo.register_remote_engine.side_effect = _StopAfterSlice()

888 +

889 + meta = NixlAgentMetadata(

890 + engine_id="p-engine",

891 + agent_metadata=b"agent",

892 + kv_caches_base_addr=[10, 11, 12, 13], # full 4-layer model

893 + device_id=0,

894 + num_blocks=4,

895 + block_lens=[block_len] * 4,

896 + kv_cache_layout="HND",

897 + block_size=16,

898 + ssm_sizes=(0, 0),

899 + attn_backend_name="FLASH_ATTN",

900 + physical_blocks_per_logical_kv_block=1,

901 + )

902 +

903 + with pytest.raises(_StopAfterSlice):

904 + w.add_remote_agent(meta, remote_tp_rank=0, remote_tp_size=1)

905 +

906 + # Sliced down to this worker's layer window (regions [2:4]).

907 + assert meta.kv_caches_base_addr == [12, 13]

908 + assert meta.block_lens == [block_len, block_len]

909 +

910 +

911 +class TestPushWriterMlaReplication:

912 + """MLA latent KV is replicated across D's TP ranks, so when D_TP > P_TP

913 + (``tp_ratio < 0``) the producer must *WRITE* the latent into every D rank

914 + it handshook -- not write one and merely notify the rest, which would

915 + leave the un-written ranks decoding against stale KV."""

916 +

917 + @staticmethod

918 + def _mla_worker_writing_to(d_ranks):

919 + from types import SimpleNamespace

920 +

921 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (

922 + TPMapping,

923 + )

924 +

925 + engine_id = "decode-engine"

926 + w = _StubWriterWorker.fresh()

927 + w.use_mla = True

928 + w.nixl_wrapper = MagicMock()

929 + w.transfer_topo = MagicMock()

930 + w.transfer_topo.get_engine_info.return_value = SimpleNamespace(

931 + remote_tp_size=len(d_ranks),

932 + remote_block_size=16,

933 + remote_physical_blocks_per_logical=1,

934 + )

935 + # D_TP > P_TP => negative ratio (one P rank feeds |ratio| D ranks).

936 + w.transfer_topo.tp_ratio.return_value = -len(d_ranks)

937 + # The tp-mapping collapses MLA to a single source rank (correct for

938 + # the pull/read direction); the push path must fan it back out.

939 + w.tp_mappings = {

940 + engine_id: TPMapping(

941 + source_ranks_per_group=((0,),),

942 + all_source_ranks=(0,),

943 + rank_to_attention_slot={0: 0},

944 + rank_offset_factor=0,

945 + )

946 + }

947 + w._logical_to_remote_kernel_block_ids = lambda block_ids, ratio: block_ids

948 + w.dst_xfer_side_handles = {engine_id: {r: 1000 + r for r in d_ranks}}

949 + w.src_xfer_handles_by_block_size = {16: 2000}

950 + w._remote_agents = {engine_id: {(0, r): f"agent-{r}" for r in d_ranks}}

951 + return w, engine_id

952 +

953 + def test_mla_hetero_tp_writes_every_d_rank(self):

954 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (

955 + RemoteMeta,

956 + ReqMeta,

957 + )

958 +

959 + w, engine_id = self._mla_worker_writing_to(d_ranks=(0, 1))

960 +

961 + written: list[int] = []

962 +

963 + def _fake_xfer(**kw):

964 + rank = kw["read_spec"].remote_rank

965 + written.append(rank)

966 + return 1000 + rank # in-flight handle, as the real method returns

967 +

968 + w._xfer_blocks = _fake_xfer

969 +

970 + meta = ReqMeta(

971 + local_block_ids=([100, 101],),

972 + local_physical_block_ids=([100, 101],),

973 + tp_size=1,

974 + remote=RemoteMeta(

975 + block_ids=([200, 201],),

976 + host="",

977 + port=0,

978 + engine_id=engine_id,

979 + request_id="d-req",

980 + ),

981 + )

982 +

983 + w._xfer_blocks_for_req(req_id="p-req", meta=meta)

984 +

985 + # Every handshook D rank must receive a real WRITE of the latent...

986 + assert sorted(written) == [0, 1]

987 + # ...and no rank may be fobbed off with a bare completion notif.

988 + assert w.nixl_wrapper.send_notif.call_count == 0

989 + # All of the request's WRITE handles must be tracked together, so the

990 + # engine thread never sees a partial set and double-frees the request.

991 + assert sorted(w._sending_transfers["p-req"]) == [1000, 1001]

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

+17 -11

@@ -185,6 +185,7 @@ class NixlBaseConnectorScheduler:

185 host = params.get("remote_host")

186 port = params.get("remote_port")

187 tp_size = params.get("tp_size")

188 + pp_size = params.get("pp_size", 1)

189 if (

190 remote_engine_id is None

191 or remote_request_id is None

@@ -199,6 +200,7 @@ class NixlBaseConnectorScheduler:

200 host=host,

201 port=port,

202 tp_size=tp_size,

203 + pp_size=pp_size,

204 )

205 self._heartbeat_by_engine[remote_engine_id].req_ids.add(remote_request_id)

206 self._heartbeat_req_engine[request.request_id] = (

@@ -244,7 +246,7 @@ class NixlBaseConnectorScheduler:

246 )

247

248 def set_xfer_handshake_metadata(

247 - self, metadata: dict[int, KVConnectorHandshakeMetadata]

249 + self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]

250 ) -> None:

251 """

252 Set the KV connector handshake metadata for this connector.

@@ -252,19 +254,20 @@ class NixlBaseConnectorScheduler:

254 Args:

255 metadata (dict): the handshake metadata to set.

256 """

255 - encoded_data: dict[int, bytes] = {}

257 + encoded_data: dict[tuple[int, int], bytes] = {}

258 encoder = msgspec.msgpack.Encoder()

257 - for tp_rank, rank_metadata in metadata.items():

259 + for (pp_rank, tp_rank), rank_metadata in metadata.items():

260 if not isinstance(rank_metadata, NixlHandshakePayload):

261 raise ValueError(

262 "NixlConnectorScheduler expects NixlHandshakePayload for "

263 "handshake metadata."

264 )

263 - encoded_data[tp_rank] = encoder.encode(rank_metadata)

265 + encoded_data[(pp_rank, tp_rank)] = encoder.encode(rank_metadata)

266 logger.debug(

265 - "Tp rank %d: encoded NixlHandshakePayload size: %s bytes",

267 + "PP rank %d, TP rank %d: encoded NixlHandshakePayload size: %s bytes",

268 + pp_rank,

269 tp_rank,

267 - str(len(encoded_data[tp_rank])),

270 + str(len(encoded_data[(pp_rank, tp_rank)])),

271 )

272

273 # Only start the listener when we have metadata to serve.

@@ -287,7 +290,7 @@ class NixlBaseConnectorScheduler:

290

291 @staticmethod

292 def _nixl_handshake_listener(

290 - encoded_data: dict[int, Any],

293 + encoded_data: dict[tuple[int, int], Any],

294 ready_event: threading.Event,

295 stop_event: threading.Event,

296 host: str,

@@ -310,15 +313,18 @@ class NixlBaseConnectorScheduler:

313 if stop_event.is_set():

314 break

315 continue

313 - # Decode the message which contains (GET_META_MSG, rank)

314 - msg, target_tp_rank = msgspec.msgpack.decode(msg)

316 + # Decode (GET_META_MSG, pp_rank, tp_rank).

317 + msg, target_pp_rank, target_tp_rank = msgspec.msgpack.decode(msg)

318 logger.debug(

316 - "Received message for tp rank %s",

319 + "Received message for pp rank %s, tp rank %s",

320 + target_pp_rank,

321 target_tp_rank,

322 )

323 if msg != GET_META_MSG:

324 logger.warning("Connection listener got unexpected message %s", msg)

321 - sock.send_multipart((identity, b"", encoded_data[target_tp_rank]))

325 + sock.send_multipart(

326 + (identity, b"", encoded_data[(target_pp_rank, target_tp_rank)])

327 + )

328

329 def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int:

330 """D-side only. Returns N-1 for Mamba models since the decoder

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

+120 -23

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Base worker-side logic for the NIXL connector."""

5 +import itertools

6 import logging

7 import os

8 import queue

@@ -336,8 +337,11 @@ class NixlBaseConnectorWorker:

337 )

338

339 self.nixl_wrapper = nixl_wrapper_cls(str(uuid.uuid4()), config)

339 - # Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}.

340 - self._remote_agents: dict[EngineId, dict[int, str]] = defaultdict(dict)

340 + # Map of engine_id -> {(pp_rank, tp_rank): agent_name, ...}.

341 + # non-PP remote uses pp_rank 0, i.e. (0, tp_rank).

342 + self._remote_agents: dict[EngineId, dict[tuple[int, int], str]] = defaultdict(

343 + dict

344 + )

345

346 # Metadata.

347 self.engine_id: EngineId = engine_id

@@ -411,6 +415,25 @@ class NixlBaseConnectorWorker:

415 # (so 1 per layer for MLA, otherwise 2 per layer)

416 self.num_regions = 0

417

418 + # PP>1 (push mode): this worker holds a contiguous layer slice and

419 + # transfers into the matching sub-range of a PP=1 remote's regions.

420 + self.pp_size = vllm_config.parallel_config.pipeline_parallel_size

421 + self._remote_region_offset = 0

422 + # PP push slices regions per layer (uniform count); HMA breaks that.

423 + if self.pp_size > 1 and self._is_hma_required:

424 + raise NotImplementedError(

425 + "NixlPushConnector does not support pipeline_parallel_size > 1 "

426 + "with hybrid KV cache layouts (HMA) yet."

427 + )

428 + # Decode-side PP is unsupported (completions counted per consumer rank).

429 + if vllm_config.kv_transfer_config.kv_role == "kv_consumer" and self.pp_size > 1:

430 + raise NotImplementedError(

431 + "NixlPushConnector consumer (decode) does not support "

432 + "pipeline_parallel_size > 1."

433 + )

434 + # Keep heartbeat handshakes to a PP-sharded producer notif-only.

435 + self._hb_handshake_notif_only = False

436 +

437 # nixl_prepped_dlist_handle.

438 self.src_xfer_handles_by_block_size: dict[int, int] = {}

439 # Populated dynamically during handshake based on remote configuration.

@@ -449,7 +472,7 @@ class NixlBaseConnectorWorker:

472 thread_name_prefix="vllm-nixl-handshake-initiator",

473 )

474 self._ready_requests = queue.Queue[tuple[ReqId, ReqMeta]]()

452 - self._handshake_futures: dict[EngineId, Future[dict[int, str]]] = {}

475 + self._handshake_futures: dict[EngineId, Future[dict[tuple[int, int], str]]] = {}

476 # Protects _handshake_futures and _remote_agents.

477 self._handshake_lock = threading.RLock()

478

@@ -536,7 +559,9 @@ class NixlBaseConnectorWorker:

559 port: int,

560 remote_tp_size: int,

561 expected_engine_id: str,

539 - ) -> dict[int, str]:

562 + remote_pp_size: int = 1,

563 + notif_agents_only: bool = False,

564 + ) -> dict[tuple[int, int], str]:

565 """Do a NIXL handshake with a remote instance."""

566

567 # the first time we connect to a remote agent.

@@ -558,20 +583,25 @@ class NixlBaseConnectorWorker:

583 # this happens to be the same single rank_i.

584 assert self.transfer_topo is not None

585 p_remote_ranks = self.transfer_topo.handshake_target_ranks(remote_tp_size)

561 - remote_rank_to_agent_name = {}

586 + remote_rank_to_agent_name: dict[tuple[int, int], str] = {}

587 path = make_zmq_path("tcp", host, port)

588

589 with zmq_ctx(zmq.REQ, path) as sock:

565 - for remote_rank in p_remote_ranks:

590 + for remote_pp_rank, remote_rank in itertools.product(

591 + range(remote_pp_size), p_remote_ranks

592 + ):

593 logger.debug(

567 - "Querying metadata on path: %s at remote tp rank %s",

594 + "Querying metadata on path: %s at remote pp rank %s, tp rank %s",

595 path,

596 + remote_pp_rank,

597 remote_rank,

598 )

599

600 start_time = time.perf_counter()

601 # Send query for the request.

574 - msg = msgspec.msgpack.encode((GET_META_MSG, remote_rank))

602 + msg = msgspec.msgpack.encode(

603 + (GET_META_MSG, remote_pp_rank, remote_rank)

604 + )

605 # Set receive timeout to 5 seconds to avoid hanging on dead server

606 sock.setsockopt(zmq.RCVTIMEO, 5000) # milliseconds

607 sock.send(msg)

@@ -638,17 +668,44 @@ class NixlBaseConnectorWorker:

668 )

669

670 # Register Remote agent.

641 - remote_agent_name = self.add_remote_agent(

642 - metadata, remote_rank, remote_tp_size

643 - )

671 + if notif_agents_only:

672 + remote_agent_name = self._add_notif_only_remote_agent(

673 + metadata, remote_tp_size

674 + )

675 + else:

676 + remote_agent_name = self.add_remote_agent(

677 + metadata, remote_rank, remote_tp_size

678 + )

679 setup_agent_time = time.perf_counter()

680 logger.debug(

681 "NIXL handshake: add agent took: %s",

682 setup_agent_time - got_metadata_time,

683 )

649 - remote_rank_to_agent_name[remote_rank] = remote_agent_name

684 + remote_ranks = (remote_pp_rank, remote_rank)

685 + remote_rank_to_agent_name[remote_ranks] = remote_agent_name

686 return remote_rank_to_agent_name

687

688 + def _add_notif_only_remote_agent(

689 + self, metadata: NixlAgentMetadata, remote_tp_size: int

690 + ) -> str:

691 + """Load a remote agent for notifs only on the push-mode decode side.

692 +

693 + Skips descriptor setup but records engine info for block accounting.

694 + """

695 + assert self.transfer_topo is not None

696 + self.transfer_topo.register_remote_engine(

697 + metadata.engine_id,

698 + EngineTransferInfo(

699 + remote_tp_size=remote_tp_size,

700 + remote_block_size=metadata.block_size,

701 + remote_block_len=metadata.block_lens[0],

702 + remote_physical_blocks_per_logical=(

703 + metadata.physical_blocks_per_logical_kv_block

704 + ),

705 + ),

706 + )

707 + return self.nixl_wrapper.add_remote_agent(metadata.agent_metadata)

708 +

709 def initialize_host_xfer_buffer(self, kv_caches: dict[str, torch.Tensor]) -> None:

710 """

711 Initialize transfer buffer in CPU mem for accelerators

@@ -761,7 +818,9 @@ class NixlBaseConnectorWorker:

818 host: str,

819 port: int,

820 tp_size: int,

764 - ) -> Future[dict[int, str]] | None:

821 + pp_size: int = 1,

822 + notif_agents_only: bool = False,

823 + ) -> Future[dict[tuple[int, int], str]] | None:

824 """

825 Ensure a handshake is in-flight (or already done) for *engine_id*.

826

@@ -784,10 +843,12 @@ class NixlBaseConnectorWorker:

843 port,

844 tp_size,

845 engine_id,

846 + pp_size,

847 + notif_agents_only,

848 )

849 self._handshake_futures[engine_id] = fut

850

790 - def done_callback(f: Future[dict[int, str]], eid=engine_id):

851 + def done_callback(f: Future[dict[tuple[int, int], str]], eid=engine_id):

852 with self._handshake_lock:

853 del self._handshake_futures[eid]

854 try:

@@ -1114,6 +1175,15 @@ class NixlBaseConnectorWorker:

1175 self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses

1176 self.num_regions = len(caches_data)

1177

1178 + if self.pp_size > 1:

1179 + start_layer, end_layer = self.model_config.get_layers_start_end_indices(

1180 + self.vllm_config.parallel_config

1181 + )

1182 + num_local_layers = end_layer - start_layer

1183 + assert num_local_layers > 0 and self.num_regions % num_local_layers == 0

1184 + regions_per_layer = self.num_regions // num_local_layers

1185 + self._remote_region_offset = regions_per_layer * start_layer

1186 +

1187 if self.transfer_topo.virtually_split_kv_in_blocks:

1188 # NOTE (NickLucche) When FlashInfer is used, memory is registered

1189 # with joint KV for each block. This minimizes the overhead in

@@ -1464,14 +1534,32 @@ class NixlBaseConnectorWorker:

1534 """ # noqa: E501

1535 engine_id = nixl_agent_meta.engine_id

1536 # TODO re-evaluate refreshing for scaling/recovery

1467 - if remote_tp_rank in self._remote_agents.get(engine_id, {}):

1537 + if (0, remote_tp_rank) in self._remote_agents.get(engine_id, {}):

1538 logger.debug(

1539 "Remote agent with engine_id %s and rank"

1540 "%s already exchanged metadata, skip handshake.",

1541 engine_id,

1542 remote_tp_rank,

1543 )

1474 - return self._remote_agents[engine_id][remote_tp_rank]

1544 + return self._remote_agents[engine_id][(0, remote_tp_rank)]

1545 +

1546 + # Compare physical regions, not self.num_regions (doubled by

1547 + # FlashInfer's virtual K/V split).

1548 + num_local_regions = len(self.block_len_per_layer)

1549 + if (

1550 + self.pp_size > 1

1551 + and len(nixl_agent_meta.kv_caches_base_addr) > num_local_regions

1552 + ):

1553 + # This worker holds a PP layer-slice; the PP=1 remote registered

1554 + # the full model. Slice its regions to our layer window so the

1555 + # logic below sees congruent local/remote lists.

1556 + start = self._remote_region_offset

1557 + end = start + num_local_regions

1558 + assert len(nixl_agent_meta.kv_caches_base_addr) >= end

1559 + nixl_agent_meta.kv_caches_base_addr = nixl_agent_meta.kv_caches_base_addr[

1560 + start:end

1561 + ]

1562 + nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end]

1563

1564 ### Register remote engine in TransferTopology (idempotent).

1565 assert self.transfer_topo is not None

@@ -2098,7 +2186,12 @@ class NixlBaseConnectorWorker:

2186 # the **next** heartbeat for this remote can go through.

2187 if (

2188 self._ensure_handshake(

2101 - engine_id, hb_info.host, hb_info.port, hb_info.tp_size

2189 + engine_id,

2190 + hb_info.host,

2191 + hb_info.port,

2192 + hb_info.tp_size,

2193 + hb_info.pp_size,

2194 + self._hb_handshake_notif_only and hb_info.pp_size > 1,

2195 )

2196 is not None

2197 ):

@@ -2374,19 +2467,23 @@ class NixlBaseConnectorWorker:

2467 """

2468 assert engine_id in self._remote_agents

2469

2377 - for handle in self.dst_xfer_side_handles.pop(engine_id).values():

2470 + # Notif-only engines (push-mode D side) have no descriptor state.

2471 + for handle in self.dst_xfer_side_handles.pop(engine_id, {}).values():

2472 self.nixl_wrapper.release_dlist_handle(handle)

2473 for agent_name in self._remote_agents.pop(engine_id).values():

2474 self.nixl_wrapper.remove_remote_agent(agent_name)

2475

2382 - del self.kv_caches_base_addr[engine_id]

2383 - del self.dst_num_blocks[engine_id]

2384 - del self.tp_mappings[engine_id]

2476 + self.kv_caches_base_addr.pop(engine_id, None)

2477 + self.dst_num_blocks.pop(engine_id, None)

2478 + self.tp_mappings.pop(engine_id, None)

2479 if self.transfer_topo is not None:

2480 self.transfer_topo.unregister_remote_engine(engine_id)

2481

2388 - last_active = self._engine_last_active.pop(engine_id)

2389 - if log_eviction:

2482 + # Push P-side engines are tracked in _remote_agents but not in

2483 + # _engine_last_active (they don't participate in stale eviction), so

2484 + # tolerate a missing entry.

2485 + last_active = self._engine_last_active.pop(engine_id, None)

2486 + if log_eviction and last_active is not None:

2487 logger.info(

2488 "Evicted stale remote engine %s (inactive for %.1fs).",

2489 engine_id,

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

+4 -3

@@ -207,11 +207,12 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):

207 assert self.connector_scheduler is not None

208 return self.connector_scheduler.request_finished(request, block_ids)

209

210 - def set_xfer_handshake_metadata(

211 - self, metadata: dict[int, KVConnectorHandshakeMetadata]

210 + def set_xfer_handshake_metadata_pp_aware(

211 + self, metadata: dict[tuple[int, int], KVConnectorHandshakeMetadata]

212 ) -> None:

213 """

214 - Set the KV connector handshake metadata for this connector.

214 + Set handshake metadata keyed by (pp_rank, tp_rank) so the side

215 + channel can serve every PP stage's agent metadata.

216

217 Args:

218 metadata (dict): the handshake metadata to set.

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

@@ -147,6 +147,7 @@ class HeartbeatInfo:

147 host: str

148 port: int

149 tp_size: int

150 + pp_size: int = 1

151

152

153 @dataclass

@@ -167,6 +168,8 @@ class ReqMeta:

168 remote: RemoteMeta | None = None

169 # Remote block size, discovered during NIXL handshake (push mode).

170 remote_block_size: int | None = None

171 + # Remote producer pipeline-parallel size (push mode, D side).

172 + pp_size: int = 1

173

174

175 class NixlConnectorMetadata(KVConnectorMetadata):

@@ -196,6 +199,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):

199 # P workers don't need to receive tp_size from proxy here.

200 tp_size=kv_transfer_params.get("tp_size", 1),

201 remote_block_size=kv_transfer_params.get("remote_block_size"),

202 + pp_size=kv_transfer_params.get("pp_size", 1),

203 )

204

205 def add_new_req_to_save(

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

+2 -2

@@ -183,7 +183,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

183 notif_id = f"{meta.remote.request_id}:{self.world_size}".encode()

184 remote_agents = self._remote_agents[meta.remote.engine_id]

185 for rank_to_notify, agent in remote_agents.items():

186 - if rank_to_notify != read_specs[0].remote_rank:

186 + if rank_to_notify != (0, read_specs[0].remote_rank):

187 self.nixl_wrapper.send_notif(agent, notif_msg=notif_id)

188

189 def _read_blocks(

@@ -250,7 +250,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

250 # just notify P worker that we have the blocks we need.

251 if len(local_block_ids) == 0:

252 # A full prefix cache hit is indicated with an empty list.

253 - agent_name = self._remote_agents[dst_engine_id][remote_rank]

253 + agent_name = self._remote_agents[dst_engine_id][(0, remote_rank)]

254 try:

255 self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id)

256 except Exception as e:

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

@@ -186,6 +186,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):

186 "remote_host": params["remote_host"],

187 "remote_port": params["remote_port"],

188 "remote_tp_size": params["tp_size"],

189 + "remote_pp_size": params.get("pp_size", 1),

190 }

191 self._push_registration_deadlines[request.request_id] = (

192 time.perf_counter() + self._push_registration_timeout

@@ -238,6 +239,8 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):

239 # serving layer via abort_immediately. To keep P from

240 # stranding the prefill blocks, we still register an empty

241 # recv so the worker emits a notif that lets P free them.

242 + # Seed remote_block_ids so add_new_req_to_recv won't KeyError.

243 + params["remote_block_ids"] = ()

244 self._reqs_need_recv[request.request_id] = (request, [])

245 params["do_remote_prefill"] = False

246 return False, None

@@ -285,6 +288,7 @@ class NixlPushConnectorScheduler(NixlBaseConnectorScheduler):

288 remote_host=self.side_channel_host,

289 remote_port=self.side_channel_port,

290 tp_size=self.vllm_config.parallel_config.tensor_parallel_size,

291 + pp_size=self.vllm_config.parallel_config.pipeline_parallel_size,

292 remote_num_tokens=remote_num_tokens,

293 )

294

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

+77 -39

@@ -80,6 +80,10 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

80 ):

81 super().__init__(vllm_config, engine_id, kv_cache_config)

83 + # Heartbeat handshakes to a PP-sharded producer must be notif-only,

84 + # like the PUSH_REG path.

85 + self._hb_handshake_notif_only = True

86 +

87 # Push-specific state.

88 # P-side: outgoing WRITE handles awaiting completion, keyed by

89 # request_id. Mutated by writer (submit) and main thread

@@ -275,18 +279,22 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

279 and the request is re-queued onto ``_reg_send_inbox`` once it

280 completes (at which point ``_ensure_handshake`` returns ``None`` and we

281 send directly)."""

282 + remote_pp_size = reg_data.get("remote_pp_size", 1)

283 fut = self._ensure_handshake(

284 reg_data["remote_engine_id"],

285 reg_data["remote_host"],

286 reg_data["remote_port"],

287 reg_data["remote_tp_size"],

288 + pp_size=remote_pp_size,

289 + # D never addresses P memory in push mode; just load P's agents.

290 + notif_agents_only=remote_pp_size > 1,

291 )

292 if fut is None:

293 self._do_send_reg_notif(req_id, reg_data)

294 return

295

296 def _on_handshake(

289 - f: Future[dict[int, str]],

297 + f: Future[dict[tuple[int, int], str]],

298 rid: str = req_id,

299 rd: dict[str, Any] = reg_data,

300 ) -> None:

@@ -500,28 +508,43 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

508 remote_block_ids = meta.remote.block_ids

509 local_block_ids = meta.local_physical_block_ids

510 num_groups = len(local_block_ids)

503 - read_specs = [

504 - ReadSpec(

505 - remote_rank=rank,

506 - local_block_ids=[

507 - list(local_block_ids[g])

508 - if rank in plan.source_ranks_per_group[g]

509 - else []

510 - for g in range(num_groups)

511 - ],

512 - remote_block_ids=[

513 - list(remote_block_ids[g])

514 - if rank in plan.source_ranks_per_group[g]

515 - else []

516 - for g in range(num_groups)

517 - ],

518 - )

519 - for rank in plan.all_source_ranks

520 - ]

511

512 if self.use_mla and tp_ratio < 0:

523 - assert len(read_specs) == 1

513 + # MLA latent is replicated across D's TP ranks: the tp-mapping

514 + # collapses to one rank (fine for reads), but push must WRITE every

515 + # D rank or the rest decode stale KV; only the dst differs per rank.

516 + assert len(plan.all_source_ranks) == 1

517 + mla_local_ids = [list(ids) for ids in local_block_ids]

518 + mla_remote_ids = [list(ids) for ids in remote_block_ids]

519 + read_specs = [

520 + ReadSpec(

521 + remote_rank=rank,

522 + local_block_ids=mla_local_ids,

523 + remote_block_ids=mla_remote_ids,

524 + )

525 + for rank in self.dst_xfer_side_handles[engine_id]

526 + ]

527 + else:

528 + read_specs = [

529 + ReadSpec(

530 + remote_rank=rank,

531 + local_block_ids=[

532 + list(local_block_ids[g])

533 + if rank in plan.source_ranks_per_group[g]

534 + else []

535 + for g in range(num_groups)

536 + ],

537 + remote_block_ids=[

538 + list(remote_block_ids[g])

539 + if rank in plan.source_ranks_per_group[g]

540 + else []

541 + for g in range(num_groups)

542 + ],

543 + )

544 + for rank in plan.all_source_ranks

545 + ]

546

547 + handles: list[int] = []

548 for i, spec in enumerate(read_specs):

549 remote_block_size = remote_info.remote_block_size

550 logger.debug(

@@ -544,7 +567,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

567 spec.remote_rank

568 ]

569

547 - self._xfer_blocks(

570 + handle = self._xfer_blocks(

571 read_spec=spec,

572 request_id=req_id,

573 dst_engine_id=meta.remote.engine_id,

@@ -552,13 +575,15 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

575 local_xfer_side_handle=local_xfer_side_handle,

576 remote_xfer_side_handle=remote_xfer_side_handle,

577 )

578 + if handle is not None:

579 + handles.append(handle)

580

556 - if self.use_mla and tp_ratio < 0 and read_specs:

557 - notif_id = f"{meta.remote.request_id}:{self.world_size}".encode()

558 - remote_agents = self._remote_agents[meta.remote.engine_id]

559 - for rank_to_notify, agent in remote_agents.items():

560 - if rank_to_notify != read_specs[0].remote_rank:

561 - self.nixl_wrapper.send_notif(agent, notif_msg=notif_id)

581 + # Publish all the request's WRITE handles in one locked update: a

582 + # partial set would let ``_pop_done_transfers`` finish the request

583 + # early, then double-report it as the remaining writes land.

584 + if handles:

585 + with self._sending_transfers_lock:

586 + self._sending_transfers[req_id].extend(handles)

587

588 def _xfer_blocks(

589 self,

@@ -568,8 +593,12 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

593 remote_request_id: str,

594 local_xfer_side_handle: int,

595 remote_xfer_side_handle: int,

571 - ):

572 - """Post a WRITE point-to-point xfer request."""

596 + ) -> int | None:

597 + """Post a WRITE point-to-point xfer request.

598 +

599 + Returns the in-flight transfer handle (so the caller can track all of

600 + a request's handles atomically), or ``None`` if nothing was submitted.

601 + """

602 assert self.transfer_topo is not None

603 remote_rank = read_spec.remote_rank

604 local_block_ids = read_spec.local_block_ids

@@ -597,7 +626,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

626

627 if len(local_block_ids) == 0:

628 logger.warning("No blocks to push for request %s", request_id)

600 - return

629 + return None

630

631 # Align per-group block counts for push.

632 local_block_ids = list(local_block_ids)

@@ -637,9 +666,9 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

666 notif_msg=notif_id,

667 )

668 self.nixl_wrapper.transfer(handle)

640 - # Track push WRITE handles so P can free blocks once done.

641 - with self._sending_transfers_lock:

642 - self._sending_transfers[request_id].append(handle)

669 + # Caller tracks the handle (atomically with the request's other

670 + # writes) so P can free blocks once all of them are done.

671 + return handle

672 except Exception as e:

673 self._log_failure(

674 failure_type="transfer_setup_failed",

@@ -656,6 +685,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

685 if handle is not None:

686 self.nixl_wrapper.release_xfer_handle(handle)

687 self.xfer_stats.record_failed_transfer()

688 + return None

689

690 # --- Notification handling on engine main thread ------------------ #

691

@@ -683,12 +713,20 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

713

714 # Not tracked as a P-side send/process for this notif.

715 if req_id not in self._reqs_to_send and req_id not in self._reqs_to_process:

686 - if req_id in self._recving_metadata:

687 - # D-side: P signalled push completion. The transfer was

688 - # driven entirely by P (we don't own a NIXL handle here),

689 - # so materialise an empty entry in ``_recving_transfers``

690 - # and let ``_pop_done_transfers`` report it done on the

691 - # next ``get_finished``.

716 + if (meta := self._recving_metadata.get(req_id)) is not None:

717 + # Consumer waits for one notif per producer rank writing

718 + # here: pp_size stages * producers-per-consumer (>1 when

719 + # producer TP > consumer TP; tp_size is the producer TP).

720 + producers_per_consumer = max(1, int(tp_size) // self.world_size)

721 + expected_notifs = meta.pp_size * producers_per_consumer

722 + self.consumer_notification_counts_by_req[req_id] += 1

723 + notifs = self.consumer_notification_counts_by_req[req_id]

724 + if notifs < expected_notifs:

725 + continue

726 + del self.consumer_notification_counts_by_req[req_id]

727 + # P drove the transfer (we own no NIXL handle), so

728 + # materialise an empty ``_recving_transfers`` entry for

729 + # ``_pop_done_transfers`` to report done.

730 self._recving_transfers.setdefault(req_id, [])

731 else:

732 # Not tracked on either side (lease may have expired