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

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 13, 12:00 UTC

[NIXL] Generalize block-first backend layouts (FlashInfer-like) (#28282)

Nicolò Lucchesi committed Nov 11, 2025 at 17:57 UTC a7ef3eb0cd03e729c7a29914400e0ca928767999

2 files changed +52 -12

tests/v1/kv_connector/unit/test_nixl_connector.py

+15 -2

@@ -1096,7 +1096,8 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):

1096 llm.llm_engine.engine_core.shutdown()

1097

1098

1099 -def test_register_kv_caches(dist_init):

1099 +@pytest.mark.parametrize("attn_backend", ["FLASH_ATTN", "TRITON_ATTN"])

1100 +def test_register_kv_caches(dist_init, attn_backend, monkeypatch):

1101 """

1102 Test that register_kv_caches() properly calls nixl_wrapper methods with

1103 correct data.

@@ -1108,10 +1109,22 @@ def test_register_kv_caches(dist_init):

1109 block layout info

1110 """

1111

1112 + monkeypatch.setenv("VLLM_ATTENTION_BACKEND", attn_backend)

1113 +

1114 vllm_config = create_vllm_config()

1115

1116 + # Import the appropriate backend based on the parameter

1117 + if attn_backend == "FLASH_ATTN":

1118 + from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend

1119 +

1120 + backend_cls = FlashAttentionBackend

1121 + else: # TRITON_ATTN

1122 + from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend

1123 +

1124 + backend_cls = TritonAttentionBackend

1125 +

1126 # Create test kv cache tensors using proper backend shape

1114 - kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(

1127 + kv_cache_shape = backend_cls.get_kv_cache_shape(

1128 num_blocks=2, block_size=16, num_kv_heads=4, head_size=64

1129 )

1130 shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)

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

+37 -10

@@ -21,6 +21,7 @@ import torch

21 import zmq

23 from vllm import envs

24 +from vllm.attention import AttentionBackend

25 from vllm.attention.backends.registry import AttentionBackendEnum

26 from vllm.attention.selector import get_attn_backend

27 from vllm.config import VllmConfig

@@ -669,6 +670,33 @@ class NixlConnectorWorker:

670 remote_tp_size: dict[EngineId, int]

671 is_mla: bool

672 total_num_kv_heads: int

673 + attn_backend: type[AttentionBackend]

674 +

675 + def __post_init__(self):

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

677 + # or num_blocks. This is used to register the memory regions correctly.

678 + kv_cache_shape = self.attn_backend.get_kv_cache_shape(

679 + num_blocks=1, block_size=16, num_kv_heads=1, head_size=1

680 + )

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

682 + # we just mock num_blocks to 1 for the dimension check below.

683 + self._is_kv_layout_blocks_first = (

684 + len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1

685 + )

686 +

687 + attn_backend = AttentionBackendEnum[self.attn_backend.get_name()]

688 + self._use_pallas = attn_backend == AttentionBackendEnum.PALLAS

689 +

690 + @property

691 + def is_kv_layout_blocks_first(self) -> bool:

692 + return self._is_kv_layout_blocks_first

693 +

694 + @property

695 + def split_k_and_v(self) -> bool:

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

697 + return not (

698 + self.is_mla or self._use_pallas or self.is_kv_layout_blocks_first

699 + )

700

701 def tp_ratio(

702 self,

@@ -876,9 +904,6 @@ class NixlConnectorWorker:

904 use_mla=self.use_mla,

905 )

906 self.backend_name = backend.get_name()

879 - attn_backend = AttentionBackendEnum[self.backend_name]

880 - self._use_flashinfer = attn_backend == AttentionBackendEnum.FLASHINFER

881 - self._use_pallas = attn_backend == AttentionBackendEnum.PALLAS

907 self.kv_cache_layout = get_kv_cache_layout()

908 self.host_buffer_kv_cache_layout = self.kv_cache_layout

909 logger.debug("Detected attention backend %s", self.backend_name)

@@ -896,7 +921,9 @@ class NixlConnectorWorker:

921 remote_tp_size=self._tp_size, # shared state

922 is_mla=self.use_mla,

923 total_num_kv_heads=self.model_config.get_total_num_kv_heads(),

924 + attn_backend=backend,

925 )

926 + self._use_pallas = self.kv_topo._use_pallas

927

928 def _nixl_handshake(

929 self,

@@ -1076,7 +1103,7 @@ class NixlConnectorWorker:

1103 # (roughly 8KB vs 5KB).

1104 # Conversely for FlashInfer, K and V are registered in the same region

1105 # to better exploit the memory layout (ie num_blocks is the first dim).

1079 - split_k_and_v = not (self.use_mla or self._use_pallas or self._use_flashinfer)

1106 + split_k_and_v = self.kv_topo.split_k_and_v

1107 tensor_size_bytes = None

1108 # Enable different block lengths for different layers when MLA is used.

1109 self.block_len_per_layer = list[int]()

@@ -1141,7 +1168,7 @@ class NixlConnectorWorker:

1168

1169 self.device_kv_caches = kv_caches

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

1144 - if self._use_flashinfer:

1171 + if self.kv_topo.is_kv_layout_blocks_first:

1172 for i in range(len(self.slot_size_per_layer)):

1173 assert self.slot_size_per_layer[i] % 2 == 0

1174 self.slot_size_per_layer[i] //= 2

@@ -1169,7 +1196,7 @@ class NixlConnectorWorker:

1196 # (addr, len, device id)

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

1198

1172 - if self._use_flashinfer:

1199 + if self.kv_topo.is_kv_layout_blocks_first:

1200 # Separate and interleave K/V regions to maintain the same

1201 # descs ordering. This is needed for selecting contiguous heads

1202 # when split across TP ranks.

@@ -1331,7 +1358,7 @@ class NixlConnectorWorker:

1358 # (addr, len, device id)

1359 blocks_data.append((addr, kv_block_len, nixl_agent_meta.device_id))

1360

1334 - if self._use_flashinfer:

1361 + if self.kv_topo.is_kv_layout_blocks_first:

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

1363 for block_id in range(nixl_agent_meta.num_blocks):

1364 block_offset = block_id * nixl_agent_meta.block_lens[i]

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

1441 remote_block_size = remote_block_len // (

1442 self.slot_size_per_layer[0] * tp_ratio

1443 )

1417 - if self._use_flashinfer:

1444 + if self.kv_topo.is_kv_layout_blocks_first:

1445 # With flashinfer, KV are sent in the same message.

1446 remote_block_size //= 2

1447

@@ -1494,7 +1521,7 @@ class NixlConnectorWorker:

1521 - cache.index_copy_(0, indices, permuted_blocks) # copy permuted kv back

1522

1523 """

1497 - split_k_and_v = not (self.use_mla or self._use_pallas or self._use_flashinfer)

1524 + split_k_and_v = self.kv_topo.split_k_and_v

1525 inv_order = [0, 2, 1, 3]

1526 sample_cache = list(self.device_kv_caches.values())[0][0]

1527 target_shape = list(sample_cache.shape)

@@ -1874,7 +1901,7 @@ class NixlConnectorWorker:

1901 For FlashInfer, this is half the length of the whole block, as K and V

1902 share the same region.

1903 """

1877 - if self._use_flashinfer:

1904 + if self.kv_topo.is_kv_layout_blocks_first:

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

1906 block_len = self.block_len_per_layer[layer_idx] // 2

1907 else: