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

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 6, 21:45 UTC

[NIXL][1/N] Refactor `kernel_block_size` detection (#35752)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed Mar 11, 2026 at 09:11 UTC 098d844731c535c40c30498181de8f11f4b92cbb

5 files changed +126 -96

tests/v1/kv_connector/unit/test_nixl_connector.py

+57 -23

@@ -9,7 +9,7 @@ import textwrap

9 import time

10 import uuid

11 from collections import defaultdict

12 -from typing import Any

12 +from typing import Any, cast

13 from unittest.mock import MagicMock, patch

15 import msgspec

@@ -332,14 +332,22 @@ def test_kv_transfer_handshake(dist_init):

332

333 # Prefill connector will register KV cache to populate proper handshake

334 # metadata.

335 + # TODO this must match with values used in kv cache config

336 + kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)

337 prefill_connector = NixlConnector(

336 - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)

338 + vllm_config, KVConnectorRole.WORKER, kv_cache_config

339 + )

340 + kv_cache_spec = cast(

341 + AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec

342 )

343 kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(

339 - num_blocks=2, block_size=16, num_kv_heads=4, head_size=64

344 + num_blocks=kv_cache_config.num_blocks,

345 + block_size=kv_cache_spec.block_size,

346 + num_kv_heads=kv_cache_spec.num_kv_heads,

347 + head_size=kv_cache_spec.head_size,

348 )

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

342 - unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)

349 + shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)

350 + unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)

351 kv_caches = {

352 "layer0": shared_tensor,

353 "layer1": unique_tensor,

@@ -383,7 +391,7 @@ def test_kv_transfer_handshake(dist_init):

391

392 # Decode connector will be able to create handshake with the prefill connector.

393 decode_connector = NixlConnector(

386 - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)

394 + vllm_config, KVConnectorRole.WORKER, kv_cache_config

395 )

396 decode_connector.register_kv_caches(kv_caches)

397

@@ -525,11 +533,13 @@ class TestNixlHandshake:

533 request_id = "req_id"

534

535 # Test worker role in decode server.

528 - connector = NixlConnector(

529 - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)

530 - )

536 + kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)

537 + connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)

538 connector.connector_worker = FakeNixlConnectorWorker(

532 - vllm_config, connector.engine_id, hand_shake_latency=0

539 + vllm_config,

540 + connector.engine_id,

541 + hand_shake_latency=0,

542 + kv_cache_config=kv_cache_config,

543 )

544 assert isinstance(connector.connector_worker.nixl_wrapper, FakeNixlWrapper)

545 worker = connector.connector_worker

@@ -1479,18 +1489,22 @@ def test_register_kv_caches(

1489 patch(f"{nixl_module}.threading.Event"),

1490 patch(f"{nixl_module}.threading.Thread") as mock_thread,

1491 patch(f"{nixl_module}.get_current_attn_backend") as mock_get_attn_backend,

1492 + patch(f"{nixl_module}.get_current_attn_backends") as mock_get_attn_backends,

1493 ):

1494 # Ensure get_attn_backend returns the correct value due to

1495 # _cached_get_attn_backend returning the backend from previous

1496 # test run if not mocking.

1497 mock_get_attn_backend.return_value = backend_cls

1498 + mock_get_attn_backends.return_value = [backend_cls]

1499

1500 # Create connector

1489 - connector = NixlConnector(

1490 - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)

1491 - )

1501 + kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)

1502 + connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)

1503 connector.connector_worker = FakeNixlConnectorWorker(

1493 - vllm_config, connector.engine_id, hand_shake_latency=0

1504 + vllm_config,

1505 + connector.engine_id,

1506 + hand_shake_latency=0,

1507 + kv_cache_config=kv_cache_config,

1508 )

1509

1510 # Get the mock instance

@@ -1515,6 +1529,13 @@ def test_register_kv_caches(

1529 num_layers = 32

1530 block_size = 16

1531 num_blocks = 8

1532 + # Keep the fake worker's expected num_blocks in sync with the

1533 + # cross-layer tensor we are about to register.

1534 + worker_kv_cache_config = make_kv_cache_config(

1535 + block_size=block_size, num_blocks=num_blocks

1536 + )

1537 + connector.connector_worker.kv_cache_config = worker_kv_cache_config

1538 + connector.connector_worker.num_blocks = worker_kv_cache_config.num_blocks

1539 kv_cache_spec = AttentionSpec(

1540 block_size=block_size,

1541 num_kv_heads=4,

@@ -1568,11 +1589,17 @@ def test_register_kv_caches(

1589

1590 else:

1591 # Create test kv cache tensors using proper backend shape

1592 + kv_cache_spec = cast(

1593 + AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec

1594 + )

1595 kv_cache_shape = backend_cls.get_kv_cache_shape(

1572 - num_blocks=2, block_size=16, num_kv_heads=4, head_size=64

1596 + num_blocks=kv_cache_config.num_blocks,

1597 + block_size=kv_cache_spec.block_size,

1598 + num_kv_heads=kv_cache_spec.num_kv_heads,

1599 + head_size=kv_cache_spec.head_size,

1600 )

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

1575 - unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)

1601 + shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)

1602 + unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)

1603 kv_caches = {

1604 "layer0": shared_tensor,

1605 "layer1": unique_tensor,

@@ -1606,7 +1633,7 @@ def test_register_kv_caches(

1633 unique_tensor[1].data_ptr(),

1634 ]

1635 expected_num_entries = 4

1609 - expected_blocks_count = 8

1636 + expected_blocks_count = kv_cache_config.num_blocks * 4

1637

1638 # Execute register_kv_caches

1639 connector.register_kv_caches(kv_caches)

@@ -1639,7 +1666,7 @@ def test_register_kv_caches(

1666 num_blocks = 8

1667 expected_block_len = expected_tensor_size // num_blocks

1668 else:

1642 - num_blocks = 2

1669 + num_blocks = kv_cache_config.num_blocks

1670 if is_blocks_first:

1671 expected_block_len = expected_tensor_size // num_blocks // 2

1672 else:

@@ -2226,15 +2253,22 @@ def test_compatibility_hash_validation(

2253 "enforce_handshake_compat": enforce_handshake_compat

2254 },

2255 )

2256 + kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)

2257 decode_connector = NixlConnector(

2230 - local_vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)

2258 + local_vllm_config, KVConnectorRole.WORKER, kv_cache_config

2259 )

2260 decode_worker = decode_connector.connector_worker

2261 + kv_cache_spec = cast(

2262 + AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec

2263 + )

2264 kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape(

2234 - num_blocks=2, block_size=16, num_kv_heads=4, head_size=64

2265 + num_blocks=kv_cache_config.num_blocks,

2266 + block_size=kv_cache_spec.block_size,

2267 + num_kv_heads=kv_cache_spec.num_kv_heads,

2268 + head_size=kv_cache_spec.head_size,

2269 )

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

2237 - unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)

2270 + shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)

2271 + unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)

2272 kv_caches = {

2273 "layer0": shared_tensor,

2274 "layer1": unique_tensor,

tests/v1/worker/test_gpu_model_runner.py

+4 -16

@@ -38,7 +38,7 @@ from vllm.v1.kv_cache_interface import (

38 from vllm.v1.sample.metadata import SamplingMetadata

39 from vllm.v1.worker.gpu_input_batch import InputBatch

40 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

41 -from vllm.v1.worker.utils import AttentionGroup, select_common_block_size

41 +from vllm.v1.worker.utils import select_common_block_size

43 BLOCK_SIZE = 16

44 NUM_BLOCKS = 10

@@ -203,37 +203,25 @@ def _make_kv_cache_spec() -> FullAttentionSpec:

203 def test_select_common_block_size_prefers_manager_block_size():

204 backend_a = _make_mock_backend_for_kernel_block_size([MultipleOf(32)])

205 backend_b = _make_mock_backend_for_kernel_block_size([64, MultipleOf(16)])

206 - attn_groups = [

207 - AttentionGroup(backend_a, [], [], _make_kv_cache_spec(), 0),

208 - AttentionGroup(backend_b, [], [], _make_kv_cache_spec(), 0),

209 - ]

206

211 - selected_size = select_common_block_size(128, attn_groups)

207 + selected_size = select_common_block_size(128, [backend_a, backend_b])

208 assert selected_size == 128

209

210

211 def test_select_common_block_size_uses_largest_shared_int():

212 backend_a = _make_mock_backend_for_kernel_block_size([128, 64])

213 backend_b = _make_mock_backend_for_kernel_block_size([64, 32])

218 - attn_groups = [

219 - AttentionGroup(backend_a, [], [], _make_kv_cache_spec(), 0),

220 - AttentionGroup(backend_b, [], [], _make_kv_cache_spec(), 0),

221 - ]

214

223 - selected_size = select_common_block_size(256, attn_groups)

215 + selected_size = select_common_block_size(256, [backend_a, backend_b])

216 assert selected_size == 64

217

218

219 def test_select_common_block_size_no_valid_option():

220 backend_a = _make_mock_backend_for_kernel_block_size([64])

221 backend_b = _make_mock_backend_for_kernel_block_size([MultipleOf(16)])

230 - attn_groups = [

231 - AttentionGroup(backend_a, [], [], _make_kv_cache_spec(), 0),

232 - AttentionGroup(backend_b, [], [], _make_kv_cache_spec(), 0),

233 - ]

222

223 with pytest.raises(ValueError):

236 - select_common_block_size(48, attn_groups)

224 + select_common_block_size(48, [backend_a, backend_b])

225

226

227 def test_update_states_new_request(model_runner, dist_init):

vllm/distributed/kv_transfer/kv_connector/utils.py

+35 -25

@@ -358,15 +358,6 @@ class TpKVTopology:

358 # stride_order to retrieve physical position of block_size

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

360

361 - # In the default non-cross layers layout the block_size position

362 - # is logical while in the cross layers case it is the physical

363 - # position. This matches the shape of the actual kv cache tensors

364 - # passed at register_kv_caches()/register_cross_layers_kv_cache()

365 - block_size_position = kv_cache_shape.index(_MOCK_BLOCK_SIZE)

366 -

367 - assert block_size_position is not None

368 - self._block_size_position = -(len(kv_cache_shape) - block_size_position)

369 -

361 @property

362 def is_kv_layout_blocks_first(self) -> bool:

363 return self._is_kv_layout_blocks_first

@@ -390,10 +381,6 @@ class TpKVTopology:

381 def cross_layers_blocks(self) -> bool:

382 return self._cross_layers_blocks

383

393 - @property

394 - def block_size_position(self) -> int:

395 - return self._block_size_position

396 -

384 def tp_ratio(

385 self,

386 remote_tp_size: int,

@@ -484,23 +471,46 @@ class TpKVTopology:

471 return self.get_target_remote_ranks(remote_tp_size)

472

473

487 -def get_current_attn_backend(vllm_config: VllmConfig):

474 +def get_current_attn_backends(

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

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

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

478 +

479 + Args:

480 + vllm_config: The current vLLM configuration.

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

482 + When None, all attention layers are considered.

483 +

484 + Returns:

485 + Deduplicated list of attention backend classes.

486 + """

487 layer_type = cast(type[Any], AttentionLayerBase)

489 - layers = get_layers_from_vllm_config(vllm_config, layer_type, None)

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

489 if layers:

491 - backend = next(iter(layers.values())).get_attn_backend()

492 - else:

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

494 - logger.debug(

495 - "No layers found in the vLLM config. "

496 - "Falling back to default attention backend."

497 - )

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

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

491 + for layer in layers.values():

492 + backend = layer.get_attn_backend()

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

494 + return list(seen.values())

495 +

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

497 + logger.debug(

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

499 + )

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

501

500 - backend = get_attn_backend(

502 + return [

503 + get_attn_backend(

504 head_size=vllm_config.model_config.get_head_size(),

505 dtype=vllm_config.model_config.dtype,

506 kv_cache_dtype=vllm_config.cache_config.cache_dtype,

507 use_mla=vllm_config.model_config.use_mla,

508 )

506 - return backend

509 + ]

510 +

511 +

512 +def get_current_attn_backend(

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

514 +) -> type[AttentionBackend]:

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

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

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

+25 -27

@@ -13,7 +13,7 @@ from collections import defaultdict

13 from collections.abc import Iterator

14 from concurrent.futures import Future, ThreadPoolExecutor

15 from dataclasses import dataclass

16 -from typing import TYPE_CHECKING, Any

16 +from typing import TYPE_CHECKING, Any, cast

18 import msgspec

19 import numpy as np

@@ -27,6 +27,7 @@ from vllm.distributed.kv_transfer.kv_connector.utils import (

27 EngineId,

28 TpKVTopology,

29 get_current_attn_backend,

30 + get_current_attn_backends,

31 kv_postprocess_blksize_and_layout_on_receive,

32 kv_postprocess_blksize_on_receive,

33 kv_postprocess_layout_on_receive,

@@ -61,6 +62,7 @@ from vllm.v1.attention.backends.utils import get_kv_cache_layout

62 from vllm.v1.core.sched.output import SchedulerOutput

63 from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, SlidingWindowSpec

64 from vllm.v1.worker.block_table import BlockTable

65 +from vllm.v1.worker.utils import select_common_block_size

67 if TYPE_CHECKING:

68 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

@@ -945,7 +947,8 @@ class NixlConnectorWorker:

947

948 # Config.

949 self.vllm_config = vllm_config

948 - self.block_size = vllm_config.cache_config.block_size

950 + # mypy will complain on re-assignment otherwise.

951 + self.block_size: int = cast(int, vllm_config.cache_config.block_size)

952

953 if vllm_config.kv_transfer_config is None:

954 raise ValueError("kv_transfer_config must be set for NixlConnector")

@@ -993,7 +996,7 @@ class NixlConnectorWorker:

996 self.tp_rank = get_tensor_model_parallel_rank()

997 self.world_size = get_tensor_model_parallel_world_size()

998 self.tp_group = get_tp_group()

996 - self.num_blocks = 0

999 + self.num_blocks = kv_cache_config.num_blocks

1000 self.enable_permute_local_kv = False

1001

1002 # KV Caches and nixl tracking data.

@@ -1131,11 +1134,30 @@ class NixlConnectorWorker:

1134 self.xfer_stats = NixlKVConnectorStats()

1135

1136 self._physical_blocks_per_logical_kv_block = 1

1137 + self._sync_block_size_with_kernel()

1138

1139 self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config(

1140 "enforce_handshake_compat", True

1141 )

1142

1143 + def _sync_block_size_with_kernel(self) -> None:

1144 + backends = get_current_attn_backends(self.vllm_config)

1145 + kernel_block_size = select_common_block_size(self.block_size, backends)

1146 + if self.block_size != kernel_block_size:

1147 + logger.info_once(

1148 + "User-specified logical block size (%s) does not match"

1149 + " physical kernel block size (%s). Using the latter.",

1150 + self.block_size,

1151 + kernel_block_size,

1152 + )

1153 + assert self.block_size > kernel_block_size

1154 + self._physical_blocks_per_logical_kv_block = (

1155 + self.block_size // kernel_block_size

1156 + )

1157 + self.block_size = kernel_block_size

1158 + self._block_size[self.engine_id] = kernel_block_size

1159 + self.num_blocks *= self._physical_blocks_per_logical_kv_block

1160 +

1161 def _nixl_handshake(

1162 self,

1163 host: str,

@@ -1469,7 +1491,6 @@ class NixlConnectorWorker:

1491

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

1493 self.block_len_per_layer = list[int]()

1472 - self.slot_size_per_layer = list[int]() # HD bytes in kv terms

1494 for layer_name, cache_or_caches in xfer_buffers.items():

1495 cache_list = (

1496 cache_or_caches if self.kv_topo.split_k_and_v else [cache_or_caches]

@@ -1486,26 +1507,11 @@ class NixlConnectorWorker:

1507 logger.debug(

1508 "Registering layer %s with cache shape: %s", layer_name, cache.shape

1509 )

1489 - kernel_block_size = cache.shape[self.kv_topo.block_size_position]

1490 - if self.block_size != kernel_block_size:

1491 - logger.info_once(

1492 - "User-specified logical block size (%s) does not match"

1493 - " physical kernel block size (%s). Using the latter. ",

1494 - self.block_size,

1495 - kernel_block_size,

1496 - )

1497 - self._physical_blocks_per_logical_kv_block = (

1498 - self.block_size // kernel_block_size

1499 - )

1500 - self.block_size = kernel_block_size

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

1502 -

1510 seen_base_addresses.append(base_addr)

1511 curr_tensor_size_bytes = cache.numel() * cache.element_size()

1512

1513 if tensor_size_bytes is None:

1514 tensor_size_bytes = curr_tensor_size_bytes

1508 - self.num_blocks = cache.shape[0]

1515

1516 assert cache.shape[0] == self.num_blocks, (

1517 "All kv cache tensors must have the same number of blocks"

@@ -1514,9 +1520,6 @@ class NixlConnectorWorker:

1520 self.block_len_per_layer.append(

1521 curr_tensor_size_bytes // self.num_blocks

1522 )

1517 - self.slot_size_per_layer.append(

1518 - self.block_len_per_layer[-1] // self.block_size

1519 - )

1523

1524 if not self.use_mla:

1525 # Different kv cache shape is not supported by HeteroTP

@@ -1534,7 +1537,6 @@ class NixlConnectorWorker:

1537 "Different block lengths collected: %s", set(self.block_len_per_layer)

1538 )

1539 assert len(self.block_len_per_layer) == len(seen_base_addresses)

1537 - assert self.num_blocks != 0

1540

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

1542 self.num_regions = len(caches_data)

@@ -1550,10 +1552,6 @@ class NixlConnectorWorker:

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

1553

1554 if self.kv_topo.is_kv_layout_blocks_first:

1553 - for i in range(len(self.slot_size_per_layer)):

1554 - assert self.slot_size_per_layer[i] % 2 == 0

1555 - self.slot_size_per_layer[i] //= 2

1556 -

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

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

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

vllm/v1/worker/utils.py

+5 -5

@@ -258,7 +258,8 @@ class AttentionGroup:

258

259

260 def select_common_block_size(

261 - kv_manager_block_size: int, attn_groups: list[AttentionGroup]

261 + kv_manager_block_size: int,

262 + backends: list[type[AttentionBackend]],

263 ) -> int:

264 """

265 Select a block size that is supported by all backends and is a factor of

@@ -269,7 +270,7 @@ def select_common_block_size(

270

271 Args:

272 kv_manager_block_size: Block size of KV cache.

272 - attn_groups: List of attention groups.

273 + backends: List of attention backend classes.

274

275 Returns:

276 The selected block size.

@@ -297,8 +298,6 @@ def select_common_block_size(

298 return False

299 return True

300

300 - backends = [group.backend for group in attn_groups]

301 -

301 # Case 1: if the block_size of kv cache manager is supported by all backends,

302 # return it directly.

303 if block_size_is_supported(backends, kv_manager_block_size):

@@ -356,8 +355,9 @@ def prepare_kernel_block_sizes(

355 if isinstance(kv_cache_spec, AttentionSpec):

356 # This is an attention backend that supports virtual block splitting.

357 kv_manager_block_size = kv_cache_group.kv_cache_spec.block_size

358 + group_backends = [g.backend for g in attn_groups[kv_cache_gid]]

359 selected_kernel_size = select_common_block_size(

360 - kv_manager_block_size, attn_groups[kv_cache_gid]

360 + kv_manager_block_size, group_backends

361 )

362 kernel_block_sizes.append(selected_kernel_size)

363 elif isinstance(kv_cache_spec, MambaSpec):