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

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 4, 22:30 UTC

nixl refactor [2/N]: unify TpKVTopology + HeteroTPTransferConfig into TransferTopology (#39529)

Signed-off-by: Zhanqiu Hu <zhu@redhat.com>

zhanqiuhu committed Apr 20, 2026 at 06:39 UTC cc3993b05d008be370778466865cce1959ac85a3

7 files changed +653 -620

tests/v1/kv_connector/unit/test_mooncake_connector.py

+4 -4

@@ -631,7 +631,7 @@ def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():

631 mock_thread.return_value.is_alive.return_value = False

632

633 worker.use_mla = True

634 - worker.kv_topo.is_mla = True

634 + worker.transfer_topo.is_mla = True

635

636 # MLA cache tensor: shape[-2] is the block size.

637 mla_cache = torch.zeros((2, 16, 96), dtype=torch.float16)

@@ -692,9 +692,9 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):

692 # Override TP rank/size to simulate P TP=2

693 prefill_worker.tp_rank = P_TP_RANK

694 prefill_worker.tp_size = P_TP_SIZE

695 - # Update shared dict so kv_topo sees correct TP size

695 prefill_worker._tp_size[prefill_worker.engine_id] = P_TP_SIZE

697 - prefill_worker.kv_topo.tp_rank = P_TP_RANK

696 + prefill_worker.transfer_topo.tp_rank = P_TP_RANK

697 + prefill_worker.transfer_topo.tp_size = P_TP_SIZE

698

699 prefill_worker.kv_caches_base_addr = [0x1000]

700 prefill_worker.block_len_per_layer = [local_block_len]

@@ -714,7 +714,7 @@ async def test_kv_producer_heterogeneous_tp(monkeypatch, d_tp_size):

714 send_meta.ready.set()

715

716 # Compute target D ranks using the production code path

717 - target_d_ranks = prefill_worker.kv_topo.get_target_remote_ranks(d_tp_size)

717 + target_d_ranks = prefill_worker.transfer_topo.handshake_target_ranks(d_tp_size)

718

719 mock_socket = AsyncMock(spec=zmq.asyncio.Socket)

720 mock_socket.send_multipart = AsyncMock()

tests/v1/kv_connector/unit/test_nixl_connector.py

+17 -14

@@ -21,7 +21,7 @@ from vllm import LLM

21 from vllm.config import KVTransferConfig, set_current_vllm_config

22 from vllm.distributed.kv_transfer.kv_connector.utils import (

23 KVOutputAggregator,

24 - TpKVTopology,

24 + TransferTopology,

25 get_current_attn_backend,

26 )

27 from vllm.distributed.kv_transfer.kv_connector.v1 import nixl

@@ -463,19 +463,20 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

463 test_shape = self.attn_backends[0].get_kv_cache_shape(

464 num_blocks=1, block_size=16, num_kv_heads=1, head_size=1

465 )

466 - self.kv_topo = TpKVTopology(

466 + self.transfer_topo = TransferTopology(

467 tp_rank=self.tp_rank,

468 + tp_size=self.world_size,

469 + block_size=self.block_size,

470 engine_id=self.engine_id,

469 - remote_tp_size=self._tp_size, # shared state

470 - remote_block_size=self._block_size, # shared state

471 is_mla=self.use_mla,

472 + is_mamba=False,

473 total_num_kv_heads=self.model_config.get_total_num_kv_heads(),

474 attn_backends=self.attn_backends,

475 tensor_shape=test_shape,

476 )

477

478 self.compat_hash = compute_nixl_compatibility_hash(

478 - self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks

479 + self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks

480 )

481

482 def _nixl_handshake(

@@ -496,7 +497,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

497 # Adjust remote block length metadata to satisfy heterogeneous TP

498 # invariants enforced during handshake validation.

499 remote_block_lens = list(self.block_len_per_layer)

499 - tp_ratio = self.kv_topo.tp_ratio(remote_tp_size)

500 + tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size)

501 if remote_tp_size > self.world_size:

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

503 remote_block_lens = [

@@ -731,8 +732,9 @@ class TestNixlHandshake:

732 assert set(remote_agents.keys()) == set(range(tp_ratio))

733

734 remote_engine_id = worker.REMOTE_ENGINE_ID

734 - assert worker._tp_size[remote_engine_id] == remote_tp_size

735 - assert -tp_ratio == worker.kv_topo.tp_ratio_from_engine_id(remote_engine_id)

735 + remote_info = worker.transfer_topo.get_engine_info(remote_engine_id)

736 + assert remote_info.remote_tp_size == remote_tp_size

737 + assert -tp_ratio == worker.transfer_topo.tp_ratio(remote_tp_size)

738 # ensure src_xfer_handles_by_tp_ratio is populated with tpratio chunks

739 assert -tp_ratio in worker.src_xfer_handles_by_tp_ratio

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

@@ -796,7 +798,7 @@ class TestNixlHandshake:

798 (conn_p0.connector_worker, conn_p1.connector_worker)

799 ):

800 worker.world_size = p_tp_size

799 - worker.kv_topo.remote_tp_size = {worker.engine_id: p_tp_size}

801 + worker.transfer_topo.tp_size = p_tp_size

802 worker.tp_rank = rank

803 worker.use_mla = True

804

@@ -2337,7 +2339,7 @@ def test_compatibility_hash_validation(

2339 remote_hash = compute_nixl_compatibility_hash(

2340 remote_vllm_config,

2341 decode_worker.backend_name,

2340 - decode_worker.kv_topo.cross_layers_blocks,

2342 + decode_worker.transfer_topo.cross_layers_blocks,

2343 )

2344

2345 prefill_block_size = config_overrides.get("block_size", 16)

@@ -2424,12 +2426,13 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)

2426 test_shape = backend.get_kv_cache_shape(

2427 num_blocks=1, block_size=16, num_kv_heads=1, head_size=1

2428 )

2427 - decode_worker.kv_topo = TpKVTopology(

2429 + decode_worker.transfer_topo = TransferTopology(

2430 tp_rank=decode_worker.tp_rank,

2431 + tp_size=decode_worker.world_size,

2432 + block_size=decode_worker.block_size,

2433 engine_id=decode_worker.engine_id,

2430 - remote_tp_size=decode_worker._tp_size, # shared state

2431 - remote_block_size=decode_worker._block_size, # shared state

2434 is_mla=decode_worker.use_mla,

2435 + is_mamba=False,

2436 total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(),

2437 attn_backends=[backend],

2438 tensor_shape=test_shape,

@@ -2438,7 +2441,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)

2441 decode_worker.compat_hash = compute_nixl_compatibility_hash(

2442 decode_worker.vllm_config,

2443 decode_worker.backend_name,

2441 - decode_worker.kv_topo.cross_layers_blocks,

2444 + decode_worker.transfer_topo.cross_layers_blocks,

2445 )

2446

2447 if error_scenario == "handshake_decode_error":

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

+14 -13

@@ -152,13 +152,14 @@ def test_read_blocks_for_req_expands_remote_ids(

152

153 remote_engine_id = "remote-engine"

154 if has_mamba:

155 - worker._mamba_phys_ratio = {remote_engine_id: remote_ratio}

155 + worker._physical_blocks_per_logical = {remote_engine_id: remote_ratio}

156

157 - # Mock kv_topo: empty remote ranks skips the transfer machinery entirely,

158 - # isolating the block-ID expansion logic.

159 - worker.kv_topo = MagicMock()

160 - worker.kv_topo.get_target_remote_ranks_from_engine_id.return_value = []

161 - worker.kv_topo.tp_ratio_from_engine_id.return_value = 1

157 + # Mock transfer_topo: empty remote ranks skips the transfer machinery

158 + # entirely, isolating the block-ID expansion logic.

159 + worker.transfer_topo = MagicMock()

160 + worker.transfer_topo.target_remote_ranks.return_value = []

161 + worker.transfer_topo.get_engine_info.return_value = MagicMock(remote_tp_size=1)

162 + worker.transfer_topo.tp_ratio.return_value = 1

163

164 metadata = NixlConnectorMetadata()

165 metadata.add_new_req_to_recv(

@@ -317,7 +318,7 @@ def test_get_block_descs_ids_hybrid_ssm():

318 worker._has_mamba = True

319 worker._is_mamba_group = [False, True]

320 worker._physical_blocks_per_logical_kv_block = 1

320 - worker._mamba_phys_ratio = {engine_id: 1}

321 + worker._physical_blocks_per_logical = {engine_id: 1}

322 worker.block_len_per_layer = [100]

323 # num_descs = num_regions * num_blocks (no blocks_first doubling)

324 worker.num_descs = 2 * num_blocks

@@ -355,7 +356,7 @@ def test_get_block_descs_ids_kernel_block_mismatch():

356 worker._has_mamba = True

357 worker._is_mamba_group = [False, True]

358 worker._physical_blocks_per_logical_kv_block = ratio

358 - worker._mamba_phys_ratio = {engine_id: ratio}

359 + worker._physical_blocks_per_logical = {engine_id: ratio}

360 worker.block_len_per_layer = [100]

361 worker.num_descs = 2 * num_blocks # 800

362

@@ -532,15 +533,15 @@ def test_has_mamba_init(

533 ((9216, 524288), 4096, 131),

534 ],

535 )

535 -def test_compute_mamba_phys_ratio(ssm_sizes, block_len, expected_ratio):

536 - """Verify that compute_mamba_phys_ratio is TP-dependent.

536 +def test_compute_physical_blocks_per_logical(ssm_sizes, block_len, expected_ratio):

537 + """Verify that compute_physical_blocks_per_logical is TP-dependent.

538

539 With dimension-sharded Mamba state, the ratio differs across TP sizes

540 (e.g. TP=1 → 261, TP=4 → 131 for Nemotron 30B). This is why

540 - _mamba_phys_ratio must be stored per-engine.

541 + _physical_blocks_per_logical must be stored per-engine.

542 """

543 from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (

543 - compute_mamba_phys_ratio,

544 + compute_physical_blocks_per_logical,

545 )

546

546 - assert compute_mamba_phys_ratio(ssm_sizes, block_len) == expected_ratio

547 + assert compute_physical_blocks_per_logical(ssm_sizes, block_len) == expected_ratio

vllm/distributed/kv_transfer/kv_connector/utils.py

+486 -456

@@ -5,7 +5,7 @@ KV cache helper for store.

5 """

7 from collections.abc import Iterator

8 -from dataclasses import dataclass, field

8 +from dataclasses import dataclass

9 from typing import TYPE_CHECKING, Any, Literal, cast

11 import torch

@@ -319,31 +319,139 @@ def yield_req_data(

319 )

320

321

322 -@dataclass

323 -class TpKVTopology:

322 +def get_current_attn_backends(

323 + vllm_config: VllmConfig, layer_names: list[str] | None = None

324 +) -> list[type[AttentionBackend]]:

325 + """Get all distinct attention backends for the given layers.

326 +

327 + Args:

328 + vllm_config: The current vLLM configuration.

329 + layer_names: Optional list of layer names to scope the lookup.

330 + When None, all attention layers are considered.

331 +

332 + Returns:

333 + Deduplicated list of attention backend classes.

334 + """

335 + layer_type = cast(type[Any], AttentionLayerBase)

336 + layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names)

337 + if layers:

338 + seen: dict[str, type[AttentionBackend]] = {}

339 + for layer in layers.values():

340 + backend = layer.get_attn_backend()

341 + seen[backend.full_cls_name()] = backend

342 + return list(seen.values())

343 +

344 + # Fallback for tests, when static_forward_context is empty.

345 + logger.debug(

346 + "No layers found in the vLLM config. Falling back to default attention backend."

347 + )

348 + from vllm.v1.attention.selector import get_attn_backend

349 +

350 + return [

351 + get_attn_backend(

352 + head_size=vllm_config.model_config.get_head_size(),

353 + dtype=vllm_config.model_config.dtype,

354 + kv_cache_dtype=vllm_config.cache_config.cache_dtype,

355 + use_mla=vllm_config.model_config.use_mla,

356 + )

357 + ]

358 +

359 +

360 +def get_current_attn_backend(

361 + vllm_config: VllmConfig, layer_names: list[str] | None = None

362 +) -> type[AttentionBackend]:

363 + """Get the first attention backend for the given layers."""

364 + return get_current_attn_backends(vllm_config, layer_names)[0]

365 +

366 +

367 +# ---- Per-engine transfer info ----

368 +

369 +

370 +@dataclass(frozen=True)

371 +class EngineTransferInfo:

372 + """Common per-remote-engine transfer state, computed at handshake.

373 +

374 + Stored per ``engine_id`` inside ``TransferTopology._engines``.

375 """

325 - Helper class for tensor parallel and KV topology information for

326 - mapping between local and remote TP workers.

376 +

377 + remote_tp_size: int

378 +

379 + remote_block_len: int

380 + """Block length (bytes)"""

381 +

382 + remote_block_size: int

383 + """Tokens per block."""

384 +

385 + remote_physical_blocks_per_logical: int

386 + """Physical blocks per logical block."""

387 +

388 +

389 +@dataclass(frozen=True)

390 +class MambaEngineTransferInfo(EngineTransferInfo):

391 + """Extends ``EngineTransferInfo`` with Mamba-hybrid transfer geometry.

392 +

393 + For hybrid SSM+Attention models, FA and Mamba layers may require

394 + different numbers of reads from different remote ranks. This

395 + dataclass captures that per-engine transfer plan.

396 """

397

398 + remote_fa_source_ranks: tuple[int, ...]

399 + """Remote ranks carrying unique FA heads for this local rank."""

400 +

401 + remote_all_source_ranks: tuple[int, ...]

402 + """All remote ranks this local rank reads from (FA + Mamba)."""

403 +

404 + remote_num_fa_reads: int

405 + """Number of distinct remote ranks needed for FA data."""

406 +

407 + remote_num_mamba_reads: int

408 + """Number of distinct remote ranks needed for Mamba data."""

409 +

410 + remote_fa_descriptor_bytes: int

411 + """Byte size of one FA K (or V) descriptor entry."""

412 +

413 + is_remote_replicated: bool

414 + """Whether the remote engine has replicated KV heads

415 + (remote_tp_size > total_num_kv_heads)."""

416 +

417 + remote_physical_heads: int

418 + """Physical KV heads stored per remote rank."""

419 +

420 +

421 +# ---- Transfer topology ----

422 +

423 +

424 +@dataclass

425 +class TransferTopology:

426 + """Single source of truth for local TP identity and per-engine remote info."""

427 +

428 tp_rank: int

330 - remote_tp_size: dict[EngineId, int]

429 + tp_size: int

430 + block_size: int

431 + engine_id: EngineId

432 is_mla: bool

433 + is_mamba: bool

434 total_num_kv_heads: int

435 attn_backends: list[type[AttentionBackend]]

334 - engine_id: EngineId

335 - remote_block_size: dict[EngineId, int]

436 tensor_shape: torch.Size | None = None

337 - is_mamba: bool = False

437

438 def __post_init__(self):

439 + self.local_physical_heads = max(1, self.total_num_kv_heads // self.tp_size)

440 +

441 + self._engines: dict[EngineId, EngineTransferInfo] = {}

442 + self._fa_source_sets: dict[EngineId, frozenset[int]] = {}

443 + self._fa_source_indices: dict[EngineId, dict[int, int]] = {}

444 +

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

341 - # or num_blocks. This is used to register the memory regions correctly.

446 + # or num_blocks.

447 attn_backend = self.attn_backends[0]

448 if not self.is_mamba:

449 _MOCK_BLOCK_SIZE = 16

450 kv_cache_shape: tuple[int, ...] = attn_backend.get_kv_cache_shape(

346 - num_blocks=1, block_size=_MOCK_BLOCK_SIZE, num_kv_heads=1, head_size=1

451 + num_blocks=1,

452 + block_size=_MOCK_BLOCK_SIZE,

453 + num_kv_heads=1,

454 + head_size=1,

455 )

456 logger.debug("Test kv_cache_shape: %s", kv_cache_shape)

457 # Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D],

@@ -358,11 +466,9 @@ class TpKVTopology:

466 self._cross_layers_blocks = (

467 len(self.tensor_shape) == len(kv_cache_shape) + 1

468 )

361 - self.tensor_shape: torch.Size

469

470 if self._cross_layers_blocks:

471 logger.debug("Using cross-layer KV cache")

365 - # prepend layers dimension

472 _MOCK_NUM_LAYERS = 80

473 kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape

474 try:

@@ -372,15 +478,81 @@ class TpKVTopology:

478 except (AttributeError, NotImplementedError):

479 assert self.tensor_shape is not None

480 kv_cache_stride_order = tuple(range(len(self.tensor_shape)))

375 -

376 - # In case of cross layers permute kv_cache_shape according to

377 - # stride_order to retrieve physical position of block_size

481 kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)

482

483 + # ============================================================

484 + # Engine registration

485 + # ============================================================

486 +

487 + def register_remote_engine(

488 + self,

489 + remote_engine_id: EngineId,

490 + remote_tp_size: int,

491 + remote_block_size: int,

492 + remote_block_len: int,

493 + remote_physical_blocks_per_logical: int,

494 + *,

495 + local_block_len: int = 0,

496 + ) -> EngineTransferInfo:

497 + """Register a remote engine, unifying worker dicts state.

498 +

499 + Only remote engines should be registered here — the local engine's

500 + identity (tp_size, block_size, etc.) is set via ``__init__`` params.

501 +

502 + For Mamba models, also computes the Mamba transfer plan and

503 + builds the FA source lookup caches.

504 +

505 + Args:

506 + local_block_len: Local representative block_len (bytes).

507 + Required for Mamba models to compute ``fa_descriptor_bytes``.

508 + """

509 + assert remote_engine_id != self.engine_id, (

510 + f"Cannot register local engine {self.engine_id} as remote. "

511 + f"Local identity is set via __init__ params."

512 + )

513 + if remote_engine_id in self._engines:

514 + return self._engines[remote_engine_id]

515 + info: EngineTransferInfo

516 + if self.is_mamba:

517 + info = self._build_mamba_info(

518 + remote_tp_size=remote_tp_size,

519 + remote_block_size=remote_block_size,

520 + remote_block_len=remote_block_len,

521 + remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical),

522 + local_block_len=local_block_len,

523 + )

524 + assert isinstance(info, MambaEngineTransferInfo)

525 + self._fa_source_sets[remote_engine_id] = frozenset(

526 + info.remote_fa_source_ranks

527 + )

528 + self._fa_source_indices[remote_engine_id] = {

529 + r: i for i, r in enumerate(info.remote_fa_source_ranks)

530 + }

531 + else:

532 + info = EngineTransferInfo(

533 + remote_tp_size=remote_tp_size,

534 + remote_block_len=remote_block_len,

535 + remote_block_size=remote_block_size,

536 + remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical),

537 + )

538 + self._engines[remote_engine_id] = info

539 + return info

540 +

541 + def get_engine_info(self, remote_engine_id: EngineId) -> EngineTransferInfo:

542 + return self._engines[remote_engine_id]

543 +

544 + # ============================================================

545 + # Layout properties

546 + # ============================================================

547 +

548 @property

549 def is_kv_layout_blocks_first(self) -> bool:

550 return self._is_kv_layout_blocks_first

551

552 + @property

553 + def cross_layers_blocks(self) -> bool:

554 + return self._cross_layers_blocks

555 +

556 @property

557 def split_k_and_v(self) -> bool:

558 # Whether to register regions for K and V separately (when present).

@@ -388,29 +560,16 @@ class TpKVTopology:

560 self._cross_layers_blocks or self.is_mla or self.is_kv_layout_blocks_first

561 )

562

391 - @property

392 - def tp_size(self) -> int:

393 - return self.remote_tp_size[self.engine_id]

563 + # ============================================================

564 + # Common methods

565 + # ============================================================

566

395 - @property

396 - def block_size(self) -> int:

397 - return self.remote_block_size[self.engine_id]

567 + def tp_ratio(self, remote_tp_size: int) -> int:

568 + """Calculate the tensor parallel ratio between local and remote TP.

569

399 - @property

400 - def cross_layers_blocks(self) -> bool:

401 - return self._cross_layers_blocks

402 -

403 - def tp_ratio(

404 - self,

405 - remote_tp_size: int,

406 - ) -> int:

407 - """

408 - Calculate the tensor parallel ratio between local and remote TP.

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

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

411 - groups of size `tp_ratio`.If remote tp_size > local tp_size, the

412 - ratio is flipped (remote_size/local_size) and the returned value is

413 - negative.

570 + Positive when local_tp >= remote_tp (local workers read from the

571 + same remote worker in groups of size ``tp_ratio``). Negative when

572 + remote_tp > local_tp (ratio is flipped).

573 """

574 if self.tp_size >= remote_tp_size:

575 assert self.tp_size % remote_tp_size == 0, (

@@ -418,78 +577,65 @@ class TpKVTopology:

577 f"by remote tensor parallel size {remote_tp_size}."

578 )

579 return self.tp_size // remote_tp_size

421 -

580 assert remote_tp_size % self.tp_size == 0, (

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

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

583 )

426 - # P TP > D TP case, return the ratio as negative

427 - return -remote_tp_size // self.tp_size

584 + return -(remote_tp_size // self.tp_size)

585

429 - def block_size_ratio(

430 - self,

431 - remote_block_size: int,

432 - ) -> int:

433 - """

434 - Calculate the block size ratio between local and remote TP.

435 - """

586 + def block_size_ratio(self, remote_block_size: int) -> int:

587 + """Calculate the block size ratio between local and remote."""

588 assert self.block_size % remote_block_size == 0, (

589 f"Local block size {self.block_size} is not divisible "

590 f"by remote block size {remote_block_size} or vice versa."

591 )

592 return self.block_size // remote_block_size

593

442 - def tp_ratio_from_engine_id(

443 - self,

444 - remote_engine_id: EngineId,

445 - ) -> int:

446 - remote_tp_size = self.remote_tp_size[remote_engine_id]

447 - return self.tp_ratio(remote_tp_size)

448 -

449 - def block_size_ratio_from_engine_id(

450 - self,

451 - remote_engine_id: EngineId,

452 - ) -> int:

453 - remote_block_size = self.remote_block_size[remote_engine_id]

454 - return self.block_size_ratio(remote_block_size)

455 -

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

457 - """

458 - Whether the KV cache is replicated across TP workers due to the

594 + def is_kv_replicated(self, remote_engine_id: EngineId) -> bool:

595 + """Whether the KV cache is replicated across TP workers due to the

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

460 - When they are equal, each TP rank still owns one distinct KV head,

461 - so this is not considered replication.

597 """

463 - tp_size = self.remote_tp_size[engine_id]

464 - return tp_size > self.total_num_kv_heads

598 + return self._engines[remote_engine_id].remote_tp_size > self.total_num_kv_heads

599

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

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

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

603

470 - def get_target_remote_ranks(

471 - self,

472 - remote_tp_size: int,

473 - ) -> list[int]:

474 - """

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

476 - (on D) will read from. When remote tp_size > local tp_size, we

477 - read from multiple remote ranks.

604 + @property

605 + def local_replicates_kv_cache(self) -> bool:

606 + """Whether the local engine's KV cache is replicated."""

607 + return self.is_mla or self.tp_size > self.total_num_kv_heads

608 +

609 + def handshake_target_ranks(self, remote_tp_size: int) -> list[int]:

610 + """Pre-registration: compute which remote TP ranks to handshake with.

611 +

612 + Pure math based on local/remote TP sizes — does not require

613 + the remote engine to be registered yet.

614 """

615 tp_ratio = self.tp_ratio(remote_tp_size)

616 if tp_ratio > 0:

617 return [self.tp_rank // tp_ratio]

618 + abs_ratio = -tp_ratio

619 + return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)]

620

483 - # P TP > D TP case, D reads from |tp_ratio| remote workers.

484 - tp_ratio = -tp_ratio

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

621 + def target_remote_ranks(self, remote_engine_id: EngineId) -> list[int]:

622 + """Get the remote TP rank(s) that the current local TP rank will

623 + read from. When remote tp_size > local tp_size, reads from

624 + multiple remote ranks.

625

487 - def get_target_remote_ranks_from_engine_id(

488 - self,

489 - remote_engine_id: EngineId,

490 - ) -> list[int]:

491 - remote_tp_size = self.remote_tp_size[remote_engine_id]

492 - return self.get_target_remote_ranks(remote_tp_size)

626 + For Mamba models, returns the precomputed ``all_source_ranks``

627 + (FA + Mamba union).

628 + """

629 + info = self._engines[remote_engine_id]

630 + if isinstance(info, MambaEngineTransferInfo):

631 + return list(info.remote_all_source_ranks)

632 +

633 + tp_ratio = self.tp_ratio(info.remote_tp_size)

634 + if tp_ratio > 0:

635 + return [self.tp_rank // tp_ratio]

636 + # remote TP > local TP: read from |tp_ratio| remote workers

637 + abs_ratio = -tp_ratio

638 + return [self.tp_rank * abs_ratio + i for i in range(abs_ratio)]

639

640 def get_transfer_cache_regions(

641 self, cache: torch.Tensor, layer_spec: "KVCacheSpec"

@@ -498,331 +644,139 @@ class TpKVTopology:

644 also accounting for hybrid SSM models specificities.

645 """

646 if isinstance(layer_spec, MambaSpec):

501 - # Register the whole kv cache shared tensor, including SSM/Conv. This is

502 - # similar to FI with the difference that SSM/Conv have different sizes

647 + # Register the whole kv cache shared tensor, including

648 + # SSM/Conv.

649 conv, ssm = cache

650 return [conv]

651

506 - # Check may be hacky but it's matching `_update_hybrid_attention_mamba_layout`.

652 + # Check may be hacky but it's matching

653 + # `_update_hybrid_attention_mamba_layout`.

654 if self.is_mamba and cache.shape[0] == 2:

508 - # When MAMBA is present, all backends are blocks first, so that blocks

509 - # can be shared between attention layers and mamba layers. Runner

510 - # `_update_hybrid_attention_mamba_layout` already adjusted strides

511 - # for FlashAttn-like backends so its num_blocks first.

512 - # Swap [2<>num_blocks] dims to get required layout for hybrid SSM.

655 + # When MAMBA is present, all backends are blocks first, so

656 + # that blocks can be shared between attention layers and mamba

657 + # layers. Runner already adjusted strides for FlashAttn-like

658 + # backends so its num_blocks first.

659 + # Swap [2<>num_blocks] dims for hybrid SSM layout.

660 cache = cache.transpose(0, 1)

661

662 # Regular case: backends like FA register K/V in separate regions

663 return cache if self.split_k_and_v else [cache]

664

665 + # ============================================================

666 + # Mamba-specific methods

667 + # ============================================================

668

519 -# ---- Mamba-HMA hetero-TP transfer config ----

520 -#

521 -# Key insight: with hetero-TP (P_TP > D_TP), FA KV cache may be

522 -# replicated across P ranks (when P_TP > num_kv_heads), but Mamba

523 -# conv/SSM state is almost always uniquely sharded per P rank. So the

524 -# number of P ranks D must read from can differ between FA and Mamba,

525 -# and they must be handled separately.

526 -

527 -

528 -def _physical_head_range(tp_size: int, num_heads: int, rank: int) -> range:

529 - """Physical KV head range stored in a rank's KV cache tensor.

530 -

531 - When ``tp_size <= num_heads``: sharded, K/TP contiguous heads per rank.

532 - When ``tp_size > num_heads``: 1 physical head per rank. Heads are

533 - distributed **contiguously** (matching vLLM's GQA weight partitioning):

534 - consecutive ranks share a head before moving to the next one.

535 - """

536 - if tp_size <= num_heads:

537 - assert num_heads % tp_size == 0

538 - per_rank = num_heads // tp_size

539 - return range(rank * per_rank, (rank + 1) * per_rank)

540 - else:

541 - h = rank * num_heads // tp_size

542 - return range(h, h + 1)

543 -

544 -

545 -def _range_overlap(a: range, b: range) -> range:

546 - start = max(a.start, b.start)

547 - stop = min(a.stop, b.stop)

548 - return range(start, max(start, stop))

549 -

669 + def should_skip_fa(self, remote_engine_id: EngineId, remote_rank: int) -> bool:

670 + """Whether to skip FA groups for this remote rank (mamba-only)."""

671 + return remote_rank not in self._fa_source_sets[remote_engine_id]

672

551 -@dataclass

552 -class HeteroTPTransferConfig:

553 - """Precomputed transfer plan for one (D rank, P engine) pair.

554 -

555 - Currently only instantiated for Mamba-HMA (hybrid SSM+Attention) models

556 - where FA and mamba require different splitting factors. Could be extended

557 - to other model types that need non-uniform hetero-TP transfer sizing.

558 -

559 - All descriptor sizes are computed here. The guarantee is:

560 - local_entry_size == remote_entry_size (for NIXL)

561 -

562 - Attributes that start with ``fa_`` concern FlashAttention KV cache.

563 - Attributes that start with ``mamba_`` concern Mamba conv/SSM state.

564 - """

565 -

566 - # ---- Input parameters (from handshake) ----

567 - tp_ratio: int

568 - K: int # total_num_kv_heads (before TP sharding)

569 - d_tp: int # D engine's tensor_parallel_size

570 - p_tp: int # P engine's tensor_parallel_size

571 - d_rank: int # this D worker's TP rank

572 - use_mla: bool

573 -

574 - # Per-layer block lengths (bytes, K+V combined for blocks_first).

575 - # Uniform across layers for current models.

576 - d_block_len: int # D's block_len_per_layer (representative)

577 - p_block_len: int # P's block_len_per_layer (from handshake)

578 - is_blocks_first: bool # kv_topo.is_kv_layout_blocks_first

579 -

580 - # ---- Derived: computed in __post_init__ ----

581 - #

582 - # Physical heads per rank (what the KV tensor actually stores)

583 - d_physical_heads: int = field(init=False)

584 - p_physical_heads: int = field(init=False)

585 -

586 - # How many distinct P ranks D needs for FA data

587 - physical_fa_num_reads: int = field(init=False)

588 -

589 - # Which P ranks contribute unique FA heads (ordered by head index)

590 - fa_read_targets: list[int] = field(init=False)

591 -

592 - # All P ranks needed for mamba (always abs_tp for tp_ratio < 0)

593 - mamba_num_reads: int = field(init=False)

594 -

595 - # All P ranks this D rank communicates with (FA ∪ mamba)

596 - transfer_targets: list[int] = field(init=False)

597 -

598 - # FA descriptor entry size (K or V side, for blocks_first layout)

599 - # Guaranteed: fa_entry_size is the SAME for local handle AND remote desc.

600 - fa_entry_size: int = field(init=False)

601 -

602 - # Replication flags

603 - is_d_replicated: bool = field(init=False)

604 - is_p_replicated: bool = field(init=False)

605 -

606 - # Pre-built set for fast lookup

607 - _fa_target_set: frozenset[int] = field(init=False, repr=False)

608 - # Map: P rank → index in fa_read_targets (for head slot offset)

609 - _fa_target_index: dict[int, int] = field(init=False, repr=False)

610 -

611 - def __post_init__(self) -> None:

612 - K = self.K

613 - self.is_d_replicated = self.d_tp > K

614 - self.is_p_replicated = self.p_tp > K

615 -

616 - self.d_physical_heads = max(1, K // self.d_tp)

617 - self.p_physical_heads = max(1, K // self.p_tp)

618 -

619 - abs_tp = -self.tp_ratio if self.tp_ratio < 0 else 1

620 -

621 - # ---- Mamba range (computed first so FA can prefer ranks in it) ----

622 - mamba_range: range | None = None

623 - if self.tp_ratio < 0:

624 - mamba_range = range(self.d_rank * abs_tp, (self.d_rank + 1) * abs_tp)

625 -

626 - # ---- FA read targets ----

627 - if self.use_mla or self.tp_ratio >= 0:

628 - self.physical_fa_num_reads = 1

629 - self.fa_read_targets = (

630 - [0]

631 - if self.use_mla

632 - # Must match kv_topo.get_target_remote_ranks (d_rank // tp_ratio).

633 - else [

634 - self.d_rank // self.tp_ratio if self.tp_ratio > 0 else self.d_rank

635 - ]

636 - )

637 - else:

638 - d_needs = _physical_head_range(self.d_tp, K, self.d_rank)

639 - # When mamba range exists, prefer P ranks within it so that

640 - # FA targets are a subset of mamba transfer_targets (avoids

641 - # orphaned FA targets outside the transfer loop).

642 - search_range = mamba_range if mamba_range is not None else range(self.p_tp)

643 - seen: set[tuple[int, int]] = set()

644 - targets: list[int] = []

645 - for p in search_range:

646 - p_has = _physical_head_range(self.p_tp, K, p)

647 - ov = _range_overlap(d_needs, p_has)

648 - if len(ov) > 0:

649 - key = (ov.start, ov.stop)

650 - if key not in seen:

651 - seen.add(key)

652 - targets.append(p)

653 - if not targets:

654 - # Fallback: search globally (should not happen in practice)

655 - for p in range(self.p_tp):

656 - p_has = _physical_head_range(self.p_tp, K, p)

657 - ov = _range_overlap(d_needs, p_has)

658 - if len(ov) > 0:

659 - key = (ov.start, ov.stop)

660 - if key not in seen:

661 - seen.add(key)

662 - targets.append(p)

663 - self.fa_read_targets = targets

664 - self.physical_fa_num_reads = len(targets)

665 -

666 - self._fa_target_set = frozenset(self.fa_read_targets)

667 - self._fa_target_index = {r: i for i, r in enumerate(self.fa_read_targets)}

668 -

669 - # ---- Mamba targets ----

670 - if mamba_range is not None and abs_tp > self.physical_fa_num_reads:

671 - self.mamba_num_reads = abs_tp

672 - self.transfer_targets = list(mamba_range)

673 - else:

674 - self.mamba_num_reads = self.physical_fa_num_reads

675 - self.transfer_targets = list(self.fa_read_targets)

676 -

677 - # ---- FA entry size ----

678 - # For blocks_first: block_len_per_layer includes K+V; // 2 gives K (or V).

679 - # Use min(D, P) because D indexes into P when tp_ratio > 0,

680 - # and P is the natural unit when tp_ratio < 0.

681 - effective_block_len = min(self.d_block_len, self.p_block_len)

682 - if self.is_blocks_first:

683 - self.fa_entry_size = effective_block_len // 2

684 - else:

685 - self.fa_entry_size = effective_block_len

686 -

687 - self._validate()

688 -

689 - def _validate(self) -> None:

690 - """Cross-check internal consistency."""

691 - if self.is_d_replicated and self.is_p_replicated and self.tp_ratio > 0:

692 - logger.info(

693 - "Both-replicated hetero-TP: D_TP=%d > P_TP=%d > K=%d. "

694 - "Using d_rank // tp_ratio routing with relative head offset.",

695 - self.d_tp,

696 - self.p_tp,

697 - self.K,

698 - )

699 -

700 - # FA targets must be a subset of transfer_targets

701 - tt_set = set(self.transfer_targets)

702 - for t in self.fa_read_targets:

703 - if t not in tt_set:

704 - logger.error(

705 - "FA target P rank %d is NOT in transfer_targets %s. "

706 - "This will cause missed FA reads!",

707 - t,

708 - self.transfer_targets,

709 - )

710 -

711 - # For tp_ratio < 0 with blocks_first: D_K_half / reads should == P_K_half

712 - if (

713 - self.is_blocks_first

714 - and self.tp_ratio < 0

715 - and self.physical_fa_num_reads > 0

716 - ):

717 - d_k_half = self.d_block_len // 2

718 - p_k_half = self.p_block_len // 2

719 - expected_local = d_k_half // self.physical_fa_num_reads

720 - if expected_local != p_k_half:

721 - logger.warning(

722 - "FA size mismatch: D_K_half=%d / reads=%d = %d, "

723 - "but P_K_half=%d. This may indicate a head count or "

724 - "Mamba-HMA inflation inconsistency.",

725 - d_k_half,

726 - self.physical_fa_num_reads,

727 - expected_local,

728 - p_k_half,

729 - )

673 + def fa_head_slot(self, remote_engine_id: EngineId, remote_rank: int) -> int:

674 + """Index into local FA block for this remote rank's head data.

675

731 - # ---- Query methods ----

732 -

733 - def should_skip_fa(self, p_rank: int) -> bool:

734 - """Whether to skip FA groups for this P rank (mamba-only transfer)."""

735 - return p_rank not in self._fa_target_set

736 -

737 - def fa_head_slot(self, p_rank: int) -> int:

738 - """Index into D's FA block for this P rank's head data.

739 -

740 - For P ranks in fa_read_targets, returns 0, 1, ..., reads-1.

741 - For P ranks NOT in fa_read_targets (replicated duplicates),

742 - returns the slot of the matching FA target with the same head.

676 + For remote ranks in ``fa_source_ranks``, returns 0, 1, …, reads-1.

677 + For ranks NOT in ``fa_source_ranks`` (replicated duplicates),

678 + returns the slot of the matching source rank with the same head.

679 """

744 - if p_rank in self._fa_target_index:

745 - return self._fa_target_index[p_rank]

746 - # Duplicate head: find which fa_target has the same physical head

747 - p_head = _physical_head_range(self.p_tp, self.K, p_rank)

748 - for target in self.fa_read_targets:

749 - t_head = _physical_head_range(self.p_tp, self.K, target)

750 - if _range_overlap(p_head, t_head):

751 - return self._fa_target_index[target]

752 - return 0 # fallback

753 -

754 - def fa_rank_offset(self, remote_kv_block_len: int) -> int:

755 - """Byte offset into P's FA block for this D rank.

756 -

757 - When D is replicated (D_TP > K), multiple D ranks share a head.

758 - Computes offset *relative to the target P rank's first head*

759 - so it works regardless of how many heads P has.

760 - When neither side replicates, falls back to tp_rank % tp_ratio.

761 - Returns 0 when D does not index into P's block.

680 + fa_index = self._fa_source_indices[remote_engine_id]

681 + if remote_rank in fa_index:

682 + return fa_index[remote_rank]

683 + mamba_info = self._engines[remote_engine_id]

684 + assert isinstance(mamba_info, MambaEngineTransferInfo)

685 + K = self.total_num_kv_heads

686 + remote_tp = mamba_info.remote_tp_size

687 + r_head = self._physical_head_range(remote_tp, K, remote_rank)

688 + for target in mamba_info.remote_fa_source_ranks:

689 + t_head = self._physical_head_range(remote_tp, K, target)

690 + if self._range_overlap(r_head, t_head):

691 + return fa_index[target]

692 + return 0

693 +

694 + def fa_rank_offset(

695 + self, remote_engine_id: EngineId, remote_kv_block_len: int

696 + ) -> int:

697 + """Byte offset into remote FA block for this local rank.

698 +

699 + When local TP is replicated (local_tp > K), multiple local ranks

700 + share a head. Computes offset *relative to the target remote

701 + rank's first head* so it works regardless of how many heads the

702 + remote has. Returns 0 when local does not index into remote.

703 """

763 - if self.use_mla or self.tp_ratio <= 0:

704 + mamba_info = self._engines[remote_engine_id]

705 + assert isinstance(mamba_info, MambaEngineTransferInfo)

706 + tp_ratio = self.tp_ratio(mamba_info.remote_tp_size)

707 + if self.is_mla or tp_ratio <= 0:

708 return 0

765 - if self.is_d_replicated:

766 - d_head = self.d_rank * self.K // self.d_tp

767 - p_rank = self.fa_read_targets[0]

768 - p_start = p_rank * self.K // self.p_tp

769 - return (d_head - p_start) * remote_kv_block_len

770 - return self.d_rank % self.tp_ratio * remote_kv_block_len

771 -

772 - @property

773 - def needs_split_handles(self) -> bool:

774 - """Whether per-P-rank split handles are needed.

709 + K = self.total_num_kv_heads

710 + is_local_replicated = self.tp_size > K

711 + if is_local_replicated:

712 + local_head = self.tp_rank * K // self.tp_size

713 + p_rank = mamba_info.remote_fa_source_ranks[0]

714 + p_start = p_rank * K // mamba_info.remote_tp_size

715 + return (local_head - p_start) * remote_kv_block_len

716 + return self.tp_rank % tp_ratio * remote_kv_block_len

717 +

718 + def needs_split_handles(self, remote_engine_id: EngineId) -> bool:

719 + """Whether per-remote-rank split handles are needed.

720

721 True when FA and mamba have different read counts, requiring

722 different splitting factors in the local handle.

723 """

779 - return self.tp_ratio < 0 and not self.use_mla and len(self.transfer_targets) > 1

724 + mamba_info = self._engines[remote_engine_id]

725 + assert isinstance(mamba_info, MambaEngineTransferInfo)

726 + tp_ratio = self.tp_ratio(mamba_info.remote_tp_size)

727 + return (

728 + tp_ratio < 0

729 + and not self.is_mla

730 + and len(mamba_info.remote_all_source_ranks) > 1

731 + )

732

733 def compute_split_handle_data(

734 self,

735 + remote_engine_id: EngineId,

736 src_blocks_data: list[tuple[int, int, int]],

737 num_fa_descs: int,

738 abs_tp: int,

739 ) -> list[list[tuple[int, int, int]]]:

787 - """Compute per-P-rank (addr, len, tp) triples for Mamba-HMA split handles.

740 + """Per-remote-rank (addr, len, dev) triples for Mamba-HMA split

741 + handles.

742

743 FA descriptors (indices < num_fa_descs) are sliced by

790 - ``physical_fa_num_reads``; mamba descriptors are sliced uniformly

744 + ``remote_num_fa_reads``; mamba descriptors are sliced uniformly

745 by ``abs_tp``.

792 -

793 - Returns one list of triples per transfer target.

746 """

747 + mamba_info = self._engines[remote_engine_id]

748 + assert isinstance(mamba_info, MambaEngineTransferInfo)

749 all_handle_data: list[list[tuple[int, int, int]]] = []

796 - for p_idx, p_rank in enumerate(self.transfer_targets):

750 + for p_idx, p_rank in enumerate(mamba_info.remote_all_source_ranks):

751 handle_data: list[tuple[int, int, int]] = []

798 - skip_fa = self.should_skip_fa(p_rank)

799 - fa_slot = self.fa_head_slot(p_rank) if not skip_fa else 0

800 -

801 - for j, (addr, local_len, tp) in enumerate(src_blocks_data):

752 + skip_fa = self.should_skip_fa(remote_engine_id, p_rank)

753 + fa_slot = self.fa_head_slot(remote_engine_id, p_rank) if not skip_fa else 0

754 + for j, (addr, local_len, dev) in enumerate(src_blocks_data):

755 if j < num_fa_descs:

803 - assert self.physical_fa_num_reads >= 1

804 - fa_chunk = local_len // self.physical_fa_num_reads

805 - handle_data.append((addr + fa_slot * fa_chunk, fa_chunk, tp))

756 + assert mamba_info.remote_num_fa_reads >= 1

757 + fa_chunk = local_len // mamba_info.remote_num_fa_reads

758 + handle_data.append((addr + fa_slot * fa_chunk, fa_chunk, dev))

759 else:

760 mamba_chunk = local_len // abs_tp

808 - handle_data.append((addr + p_idx * mamba_chunk, mamba_chunk, tp))

761 + handle_data.append((addr + p_idx * mamba_chunk, mamba_chunk, dev))

762 all_handle_data.append(handle_data)

763 return all_handle_data

764

765 def filter_block_ids_for_rank(

766 self,

767 + remote_engine_id: EngineId,

768 remote_rank: int,

769 local_ids: BlockIds,

770 remote_ids: BlockIds,

771 is_mamba_group: list[bool],

772 ) -> tuple[BlockIds, BlockIds]:

819 - """Zero out FA groups for P ranks outside fa_read_targets.

773 + """Zero out FA groups for remote ranks outside ``fa_source_ranks``.

774

775 Returns (filtered_local_ids, filtered_remote_ids). When the

822 - remote rank carries FA data for this D rank, returns the inputs

823 - unchanged.

776 + remote rank carries FA data for this local rank, returns the

777 + inputs unchanged.

778 """

825 - if not self.should_skip_fa(remote_rank):

779 + if not self.should_skip_fa(remote_engine_id, remote_rank):

780 return local_ids, remote_ids

781 num_groups = len(local_ids)

782 filtered_local: list[list[int]] = [

@@ -833,108 +787,184 @@ class HeteroTPTransferConfig:

787 ]

788 return filtered_local, filtered_remote

789

836 - def describe(self) -> str:

837 - """One-line summary for logging."""

838 - return (

839 - f"HeteroTPTransferConfig("

840 - f"tp_ratio={self.tp_ratio}, K={self.K}, "

841 - f"d_tp={self.d_tp}, p_tp={self.p_tp}, d_rank={self.d_rank}, "

842 - f"physical_fa_reads={self.physical_fa_num_reads}, "

843 - f"mamba_reads={self.mamba_num_reads}, "

844 - f"fa_targets={self.fa_read_targets}, "

845 - f"transfer_targets={self.transfer_targets}, "

846 - f"fa_entry_size={self.fa_entry_size}, "

847 - f"d_block_len={self.d_block_len}, p_block_len={self.p_block_len})"

790 + def describe(self, remote_engine_id: EngineId) -> str:

791 + """One-line summary of transfer config for logging."""

792 + info = self._engines[remote_engine_id]

793 + base = (

794 + f"tp_ratio={self.tp_ratio(info.remote_tp_size)}, "

795 + f"K={self.total_num_kv_heads}, "

796 + f"local_tp={self.tp_size}, "

797 + f"remote_tp={info.remote_tp_size}, "

798 + f"local_rank={self.tp_rank}, "

799 + f"remote_block_len={info.remote_block_len}"

800 )

801 + if isinstance(info, MambaEngineTransferInfo):

802 + return (

803 + f"TransferTopology.mamba({base}, "

804 + f"fa_reads={info.remote_num_fa_reads}, "

805 + f"mamba_reads={info.remote_num_mamba_reads}, "

806 + f"fa_sources={list(info.remote_fa_source_ranks)}, "

807 + f"all_sources={list(info.remote_all_source_ranks)}, "

808 + f"fa_desc_bytes={info.remote_fa_descriptor_bytes})"

809 + )

810 + return f"TransferTopology({base})"

811 +

812 + # ============================================================

813 + # Private helpers

814 + # ============================================================

815 + # Mamba-HMA hetero-TP transfer config:

816 + # With hetero-TP (P_TP > D_TP), FA KV cache may be replicated across

817 + # P ranks (when P_TP > num_kv_heads), but Mamba conv/SSM state is

818 + # almost always uniquely sharded per P rank. So the number of P

819 + # ranks D must read from can differ between FA and Mamba, and they

820 + # must be handled separately.

821 +

822 + @staticmethod

823 + def _physical_head_range(tp_size: int, num_heads: int, rank: int) -> range:

824 + """Physical KV head range stored in a rank's KV cache tensor.

825 +

826 + When ``tp_size <= num_heads``: sharded, K/TP contiguous heads per rank.

827 + When ``tp_size > num_heads``: 1 physical head per rank. Heads are

828 + distributed **contiguously** (matching vLLM's GQA weight partitioning):

829 + consecutive ranks share a head before moving to the next one.

830 + """

831 + if tp_size <= num_heads:

832 + assert num_heads % tp_size == 0

833 + per_rank = num_heads // tp_size

834 + return range(rank * per_rank, (rank + 1) * per_rank)

835 + else:

836 + h = rank * num_heads // tp_size

837 + return range(h, h + 1)

838

839 + @staticmethod

840 + def _range_overlap(a: range, b: range) -> range:

841 + start = max(a.start, b.start)

842 + stop = min(a.stop, b.stop)

843 + return range(start, max(start, stop))

844

851 -def get_current_attn_backends(

852 - vllm_config: VllmConfig, layer_names: list[str] | None = None

853 -) -> list[type[AttentionBackend]]:

854 - """Get all distinct attention backends for the given layers.

845 + # ============================================================

846 + # Private: build Mamba transfer info

847 + # ============================================================

848

856 - Args:

857 - vllm_config: The current vLLM configuration.

858 - layer_names: Optional list of layer names to scope the lookup.

859 - When None, all attention layers are considered.

849 + def _build_mamba_info(

850 + self,

851 + remote_tp_size: int,

852 + remote_block_size: int,

853 + remote_block_len: int,

854 + remote_physical_blocks_per_logical: int,

855 + local_block_len: int,

856 + ) -> MambaEngineTransferInfo:

857 + """Compute Mamba transfer plan."""

858 + K = self.total_num_kv_heads

859 + local_tp = self.tp_size

860 + local_rank = self.tp_rank

861 +

862 + is_remote_replicated = remote_tp_size > K

863 + remote_physical_heads = max(1, K // remote_tp_size)

864 +

865 + if local_tp >= remote_tp_size:

866 + assert local_tp % remote_tp_size == 0

867 + tp_ratio = local_tp // remote_tp_size

868 + else:

869 + assert remote_tp_size % local_tp == 0

870 + tp_ratio = -(remote_tp_size // local_tp)

871

861 - Returns:

862 - Deduplicated list of attention backend classes.

863 - """

864 - layer_type = cast(type[Any], AttentionLayerBase)

865 - layers = get_layers_from_vllm_config(vllm_config, layer_type, layer_names)

866 - if layers:

867 - seen: dict[str, type[AttentionBackend]] = {}

868 - for layer in layers.values():

869 - backend = layer.get_attn_backend()

870 - seen[backend.full_cls_name()] = backend

871 - return list(seen.values())

872 + abs_tp = -tp_ratio if tp_ratio < 0 else 1

873

873 - # Fallback for tests, when static_forward_context is empty.

874 - logger.debug(

875 - "No layers found in the vLLM config. Falling back to default attention backend."

876 - )

877 - from vllm.v1.attention.selector import get_attn_backend

874 + mamba_range: range | None = None

875 + if tp_ratio < 0:

876 + mamba_range = range(local_rank * abs_tp, (local_rank + 1) * abs_tp)

877

879 - return [

880 - get_attn_backend(

881 - head_size=vllm_config.model_config.get_head_size(),

882 - dtype=vllm_config.model_config.dtype,

883 - kv_cache_dtype=vllm_config.cache_config.cache_dtype,

884 - use_mla=vllm_config.model_config.use_mla,

885 - )

886 - ]

878 + # ---- FA read targets ----

879 + if self.is_mla or tp_ratio >= 0:

880 + num_fa_reads = 1

881 + fa_source_ranks: list[int] = (

882 + [0]

883 + if self.is_mla

884 + else [local_rank // tp_ratio if tp_ratio > 0 else local_rank]

885 + )

886 + else:

887 + local_needs = self._physical_head_range(local_tp, K, local_rank)

888 + search_range = (

889 + mamba_range if mamba_range is not None else range(remote_tp_size)

890 + )

891 + seen: set[tuple[int, int]] = set()

892 + fa_source_ranks = []

893 + for p in search_range:

894 + p_has = self._physical_head_range(remote_tp_size, K, p)

895 + ov = self._range_overlap(local_needs, p_has)

896 + if len(ov) > 0:

897 + key = (ov.start, ov.stop)

898 + if key not in seen:

899 + seen.add(key)

900 + fa_source_ranks.append(p)

901 + if not fa_source_ranks:

902 + for p in range(remote_tp_size):

903 + p_has = self._physical_head_range(remote_tp_size, K, p)

904 + ov = self._range_overlap(local_needs, p_has)

905 + if len(ov) > 0:

906 + key = (ov.start, ov.stop)

907 + if key not in seen:

908 + seen.add(key)

909 + fa_source_ranks.append(p)

910 + num_fa_reads = len(fa_source_ranks)

911

912 + # ---- All source ranks (mamba + FA) ----

913 + if mamba_range is not None and abs_tp > num_fa_reads:

914 + num_mamba_reads = abs_tp

915 + all_source_ranks = list(mamba_range)

916 + else:

917 + num_mamba_reads = num_fa_reads

918 + all_source_ranks = list(fa_source_ranks)

919

889 -def get_current_attn_backend(

890 - vllm_config: VllmConfig, layer_names: list[str] | None = None

891 -) -> type[AttentionBackend]:

892 - """Get the first attention backend for the given layers."""

893 - return get_current_attn_backends(vllm_config, layer_names)[0]

920 + # ---- FA descriptor bytes ----

921 + effective_block_len = min(local_block_len, remote_block_len)

922 + if self.is_kv_layout_blocks_first:

923 + fa_descriptor_bytes = effective_block_len // 2

924 + else:

925 + fa_descriptor_bytes = effective_block_len

926

927 + # ---- Validation ----

928 + is_local_replicated = local_tp > K

929 + if is_local_replicated and is_remote_replicated and tp_ratio > 0:

930 + logger.info(

931 + "Both-replicated hetero-TP: local_tp=%d > remote_tp=%d > K=%d.",

932 + local_tp,

933 + remote_tp_size,

934 + K,

935 + )

936 + tt_set = set(all_source_ranks)

937 + for t in fa_source_ranks:

938 + if t not in tt_set:

939 + logger.error(

940 + "FA source rank %d NOT in all_source_ranks %s.",

941 + t,

942 + all_source_ranks,

943 + )

944 + if self.is_kv_layout_blocks_first and tp_ratio < 0 and num_fa_reads > 0:

945 + local_k_half = local_block_len // 2

946 + remote_k_half = remote_block_len // 2

947 + expected = local_k_half // num_fa_reads

948 + if expected != remote_k_half:

949 + logger.warning(

950 + "FA size mismatch: local_k_half=%d / reads=%d = %d, "

951 + "but remote_k_half=%d.",

952 + local_k_half,

953 + num_fa_reads,

954 + expected,

955 + remote_k_half,

956 + )

957

896 -# TODO (ZhanqiuHu): Consolidate TpKVTopology and HeteroTPTransferConfig

897 -# into a single engine-agnostic TransferTopology class.

898 -# 6 of 9 HeteroTPTransferConfig init fields duplicate TpKVTopology data.

899 -#

900 -# @dataclass

901 -# class EngineTransferInfo:

902 -# """Per-remote-engine transfer state, computed at handshake."""

903 -# p_tp: int

904 -# tp_ratio: int

905 -# p_block_len: int

906 -# block_size: int

907 -# # Mamba-specific (None for non-mamba models)

908 -# fa_read_targets: list[int] | None = None

909 -# transfer_targets: list[int] | None = None

910 -# physical_fa_num_reads: int | None = None

911 -# mamba_num_reads: int | None = None

912 -# fa_entry_size: int | None = None

913 -#

914 -# class TransferTopology:

915 -# """Single source of truth for TP topology + transfer sizing."""

916 -# # Shared (set once at init, replaces duplicate fields)

917 -# tp_rank: int # == TpKVTopology.tp_rank == HeteroTP.d_rank

918 -# tp_size: int # == TpKVTopology.tp_size == HeteroTP.d_tp

919 -# total_num_kv_heads: int # == HeteroTP.K

920 -# is_mla: bool # == HeteroTP.use_mla

921 -# is_mamba: bool

922 -# is_blocks_first: bool # == HeteroTP.is_blocks_first

923 -# d_block_len: int

924 -#

925 -# # Per-engine (populated via register_engine() at handshake)

926 -# _engines: dict[EngineId, EngineTransferInfo]

927 -#

928 -# def register_engine(self, engine_id, p_tp, p_block_len, ...): ...

929 -#

930 -# # General (from TpKVTopology)

931 -# def tp_ratio(self, engine_id) -> int: ...

932 -# def target_remote_ranks(self, engine_id) -> list[int]: ...

933 -# def is_kv_replicated(self, engine_id) -> bool: ...

934 -#

935 -# # Mamba-specific (from HeteroTPTransferConfig, gated by is_mamba)

936 -# def fa_rank_offset(self, engine_id, block_len) -> int: ...

937 -# def physical_fa_num_reads(self, engine_id) -> int: ...

938 -# def transfer_targets(self, engine_id) -> list[int]: ...

939 -# def should_skip_fa(self, engine_id, p_rank) -> bool: ...

940 -# def filter_block_ids_for_rank(self, engine_id, ...) -> ...: ...

958 + return MambaEngineTransferInfo(

959 + remote_tp_size=remote_tp_size,

960 + remote_block_len=remote_block_len,

961 + remote_block_size=remote_block_size,

962 + remote_physical_blocks_per_logical=(remote_physical_blocks_per_logical),

963 + remote_fa_source_ranks=tuple(fa_source_ranks),

964 + remote_all_source_ranks=tuple(all_source_ranks),

965 + remote_num_fa_reads=num_fa_reads,

966 + remote_num_mamba_reads=num_mamba_reads,

967 + remote_fa_descriptor_bytes=fa_descriptor_bytes,

968 + is_remote_replicated=is_remote_replicated,

969 + remote_physical_heads=remote_physical_heads,

970 + )

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py

+11 -11

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

21 from vllm.config import VllmConfig

22 from vllm.distributed.kv_transfer.kv_connector.utils import (

23 EngineId,

24 - TpKVTopology,

24 + TransferTopology,

25 get_current_attn_backend,

26 get_current_attn_backends,

27 )

@@ -764,13 +764,13 @@ class MooncakeConnectorWorker:

764 logger.debug("Detected kv cache layout %s", self.kv_cache_layout)

765

766 self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size}

767 - self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size}

768 - self.kv_topo = TpKVTopology(

767 + self.transfer_topo = TransferTopology(

768 tp_rank=self.tp_rank,

769 + tp_size=self.tp_size,

770 + block_size=self.block_size,

771 engine_id=self.engine_id,

771 - remote_tp_size=self._tp_size, # shared state

772 - remote_block_size=self._block_size, # shared state

772 is_mla=self.use_mla,

773 + is_mamba=False,

774 total_num_kv_heads=self.model_config.get_total_num_kv_heads(),

775 attn_backends=[backend],

776 )

@@ -911,7 +911,7 @@ class MooncakeConnectorWorker:

911 self, identity: bytes, sock: zmq.asyncio.Socket, meta: MooncakeXferMetadata

912 ):

913 pending_reqs: dict[ReqId, SendBlockMeta] = {}

914 - remote_tp_ranks = self.kv_topo.get_target_remote_ranks(meta.remote_tp_size)

914 + remote_tp_ranks = self.transfer_topo.handshake_target_ranks(meta.remote_tp_size)

915 if meta.remote_tp_rank not in remote_tp_ranks:

916 # This D worker does not pair with the P worker.

917 msg = (

@@ -1256,7 +1256,7 @@ class MooncakeConnectorWorker:

1256 seen_base_addresses = []

1257 self.block_len_per_layer = []

1258

1259 - split_k_and_v = self.kv_topo.split_k_and_v

1259 + split_k_and_v = self.transfer_topo.split_k_and_v

1260 tensor_size_bytes = None

1261 for layer_name, cache_or_caches in kv_caches.items():

1262 cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]

@@ -1495,8 +1495,8 @@ class MooncakeConnectorWorker:

1495 remote_engine_id: EngineId,

1496 pull_metas: dict[ReqId, PullReqMeta],

1497 ):

1498 - remote_tp_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id(

1499 - remote_engine_id

1498 + remote_tp_ranks = self.transfer_topo.handshake_target_ranks(

1499 + self._tp_size[remote_engine_id]

1500 )

1501 count = len(remote_tp_ranks)

1502 logger.debug(

@@ -1587,7 +1587,7 @@ class MooncakeConnectorWorker:

1587 )

1588

1589 def _producer_cache_is_replicated(self) -> bool:

1590 - return self.kv_topo.replicates_kv_cache(self.engine_id)

1590 + return self.transfer_topo.local_replicates_kv_cache

1591

1592 def _get_transfer_regions(

1593 self, base_addrs: list[int], block_lens: list[int]

@@ -1595,7 +1595,7 @@ class MooncakeConnectorWorker:

1595 return _expand_transfer_regions(

1596 base_addrs=base_addrs,

1597 block_lens=block_lens,

1598 - is_kv_layout_blocks_first=self.kv_topo.is_kv_layout_blocks_first,

1598 + is_kv_layout_blocks_first=self.transfer_topo.is_kv_layout_blocks_first,

1599 )

1600

1601 def _get_sender_transfer_plan(

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

+118 -121

@@ -21,8 +21,8 @@ from vllm import envs

21 from vllm.distributed.kv_transfer.kv_connector.utils import (

22 BlockIds,

23 EngineId,

24 - HeteroTPTransferConfig,

25 - TpKVTopology,

24 + MambaEngineTransferInfo,

25 + TransferTopology,

26 get_current_attn_backends,

27 kv_postprocess_blksize_and_layout_on_receive,

28 kv_postprocess_blksize_on_receive,

@@ -49,7 +49,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (

49 )

50 from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (

51 MambaConvSplitInfo,

52 - compute_mamba_phys_ratio,

52 + compute_physical_blocks_per_logical,

53 derive_mamba_conv_split,

54 )

55 from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config

@@ -269,14 +269,12 @@ class NixlConnectorWorker:

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

270

271 # ---- Mamba-HMA per-engine state (only used when self._has_mamba) ----

272 - # Per-engine transfer config (source of truth for FA/mamba sizing).

273 - self._transfer_configs: dict[str, HeteroTPTransferConfig] = {}

274 - # NOTE (ZhanqiuHu): _mamba_phys_ratio MUST be per-engine.

275 - # compute_mamba_phys_ratio = ceil((conv_bytes + ssm_bytes) / block_len)

272 + # NOTE (ZhanqiuHu): _physical_blocks_per_logical MUST be per-engine.

273 + # physical_blocks_per_logical = ceil((conv_bytes + ssm_bytes) / block_len)

274 # where conv/ssm bytes are per-TP-rank (dimension-sharded). With

275 # heterogeneous TP the per-rank sizes differ, so the ratio differs:

276 # e.g. Nemotron 30B: P(TP=4) → 131, D(TP=1) → 261.

279 - self._mamba_phys_ratio: dict[EngineId, int] = {}

277 + self._physical_blocks_per_logical: dict[EngineId, int] = {}

278

279 # In progress transfers.

280 # [req_id -> list[handle]]

@@ -322,10 +320,8 @@ class NixlConnectorWorker:

320

321 # lazy initialized in register_kv_caches

322 self.compat_hash: str | None = None

325 - self.kv_topo: TpKVTopology | None = None

323 + self.transfer_topo: TransferTopology | None = None

324

327 - self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size}

328 - self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size}

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

326 # finish reading before safely freeing the blocks.

327 self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)

@@ -355,7 +351,6 @@ class NixlConnectorWorker:

351 self.block_size // kernel_block_size

352 )

353 self.block_size = kernel_block_size

358 - self._block_size[self.engine_id] = kernel_block_size

354 self.num_blocks *= self._physical_blocks_per_logical_kv_block

355

356 def _nixl_handshake(

@@ -384,8 +379,8 @@ class NixlConnectorWorker:

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

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

381 # this happens to be the same single rank_i.

387 - assert self.kv_topo is not None

388 - p_remote_ranks = self.kv_topo.get_target_remote_ranks(remote_tp_size)

382 + assert self.transfer_topo is not None

383 + p_remote_ranks = self.transfer_topo.handshake_target_ranks(remote_tp_size)

384 remote_rank_to_agent_name = {}

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

386

@@ -649,11 +644,11 @@ class NixlConnectorWorker:

644

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

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

652 - self.kv_topo = TpKVTopology(

647 + self.transfer_topo = TransferTopology(

648 tp_rank=self.tp_rank,

649 + tp_size=self.world_size,

650 + block_size=self.block_size,

651 engine_id=self.engine_id,

655 - remote_tp_size=self._tp_size, # shared state

656 - remote_block_size=self._block_size, # shared state

652 is_mla=self.use_mla,

653 total_num_kv_heads=self.model_config.get_total_num_kv_heads(),

654 attn_backends=self.attn_backends,

@@ -664,7 +659,7 @@ class NixlConnectorWorker:

659 is_mamba=self._has_mamba,

660 )

661 self.compat_hash = compute_nixl_compatibility_hash(

667 - self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks

662 + self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks

663 )

664

665 if self.use_host_buffer:

@@ -716,7 +711,7 @@ class NixlConnectorWorker:

711 if isinstance(layer_spec, UniformTypeKVCacheSpecs):

712 # MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs

713 layer_spec = layer_spec.kv_cache_specs[layer_name]

719 - cache_list = self.kv_topo.get_transfer_cache_regions(

714 + cache_list = self.transfer_topo.get_transfer_cache_regions(

715 cache_or_caches, layer_spec

716 )

717 # `layer_spec.page_size_bytes` only accounts for logical page_size, that is

@@ -729,7 +724,7 @@ class NixlConnectorWorker:

724 )

725 # For when registering multiple tensors eg K/V in separate regions.

726 physical_page_size = physical_page_size // len(cache_list)

732 - if self.kv_topo._cross_layers_blocks:

727 + if self.transfer_topo._cross_layers_blocks:

728 # When cross-layers blocks are used, multiply by number of layers

729 physical_page_size = physical_page_size * len(

730 self.kv_cache_config.kv_cache_tensors

@@ -793,7 +788,7 @@ class NixlConnectorWorker:

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

789 self.num_regions = len(caches_data)

790

796 - if self.kv_topo.is_kv_layout_blocks_first:

791 + if self.transfer_topo.is_kv_layout_blocks_first:

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

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

794 # registerMem allowing faster descs queries. In order to be able to

@@ -817,7 +812,7 @@ class NixlConnectorWorker:

812 self.dst_num_blocks[self.engine_id] = self.num_blocks

813

814 if self._has_mamba:

820 - self._mamba_phys_ratio[self.engine_id] = (

815 + self._physical_blocks_per_logical[self.engine_id] = (

816 self._physical_blocks_per_logical_kv_block

817 )

818 logger.info(

@@ -876,11 +871,13 @@ class NixlConnectorWorker:

871 conv_offsets = self._conv_decomp.local_conv_offsets

872 conv_size, ssm_size = self._mamba_ssm_size

873 num_blocks = self._logical_num_blocks * block_size_ratio

879 - phys_ratio = self._physical_blocks_per_logical_kv_block

874 + physical_per_logical = self._physical_blocks_per_logical_kv_block

875

876 result: list[tuple[int, int, int]] = []

877 for i, base_addr in enumerate(base_addresses):

883 - page_stride = self.block_len_per_layer[i] // block_size_ratio * phys_ratio

878 + page_stride = (

879 + self.block_len_per_layer[i] // block_size_ratio * physical_per_logical

880 + )

881 for off, sz in conv_offsets:

882 for blk in range(num_blocks):

883 result.append(

@@ -900,14 +897,14 @@ class NixlConnectorWorker:

897 def _build_fa_remote_for_mamba(

898 self,

899 nixl_agent_meta: NixlAgentMetadata,

903 - transfer_cfg: HeteroTPTransferConfig,

900 block_size_ratio: int,

905 - kv_topo: TpKVTopology,

901 + transfer_topo: TransferTopology,

902 + remote_engine_id: EngineId,

903 ) -> list[tuple[int, int, int]]:

904 """Build remote FA descriptors for mamba models.

905

909 - Uses transfer_cfg for GQA-aware FA divisor and head-based rank offset

910 - instead of the standard uniform tp_ratio split.

906 + Uses TransferTopology for GQA-aware FA divisor and head-based rank

907 + offset instead of the standard uniform tp_ratio split.

908 """

909 assert block_size_ratio == 1, (

910 "Mamba 3-read transfer with block_size_ratio != 1 is not tested. "

@@ -915,7 +912,9 @@ class NixlConnectorWorker:

912 )

913 # TODO (ZhanqiuHu): unify with register_remote_blocks when Mamba-HMA

914 # hetero-TP logic stabilizes.

918 - tp_ratio = transfer_cfg.tp_ratio

915 + mamba_info = transfer_topo.get_engine_info(remote_engine_id)

916 + assert isinstance(mamba_info, MambaEngineTransferInfo)

917 + tp_ratio = transfer_topo.tp_ratio(mamba_info.remote_tp_size)

918 result: list[tuple[int, int, int]] = []

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

920 local_block_len = self.get_backend_aware_kv_block_len(

@@ -926,9 +925,11 @@ class NixlConnectorWorker:

925 local_block_len = remote_kv_block_len

926

927 if tp_ratio < 0 and not self.use_mla:

929 - local_block_len = local_block_len // transfer_cfg.physical_fa_num_reads

928 + local_block_len = local_block_len // mamba_info.remote_num_fa_reads

929

931 - rank_offset = transfer_cfg.fa_rank_offset(remote_kv_block_len)

930 + rank_offset = transfer_topo.fa_rank_offset(

931 + remote_engine_id, remote_kv_block_len

932 + )

933

934 num_blocks = nixl_agent_meta.num_blocks

935 page_size = nixl_agent_meta.block_lens[i]

@@ -937,12 +938,12 @@ class NixlConnectorWorker:

938 addr = base_addr + block_offset + rank_offset

939 result.append((addr, local_block_len, nixl_agent_meta.device_id))

940

940 - if kv_topo.is_kv_layout_blocks_first:

941 + if transfer_topo.is_kv_layout_blocks_first:

942 second_split = self.get_backend_aware_kv_block_len(

943 layer_idx=i, first_split=False, mamba_view=False

944 )

945 if tp_ratio < 0 and not self.use_mla:

945 - second_split = second_split // transfer_cfg.physical_fa_num_reads

946 + second_split = second_split // mamba_info.remote_num_fa_reads

947 for block_id in range(num_blocks):

948 block_offset = block_id * page_size

949 addr = base_addr + block_offset + rank_offset

@@ -981,15 +982,17 @@ class NixlConnectorWorker:

982 conv_offsets = [(0, xb_p), (xb_p, bb_p), (xb_p + bb_p, bb_p)]

983 ssm_read_size = nixl_agent_meta.ssm_sizes[1]

984

984 - remote_ratio = self._mamba_phys_ratio[nixl_agent_meta.engine_id]

985 - num_blocks = nixl_agent_meta.num_blocks // remote_ratio

985 + remote_physical_per_logical = self._physical_blocks_per_logical[

986 + nixl_agent_meta.engine_id

987 + ]

988 + num_blocks = nixl_agent_meta.num_blocks // remote_physical_per_logical

989 device_id = nixl_agent_meta.device_id

990

991 result: list[tuple[int, int, int]] = []

992 # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case

993 # block lengths vary across layers (e.g. MLA).

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

992 - page_stride = nixl_agent_meta.block_lens[i] * remote_ratio

995 + page_stride = nixl_agent_meta.block_lens[i] * remote_physical_per_logical

996 for off, sz in conv_offsets:

997 for blk in range(num_blocks):

998 result.append((base_addr + blk * page_stride + off, sz, device_id))

@@ -1019,8 +1022,8 @@ class NixlConnectorWorker:

1022 register another local_xfer_handler using remote block len to ensure

1023 data copy correctness.

1024 """

1022 - assert self.kv_topo is not None

1023 - kv_topo = self.kv_topo

1025 + assert self.transfer_topo is not None

1026 + transfer_topo = self.transfer_topo

1027

1028 block_size_ratio = self.block_size // block_size

1029 blocks_data: list[tuple[int, int, int]] = []

@@ -1051,7 +1054,7 @@ class NixlConnectorWorker:

1054 # (addr, len, device id)

1055 blocks_data.append((addr, kv_block_len, self.device_id))

1056

1054 - if kv_topo.is_kv_layout_blocks_first:

1057 + if transfer_topo.is_kv_layout_blocks_first:

1058 second_split = self.get_backend_aware_kv_block_len(

1059 layer_idx=i, first_split=False, mamba_view=mamba

1060 )

@@ -1153,11 +1156,29 @@ class NixlConnectorWorker:

1156 )

1157 return self._remote_agents[engine_id][remote_tp_rank]

1158

1156 - ### Register remote agent metadata

1157 - if engine_id not in self._tp_size:

1158 - self._tp_size[engine_id] = remote_tp_size

1159 - if engine_id not in self._block_size:

1160 - self._block_size[engine_id] = nixl_agent_meta.block_size

1159 + ### Register remote engine in TransferTopology (idempotent).

1160 + assert self.transfer_topo is not None

1161 + transfer_topo = self.transfer_topo

1162 + physical_blocks_per_logical = (

1163 + compute_physical_blocks_per_logical(

1164 + nixl_agent_meta.ssm_sizes,

1165 + nixl_agent_meta.block_lens[0],

1166 + )

1167 + if self._has_mamba

1168 + else 1

1169 + )

1170 + transfer_topo.register_remote_engine(

1171 + remote_engine_id=engine_id,

1172 + remote_tp_size=remote_tp_size,

1173 + remote_block_size=nixl_agent_meta.block_size,

1174 + remote_block_len=nixl_agent_meta.block_lens[0],

1175 + remote_physical_blocks_per_logical=physical_blocks_per_logical,

1176 + local_block_len=self.block_len_per_layer[0],

1177 + )

1178 + if self._has_mamba and engine_id not in self._physical_blocks_per_logical:

1179 + self._physical_blocks_per_logical[engine_id] = physical_blocks_per_logical

1180 +

1181 + logger.info("Transfer plan: %s", transfer_topo.describe(engine_id))

1182

1183 remote_agent_name = self.nixl_wrapper.add_remote_agent(

1184 nixl_agent_meta.agent_metadata

@@ -1170,16 +1191,10 @@ class NixlConnectorWorker:

1191 # remote: | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|

1192 # local origin:| 0| 1| 8| 12|

1193 # local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15|

1173 - assert self.kv_topo is not None

1174 - kv_topo = self.kv_topo

1175 - block_size_ratio = kv_topo.block_size_ratio_from_engine_id(engine_id)

1194 + block_size_ratio = transfer_topo.block_size_ratio(nixl_agent_meta.block_size)

1195

1196 if engine_id not in self.dst_num_blocks:

1197 self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks

1179 - if self._has_mamba:

1180 - self._mamba_phys_ratio[engine_id] = compute_mamba_phys_ratio(

1181 - nixl_agent_meta.ssm_sizes, nixl_agent_meta.block_lens[0]

1182 - )

1198

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

1200 self.kv_caches_base_addr[engine_id][remote_tp_rank] = (

@@ -1189,28 +1204,13 @@ class NixlConnectorWorker:

1204

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

1206 # this is the ratio between the two sizes.

1192 - tp_ratio = self.kv_topo.tp_ratio_from_engine_id(engine_id)

1207 + tp_ratio = transfer_topo.tp_ratio(remote_tp_size)

1208

1209 # Handle tp_size>num_kv_heads: replicate KV cache.

1210 indexes_into_remote = (

1196 - not self.kv_topo.replicates_kv_cache(engine_id) and tp_ratio > 0

1211 + not transfer_topo.replicates_kv_cache(engine_id) and tp_ratio > 0

1212 )

1213

1199 - # Create transfer config (single source of truth for descriptor sizes).

1200 - if self._has_mamba and engine_id not in self._transfer_configs:

1201 - self._transfer_configs[engine_id] = HeteroTPTransferConfig(

1202 - tp_ratio=tp_ratio,

1203 - K=kv_topo.total_num_kv_heads,

1204 - d_tp=self.world_size,

1205 - p_tp=remote_tp_size,

1206 - d_rank=self.tp_rank,

1207 - use_mla=self.use_mla,

1208 - d_block_len=self.block_len_per_layer[0],

1209 - p_block_len=nixl_agent_meta.block_lens[0],

1210 - is_blocks_first=kv_topo.is_kv_layout_blocks_first,

1211 - )

1212 - logger.info("Created %s", self._transfer_configs[engine_id].describe())

1213 -

1214 logger.debug(

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

1216 engine_id,

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

1231 self.src_xfer_handles_by_tp_ratio[tp_ratio] = []

1232

1233 if self._has_mamba:

1234 - transfer_cfg = self._transfer_configs.get(engine_id)

1235 - assert transfer_cfg is not None

1236 - if transfer_cfg.needs_split_handles:

1234 + if transfer_topo.needs_split_handles(engine_id):

1235 # Mamba-HMA: FA and Mamba use different split factors.

1238 - for handle_data in transfer_cfg.compute_split_handle_data(

1239 - self.src_blocks_data, self.num_descs, abs_tp

1236 + for handle_data in transfer_topo.compute_split_handle_data(

1237 + engine_id, self.src_blocks_data, self.num_descs, abs_tp

1238 ):

1239 descs = self.nixl_wrapper.get_xfer_descs(

1240 handle_data, self.nixl_memory_type

@@ -1247,12 +1245,8 @@ class NixlConnectorWorker:

1245 self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)

1246

1247 logger.info(

1250 - "Mamba-HMA split handles: targets=%s, fa_reads=%s, "

1251 - "fa_entry=%s, mamba_reads=%s, num_descs=%s",

1252 - transfer_cfg.transfer_targets,

1253 - transfer_cfg.physical_fa_num_reads,

1254 - transfer_cfg.fa_entry_size,

1255 - transfer_cfg.mamba_num_reads,

1248 + "Mamba-HMA split handles: %s, num_descs=%s",

1249 + transfer_topo.describe(engine_id),

1250 self.num_descs,

1251 )

1252 else:

@@ -1321,7 +1315,7 @@ class NixlConnectorWorker:

1315 (addr, local_block_len, nixl_agent_meta.device_id)

1316 )

1317

1324 - if kv_topo.is_kv_layout_blocks_first:

1318 + if transfer_topo.is_kv_layout_blocks_first:

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

1320 second_split = self.get_backend_aware_kv_block_len(

1321 layer_idx=i, first_split=False, mamba_view=mamba

@@ -1360,14 +1354,12 @@ class NixlConnectorWorker:

1354 engine_id,

1355 remote_tp_rank,

1356 )

1363 - transfer_cfg = self._transfer_configs.get(engine_id)

1364 - assert transfer_cfg is not None

1357 blocks_data.extend(

1358 self._build_fa_remote_for_mamba(

1359 nixl_agent_meta,

1368 - transfer_cfg,

1360 block_size_ratio,

1370 - kv_topo,

1361 + transfer_topo,

1362 + engine_id,

1363 )

1364 )

1365 blocks_data.extend(

@@ -1403,18 +1395,19 @@ class NixlConnectorWorker:

1395 """

1396 remote_engine_id = nixl_agent_meta.engine_id

1397

1406 - assert self._tp_size[remote_engine_id] == remote_tp_size

1407 - assert self.kv_topo is not None

1398 + assert self.transfer_topo is not None

1399 + remote_info = self.transfer_topo.get_engine_info(remote_engine_id)

1400 + assert remote_info.remote_tp_size == remote_tp_size

1401

1409 - tp_ratio = self.kv_topo.tp_ratio_from_engine_id(remote_engine_id)

1410 - block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(

1411 - remote_engine_id

1402 + tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size)

1403 + block_size_ratio = self.transfer_topo.block_size_ratio(

1404 + nixl_agent_meta.block_size

1405 )

1406 # num_kv_heads > tp_size with P_TP > D_TP not supported for non-mamba.

1407 # Mamba models can have replicated FA KV with tp_ratio < 0.

1408 if not self._has_mamba:

1409 assert not (

1417 - tp_ratio < 0 and self.kv_topo.is_kv_replicated(remote_engine_id)

1410 + tp_ratio < 0 and self.transfer_topo.is_kv_replicated(remote_engine_id)

1411 )

1412

1413 if self._is_hma_required:

@@ -1467,7 +1460,7 @@ class NixlConnectorWorker:

1460 if (

1461 abs(tp_ratio) != 1

1462 and not self.use_mla

1470 - and not self.kv_topo.is_kv_replicated(remote_engine_id)

1463 + and not self.transfer_topo.is_kv_replicated(remote_engine_id)

1464 and kv_cache_layout != "HND"

1465 and not self.enable_permute_local_kv

1466 ):

@@ -1478,7 +1471,7 @@ class NixlConnectorWorker:

1471

1472 # Block len can only vary across layers when using MLA.

1473 remote_block_len = nixl_agent_meta.block_lens[0]

1481 - if self.use_mla or self.kv_topo.is_kv_replicated(remote_engine_id):

1474 + if self.use_mla or self.transfer_topo.is_kv_replicated(remote_engine_id):

1475 # With replicated KV cache, only the number of blocks can differ.

1476 # TODO (ZhanqiuHu): For mamba models, validate FA and mamba

1477 # block_lens separately.

@@ -1594,7 +1587,7 @@ class NixlConnectorWorker:

1587 if len(self.device_kv_caches) == 0:

1588 return

1589 assert block_size_ratio >= 1, "Only nP < nD supported currently."

1597 - assert self.kv_topo is not None

1590 + assert self.transfer_topo is not None

1591 if self.enable_permute_local_kv and block_size_ratio > 1:

1592 logger.debug(

1593 "Post-processing device kv cache on receive by converting "

@@ -1614,7 +1607,7 @@ class NixlConnectorWorker:

1607 block_size_ratio,

1608 )

1609

1617 - split_k_and_v = self.kv_topo.split_k_and_v

1610 + split_k_and_v = self.transfer_topo.split_k_and_v

1611

1612 for block_ids in block_ids_list:

1613 indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)

@@ -1661,7 +1654,7 @@ class NixlConnectorWorker:

1654 The scheduler process (via the MultiprocExecutor) will use this output

1655 to track which workers are done.

1656 """

1664 - assert self.kv_topo is not None

1657 + assert self.transfer_topo is not None

1658 done_sending = self._get_new_notifs()

1659 done_recving = self._pop_done_transfers(self._recving_transfers)

1660

@@ -1689,8 +1682,9 @@ class NixlConnectorWorker:

1682 self.sync_recved_kv_to_device(req_id, meta)

1683

1684 # post processing for heteroblocksize

1692 - block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(

1693 - meta.remote.engine_id

1685 + remote_info = self.transfer_topo.get_engine_info(meta.remote.engine_id)

1686 + block_size_ratio = self.transfer_topo.block_size_ratio(

1687 + remote_info.remote_block_size

1688 )

1689 if not self.use_mla and (

1690 block_size_ratio > 1 or self.enable_permute_local_kv

@@ -1741,7 +1735,7 @@ class NixlConnectorWorker:

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

1736 for all consumers to be done pulling.

1737 """

1744 - assert self.kv_topo is not None

1738 + assert self.transfer_topo is not None

1739 notified_req_ids: set[str] = set()

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

1741 for notif in notifs:

@@ -1760,7 +1754,7 @@ class NixlConnectorWorker:

1754

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

1756 n_consumers = int(tp_size)

1763 - tp_ratio = self.kv_topo.tp_ratio(n_consumers)

1757 + tp_ratio = self.transfer_topo.tp_ratio(n_consumers)

1758

1759 # Number of reads *per producer* to wait for.

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

@@ -1901,17 +1895,17 @@ class NixlConnectorWorker:

1895 self._reqs_to_send[req_id] = expiration_time

1896

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

1904 - assert meta.remote is not None and self.kv_topo is not None

1905 - remote_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id(

1906 - meta.remote.engine_id

1907 - )

1908 - tp_ratio = self.kv_topo.tp_ratio_from_engine_id(meta.remote.engine_id)

1898 + assert meta.remote is not None and self.transfer_topo is not None

1899 + engine_id = meta.remote.engine_id

1900 + remote_ranks = self.transfer_topo.target_remote_ranks(engine_id)

1901 + remote_info = self.transfer_topo.get_engine_info(engine_id)

1902 + tp_ratio = self.transfer_topo.tp_ratio(remote_info.remote_tp_size)

1903

1904 if self._has_mamba:

1905 # Expand remote logical → kernel block IDs.

1906 meta.remote.block_ids = self._logical_to_remote_kernel_block_ids(

1907 meta.remote.block_ids,

1914 - self._mamba_phys_ratio[meta.remote.engine_id],

1908 + self._physical_blocks_per_logical[meta.remote.engine_id],

1909 )

1910 else:

1911 meta.remote.block_ids = self._logical_to_kernel_block_ids(

@@ -1924,7 +1918,7 @@ class NixlConnectorWorker:

1918 # the first remote rank (cache is duplicated)..

1919 break

1920

1927 - remote_block_size = self.kv_topo.remote_block_size[meta.remote.engine_id]

1921 + remote_block_size = remote_info.remote_block_size

1922 logger.debug(

1923 "Remote agent %s available, calling _read_blocks"

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

@@ -1955,9 +1949,8 @@ class NixlConnectorWorker:

1949 remote_ids: BlockIds = meta.remote.block_ids

1950 if self._has_mamba:

1951 # Mamba-HMA: zero out FA groups for P ranks outside fa_read_targets.

1958 - transfer_cfg = self._transfer_configs.get(meta.remote.engine_id)

1959 - assert transfer_cfg is not None

1960 - local_ids, remote_ids = transfer_cfg.filter_block_ids_for_rank(

1952 + local_ids, remote_ids = self.transfer_topo.filter_block_ids_for_rank(

1953 + engine_id,

1954 remote_rank,

1955 local_ids,

1956 remote_ids,

@@ -1999,8 +1992,11 @@ class NixlConnectorWorker:

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

1993 a single remote worker.

1994 """

2002 - assert self.kv_topo is not None

2003 - block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(dst_engine_id)

1995 + assert self.transfer_topo is not None

1996 + remote_info = self.transfer_topo.get_engine_info(dst_engine_id)

1997 + block_size_ratio = self.transfer_topo.block_size_ratio(

1998 + remote_info.remote_block_size

1999 + )

2000 if block_size_ratio > 1:

2001 # TODO (NickLucche) assume HMA is off. Change to handle multiple KV groups.

2002 assert not self._is_hma_required

@@ -2190,8 +2186,8 @@ class NixlConnectorWorker:

2186 # This is like having two "low-level views" of the same storage.

2187 # `num_fa_descs` offset must be computed per-engine since P and D can

2188 # have different num_blocks (and thus different FA descs counts).

2193 - ratio = self._mamba_phys_ratio[engine_id]

2194 - logical_blocks = num_blocks // ratio

2189 + physical_per_logical = self._physical_blocks_per_logical[engine_id]

2190 + logical_blocks = num_blocks // physical_per_logical

2191 num_fa_descs = self.num_regions * num_blocks

2192 # 3-read mamba: 4 regions per unique cache tensor (x, B, C, ssm).

2193 mamba_region_ids = np.arange(len(self.block_len_per_layer) * 4)[:, None]

@@ -2234,21 +2230,22 @@ class NixlConnectorWorker:

2230 ]

2231

2232 def _logical_to_remote_kernel_block_ids(

2237 - self, block_ids: BlockIds, remote_ratio: int

2233 + self, block_ids: BlockIds, remote_physical_per_logical: int

2234 ) -> BlockIds:

2235 """Map logical block IDs to physical kernel block IDs on the remote.

2236

2237 Args:

2238 block_ids: per-group lists of logical block IDs.

2243 - remote_ratio: remote engine's physical blocks per logical block.

2239 + remote_physical_per_logical: remote engine's physical blocks

2240 + per logical block.

2241

2242 Returns:

2243 Same structure with FA groups expanded (each logical block L

2247 - becomes kernel blocks [L*remote_ratio .. L*remote_ratio +

2248 - local_ratio - 1]). Mamba groups are passed through unchanged.

2244 + becomes kernel blocks [L*ratio .. L*ratio + local_ratio - 1]).

2245 + Mamba groups are passed through unchanged.

2246 """

2247 local_ratio = self._physical_blocks_per_logical_kv_block

2251 - if remote_ratio == 1:

2248 + if remote_physical_per_logical == 1:

2249 return block_ids

2250 local_arange = np.arange(local_ratio).reshape(1, -1)

2251 group_specs = self.kv_cache_config.kv_cache_groups

@@ -2256,7 +2253,7 @@ class NixlConnectorWorker:

2253 for i, group in enumerate(block_ids):

2254 if not isinstance(group_specs[i].kv_cache_spec, MambaSpec):

2255 arr = np.array(group).reshape(-1, 1)

2259 - expanded = (arr * remote_ratio + local_arange).flatten()

2256 + expanded = (arr * remote_physical_per_logical + local_arange).flatten()

2257 result.append(expanded.tolist())

2258 else:

2259 # Mamba blocks are 1:1 logical-to-physical (no expansion).

@@ -2296,8 +2293,8 @@ class NixlConnectorWorker:

2293 +-------------------+ +--------------------+

2294 |1st_split-2nd_split| |1st_split-2nd_split |

2295 """

2299 - assert self.kv_topo is not None

2300 - if self.kv_topo.is_kv_layout_blocks_first:

2296 + assert self.transfer_topo is not None

2297 + if self.transfer_topo.is_kv_layout_blocks_first:

2298 # For indexing only half (either just the K or V part).

2299 if mamba_view:

2300 # NOTE (NickLucche) Mamba Opt: this is already skipping the padding so

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

+3 -1

@@ -151,7 +151,9 @@ def derive_mamba_conv_split(

151 )

152

153

154 -def compute_mamba_phys_ratio(ssm_sizes: tuple[int, ...], block_len: int) -> int:

154 +def compute_physical_blocks_per_logical(

155 + ssm_sizes: tuple[int, ...], block_len: int

156 +) -> int:

157 """Derive _physical_blocks_per_logical_kv_block from remote metadata.

158

159 The remote engine's ratio is not sent directly in the handshake, so we