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 14, 05:30 UTC
[P/D][Nixl] Make kv cache register compatible with hybrid memory allocator (#23079)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
Flora Feng committed
Aug 21, 2025 at 22:30 UTC
53415653ff24be03e7c90f5b42ef9cb3f72aad71
3 files changed
+150
-95
tests/v1/kv_connector/unit/test_nixl_connector.py
+85
-1
@@ -14,6 +14,7 @@ from unittest.mock import patch
14
15
import pytest
16
import ray
17
+import torch
18
19
from vllm import LLM
20
from vllm.config import KVTransferConfig
@@ -22,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
23
NixlConnectorWorker)
24
from vllm.forward_context import ForwardContext
25
from vllm.sampling_params import SamplingParams
26
+from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
27
28
from .utils import create_request, create_scheduler, create_vllm_config
29
@@ -98,7 +100,6 @@ class FakeNixlWrapper:
100
101
def set_cycles_before_xfer_done(self, cycles: int):
102
"""Set the number of cycles before a transfer is considered done."""
101
- self._cycles_before_xfer_done = cycles
103
104
105
@contextlib.contextmanager
@@ -562,3 +563,86 @@ def _run_abort_timeout_test(llm_kwargs: dict, timeout: int):
563
sampling_params)
564
# Request-0 times out and is cleared!
565
assert '0' not in req_to_blocks
566
+
567
+
568
+def test_register_kv_caches(dist_init):
569
+ """
570
+ Test that register_kv_caches() properly calls nixl_wrapper methods with
571
+ correct data.
572
+
573
+ This test verifies:
574
+ 1. nixl_wrapper.get_reg_descs() is called with caches_data containing
575
+ tensor metadata
576
+ 2. nixl_wrapper.get_xfer_descs() is called with blocks_data containing
577
+ block layout info
578
+ """
579
+
580
+ vllm_config = create_vllm_config()
581
+
582
+ # Create test kv cache tensors using proper backend shape
583
+ kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(num_blocks=2,
584
+ block_size=16,
585
+ num_kv_heads=4,
586
+ head_size=64)
587
+ shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
588
+ unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
589
+ kv_caches = {
590
+ "layer0": shared_tensor,
591
+ "layer1": unique_tensor,
592
+ "layer2": shared_tensor,
593
+ }
594
+
595
+ # Store tensor info for validation
596
+ expected_tensor_size = shared_tensor[0].element_size(
597
+ ) * shared_tensor[0].numel()
598
+ expected_base_addrs = [
599
+ shared_tensor[0].data_ptr(), shared_tensor[1].data_ptr(),
600
+ unique_tensor[0].data_ptr(), unique_tensor[1].data_ptr()
601
+ ]
602
+
603
+ with patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper") as mock_nixl_wrapper, \
604
+ patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.threading.Event"), \
605
+ patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.threading.Thread"): # noqa: E501
606
+
607
+ # Create connector
608
+ connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
609
+ connector.connector_worker = FakeNixlConnectorWorker(
610
+ vllm_config, connector.engine_id, hand_shake_latency=0)
611
+
612
+ # Get the mock instance
613
+ mock_wrapper_instance = mock_nixl_wrapper.return_value
614
+ connector.connector_worker.nixl_wrapper = mock_wrapper_instance
615
+
616
+ # Execute register_kv_caches
617
+ connector.register_kv_caches(kv_caches)
618
+
619
+ # Verify get_reg_descs was called with caches_data
620
+ assert mock_wrapper_instance.get_reg_descs.called
621
+ caches_data, _ = mock_wrapper_instance.get_reg_descs.call_args[0]
622
+ assert len(caches_data) == 4
623
+
624
+ for i, cache_entry in enumerate(caches_data):
625
+ base_addr, size, _tp_rank, _ = cache_entry
626
+ assert size == expected_tensor_size, \
627
+ f"Entry {i}: Expected tensor size {expected_tensor_size}, " \
628
+ f"got {size}"
629
+ assert base_addr == expected_base_addrs[i], \
630
+ f"Entry {i}: Expected base address {expected_base_addrs[i]}, " \
631
+ f"got {base_addr}"
632
+
633
+ # Verify get_xfer_descs was called with blocks_data
634
+ assert mock_wrapper_instance.get_xfer_descs.called
635
+ blocks_data, _ = mock_wrapper_instance.get_xfer_descs.call_args[0]
636
+
637
+ # Validate blocks_data structure and size
638
+ expected_blocks_count = 8
639
+ assert len(blocks_data) == expected_blocks_count, \
640
+ f"Expected {expected_blocks_count} blocks, " \
641
+ f"got {len(blocks_data)}"
642
+
643
+ expected_block_len = expected_tensor_size // 2
644
+ for i, block_entry in enumerate(blocks_data):
645
+ block_start_addr, block_len, tp_rank = block_entry
646
+ assert block_len == expected_block_len, \
647
+ f"Block entry {i}: Expected block len {expected_block_len}, " \
648
+ f"got {block_len}"
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+2
-2
@@ -131,8 +131,8 @@ class KVConnectorBase_V1(ABC):
131
Initialize with the KV caches. Useful for pre-registering the
132
KV Caches in the KVConnector (e.g. for NIXL).
133
134
- Args: kv_caches:
135
- dictionary of layer names, kv cache
134
+ Args:
135
+ kv_caches: dictionary of layer names, kv cache
136
"""
137
return
138
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+63
-92
@@ -686,9 +686,6 @@ class NixlConnectorWorker:
686
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
687
"""Register the KV Cache data in nixl."""
688
689
- _, first_kv_cache = next(iter(kv_caches.items()))
690
- kv_elem_size = first_kv_cache.element_size()
691
-
689
if self.use_host_buffer:
690
self.initialize_host_xfer_buffer(kv_caches=kv_caches)
691
assert len(self.host_xfer_buffers) == len(kv_caches), (
@@ -701,66 +698,16 @@ class NixlConnectorWorker:
698
"host_xfer_buffer should not be initialized when "
699
f"kv_buffer_device is {self.kv_buffer_device}")
700
704
- # TODO(tms): Find a more robust way to detect and handle MLA
705
- # NOTE (NickLucche) To move blocks efficiently with NIXL, the expected
706
- # KV memory layout is HND, as opposed to the default NHD. Note that it
707
- # will only affects the strides. For MLA instead, we make require no
708
- # such thing and resort to the standard layout.
709
- use_mla = len(first_kv_cache.shape) == 3
710
- if self.device_type == "tpu":
711
- assert not use_mla, f"{self.kv_buffer_device} does not support MLA."
712
- assert self._use_pallas_v1, f"attn backend: {self.backend_name}"
713
- # tpu (v1) kv shape per layer:
714
- # (num_blocks, block_size, num_kv_heads * 2, head_size)
715
- self.num_blocks = first_kv_cache.shape[0]
716
- block_rank = 3 # [block_size, kv_heads, head_dim]
717
- block_shape = first_kv_cache.shape[-block_rank:]
718
- block_size, n_kv_heads_x_2, head_dim = block_shape
719
- self.slot_size_bytes = kv_elem_size * n_kv_heads_x_2 * head_dim
720
- elif self.device_type == "cuda":
721
- assert use_mla == self.use_mla
722
- # TODO (NickLucche) not compatible with hybrid allocator.
723
- # Enforce check once it goes live, as a single kv layout
724
- # is expected for xfers.
725
- if use_mla:
726
- # MLA case.
727
- self.num_blocks = first_kv_cache.shape[0]
728
- block_rank = 2 # [block_size, latent_dim]
729
- block_shape = first_kv_cache.shape[-block_rank:]
730
- block_size, kv_latent_dim = block_shape
731
- self.slot_size_bytes = kv_elem_size * kv_latent_dim
732
- else:
733
- # [2 (k and v), num_blocks, ...]
734
- if self._use_flashinfer:
735
- # FlashInfer swaps 2<->num_blocks dimensions.
736
- self.num_blocks = first_kv_cache.shape[0]
737
- block_rank = 4 # [2, block_size, kv_heads, head_dim]
738
- else:
739
- self.num_blocks = first_kv_cache.shape[1]
740
- block_rank = 3 # [block_size, kv_heads, head_dim]
741
- block_shape = first_kv_cache.shape[-block_rank:]
742
- block_size, n_kv_heads, head_dim = block_shape[-3:]
743
- # head size in bytes.
744
- self.slot_size_bytes = kv_elem_size * n_kv_heads * head_dim
745
- assert block_size == self.block_size
746
- else:
747
- raise RuntimeError(
748
- f"{self.device_type} ({self.backend_name}) is not supported.")
749
-
750
- # TODO(tms): self.block_len needs to be per-layer for sliding window,
751
- # hybrid attn, etc
752
- # block size in bytes
753
- self.block_len = kv_elem_size * math.prod(block_shape)
701
logger.info(
702
"Registering KV_Caches. use_mla: %s, kv_buffer_device: %s, "
756
- "use_host_buffer: %s, num_blocks: %s, block_shape: %s, "
757
- "per_layer_kv_cache_shape: %s", use_mla, self.kv_buffer_device,
758
- self.use_host_buffer, self.num_blocks, block_shape,
759
- first_kv_cache.shape)
760
- self.dst_num_blocks[self.engine_id] = self.num_blocks
761
- self.device_kv_caches = kv_caches
762
- kv_caches_base_addr = []
703
+ "use_host_buffer: %s", self.use_mla, self.kv_buffer_device,
704
+ self.use_host_buffer)
705
+
706
caches_data = []
707
+ # With hybrid allocator, layers can share a kv cache tensor
708
+ seen_base_addresses = []
709
+ xfer_buffers = (self.host_xfer_buffers
710
+ if self.use_host_buffer else kv_caches)
711
712
# Note(tms): I modified this from the original region setup code.
713
# K and V are now in different regions. Advantage is that we can
@@ -770,42 +717,35 @@ class NixlConnectorWorker:
717
# (roughly 8KB vs 5KB).
718
# Conversely for FlashInfer, K and V are transferred in the same tensor
719
# to better exploit the memory layout (ie num_blocks is the first dim).
773
- for cache_or_caches in xfer_buffers.values():
774
- # Normalize to always be a list of caches
775
- cache_list = [cache_or_caches] if use_mla \
776
- or self._use_pallas_v1 or self._use_flashinfer \
777
- else cache_or_caches
720
+ split_k_and_v = not (self.use_mla or self._use_pallas_v1
721
+ or self._use_flashinfer)
722
+ tensor_size_bytes = None
723
+ for layer_name, cache_or_caches in xfer_buffers.items():
724
+ cache_list = cache_or_caches if split_k_and_v else [
725
+ cache_or_caches
726
+ ]
727
+
728
for cache in cache_list:
729
base_addr = cache.data_ptr()
780
- region_len = self.num_blocks * self.block_len
781
- # NOTE: use tp_rank for device_id since multi-node TP
782
- # is rarely used.
783
- caches_data.append((base_addr, region_len, self.tp_rank, ""))
784
- kv_caches_base_addr.append(base_addr)
785
- self.kv_caches_base_addr[self.engine_id] = kv_caches_base_addr
730
+ if base_addr in seen_base_addresses:
731
+ continue
732
+
733
+ seen_base_addresses.append(base_addr)
734
+ curr_tensor_size_bytes = cache.numel() * cache.element_size()
735
+
736
+ if tensor_size_bytes is None:
737
+ tensor_size_bytes = curr_tensor_size_bytes
738
+ self.num_blocks = cache.shape[0]
739
+
740
+ assert tensor_size_bytes == curr_tensor_size_bytes, \
741
+ "All kv cache tensors must have the same size"
742
+ caches_data.append(
743
+ (base_addr, tensor_size_bytes, self.tp_rank, ""))
744
+
745
+ self.kv_caches_base_addr[self.engine_id] = seen_base_addresses
746
self.num_regions = len(caches_data)
747
self.num_layers = len(xfer_buffers.keys())
748
789
- # TODO(mgoin): remove this once we have hybrid memory allocator
790
- # Optimization for models with local attention (Llama 4)
791
- if self.vllm_config.model_config.hf_config.model_type == "llama4":
792
- from transformers import Llama4TextConfig
793
- assert isinstance(self.vllm_config.model_config.hf_text_config,
794
- Llama4TextConfig)
795
- llama4_config = self.vllm_config.model_config.hf_text_config
796
- no_rope_layers = llama4_config.no_rope_layers
797
- chunk_size = llama4_config.attention_chunk_size
798
- chunk_block_size = math.ceil(chunk_size / self.block_size)
799
- for layer_idx in range(self.num_layers):
800
- # no_rope_layers[layer_idx] == 0 means NoPE (global)
801
- # Any other value means RoPE (local chunked)
802
- is_local_attention = no_rope_layers[layer_idx] != 0
803
- block_window = chunk_block_size if is_local_attention else None
804
- self.block_window_per_layer.append(block_window)
805
- logger.debug("Llama 4 block window per layer mapping: %s",
806
- self.block_window_per_layer)
807
- assert len(self.block_window_per_layer) == self.num_layers
808
-
749
descs = self.nixl_wrapper.get_reg_descs(caches_data,
750
self.nixl_memory_type)
751
logger.debug("Registering descs: %s", caches_data)
@@ -813,9 +753,20 @@ class NixlConnectorWorker:
753
logger.debug("Done registering descs")
754
self._registered_descs.append(descs)
755
756
+ assert tensor_size_bytes is not None
757
+ assert self.num_blocks != 0
758
+ assert tensor_size_bytes % self.num_blocks == 0
759
+ self.block_len = tensor_size_bytes // self.num_blocks
760
+ self.slot_size_bytes = self.block_len // self.block_size
761
+ if self._use_flashinfer:
762
+ assert self.slot_size_bytes % 2 == 0
763
+ self.slot_size_bytes /= 2
764
+ self.device_kv_caches = kv_caches
765
+ self.dst_num_blocks[self.engine_id] = self.num_blocks
766
+
767
# Register local/src descr for NIXL xfer.
768
blocks_data = []
818
- for base_addr in self.kv_caches_base_addr[self.engine_id]:
769
+ for base_addr in seen_base_addresses:
770
# NOTE With heter-TP, more blocks are prepared than what are
771
# needed as self.num_blocks >= nixl_agent_meta.num_blocks. We
772
# could create fewer, but then _get_block_descs_ids needs to
@@ -836,6 +787,26 @@ class NixlConnectorWorker:
787
self.src_xfer_side_handle = self.nixl_wrapper.prep_xfer_dlist(
788
"NIXL_INIT_AGENT", descs)
789
790
+ # TODO(mgoin): Hybrid memory allocator is currently diabled for
791
+ # models with local attention (Llama 4). Can remove this once enabled.
792
+ if self.vllm_config.model_config.hf_config.model_type == "llama4":
793
+ from transformers import Llama4TextConfig
794
+ assert isinstance(self.vllm_config.model_config.hf_text_config,
795
+ Llama4TextConfig)
796
+ llama4_config = self.vllm_config.model_config.hf_text_config
797
+ no_rope_layers = llama4_config.no_rope_layers
798
+ chunk_size = llama4_config.attention_chunk_size
799
+ chunk_block_size = math.ceil(chunk_size / self.block_size)
800
+ for layer_idx in range(self.num_layers):
801
+ # no_rope_layers[layer_idx] == 0 means NoPE (global)
802
+ # Any other value means RoPE (local chunked)
803
+ is_local_attention = no_rope_layers[layer_idx] != 0
804
+ block_window = chunk_block_size if is_local_attention else None
805
+ self.block_window_per_layer.append(block_window)
806
+ logger.debug("Llama 4 block window per layer mapping: %s",
807
+ self.block_window_per_layer)
808
+ assert len(self.block_window_per_layer) == self.num_layers
809
+
810
# After KV Caches registered, listen for new connections.
811
metadata = NixlAgentMetadata(
812
engine_id=self.engine_id,