@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d4ecb75ba2

[PD][NixlPush][Bugfix] Fix prefix caching (#48758)

Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>

Nicolò Lucchesi committed Aug 7, 2026 at 16:21 UTC d4ecb75ba2f53a1e445cf5ac277ea8a5e78d516b

5 files changed +189 -73

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

+10 -3

@@ -314,7 +314,10 @@ def test_apply_prefix_caching_mamba_hybrid(

314 worker.kv_cache_config = make_kv_cache_config(block_size=16, mamba_enabled=True)

315

316 aligned_local, aligned_remote = worker._apply_prefix_caching(

317 - local_block_ids, remote_block_ids, remote_physical_per_logical

317 + local_block_ids,

318 + remote_block_ids,

319 + local_physical_per_logical,

320 + remote_physical_per_logical,

321 )

322

323 assert aligned_local == expected_local, (

@@ -413,7 +416,10 @@ def test_apply_prefix_caching_ssm_prefix_cache_hit(

416 worker.kv_cache_config = make_kv_cache_config(block_size=16, mamba_enabled=True)

417

418 aligned_local, aligned_remote = worker._apply_prefix_caching(

416 - local_block_ids, remote_block_ids, remote_physical_per_logical

419 + local_block_ids,

420 + remote_block_ids,

421 + local_physical_per_logical,

422 + remote_physical_per_logical,

423 )

424

425 assert aligned_local == expected_local, (

@@ -442,7 +448,7 @@ def test_apply_prefix_caching_ssm_unpairable_slots_rejected():

448

449 with pytest.raises(AssertionError, match="unpairable SSM state slots"):

450 worker._apply_prefix_caching(

445 - [list(range(10)), [4, 5, 6, 7]], [list(range(10)), [8, 9]], 10

451 + [list(range(10)), [4, 5, 6, 7]], [list(range(10)), [8, 9]], 10, 10

452 )

453

454

@@ -527,6 +533,7 @@ def test_mismatched_physical_per_logical_fails_with_prefix_caching(

533 aligned_local, aligned_remote = worker._apply_prefix_caching(

534 local_block_ids,

535 remote_block_ids,

536 + local_physical_per_logical,

537 remote_physical_per_logical,

538 )

539

tests/v1/kv_connector/unit/test_nixl_push_connector.py

+92 -1

@@ -341,6 +341,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):

341 # Single non-hybrid attention group, matching the stub block id lists.

342 w._has_mamba = False

343 w._group_spec_types = (FullAttentionSpec,)

344 + w._engine_ttl = 0.0

345 w._engine_last_active = {}

346

347 # Track _do_start_push_kv invocations.

@@ -569,7 +570,7 @@ def test_do_start_push_kv_drops_request_on_handshake_failure():

570 the failure is logged. Blocks are reclaimed by the lease/watchdog, matching

571 the old blocking behaviour."""

572 w = _StubWriterWorker.fresh()

572 - w._logical_to_kernel_block_ids = lambda x: x

573 + w._logical_to_kernel_block_ids = lambda x, ratio: x

574 xfer_calls: list[dict[str, Any]] = []

575 w._xfer_blocks_for_req = lambda **kw: xfer_calls.append(kw)

576 failures: list[dict[str, Any]] = []

@@ -1157,3 +1158,93 @@ class TestPushWriterMlaReplication:

1158 # All of the request's WRITE handles must be tracked together, so the

1159 # engine thread never sees a partial set and double-frees the request.

1160 assert sorted(w._sending_transfers["p-req"]) == [1000, 1001]

1161 +

1162 +

1163 +class TestPushPrefixCaching:

1164 + """Partial prefix-cache hit on D: D preallocates only its *uncomputed*

1165 + blocks, so on a partial hit it registers fewer blocks than P's full

1166 + sequence. P must WRITE only the *tail* of its sequence into D's slots --

1167 + mirroring pull-mode ``_apply_prefix_caching`` (end-trim), never a front-trim

1168 + which would write P's cached prefix into D's uncomputed suffix slots.

1169 +

1170 + The trim runs inside ``_xfer_blocks``, so these tests drive the real

1171 + ``_xfer_blocks_for_req`` path and stub only the NIXL WRITE. With one region

1172 + and 1:1 ratios ``_compute_desc_ids`` is the identity map, so the descs

1173 + captured from ``make_prepped_xfer`` are exactly the (trimmed) block IDs.

1174 + """

1175 +

1176 + @staticmethod

1177 + def _worker_driving_xfer(engine_id: str = "decode-engine"):

1178 + from types import SimpleNamespace

1179 +

1180 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (

1181 + TPMapping,

1182 + )

1183 +

1184 + w = _StubWriterWorker.fresh()

1185 + w._has_mamba = False

1186 + w.use_mla = False

1187 + w.block_size = 16

1188 + w.num_regions = 1

1189 + w._physical_blocks_per_logical_kv_block = 1

1190 +

1191 + w.transfer_topo = MagicMock()

1192 + w.transfer_topo.get_engine_info.return_value = SimpleNamespace(

1193 + remote_physical_blocks_per_logical=1,

1194 + remote_block_size=16,

1195 + remote_tp_size=1,

1196 + )

1197 + w.transfer_topo.tp_ratio.return_value = 1

1198 + w.transfer_topo.block_size_ratio.return_value = 1

1199 + w.tp_mappings = {

1200 + engine_id: TPMapping(

1201 + source_ranks_per_group=((0,),),

1202 + all_source_ranks=(0,),

1203 + rank_to_attention_slot={0: 0},

1204 + rank_offset_factor=0,

1205 + )

1206 + }

1207 + w.dst_num_blocks = {engine_id: 10_000, w.engine_id: 10_000}

1208 + w.dst_xfer_side_handles = {engine_id: {0: 5000}}

1209 + w.src_xfer_handles_by_block_size = {16: 2000}

1210 +

1211 + # Stub only the NIXL WRITE; kernel expansion, prefix-cache trim, desc

1212 + # computation and count assertions all run for real.

1213 + w.nixl_wrapper = MagicMock()

1214 + w.nixl_wrapper.make_prepped_xfer.return_value = 7

1215 + w._ensure_handshake = lambda *a, **k: None

1216 + w._logical_to_kernel_block_ids = lambda x, ratio: x

1217 + return w, engine_id

1218 +

1219 + @staticmethod

1220 + def _written_block_ids(w) -> tuple[list[int], list[int]]:

1221 + """Return the (local, remote) block IDs handed to the NIXL WRITE."""

1222 + args, _ = w.nixl_wrapper.make_prepped_xfer.call_args

1223 + # ("WRITE", local_handle, local_descs, remote_handle, remote_descs)

1224 + return list(args[2]), list(args[4])

1225 +

1226 + def test_partial_prefix_hit_end_trims_producer_blocks(self):

1227 + """D registered only its 2 uncomputed suffix blocks; P finished the

1228 + full 5-block sequence. P must WRITE its LAST 2 blocks into D's slots."""

1229 + w, _ = self._worker_driving_xfer()

1230 + reg = _registration_data("req-pc", local_block_ids=([500, 501],))

1231 +

1232 + NixlPushConnectorWorker._do_start_push_kv(

1233 + w, "req-pc", ([10, 11, 12, 13, 14],), reg

1234 + )

1235 +

1236 + local, remote = self._written_block_ids(w)

1237 + # End-trim (suffix), NOT front-trim: [13, 14], not [10, 11].

1238 + assert local == [13, 14]

1239 + assert remote == [500, 501]

1240 +

1241 + def test_no_prefix_hit_leaves_blocks_untrimmed(self):

1242 + """Equal counts (no prefix cache hit on D): nothing is trimmed."""

1243 + w, _ = self._worker_driving_xfer()

1244 + reg = _registration_data("req-full", local_block_ids=([500, 501, 502],))

1245 +

1246 + NixlPushConnectorWorker._do_start_push_kv(w, "req-full", ([10, 11, 12],), reg)

1247 +

1248 + local, remote = self._written_block_ids(w)

1249 + assert local == [10, 11, 12]

1250 + assert remote == [500, 501, 502]

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

+63 -55

@@ -2397,93 +2397,101 @@ class NixlBaseConnectorWorker:

2397

2398 def _apply_prefix_caching(

2399 self,

2400 - local_block_ids: BlockIds,

2401 - remote_block_ids: BlockIds,

2402 - remote_physical_per_logical: int,

2403 - ) -> tuple[BlockIds, list]:

2404 - """Apply prefix caching by trimming local/remote block ID lists.

2400 + decode_block_ids: BlockIds,

2401 + prefill_block_ids: BlockIds,

2402 + decode_physical_per_logical: int,

2403 + prefill_physical_per_logical: int,

2404 + ) -> tuple[BlockIds, BlockIds]:

2405 + """Trim block ID lists so only the uncomputed suffix is transferred.

2406 +

2407 + Inputs are *kernel* (physical) block IDs, already expanded from logical IDs

2408 + with each side's physical-per-logical ratio. Both pull and push call this after

2409 + that expansion so the trim happens at kernel granularity.

2410 +

2411 + The prefix-cache hit is always on the decode (D) side, so ``decode``

2412 + holds only its uncomputed blocks while prefill (P) holds the full

2413 + sequence. This is mode-independent: pull passes its own (D) blocks as

2414 + ``decode``, push passes the remote D registration as ``decode``.

2415

2406 - For non-Mamba models: end-trim remote to match local count, so that

2416 + For non-Mamba models: end-trim ``prefill`` to match ``decode`` count, so

2417 already-cached prefix blocks are skipped in the transfer.

2418

2409 - For Mamba hybrid (prefix caching not yet supported): front-trim both

2410 - to the minimum count to handle kernel block count discrepancies from

2411 - logical block rounding in heterogeneous TP.

2419 + For Mamba hybrid: SSM groups pair state slots by position and FA

2420 + groups end-trim to the uncomputed suffix when physical-per-logical

2421 + matches.

2422 """

2413 - # Partial prefix cache hit: just read uncomputed blocks.

2423 + # Partial prefix cache hit: just transfer uncomputed blocks.

2424 # Skip mamba groups — their blocks represent full state (conv+ssm),

2425 # not per-token data, so trimming would corrupt the transfer.

2416 - remote_block_ids = list(remote_block_ids)

2426 + prefill_block_ids = list(prefill_block_ids)

2427 if not self._has_mamba:

2418 - for i, remote_group in enumerate(remote_block_ids):

2419 - num_local_blocks = len(local_block_ids[i])

2420 - assert num_local_blocks <= len(remote_group)

2421 - if num_local_blocks < len(remote_group):

2422 - remote_block_ids[i] = remote_group[-num_local_blocks:]

2428 + for i, prefill_group in enumerate(prefill_block_ids):

2429 + num_decode_blocks = len(decode_block_ids[i])

2430 + assert num_decode_blocks <= len(prefill_group)

2431 + if num_decode_blocks < len(prefill_group):

2432 + prefill_block_ids[i] = prefill_group[-num_decode_blocks:]

2433 else:

2424 - # (NOTE: ZhanqiuHu) Mamba hybrid: no prefix caching support so far.HeteroTP

2425 - # can cause different kernel block counts due to logical block rounding.

2434 + # (NOTE: ZhanqiuHu) HeteroTP can cause different kernel block counts

2435 + # due to logical block rounding.

2436 # Example: 640 prompt tokens, kernel_block_size=64

2427 - # remote physical_per_logical=10, local physical_per_logical=6

2428 - # remote logical ids from kv_transfer_params = [0]

2429 - # local logical ids allocated = [0, 1]

2430 - # remote kernel blocks: [0..9] (1*10=10)

2431 - # local kernel blocks: [0..11] (2*6=12)

2437 + # prefill physical_per_logical=10, decode physical_per_logical=6

2438 + # prefill logical ids from kv_transfer_params = [0]

2439 + # decode logical ids allocated = [0, 1]

2440 + # prefill kernel blocks: [0..9] (1*10=10)

2441 + # decode kernel blocks: [0..11] (2*6=12)

2442 # actual data blocks = ceil(640/64) = 10, trim both to 10

2433 - # Vice versa (remote physical_per_logical=6, local=10):

2434 - # remote logical ids = [0, 1], local logical ids = [0]

2435 - # remote kernel blocks: [0..11] (2*6=12)

2436 - # local kernel blocks: [0..9] (1*10=10)

2443 + # Vice versa (prefill physical_per_logical=6, decode=10):

2444 + # prefill logical ids = [0, 1], decode logical ids = [0]

2445 + # prefill kernel blocks: [0..11] (2*6=12)

2446 + # decode kernel blocks: [0..9] (1*10=10)

2447 # actual data blocks = ceil(640/64) = 10, trim both to 10

2438 - local_block_ids = list(local_block_ids)

2439 - for i, remote_group in enumerate(remote_block_ids):

2440 - num_local_blocks = len(local_block_ids[i])

2441 - num_remote_blocks = len(remote_group)

2448 + decode_block_ids = list(decode_block_ids)

2449 + for i, prefill_group in enumerate(prefill_block_ids):

2450 + num_decode_blocks = len(decode_block_ids[i])

2451 + num_prefill_blocks = len(prefill_group)

2452 if _is_ssm_spec(self._group_spec_types[i]):

2443 - if num_local_blocks == num_remote_blocks:

2453 + if num_decode_blocks == num_prefill_blocks:

2454 continue

2455 # Only state-bearing slots reach here, single-state modes

2456 # just one (see get_exchange_clipped_blocks), so differing

2457 # counts mean position-indexed "all"-mode lists. A longer

2448 - # remote list carries earlier positions the local side

2449 - # already has (prefix hit) -> read its tail; a longer local

2458 + # prefill list carries earlier positions the decode side

2459 + # already has (prefix hit) -> take its tail; a longer decode

2460 # list holds the position D recomputes itself, which gets

2451 - # no remote state.

2452 - assert num_local_blocks - num_remote_blocks <= 1, (

2461 + # no prefill state.

2462 + assert num_decode_blocks - num_prefill_blocks <= 1, (

2463 f"Group {i}: unpairable SSM state slots, "

2454 - f"local={num_local_blocks} remote={num_remote_blocks}"

2464 + f"decode={num_decode_blocks} prefill={num_prefill_blocks}"

2465 )

2456 - num_blocks = min(num_local_blocks, num_remote_blocks)

2457 - if num_local_blocks < num_remote_blocks:

2458 - remote_block_ids[i] = remote_group[-num_blocks:]

2466 + num_blocks = min(num_decode_blocks, num_prefill_blocks)

2467 + if num_decode_blocks < num_prefill_blocks:

2468 + prefill_block_ids[i] = prefill_group[-num_blocks:]

2469 else:

2460 - local_block_ids[i] = local_block_ids[i][:num_blocks]

2470 + decode_block_ids[i] = decode_block_ids[i][:num_blocks]

2471 elif (

2462 - self._physical_blocks_per_logical_kv_block

2463 - == remote_physical_per_logical

2464 - and num_local_blocks < num_remote_blocks

2472 + decode_physical_per_logical == prefill_physical_per_logical

2473 + and num_decode_blocks < num_prefill_blocks

2474 ):

2475 # Partial prefix cache hit for FA group.

2467 - remote_block_ids[i] = remote_group[-num_local_blocks:]

2476 + prefill_block_ids[i] = prefill_group[-num_decode_blocks:]

2477 else:

2478 # TODO Handle prefix caching with different block_sizes

2479 # Allocation rounding legitimately leaves up to

2480 # ppl - 1 trailing dead kernel blocks per side (plus one

2472 - # extra local block for the recomputed final token), so

2481 + # extra decode block for the recomputed final token), so

2482 # the counts may differ by up to the sum of the two

2483 # ratios; anything larger indicates mismatched lists.

2484 max_padding = (

2476 - self._physical_blocks_per_logical_kv_block

2477 - + remote_physical_per_logical

2485 + decode_physical_per_logical + prefill_physical_per_logical

2486 )

2479 - assert abs(num_local_blocks - num_remote_blocks) <= max_padding, (

2480 - f"Group {i}: |{num_local_blocks} - "

2481 - f"{num_remote_blocks}| > {max_padding}"

2487 + assert abs(num_decode_blocks - num_prefill_blocks) <= max_padding, (

2488 + f"Group {i}: |{num_decode_blocks} - "

2489 + f"{num_prefill_blocks}| > {max_padding}"

2490 )

2483 - num_blocks = min(num_local_blocks, num_remote_blocks)

2484 - local_block_ids[i] = local_block_ids[i][:num_blocks]

2485 - remote_block_ids[i] = remote_group[:num_blocks]

2486 - return local_block_ids, remote_block_ids

2491 + num_blocks = min(num_decode_blocks, num_prefill_blocks)

2492 + decode_block_ids[i] = decode_block_ids[i][:num_blocks]

2493 + prefill_block_ids[i] = prefill_group[:num_blocks]

2494 + return decode_block_ids, prefill_block_ids

2495

2496 def get_kv_connector_stats(self) -> KVConnectorStats | None:

2497 """

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py

+4 -2

@@ -289,9 +289,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

289 == len(local_block_ids)

290 == len(self.kv_cache_config.kv_cache_groups)

291 )

292 - remote_physical_per_logical = remote_info.remote_physical_blocks_per_logical

292 local_block_ids, remote_block_ids = self._apply_prefix_caching(

294 - local_block_ids, remote_block_ids, remote_physical_per_logical

293 + decode_block_ids=local_block_ids,

294 + prefill_block_ids=remote_block_ids,

295 + decode_physical_per_logical=self._physical_blocks_per_logical_kv_block,

296 + prefill_physical_per_logical=remote_info.remote_physical_blocks_per_logical,

297 )

298

299 # NOTE (nicolo) With homogeneous TP, each TP worker loads KV from

vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py

+20 -12

@@ -454,10 +454,6 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

454 # how pull-mode transfers touch _engine_last_active in start_load_kv.

455 self._engine_last_active[decode_engine_id] = time.perf_counter()

456

457 - # Both sides stay logical here; ``_xfer_blocks_for_req`` converts each

458 - # to physical with its own physical-blocks-per-logical ratio -- P uses

459 - # ``self._physical_blocks_per_logical_kv_block``, D's is learned during

460 - # the NIXL handshake.

457 logical_local = self._as_grouped_block_ids(local_block_ids)

458 logical_remote = self._as_grouped_block_ids(remote_block_ids)

459 physical_local = self._logical_to_kernel_block_ids(

@@ -629,16 +625,28 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

625 logger.warning("No blocks to push for request %s", request_id)

626 return None

627

632 - # Align per-group block counts for push.

628 + # Prefix caching: D allocated only uncached blocks, so on a partial hit it

629 + # sends fewer than P's. End-trim P's blocks to that same suffix so we WRITE only

630 + # the uncomputed tail into D's slots. Runs on kernel ids, post-expansion.

631 + remote_block_ids, local_block_ids = self._apply_prefix_caching(

632 + decode_block_ids=remote_block_ids,

633 + prefill_block_ids=local_block_ids,

634 + decode_physical_per_logical=remote_info.remote_physical_blocks_per_logical,

635 + prefill_physical_per_logical=self._physical_blocks_per_logical_kv_block,

636 + )

637 +

638 local_block_ids = list(local_block_ids)

639 remote_block_ids = list(remote_block_ids)

635 - for i in range(min(len(local_block_ids), len(remote_block_ids))):

636 - num_local = len(local_block_ids[i])

637 - num_remote = len(remote_block_ids[i])

638 - if num_local > num_remote:

639 - local_block_ids[i] = local_block_ids[i][:num_remote]

640 - elif num_local < num_remote:

641 - remote_block_ids[i] = remote_block_ids[i][:num_local]

640 + assert len(local_block_ids) == len(remote_block_ids), (

641 + f"push group-count mismatch for {request_id}: {len(local_block_ids)} "

642 + f"local vs {len(remote_block_ids)} remote groups"

643 + )

644 + for i in range(len(local_block_ids)):

645 + assert len(local_block_ids[i]) == len(remote_block_ids[i]), (

646 + f"push block-count mismatch for {request_id} group {i}: "

647 + f"{len(local_block_ids[i])} local vs "

648 + f"{len(remote_block_ids[i])} remote blocks"

649 + )

650

651 # Get descs ids.

652 remote_block_descs_ids = self._compute_desc_ids(