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

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 15, 07:15 UTC

[P/D] Heterogeneous TP (#18833)

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

Nicolò Lucchesi committed Jun 5, 2025 at 01:25 UTC b2fac67130b1cff07b175c4859c297c746e06601

6 files changed +287 -100

tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

+8 -3

@@ -8,7 +8,9 @@ MODELS=(

9 # Number of prefill and decode instances to create

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

11 -NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-2} # Default to 2

11 +NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-1} # Default to 1

12 +PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1}

13 +DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}

15 # Find the git repository root directory

16 GIT_ROOT=$(git rev-parse --show-toplevel)

@@ -74,9 +76,10 @@ run_tests_for_model() {

76 for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do

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

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

79 +

80 # Calculate port number (base port + instance number)

81 PORT=$((8100 + i))

79 - # Calculate side channel port

82 + # Calculate side channel port. Avoid clash with with TP workers.

83 SIDE_CHANNEL_PORT=$((5559 + i))

85 echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"

@@ -87,6 +90,7 @@ run_tests_for_model() {

90 --enforce-eager \

91 --disable-log-requests \

92 --gpu-memory-utilization 0.2 \

93 + --tensor-parallel-size $PREFILLER_TP_SIZE \

94 --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\"}'"

96 if [ -n "$model_args" ]; then

@@ -109,7 +113,7 @@ run_tests_for_model() {

113 # Calculate port number (base port + instance number)

114 PORT=$((8200 + i))

115 # Calculate side channel port

112 - SIDE_CHANNEL_PORT=$((5659 + i))

116 + SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE))

117

118 echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"

119

@@ -119,6 +123,7 @@ run_tests_for_model() {

123 --enforce-eager \

124 --disable-log-requests \

125 --gpu-memory-utilization 0.2 \

126 + --tensor-parallel-size $DECODER_TP_SIZE \

127 --kv-transfer-config '{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\"}'"

128

129 if [ -n "$model_args" ]; then

tests/v1/kv_connector/nixl_integration/test_accuracy.py

@@ -14,6 +14,7 @@ RTOL = 0.03

14 # Model-specific expected values

15 EXPECTED_VALUES = {

16 "Qwen/Qwen3-0.6B": 0.41,

17 + "deepseek-ai/deepseek-vl2-small": 0.59

18 }

20 SIMPLE_PROMPT = "The best part about working on vLLM is that I got to meet so many people across various different organizations like UCB, Google, and Meta which means", # noqa: E501

vllm/distributed/kv_transfer/kv_connector/utils.py

+17 -1

@@ -3,11 +3,12 @@

3 """

4 KV cache helper for store.

5 """

6 +

7 import torch

9 import vllm.envs as envs

10 from vllm import _custom_ops as ops

10 -from vllm.config import VllmConfig

11 +from vllm.config import VllmConfig, get_current_vllm_config

12 from vllm.logger import init_logger

14 logger = init_logger(__name__)

@@ -90,3 +91,18 @@ class model_aware_kv_ops_helper:

91 layer.self_attn.attn._k_scale,

92 layer.self_attn.attn._v_scale,

93 )

94 +

95 +

96 +def get_kv_connector_cache_layout():

97 + vllm_config = get_current_vllm_config()

98 + kv_config = vllm_config.kv_transfer_config

99 + if vllm_config.model_config is None:

100 + logger.warning("Unable to detect current VLLM config. " \

101 + "Defaulting to NHD kv cache layout.")

102 + else:

103 + use_mla = vllm_config.model_config.use_mla

104 + if not use_mla and kv_config.kv_connector == "NixlConnector":

105 + logger.info("NixlConnector detected. Setting KV cache " \

106 + "layout to HND for better xfer performance.")

107 + return "HND"

108 + return "NHD"

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

+243 -95

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

32 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

33 from vllm.v1.request import Request

35 +Transfer = tuple[int, float] # (xfer_handle, start_time)

36 GET_META_MSG = b"get_meta_msg"

38 logger = init_logger(__name__)

@@ -54,6 +55,8 @@ class NixlAgentMetadata(

55 agent_metadata: bytes

56 kv_caches_base_addr: list[int]

57 num_blocks: int

58 + tp_size: int

59 + block_len: int

62 @dataclass

@@ -331,10 +334,14 @@ class NixlConnectorWorker:

334 logger.info("Initializing NIXL wrapper")

335 logger.info("Initializing NIXL worker %s", engine_id)

336

337 + # Config.

338 + self.vllm_config = vllm_config

339 + self.block_size = vllm_config.cache_config.block_size

340 +

341 # Agent.

342 self.nixl_wrapper = NixlWrapper(str(uuid.uuid4()), None)

336 - # Map of engine_id -> agent_name.

337 - self._remote_agents: dict[str, str] = {}

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

344 + self._remote_agents: dict[str, dict[int, str]] = defaultdict(dict)

345

346 # NIXL handshake port.

347 # NOTE(rob): Within a DP group, each DP rank gets its own

@@ -354,7 +361,8 @@ class NixlConnectorWorker:

361 # KV Caches and nixl tracking data.

362 self.kv_caches: dict[str, torch.Tensor] = {}

363

357 - # Map of engine_id -> kv_caches_base_addr

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

365 + # rank will still only pull from a single remote TP worker.

366 self.kv_caches_base_addr: dict[str, list[int]] = {}

367

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

@@ -362,19 +370,19 @@ class NixlConnectorWorker:

370 self.num_regions = 0

371 self.num_layers = 0

372

365 - # nixl_prepped_dlist_handle (int).

373 + # nixl_prepped_dlist_handle.

374 self.src_xfer_side_handle: int = 0

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

376 self.dst_xfer_side_handles: dict[str, int] = {}

377

370 - # Map of engine_id -> num_blocks.

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

379 + # have the same number of blocks.

380 self.dst_num_blocks: dict[str, int] = {}

381 self._registered_descs: list[Any] = []

382

383 # In progress transfers.

384 # [req_id -> list[handle]]

376 - self._recving_transfers: defaultdict[str, list[Any]] = defaultdict(

377 - list[Any])

385 + self._recving_transfers = defaultdict[str, list[Transfer]](list)

386

387 # Complete transfer tracker. Used by the rank 0 to track finished

388 # transactions on ranks 1 to N-1.

@@ -395,6 +403,11 @@ class NixlConnectorWorker:

403 # List of block window sizes for each layer for local attention

404 self.block_window_per_layer: list[Optional[int]] = []

405

406 + self._tp_size: dict[str, int] = {self.engine_id: self.world_size}

407 + # With heterogeneous TP, P must wait for all assigned D TP workers to

408 + # finish reading before safely freeing the blocks.

409 + self.consumer_notification_counts_by_req = defaultdict[str, int](int)

410 +

411 @staticmethod

412 def _nixl_handshake_listener(metadata: NixlAgentMetadata,

413 ready_event: threading.Event, base_port: int,

@@ -426,27 +439,44 @@ class NixlConnectorWorker:

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

440

441 start_time = time.perf_counter()

429 - # NOTE(rob): we need each tp_rank to have a unique port.

430 - # This is a hack to keep us moving. We will switch when

431 - # we switch to HTTP-based NIXL metadata exchange.

432 - path = make_zmq_path("tcp", host, port + self.tp_rank)

433 - logger.debug("Querying metadata on path: %s", path)

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

435 - # Send query for the request.

436 - sock.send(GET_META_MSG)

437 - metadata_bytes = sock.recv()

438 - decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)

439 - metadata = decoder.decode(metadata_bytes)

440 - got_metadata_time = time.perf_counter()

442

442 - # Register Remote agent.

443 - self.add_remote_agent(metadata)

444 - setup_agent_time = time.perf_counter()

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

444 + # a hack to keep us moving. We will switch when moving to etcd

445 + # or where we have a single ZMQ socket in the scheduler.

446

446 - logger.debug("NIXL handshake: get metadata took: %s",

447 - got_metadata_time - start_time)

448 - logger.debug("NIXL handshake: add agent took: %s",

449 - setup_agent_time - got_metadata_time)

447 + def handshake(path: str, rank: int) -> NixlAgentMetadata:

448 + # Send query for the request.

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

450 + sock.send(GET_META_MSG)

451 + metadata_bytes = sock.recv()

452 + decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)

453 + metadata = decoder.decode(metadata_bytes)

454 + got_metadata_time = time.perf_counter()

455 +

456 + # Register Remote agent.

457 + self.add_remote_agent(metadata, rank)

458 + setup_agent_time = time.perf_counter()

459 +

460 + logger.debug("NIXL handshake: get metadata took: %s",

461 + got_metadata_time - start_time)

462 + logger.debug("NIXL handshake: add agent took: %s",

463 + setup_agent_time - got_metadata_time)

464 + return metadata

465 +

466 + # Handshake with remote agent-rank0 first to get the tp_size of remote

467 + path = make_zmq_path("tcp", host, port)

468 + logger.debug("Querying master rank metadata on path: %s", path)

469 + metadata = handshake(path, 0)

470 +

471 + # Handshake only with the other TP remote the current local rank will

472 + # pull from. With homogeneous TP it happens to be the same rank_i.

473 + tp_ratio = self._tp_size[self.engine_id] // metadata.tp_size

474 + p_remote_rank = self.tp_rank // tp_ratio

475 + if p_remote_rank > 0:

476 + path = make_zmq_path("tcp", host, port + p_remote_rank)

477 + logger.debug("Querying metadata on path: %s at remote rank %s",

478 + path, p_remote_rank)

479 + _ = handshake(path, p_remote_rank)

480

481 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):

482 """Register the KV Cache data in nixl."""

@@ -455,24 +485,34 @@ class NixlConnectorWorker:

485 kv_elem_size = first_kv_cache.element_size()

486

487 # TODO(tms): Find a more robust way to detect and handle MLA

458 - use_mla = len(first_kv_cache.shape) == 3

459 - if use_mla:

488 + self.use_mla = len(first_kv_cache.shape) == 3

489 + # NOTE (NickLucche) To move blocks efficiently with NIXL, the expected

490 + # KV memory layout is HND, as opposed to the default NHD. Note that it

491 + # will only affects the strides. For MLA instead, we make require no

492 + # such thing and resort to the standard layout.

493 + if self.use_mla:

494 # MLA case.

495 self.num_blocks = first_kv_cache.shape[0]

496 block_rank = 2 # [block_size, latent_dim]

497 block_shape = first_kv_cache.shape[-block_rank:]

498 + block_size, kv_latent_dim = block_shape

499 + self.slot_size_bytes = kv_elem_size * kv_latent_dim

500 else:

465 - # [2 (k and v), num_blocks, ...]

501 + # [2 (k and v), num_blocks, block_size, kv_heads, head_dim]

502 self.num_blocks = first_kv_cache.shape[1]

503 block_rank = 3 # [block_size, kv_heads, head_dim]

504 block_shape = first_kv_cache.shape[-block_rank:]

469 -

505 + block_size, n_kv_heads, head_dim = block_shape

506 + # head size in bytes.

507 + self.slot_size_bytes = kv_elem_size * n_kv_heads * head_dim

508 + assert block_size == self.block_size

509 # TODO(tms): self.block_len needs to be per-layer for sliding window,

510 # hybrid attn, etc

511 + # block size in bytes

512 self.block_len = kv_elem_size * math.prod(block_shape)

513

474 - logger.debug("Registering KV_Caches. use_mla: %s, shape %s", use_mla,

475 - first_kv_cache.shape)

514 + logger.debug("Registering KV_Caches. use_mla: %s, shape %s",

515 + self.use_mla, first_kv_cache.shape)

516 logger.debug("num_blocks: %s, block_shape: %s", self.num_blocks,

517 block_shape)

518 logger.debug("Per layer kv cache size: %s", first_kv_cache.shape)

@@ -489,7 +529,7 @@ class NixlConnectorWorker:

529 # (roughly 8KB vs 5KB).

530 for cache_or_caches in kv_caches.values():

531 # Normalize to always be a list of caches

492 - cache_list = [cache_or_caches] if use_mla else cache_or_caches

532 + cache_list = [cache_or_caches] if self.use_mla else cache_or_caches

533 for cache in cache_list:

534 base_addr = cache.data_ptr()

535 region_len = self.num_blocks * self.block_len

@@ -524,16 +564,37 @@ class NixlConnectorWorker:

564 logger.debug("Registering descs: %s", caches_data)

565 self.nixl_wrapper.register_memory(descs)

566 logger.debug("Done registering descs")

527 -

567 self._registered_descs.append(descs)

568

569 + # Register local/src descr for NIXL xfer.

570 + blocks_data = []

571 + for base_addr in self.kv_caches_base_addr[self.engine_id]:

572 + # NOTE With heter-TP, more blocks are prepared than what are

573 + # needed as self.num_blocks >= nixl_agent_meta.num_blocks. We

574 + # could create fewer, but then _get_block_descs_ids needs to

575 + # select agent_meta.num_blocks instead of self.num_blocks for

576 + # local descr, and that makes handling regular flow less clean.

577 + for block_id in range(self.num_blocks):

578 + block_offset = block_id * self.block_len

579 + addr = base_addr + block_offset

580 + # (addr, len, device id)

581 + blocks_data.append((addr, self.block_len, self.tp_rank))

582 + logger.debug("Created %s blocks for src engine %s and rank %s",

583 + len(blocks_data), self.engine_id, self.tp_rank)

584 +

585 + descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")

586 + # NIXL_INIT_AGENT to be used for preparations of local descs.

587 + self.src_xfer_side_handle = self.nixl_wrapper.prep_xfer_dlist(

588 + "NIXL_INIT_AGENT", descs)

589 +

590 # After KV Caches registered, listen for new connections.

591 metadata = NixlAgentMetadata(

592 engine_id=self.engine_id,

593 agent_metadata=self.nixl_wrapper.get_agent_metadata(),

594 kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id],

595 num_blocks=self.num_blocks,

536 - )

596 + tp_size=self.world_size,

597 + block_len=self.block_len)

598 ready_event = threading.Event()

599 self._nixl_handshake_listener_t = threading.Thread(

600 target=self._nixl_handshake_listener,

@@ -543,50 +604,123 @@ class NixlConnectorWorker:

604 self._nixl_handshake_listener_t.start()

605 ready_event.wait()

606

546 - def add_remote_agent(self, nixl_agent_meta: NixlAgentMetadata):

607 + def add_remote_agent(self,

608 + nixl_agent_meta: NixlAgentMetadata,

609 + remote_tp_rank: int = 0):

610 + """

611 + Add the remote NIXL agent and prepare the descriptors for reading cache

612 + blocks from remote.

613 +

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

615 + requires local rank_i to read from remote rank_i.

616 + The latter, assuming D.world_size > P.world_size, requires that two or

617 + more local TP worker share the xfer from a single TP worker.

618 +

619 + Here's an example:

620 +

621 + rank_offset p_remote_tp_rank

622 + (kv split no)

623 + --------------------------------

624 + 0 0 Worker0 ---- 1st half of KV ----> Worker0 [ KV Cache ]

625 + /

626 + 1 0 Worker1 ---- 2nd half of KV -----/

627 +

628 + 0 1 Worker2 ---- 1st half of KV ----> Worker1 [ KV Cache ]

629 + /

630 + 1 1 Worker3 ---- 2nd half of KV -----/

631 +

632 +

633 + Decoder TP workers Prefix TP workers

634 + (world_size=4) (world_size=2)

635 + tp_ratio = 4 // 2 = 2

636 +

637 + Considering the KV Caches, if P-Worker_i has cache size [2, num_blocksP, kv_heads, block_size, head_dim]

638 + then D-Worker_j has [2, num_blocksD, kv_heads//tp_ratio, block_size, head_dim]. Mind the "HND" layout format.

639 + Assuming num_blocksD >= num_blocksP, D-Worker0 reads from P-Worker0 by preparing the kv_heads//tp_ratio

640 + first heads from all the slots of all the blocks. D-Worker1 will do the same, but reading the second split

641 + along the kv_heads dimension, and so forth until "tp_ratio" D TP workers have pulled from P-Worker0.

642 +

643 + Note that the above will also hold true for the homogeneous TP case, where tp_ratio evaluates to 1.

644 +

645 + Regarding MLA case, the cache is replicated across TP workers so the rank_offset will just always be 0

646 + so that the whole cache is shared by "tp_ratio" D TP workers.

647 + """ # noqa: E501

648 engine_id = nixl_agent_meta.engine_id

548 - assert engine_id != self.engine_id, "Conflict engine id found!"

549 - if engine_id in self._remote_agents:

649 + # TODO re-evaluate refreshing for scaling/recovery

650 + if remote_tp_rank in self._remote_agents.get(engine_id, ()):

651 return

652

552 - self._remote_agents[engine_id] = self.nixl_wrapper.add_remote_agent(

553 - nixl_agent_meta.agent_metadata)

554 - self.kv_caches_base_addr[

555 - engine_id] = nixl_agent_meta.kv_caches_base_addr

653 + if engine_id in self._tp_size:

654 + assert self._tp_size[engine_id] == nixl_agent_meta.tp_size

655 + else:

656 + self._tp_size[engine_id] = nixl_agent_meta.tp_size

657 + self._remote_agents[engine_id][

658 + remote_tp_rank] = self.nixl_wrapper.add_remote_agent(

659 + nixl_agent_meta.agent_metadata)

660 +

661 + # Number of D TP workers reading from a single P TP worker. This is

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

663 + assert self._tp_size[self.engine_id] % self._tp_size[engine_id] == 0, \

664 + "Local TP size must be divisible by remote TP size."

665 + tp_ratio = self._tp_size[self.engine_id] // self._tp_size[engine_id]

666 + assert tp_ratio > 0, "Decode TP cannot be smaller than"

667 + " prefill TP"

668 + if self.use_mla:

669 + # With MLA the only difference is in the number of blocks.

670 + remote_block_size = nixl_agent_meta.block_len / (

671 + self.slot_size_bytes)

672 + assert self.block_len == nixl_agent_meta.block_len

673 + else:

674 + remote_block_size = nixl_agent_meta.block_len / (

675 + self.slot_size_bytes * tp_ratio)

676 +

677 + assert nixl_agent_meta.block_len == self.block_len * tp_ratio, \

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

679 + local_kv_heads*tp_ratio, block_size, head_dim] and same dtype."

680

557 - # Create src descs and xfer side handles.

558 - blocks_data = []

559 - for base_addr in self.kv_caches_base_addr[self.engine_id]:

560 - for block_id in range(self.num_blocks):

561 - block_offset = block_id * self.block_len

562 - # (addr, len, device id)

563 - blocks_data.append(

564 - (base_addr + block_offset, self.block_len, self.tp_rank))

565 - logger.debug("Created %s blocks for src engine %s and tp_rank %s",

566 - len(blocks_data), self.engine_id, self.tp_rank)

681 + assert self.block_size == remote_block_size, "Remote P worker with \

682 + different block size is not supported"

683

568 - # Register with NIXL.

569 - descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")

570 - self.src_xfer_side_handle = self.nixl_wrapper.prep_xfer_dlist(

571 - "NIXL_INIT_AGENT", descs)

684 + assert self.num_blocks >= nixl_agent_meta.num_blocks

685 +

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

687 + if engine_id in self.dst_num_blocks:

688 + assert self.dst_num_blocks[engine_id] == nixl_agent_meta.num_blocks

689 + else:

690 + self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks

691

573 - # Create dst descs and xfer side handles.

574 - self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks

692 blocks_data = []

576 - for base_addr in self.kv_caches_base_addr[engine_id]:

577 - for block_id in range(nixl_agent_meta.num_blocks):

578 - block_offset = block_id * self.block_len

579 - # (addr, len, device id)

580 - blocks_data.append(

581 - (base_addr + block_offset, self.block_len, self.tp_rank))

582 - logger.debug("Created %s blocks for dst engine %s and tp_rank %s",

583 - len(blocks_data), engine_id, self.tp_rank)

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

694 + # rank. With heterogeneous TP, prepare the descriptors by splitting the

695 + # P KV cache along kv_head dim, of D worker's kv_head size (D>P).

696 + # Eg. PTP1 DTP2 => P0 KV:[block0-KV_0 | block0-KV_1..].

697 + p_remote_tp_rank = self.tp_rank // tp_ratio

698 + # Only register the remote's descriptors if current rank pulls from it.

699 + if p_remote_tp_rank == remote_tp_rank:

700 + self.kv_caches_base_addr[

701 + engine_id] = nixl_agent_meta.kv_caches_base_addr

702 + rank_offset = self.tp_rank % tp_ratio * self.block_len \

703 + if not self.use_mla else 0

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

705 + for base_addr in nixl_agent_meta.kv_caches_base_addr:

706 + for block_id in range(nixl_agent_meta.num_blocks):

707 + block_offset = block_id * nixl_agent_meta.block_len

708 + # For each block, grab the heads chunk belonging to rank_i

709 + # of size remote_nheads // tp_ratio, which correspond to

710 + # self.block_len == remote_block_len//tp_ratio bytes.

711 + addr = base_addr + block_offset + rank_offset

712 + # (addr, len, device id)

713 + blocks_data.append((addr, self.block_len, remote_tp_rank))

714 + logger.debug(

715 + "Created %s blocks for dst engine %s with remote rank %s and " \

716 + "local rank %s",

717 + len(blocks_data), engine_id, remote_tp_rank, self.tp_rank)

718

585 - # Register with NIXL.

586 - descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")

587 - self.dst_xfer_side_handles[

588 - engine_id] = self.nixl_wrapper.prep_xfer_dlist(

589 - self._remote_agents[engine_id], descs)

719 + # Register with NIXL.

720 + descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")

721 + self.dst_xfer_side_handles[

722 + engine_id] = self.nixl_wrapper.prep_xfer_dlist(

723 + self._remote_agents[engine_id][remote_tp_rank], descs)

724

725 def get_finished(self) -> tuple[set[str], set[str]]:

726 """

@@ -654,16 +788,25 @@ class NixlConnectorWorker:

788 return done_sending, done_recving

789

790 def _get_new_notifs(self) -> set[str]:

657 - """Get req_ids which got a remote xfer message."""

658 -

791 + """

792 + Get req_ids which got a remote xfer message. When multiple consumers

793 + are reading from the same producer (heterogeneous TP scenario), wait

794 + for all consumers to be done pulling.

795 + """

796 notified_req_ids: set[str] = set()

660 - for req_ids in self.nixl_wrapper.get_new_notifs().values():

661 - for req_id in req_ids:

662 - assert req_id not in notified_req_ids

663 - notified_req_ids.add(req_id.decode("utf-8"))

797 + for notifs in self.nixl_wrapper.get_new_notifs().values():

798 + for notif in notifs:

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

800 + self.consumer_notification_counts_by_req[req_id] += 1

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

802 + if self.consumer_notification_counts_by_req[req_id] == int(

803 + tp_ratio):

804 + notified_req_ids.add(req_id)

805 + del self.consumer_notification_counts_by_req[req_id]

806 return notified_req_ids

807

666 - def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]:

808 + def _pop_done_transfers(

809 + self, transfers: dict[str, list[tuple[int, float]]]) -> set[str]:

810 """

811 Pop completed xfers by checking for DONE state.

812 Args:

@@ -673,23 +816,17 @@ class NixlConnectorWorker:

816 """

817 done_req_ids: set[str] = set()

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

676 - running_reqs = []

677 - for handle in handles:

819 + for handle, xfer_stime in handles:

820 xfer_state = self.nixl_wrapper.check_xfer_state(handle)

821 if xfer_state == "DONE":

680 - # TODO ptarasiewicz: why abort is throwing errors?

681 - # self.nixl_wrapper.release_xfer_handle(handle)

822 + self.nixl_wrapper.release_xfer_handle(handle)

823 + done_req_ids.add(req_id)

824 + del transfers[req_id]

825 + elif xfer_state == "PROC":

826 continue

683 - if xfer_state == "PROC":

684 - running_reqs.append(handle)

827 else:

828 raise RuntimeError("Transfer failed with state %s",

829 xfer_state)

688 - if len(running_reqs) == 0:

689 - done_req_ids.add(req_id)

690 - del transfers[req_id]

691 - else:

692 - transfers[req_id] = running_reqs

830 return done_req_ids

831

832 def start_load_kv(self, metadata: NixlConnectorMetadata):

@@ -735,13 +872,19 @@ class NixlConnectorWorker:

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

873 # blocks until we are ready.

874

875 + # Number of D TP workers that will read from dst P. Propagate tp_ratio

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

877 + tp_ratio = self._tp_size[

878 + self.engine_id] // self._tp_size[dst_engine_id]

879 + notif_id = f"{request_id}:{tp_ratio}".encode()

880 +

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

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

883 num_local_blocks = len(local_block_ids)

884 if num_local_blocks == 0:

742 - agent_name = self._remote_agents[dst_engine_id]

743 - self.nixl_wrapper.send_notif(agent_name,

744 - notif_msg=request_id.encode("utf-8"))

885 + remote_rank = self.tp_rank // tp_ratio

886 + agent_name = self._remote_agents[dst_engine_id][remote_rank]

887 + self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id)

888 return

889

890 # Partial prefix cache hit: just read uncomputed blocks.

@@ -754,6 +897,10 @@ class NixlConnectorWorker:

897 local_xfer_side_handle = self.src_xfer_side_handle

898 remote_xfer_side_handle = self.dst_xfer_side_handles[dst_engine_id]

899

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

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

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

903 +

904 # Get descs ids.

905 local_block_descs_ids: list[int] = []

906 remote_block_descs_ids: list[int] = []

@@ -797,14 +944,16 @@ class NixlConnectorWorker:

944 local_block_descs_ids,

945 remote_xfer_side_handle,

946 remote_block_descs_ids,

800 - notif_msg=request_id.encode("utf-8"),

947 + notif_msg=notif_id,

948 )

949

950 # Begin async xfer.

951 self.nixl_wrapper.transfer(handle)

952

953 # Use handle to check completion in future step().

807 - self._recving_transfers[request_id].append(handle)

954 + # TODO (NickLucche) surface xfer elapsed time

955 + self._recving_transfers[request_id].append(

956 + (handle, time.perf_counter()))

957

958 def _get_block_descs_ids(self,

959 engine_id: str,

@@ -815,7 +964,6 @@ class NixlConnectorWorker:

964 If layer_idx is provided, we use the region_ids for the given layer.

965 Otherwise, we use all regions.

966 """

818 -

967 if layer_idx is None:

968 region_ids = range(self.num_regions)

969 else:

vllm/v1/attention/backends/flash_attn.py

+16

@@ -16,6 +16,8 @@ from vllm.attention.ops.merge_attn_states import merge_attn_states

16 from vllm.attention.utils.fa_utils import (flash_attn_supports_fp8,

17 get_flash_attn_version)

18 from vllm.config import VllmConfig, get_layers_from_vllm_config

19 +from vllm.distributed.kv_transfer.kv_connector.utils import (

20 + get_kv_connector_cache_layout)

21 from vllm.logger import init_logger

22 from vllm.platforms import current_platform

23 from vllm.utils import cdiv

@@ -70,6 +72,20 @@ class FlashAttentionBackend(AttentionBackend):

72 raise ValueError("Block size must be a multiple of 16.")

73 return (2, num_blocks, block_size, num_kv_heads, head_size)

75 + @staticmethod

76 + def get_kv_cache_stride_order() -> tuple[int, ...]:

77 + # NOTE When running disaggregated PD with NIXL, HND layout is used for

78 + # faster transfer. `stride_order` indicates the permutation that gets

79 + # us from `get_kv_cache_shape` to the actual memory layout we want.

80 + cache_layout = get_kv_connector_cache_layout()

81 + if cache_layout == "NHD":

82 + stride_order = (0, 1, 2, 3, 4)

83 + elif cache_layout == "HND":

84 + stride_order = (0, 1, 3, 2, 4)

85 + else:

86 + raise ValueError("Unknown cache layout format %s.", cache_layout)

87 + return stride_order

88 +

90 @dataclass

91 class FlashAttentionMetadata:

vllm/worker/worker_base.py

+2 -1

@@ -597,7 +597,8 @@ class WorkerWrapperBase:

597

598 def initialize_from_config(self, kv_cache_configs: List[Any]) -> None:

599 kv_cache_config = kv_cache_configs[self.rpc_rank]

600 - self.worker.initialize_from_config(kv_cache_config) # type: ignore

600 + with set_current_vllm_config(self.vllm_config):

601 + self.worker.initialize_from_config(kv_cache_config) # type: ignore

602

603 def init_device(self):

604 with set_current_vllm_config(self.vllm_config):