iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
This commit is contained in:
Nicolò Lucchesi
2026-08-07 16:21:17 +02:00
committed by GitHub
parent f2bfad9167
commit d4ecb75ba2
5 changed files with 189 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+10 -3
View File
@@ -314,7 +314,10 @@ def test_apply_prefix_caching_mamba_hybrid(
worker.kv_cache_config = make_kv_cache_config(block_size=16, mamba_enabled=True)
aligned_local, aligned_remote = worker._apply_prefix_caching(
local_block_ids, remote_block_ids, remote_physical_per_logical
local_block_ids,
remote_block_ids,
local_physical_per_logical,
remote_physical_per_logical,
)
assert aligned_local == expected_local, (
@@ -413,7 +416,10 @@ def test_apply_prefix_caching_ssm_prefix_cache_hit(
worker.kv_cache_config = make_kv_cache_config(block_size=16, mamba_enabled=True)
aligned_local, aligned_remote = worker._apply_prefix_caching(
local_block_ids, remote_block_ids, remote_physical_per_logical
local_block_ids,
remote_block_ids,
local_physical_per_logical,
remote_physical_per_logical,
)
assert aligned_local == expected_local, (
@@ -442,7 +448,7 @@ def test_apply_prefix_caching_ssm_unpairable_slots_rejected():
with pytest.raises(AssertionError, match="unpairable SSM state slots"):
worker._apply_prefix_caching(
[list(range(10)), [4, 5, 6, 7]], [list(range(10)), [8, 9]], 10
[list(range(10)), [4, 5, 6, 7]], [list(range(10)), [8, 9]], 10, 10
)
@@ -527,6 +533,7 @@ def test_mismatched_physical_per_logical_fails_with_prefix_caching(
aligned_local, aligned_remote = worker._apply_prefix_caching(
local_block_ids,
remote_block_ids,
local_physical_per_logical,
remote_physical_per_logical,
)
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+92 -1
View File
@@ -341,6 +341,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):
# Single non-hybrid attention group, matching the stub block id lists.
w._has_mamba = False
w._group_spec_types = (FullAttentionSpec,)
w._engine_ttl = 0.0
w._engine_last_active = {}
# Track _do_start_push_kv invocations.
@@ -569,7 +570,7 @@ def test_do_start_push_kv_drops_request_on_handshake_failure():
the failure is logged. Blocks are reclaimed by the lease/watchdog, matching
the old blocking behaviour."""
w = _StubWriterWorker.fresh()
w._logical_to_kernel_block_ids = lambda x: x
w._logical_to_kernel_block_ids = lambda x, ratio: x
xfer_calls: list[dict[str, Any]] = []
w._xfer_blocks_for_req = lambda **kw: xfer_calls.append(kw)
failures: list[dict[str, Any]] = []
@@ -1157,3 +1158,93 @@ class TestPushWriterMlaReplication:
# All of the request's WRITE handles must be tracked together, so the
# engine thread never sees a partial set and double-frees the request.
assert sorted(w._sending_transfers["p-req"]) == [1000, 1001]
class TestPushPrefixCaching:
"""Partial prefix-cache hit on D: D preallocates only its *uncomputed*
blocks, so on a partial hit it registers fewer blocks than P's full
sequence. P must WRITE only the *tail* of its sequence into D's slots --
mirroring pull-mode ``_apply_prefix_caching`` (end-trim), never a front-trim
which would write P's cached prefix into D's uncomputed suffix slots.
The trim runs inside ``_xfer_blocks``, so these tests drive the real
``_xfer_blocks_for_req`` path and stub only the NIXL WRITE. With one region
and 1:1 ratios ``_compute_desc_ids`` is the identity map, so the descs
captured from ``make_prepped_xfer`` are exactly the (trimmed) block IDs.
"""
@staticmethod
def _worker_driving_xfer(engine_id: str = "decode-engine"):
from types import SimpleNamespace
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (
TPMapping,
)
w = _StubWriterWorker.fresh()
w._has_mamba = False
w.use_mla = False
w.block_size = 16
w.num_regions = 1
w._physical_blocks_per_logical_kv_block = 1
w.transfer_topo = MagicMock()
w.transfer_topo.get_engine_info.return_value = SimpleNamespace(
remote_physical_blocks_per_logical=1,
remote_block_size=16,
remote_tp_size=1,
)
w.transfer_topo.tp_ratio.return_value = 1
w.transfer_topo.block_size_ratio.return_value = 1
w.tp_mappings = {
engine_id: TPMapping(
source_ranks_per_group=((0,),),
all_source_ranks=(0,),
rank_to_attention_slot={0: 0},
rank_offset_factor=0,
)
}
w.dst_num_blocks = {engine_id: 10_000, w.engine_id: 10_000}
w.dst_xfer_side_handles = {engine_id: {0: 5000}}
w.src_xfer_handles_by_block_size = {16: 2000}
# Stub only the NIXL WRITE; kernel expansion, prefix-cache trim, desc
# computation and count assertions all run for real.
w.nixl_wrapper = MagicMock()
w.nixl_wrapper.make_prepped_xfer.return_value = 7
w._ensure_handshake = lambda *a, **k: None
w._logical_to_kernel_block_ids = lambda x, ratio: x
return w, engine_id
@staticmethod
def _written_block_ids(w) -> tuple[list[int], list[int]]:
"""Return the (local, remote) block IDs handed to the NIXL WRITE."""
args, _ = w.nixl_wrapper.make_prepped_xfer.call_args
# ("WRITE", local_handle, local_descs, remote_handle, remote_descs)
return list(args[2]), list(args[4])
def test_partial_prefix_hit_end_trims_producer_blocks(self):
"""D registered only its 2 uncomputed suffix blocks; P finished the
full 5-block sequence. P must WRITE its LAST 2 blocks into D's slots."""
w, _ = self._worker_driving_xfer()
reg = _registration_data("req-pc", local_block_ids=([500, 501],))
NixlPushConnectorWorker._do_start_push_kv(
w, "req-pc", ([10, 11, 12, 13, 14],), reg
)
local, remote = self._written_block_ids(w)
# End-trim (suffix), NOT front-trim: [13, 14], not [10, 11].
assert local == [13, 14]
assert remote == [500, 501]
def test_no_prefix_hit_leaves_blocks_untrimmed(self):
"""Equal counts (no prefix cache hit on D): nothing is trimmed."""
w, _ = self._worker_driving_xfer()
reg = _registration_data("req-full", local_block_ids=([500, 501, 502],))
NixlPushConnectorWorker._do_start_push_kv(w, "req-full", ([10, 11, 12],), reg)
local, remote = self._written_block_ids(w)
assert local == [10, 11, 12]
assert remote == [500, 501, 502]
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+63 -55
View File
@@ -2397,93 +2397,101 @@ class NixlBaseConnectorWorker:
def _apply_prefix_caching(
self,
local_block_ids: BlockIds,
remote_block_ids: BlockIds,
remote_physical_per_logical: int,
) -> tuple[BlockIds, list]:
"""Apply prefix caching by trimming local/remote block ID lists.
decode_block_ids: BlockIds,
prefill_block_ids: BlockIds,
decode_physical_per_logical: int,
prefill_physical_per_logical: int,
) -> tuple[BlockIds, BlockIds]:
"""Trim block ID lists so only the uncomputed suffix is transferred.
For non-Mamba models: end-trim remote to match local count, so that
Inputs are *kernel* (physical) block IDs, already expanded from logical IDs
with each side's physical-per-logical ratio. Both pull and push call this after
that expansion so the trim happens at kernel granularity.
The prefix-cache hit is always on the decode (D) side, so ``decode``
holds only its uncomputed blocks while prefill (P) holds the full
sequence. This is mode-independent: pull passes its own (D) blocks as
``decode``, push passes the remote D registration as ``decode``.
For non-Mamba models: end-trim ``prefill`` to match ``decode`` count, so
already-cached prefix blocks are skipped in the transfer.
For Mamba hybrid (prefix caching not yet supported): front-trim both
to the minimum count to handle kernel block count discrepancies from
logical block rounding in heterogeneous TP.
For Mamba hybrid: SSM groups pair state slots by position and FA
groups end-trim to the uncomputed suffix when physical-per-logical
matches.
"""
# Partial prefix cache hit: just read uncomputed blocks.
# Partial prefix cache hit: just transfer uncomputed blocks.
# Skip mamba groups — their blocks represent full state (conv+ssm),
# not per-token data, so trimming would corrupt the transfer.
remote_block_ids = list(remote_block_ids)
prefill_block_ids = list(prefill_block_ids)
if not self._has_mamba:
for i, remote_group in enumerate(remote_block_ids):
num_local_blocks = len(local_block_ids[i])
assert num_local_blocks <= len(remote_group)
if num_local_blocks < len(remote_group):
remote_block_ids[i] = remote_group[-num_local_blocks:]
for i, prefill_group in enumerate(prefill_block_ids):
num_decode_blocks = len(decode_block_ids[i])
assert num_decode_blocks <= len(prefill_group)
if num_decode_blocks < len(prefill_group):
prefill_block_ids[i] = prefill_group[-num_decode_blocks:]
else:
# (NOTE: ZhanqiuHu) Mamba hybrid: no prefix caching support so far.HeteroTP
# can cause different kernel block counts due to logical block rounding.
# (NOTE: ZhanqiuHu) HeteroTP can cause different kernel block counts
# due to logical block rounding.
# Example: 640 prompt tokens, kernel_block_size=64
# remote physical_per_logical=10, local physical_per_logical=6
# remote logical ids from kv_transfer_params = [0]
# local logical ids allocated = [0, 1]
# remote kernel blocks: [0..9] (1*10=10)
# local kernel blocks: [0..11] (2*6=12)
# prefill physical_per_logical=10, decode physical_per_logical=6
# prefill logical ids from kv_transfer_params = [0]
# decode logical ids allocated = [0, 1]
# prefill kernel blocks: [0..9] (1*10=10)
# decode kernel blocks: [0..11] (2*6=12)
# actual data blocks = ceil(640/64) = 10, trim both to 10
# Vice versa (remote physical_per_logical=6, local=10):
# remote logical ids = [0, 1], local logical ids = [0]
# remote kernel blocks: [0..11] (2*6=12)
# local kernel blocks: [0..9] (1*10=10)
# Vice versa (prefill physical_per_logical=6, decode=10):
# prefill logical ids = [0, 1], decode logical ids = [0]
# prefill kernel blocks: [0..11] (2*6=12)
# decode kernel blocks: [0..9] (1*10=10)
# actual data blocks = ceil(640/64) = 10, trim both to 10
local_block_ids = list(local_block_ids)
for i, remote_group in enumerate(remote_block_ids):
num_local_blocks = len(local_block_ids[i])
num_remote_blocks = len(remote_group)
decode_block_ids = list(decode_block_ids)
for i, prefill_group in enumerate(prefill_block_ids):
num_decode_blocks = len(decode_block_ids[i])
num_prefill_blocks = len(prefill_group)
if _is_ssm_spec(self._group_spec_types[i]):
if num_local_blocks == num_remote_blocks:
if num_decode_blocks == num_prefill_blocks:
continue
# Only state-bearing slots reach here, single-state modes
# just one (see get_exchange_clipped_blocks), so differing
# counts mean position-indexed "all"-mode lists. A longer
# remote list carries earlier positions the local side
# already has (prefix hit) -> read its tail; a longer local
# prefill list carries earlier positions the decode side
# already has (prefix hit) -> take its tail; a longer decode
# list holds the position D recomputes itself, which gets
# no remote state.
assert num_local_blocks - num_remote_blocks <= 1, (
# no prefill state.
assert num_decode_blocks - num_prefill_blocks <= 1, (
f"Group {i}: unpairable SSM state slots, "
f"local={num_local_blocks} remote={num_remote_blocks}"
f"decode={num_decode_blocks} prefill={num_prefill_blocks}"
)
num_blocks = min(num_local_blocks, num_remote_blocks)
if num_local_blocks < num_remote_blocks:
remote_block_ids[i] = remote_group[-num_blocks:]
num_blocks = min(num_decode_blocks, num_prefill_blocks)
if num_decode_blocks < num_prefill_blocks:
prefill_block_ids[i] = prefill_group[-num_blocks:]
else:
local_block_ids[i] = local_block_ids[i][:num_blocks]
decode_block_ids[i] = decode_block_ids[i][:num_blocks]
elif (
self._physical_blocks_per_logical_kv_block
== remote_physical_per_logical
and num_local_blocks < num_remote_blocks
decode_physical_per_logical == prefill_physical_per_logical
and num_decode_blocks < num_prefill_blocks
):
# Partial prefix cache hit for FA group.
remote_block_ids[i] = remote_group[-num_local_blocks:]
prefill_block_ids[i] = prefill_group[-num_decode_blocks:]
else:
# TODO Handle prefix caching with different block_sizes
# Allocation rounding legitimately leaves up to
# ppl - 1 trailing dead kernel blocks per side (plus one
# extra local block for the recomputed final token), so
# extra decode block for the recomputed final token), so
# the counts may differ by up to the sum of the two
# ratios; anything larger indicates mismatched lists.
max_padding = (
self._physical_blocks_per_logical_kv_block
+ remote_physical_per_logical
decode_physical_per_logical + prefill_physical_per_logical
)
assert abs(num_local_blocks - num_remote_blocks) <= max_padding, (
f"Group {i}: |{num_local_blocks} - "
f"{num_remote_blocks}| > {max_padding}"
assert abs(num_decode_blocks - num_prefill_blocks) <= max_padding, (
f"Group {i}: |{num_decode_blocks} - "
f"{num_prefill_blocks}| > {max_padding}"
)
num_blocks = min(num_local_blocks, num_remote_blocks)
local_block_ids[i] = local_block_ids[i][:num_blocks]
remote_block_ids[i] = remote_group[:num_blocks]
return local_block_ids, remote_block_ids
num_blocks = min(num_decode_blocks, num_prefill_blocks)
decode_block_ids[i] = decode_block_ids[i][:num_blocks]
prefill_block_ids[i] = prefill_group[:num_blocks]
return decode_block_ids, prefill_block_ids
def get_kv_connector_stats(self) -> KVConnectorStats | None:
"""
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+4 -2
View File
@@ -289,9 +289,11 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
== len(local_block_ids)
== len(self.kv_cache_config.kv_cache_groups)
)
remote_physical_per_logical = remote_info.remote_physical_blocks_per_logical
local_block_ids, remote_block_ids = self._apply_prefix_caching(
local_block_ids, remote_block_ids, remote_physical_per_logical
decode_block_ids=local_block_ids,
prefill_block_ids=remote_block_ids,
decode_physical_per_logical=self._physical_blocks_per_logical_kv_block,
prefill_physical_per_logical=remote_info.remote_physical_blocks_per_logical,
)
# NOTE (nicolo) With homogeneous TP, each TP worker loads KV from
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+20 -12
View File
@@ -454,10 +454,6 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
# how pull-mode transfers touch _engine_last_active in start_load_kv.
self._engine_last_active[decode_engine_id] = time.perf_counter()
# Both sides stay logical here; ``_xfer_blocks_for_req`` converts each
# to physical with its own physical-blocks-per-logical ratio -- P uses
# ``self._physical_blocks_per_logical_kv_block``, D's is learned during
# the NIXL handshake.
logical_local = self._as_grouped_block_ids(local_block_ids)
logical_remote = self._as_grouped_block_ids(remote_block_ids)
physical_local = self._logical_to_kernel_block_ids(
@@ -629,16 +625,28 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
logger.warning("No blocks to push for request %s", request_id)
return None
# Align per-group block counts for push.
# Prefix caching: D allocated only uncached blocks, so on a partial hit it
# sends fewer than P's. End-trim P's blocks to that same suffix so we WRITE only
# the uncomputed tail into D's slots. Runs on kernel ids, post-expansion.
remote_block_ids, local_block_ids = self._apply_prefix_caching(
decode_block_ids=remote_block_ids,
prefill_block_ids=local_block_ids,
decode_physical_per_logical=remote_info.remote_physical_blocks_per_logical,
prefill_physical_per_logical=self._physical_blocks_per_logical_kv_block,
)
local_block_ids = list(local_block_ids)
remote_block_ids = list(remote_block_ids)
for i in range(min(len(local_block_ids), len(remote_block_ids))):
num_local = len(local_block_ids[i])
num_remote = len(remote_block_ids[i])
if num_local > num_remote:
local_block_ids[i] = local_block_ids[i][:num_remote]
elif num_local < num_remote:
remote_block_ids[i] = remote_block_ids[i][:num_local]
assert len(local_block_ids) == len(remote_block_ids), (
f"push group-count mismatch for {request_id}: {len(local_block_ids)} "
f"local vs {len(remote_block_ids)} remote groups"
)
for i in range(len(local_block_ids)):
assert len(local_block_ids[i]) == len(remote_block_ids[i]), (
f"push block-count mismatch for {request_id} group {i}: "
f"{len(local_block_ids[i])} local vs "
f"{len(remote_block_ids[i])} remote blocks"
)
# Get descs ids.
remote_block_descs_ids = self._compute_desc_ids(