@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / bfdc0a3a99
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 5, 12:15 UTC

[NIXL][Mamba][3/N] Heterogeneous TP: 3-read conv state transfer (#37635)

zhanqiuhu committed Apr 6, 2026 at 13:07 UTC bfdc0a3a996f82f592f43ff803adad173d3fbe11
5 files changed +970 -75
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
+2 -2
@@ -19,9 +19,9 @@ dp_ep_configs=(
19 "DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1)
20 )
21 hybrid_ssm_configs=(
22 - "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code"
22 + "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code"
23 # TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models.
24 - "ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"
24 + "VLLM_SSM_CONV_STATE_LAYOUT=DS ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=ibm-granite/granite-4.0-h-tiny VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"
25 )
26 sw_attn_configs=(
27 "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=google/gemma-3-4b-it PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192"
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+39 -6
@@ -224,6 +224,8 @@ def test_get_block_descs_ids_hybrid_ssm():
224 worker._has_mamba = True
225 worker._is_mamba_group = [False, True]
226 worker._physical_blocks_per_logical_kv_block = 1
227 + worker._mamba_phys_ratio = {engine_id: 1}
228 + worker.block_len_per_layer = [100]
229 # num_descs = num_regions * num_blocks (no blocks_first doubling)
230 worker.num_descs = 2 * num_blocks
231
@@ -234,9 +236,10 @@ def test_get_block_descs_ids_hybrid_ssm():
236 # FA group: stride=num_blocks=100, offset=0
237 # region0: [3, 5], region1: [103, 105]
238 # SSM group: stride=logical_blocks=100 (=num_blocks/ratio=100/1),
237 - # offset=num_descs=200
238 - # region0: [201, 202], region1: [301, 302]
239 - expected = [3, 5, 103, 105, 201, 202, 301, 302]
239 + # offset=num_fa_descs=200, 4 regions per Mamba layer (x, B, C, ssm)
240 + # region0: [201, 202], region1: [301, 302],
241 + # region2: [401, 402], region3: [501, 502]
242 + expected = [3, 5, 103, 105, 201, 202, 301, 302, 401, 402, 501, 502]
243 assert list(result) == expected, f"Expected {expected}, got {list(result)}"
244
245
@@ -259,6 +262,8 @@ def test_get_block_descs_ids_kernel_block_mismatch():
262 worker._has_mamba = True
263 worker._is_mamba_group = [False, True]
264 worker._physical_blocks_per_logical_kv_block = ratio
265 + worker._mamba_phys_ratio = {engine_id: ratio}
266 + worker.block_len_per_layer = [100]
267 worker.num_descs = 2 * num_blocks # 800
268
269 fa_blocks = [3, 7] # kernel-level block IDs
@@ -267,9 +272,11 @@ def test_get_block_descs_ids_kernel_block_mismatch():
272
273 # FA group: stride=num_blocks=400, offset=0
274 # region0: [3, 7], region1: [403, 407]
270 - # SSM group: stride=logical_blocks=400//4=100, offset=num_descs=800
271 - # region0: [801, 802], region1: [901, 902]
272 - expected = [3, 7, 403, 407, 801, 802, 901, 902]
275 + # SSM group: stride=logical_blocks=400//4=100, offset=num_fa_descs=800,
276 + # 4 regions per Mamba layer (x, B, C, ssm)
277 + # region0: [801, 802], region1: [901, 902],
278 + # region2: [1001, 1002], region3: [1101, 1102]
279 + expected = [3, 7, 403, 407, 801, 802, 901, 902, 1001, 1002, 1101, 1102]
280 assert list(result) == expected, f"Expected {expected}, got {list(result)}"
281
282
@@ -418,3 +425,29 @@ def test_has_mamba_init(
425 )
426 assert scheduler._has_mamba is expected_has_mamba
427 assert scheduler._is_hma_required is expected_is_hma
428 +
429 +
430 +@pytest.mark.cpu_test
431 +@pytest.mark.parametrize(
432 + "ssm_sizes,block_len,expected_ratio",
433 + [
434 + # Nemotron 30B TP=1: ceil((36864 + 2097152) / 8192) = 261
435 + ((36864, 2097152), 8192, 261),
436 + # Nemotron 30B TP=2: ceil((18432 + 1048576) / 4096) = 261
437 + ((18432, 1048576), 4096, 261),
438 + # Nemotron 30B TP=4: ceil((9216 + 524288) / 4096) = 131
439 + ((9216, 524288), 4096, 131),
440 + ],
441 +)
442 +def test_compute_mamba_phys_ratio(ssm_sizes, block_len, expected_ratio):
443 + """Verify that compute_mamba_phys_ratio is TP-dependent.
444 +
445 + With dimension-sharded Mamba state, the ratio differs across TP sizes
446 + (e.g. TP=1 → 261, TP=4 → 131 for Nemotron 30B). This is why
447 + _mamba_phys_ratio must be stored per-engine.
448 + """
449 + from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
450 + compute_mamba_phys_ratio,
451 + )
452 +
453 + assert compute_mamba_phys_ratio(ssm_sizes, block_len) == expected_ratio
vllm/distributed/kv_transfer/kv_connector/utils.py
+380 -1
@@ -5,7 +5,7 @@ KV cache helper for store.
5 """
6
7 from collections.abc import Iterator
8 -from dataclasses import dataclass
8 +from dataclasses import dataclass, field
9 from typing import TYPE_CHECKING, Any, Literal, cast
10
11 import torch
@@ -516,6 +516,338 @@ class TpKVTopology:
516 return cache if self.split_k_and_v else [cache]
517
518
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 +
550 +
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 + )
730 +
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.
743 + """
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.
762 + """
763 + if self.use_mla or self.tp_ratio <= 0:
764 + 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.
775 +
776 + True when FA and mamba have different read counts, requiring
777 + different splitting factors in the local handle.
778 + """
779 + return self.tp_ratio < 0 and not self.use_mla and len(self.transfer_targets) > 1
780 +
781 + def compute_split_handle_data(
782 + self,
783 + src_blocks_data: list[tuple[int, int, int]],
784 + num_fa_descs: int,
785 + abs_tp: int,
786 + ) -> list[list[tuple[int, int, int]]]:
787 + """Compute per-P-rank (addr, len, tp) triples for Mamba-HMA split handles.
788 +
789 + FA descriptors (indices < num_fa_descs) are sliced by
790 + ``physical_fa_num_reads``; mamba descriptors are sliced uniformly
791 + by ``abs_tp``.
792 +
793 + Returns one list of triples per transfer target.
794 + """
795 + all_handle_data: list[list[tuple[int, int, int]]] = []
796 + for p_idx, p_rank in enumerate(self.transfer_targets):
797 + 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):
802 + 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))
806 + else:
807 + mamba_chunk = local_len // abs_tp
808 + handle_data.append((addr + p_idx * mamba_chunk, mamba_chunk, tp))
809 + all_handle_data.append(handle_data)
810 + return all_handle_data
811 +
812 + def filter_block_ids_for_rank(
813 + self,
814 + remote_rank: int,
815 + local_ids: BlockIds,
816 + remote_ids: BlockIds,
817 + is_mamba_group: list[bool],
818 + ) -> tuple[BlockIds, BlockIds]:
819 + """Zero out FA groups for P ranks outside fa_read_targets.
820 +
821 + Returns (filtered_local_ids, filtered_remote_ids). When the
822 + remote rank carries FA data for this D rank, returns the inputs
823 + unchanged.
824 + """
825 + if not self.should_skip_fa(remote_rank):
826 + return local_ids, remote_ids
827 + num_groups = len(local_ids)
828 + filtered_local: list[list[int]] = [
829 + [] if not is_mamba_group[g] else local_ids[g] for g in range(num_groups)
830 + ]
831 + filtered_remote: list[list[int]] = [
832 + [] if not is_mamba_group[g] else remote_ids[g] for g in range(num_groups)
833 + ]
834 + return filtered_local, filtered_remote
835 +
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})"
848 + )
849 +
850 +
851 def get_current_attn_backends(
852 vllm_config: VllmConfig, layer_names: list[str] | None = None
853 ) -> list[type[AttentionBackend]]:
@@ -559,3 +891,50 @@ def get_current_attn_backend(
891 ) -> type[AttentionBackend]:
892 """Get the first attention backend for the given layers."""
893 return get_current_attn_backends(vllm_config, layer_names)[0]
894 +
895 +
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, ...) -> ...: ...
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+385 -66
@@ -25,6 +25,7 @@ from vllm.config import VllmConfig
25 from vllm.distributed.kv_transfer.kv_connector.utils import (
26 BlockIds,
27 EngineId,
28 + HeteroTPTransferConfig,
29 TpKVTopology,
30 get_current_attn_backend,
31 get_current_attn_backends,
@@ -47,12 +48,18 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
48 PromMetric,
49 PromMetricT,
50 )
51 +from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import (
52 + MambaConvSplitInfo,
53 + compute_mamba_phys_ratio,
54 + derive_mamba_conv_split,
55 +)
56 from vllm.distributed.parallel_state import (
57 get_tensor_model_parallel_rank,
58 get_tensor_model_parallel_world_size,
59 )
60 from vllm.forward_context import ForwardContext
61 from vllm.logger import init_logger
62 +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
63 from vllm.platforms import current_platform
64 from vllm.utils.math_utils import cdiv
65 from vllm.utils.network_utils import make_zmq_path, make_zmq_socket
@@ -1038,7 +1045,7 @@ class NixlConnectorWorker:
1045 }
1046 self.hma_group_size = len(kv_cache_config.kv_cache_tensors)
1047
1041 - # Mamba metadata
1048 + # ---- Mamba model state (derived from model config) ----
1049 self._is_mamba_group = [
1050 isinstance(group.kv_cache_spec, MambaSpec)
1051 for group in kv_cache_config.kv_cache_groups
@@ -1065,6 +1072,17 @@ class NixlConnectorWorker:
1072 ssm_shape.numel() * ssm_nbytes,
1073 )
1074 self._mamba_ssm_size = mamba_ssm_size
1075 + # Conv state sub-projection decomposition (None when no Mamba).
1076 + # The 3-read transfer requires DS (dim, state_len) conv layout so
1077 + # that x/B/C sub-projections are contiguous in memory.
1078 + self._conv_decomp: MambaConvSplitInfo | None = None
1079 + if self._has_mamba:
1080 + assert is_conv_state_dim_first(), (
1081 + "3-read Mamba conv transfer requires DS conv state layout. "
1082 + "Set VLLM_SSM_CONV_STATE_LAYOUT=DS"
1083 + )
1084 + local_tp = vllm_config.parallel_config.tensor_parallel_size
1085 + self._conv_decomp = derive_mamba_conv_split(mamba_spec, local_tp)
1086
1087 # Agent.
1088 non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"]
@@ -1175,6 +1193,16 @@ class NixlConnectorWorker:
1193 self.dst_num_blocks: dict[EngineId, int] = {}
1194 self._registered_descs: list[Any] = []
1195
1196 + # ---- Mamba-HMA per-engine state (only used when self._has_mamba) ----
1197 + # Per-engine transfer config (source of truth for FA/mamba sizing).
1198 + self._transfer_configs: dict[str, HeteroTPTransferConfig] = {}
1199 + # NOTE (ZhanqiuHu): _mamba_phys_ratio MUST be per-engine.
1200 + # compute_mamba_phys_ratio = ceil((conv_bytes + ssm_bytes) / block_len)
1201 + # where conv/ssm bytes are per-TP-rank (dimension-sharded). With
1202 + # heterogeneous TP the per-rank sizes differ, so the ratio differs:
1203 + # e.g. Nemotron 30B: P(TP=4) → 131, D(TP=1) → 261.
1204 + self._mamba_phys_ratio: dict[EngineId, int] = {}
1205 +
1206 # In progress transfers.
1207 # [req_id -> list[handle]]
1208 self._recving_metadata: dict[ReqId, ReqMeta] = {}
@@ -1701,8 +1729,7 @@ class NixlConnectorWorker:
1729 # then duplicate it logically to be able to index SSM/Conv separately.
1730 self.num_regions *= 2
1731
1704 - # TODO (NickLucche) Adapt to different descs views (engine_id->tp_rank) to
1705 - # support heterogeneous TP.
1732 + # Total local FA descriptors (boundary between FA and mamba descs).
1733 self.num_descs = self.num_regions * self.num_blocks
1734
1735 descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type)
@@ -1715,6 +1742,9 @@ class NixlConnectorWorker:
1742 self.dst_num_blocks[self.engine_id] = self.num_blocks
1743
1744 if self._has_mamba:
1745 + self._mamba_phys_ratio[self.engine_id] = (
1746 + self._physical_blocks_per_logical_kv_block
1747 + )
1748 logger.info(
1749 "Hybrid SSM registration: num_blocks=%s, "
1750 "logical_num_blocks=%s, ratio=%s, num_regions=%s, "
@@ -1755,6 +1785,149 @@ class NixlConnectorWorker:
1785 agent_metadata_bytes=encoder.encode(agent_metadata),
1786 )
1787
1788 + def _build_mamba_local(
1789 + self,
1790 + base_addresses: list[int],
1791 + block_size_ratio: int,
1792 + ) -> list[tuple[int, int, int]]:
1793 + """Build 4 desc regions (x, B, C, ssm) per layer for local mamba
1794 + blocks, enabling the 3-read transfer with DS conv layout."""
1795 + assert block_size_ratio == 1, (
1796 + "Mamba 3-read transfer with block_size_ratio != 1 is not tested. "
1797 + f"Got block_size_ratio={block_size_ratio}."
1798 + )
1799 + assert self._conv_decomp is not None
1800 + conv_offsets = self._conv_decomp.local_conv_offsets
1801 + conv_size, ssm_size = self._mamba_ssm_size
1802 + num_blocks = self._logical_num_blocks * block_size_ratio
1803 + phys_ratio = self._physical_blocks_per_logical_kv_block
1804 +
1805 + result: list[tuple[int, int, int]] = []
1806 + for i, base_addr in enumerate(base_addresses):
1807 + page_stride = self.block_len_per_layer[i] // block_size_ratio * phys_ratio
1808 + for off, sz in conv_offsets:
1809 + for blk in range(num_blocks):
1810 + result.append(
1811 + (base_addr + blk * page_stride + off, sz, self.device_id)
1812 + )
1813 + # SSM temporal state follows the conv state.
1814 + for blk in range(num_blocks):
1815 + result.append(
1816 + (
1817 + base_addr + blk * page_stride + conv_size,
1818 + ssm_size,
1819 + self.device_id,
1820 + )
1821 + )
1822 + return result
1823 +
1824 + def _build_fa_remote_for_mamba(
1825 + self,
1826 + nixl_agent_meta: NixlAgentMetadata,
1827 + transfer_cfg: HeteroTPTransferConfig,
1828 + block_size_ratio: int,
1829 + kv_topo: TpKVTopology,
1830 + ) -> list[tuple[int, int, int]]:
1831 + """Build remote FA descriptors for mamba models.
1832 +
1833 + Uses transfer_cfg for GQA-aware FA divisor and head-based rank offset
1834 + instead of the standard uniform tp_ratio split.
1835 + """
1836 + assert block_size_ratio == 1, (
1837 + "Mamba 3-read transfer with block_size_ratio != 1 is not tested. "
1838 + f"Got block_size_ratio={block_size_ratio}."
1839 + )
1840 + # TODO (ZhanqiuHu): unify with register_remote_blocks when Mamba-HMA
1841 + # hetero-TP logic stabilizes.
1842 + tp_ratio = transfer_cfg.tp_ratio
1843 + result: list[tuple[int, int, int]] = []
1844 + for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
1845 + local_block_len = self.get_backend_aware_kv_block_len(
1846 + layer_idx=i, first_split=True, mamba_view=False
1847 + )
1848 + remote_kv_block_len = local_block_len // block_size_ratio
1849 + if block_size_ratio > 1:
1850 + local_block_len = remote_kv_block_len
1851 +
1852 + if tp_ratio < 0 and not self.use_mla:
1853 + local_block_len = local_block_len // transfer_cfg.physical_fa_num_reads
1854 +
1855 + rank_offset = transfer_cfg.fa_rank_offset(remote_kv_block_len)
1856 +
1857 + num_blocks = nixl_agent_meta.num_blocks
1858 + page_size = nixl_agent_meta.block_lens[i]
1859 + for block_id in range(num_blocks):
1860 + block_offset = block_id * page_size
1861 + addr = base_addr + block_offset + rank_offset
1862 + result.append((addr, local_block_len, nixl_agent_meta.device_id))
1863 +
1864 + if kv_topo.is_kv_layout_blocks_first:
1865 + second_split = self.get_backend_aware_kv_block_len(
1866 + layer_idx=i, first_split=False, mamba_view=False
1867 + )
1868 + if tp_ratio < 0 and not self.use_mla:
1869 + second_split = second_split // transfer_cfg.physical_fa_num_reads
1870 + for block_id in range(num_blocks):
1871 + block_offset = block_id * page_size
1872 + addr = base_addr + block_offset + rank_offset
1873 + v_addr = addr + nixl_agent_meta.block_lens[i] // 2
1874 + result.append((v_addr, second_split, nixl_agent_meta.device_id))
1875 + return result
1876 +
1877 + def _build_mamba_remote(
1878 + self,
1879 + nixl_agent_meta: NixlAgentMetadata,
1880 + tp_ratio: int,
1881 + ) -> list[tuple[int, int, int]]:
1882 + """Build 4 remote desc regions (x, B, C, ssm) per layer for
1883 + the 3-read transfer. For hetero-TP, each D rank reads only its
1884 + sub-projection slice from the P rank."""
1885 + assert self._conv_decomp is not None
1886 + effective_ratio = max(tp_ratio, 1)
1887 + # Mamba conv state is always TP-sharded, even when attention KV
1888 + # is replicated (num_kv_heads < tp_size).
1889 + local_offset = self.tp_rank % effective_ratio
1890 + conv_size_remote = nixl_agent_meta.ssm_sizes[0]
1891 +
1892 + if tp_ratio >= 1:
1893 + # D_TP >= P_TP: P page is larger, D reads its slice.
1894 + conv_offsets = self._conv_decomp.remote_conv_offsets(
1895 + local_offset, effective_ratio
1896 + )
1897 + ssm_read_size = self._mamba_ssm_size[1]
1898 + else:
1899 + # NOTE (ZhanqiuHu): tp_ratio < 0 means P_TP > D_TP, so P pages
1900 + # are smaller than D's. self._conv_decomp has D-sized dimensions,
1901 + # but we need P-sized offsets. Scale down by |tp_ratio|.
1902 + abs_ratio = -tp_ratio
1903 + xb_p = self._conv_decomp.x_bytes // abs_ratio
1904 + bb_p = self._conv_decomp.b_bytes // abs_ratio
1905 + conv_offsets = [(0, xb_p), (xb_p, bb_p), (xb_p + bb_p, bb_p)]
1906 + ssm_read_size = nixl_agent_meta.ssm_sizes[1]
1907 +
1908 + remote_ratio = self._mamba_phys_ratio[nixl_agent_meta.engine_id]
1909 + num_blocks = nixl_agent_meta.num_blocks // remote_ratio
1910 + device_id = nixl_agent_meta.device_id
1911 +
1912 + result: list[tuple[int, int, int]] = []
1913 + # NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case
1914 + # block lengths vary across layers (e.g. MLA).
1915 + for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
1916 + page_stride = nixl_agent_meta.block_lens[i] * remote_ratio
1917 + for off, sz in conv_offsets:
1918 + for blk in range(num_blocks):
1919 + result.append((base_addr + blk * page_stride + off, sz, device_id))
1920 + # SSM temporal state is also TP-sharded on the heads dimension.
1921 + for blk in range(num_blocks):
1922 + ssm_addr = (
1923 + base_addr
1924 + + blk * page_stride
1925 + + conv_size_remote
1926 + + local_offset * ssm_read_size
1927 + )
1928 + result.append((ssm_addr, ssm_read_size, device_id))
1929 + return result
1930 +
1931 def register_local_xfer_handler(
1932 self,
1933 block_size: int,
@@ -1823,13 +1996,22 @@ class NixlConnectorWorker:
1996 self.device_id,
1997 )
1998
1999 + # NOTE (ZhanqiuHu): mamba=True path in register_blocks is not used
2000 + # right now — we use _build_mamba_local instead for the 3-read
2001 + # approach. However, we might still need this as a fallback for homogeneous TP.
2002 register_blocks(blocks_data, mamba=False)
2003 if self._has_mamba:
2004 assert self.num_descs == len(blocks_data)
1829 - logger.debug(
1830 - "Registering additional %s local Mamba blocks", len(blocks_data)
2005 + # TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-read split is
2006 + # unnecessary — a single conv desc per block suffices. Consider
2007 + # adding a fast path that falls back to the standard 2-region
2008 + # registration (register_blocks mamba=True) when no hetero-TP
2009 + # remote has been seen. Currently we always register 4 regions
2010 + # because local descs are created before knowing the remote TP.
2011 + logger.debug("Registering local Mamba descriptors (4 regions/layer)")
2012 + blocks_data.extend(
2013 + self._build_mamba_local(local_base_addresses, block_size_ratio)
2014 )
1832 - register_blocks(blocks_data, mamba=True)
2015
2016 descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
2017 # NIXL_INIT_AGENT to be used for preparations of local descs.
@@ -1880,6 +2062,9 @@ class NixlConnectorWorker:
2062
2063 Regarding MLA case, the cache is replicated across TP workers so the rank_offset will just always be 0
2064 so that the whole cache is shared by "tp_ratio" D TP workers.
2065 +
2066 + For Mamba hetero-TP, both tp_ratio > 0 (D_TP > P_TP) and
2067 + tp_ratio < 0 (P_TP > D_TP) are supported by the 3-read transfer.
2068 """ # noqa: E501
2069 engine_id = nixl_agent_meta.engine_id
2070 # TODO re-evaluate refreshing for scaling/recovery
@@ -1915,6 +2100,10 @@ class NixlConnectorWorker:
2100
2101 if engine_id not in self.dst_num_blocks:
2102 self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks
2103 + if self._has_mamba:
2104 + self._mamba_phys_ratio[engine_id] = compute_mamba_phys_ratio(
2105 + nixl_agent_meta.ssm_sizes, nixl_agent_meta.block_lens[0]
2106 + )
2107
2108 # Keep track of remote agent kv caches base addresses.
2109 self.kv_caches_base_addr[engine_id][remote_tp_rank] = (
@@ -1931,6 +2120,21 @@ class NixlConnectorWorker:
2120 not self.kv_topo.replicates_kv_cache(engine_id) and tp_ratio > 0
2121 )
2122
2123 + # Create transfer config (single source of truth for descriptor sizes).
2124 + if self._has_mamba and engine_id not in self._transfer_configs:
2125 + self._transfer_configs[engine_id] = HeteroTPTransferConfig(
2126 + tp_ratio=tp_ratio,
2127 + K=kv_topo.total_num_kv_heads,
2128 + d_tp=self.world_size,
2129 + p_tp=remote_tp_size,
2130 + d_rank=self.tp_rank,
2131 + use_mla=self.use_mla,
2132 + d_block_len=self.block_len_per_layer[0],
2133 + p_block_len=nixl_agent_meta.block_lens[0],
2134 + is_blocks_first=kv_topo.is_kv_layout_blocks_first,
2135 + )
2136 + logger.info("Created %s", self._transfer_configs[engine_id].describe())
2137 +
2138 logger.debug(
2139 "Registering remote agent (%s, rank %s) memory regions with tp_ratio %s",
2140 engine_id,
@@ -1947,21 +2151,48 @@ class NixlConnectorWorker:
2151 # Remote tp_size > local tp_size: read from multiple remote ranks.
2152 # Logically "split" own regions into |tp_ratio| chunks. Mind that
2153 # we only do this once per remote tp_size (replica-friendly).
2154 + abs_tp = -tp_ratio
2155 self.src_xfer_handles_by_tp_ratio[tp_ratio] = []
1951 - for i in range(-tp_ratio):
1952 - blocks_data = []
1953 - for memory_region in self.src_blocks_data:
1954 - addr, local_block_len, own_tp_rank = memory_region
1955 - # Computing block len layer by layer allows for different
1956 - # block sizes to be used.
1957 - remote_block_len = local_block_len // (-tp_ratio)
1958 - addr = addr + i * remote_block_len
1959 - blocks_data.append((addr, remote_block_len, own_tp_rank))
1960 - descs = self.nixl_wrapper.get_xfer_descs(
1961 - blocks_data, self.nixl_memory_type
1962 - )
1963 - handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)
1964 - self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)
2156 +
2157 + if self._has_mamba:
2158 + transfer_cfg = self._transfer_configs.get(engine_id)
2159 + assert transfer_cfg is not None
2160 + if transfer_cfg.needs_split_handles:
2161 + # Mamba-HMA: FA and Mamba use different split factors.
2162 + for handle_data in transfer_cfg.compute_split_handle_data(
2163 + self.src_blocks_data, self.num_descs, abs_tp
2164 + ):
2165 + descs = self.nixl_wrapper.get_xfer_descs(
2166 + handle_data, self.nixl_memory_type
2167 + )
2168 + handle = self.nixl_wrapper.prep_xfer_dlist(
2169 + "NIXL_INIT_AGENT", descs
2170 + )
2171 + self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)
2172 +
2173 + logger.info(
2174 + "Mamba-HMA split handles: targets=%s, fa_reads=%s, "
2175 + "fa_entry=%s, mamba_reads=%s, num_descs=%s",
2176 + transfer_cfg.transfer_targets,
2177 + transfer_cfg.physical_fa_num_reads,
2178 + transfer_cfg.fa_entry_size,
2179 + transfer_cfg.mamba_num_reads,
2180 + self.num_descs,
2181 + )
2182 + else:
2183 + # Original path: uniform divide by abs_tp (non-Mamba-HMA).
2184 + for i in range(abs_tp):
2185 + blocks_data = []
2186 + for memory_region in self.src_blocks_data:
2187 + addr, local_block_len, own_tp_rank = memory_region
2188 + remote_block_len = local_block_len // abs_tp
2189 + addr = addr + i * remote_block_len
2190 + blocks_data.append((addr, remote_block_len, own_tp_rank))
2191 + descs = self.nixl_wrapper.get_xfer_descs(
2192 + blocks_data, self.nixl_memory_type
2193 + )
2194 + handle = self.nixl_wrapper.prep_xfer_dlist("NIXL_INIT_AGENT", descs)
2195 + self.src_xfer_handles_by_tp_ratio[tp_ratio].append(handle)
2196
2197 ### Register remote agent memory regions
2198 blocks_data = []
@@ -2044,13 +2275,33 @@ class NixlConnectorWorker:
2275 self.tp_rank,
2276 )
2277
2047 - register_remote_blocks(blocks_data, mamba=False)
2278 if self._has_mamba:
2049 - # Create extra descs for the Mamba "view" of the same KV cache tensors.
2279 + # Mamba-HMA: separate FA registration with GQA-aware sizing,
2280 + # plus mamba 3-read registration for the Mamba "view" of the
2281 + # same KV cache tensors.
2282 logger.debug(
2051 - "Registering additional %s remote Mamba blocks", len(blocks_data)
2283 + "Registering remote Mamba blocks for engine %s rank %s",
2284 + engine_id,
2285 + remote_tp_rank,
2286 + )
2287 + transfer_cfg = self._transfer_configs.get(engine_id)
2288 + assert transfer_cfg is not None
2289 + blocks_data.extend(
2290 + self._build_fa_remote_for_mamba(
2291 + nixl_agent_meta,
2292 + transfer_cfg,
2293 + block_size_ratio,
2294 + kv_topo,
2295 + )
2296 )
2053 - register_remote_blocks(blocks_data, mamba=True)
2297 + blocks_data.extend(
2298 + self._build_mamba_remote(
2299 + nixl_agent_meta,
2300 + tp_ratio,
2301 + )
2302 + )
2303 + else:
2304 + register_remote_blocks(blocks_data, mamba=False)
2305
2306 # Register with NIXL.
2307 descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)
@@ -2083,17 +2334,17 @@ class NixlConnectorWorker:
2334 block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(
2335 remote_engine_id
2336 )
2086 - # Num kv_heads > tp_size and P TP > D TP case, not supported
2087 - assert not (tp_ratio < 0 and self.kv_topo.is_kv_replicated(remote_engine_id))
2337 + # num_kv_heads > tp_size with P_TP > D_TP not supported for non-mamba.
2338 + # Mamba models can have replicated FA KV with tp_ratio < 0.
2339 + if not self._has_mamba:
2340 + assert not (
2341 + tp_ratio < 0 and self.kv_topo.is_kv_replicated(remote_engine_id)
2342 + )
2343
2344 if self._is_hma_required:
2345 assert block_size_ratio == 1, (
2346 "HMA does not support different remote block size yet"
2347 )
2093 - # Mamba additional constraints
2094 - if self._has_mamba:
2095 - assert tp_ratio == 1, "Mamba does not support heterogeneous TP yet"
2096 -
2348 kv_cache_layout = (
2349 self.kv_cache_layout
2350 if not self.use_host_buffer
@@ -2138,11 +2389,14 @@ class NixlConnectorWorker:
2389 remote_block_len = nixl_agent_meta.block_lens[0]
2390 if self.use_mla or self.kv_topo.is_kv_replicated(remote_engine_id):
2391 # With replicated KV cache, only the number of blocks can differ.
2141 - for i in range(len(self.block_len_per_layer)):
2142 - assert (
2143 - self.block_len_per_layer[i] // block_size_ratio
2144 - == nixl_agent_meta.block_lens[i]
2145 - ), "KV cache sizes must match between P and D when replicated"
2392 + # TODO (ZhanqiuHu): For mamba models, validate FA and mamba
2393 + # block_lens separately.
2394 + if not self._has_mamba:
2395 + for i in range(len(self.block_len_per_layer)):
2396 + assert (
2397 + self.block_len_per_layer[i] // block_size_ratio
2398 + == nixl_agent_meta.block_lens[i]
2399 + ), "KV cache sizes must match between P and D when replicated"
2400 else:
2401 # When MLA is not used, this is a list of the same block length
2402 for block_len in nixl_agent_meta.block_lens:
@@ -2150,25 +2404,31 @@ class NixlConnectorWorker:
2404 "All remote layers must have the same block size"
2405 )
2406
2153 - if tp_ratio > 0:
2154 - # Remote tp is smaller: remote block_len size is bigger
2155 - assert (
2156 - remote_block_len
2157 - == (self.block_len_per_layer[0] * tp_ratio) // block_size_ratio
2158 - ), (
2159 - "Remote P worker KV layer cache must be of shape [2, N, "
2160 - "local_kv_heads*tp_ratio, page_size, head_dim] and same dtype."
2161 - ) # noqa: E501
2162 - else:
2163 - assert block_size_ratio == 1, (
2164 - "Different local/remote block sizes are not supported when"
2165 - " P TP > D TP."
2166 - )
2167 - # Remote tp is bigger: remote block_len size is smaller
2168 - assert remote_block_len == self.block_len_per_layer[0] // (-tp_ratio), (
2169 - "Remote P worker KV layer cache must be of shape [2, N, "
2170 - "local_kv_heads/tp_ratio, page_size, head_dim] and same dtype."
2171 - ) # noqa: E501
2407 + # HMA hybrid models (mamba+attention) pad block_len to
2408 + # max(attn_page, mamba_page), so the linear tp_ratio scaling
2409 + # assumption only holds for pure-attention models.
2410 + if not self._has_mamba:
2411 + if tp_ratio > 0:
2412 + assert (
2413 + remote_block_len
2414 + == (self.block_len_per_layer[0] * tp_ratio) // block_size_ratio
2415 + ), (
2416 + "Remote P worker KV layer cache must be of shape [2, N,"
2417 + " local_kv_heads*tp_ratio, page_size, head_dim] and "
2418 + "same dtype."
2419 + )
2420 + else:
2421 + assert block_size_ratio == 1, (
2422 + "Different local/remote block sizes are not supported"
2423 + " when P TP > D TP."
2424 + )
2425 + assert remote_block_len == self.block_len_per_layer[0] // (
2426 + -tp_ratio
2427 + ), (
2428 + "Remote P worker KV layer cache must be of shape [2, N,"
2429 + " local_kv_heads/tp_ratio, page_size, head_dim] and "
2430 + "same dtype."
2431 + )
2432
2433 # TP workers that handhshake with same remote have same #blocks.
2434 assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks
@@ -2471,9 +2731,8 @@ class NixlConnectorWorker:
2731 meta.local_block_ids
2732 )
2733 assert meta.remote is not None
2474 - meta.remote.block_ids = self._logical_to_kernel_block_ids(
2475 - meta.remote.block_ids
2476 - )
2734 + # Remote block IDs are kept logical here; expanded in
2735 + # _read_blocks_for_req using the remote engine's phys ratio.
2736 remote_engine_id = meta.remote.engine_id
2737 logger.debug(
2738 "start_load_kv for request %s from remote engine %s. "
@@ -2525,6 +2784,13 @@ class NixlConnectorWorker:
2784 meta.remote.engine_id
2785 )
2786 tp_ratio = self.kv_topo.tp_ratio_from_engine_id(meta.remote.engine_id)
2787 +
2788 + if self._has_mamba:
2789 + # Expand remote logical → kernel block IDs.
2790 + meta.remote.block_ids = self._logical_to_remote_kernel_block_ids(
2791 + meta.remote.block_ids,
2792 + self._mamba_phys_ratio[meta.remote.engine_id],
2793 + )
2794 # D may have to perform multiple reads from different remote ranks.
2795 for i, remote_rank in enumerate(remote_ranks):
2796 if self.use_mla and tp_ratio < 0 and i > 0:
@@ -2558,12 +2824,26 @@ class NixlConnectorWorker:
2824 remote_xfer_side_handle = self.dst_xfer_side_handles[meta.remote.engine_id][
2825 remote_rank
2826 ]
2827 +
2828 + local_ids: BlockIds = meta.local_physical_block_ids
2829 + remote_ids: BlockIds = meta.remote.block_ids
2830 + if self._has_mamba:
2831 + # Mamba-HMA: zero out FA groups for P ranks outside fa_read_targets.
2832 + transfer_cfg = self._transfer_configs.get(meta.remote.engine_id)
2833 + assert transfer_cfg is not None
2834 + local_ids, remote_ids = transfer_cfg.filter_block_ids_for_rank(
2835 + remote_rank,
2836 + local_ids,
2837 + remote_ids,
2838 + self._is_mamba_group,
2839 + )
2840 +
2841 self._read_blocks(
2842 request_id=req_id,
2843 dst_engine_id=meta.remote.engine_id,
2844 remote_request_id=meta.remote.request_id,
2565 - local_block_ids=meta.local_physical_block_ids,
2566 - remote_block_ids=meta.remote.block_ids,
2845 + local_block_ids=local_ids,
2846 + remote_block_ids=remote_ids,
2847 remote_rank=remote_rank,
2848 local_xfer_side_handle=local_xfer_side_handle,
2849 remote_xfer_side_handle=remote_xfer_side_handle,
@@ -2663,9 +2943,12 @@ class NixlConnectorWorker:
2943 for i, remote_group in enumerate(remote_block_ids):
2944 num_remote_blocks = len(remote_group)
2945 num_local_blocks = len(local_block_ids[i])
2666 - assert num_local_blocks <= num_remote_blocks
2946 + if not self._is_mamba_group[i]:
2947 + assert num_local_blocks <= num_remote_blocks
2948 # Partial prefix cache hit: just read uncomputed blocks.
2668 - if num_local_blocks < num_remote_blocks:
2949 + # Skip mamba groups — their blocks represent full state (conv+ssm),
2950 + # not per-token data, so trimming would corrupt the transfer.
2951 + if num_local_blocks < num_remote_blocks and not self._is_mamba_group[i]:
2952 remote_block_ids[i] = remote_group[-num_local_blocks:]
2953
2954 # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from
@@ -2781,16 +3064,22 @@ class NixlConnectorWorker:
3064 # This is like having two "low-level views" of the same storage.
3065 # `num_fa_descs` offset must be computed per-engine since P and D can
3066 # have different num_blocks (and thus different FA descs counts).
2784 - ratio = self._physical_blocks_per_logical_kv_block
2785 - # SSM may register fewer num_blocks than FA
3067 + ratio = self._mamba_phys_ratio[engine_id]
3068 logical_blocks = num_blocks // ratio
3069 num_fa_descs = self.num_regions * num_blocks
3070 + # 3-read mamba: 4 regions per unique cache tensor (x, B, C, ssm).
3071 + mamba_region_ids = np.arange(len(self.block_len_per_layer) * 4)[:, None]
3072 all_descs = []
3073 for i, group in enumerate(block_ids):
2790 - stride = logical_blocks if self._is_mamba_group[i] else num_blocks
3074 group_arr = np.asarray(group)[None, :]
2792 - offset = num_fa_descs if self._is_mamba_group[i] else 0
2793 - all_descs.append((region_ids * stride + group_arr + offset).flatten())
3075 + if self._is_mamba_group[i]:
3076 + all_descs.append(
3077 + (
3078 + mamba_region_ids * logical_blocks + group_arr + num_fa_descs
3079 + ).flatten()
3080 + )
3081 + else:
3082 + all_descs.append((region_ids * num_blocks + group_arr).flatten())
3083 return np.concatenate(all_descs)
3084
3085 def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds:
@@ -2818,6 +3107,36 @@ class NixlConnectorWorker:
3107 for i, group in enumerate(block_ids)
3108 ]
3109
3110 + def _logical_to_remote_kernel_block_ids(
3111 + self, block_ids: BlockIds, remote_ratio: int
3112 + ) -> BlockIds:
3113 + """Map logical block IDs to physical kernel block IDs on the remote.
3114 +
3115 + Args:
3116 + block_ids: per-group lists of logical block IDs.
3117 + remote_ratio: remote engine's physical blocks per logical block.
3118 +
3119 + Returns:
3120 + Same structure with FA groups expanded (each logical block L
3121 + becomes kernel blocks [L*remote_ratio .. L*remote_ratio +
3122 + local_ratio - 1]). Mamba groups are passed through unchanged.
3123 + """
3124 + local_ratio = self._physical_blocks_per_logical_kv_block
3125 + if remote_ratio == 1:
3126 + return block_ids
3127 + local_arange = np.arange(local_ratio).reshape(1, -1)
3128 + group_specs = self.kv_cache_config.kv_cache_groups
3129 + result: list[list[int]] = []
3130 + for i, group in enumerate(block_ids):
3131 + if not isinstance(group_specs[i].kv_cache_spec, MambaSpec):
3132 + arr = np.array(group).reshape(-1, 1)
3133 + expanded = (arr * remote_ratio + local_arange).flatten()
3134 + result.append(expanded.tolist())
3135 + else:
3136 + # Mamba blocks are 1:1 logical-to-physical (no expansion).
3137 + result.append(group)
3138 + return result
3139 +
3140 def get_backend_aware_kv_block_len(
3141 self, layer_idx: int, first_split: bool = True, mamba_view: bool = False
3142 ) -> int:
vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py new
+164
@@ -0,0 +1,164 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Mamba conv-state sub-projection decomposition for the 3-read transfer.
4 +
5 +With DS conv state layout (dim, state_len), x/B/C sub-projections are
6 +contiguous in memory. Each D rank reads its x, B, C slices via 3
7 +separate RDMA transfers — no P-side permutation needed.
8 +"""
9 +
10 +import math
11 +from dataclasses import dataclass
12 +
13 +import torch
14 +
15 +from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
16 +from vllm.v1.kv_cache_interface import MambaSpec
17 +
18 +
19 +@dataclass(frozen=True)
20 +class MambaConvSplitInfo:
21 + """Per-rank byte sizes of x, B, C sub-projections in the Mamba conv state.
22 +
23 + Used by both P and D sides for NIXL descriptor registration.
24 + All fields are LOCAL to this engine's TP (already divided by TP size).
25 +
26 + DS memory layout within one page (contiguous in memory):
27 + |--- x (x_local * conv_rows) ---|- B (b_local * conv_rows) -|- C -|
28 + """
29 +
30 + conv_rows: int # conv_kernel - 1 (typically 3)
31 + x_local: int # intermediate_size / TP (columns for x)
32 + b_local: int # groups_ss / TP (columns for B; C is same size)
33 + conv_dtype_size: int # bytes per element (e.g. 2 for float16)
34 +
35 + @property
36 + def conv_dim_local(self) -> int:
37 + """Total conv columns per rank: x + B + C."""
38 + return self.x_local + 2 * self.b_local
39 +
40 + @property
41 + def x_bytes(self) -> int:
42 + """Byte size of the x sub-projection for one rank."""
43 + return self.x_local * self.conv_rows * self.conv_dtype_size
44 +
45 + @property
46 + def b_bytes(self) -> int:
47 + """Byte size of the B (or C) sub-projection for one rank."""
48 + return self.b_local * self.conv_rows * self.conv_dtype_size
49 +
50 + @property
51 + def local_conv_offsets(self) -> list[tuple[int, int]]:
52 + """(byte_offset, byte_size) of x, B, C within this engine's page.
53 +
54 + Used by both P and D for local descriptor registration.
55 + """
56 + xb = self.x_bytes
57 + bb = self.b_bytes
58 + return [(0, xb), (xb, bb), (xb + bb, bb)]
59 +
60 + def remote_conv_offsets(
61 + self, local_rank_offset: int, tp_ratio: int
62 + ) -> list[tuple[int, int]]:
63 + """(byte_offset, byte_size) of this D rank's x, B, C slice within
64 + one P page.
65 +
66 + Used by D side only, during remote descriptor registration.
67 +
68 + Args:
69 + local_rank_offset: which slice this D rank reads.
70 + tp_ratio > 0: tp_rank % tp_ratio (selects slice of P's page).
71 + tp_ratio < 0: always 0 (read P's full page).
72 + tp_ratio: effective ratio (>= 1 when D_TP > P_TP, 1 when
73 + P_TP > D_TP since each P rank is read in full).
74 + """
75 + xb = self.x_bytes
76 + bb = self.b_bytes
77 + xr = xb * tp_ratio # full remote x section in bytes
78 + br = bb * tp_ratio # full remote B section in bytes
79 + return [
80 + (local_rank_offset * xb, xb),
81 + (xr + local_rank_offset * bb, bb),
82 + (xr + br + local_rank_offset * bb, bb),
83 + ]
84 +
85 +
86 +def derive_mamba_conv_split(
87 + mamba_spec: MambaSpec,
88 + local_tp: int,
89 +) -> MambaConvSplitInfo:
90 + """Derive per-rank x/B/C byte sizes from a MambaSpec.
91 +
92 + Called once at init on both P and D. Decomposes the conv dimension
93 + (= intermediate_size + 2 * groups_ss) into its x, B, C parts.
94 +
95 + Args:
96 + mamba_spec: MambaSpec whose shapes are:
97 + shapes[0] = conv state: (conv_dim_local, conv_rows) in DS layout.
98 + shapes[1] = SSM temporal: (local_num_heads, head_dim).
99 + local_tp: this engine's tensor-parallel size.
100 +
101 + Returns:
102 + MambaConvSplitInfo with per-rank x_local, b_local, conv_rows, and
103 + conv_dtype_size.
104 + """
105 + if mamba_spec.mamba_type != "mamba2":
106 + raise NotImplementedError(
107 + f"3-read conv transfer only supports Mamba2 models, "
108 + f"got mamba_type={mamba_spec.mamba_type!r}. "
109 + f"Mamba1 SSM temporal shape is (intermediate_size // tp, state_size) "
110 + f"which cannot be used to reconstruct intermediate_size."
111 + )
112 +
113 + conv_shape = mamba_spec.shapes[0]
114 + assert len(conv_shape) == 2, f"Expected 2D conv state shape, got {conv_shape}"
115 +
116 + # NOTE (ZhanqiuHu): 3-read requires DS layout, which is already asserted
117 + # in nixl_connector __init__. Use it directly instead of heuristic detection.
118 + assert is_conv_state_dim_first(), "3-read requires DS conv state layout"
119 + local_conv_dim = conv_shape[0] # DS: (conv_dim_local, conv_rows)
120 + conv_rows = conv_shape[1]
121 +
122 + # NOTE (ZhanqiuHu): intermediate_size (= global x dim) is not stored
123 + # in MambaSpec, so we reconstruct it from the SSM temporal state shape:
124 + # shapes[1] = (local_num_heads, head_dim), already divided by TP.
125 + head_dim = mamba_spec.shapes[1][1]
126 + local_num_heads = mamba_spec.shapes[1][0]
127 + intermediate_size = local_num_heads * local_tp * head_dim
128 +
129 + # NOTE (ZhanqiuHu): global conv dim = intermediate_size + 2 * groups_ss,
130 + # where groups_ss is the B (= C) dimension. B and C are always the same
131 + # size, so we recover groups_ss from the remainder after subtracting x.
132 + remainder = local_conv_dim * local_tp - intermediate_size
133 + assert remainder > 0 and remainder % 2 == 0, (
134 + f"Conv dim ({local_conv_dim}*tp={local_tp}) doesn't decompose into "
135 + f"intermediate_size={intermediate_size} + 2*groups_ss. "
136 + f"remainder={remainder}"
137 + )
138 + groups_ss = remainder // 2
139 +
140 + conv_dtype_size = torch.tensor(
141 + [],
142 + dtype=mamba_spec.dtypes[0], # type: ignore[misc]
143 + ).element_size()
144 +
145 + # Divide by TP to get per-rank column counts.
146 + return MambaConvSplitInfo(
147 + conv_rows=conv_rows,
148 + x_local=intermediate_size // local_tp,
149 + b_local=groups_ss // local_tp,
150 + conv_dtype_size=conv_dtype_size,
151 + )
152 +
153 +
154 +def compute_mamba_phys_ratio(ssm_sizes: tuple[int, ...], block_len: int) -> int:
155 + """Derive _physical_blocks_per_logical_kv_block from remote metadata.
156 +
157 + The remote engine's ratio is not sent directly in the handshake, so we
158 + reconstruct it: total mamba state per logical block / block_len.
159 +
160 + Args:
161 + ssm_sizes: (conv_state_bytes, ssm_state_bytes) from NixlAgentMetadata.
162 + block_len: the engine's block_len in bytes (from block_lens[0]).
163 + """
164 + return math.ceil((ssm_sizes[0] + ssm_sizes[1]) / block_len)