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 10, 10:15 UTC
[NIXL] refine decoder side post process for heterogeneous BlockSize and kv_layout (#30275)
Chendi.Xue committed
Jan 9, 2026 at 15:22 UTC
94578127a488d186c38665dd82116767fc34fa6d
2 files changed
+139
-87
vllm/distributed/kv_transfer/kv_connector/utils.py
+78
@@ -203,6 +203,84 @@ def copy_kv_blocks(
203
copy_fn(src_tensor, dst_tensor, src_indices, dst_indices)
204
205
206
+def kv_postprocess_blksize_on_receive(cache, indices, block_size_ratio):
207
+ """
208
+ Transforms the layout of received KV cache blocks to the local block_size.
209
+ (Only works for local blocksize > remote blocksize)
210
+
211
+ example:
212
+ local blocksize = 16 tokens, remote blocksize = 4 tokens
213
+ local block[0] = remote block[0, 1, 2, 3]
214
+ remote is |h0-b0|h1-b0|h2-b0|h3-b0|h0-b1|h1-b1|h2-b1|h3-b1|...
215
+ local is |h0-b0..................|h1-b0..................|...
216
+ permute is to:
217
+ 1. view => view remote as n_blocks * remote_shape(H,remoteN,D)
218
+ 2. permute => (H, nblocks, remoteN, D)
219
+ 3. flatten => (H, localN, D)
220
+ """
221
+ blocks_to_update = cache.index_select(0, indices)
222
+ # use physical order
223
+ blocks_to_update = blocks_to_update.permute(0, 2, 1, 3)
224
+ n_kv_heads, block_size, head_size = blocks_to_update.shape[1:]
225
+ remote_block_size = block_size // block_size_ratio
226
+ n_blocks = block_size_ratio
227
+
228
+ permuted_blocks = (
229
+ blocks_to_update.reshape(-1, n_blocks, n_kv_heads, remote_block_size, head_size)
230
+ .permute(0, 2, 1, 3, 4)
231
+ .flatten(2, 3)
232
+ )
233
+ permuted_blocks = permuted_blocks.permute(0, 2, 1, 3)
234
+ cache.index_copy_(0, indices, permuted_blocks)
235
+
236
+
237
+def kv_postprocess_layout_on_receive(cache, indices):
238
+ """Transforms the layout of received KV cache blocks to the local format.
239
+
240
+ This method corrects layout mismatches from direct memory copies by
241
+ permuting the tensor dimensions.
242
+
243
+ - **Source Layout:** `[num_blocks, n_kv_head, block_size, head_dim]`
244
+ - **Target Layout:** `[num_blocks, block_size, n_kv_head, head_dim]`
245
+
246
+ Implementation:
247
+ - x = blocks_to_update.reshape(src_shape) # view local kv with sender layout
248
+ - permuted_blocks = x.permute(*inv_order) # transpose n_kv_heads, block_size
249
+ - cache.index_copy_(0, indices, permuted_blocks) # copy permuted kv back
250
+
251
+ """
252
+ blocks_to_update = cache.index_select(0, indices)
253
+ target_shape = list(blocks_to_update.shape)
254
+ target_shape[0] = -1
255
+ inv_order = [0, 2, 1, 3]
256
+ src_shape = tuple(target_shape[i] for i in inv_order)
257
+ blocks_to_update = cache.index_select(0, indices)
258
+ permuted_blocks = blocks_to_update.reshape(src_shape).permute(*inv_order)
259
+ cache.index_copy_(0, indices, permuted_blocks)
260
+
261
+
262
+def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_ratio):
263
+ """
264
+ Transforms the layout of received KV cache to the local block_size and HND.
265
+ (Only works for local blocksize > remote blocksize)
266
+
267
+ prefill is HND, smaller block_size
268
+ decode(local) is NHD, larger block_size
269
+ """
270
+ blocks_to_update = cache.index_select(0, indices)
271
+
272
+ block_size, n_kv_heads, head_size = blocks_to_update.shape[1:]
273
+ remote_block_size = block_size // block_size_ratio
274
+ n_blocks = block_size_ratio
275
+
276
+ permuted_blocks = (
277
+ blocks_to_update.reshape(-1, n_blocks, n_kv_heads, remote_block_size, head_size)
278
+ .permute(0, 1, 3, 2, 4)
279
+ .flatten(1, 2)
280
+ )
281
+ cache.index_copy_(0, indices, permuted_blocks)
282
+
283
+
284
def yield_req_data(
285
scheduler_output,
286
) -> Iterator[tuple[str, tuple[list[int], ...], bool]]:
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+61
-87
@@ -24,6 +24,9 @@ from vllm.config import VllmConfig
24
from vllm.distributed.kv_transfer.kv_connector.utils import (
25
EngineId,
26
TpKVTopology,
27
+ kv_postprocess_blksize_and_layout_on_receive,
28
+ kv_postprocess_blksize_on_receive,
29
+ kv_postprocess_layout_on_receive,
30
yield_req_data,
31
)
32
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
@@ -1749,88 +1752,62 @@ class NixlConnectorWorker:
1752
"d2h",
1753
)
1754
1752
- def permute_device_kv(self, block_ids: list[int]):
1753
- """Transforms the layout of received KV cache blocks to the local format.
1754
-
1755
- This method corrects layout mismatches from direct memory copies by
1756
- permuting the tensor dimensions.
1757
-
1758
- - **Source Layout:** `[num_blocks, n_kv_head, block_size, head_dim]`
1759
- - **Target Layout:** `[num_blocks, block_size, n_kv_head, head_dim]`
1760
-
1761
- Args:
1762
- block_ids: A list of block IDs to update and permute.
1755
+ def post_process_device_kv_on_receive(
1756
+ self,
1757
+ block_size_ratio: int,
1758
+ block_ids_list: list[list[int]],
1759
+ ):
1760
+ """
1761
+ Post process device kv cache after receiving from remote.
1762
1764
- Implementation:
1765
- - x = blocks_to_update.reshape(src_shape) # view local kv with sender layout
1766
- - permuted_blocks = x.permute(*inv_order) # transpose n_kv_heads, block_size
1767
- - cache.index_copy_(0, indices, permuted_blocks) # copy permuted kv back
1763
+ 3 types of post processing supported:
1764
+ * kv_cache_postprocess_layout => convert from HND to NHD
1765
+ * kv_cache_postprocess_blksize => convert from small block size
1766
+ to large block size
1767
+ * kv_cache_postprocess_blksize_and_layout => convert from small
1768
+ block size to large block size and convert from HND to NHD
1769
1770
"""
1770
- split_k_and_v = self.kv_topo.split_k_and_v
1771
- inv_order = [0, 2, 1, 3]
1772
- sample_cache = list(self.device_kv_caches.values())[0][0]
1773
- target_shape = list(sample_cache.shape)
1774
- target_shape[0] = -1
1775
- src_shape = tuple(target_shape[i] for i in inv_order)
1776
- indices = torch.tensor(block_ids, device=sample_cache.device)
1777
-
1778
- for _, cache_or_caches in self.device_kv_caches.items():
1779
- cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
1780
- for cache in cache_list:
1781
- blocks_to_update = cache.index_select(0, indices)
1782
- permuted_blocks = blocks_to_update.reshape(src_shape).permute(
1783
- *inv_order
1784
- )
1785
- cache.index_copy_(0, indices, permuted_blocks)
1786
-
1787
- def blocksize_post_process(self, block_ids_per_ratio: dict[int, list[list[int]]]):
1788
- def _process_local_gt_remote(blocks_to_update, block_size_ratio):
1789
- n_kv_heads, block_size, head_size = blocks_to_update.shape[1:]
1790
- remote_block_size = block_size // block_size_ratio
1791
- n_blocks = block_size_ratio
1792
- # actual permute is to convert
1793
- # for local blocksize > remote blocksize
1794
- # ex: local blocksize = 16 tokens, remote blocksize = 4 tokens
1795
- # local block[0] = remote block[0, 1, 2, 3]
1796
- # remote is |h0-b0|h1-b0|h2-b0|h3-b0|h0-b1|h1-b1|h2-b1|h3-b1|...
1797
- # local is |h0-b0..................|h1-b0..................|...
1798
- # permute is to:
1799
- # 1. view => view remote as n_blocks * remote_shape(H,remoteN,D)
1800
- # 2. permute => (H, nblocks, remoteN, D)
1801
- # 3. flatten => (H, localN, D)
1802
- permuted_blocks = (
1803
- blocks_to_update.reshape(
1804
- -1, n_blocks, n_kv_heads, remote_block_size, head_size
1805
- )
1806
- .permute(0, 2, 1, 3, 4)
1807
- .flatten(2, 3)
1808
- )
1809
- return permuted_blocks
1810
-
1771
if len(self.device_kv_caches) == 0:
1772
return
1773
+ assert block_size_ratio >= 1, "Only nP < nD supported currently."
1774
+ if self.enable_permute_local_kv and block_size_ratio > 1:
1775
+ logger.debug(
1776
+ "Post-processing device kv cache on receive by converting "
1777
+ "block_size with %sx bigger and permuting layout from HND"
1778
+ " to NHD.",
1779
+ block_size_ratio,
1780
+ )
1781
+ elif self.enable_permute_local_kv:
1782
+ logger.debug(
1783
+ "Post-processing device kv cache on receive by permuting layout"
1784
+ "from HND to NHD."
1785
+ )
1786
+ else:
1787
+ logger.debug(
1788
+ "Post-processing device kv cache on receive by converting "
1789
+ "block_size with %sx bigger.",
1790
+ block_size_ratio,
1791
+ )
1792
+
1793
split_k_and_v = not (self.use_mla or self.kv_topo.is_kv_layout_blocks_first)
1814
- sample_cache = list(self.device_kv_caches.values())[0][0]
1815
- for block_size_ratio, block_ids_list in block_ids_per_ratio.items():
1816
- assert block_size_ratio > 1, "Only nP < nD supported currently."
1817
- block_ids_list = [[item for sublist in block_ids_list for item in sublist]]
1818
-
1819
- for block_ids in block_ids_list:
1820
- indices = torch.tensor(block_ids, device=sample_cache.device)
1821
-
1822
- for _, cache_or_caches in self.device_kv_caches.items():
1823
- cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
1824
- for cache in cache_list:
1825
- blocks_to_update = cache.index_select(0, indices)
1826
- # because kv_cache is always using original layout NHD as
1827
- # virtual shape while stride can be either HND / NHD at
1828
- # initialization.
1829
- # we need to firstly get physical view of the tensor
1830
- permuted_blocks = _process_local_gt_remote(
1831
- blocks_to_update.permute(0, 2, 1, 3), block_size_ratio
1832
- ).permute(0, 2, 1, 3)
1833
- cache.index_copy_(0, indices, permuted_blocks)
1794
+
1795
+ for block_ids in block_ids_list:
1796
+ indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
1797
+
1798
+ for _, cache_or_caches in self.device_kv_caches.items():
1799
+ cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
1800
+ for cache in cache_list:
1801
+ if self.enable_permute_local_kv and block_size_ratio > 1:
1802
+ kv_postprocess_blksize_and_layout_on_receive(
1803
+ cache, indices, block_size_ratio
1804
+ )
1805
+ elif self.enable_permute_local_kv:
1806
+ kv_postprocess_layout_on_receive(cache, indices)
1807
+ else:
1808
+ kv_postprocess_blksize_on_receive(
1809
+ cache, indices, block_size_ratio
1810
+ )
1811
1812
def get_finished(self) -> tuple[set[str], set[str]]:
1813
"""
@@ -1854,7 +1831,6 @@ class NixlConnectorWorker:
1831
len(done_recving),
1832
)
1833
1857
- block_ids_to_permute = []
1834
block_ids_for_blocksize_post_process = defaultdict(list)
1835
for req_id in done_recving:
1836
# clean up metadata for completed requests
@@ -1863,24 +1839,22 @@ class NixlConnectorWorker:
1839
assert meta.remote is not None
1840
if self.use_host_buffer:
1841
self.sync_recved_kv_to_device(req_id, meta)
1866
- if self.enable_permute_local_kv:
1867
- block_ids_to_permute += meta.local_physical_block_ids
1842
1843
# post processing for heteroblocksize
1844
block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(
1845
meta.remote.engine_id
1846
)
1873
- if (
1874
- not self.use_mla
1875
- and block_size_ratio > 1
1876
- and self.kv_cache_layout == "HND"
1847
+ if not self.use_mla and (
1848
+ block_size_ratio > 1 or self.enable_permute_local_kv
1849
):
1850
block_ids_for_blocksize_post_process[block_size_ratio].append(
1879
- meta.local_block_ids
1851
+ meta.local_physical_block_ids
1852
)
1881
- self.blocksize_post_process(block_ids_for_blocksize_post_process)
1882
- if len(block_ids_to_permute) > 0:
1883
- self.permute_device_kv(block_ids_to_permute)
1853
+ for (
1854
+ block_size_ratio,
1855
+ block_ids_list,
1856
+ ) in block_ids_for_blocksize_post_process.items():
1857
+ self.post_process_device_kv_on_receive(block_size_ratio, block_ids_list)
1858
1859
# Handle timeout to avoid stranding blocks on remote.
1860
now = time.perf_counter()