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 1, 23:45 UTC
[PD][Nixl] Mamba prefix caching mode support (#42554)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Jun 4, 2026 at 15:41 UTC
68f5e565c90b670911024a2558aa265575b3e71a
3 files changed
+97
-6
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+77
@@ -275,6 +275,83 @@ def test_apply_prefix_caching_mamba_hybrid(
275
)
276
277
278
+@pytest.mark.cpu_test
279
+@pytest.mark.parametrize(
280
+ "local_physical_per_logical,remote_physical_per_logical,"
281
+ "local_block_ids,remote_block_ids,"
282
+ "expected_local,expected_remote",
283
+ [
284
+ # SSM prefix caching: remote has 3 placeholder + 1 real block,
285
+ # local has only the 1 real block. FA blocks are equal (no trim).
286
+ pytest.param(
287
+ 10,
288
+ 10,
289
+ [list(range(10)), [42]],
290
+ [list(range(10)), [40, 41, 42, 43]],
291
+ [list(range(10)), [42]],
292
+ [list(range(10)), [43]],
293
+ id="ssm_prefix_trim_only",
294
+ ),
295
+ # FA partial prefix cache hit with homogeneous TP: local has 4 FA
296
+ # blocks (prefix cached), remote has full 10. SSM equal (no trim).
297
+ pytest.param(
298
+ 10,
299
+ 10,
300
+ [list(range(6, 10)), [42]],
301
+ [list(range(10)), [42]],
302
+ [list(range(6, 10)), [42]],
303
+ [list(range(6, 10)), [42]],
304
+ id="fa_prefix_hit_homo_tp",
305
+ ),
306
+ # Both: FA partial prefix hit + SSM placeholder trim.
307
+ # local FA=[6..9] (4 blocks, prefix cached), remote FA=[0..9]
308
+ # local SSM=[99], remote SSM=[10, 20, 99] (2 placeholders + real)
309
+ pytest.param(
310
+ 10,
311
+ 10,
312
+ [[6, 7, 8, 9], [99]],
313
+ [list(range(10)), [10, 20, 99]],
314
+ [[6, 7, 8, 9], [99]],
315
+ [[6, 7, 8, 9], [99]],
316
+ id="fa_prefix_hit_and_ssm_trim",
317
+ ),
318
+ ],
319
+)
320
+def test_apply_prefix_caching_ssm_prefix_cache_hit(
321
+ local_physical_per_logical,
322
+ remote_physical_per_logical,
323
+ local_block_ids,
324
+ remote_block_ids,
325
+ expected_local,
326
+ expected_remote,
327
+):
328
+ """_apply_prefix_caching end-trims SSM remote blocks to match the single
329
+ local block (placeholders dropped) and end-trims FA remote blocks on
330
+ partial prefix cache hits when physical_per_logical matches.
331
+ """
332
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
333
+ NixlConnectorWorker,
334
+ )
335
+ from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
336
+
337
+ worker = object.__new__(NixlConnectorWorker)
338
+ worker._has_mamba = True
339
+ worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
340
+ worker._group_spec_types = (FullAttentionSpec, MambaSpec)
341
+ worker.kv_cache_config = make_kv_cache_config(block_size=16, mamba_enabled=True)
342
+
343
+ aligned_local, aligned_remote = worker._apply_prefix_caching(
344
+ local_block_ids, remote_block_ids, remote_physical_per_logical
345
+ )
346
+
347
+ assert aligned_local == expected_local, (
348
+ f"Expected local {expected_local}, got {aligned_local}"
349
+ )
350
+ assert aligned_remote == expected_remote, (
351
+ f"Expected remote {expected_remote}, got {aligned_remote}"
352
+ )
353
+
354
+
355
@pytest.mark.cpu_test
356
@pytest.mark.parametrize(
357
"local_physical_per_logical,remote_physical_per_logical,"
vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py
+18
-2
@@ -2333,9 +2333,25 @@ class NixlConnectorWorker:
2333
for i, remote_group in enumerate(remote_block_ids):
2334
num_local_blocks = len(local_block_ids[i])
2335
num_remote_blocks = len(remote_group)
2336
- if _is_ssm_spec(self._group_spec_types[i]):
2337
- assert num_local_blocks == num_remote_blocks
2336
+ if (
2337
+ _is_ssm_spec(self._group_spec_types[i])
2338
+ and num_local_blocks < num_remote_blocks
2339
+ ):
2340
+ # NOTE (NickLucche): With prefix caching on SSM, (remote) blocks
2341
+ # prior to the last one are placeholders (null blocks). Mind that
2342
+ # this doesn't really impact transfer, as we only still care about
2343
+ # the last "block", the full in-place state.
2344
+ assert num_local_blocks == 1, "SSM can only have one local block"
2345
+ remote_block_ids[i] = remote_group[-num_local_blocks:]
2346
+ elif (
2347
+ self._physical_blocks_per_logical_kv_block
2348
+ == remote_physical_per_logical
2349
+ and num_local_blocks < num_remote_blocks
2350
+ ):
2351
+ # Partial prefix cache hit for FA group.
2352
+ remote_block_ids[i] = remote_group[-num_local_blocks:]
2353
else:
2354
+ # TODO Handle prefix caching with different block_sizes
2355
max_padding = max(
2356
self._physical_blocks_per_logical_kv_block,
2357
remote_physical_per_logical,
vllm/v1/core/sched/scheduler.py
+2
-4
@@ -293,9 +293,6 @@ class Scheduler(SchedulerInterface):
293
num_new_local_computed_tokens: int = 0,
294
num_external_computed_tokens: int = 0,
295
) -> int:
296
- assert num_external_computed_tokens == 0, (
297
- "External KV connector is not verified yet"
298
- )
296
num_computed_tokens = (
297
request.num_computed_tokens
298
+ num_new_local_computed_tokens
@@ -716,7 +713,8 @@ class Scheduler(SchedulerInterface):
713
# The request cannot be scheduled.
714
break
715
719
- if self.need_mamba_block_aligned_split:
716
+ # Skip block alignment when setting up async receive (no local work).
717
+ if self.need_mamba_block_aligned_split and not load_kv_async:
718
num_new_tokens = self._mamba_block_aligned_split(
719
request,
720
num_new_tokens,