@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / bc3700e0cd

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 11, 06:15 UTC

[NIXL] Support P tensor-parallel-size > D tensor-parallel-size (#27274)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed Dec 18, 2025 at 04:53 UTC bc3700e0cd8875951b87b17accd7aa4d80ddca50

4 files changed +556 -212

tests/v1/kv_connector/nixl_integration/tp_config_sweep_accuracy_test.sh

@@ -8,9 +8,12 @@ SCRIPT="v1/kv_connector/nixl_integration/run_accuracy_test.sh"

8 configs=(

9 "GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2"

10 "GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2"

11 + "GPU_MEMORY_UTILIZATION=0.6 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1"

12 "GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA case

13 "GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny"

14 + "GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=1 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny"

15 "DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1)

16 + "DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1)

17 )

19 run_tests() {

tests/v1/kv_connector/unit/test_nixl_connector.py

+223 -22

@@ -391,6 +391,8 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

391 super().__init__(*args, **kwargs)

392 self._hand_shake_latency = hand_shake_latency

393 self.kv_cache_layout = kv_cache_layout

394 + # Mock register_kv_caches attribute needed for tests that do not call it.

395 + self.src_xfer_handles_by_block_size = {self.block_size: 1}

396

397 def _nixl_handshake(

398 self, host: str, port: int, remote_tp_size: int, expected_engine_id: str

@@ -407,22 +409,43 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

409

410 assert expected_engine_id == self.REMOTE_ENGINE_ID

411

410 - remote_agent_name = self.add_remote_agent(

411 - NixlAgentMetadata(

412 - engine_id=self.REMOTE_ENGINE_ID,

413 - agent_metadata=FakeNixlWrapper.AGENT_METADATA,

414 - kv_caches_base_addr=[0],

415 - device_id=0,

416 - num_blocks=1,

417 - block_lens=self.block_len_per_layer,

418 - # `self.kv_cache_layout` is only forced to HND when vllm engine

419 - # is started. We mock HND here.

420 - kv_cache_layout="HND",

421 - block_size=self.block_size,

422 - ),

423 - remote_tp_size=remote_tp_size,

424 - )

425 - return {0: remote_agent_name}

412 + # Adjust remote block length metadata to satisfy heterogeneous TP

413 + # invariants enforced during handshake validation.

414 + remote_block_lens = list(self.block_len_per_layer)

415 + tp_ratio = self.kv_topo.tp_ratio(remote_tp_size)

416 + if remote_tp_size > self.world_size:

417 + # P TP > D TP case, block_len of remote is smaller

418 + remote_block_lens = [

419 + block_len // (-tp_ratio) for block_len in remote_block_lens

420 + ]

421 + elif remote_tp_size < self.world_size:

422 + remote_block_lens = [

423 + block_len * tp_ratio for block_len in remote_block_lens

424 + ]

425 +

426 + # When remote tp_size > local tp_size, handshake with multiple

427 + # remote ranks.

428 + num_hanshakes = 1 if tp_ratio > 0 else -tp_ratio

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

430 + for remote_tp_rank in range(num_hanshakes):

431 + remote_agent_name = self.add_remote_agent(

432 + NixlAgentMetadata(

433 + engine_id=self.REMOTE_ENGINE_ID,

434 + agent_metadata=FakeNixlWrapper.AGENT_METADATA,

435 + kv_caches_base_addr=[0],

436 + device_id=remote_tp_rank,

437 + num_blocks=1,

438 + block_lens=remote_block_lens,

439 + # `self.kv_cache_layout` is only forced to HND when vllm engine

440 + # is started. We mock HND here.

441 + kv_cache_layout="HND",

442 + block_size=self.block_size,

443 + ),

444 + remote_tp_rank=remote_tp_rank,

445 + remote_tp_size=remote_tp_size,

446 + )

447 + remote_agents[remote_tp_rank] = remote_agent_name

448 + return remote_agents

449

450

451 class TestNixlHandshake:

@@ -453,7 +476,13 @@ class TestNixlHandshake:

476 vllm_config, connector.engine_id, hand_shake_latency=0

477 )

478 assert isinstance(connector.connector_worker.nixl_wrapper, FakeNixlWrapper)

456 - connector.connector_worker.nixl_wrapper.set_cycles_before_xfer_done(3)

479 + worker = connector.connector_worker

480 + worker.nixl_wrapper.set_cycles_before_xfer_done(3)

481 + # simulate handshake

482 + worker.dst_xfer_side_handles = {

483 + FakeNixlConnectorWorker.REMOTE_ENGINE_ID: {0: 1}

484 + }

485 + worker.kv_cache_layout = "HND"

486 num_xfers = 4

487 while True:

488 # For the same request_id, initiate multiple xfers across different

@@ -567,6 +596,171 @@ class TestNixlHandshake:

596 return

597 raise TimeoutError("Took too long to complete async handshake.")

598

599 + @patch(

600 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

601 + FakeNixlWrapper,

602 + )

603 + @pytest.mark.parametrize("local_tp_size", [1, 2])

604 + def test_prefill_tp_size_greater_than_decode_tp_size(

605 + self, local_tp_size: int, dist_init

606 + ):

607 + """

608 + Verify remote TP > local TP handshake succeeds with different

609 + remote configurations.

610 + """

611 +

612 + vllm_config = create_vllm_config()

613 + local_tp_size = 1

614 + vllm_config.parallel_config.tensor_parallel_size = local_tp_size

615 +

616 + connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)

617 + connector.connector_worker = FakeNixlConnectorWorker(

618 + vllm_config, connector.engine_id, hand_shake_latency=0

619 + )

620 + worker = connector.connector_worker

621 +

622 + # Minimal local registration params used by add_remote_agent

623 + worker.slot_size_per_layer = [4096]

624 + worker.block_len_per_layer = [4096 * worker.block_size]

625 + worker.num_blocks = 1

626 + worker.dst_num_blocks[worker.engine_id] = worker.num_blocks

627 + worker.src_blocks_data = [(0, worker.block_len_per_layer[0], worker.tp_rank)]

628 +

629 + def check_handshake(remote_tp_size: int):

630 + tp_ratio = remote_tp_size // local_tp_size

631 + assert set(remote_agents.keys()) == set(range(tp_ratio))

632 +

633 + remote_engine_id = worker.REMOTE_ENGINE_ID

634 + assert worker._tp_size[remote_engine_id] == remote_tp_size

635 + assert -tp_ratio == worker.kv_topo.tp_ratio_from_engine_id(remote_engine_id)

636 + # ensure src_xfer_handles_by_tp_ratio is populated with tpratio chunks

637 + assert -tp_ratio in worker.src_xfer_handles_by_tp_ratio

638 + assert len(worker.src_xfer_handles_by_tp_ratio[-tp_ratio]) == tp_ratio

639 + assert remote_engine_id in worker.dst_xfer_side_handles

640 + assert set(worker.dst_xfer_side_handles[remote_engine_id].keys()) == set(

641 + range(tp_ratio)

642 + )

643 +

644 + remote_agents = worker._nixl_handshake(

645 + host="localhost",

646 + port=1234,

647 + remote_tp_size=2,

648 + expected_engine_id=worker.REMOTE_ENGINE_ID,

649 + )

650 + check_handshake(2)

651 +

652 + # NOTE flexiblity: a second remote with higher number of ranks is

653 + # discovered. This is not a scenario we actively support right now, but

654 + # the connector allows it.

655 + worker.REMOTE_ENGINE_ID = "remote_engine_2"

656 + remote_agents = worker._nixl_handshake(

657 + host="localhost",

658 + port=1234,

659 + remote_tp_size=6,

660 + expected_engine_id=worker.REMOTE_ENGINE_ID,

661 + )

662 + check_handshake(6)

663 +

664 + @patch(

665 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

666 + FakeNixlWrapper,

667 + )

668 + @pytest.mark.parametrize("local_tp_size", [1, 2])

669 + def test_prefill_tp_size_greater_than_decode_tp_size_mla(

670 + self, local_tp_size: int, dist_init

671 + ):

672 + """

673 + Verify remote TP > local TP handshake succeeds with different

674 + remote configurations for an MLA model.

675 + """

676 + vllm_config = create_vllm_config()

677 + d_tp_size = 1

678 + p_tp_size = 2

679 +

680 + # Build two separate connectors/workers to emulate P TP=2 ranks.

681 + conn_p0 = NixlConnector(vllm_config, KVConnectorRole.WORKER)

682 + conn_p1 = NixlConnector(vllm_config, KVConnectorRole.WORKER)

683 + conn_p0.connector_worker = FakeNixlConnectorWorker(

684 + vllm_config, conn_p0.engine_id, hand_shake_latency=0

685 + )

686 + conn_p1.connector_worker = FakeNixlConnectorWorker(

687 + vllm_config, conn_p1.engine_id, hand_shake_latency=0

688 + )

689 +

690 + # Force P world size to 2 for both workers and emulate distinct tp_ranks.

691 + # Also enable MLA path so that expected_finished_count is updated.

692 + for rank, worker in enumerate(

693 + (conn_p0.connector_worker, conn_p1.connector_worker)

694 + ):

695 + worker.world_size = p_tp_size

696 + worker.kv_topo.remote_tp_size = {worker.engine_id: p_tp_size}

697 + worker.tp_rank = rank

698 + worker.use_mla = True

699 +

700 + req_id = "req-ep-dp2-p0"

701 + now = time.perf_counter()

702 + # Register a request on P that is waiting for consumers to read

703 + # (both workers track it).

704 + conn_p0.connector_worker._reqs_to_send[req_id] = now + 10.0

705 + conn_p0.connector_worker._reqs_to_process.add(req_id)

706 + conn_p1.connector_worker._reqs_to_send[req_id] = now + 10.0

707 + conn_p1.connector_worker._reqs_to_process.add(req_id)

708 +

709 + # Simulate a read notification coming from D with (tp=1, dp=2).

710 + notif = f"{req_id}:{d_tp_size}".encode()

711 + # D0-0->P0 notif

712 + conn_p0.connector_worker.nixl_wrapper.get_new_notifs = lambda: {

713 + "agent": [notif]

714 + } # type: ignore[method-assign]

715 + conn_p1.connector_worker.nixl_wrapper.get_new_notifs = lambda: {

716 + "agent": [notif]

717 + } # type: ignore[method-assign]

718 +

719 + # Trigger notification processing via get_finished().

720 + done_sending0, _ = conn_p0.get_finished(finished_req_ids=set())

721 + done_sending1, _ = conn_p1.get_finished(finished_req_ids=set())

722 + assert req_id in done_sending0 and req_id in done_sending1

723 +

724 + # E2E aggregation: ensure the aggregated output marks the request

725 + # as finished using the connector's expected_finished_count.

726 + from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput

727 +

728 + aggregator = KVOutputAggregator.from_connector(conn_p0, world_size=2)

729 +

730 + out0 = ModelRunnerOutput(

731 + req_ids=[req_id],

732 + req_id_to_index={req_id: 0},

733 + sampled_token_ids=[[0]],

734 + logprobs=None,

735 + prompt_logprobs_dict={},

736 + pooler_output=[None],

737 + kv_connector_output=KVConnectorOutput(

738 + finished_sending=done_sending0,

739 + finished_recving=None,

740 + ),

741 + )

742 + out1 = ModelRunnerOutput(

743 + req_ids=[req_id],

744 + req_id_to_index={req_id: 0},

745 + sampled_token_ids=[[0]],

746 + logprobs=None,

747 + prompt_logprobs_dict={},

748 + pooler_output=[None],

749 + kv_connector_output=KVConnectorOutput(

750 + finished_sending=done_sending1,

751 + finished_recving=None,

752 + ),

753 + )

754 + aggregated = aggregator.aggregate([out0, out1], output_rank=0)

755 + assert aggregated.kv_connector_output is not None

756 + assert aggregated.kv_connector_output.finished_sending == {req_id}

757 +

758 + # Producers cleaned up state for the finished request.

759 + assert req_id not in conn_p0.connector_worker._reqs_to_send

760 + assert req_id not in conn_p0.connector_worker._reqs_to_process

761 + assert req_id not in conn_p1.connector_worker._reqs_to_send

762 + assert req_id not in conn_p1.connector_worker._reqs_to_process

763 +

764 @patch(

765 "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

766 FakeNixlWrapper,

@@ -585,6 +779,9 @@ class TestNixlHandshake:

779 connector.connector_worker = FakeNixlConnectorWorker(

780 vllm_config, connector.engine_id

781 )

782 + # Register (mocked) local xfer handler

783 + # worker = connector.connector_worker

784 + # worker.src_xfer_handles_by_block_size = {worker.block_size: 1}

785 metadata = NixlConnectorMetadata()

786 total_reqs = 5

787 for i in range(total_reqs):

@@ -672,7 +869,6 @@ class TestNixlHandshake:

869 with pytest.raises(RuntimeError):

870 # mismatched layout is expected to fail

871 worker.add_remote_agent(meta, remote_tp_size=2)

675 - with pytest.raises(AssertionError):

872 worker.add_remote_agent(meta, remote_tp_size=1)

873

874 @patch(

@@ -1357,8 +1553,11 @@ def test_shutdown_cleans_up_resources(dist_init):

1553 patch.object(nixl_wrapper, "deregister_memory") as mock_dereg,

1554 ):

1555 worker._recving_transfers = {"req1": [123]}

1360 - worker.src_xfer_side_handle = 456

1361 - worker.dst_xfer_side_handles = {"engine1": 789}

1556 + # Mock register_kv_cache which registers local handle

1557 + worker.src_xfer_handles_by_block_size = {worker.block_size: 455}

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

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

1560 + worker.dst_xfer_side_handles = {"engine1": {0: 789}}

1561 worker._remote_agents = {"engine1": {0: "agent1"}}

1562 worker._registered_descs = ["desc1", "desc2"]

1563

@@ -1379,8 +1578,10 @@ def test_shutdown_cleans_up_resources(dist_init):

1578 mock_listener.join.assert_called_once()

1579

1580 mock_rel_xfer.assert_called_once_with(123)

1382 - assert mock_rel_dlist.call_count == 2

1383 - mock_rel_dlist.assert_any_call(456) # src handle

1581 + assert mock_rel_dlist.call_count == 4

1582 + mock_rel_dlist.assert_any_call(455) # src handle (whole region)

1583 + mock_rel_dlist.assert_any_call(456) # src handle (1st chunk)

1584 + mock_rel_dlist.assert_any_call(457) # src handle (2nd chunk)

1585 mock_rel_dlist.assert_any_call(789) # dst handle

1586 mock_rem_agent.assert_called_once_with("agent1")

1587 assert mock_dereg.call_count == 2

vllm/distributed/kv_transfer/kv_connector/utils.py

+40 -22

@@ -21,6 +21,8 @@ if TYPE_CHECKING:

22 logger = init_logger(__name__)

24 +EngineId = str

25 +

27 def get_kv_connector_cache_layout():

28 # NOTE (NickLucche) When running disaggregated PD with NIXL, HND layout is

@@ -209,12 +211,12 @@ class TpKVTopology:

211 """

212

213 tp_rank: int

212 - remote_tp_size: dict[str, int]

214 + remote_tp_size: dict[EngineId, int]

215 is_mla: bool

216 total_num_kv_heads: int

217 attn_backend: type[AttentionBackend]

216 - engine_id: str

217 - remote_block_size: dict[str, int]

218 + engine_id: EngineId

219 + remote_block_size: dict[EngineId, int]

220

221 def __post_init__(self):

222 # Figure out whether the first dimension of the cache is K/V

@@ -256,18 +258,28 @@ class TpKVTopology:

258 Calculate the tensor parallel ratio between local and remote TP.

259 We can think of it as the number of local TP workers-per-remote TP

260 workers. Local workers will read from the same remote TP worker in

259 - groups of size `tp_ratio`.

261 + groups of size `tp_ratio`.If remote tp_size > local tp_size, the

262 + ratio is flipped (remote_size/local_size) and the returned value is

263 + negative.

264 """

261 - assert self.tp_size % remote_tp_size == 0, (

262 - f"Local tensor parallel size {self.tp_size} is not divisible "

263 - f"by remote tensor parallel size {remote_tp_size}."

265 + if self.tp_size >= remote_tp_size:

266 + assert self.tp_size % remote_tp_size == 0, (

267 + f"Local tensor parallel size {self.tp_size} is not divisible "

268 + f"by remote tensor parallel size {remote_tp_size}."

269 + )

270 + return self.tp_size // remote_tp_size

271 +

272 + assert remote_tp_size % self.tp_size == 0, (

273 + f"Remote tensor parallel size {remote_tp_size} is not divisible "

274 + f"by local tensor parallel size {self.tp_size}."

275 )

265 - return self.tp_size // remote_tp_size

276 + # P TP > D TP case, return the ratio as negative

277 + return -remote_tp_size // self.tp_size

278

279 def block_size_ratio(

280 self,

281 remote_block_size: int,

270 - ) -> float:

282 + ) -> int:

283 """

284 Calculate the block size ratio between local and remote TP.

285 """

@@ -279,19 +291,19 @@ class TpKVTopology:

291

292 def tp_ratio_from_engine_id(

293 self,

282 - remote_engine_id: str,

294 + remote_engine_id: EngineId,

295 ) -> int:

296 remote_tp_size = self.remote_tp_size[remote_engine_id]

297 return self.tp_ratio(remote_tp_size)

298

299 def block_size_ratio_from_engine_id(

300 self,

289 - remote_engine_id: str,

290 - ) -> float:

301 + remote_engine_id: EngineId,

302 + ) -> int:

303 remote_block_size = self.remote_block_size[remote_engine_id]

304 return self.block_size_ratio(remote_block_size)

305

294 - def is_kv_replicated(self, engine_id: str) -> bool:

306 + def is_kv_replicated(self, engine_id: EngineId) -> bool:

307 """

308 Whether the KV cache is replicated across TP workers due to the

309 number of TP workers being greater than the number of KV heads.

@@ -299,24 +311,30 @@ class TpKVTopology:

311 tp_size = self.remote_tp_size[engine_id]

312 return tp_size // self.total_num_kv_heads >= 1

313

302 - def replicates_kv_cache(self, remote_engine_id: str) -> bool:

314 + def replicates_kv_cache(self, remote_engine_id: EngineId) -> bool:

315 # MLA is always replicated as the hidden dim can't be split.

316 return self.is_mla or self.is_kv_replicated(remote_engine_id)

317

306 - def get_target_remote_rank(

318 + def get_target_remote_ranks(

319 self,

320 remote_tp_size: int,

309 - ) -> int:

321 + ) -> list[int]:

322 """

323 Get the remote TP rank (on P) that the current local TP rank

312 - (on D) will read from.

324 + (on D) will read from. When remote tp_size > local tp_size, we

325 + read from multiple remote ranks.

326 """

327 tp_ratio = self.tp_ratio(remote_tp_size)

315 - return self.tp_rank // tp_ratio

328 + if tp_ratio > 0:

329 + return [self.tp_rank // tp_ratio]

330

317 - def get_target_remote_rank_from_engine_id(

331 + # P TP > D TP case, D reads from |tp_ratio| remote workers.

332 + tp_ratio = -tp_ratio

333 + return [self.tp_rank * tp_ratio + i for i in range(tp_ratio)]

334 +

335 + def get_target_remote_ranks_from_engine_id(

336 self,

319 - remote_engine_id: str,

320 - ) -> int:

337 + remote_engine_id: EngineId,

338 + ) -> list[int]:

339 remote_tp_size = self.remote_tp_size[remote_engine_id]

322 - return self.get_target_remote_rank(remote_tp_size)

340 + return self.get_target_remote_ranks(remote_tp_size)

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+290 -168

@@ -23,7 +23,7 @@ from vllm import envs

23 from vllm.attention.backends.abstract import AttentionMetadata

24 from vllm.attention.selector import get_attn_backend

25 from vllm.config import VllmConfig

26 -from vllm.distributed.kv_transfer.kv_connector.utils import TpKVTopology

26 +from vllm.distributed.kv_transfer.kv_connector.utils import EngineId, TpKVTopology

27 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

28 CopyBlocksOp,

29 KVConnectorBase_V1,

@@ -56,7 +56,6 @@ if TYPE_CHECKING:

56 from vllm.v1.request import Request

58 TransferHandle = int

59 -EngineId = str

59 ReqId = str

61 #

@@ -873,9 +872,10 @@ class NixlConnectorWorker:

872 self.copy_blocks: CopyBlocksOp | None = None

873

874 # Map of engine_id -> kv_caches_base_addr. For TP case, each local

876 - # rank will still only pull from a single remote TP worker.

877 - self.kv_caches_base_addr: dict[EngineId, list[int]] = {}

875 self.device_id: int = 0

876 + # Current rank may pull from multiple remote TP workers.

877 + # EngineId, dict[int, list[int]] -> engine_id, tp_rank, base_addr_for_layer

878 + self.kv_caches_base_addr = defaultdict[EngineId, dict[int, list[int]]](dict)

879

880 # Number of NIXL regions. Currently one region per cache

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

@@ -883,10 +883,12 @@ class NixlConnectorWorker:

883 self.num_layers = 0

884

885 # nixl_prepped_dlist_handle.

886 - self.src_xfer_side_handle: int = 0

887 - self.src_xfer_side_handles: dict[int, int] = {}

888 - # Map of engine_id -> nixl_prepped_dlist_handle (int)].

889 - self.dst_xfer_side_handles: dict[EngineId, int] = {}

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

887 + # Populated dynamically during handshake based on remote configuration.

888 + # Keep track of regions at different tp_ratio values. tp_ratio->handles

889 + self.src_xfer_handles_by_tp_ratio: dict[int, list[int]] = {}

890 + # Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}.

891 + self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict)

892

893 # Map of engine_id -> num_blocks. All ranks in the same deployment will

894 # have the same number of blocks.

@@ -977,103 +979,108 @@ class NixlConnectorWorker:

979 expected_engine_id: str,

980 ) -> dict[int, str]:

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

980 -

981 - start_time = time.perf_counter()

982 -

983 - # NOTE(rob): we need each rank to have a unique port. This is

984 - # a hack to keep us moving. We will switch when moving to etcd

985 - # or where we have a single ZMQ socket in the scheduler.

986 -

987 - # Handshake only with the remote TP rank that current local rank will

988 - # pull from. With homogeneous TP it happens to be the same rank_i.

989 - p_remote_rank = self.kv_topo.get_target_remote_rank(remote_tp_size)

982 + # When target instance TP > local TP, we need to perform multiple

983 + # handshakes. Do it in a single background job for simplicity.

984 + # Regardless, only handshake with the remote TP rank(s) that current

985 + # local rank will read from. Note that With homogeneous TP,

986 + # this happens to be the same single rank_i.

987 + p_remote_ranks = self.kv_topo.get_target_remote_ranks(remote_tp_size)

988 + remote_rank_to_agent_name = {}

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

991 - logger.debug(

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

993 - )

990

995 - # Send query for the request.

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

997 - msg = msgspec.msgpack.encode((GET_META_MSG, p_remote_rank))

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

999 - sock.setsockopt(zmq.RCVTIMEO, 5000) # milliseconds

1000 - sock.send(msg)

1001 - handshake_bytes = sock.recv()

1002 -

1003 - # Decode handshake payload to get compatibility hash

1004 - handshake_decoder = msgspec.msgpack.Decoder(NixlHandshakePayload)

1005 - try:

1006 - handshake_payload = handshake_decoder.decode(handshake_bytes)

1007 - except (msgspec.DecodeError, msgspec.ValidationError) as e:

1008 - raise RuntimeError(

1009 - f"Failed to decode NixlHandshakePayload. This likely indicates "

1010 - f"an incompatibility between connector version. Error: {e}"

1011 - ) from e

1012 -

1013 - got_metadata_time = time.perf_counter()

1014 - logger.debug(

1015 - "NIXL handshake: get metadata took: %s", got_metadata_time - start_time

1016 - )

1017 -

1018 - # Check compatibility hash BEFORE decoding agent metadata

1019 - if (

1020 - self.enforce_compat_hash

1021 - and handshake_payload.compatibility_hash != self.compat_hash

1022 - ):

1023 - raise RuntimeError(

1024 - f"NIXL compatibility hash mismatch. "

1025 - f"Local: {self.compat_hash}, "

1026 - f"Remote: {handshake_payload.compatibility_hash}. "

1027 - f"Prefill and decode instances have incompatible configurations. "

1028 - f"This may be due to: different vLLM versions, models, dtypes, "

1029 - f"KV cache layouts, attention backends, etc. "

1030 - f"Both instances must use identical configurations."

1031 - f"Disable this check using "

1032 - f'--kv-transfer-config \'{{"kv_connector_extra_config": '

1033 - f'{{"enforce_handshake_compat": false}}}}\''

992 + for remote_rank in p_remote_ranks:

993 + logger.debug(

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

995 + path,

996 + remote_rank,

997 )

998

1036 - logger.info(

1037 - "NIXL compatibility check passed (hash: %s)",

1038 - handshake_payload.compatibility_hash,

1039 - )

999 + start_time = time.perf_counter()

1000 + # Send query for the request.

1001 + msg = msgspec.msgpack.encode((GET_META_MSG, remote_rank))

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

1003 + sock.setsockopt(zmq.RCVTIMEO, 5000) # milliseconds

1004 + sock.send(msg)

1005 + handshake_bytes = sock.recv()

1006

1041 - # Decode agent metadata

1042 - metadata_decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)

1043 - try:

1044 - metadata = metadata_decoder.decode(

1045 - handshake_payload.agent_metadata_bytes

1007 + # Decode handshake payload to get compatibility hash

1008 + handshake_decoder = msgspec.msgpack.Decoder(NixlHandshakePayload)

1009 + try:

1010 + handshake_payload = handshake_decoder.decode(handshake_bytes)

1011 + except (msgspec.DecodeError, msgspec.ValidationError) as e:

1012 + raise RuntimeError(

1013 + f"Failed to decode NixlHandshakePayload. This likely indicates "

1014 + f"an incompatibility between connector version. Error: {e}"

1015 + ) from e

1016 +

1017 + got_metadata_time = time.perf_counter()

1018 + logger.debug(

1019 + "NIXL handshake: get metadata took: %s",

1020 + got_metadata_time - start_time,

1021 )

1047 - except (msgspec.DecodeError, msgspec.ValidationError) as e:

1048 - # This should not happen if hash matched

1049 - raise RuntimeError(

1050 - f"Failed to decode NixlAgentMetadata. Error: {e}"

1051 - ) from e

1022

1053 - # Ensure engine id matches.

1054 - if metadata.engine_id != expected_engine_id:

1055 - raise RuntimeError(

1056 - f"Remote NIXL agent engine ID mismatch. "

1057 - f"Expected {expected_engine_id},"

1058 - f"received {metadata.engine_id}."

1059 - )

1023 + # Check compatibility hash BEFORE decoding agent metadata

1024 + if (

1025 + self.enforce_compat_hash

1026 + and handshake_payload.compatibility_hash != self.compat_hash

1027 + ):

1028 + raise RuntimeError(

1029 + f"NIXL compatibility hash mismatch. "

1030 + f"Local: {self.compat_hash}, "

1031 + f"Remote: {handshake_payload.compatibility_hash}. "

1032 + f"Prefill and decode instances have incompatible "

1033 + f"configurations. This may be due to: different vLLM versions,"

1034 + f" models, dtypes, KV cache layouts, attention backends, etc. "

1035 + f"Both instances must use identical configurations."

1036 + f"Disable this check using "

1037 + f'--kv-transfer-config \'{{"kv_connector_extra_config": '

1038 + f'{{"enforce_handshake_compat": false}}}}\''

1039 + )

1040

1061 - # Register Remote agent.

1062 - assert metadata.block_size <= self.block_size, (

1063 - "nP > nD is not supported yet."

1064 - )

1065 - remote_agent_name = self.add_remote_agent(

1066 - metadata, p_remote_rank, remote_tp_size

1067 - )

1041 + logger.info(

1042 + "NIXL compatibility check passed (hash: %s)",

1043 + handshake_payload.compatibility_hash,

1044 + )

1045

1069 - setup_agent_time = time.perf_counter()

1070 - logger.debug(

1071 - "NIXL handshake: add agent took: %s",

1072 - setup_agent_time - got_metadata_time,

1073 - )

1046 + # Decode agent metadata

1047 + metadata_decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)

1048 + try:

1049 + metadata = metadata_decoder.decode(

1050 + handshake_payload.agent_metadata_bytes

1051 + )

1052 + except (msgspec.DecodeError, msgspec.ValidationError) as e:

1053 + # This should not happen if hash matched

1054 + raise RuntimeError(

1055 + f"Failed to decode NixlAgentMetadata. Error: {e}"

1056 + ) from e

1057 +

1058 + # Ensure engine id matches.

1059 + if metadata.engine_id != expected_engine_id:

1060 + raise RuntimeError(

1061 + f"Remote NIXL agent engine ID mismatch. "

1062 + f"Expected {expected_engine_id},"

1063 + f"received {metadata.engine_id}."

1064 + )

1065 + # Ensure engine id matches.

1066 + if metadata.engine_id != expected_engine_id:

1067 + raise RuntimeError(

1068 + f"Remote NIXL agent engine ID mismatch. "

1069 + f"Expected {expected_engine_id},"

1070 + f"received {metadata.engine_id}."

1071 + )

1072 + setup_agent_time = time.perf_counter()

1073

1075 - # Remote rank -> agent name.

1076 - return {p_remote_rank: remote_agent_name}

1074 + # Register Remote agent.

1075 + remote_agent_name = self.add_remote_agent(

1076 + metadata, remote_rank, remote_tp_size

1077 + )

1078 + logger.debug(

1079 + "NIXL handshake: add agent took: %s",

1080 + setup_agent_time - got_metadata_time,

1081 + )

1082 + remote_rank_to_agent_name[remote_rank] = remote_agent_name

1083 + return remote_rank_to_agent_name

1084

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

1086 """

@@ -1283,7 +1290,7 @@ class NixlConnectorWorker:

1290 assert len(self.block_len_per_layer) == len(seen_base_addresses)

1291 assert self.num_blocks != 0

1292

1286 - self.kv_caches_base_addr[self.engine_id] = seen_base_addresses

1293 + self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses

1294 self.num_regions = len(caches_data)

1295 self.num_layers = len(xfer_buffers.keys())

1296

@@ -1310,9 +1317,9 @@ class NixlConnectorWorker:

1317

1318 # Register local/src descr for NIXL xfer.

1319 self.seen_base_addresses = seen_base_addresses

1313 - self.src_xfer_side_handle = self.register_local_xfer_handler(self.block_size)

1314 -

1315 - self.src_xfer_side_handles[self.block_size] = self.src_xfer_side_handle

1320 + self.src_xfer_handles_by_block_size[self.block_size], self.src_blocks_data = (

1321 + self.register_local_xfer_handler(self.block_size)

1322 + )

1323

1324 # TODO(mgoin): Hybrid memory allocator is currently disabled for

1325 # models with local attention (Llama 4). Can remove this once enabled.

@@ -1340,8 +1347,8 @@ class NixlConnectorWorker:

1347 agent_metadata = NixlAgentMetadata(

1348 engine_id=self.engine_id,

1349 agent_metadata=self.nixl_wrapper.get_agent_metadata(),

1343 - kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id],

1350 device_id=self.device_id,

1351 + kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id][self.tp_rank],

1352 num_blocks=self.num_blocks,

1353 block_lens=self.block_len_per_layer,

1354 kv_cache_layout=self.kv_cache_layout

@@ -1359,7 +1366,7 @@ class NixlConnectorWorker:

1366 def register_local_xfer_handler(

1367 self,

1368 block_size: int,

1362 - ) -> int:

1369 + ) -> tuple[int, list[tuple[int, int, int]]]:

1370 """

1371 Function used for register local xfer handler with local block_size or

1372 Remote block_size.

@@ -1407,7 +1414,7 @@ class NixlConnectorWorker:

1414

1415 descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)

1416 # NIXL_INIT_AGENT to be used for preparations of local descs.

1410 - return self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)

1417 + return self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs), blocks_data

1418

1419 def add_remote_agent(

1420 self,

@@ -1421,10 +1428,12 @@ class NixlConnectorWorker:

1428

1429 In particular, handle both homogeneous and heterogeneous TP. The former

1430 requires local rank_i to read from remote rank_i.

1424 - The latter, assuming D.world_size > P.world_size, requires that two or

1425 - more local TP worker share the xfer from a single TP worker.

1431 + The latter, in the case of D.world_size < P.world_size, requires that a

1432 + local (D) TP worker reads from multiple remote (P) TP workers.

1433 + Conversely, assuming D.world_size > P.world_size, two or more local TP

1434 + workers will read from a single remote TP worker.

1435

1427 - Here's an example (non-MLA case):

1436 + Here's an example for the last case described above (non-MLA):

1437

1438 rank_offset p_remote_tp_rank

1439 (kv split no)

@@ -1474,9 +1483,6 @@ class NixlConnectorWorker:

1483 nixl_agent_meta.agent_metadata

1484 )

1485

1477 - # Handle tp_size>num_kv_heads: replicate KV cache.

1478 - replicates_kv_cache = self.kv_topo.replicates_kv_cache(engine_id)

1479 -

1486 # Create dst descs and xfer side handles. TP workers have same #blocks

1487 # so we only register once per engine_id.

1488 # Example:

@@ -1490,14 +1496,52 @@ class NixlConnectorWorker:

1496 self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks

1497

1498 # Keep track of remote agent kv caches base addresses.

1493 - self.kv_caches_base_addr[engine_id] = nixl_agent_meta.kv_caches_base_addr

1494 -

1499 + self.kv_caches_base_addr[engine_id][remote_tp_rank] = (

1500 + nixl_agent_meta.kv_caches_base_addr

1501 + )

1502 self._validate_remote_agent_handshake(nixl_agent_meta, remote_tp_size)

1503

1497 - # Number of D TP workers reading from a single P TP worker. This is

1498 - # 1 when P and D `--tensor-parallel-size` match.

1504 + # This is 1 when P and D `--tensor-parallel-size` match. Otherwise,

1505 + # this is the ratio between the two sizes.

1506 tp_ratio = self.kv_topo.tp_ratio_from_engine_id(engine_id)

1507

1508 + # Handle tp_size>num_kv_heads: replicate KV cache.

1509 + indexes_into_remote = (

1510 + not self.kv_topo.replicates_kv_cache(engine_id) and tp_ratio > 0

1511 + )

1512 +

1513 + logger.debug(

1514 + "Registering remote agent (%s, rank %s) memory regions with tp_ratio %s",

1515 + engine_id,

1516 + remote_tp_rank,

1517 + tp_ratio,

1518 + )

1519 +

1520 + ### (Optional) Register local agent memory regions. MLA is not split.

1521 + if (

1522 + tp_ratio < 0

1523 + and not self.use_mla

1524 + and tp_ratio not in self.src_xfer_handles_by_tp_ratio

1525 + ):

1526 + # Remote tp_size > local tp_size: read from multiple remote ranks.

1527 + # Logically "split" own regions into |tp_ratio| chunks. Mind that

1528 + # we only do this once per remote tp_size (replica-friendly).

1529 + self.src_xfer_handles_by_tp_ratio[tp_ratio] = []

1530 + for i in range(-tp_ratio):

1531 + blocks_data = []

1532 + for memory_region in self.src_blocks_data:

1533 + addr, local_block_len, own_tp_rank = memory_region

1534 + # Computing block len layer by layer allows for different

1535 + # block sizes to be used.

1536 + remote_block_len = local_block_len // (-tp_ratio)

1537 + addr = addr + i * remote_block_len

1538 + blocks_data.append((addr, remote_block_len, own_tp_rank))

1539 + descs = self.nixl_wrapper.get_xfer_descs(

1540 + blocks_data, self.nixl_memory_type

1541 + )

1542 + handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)

1543 + self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)

1544 +

1545 ### Register remote agent memory regions

1546 blocks_data = []

1547 # With homogeneous TP, D pulls the whole kv cache from corresponding

@@ -1507,14 +1551,19 @@ class NixlConnectorWorker:

1551

1552 # Register all remote blocks, but only the corresponding kv heads.

1553 for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):

1510 - kv_block_len = self.get_backend_aware_kv_block_len(layer_idx=i)

1511 - remote_kv_block_len = kv_block_len // block_size_ratio

1554 + # Read our whole local region size from remote.

1555 + local_block_len = self.get_backend_aware_kv_block_len(layer_idx=i)

1556 + remote_kv_block_len = local_block_len // block_size_ratio

1557 if block_size_ratio > 1:

1558 # using remote kv_block_len as transfer unit

1514 - kv_block_len = remote_kv_block_len

1559 + local_block_len = remote_kv_block_len

1560 +

1561 + if tp_ratio < 0 and not self.use_mla:

1562 + # Remote tp is bigger: read a chunk of local region from remote

1563 + local_block_len = local_block_len // (-tp_ratio)

1564 rank_offset = (

1565 self.tp_rank % tp_ratio * remote_kv_block_len

1517 - if not replicates_kv_cache

1566 + if indexes_into_remote

1567 else 0

1568 )

1569 for block_id in range(nixl_agent_meta.num_blocks):

@@ -1524,7 +1573,7 @@ class NixlConnectorWorker:

1573 # self.block_len == remote_block_len//tp_ratio bytes.

1574 addr = base_addr + block_offset + rank_offset

1575 # (addr, len, device id)

1527 - blocks_data.append((addr, kv_block_len, nixl_agent_meta.device_id))

1576 + blocks_data.append((addr, local_block_len, nixl_agent_meta.device_id))

1577

1578 if self.kv_topo.is_kv_layout_blocks_first:

1579 # With FlashInfer index V separately to allow head splitting.

@@ -1533,7 +1582,7 @@ class NixlConnectorWorker:

1582 addr = base_addr + block_offset + rank_offset

1583 v_addr = addr + nixl_agent_meta.block_lens[i] // 2

1584 blocks_data.append(

1536 - (v_addr, kv_block_len, nixl_agent_meta.device_id)

1585 + (v_addr, local_block_len, nixl_agent_meta.device_id)

1586 )

1587

1588 logger.debug(

@@ -1546,15 +1595,15 @@ class NixlConnectorWorker:

1595

1596 # Register with NIXL.

1597 descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)

1549 - self.dst_xfer_side_handles[engine_id] = self.nixl_wrapper.prep_xfer_dlist(

1550 - remote_agent_name, descs

1598 + self.dst_xfer_side_handles[engine_id][remote_tp_rank] = (

1599 + self.nixl_wrapper.prep_xfer_dlist(remote_agent_name, descs)

1600 )

1601

1602 if block_size_ratio > 1:

1603 # when prefill with smaller block_size, we need to init a

1604 # new handler with same block_len to match

1556 - self.src_xfer_side_handles[nixl_agent_meta.block_size] = (

1557 - self.register_local_xfer_handler(nixl_agent_meta.block_size)

1605 + self.src_xfer_handles_by_block_size[nixl_agent_meta.block_size] = (

1606 + self.register_local_xfer_handler(nixl_agent_meta.block_size)[0]

1607 )

1608

1609 return remote_agent_name

@@ -1574,7 +1623,9 @@ class NixlConnectorWorker:

1623 block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(

1624 remote_engine_id

1625 )

1577 - assert tp_ratio > 0, "Decode TP cannot be smaller than prefill TP"

1626 + # Num kv_heads > tp_size and P TP > D TP case, not supported

1627 + assert not (tp_ratio < 0 and self.kv_topo.is_kv_replicated(remote_engine_id))

1628 +

1629 assert not self._use_pallas or tp_ratio == 1, (

1630 "TPU (pallas_v1) DOES NOT support heterogeneous TP yet."

1631 )

@@ -1616,17 +1667,29 @@ class NixlConnectorWorker:

1667 "All remote layers must have the same block size"

1668 )

1669

1619 - assert (

1620 - remote_block_len

1621 - == (self.block_len_per_layer[0] * tp_ratio) // block_size_ratio

1622 - ), (

1623 - "Remote P worker KV layer cache must be of shape [2, N, "

1624 - "local_kv_heads*tp_ratio, block_size, head_dim] and same dtype."

1625 - )

1670 + if tp_ratio > 0:

1671 + # Remote tp is smaller: remote block_len size is bigger

1672 + assert (

1673 + remote_block_len

1674 + == (self.block_len_per_layer[0] * tp_ratio) // block_size_ratio

1675 + ), (

1676 + "Remote P worker KV layer cache must be of shape [2, N, "

1677 + "local_kv_heads*tp_ratio, page_size, head_dim] and same dtype."

1678 + ) # noqa: E501

1679 + else:

1680 + assert block_size_ratio == 1, (

1681 + "Different local/remote block sizes are not supported when"

1682 + " P TP > D TP."

1683 + )

1684 + # Remote tp is bigger: remote block_len size is smaller

1685 + assert remote_block_len == self.block_len_per_layer[0] // (-tp_ratio), (

1686 + "Remote P worker KV layer cache must be of shape [2, N, "

1687 + "local_kv_heads/tp_ratio, page_size, head_dim] and same dtype."

1688 + ) # noqa: E501

1689

1627 - # TP workers have same #blocks.

1690 + # TP workers that handhshake with same remote have same #blocks.

1691 assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks

1629 -

1692 + # Same number of regions/~layers.

1693 assert len(nixl_agent_meta.kv_caches_base_addr) == len(self.block_len_per_layer)

1694

1695 def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta):

@@ -1710,7 +1773,7 @@ class NixlConnectorWorker:

1773 )

1774 cache.index_copy_(0, indices, permuted_blocks)

1775

1713 - def blocksize_post_process(self, block_ids_per_ratio: dict[float, list[list[int]]]):

1776 + def blocksize_post_process(self, block_ids_per_ratio: dict[int, list[list[int]]]):

1777 def _process_local_gt_remote(blocks_to_update, block_size_ratio):

1778 n_kv_heads, block_size, head_size = blocks_to_update.shape[1:]

1779 remote_block_size = block_size // block_size_ratio

@@ -1840,7 +1903,7 @@ class NixlConnectorWorker:

1903 notified_req_ids: set[str] = set()

1904 for notifs in self.nixl_wrapper.get_new_notifs().values():

1905 for notif in notifs:

1843 - req_id, tp_ratio = notif.decode("utf-8").rsplit(":", 1)

1906 + req_id, tp_size = notif.decode("utf-8").rsplit(":", 1)

1907 if (

1908 req_id not in self._reqs_to_send

1909 and req_id not in self._reqs_to_process

@@ -1853,9 +1916,22 @@ class NixlConnectorWorker:

1916 )

1917 continue

1918

1919 + # NOTE: `tp_ratio` is the opposite when swapping local<>remote

1920 + n_consumers = int(tp_size)

1921 + tp_ratio = self.kv_topo.tp_ratio(n_consumers)

1922 +

1923 + # Number of reads *per producer* to wait for.

1924 + # When remote D TP > local P TP we expect `tp_ratio` reads.

1925 + consumers_per_producer = (

1926 + -tp_ratio if n_consumers > self.world_size else 1

1927 + )

1928 +

1929 self.consumer_notification_counts_by_req[req_id] += 1

1930 # Wait all consumers (D) to be done reading before freeing.

1858 - if self.consumer_notification_counts_by_req[req_id] == int(tp_ratio):

1931 + if (

1932 + self.consumer_notification_counts_by_req[req_id]

1933 + == consumers_per_producer

1934 + ):

1935 notified_req_ids.add(req_id)

1936 del self.consumer_notification_counts_by_req[req_id]

1937 self._reqs_to_process.remove(req_id)

@@ -1872,7 +1948,7 @@ class NixlConnectorWorker:

1948 """

1949 done_req_ids: set[str] = set()

1950 for req_id, handles in list(transfers.items()):

1875 - in_progress = False

1951 + in_progress = []

1952 for handle in handles:

1953 try:

1954 xfer_state = self.nixl_wrapper.check_xfer_state(handle)

@@ -1882,7 +1958,7 @@ class NixlConnectorWorker:

1958 self.xfer_stats.record_transfer(res)

1959 self.nixl_wrapper.release_xfer_handle(handle)

1960 elif xfer_state == "PROC":

1885 - in_progress = True

1961 + in_progress.append(handle)

1962 continue

1963 else:

1964 logger.error(

@@ -1892,7 +1968,6 @@ class NixlConnectorWorker:

1968 xfer_state,

1969 )

1970 self._handle_failed_transfer(req_id, handle)

1895 - in_progress = False

1971 except Exception:

1972 logger.exception(

1973 "NIXL transfer exception for request %s. "

@@ -1900,11 +1975,13 @@ class NixlConnectorWorker:

1975 req_id,

1976 )

1977 self._handle_failed_transfer(req_id, handle)

1903 - in_progress = False

1978

1979 if not in_progress:

1980 + # Only report request as completed when all transfers are done.

1981 done_req_ids.add(req_id)

1982 del transfers[req_id]

1983 + else:

1984 + transfers[req_id] = in_progress

1985 return done_req_ids

1986

1987 def _handle_failed_transfer(self, req_id: str, handle: int):

@@ -1982,18 +2059,62 @@ class NixlConnectorWorker:

2059

2060 def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):

2061 assert meta.remote is not None

1985 - logger.debug(

1986 - "Remote agent %s available, calling _read_blocks for req %s",

1987 - meta.remote.engine_id,

1988 - req_id,

1989 - )

1990 - self._read_blocks(

1991 - request_id=req_id,

1992 - dst_engine_id=meta.remote.engine_id,

1993 - remote_request_id=meta.remote.request_id,

1994 - local_block_ids=meta.local_physical_block_ids,

1995 - remote_block_ids=meta.remote.block_ids,

2062 + remote_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id(

2063 + meta.remote.engine_id

2064 )

2065 + tp_ratio = self.kv_topo.tp_ratio_from_engine_id(meta.remote.engine_id)

2066 + # D may have to perform multiple reads from different remote ranks.

2067 + for i, remote_rank in enumerate(remote_ranks):

2068 + if self.use_mla and tp_ratio < 0 and i > 0:

2069 + # MLA opt: when P TP > D TP, only a single read is executed for

2070 + # the first remote rank (cache is duplicated)..

2071 + break

2072 +

2073 + remote_block_size = self.kv_topo.remote_block_size[meta.remote.engine_id]

2074 + logger.debug(

2075 + "Remote agent %s available, calling _read_blocks"

2076 + " on remote rank %s with remote block size %s for req %s",

2077 + meta.remote.engine_id,

2078 + remote_rank,

2079 + remote_block_size,

2080 + req_id,

2081 + )

2082 + # Get side handles.

2083 + if tp_ratio < 0 and not self.use_mla:

2084 + assert remote_block_size == self.block_size

2085 + # Remote tp_size > local tp_size: we must perform multiple

2086 + # reads. Get the memory chunk onto which we will write to.

2087 + local_xfer_side_handle = self.src_xfer_handles_by_tp_ratio[tp_ratio][i]

2088 + else:

2089 + # Single read from remote, we write to the whole memory region.

2090 + # Also handle remote block size different from local block size.

2091 + local_xfer_side_handle = self.src_xfer_handles_by_block_size[

2092 + remote_block_size

2093 + ]

2094 +

2095 + # Destination handle: remote_engine_id -> remote_rank -> handle.

2096 + remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][

2097 + remote_rank

2098 + ]

2099 + self._read_blocks(

2100 + request_id=req_id,

2101 + dst_engine_id=meta.remote.engine_id,

2102 + remote_request_id=meta.remote.request_id,

2103 + local_block_ids=meta.local_physical_block_ids,

2104 + remote_block_ids=meta.remote.block_ids,

2105 + remote_rank=remote_rank,

2106 + local_xfer_side_handle=local_xfer_side_handle,

2107 + remote_xfer_side_handle=remote_xfer_side_handle,

2108 + )

2109 +

2110 + if self.use_mla and tp_ratio < 0:

2111 + # ..but we still need to notify the other remote ranks that we

2112 + # have the blocks we need so they can update the request state.

2113 + notif_id = f"{req_id}:{self.world_size}".encode()

2114 + remote_agents = self._remote_agents[meta.remote.engine_id]

2115 + for rank_to_notify, agent in remote_agents.items():

2116 + if rank_to_notify != remote_rank:

2117 + self.nixl_wrapper.send_notif(agent, notif_msg=notif_id)

2118

2119 def _read_blocks(

2120 self,

@@ -2002,7 +2123,14 @@ class NixlConnectorWorker:

2123 dst_engine_id: str,

2124 request_id: str,

2125 remote_request_id: str,

2126 + remote_rank: int,

2127 + local_xfer_side_handle: int,

2128 + remote_xfer_side_handle: int,

2129 ):

2130 + """

2131 + Post a READ point-to-point xfer request from a single local worker to

2132 + a single remote worker.

2133 + """

2134 block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(dst_engine_id)

2135 if block_size_ratio > 1:

2136 local_block_ids = self.get_mapped_blocks(

@@ -2031,18 +2159,14 @@ class NixlConnectorWorker:

2159 # saturate IB with heterogeneous TP sizes. We should remove the staging

2160 # blocks until we are ready.

2161

2034 - # Number of D TP workers that will read from dst P. Propagate tp_ratio

2162 + # Number of D TP workers that will read from dst P. Propagate info

2163 # on notification so that dst worker can wait before freeing blocks.

2036 - tp_ratio = self.kv_topo.tp_ratio_from_engine_id(dst_engine_id)

2037 - notif_id = f"{remote_request_id}:{tp_ratio}".encode()

2164 + notif_id = f"{remote_request_id}:{self.world_size}".encode()

2165

2166 # Full prefix cache hit: do not need to read remote blocks,

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

2168 num_local_blocks = len(local_block_ids)

2169 if num_local_blocks == 0:

2043 - remote_rank = self.kv_topo.get_target_remote_rank_from_engine_id(

2044 - dst_engine_id

2045 - )

2170 agent_name = self._remote_agents[dst_engine_id][remote_rank]

2171 try:

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

@@ -2062,13 +2186,6 @@ class NixlConnectorWorker:

2186 if num_local_blocks < num_remote_blocks:

2187 remote_block_ids = remote_block_ids[-num_local_blocks:]

2188

2065 - # Get side handles.

2066 - remote_block_size = self.kv_topo.remote_block_size[dst_engine_id]

2067 - local_xfer_side_handle = self.src_xfer_side_handles.get(

2068 - remote_block_size, self.src_xfer_side_handle

2069 - )

2070 - remote_xfer_side_handle = self.dst_xfer_side_handles[dst_engine_id]

2071 -

2189 # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from

2190 # corresponding rank. With heterogeneous TP, fixing D>P, the D tp

2191 # workers will issue xfers to parts of the P worker remote kv caches.

@@ -2230,7 +2347,7 @@ class NixlConnectorWorker:

2347 block_ids_np, self._physical_blocks_per_logical_kv_block, block_arange

2348 ).tolist()

2349

2233 - def get_backend_aware_kv_block_len(self, layer_idx: int):

2350 + def get_backend_aware_kv_block_len(self, layer_idx: int) -> int:

2351 """

2352 Get the block length for one K/V element (K and V have the same size).

2353

@@ -2276,11 +2393,16 @@ class NixlConnectorWorker:

2393 for handle in handles:

2394 self.nixl_wrapper.release_xfer_handle(handle)

2395 self._recving_transfers.clear()

2279 - if self.src_xfer_side_handle:

2280 - self.nixl_wrapper.release_dlist_handle(self.src_xfer_side_handle)

2281 - self.src_xfer_side_handle = 0

2282 - for dst_xfer_side_handle in self.dst_xfer_side_handles.values():

2283 - self.nixl_wrapper.release_dlist_handle(dst_xfer_side_handle)

2396 + for handle in self.src_xfer_handles_by_block_size.values():

2397 + self.nixl_wrapper.release_dlist_handle(handle)

2398 + self.src_xfer_handles_by_block_size.clear()

2399 + for handles in self.src_xfer_handles_by_tp_ratio.values():

2400 + for handle in handles:

2401 + self.nixl_wrapper.release_dlist_handle(handle)

2402 + self.src_xfer_handles_by_tp_ratio.clear()

2403 + for dst_xfer_side_handles in self.dst_xfer_side_handles.values():

2404 + for dst_xfer_side_handle in dst_xfer_side_handles.values():

2405 + self.nixl_wrapper.release_dlist_handle(dst_xfer_side_handle)

2406 self.dst_xfer_side_handles.clear()

2407 for remote_agents in self._remote_agents.values():

2408 for agent_name in remote_agents.values():