@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a6183563b6
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 1, 23:45 UTC

[Prefix Caching] DeepSeekv4 - Support selective prefix-cache retention for sliding-window KV cache (#43447)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>

Wei Zhao committed Jun 4, 2026 at 03:48 UTC a6183563b6f604ef7b481ce8ce7af359c6dc1b74
7 files changed +792 -45
tests/v1/core/test_kv_cache_utils.py
+37
@@ -358,6 +358,43 @@ def test_free_kv_cache_block_queue_append_n():
358 )
359
360
361 +def test_free_kv_cache_block_queue_prepend_n():
362 + # Seed the queue with one block so prepend has an existing head to splice
363 + # in front of (fake_head->b0->fake_tail).
364 + blocks = [KVCacheBlock(block_id=i) for i in range(6)]
365 + queue = FreeKVCacheBlockQueue(blocks[0:1])
366 +
367 + # Prepend 0 blocks is a no-op.
368 + queue.prepend_n([])
369 + assert queue.num_free_blocks == 1
370 + assert queue.fake_free_list_head.next_free_block is blocks[0]
371 +
372 + # Prepend 2 blocks; they land in front of the existing head, in order.
373 + # fake_head->b4->b5->b0->fake_tail
374 + queue.prepend_n(blocks[4:6])
375 + assert queue.num_free_blocks == 3
376 + assert queue.fake_free_list_head.next_free_block is blocks[4]
377 + assert blocks[4].prev_free_block is queue.fake_free_list_head
378 + assert blocks[4].next_free_block is blocks[5]
379 + assert blocks[5].prev_free_block is blocks[4]
380 + assert blocks[5].next_free_block is blocks[0]
381 + assert blocks[0].prev_free_block is blocks[5]
382 + assert blocks[0].next_free_block is queue.fake_free_list_tail
383 + assert queue.fake_free_list_tail.prev_free_block is blocks[0]
384 +
385 + # A second prepend goes ahead of everything previously prepended.
386 + # fake_head->b1->b2->b4->b5->b0->fake_tail
387 + queue.prepend_n(blocks[1:3])
388 + assert queue.num_free_blocks == 5
389 + assert queue.fake_free_list_head.next_free_block is blocks[1]
390 + assert blocks[1].next_free_block is blocks[2]
391 + assert blocks[2].next_free_block is blocks[4]
392 +
393 + # The popleft order reflects the front-to-back queue order.
394 + assert [queue.popleft().block_id for _ in range(5)] == [1, 2, 4, 5, 0]
395 + assert queue.num_free_blocks == 0
396 +
397 +
398 def test_free_kv_cache_block_queue_popleft_n():
399 blocks = [KVCacheBlock(block_id=i) for i in range(6)]
400 # Create an empty FreeKVCacheBlockQueue with these blocks
tests/v1/core/test_prefix_caching.py
+557
@@ -39,6 +39,7 @@ from vllm.v1.kv_cache_interface import (
39 KVCacheGroupSpec,
40 KVCacheSpecKind,
41 MambaSpec,
42 + MLAAttentionSpec,
43 SlidingWindowSpec,
44 )
45
@@ -2875,6 +2876,350 @@ def test_hybrid_cache_blocks_clamped_to_lcm():
2876 )
2877
2878
2879 +def test_hybrid_local_kv_retention_interval_aligns_in_manager(monkeypatch):
2880 + """Verify fixed intervals retain sparse tails plus the latest replay tail."""
2881 + monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "64")
2882 + block_size = 8
2883 + kv_cache_config = KVCacheConfig(
2884 + num_blocks=100,
2885 + kv_cache_tensors=[],
2886 + kv_cache_groups=[
2887 + KVCacheGroupSpec(
2888 + ["layer1"],
2889 + FullAttentionSpec(
2890 + block_size=4 * block_size,
2891 + num_kv_heads=1,
2892 + head_size=1,
2893 + dtype=torch.float16,
2894 + ),
2895 + ),
2896 + KVCacheGroupSpec(
2897 + ["layer2"],
2898 + SlidingWindowSpec(
2899 + block_size=block_size,
2900 + num_kv_heads=1,
2901 + head_size=1,
2902 + dtype=torch.float32,
2903 + sliding_window=block_size,
2904 + ),
2905 + ),
2906 + ],
2907 + )
2908 + manager = make_kv_cache_manager(
2909 + kv_cache_config=kv_cache_config,
2910 + max_model_len=8192,
2911 + enable_caching=True,
2912 + hash_block_size=block_size,
2913 + )
2914 +
2915 + # The SWA manager uses the configured 64-token interval (a multiple of the
2916 + # 32-token lcm_block_size) as its retention segment. For this 128-token
2917 + # prompt, the retained SWA tails are the 64-token interval boundary, the
2918 + # 96-token replay boundary, and the 128-token interval boundary.
2919 + token_ids = [i for i in range(16) for _ in range(block_size)]
2920 + req = make_request("0", token_ids, block_size, sha256)
2921 + computed_blocks, _ = manager.get_computed_blocks(req)
2922 + blocks = manager.allocate_slots(
2923 + req,
2924 + len(token_ids),
2925 + len(computed_blocks.blocks[0]) * block_size,
2926 + computed_blocks,
2927 + )
2928 + assert blocks is not None
2929 +
2930 + pool = manager.block_pool
2931 + expected_swa_cached = {7, 11, 15}
2932 + for i in range(16):
2933 + cached = pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[1])
2934 + if i in expected_swa_cached:
2935 + assert cached is not None, f"SWA hash {i} should be cached"
2936 + else:
2937 + assert cached is None, f"SWA hash {i} should not be cached"
2938 +
2939 +
2940 +@pytest.mark.parametrize(
2941 + "interval, expected_match",
2942 + [
2943 + # scheduler_block_size is 32 (= lcm(4*8, 8)); 33 is not a multiple of it.
2944 + ("33", "multiple of scheduler_block_size"),
2945 + # A negative multiple (-32 % 32 == 0) must still be rejected explicitly,
2946 + # otherwise it would pass the modulo check and silently degrade to dense.
2947 + ("-32", "non-negative"),
2948 + ],
2949 +)
2950 +def test_hybrid_local_kv_retention_interval_rejects_invalid(
2951 + monkeypatch, interval, expected_match
2952 +):
2953 + """A retention interval that is negative or not a multiple of
2954 + scheduler_block_size errors out at construction time."""
2955 + monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", interval)
2956 + block_size = 8
2957 + kv_cache_config = KVCacheConfig(
2958 + num_blocks=100,
2959 + kv_cache_tensors=[],
2960 + kv_cache_groups=[
2961 + KVCacheGroupSpec(
2962 + ["layer1"],
2963 + FullAttentionSpec(
2964 + block_size=4 * block_size,
2965 + num_kv_heads=1,
2966 + head_size=1,
2967 + dtype=torch.float16,
2968 + ),
2969 + ),
2970 + KVCacheGroupSpec(
2971 + ["layer2"],
2972 + SlidingWindowSpec(
2973 + block_size=block_size,
2974 + num_kv_heads=1,
2975 + head_size=1,
2976 + dtype=torch.float32,
2977 + sliding_window=block_size,
2978 + ),
2979 + ),
2980 + ],
2981 + )
2982 + with pytest.raises(ValueError, match=expected_match):
2983 + make_kv_cache_manager(
2984 + kv_cache_config=kv_cache_config,
2985 + max_model_len=8192,
2986 + enable_caching=True,
2987 + hash_block_size=block_size,
2988 + )
2989 +
2990 +
2991 +def test_hybrid_local_kv_retention_interval_survives_recycling(monkeypatch):
2992 + """Verify retained local checkpoints are reused after block recycling."""
2993 + monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "1024")
2994 + hash_block_size = 4
2995 + kv_cache_config = KVCacheConfig(
2996 + num_blocks=800,
2997 + kv_cache_tensors=[],
2998 + kv_cache_groups=[
2999 + KVCacheGroupSpec(
3000 + ["full"],
3001 + MLAAttentionSpec(
3002 + block_size=64 * hash_block_size,
3003 + num_kv_heads=1,
3004 + head_size=1,
3005 + dtype=torch.uint8,
3006 + compress_ratio=4,
3007 + ),
3008 + ),
3009 + KVCacheGroupSpec(
3010 + ["swa"],
3011 + SlidingWindowSpec(
3012 + block_size=16 * hash_block_size,
3013 + num_kv_heads=1,
3014 + head_size=1,
3015 + dtype=torch.float32,
3016 + sliding_window=512,
3017 + ),
3018 + ),
3019 + KVCacheGroupSpec(
3020 + ["c128"],
3021 + SlidingWindowSpec(
3022 + block_size=2 * hash_block_size,
3023 + num_kv_heads=1,
3024 + head_size=1,
3025 + dtype=torch.float32,
3026 + sliding_window=128,
3027 + ),
3028 + ),
3029 + KVCacheGroupSpec(
3030 + ["c4"],
3031 + SlidingWindowSpec(
3032 + block_size=hash_block_size,
3033 + num_kv_heads=1,
3034 + head_size=1,
3035 + dtype=torch.float32,
3036 + sliding_window=8,
3037 + ),
3038 + ),
3039 + ],
3040 + )
3041 + manager = make_kv_cache_manager(
3042 + kv_cache_config=kv_cache_config,
3043 + max_model_len=4096,
3044 + enable_caching=True,
3045 + hash_block_size=hash_block_size,
3046 + )
3047 +
3048 + def fill_request(request_id: str, token_offset: int) -> list[int]:
3049 + token_ids = [
3050 + token_offset + i for i in range(1024) for _ in range(hash_block_size)
3051 + ]
3052 + fill_req = make_request(request_id, token_ids, hash_block_size, sha256)
3053 + while fill_req.num_computed_tokens < len(token_ids):
3054 + num_new_tokens = min(512, len(token_ids) - fill_req.num_computed_tokens)
3055 + blocks = manager.allocate_slots(fill_req, num_new_tokens)
3056 + assert blocks is not None
3057 + fill_req.num_computed_tokens += num_new_tokens
3058 + manager.free(fill_req)
3059 + return token_ids
3060 +
3061 + token_ids = fill_request("fill_0", 0)
3062 + replay_req = make_request("replay", token_ids[:1800], hash_block_size, sha256)
3063 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(replay_req)
3064 + assert num_computed_tokens == 1024
3065 + assert [len(blocks) for blocks in computed_blocks.blocks] == [4, 16, 128, 256]
3066 +
3067 + fill_request("fill_1", 100_000)
3068 + replay_req = make_request("replay_again", token_ids[:1800], hash_block_size, sha256)
3069 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(replay_req)
3070 + assert num_computed_tokens == 1024
3071 + assert [len(blocks) for blocks in computed_blocks.blocks] == [4, 16, 128, 256]
3072 +
3073 +
3074 +def test_hybrid_local_kv_retention_latest_only_reuses_replay_boundary(monkeypatch):
3075 + """Verify latest-only retention reuses only the replayable prompt boundary."""
3076 + monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "0")
3077 + block_size = 8
3078 + kv_cache_config = KVCacheConfig(
3079 + num_blocks=100,
3080 + kv_cache_tensors=[],
3081 + kv_cache_groups=[
3082 + KVCacheGroupSpec(
3083 + ["layer1"],
3084 + FullAttentionSpec(
3085 + block_size=4 * block_size,
3086 + num_kv_heads=1,
3087 + head_size=1,
3088 + dtype=torch.float16,
3089 + ),
3090 + ),
3091 + KVCacheGroupSpec(
3092 + ["layer2"],
3093 + SlidingWindowSpec(
3094 + block_size=block_size,
3095 + num_kv_heads=1,
3096 + head_size=1,
3097 + dtype=torch.float32,
3098 + sliding_window=block_size,
3099 + ),
3100 + ),
3101 + ],
3102 + )
3103 + manager = make_kv_cache_manager(
3104 + kv_cache_config=kv_cache_config,
3105 + max_model_len=8192,
3106 + enable_caching=True,
3107 + hash_block_size=block_size,
3108 + )
3109 +
3110 + token_ids = [i for i in range(16) for _ in range(block_size)]
3111 + req0 = make_request("0", token_ids, block_size, sha256)
3112 + computed_blocks, _ = manager.get_computed_blocks(req0)
3113 + blocks = manager.allocate_slots(
3114 + req0,
3115 + len(token_ids),
3116 + len(computed_blocks.blocks[0]) * block_size,
3117 + computed_blocks,
3118 + )
3119 + assert blocks is not None
3120 +
3121 + pool = manager.block_pool
3122 + expected_swa_cached = {11}
3123 + for i in range(16):
3124 + cached = pool.get_cached_block(req0.block_hashes[i], kv_cache_group_ids=[1])
3125 + if i in expected_swa_cached:
3126 + assert cached is not None, f"SWA hash {i} should be cached"
3127 + else:
3128 + assert cached is None, f"SWA hash {i} should not be cached"
3129 +
3130 + manager.free(req0)
3131 + retained_swa_block = pool.get_cached_block(req0.block_hashes[11], [1])
3132 + assert retained_swa_block is not None
3133 + assert retained_swa_block[0].ref_cnt == 0
3134 +
3135 + req1 = make_request("1", token_ids, block_size, sha256)
3136 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
3137 + # Full prompt hits intentionally recompute the final block for logits, so
3138 + # the longest usable hit is the previous LCM boundary: 96 tokens.
3139 + assert num_computed_tokens == 12 * block_size
3140 + assert len(computed_blocks.blocks[1]) == 12
3141 +
3142 + shorter_req = make_request("2", token_ids[: 12 * block_size], block_size, sha256)
3143 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(shorter_req)
3144 + assert num_computed_tokens == 0
3145 + assert len(computed_blocks.blocks[1]) == 0
3146 +
3147 +
3148 +def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary(monkeypatch):
3149 + """Verify MTP/EAGLE SWA retention keeps the extra proof block.
3150 +
3151 + EAGLE/MTP lookup matches one additional local block after the returned
3152 + prefix and then drops it. Sparse retention must therefore cache the normal
3153 + local tail at the latest replay boundary plus one extra SWA block.
3154 + """
3155 + monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "0")
3156 + block_size = 8
3157 + kv_cache_config = KVCacheConfig(
3158 + num_blocks=100,
3159 + kv_cache_tensors=[],
3160 + kv_cache_groups=[
3161 + KVCacheGroupSpec(
3162 + ["full"],
3163 + FullAttentionSpec(
3164 + block_size=4 * block_size,
3165 + num_kv_heads=1,
3166 + head_size=1,
3167 + dtype=torch.float16,
3168 + ),
3169 + ),
3170 + KVCacheGroupSpec(
3171 + ["swa_mtp"],
3172 + SlidingWindowSpec(
3173 + block_size=block_size,
3174 + num_kv_heads=1,
3175 + head_size=1,
3176 + dtype=torch.float32,
3177 + sliding_window=block_size,
3178 + ),
3179 + is_eagle_group=True,
3180 + ),
3181 + ],
3182 + )
3183 + manager = make_kv_cache_manager(
3184 + kv_cache_config=kv_cache_config,
3185 + max_model_len=8192,
3186 + enable_caching=True,
3187 + hash_block_size=block_size,
3188 + use_eagle=True,
3189 + )
3190 +
3191 + # 127 tokens: latest replay boundary is floor((127 - 1) / 32) * 32 = 96.
3192 + # The EAGLE/MTP SWA lookup group must cache the local tail ending at
3193 + # 104 tokens, and that tail is two 8-token blocks wide: hashes 11 and 12.
3194 + token_ids = [i for i in range(15) for _ in range(block_size)] + [15] * 7
3195 + req0 = make_request("0", token_ids, block_size, sha256)
3196 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
3197 + assert num_computed_tokens == 0
3198 + blocks = manager.allocate_slots(
3199 + req0,
3200 + len(token_ids),
3201 + num_computed_tokens,
3202 + computed_blocks,
3203 + )
3204 + assert blocks is not None
3205 +
3206 + pool = manager.block_pool
3207 + expected_swa_cached = {11, 12}
3208 + for i in range(15):
3209 + cached = pool.get_cached_block(req0.block_hashes[i], kv_cache_group_ids=[1])
3210 + if i in expected_swa_cached:
3211 + assert cached is not None, f"SWA hash {i} should be cached"
3212 + else:
3213 + assert cached is None, f"SWA hash {i} should not be cached"
3214 +
3215 + manager.free(req0)
3216 +
3217 + req1 = make_request("1", token_ids, block_size, sha256)
3218 + computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
3219 + assert num_computed_tokens == 12 * block_size
3220 + assert [len(blocks) for blocks in computed_blocks.blocks] == [3, 12]
3221 +
3222 +
3223 def test_block_lookup_cache_single_block_per_key():
3224 cache = BlockHashToBlockMap()
3225 key0 = BlockHashWithGroupId(b"hash0")
@@ -3058,3 +3403,215 @@ def test_can_fit_full_sequence_full_attention_still_gates_oversized():
3403 req = make_request("oversized", list(range(prompt_len)), block_size, sha256)
3404
3405 assert manager.allocate_slots(req, block_size, full_sequence_must_fit=True) is None
3406 +
3407 +
3408 +def test_swa_free_split_keeps_cached_tail_ahead_of_scratch(monkeypatch):
3409 + """Default path (no retention): freeing an SWA request must place its
3410 + uncached scratch blocks at the front of the free queue (recycled first)
3411 + and keep its cached checkpoint blocks at the back (retained for prefix
3412 + hits). This split is always-on, independent of the retention interval."""
3413 + monkeypatch.delenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", raising=False)
3414 + block_size = 8
3415 + kv_cache_config = KVCacheConfig(
3416 + num_blocks=100,
3417 + kv_cache_tensors=[],
3418 + kv_cache_groups=[
3419 + KVCacheGroupSpec(
3420 + ["layer1"],
3421 + FullAttentionSpec(
3422 + block_size=4 * block_size,
3423 + num_kv_heads=1,
3424 + head_size=1,
3425 + dtype=torch.float16,
3426 + ),
3427 + ),
3428 + KVCacheGroupSpec(
3429 + ["layer2"],
3430 + SlidingWindowSpec(
3431 + block_size=block_size,
3432 + num_kv_heads=1,
3433 + head_size=1,
3434 + dtype=torch.float32,
3435 + sliding_window=block_size,
3436 + ),
3437 + ),
3438 + ],
3439 + )
3440 + manager = make_kv_cache_manager(
3441 + kv_cache_config=kv_cache_config,
3442 + max_model_len=8192,
3443 + enable_caching=True,
3444 + hash_block_size=block_size,
3445 + )
3446 +
3447 + token_ids = [i for i in range(16) for _ in range(block_size)]
3448 + req = make_request("0", token_ids, block_size, sha256)
3449 + computed_blocks, _ = manager.get_computed_blocks(req)
3450 + blocks = manager.allocate_slots(
3451 + req,
3452 + len(token_ids),
3453 + len(computed_blocks.blocks[0]) * block_size,
3454 + computed_blocks,
3455 + )
3456 + assert blocks is not None
3457 +
3458 + swa_manager = manager.coordinator.single_type_managers[1]
3459 + null_block = manager.block_pool.null_block
3460 + cached_ids: set[int] = set()
3461 + uncached_ids: set[int] = set()
3462 + cached_hash_indices: list[int] = []
3463 + for i, block in enumerate(swa_manager.req_to_blocks[req.request_id]):
3464 + if block is null_block:
3465 + continue
3466 + if block.block_hash is None:
3467 + uncached_ids.add(block.block_id)
3468 + else:
3469 + cached_ids.add(block.block_id)
3470 + cached_hash_indices.append(i)
3471 + # The dense default mask caches only the per-segment tails, so a 16-block
3472 + # SWA prompt must produce a mix of retained and scratch blocks.
3473 + assert cached_ids, "expected some retained (cached) SWA tail blocks"
3474 + assert uncached_ids, "expected some scratch (uncached) SWA blocks"
3475 +
3476 + manager.free(req)
3477 +
3478 + order = [
3479 + b.block_id for b in manager.block_pool.free_block_queue.get_all_free_blocks()
3480 + ]
3481 + pos = {bid: i for i, bid in enumerate(order)}
3482 + # Every scratch block is recycled before every retained block.
3483 + assert max(pos[bid] for bid in uncached_ids) < min(pos[bid] for bid in cached_ids)
3484 + # The retained tails survive the free and still serve a prefix-cache hit.
3485 + for i in cached_hash_indices:
3486 + assert (
3487 + manager.block_pool.get_cached_block(
3488 + req.block_hashes[i], kv_cache_group_ids=[1]
3489 + )
3490 + is not None
3491 + )
3492 +
3493 +
3494 +def _make_pure_swa_manager(block_size, sliding_window, num_blocks=100, **kwargs):
3495 + """Single sliding-window group (UnitaryKVCacheCoordinator)."""
3496 + kv_cache_config = KVCacheConfig(
3497 + num_blocks=num_blocks,
3498 + kv_cache_tensors=[],
3499 + kv_cache_groups=[
3500 + KVCacheGroupSpec(
3501 + ["layer"],
3502 + SlidingWindowSpec(
3503 + block_size=block_size,
3504 + num_kv_heads=1,
3505 + head_size=1,
3506 + dtype=torch.float32,
3507 + sliding_window=sliding_window,
3508 + ),
3509 + ),
3510 + ],
3511 + )
3512 + return make_kv_cache_manager(
3513 + kv_cache_config=kv_cache_config,
3514 + max_model_len=8192,
3515 + enable_caching=True,
3516 + hash_block_size=block_size,
3517 + **kwargs,
3518 + )
3519 +
3520 +
3521 +def test_pure_swa_retention_interval_caches_sparse_tails(monkeypatch):
3522 + """Sparse retention must work for a pure-SWA single-group model, not just
3523 + hybrid models: only the per-interval tails plus the latest replay tail are
3524 + cached, and a replay still hits the latest replayable boundary."""
3525 + monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "64")
3526 + block_size = 16
3527 + manager = _make_pure_swa_manager(block_size, sliding_window=block_size)
3528 + assert type(manager.coordinator).__name__ == "UnitaryKVCacheCoordinator"
3529 +
3530 + token_ids = [i for i in range(16) for _ in range(block_size)]
3531 + req = make_request("0", token_ids, block_size, sha256)
3532 + computed_blocks, _ = manager.get_computed_blocks(req)
3533 + blocks = manager.allocate_slots(
3534 + req,
3535 + len(token_ids),
3536 + len(computed_blocks.blocks[0]) * block_size,
3537 + computed_blocks,
3538 + )
3539 + assert blocks is not None
3540 +
3541 + pool = manager.block_pool
3542 + cached = {
3543 + i
3544 + for i in range(16)
3545 + if pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[0])
3546 + is not None
3547 + }
3548 + # per_segment = 64 / 16 = 4, need = cdiv(16-1, 16) = 1 -> segment tails at
3549 + # i%4==3 -> {3,7,11,15}; latest replay boundary (255//16*16 = 240) -> tail
3550 + # block 14. Crucially this is a strict subset of all 16 blocks: retention
3551 + # is actually sparse for pure SWA (not silently dense).
3552 + assert cached == {3, 7, 11, 14, 15}
3553 +
3554 + # A replay of the same prompt hits the latest replayable boundary (240).
3555 + replay = make_request("1", token_ids, block_size, sha256)
3556 + _, num_computed = manager.get_computed_blocks(replay)
3557 + assert num_computed == 240
3558 +
3559 +
3560 +def test_pure_swa_retention_latest_only(monkeypatch):
3561 + """`=0` on a pure-SWA model keeps only the latest replay tail."""
3562 + monkeypatch.setenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", "0")
3563 + block_size = 16
3564 + manager = _make_pure_swa_manager(block_size, sliding_window=block_size)
3565 +
3566 + token_ids = [i for i in range(16) for _ in range(block_size)]
3567 + req = make_request("0", token_ids, block_size, sha256)
3568 + computed_blocks, _ = manager.get_computed_blocks(req)
3569 + blocks = manager.allocate_slots(
3570 + req,
3571 + len(token_ids),
3572 + len(computed_blocks.blocks[0]) * block_size,
3573 + computed_blocks,
3574 + )
3575 + assert blocks is not None
3576 +
3577 + pool = manager.block_pool
3578 + cached = {
3579 + i
3580 + for i in range(16)
3581 + if pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[0])
3582 + is not None
3583 + }
3584 + # No segment tails (interval 0); only the latest replay tail (block 14).
3585 + assert cached == {14}
3586 +
3587 + replay = make_request("1", token_ids, block_size, sha256)
3588 + _, num_computed = manager.get_computed_blocks(replay)
3589 + assert num_computed == 240
3590 +
3591 +
3592 +def test_pure_swa_retention_dense_default_caches_all(monkeypatch):
3593 + """With retention unset, a pure-SWA model must keep the dense behavior:
3594 + every block boundary is a potential hit, so all blocks are cached."""
3595 + monkeypatch.delenv("VLLM_PREFIX_CACHE_RETENTION_INTERVAL", raising=False)
3596 + block_size = 16
3597 + manager = _make_pure_swa_manager(block_size, sliding_window=block_size)
3598 +
3599 + token_ids = [i for i in range(16) for _ in range(block_size)]
3600 + req = make_request("0", token_ids, block_size, sha256)
3601 + computed_blocks, _ = manager.get_computed_blocks(req)
3602 + blocks = manager.allocate_slots(
3603 + req,
3604 + len(token_ids),
3605 + len(computed_blocks.blocks[0]) * block_size,
3606 + computed_blocks,
3607 + )
3608 + assert blocks is not None
3609 +
3610 + pool = manager.block_pool
3611 + cached = {
3612 + i
3613 + for i in range(16)
3614 + if pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[0])
3615 + is not None
3616 + }
3617 + assert cached == set(range(16))
vllm/envs.py
+12
@@ -279,6 +279,7 @@ if TYPE_CHECKING:
279 VLLM_LORA_ENABLE_DUAL_STREAM: bool = False
280 VLLM_GPU_NIC_PCIE_MAPPING: str = ""
281 VLLM_NIC_SELECTION_VARS: str = ""
282 + VLLM_PREFIX_CACHE_RETENTION_INTERVAL: int | None = None
283
284
285 def get_default_cache_root():
@@ -1032,6 +1033,17 @@ environment_variables: dict[str, Callable[[], Any]] = {
1033 if "VLLM_PLUGINS" not in os.environ
1034 else os.environ["VLLM_PLUGINS"].split(",")
1035 ),
1036 + # Retain local sliding-window KV checkpoints for prefix caching.
1037 + # Unset (default) preserves the dense local checkpointing behavior. `0`
1038 + # retains only the latest completed prompt boundary. Positive values retain
1039 + # checkpoints at the specified interval boundaries (rounded up to the
1040 + # prefix-cache alignment).
1041 + # Applies to sliding-window attention for now but not yet Mamba/linear attention.
1042 + "VLLM_PREFIX_CACHE_RETENTION_INTERVAL": lambda: (
1043 + int(os.environ["VLLM_PREFIX_CACHE_RETENTION_INTERVAL"])
1044 + if "VLLM_PREFIX_CACHE_RETENTION_INTERVAL" in os.environ
1045 + else None
1046 + ),
1047 # a local directory to look in for unrecognized LoRA adapters.
1048 # only works if plugins are enabled and
1049 # VLLM_ALLOW_RUNTIME_LORA_UPDATING is enabled.
vllm/v1/core/block_pool.py
+12 -4
@@ -416,21 +416,29 @@ class BlockPool:
416 if self.metrics_collector:
417 self.metrics_collector.on_block_accessed(block)
418
419 - def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None:
419 + def free_blocks(
420 + self, ordered_blocks: Iterable[KVCacheBlock], prepend: bool = False
421 + ) -> None:
422 """Free a list of blocks. The blocks should be ordered by their
423 eviction priority, where the first block will be evicted first.
424
425 Args:
426 ordered_blocks: A list of blocks to free ordered by their eviction
427 priority.
428 + prepend: Whether to put newly-free blocks at the front of the free
429 + queue to be prioritized for reuse.
430 """
431 # Materialize the iterable to allow multiple passes.
432 blocks_list = list(ordered_blocks)
433 for block in blocks_list:
434 block.ref_cnt -= 1
431 - self.free_block_queue.append_n(
432 - [block for block in blocks_list if block.ref_cnt == 0 and not block.is_null]
433 - )
435 + freed_blocks = [
436 + block for block in blocks_list if block.ref_cnt == 0 and not block.is_null
437 + ]
438 + if prepend:
439 + self.free_block_queue.prepend_n(freed_blocks)
440 + else:
441 + self.free_block_queue.append_n(freed_blocks)
442
443 def evict_blocks(self, block_ids: set[int]) -> None:
444 """evict blocks from the prefix cache by their block IDs.
vllm/v1/core/kv_cache_coordinator.py
+52 -2
@@ -4,6 +4,7 @@ from abc import ABC, abstractmethod
4 from collections.abc import Sequence
5 from typing import NamedTuple
6
7 +from vllm import envs
8 from vllm.v1.core.block_pool import BlockPool
9 from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
10 from vllm.v1.core.kv_cache_utils import (
@@ -21,10 +22,41 @@ from vllm.v1.kv_cache_interface import (
22 FullAttentionSpec,
23 KVCacheConfig,
24 KVCacheSpec,
25 + SlidingWindowSpec,
26 )
27 from vllm.v1.request import Request
28
29
30 +def _validate_prefix_cache_retention_interval(
31 + retention_interval: int | None,
32 + scheduler_block_size: int,
33 + kv_cache_config: KVCacheConfig,
34 +) -> None:
35 + if retention_interval is None:
36 + return
37 +
38 + # Retention only sparsifies sliding-window checkpoints for now; every other
39 + # manager (full attention, Mamba, chunked-local) caches densely and
40 + # ignores it to be conservative.
41 + # TODO: Support Mamba/linear attention.
42 + if not any(
43 + isinstance(g.kv_cache_spec, SlidingWindowSpec)
44 + for g in kv_cache_config.kv_cache_groups
45 + ):
46 + raise ValueError(
47 + "VLLM_PREFIX_CACHE_RETENTION_INTERVAL is set but this model has "
48 + "no sliding-window KV cache group, so retention has no effect. "
49 + "Unset it (the feature only applies to sliding-window attention)."
50 + )
51 +
52 + if retention_interval < 0 or retention_interval % scheduler_block_size != 0:
53 + raise ValueError(
54 + f"VLLM_PREFIX_CACHE_RETENTION_INTERVAL ({retention_interval}) "
55 + "must be non-negative and a multiple of scheduler_block_size "
56 + f"({scheduler_block_size})."
57 + )
58 +
59 +
60 class KVCacheCoordinator(ABC):
61 """
62 Coordinate the KV cache of different KV cache groups.
@@ -86,6 +118,14 @@ class KVCacheCoordinator(ABC):
118 for i, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups)
119 )
120
121 + # A positive retention interval must be a multiple of the base hit granularity
122 + # (``scheduler_block_size``) to land on real cache-hit boundaries.
123 + # 0 = keep only the latest replay boundary; None = dense;
124 + self.retention_interval = envs.VLLM_PREFIX_CACHE_RETENTION_INTERVAL
125 + _validate_prefix_cache_retention_interval(
126 + self.retention_interval, self.scheduler_block_size, kv_cache_config
127 + )
128 +
129 def get_num_blocks_to_allocate(
130 self,
131 request_id: str,
@@ -215,7 +255,11 @@ class KVCacheCoordinator(ABC):
255 (including tokens that are already cached).
256 """
257 for manager in self.single_type_managers:
218 - manager.cache_blocks(request, num_computed_tokens)
258 + manager.cache_blocks(
259 + request,
260 + num_computed_tokens,
261 + retention_interval=self.retention_interval,
262 + )
263
264 def free(self, request_id: str) -> None:
265 """
@@ -525,8 +569,14 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
569 num_computed_tokens,
570 aligned_num_computed_tokens + manager.block_size,
571 )
572 + # The manager already knows the fine hit granularity
573 + # (``scheduler_block_size``); retention is passed separately so it
574 + # can keep both the coarse segment tails and the fine replay
575 + # boundary (which needs the fine value).
576 manager.cache_blocks(
529 - request, num_tokens_to_cache, alignment_tokens=self.scheduler_block_size
577 + request,
578 + num_tokens_to_cache,
579 + retention_interval=self.retention_interval,
580 )
581
582 def find_longest_cache_hit(
vllm/v1/core/kv_cache_utils.py
+21
@@ -327,6 +327,27 @@ class FreeKVCacheBlockQueue:
327
328 self.num_free_blocks += 1
329
330 + def prepend_n(self, blocks: list[KVCacheBlock]) -> None:
331 + """Put a list of blocks at the front of the free list."""
332 + if len(blocks) == 0:
333 + return
334 +
335 + first_block = self.fake_free_list_head.next_free_block
336 + assert first_block is not None, (
337 + "next_free_block of fake_free_list_head should always exist"
338 + )
339 +
340 + prev_block = self.fake_free_list_head
341 + for block in blocks:
342 + block.prev_free_block = prev_block
343 + prev_block.next_free_block = block
344 + prev_block = block
345 +
346 + prev_block.next_free_block = first_block
347 + first_block.prev_free_block = prev_block
348 +
349 + self.num_free_blocks += len(blocks)
350 +
351 def append_n(self, blocks: list[KVCacheBlock]) -> None:
352 """Put a list of blocks back into the free list
353
vllm/v1/core/single_type_kv_cache_manager.py
+101 -39
@@ -62,6 +62,7 @@ class SingleTypeKVCacheManager(ABC):
62 block until the request finishes.
63 """
64 self.scheduler_block_size = scheduler_block_size
65 + # The block size for this manager; used for actual block allocation.
66 self.block_size = kv_cache_spec.block_size
67 self.dcp_world_size = dcp_world_size
68 self.pcp_world_size = pcp_world_size
@@ -298,7 +299,7 @@ class SingleTypeKVCacheManager(ABC):
299 self,
300 request: Request,
301 num_tokens: int,
301 - alignment_tokens: int | None = None,
302 + retention_interval: int | None = None,
303 ) -> None:
304 """
305 Cache the blocks for the request.
@@ -307,12 +308,10 @@ class SingleTypeKVCacheManager(ABC):
308 request: The request.
309 num_tokens: The total number of tokens that need to be cached
310 (including tokens that are already cached).
310 - alignment_tokens: The cache-hit alignment (in tokens) used by the
311 - coordinator's ``find_longest_cache_hit``. When greater than
312 - this group's ``block_size``, managers whose hit logic only
313 - returns a subset of blocks per alignment-aligned segment
314 - (SWA) skip the rest since they can never participate in a
315 - future cache hit.
311 + retention_interval: Sparse local-checkpoint granularity. ``None``
312 + keeps dense checkpointing; ``0`` keeps only the latest replay
313 + boundary; a positive multiple of ``scheduler_block_size`` keeps
314 + a tail once per that-sized segment. Only SWA acts on it.
315 """
316 num_cached_blocks = self.num_cached_block.get(request.request_id, 0)
317 num_full_blocks = num_tokens // self.block_size
@@ -320,17 +319,15 @@ class SingleTypeKVCacheManager(ABC):
319 if num_cached_blocks >= num_full_blocks:
320 return
321
323 - # Fast path: when the coordinator imposes no alignment constraint
324 - if alignment_tokens is None or alignment_tokens <= self.block_size:
325 - block_mask = None
326 - else:
327 - block_mask = self.reachable_block_mask(
328 - num_cached_blocks,
329 - num_full_blocks,
330 - alignment_tokens,
331 - self.kv_cache_spec,
332 - self.use_eagle,
333 - )
322 + block_mask = self.reachable_block_mask(
323 + start_block=num_cached_blocks,
324 + end_block=num_full_blocks,
325 + alignment_tokens=self.scheduler_block_size,
326 + kv_cache_spec=self.kv_cache_spec,
327 + use_eagle=self.use_eagle,
328 + retention_interval=retention_interval,
329 + num_prompt_tokens=request.num_prompt_tokens,
330 + )
331 self.block_pool.cache_full_blocks(
332 request=request,
333 blocks=self.req_to_blocks[request.request_id],
@@ -347,10 +344,12 @@ class SingleTypeKVCacheManager(ABC):
344 def reachable_block_mask(
345 cls,
346 start_block: int,
350 - num_blocks: int,
351 - alignment_tokens: int,
347 + end_block: int,
348 + alignment_tokens: int | None,
349 kv_cache_spec: KVCacheSpec,
350 use_eagle: bool,
351 + retention_interval: int | None = None,
352 + num_prompt_tokens: int | None = None,
353 ) -> list[bool] | None:
354 """Per-block mask for ``cache_full_blocks``. ``None`` means cache
355 every (non-null) block — the default for full attention.
@@ -476,7 +475,12 @@ class SingleTypeKVCacheManager(ABC):
475 # range), so we must cap to the number of blocks that currently exist for
476 # this request.
477 num_skipped_blocks = min(num_skipped_blocks, len(blocks))
479 - removed_blocks: list[KVCacheBlock] = []
478 +
479 + # Reuse skipped local blocks in order:
480 + # scratch blocks: no prefix-cache value, reuse first.
481 + # cached blocks: reusable prefix-cache value, reuse last.
482 + removed_cached_blocks: list[KVCacheBlock] = []
483 + removed_uncached_blocks: list[KVCacheBlock] = []
484 # Because the block starts from index 0, the num_skipped_block-th block
485 # corresponds to index num_skipped_blocks - 1.
486 for i in range(num_skipped_blocks - 1, -1, -1):
@@ -485,9 +489,16 @@ class SingleTypeKVCacheManager(ABC):
489 # should also have been set to null blocks by the previous calls
490 # to this function.
491 break
488 - removed_blocks.append(blocks[i])
492 + if blocks[i].block_hash is None:
493 + removed_uncached_blocks.append(blocks[i])
494 + else:
495 + removed_cached_blocks.append(blocks[i])
496 blocks[i] = self._null_block
490 - self.block_pool.free_blocks(removed_blocks)
497 + # `prepend=True` makes uncached scratch blocks the next allocation
498 + # candidates, while cached blocks stay behind them as best-effort
499 + # prefix-cache entries.
500 + self.block_pool.free_blocks(removed_cached_blocks)
501 + self.block_pool.free_blocks(removed_uncached_blocks, prepend=True)
502
503 def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
504 """
@@ -677,30 +688,81 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
688 def reachable_block_mask(
689 cls,
690 start_block: int,
680 - num_blocks: int,
681 - alignment_tokens: int,
691 + end_block: int,
692 + alignment_tokens: int | None,
693 kv_cache_spec: KVCacheSpec,
694 use_eagle: bool,
695 + retention_interval: int | None = None,
696 + num_prompt_tokens: int | None = None,
697 ) -> list[bool] | None:
685 - assert alignment_tokens > kv_cache_spec.block_size
698 assert isinstance(kv_cache_spec, SlidingWindowSpec)
687 - per_segment = alignment_tokens // kv_cache_spec.block_size
699 + if alignment_tokens is None:
700 + # Fast path: when the coordinator imposes no alignment constraint.
701 + return None
702 + assert alignment_tokens % kv_cache_spec.block_size == 0
703 +
704 + block_size = kv_cache_spec.block_size
705 + # Contiguous blocks a hit needs at a boundary (incl. the EAGLE peek).
706 need = cls._contiguous_blocks_for_hit(
707 window_size=kv_cache_spec.sliding_window,
690 - block_size=kv_cache_spec.block_size,
708 + block_size=block_size,
709 use_eagle=use_eagle,
710 )
693 - if need >= per_segment:
694 - return None
711 # The matched run's right edge sits on the aligned boundary block when
712 # EAGLE peeks one block past it (shift=1), otherwise on the last block
697 - # before the boundary (shift=0). A block is reachable iff it falls in
698 - # the ``need``-wide run ending at some boundary's right edge.
713 + # before the boundary (shift=0).
714 shift = 1 if use_eagle else 0
700 - return [
701 - i >= shift and (i - shift) % per_segment >= per_segment - need
702 - for i in range(start_block, num_blocks)
703 - ]
715 +
716 + mask = [False] * (end_block - start_block)
717 +
718 + # (1) Segment-boundary tails. ``retention_interval``:
719 + # None -> dense (a tail at every ``alignment_tokens`` boundary);
720 + # 0 -> no dense tails (only the replay boundary below);
721 + # >0 -> a tail once per ``retention_interval``-sized segment.
722 + segment_tokens = (
723 + alignment_tokens
724 + if retention_interval is None
725 + else (None if retention_interval == 0 else retention_interval)
726 + )
727 + if segment_tokens is not None:
728 + per_segment = segment_tokens // block_size
729 + if need >= per_segment:
730 + # Every block is reachable; cache them all.
731 + return None
732 + for i in range(start_block, end_block):
733 + if i >= shift and (i - shift) % per_segment >= per_segment - need:
734 + mask[i - start_block] = True
735 +
736 + # (2) Replay-boundary tail. ``get_computed_blocks`` caps hits at
737 + # ``num_prompt - 1`` (to recompute the last token's logits), so an exact
738 + # prompt replay can only land on the latest *fine*-aligned boundary.
739 + # Sparse retention would otherwise skip it, so keep its tail explicitly.
740 + if retention_interval is not None and num_prompt_tokens is not None:
741 + latest = (num_prompt_tokens - 1) // alignment_tokens * alignment_tokens
742 + prompt_end_block = latest // block_size + shift
743 + for i in range(
744 + max(start_block, prompt_end_block - need),
745 + min(end_block, prompt_end_block),
746 + ):
747 + mask[i - start_block] = True
748 +
749 + return mask
750 +
751 + def free(self, request_id: str) -> None:
752 + # similar to remove_skipped_blocks(), prepend the uncached blocks
753 + # and append the cached blocks to the free queue
754 + req_blocks = self.req_to_blocks.pop(request_id, [])
755 + if req_blocks:
756 + cached_blocks: list[KVCacheBlock] = []
757 + uncached_blocks: list[KVCacheBlock] = []
758 + for block in reversed(req_blocks):
759 + if block.block_hash is None:
760 + uncached_blocks.append(block)
761 + else:
762 + cached_blocks.append(block)
763 + self.block_pool.free_blocks(cached_blocks)
764 + self.block_pool.free_blocks(uncached_blocks, prepend=True)
765 + self.num_cached_block.pop(request_id, None)
766
767 def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
768 """
@@ -1152,10 +1214,10 @@ class MambaManager(SingleTypeKVCacheManager):
1214 self,
1215 request: Request,
1216 num_tokens: int,
1155 - alignment_tokens: int | None = None,
1217 + retention_interval: int | None = None,
1218 ) -> None:
1219 num_cached_blocks_before = self.num_cached_block.get(request.request_id, 0)
1158 - super().cache_blocks(request, num_tokens, alignment_tokens=alignment_tokens)
1220 + super().cache_blocks(request, num_tokens, retention_interval=retention_interval)
1221 num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0)
1222 if num_cached_blocks_after > num_cached_blocks_before:
1223 for block in self.req_to_blocks[request.request_id][
@@ -1188,7 +1250,7 @@ class CrossAttentionManager(SingleTypeKVCacheManager):
1250 self,
1251 request: Request,
1252 num_tokens: int,
1191 - alignment_tokens: int | None = None,
1253 + retention_interval: int | None = None,
1254 ) -> None:
1255 # We do not cache blocks for cross-attention to be shared between
1256 # requests, so this method is not relevant.