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

[Bugfix][KV Offload] Skip non-prefix-cacheable groups in SimpleCPUOffload (GLM-5.3-Flash kpool tail and QSA) (#56810)

Signed-off-by: Jared Wen <jaredwen@inferact.ai> Signed-off-by: Jared Wen <w13431838023@gmail.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Zhewen Li <zhewenli@inferact.ai>

Jared Wen committed Sep 20, 2026 at 18:58 UTC f648eed23dc48fcc8ba64be0c4182f8e775b5bfa

3 files changed +267 -2

tests/v1/simple_kv_offload/test_kv_events.py

+36

@@ -19,8 +19,10 @@ from tests.v1.simple_kv_offload.test_scheduler import (

19 SchedulerFixture,

20 _alloc_and_register,

21 _allocate_cp_gpu_blocks,

22 + _allocate_gpu_blocks,

23 _make_cp_request,

24 _make_kv_cache_config,

25 + _make_scratch_kv_cache_config,

26 _make_vllm_config,

27 make_request,

28 make_scheduler_output,

@@ -757,3 +759,37 @@ def test_mamba_align_skips_positional_event_metadata() -> None:

759 get_block_hash(make_block_hash_with_group_id(req.block_hashes[idx], 0))

760 )

761 assert ev.block_hashes == [expected_hash]

762 +

763 +

764 +def test_scratch_group_store_emits_events_for_cacheable_groups_only() -> None:

765 + """Event metadata is resolved only for prefix-cacheable groups.

766 +

767 + A scratch group (GLM-5.3-Flash kpool tail) has a block size that does not

768 + divide the hash block size, so viewing the request hashes at its block

769 + size is impossible; the store must skip it and still emit BlockStored

770 + for the attention blocks.

771 + """

772 + fix = make_events_scheduler(kv_cache_config=_make_scratch_kv_cache_config(16))

773 + sched = fix.scheduler

774 + gpu_pool = fix.gpu_block_pool

775 + num_blocks = 2

776 +

777 + req = make_request(num_blocks=num_blocks)

778 + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)

779 + scratch_block = gpu_pool.get_new_blocks(1)

780 + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, scratch_block))

781 + req.num_computed_tokens = num_blocks * BLOCK_SIZE

782 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

783 + sched_out = make_scheduler_output(

784 + {req.request_id: num_blocks * BLOCK_SIZE},

785 + new_reqs={req.request_id: kv_blocks.get_block_ids()},

786 + )

787 + meta = sched.build_connector_meta(sched_out)

788 + assert sorted(meta.store_gpu_blocks) == sorted(b.block_id for b in fa_blocks)

789 + simulate_store_completion(sched, meta.store_event)

790 +

791 + events = list(sched.take_events())

792 + assert len(events) == num_blocks

793 + assert all(isinstance(e, BlockStored) for e in events)

794 + assert {e.group_idx for e in events} == {0}

795 + assert all(len(e.token_ids) == BLOCK_SIZE for e in events)

tests/v1/simple_kv_offload/test_scheduler.py

+222 -1

@@ -41,7 +41,9 @@ from vllm.v1.core.single_type_kv_cache_manager import (

41 register_all_kvcache_specs,

42 )

43 from vllm.v1.kv_cache_interface import (

44 + CircularBufferSpec,

45 FullAttentionSpec,

46 + KpoolTailSpec,

47 KVCacheConfig,

48 KVCacheGroupSpec,

49 KVCacheTensor,

@@ -125,6 +127,45 @@ def _make_kv_cache_config(

127 )

128

129

130 +def _make_scratch_kv_cache_config(

131 + num_blocks: int, scratch_block_size: int = 4

132 +) -> KVCacheConfig:

133 + """FullAttention group plus a non-prefix-cacheable scratch group.

134 +

135 + Mirrors GLM-5.3-Flash, whose kpool-tail group holds one per-request

136 + scratch block of ``index_kpool`` tokens: it is excluded from prefix

137 + caching, so its block size neither divides nor is divided by the

138 + hash block size in general.

139 + """

140 + register_all_kvcache_specs(vllm_config=None)

141 + fa_config = _make_kv_cache_config(num_blocks, num_groups=1)

142 + scratch_layers = ["layer_scratch"]

143 + scratch_spec = KpoolTailSpec(

144 + block_size=scratch_block_size,

145 + num_kv_heads=2,

146 + head_size=HEAD_SIZE,

147 + head_size_v=0,

148 + dtype=DTYPE,

149 + sliding_window=scratch_block_size,

150 + )

151 + assert not scratch_spec.prefix_cacheable

152 + scratch_bytes = scratch_spec.page_size_bytes * num_blocks

153 + return KVCacheConfig(

154 + num_blocks=num_blocks,

155 + kv_cache_tensors=fa_config.kv_cache_tensors

156 + + [

157 + KVCacheTensor(

158 + size=scratch_bytes,

159 + layers=scratch_layers,

160 + layer_stride=scratch_bytes,

161 + block_stride=scratch_spec.page_size_bytes,

162 + )

163 + ],

164 + kv_cache_groups=fa_config.kv_cache_groups

165 + + [KVCacheGroupSpec(scratch_layers, scratch_spec)],

166 + )

167 +

168 +

169 def _make_vllm_config(block_size: int = BLOCK_SIZE) -> VllmConfig:

170 """Minimal VllmConfig for scheduler tests (no GPU)."""

171 model_config = ModelConfig(

@@ -174,9 +215,13 @@ def make_scheduler(

215 num_gpu_blocks: int = 16,

216 num_groups: int = 1,

217 lazy: bool = False,

218 + kv_cache_config: KVCacheConfig | None = None,

219 ) -> SchedulerFixture:

220 """Build a SimpleCPUOffloadScheduler with small block pools."""

179 - kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups)

221 + if kv_cache_config is None:

222 + kv_cache_config = _make_kv_cache_config(num_gpu_blocks, num_groups)

223 + else:

224 + num_groups = len(kv_cache_config.kv_cache_groups)

225 vllm_config = _make_vllm_config()

226 cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks * num_groups

227

@@ -1392,6 +1437,63 @@ def test_multi_group_null_blocks_skipped() -> None:

1437 )

1438

1439

1440 +# ---------------------------------------------------------------------------

1441 +# Test 8b: Non-prefix-cacheable scratch groups take no part in store or load

1442 +# ---------------------------------------------------------------------------

1443 +def test_scratch_group_excluded_from_store_and_load() -> None:

1444 + """A scratch group (GLM-5.3-Flash kpool tail) is never stored or loaded.

1445 +

1446 + Its single per-request block carries no hash, so the store must offload

1447 + only the attention blocks and the load must pair only attention blocks.

1448 + """

1449 + fix = make_scheduler(

1450 + num_cpu_blocks=8,

1451 + num_gpu_blocks=16,

1452 + kv_cache_config=_make_scratch_kv_cache_config(16),

1453 + )

1454 + sched = fix.scheduler

1455 + gpu_pool = fix.gpu_block_pool

1456 + num_blocks = 2

1457 +

1458 + req = make_request(num_blocks=num_blocks)

1459 + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)

1460 + scratch_block = gpu_pool.get_new_blocks(1)

1461 + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, scratch_block))

1462 + req.num_computed_tokens = num_blocks * BLOCK_SIZE

1463 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

1464 + sched_out = make_scheduler_output(

1465 + {req.request_id: num_blocks * BLOCK_SIZE},

1466 + new_reqs={req.request_id: kv_blocks.get_block_ids()},

1467 + )

1468 + meta = sched.build_connector_meta(sched_out)

1469 + assert sorted(meta.store_gpu_blocks) == sorted(b.block_id for b in fa_blocks)

1470 + simulate_store_completion(sched, meta.store_event)

1471 +

1472 + req2 = Request(

1473 + request_id="req-scratch-load",

1474 + prompt_token_ids=req.prompt_token_ids,

1475 + sampling_params=req.sampling_params,

1476 + pooling_params=None,

1477 + mm_features=None,

1478 + block_hasher=req._block_hasher,

1479 + )

1480 + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)

1481 + assert hit_tokens == num_blocks * BLOCK_SIZE

1482 + assert is_async is True

1483 +

1484 + fa_blocks2 = gpu_pool.get_new_blocks(num_blocks)

1485 + scratch_block2 = gpu_pool.get_new_blocks(1)

1486 + kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, scratch_block2))

1487 + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)

1488 + sched_out2 = make_scheduler_output(

1489 + {req2.request_id: 1},

1490 + new_reqs={req2.request_id: kv_blocks2.get_block_ids()},

1491 + )

1492 + meta2 = sched.build_connector_meta(sched_out2)

1493 + assert sorted(meta2.load_gpu_blocks) == sorted(b.block_id for b in fa_blocks2)

1494 + assert len(meta2.load_cpu_blocks) == num_blocks

1495 +

1496 +

1497 # ---------------------------------------------------------------------------

1498 # Test 9: Chunked prefill accumulates block_ids across steps

1499 # ---------------------------------------------------------------------------

@@ -2844,3 +2946,122 @@ def test_finished_eager_store_recovers_nulled_window_tail() -> None:

2946 hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)

2947 assert hit_tokens == num_blocks * BLOCK_SIZE

2948 assert is_async is True

2949 +

2950 +

2951 +QSA_RING_BLOCK_SIZE = 4

2952 +

2953 +

2954 +def _make_qsa_hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig:

2955 + register_all_kvcache_specs(vllm_config=None)

2956 + full_spec = FullAttentionSpec(

2957 + block_size=BLOCK_SIZE,

2958 + num_kv_heads=NUM_KV_HEADS,

2959 + head_size=HEAD_SIZE,

2960 + dtype=DTYPE,

2961 + )

2962 + ring_spec = CircularBufferSpec(

2963 + block_size=QSA_RING_BLOCK_SIZE,

2964 + num_kv_heads=NUM_KV_HEADS,

2965 + head_size=HEAD_SIZE,

2966 + head_size_v=0,

2967 + dtype=DTYPE,

2968 + )

2969 + groups = [

2970 + KVCacheGroupSpec(["layer_0"], full_spec),

2971 + KVCacheGroupSpec(["qsa_ring"], ring_spec),

2972 + ]

2973 + tensors = [

2974 + KVCacheTensor(

2975 + size=group.kv_cache_spec.page_size_bytes * num_blocks,

2976 + layers=list(group.layer_names),

2977 + layer_stride=group.kv_cache_spec.page_size_bytes * num_blocks,

2978 + block_stride=group.kv_cache_spec.page_size_bytes,

2979 + )

2980 + for group in groups

2981 + ]

2982 + return KVCacheConfig(

2983 + num_blocks=num_blocks,

2984 + kv_cache_tensors=tensors,

2985 + kv_cache_groups=groups,

2986 + )

2987 +

2988 +

2989 +def _make_qsa_scheduler(

2990 + lazy: bool = False,

2991 +) -> tuple[SimpleCPUOffloadScheduler, BlockPool]:

2992 + kv_cache_config = _make_qsa_hybrid_kv_cache_config(num_blocks=16)

2993 + sched = SimpleCPUOffloadScheduler(

2994 + vllm_config=_make_vllm_config(),

2995 + kv_cache_config=kv_cache_config,

2996 + cpu_capacity_bytes=_BYTES_PER_BLOCK * 8 * 2,

2997 + scheduler_block_size=BLOCK_SIZE,

2998 + hash_block_size=BLOCK_SIZE,

2999 + lazy_offload=lazy,

3000 + )

3001 + gpu_pool = BlockPool(

3002 + num_gpu_blocks=16, enable_caching=True, hash_block_size=BLOCK_SIZE

3003 + )

3004 + sched.bind_gpu_block_pool(gpu_pool)

3005 + return sched, gpu_pool

3006 +

3007 +

3008 +def test_qsa_ring_group_is_never_stored_or_loaded() -> None:

3009 + sched, gpu_pool = _make_qsa_scheduler()

3010 + assert sched.prefix_cacheable_group_ids == (0,)

3011 +

3012 + num_blocks = 2

3013 + req = make_request(num_blocks=num_blocks)

3014 + fa_blocks = _allocate_gpu_blocks(gpu_pool, req, num_blocks, group_id=0)

3015 + ring_block = gpu_pool.get_new_blocks(1)[0]

3016 + kv_blocks = KVCacheBlocks(blocks=(fa_blocks, [ring_block]))

3017 + req.num_computed_tokens = num_blocks * BLOCK_SIZE

3018 + sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)

3019 +

3020 + sched_out = make_scheduler_output(

3021 + {req.request_id: num_blocks * BLOCK_SIZE},

3022 + new_reqs={req.request_id: kv_blocks.get_block_ids()},

3023 + )

3024 + meta = sched.build_connector_meta(sched_out)

3025 + assert set(meta.store_gpu_blocks) == {block.block_id for block in fa_blocks}

3026 + assert ring_block.block_id not in meta.store_gpu_blocks

3027 + assert meta.store_event >= 0

3028 + simulate_store_completion(sched, meta.store_event)

3029 +

3030 + req2 = Request(

3031 + request_id="req-qsa-load",

3032 + prompt_token_ids=req.prompt_token_ids,

3033 + sampling_params=req.sampling_params,

3034 + pooling_params=None,

3035 + mm_features=None,

3036 + block_hasher=req._block_hasher,

3037 + )

3038 + hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)

3039 + assert hit_tokens == num_blocks * BLOCK_SIZE

3040 + assert is_async is True

3041 +

3042 + fa_blocks2 = gpu_pool.get_new_blocks(num_blocks)

3043 + ring_block2 = gpu_pool.get_new_blocks(1)[0]

3044 + kv_blocks2 = KVCacheBlocks(blocks=(fa_blocks2, [ring_block2]))

3045 + sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)

3046 +

3047 + meta2 = sched.build_connector_meta(make_scheduler_output({req2.request_id: 1}))

3048 + assert set(meta2.load_gpu_blocks) == {block.block_id for block in fa_blocks2}

3049 + assert ring_block2.block_id not in meta2.load_gpu_blocks

3050 + assert len(meta2.load_cpu_blocks) == num_blocks

3051 +

3052 +

3053 +def test_lazy_target_blocks_ignore_non_prefix_cacheable_groups() -> None:

3054 + with_ring = _make_qsa_hybrid_kv_cache_config(num_blocks=16)

3055 + without_ring = KVCacheConfig(

3056 + num_blocks=16,

3057 + kv_cache_tensors=with_ring.kv_cache_tensors[:1],

3058 + kv_cache_groups=with_ring.kv_cache_groups[:1],

3059 + )

3060 + max_batched = 64

3061 + target_with = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(

3062 + with_ring, max_batched

3063 + )

3064 + target_without = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(

3065 + without_ring, max_batched

3066 + )

3067 + assert target_with == target_without

vllm/v1/simple_kv_offload/manager.py

+9 -1

@@ -165,6 +165,9 @@ class SimpleCPUOffloadScheduler:

165 kv_cache_config, offload_capacity

166 )

167 self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks

168 + self.prefix_cacheable_group_ids = (

169 + self.cpu_kv_cache_config.prefix_cacheable_group_ids

170 + )

171 self.kv_event_medium = MEDIUM_STORAGE if disk_capacity_bytes > 0 else MEDIUM_CPU

172 # Find the full attention kv group for prefix cache matching.

173 self.fa_gidx = -1

@@ -297,7 +300,7 @@ class SimpleCPUOffloadScheduler:

300 """GPU blocks to keep available (free/offloaded) per step in lazy mode."""

301 WATERMARK_RATIO = 1.0 # Reserve larger space to avoid running out of GPU blocks

302 target = 0

300 - for g in kv_cache_config.kv_cache_groups:

303 + for g in kv_cache_config.prefix_cacheable_groups:

304 spec = g.kv_cache_spec

305 # Only full attention is sharded across DCP ranks; replicated specs

306 # (mamba, sliding window, chunked-local) keep their own block size.

@@ -442,6 +445,9 @@ class SimpleCPUOffloadScheduler:

445 # the rest will be released along with the temp pin below.

446 cpu_hit_blocks: list[list[KVCacheBlock]] = []

447 for g in range(num_groups):

448 + if g not in self.prefix_cacheable_group_ids:

449 + cpu_hit_blocks.append([])

450 + continue

451 g_block_size = self.group_block_sizes[g]

452 n_take_g = cdiv(num_external_tokens, g_block_size)

453 cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g])

@@ -828,6 +834,8 @@ class SimpleCPUOffloadScheduler:

834 num_free = self.cpu_block_pool.get_num_free_blocks()

835

836 for g, group_gpu_ids in enumerate(block_ids_by_group):

837 + if g not in self.prefix_cacheable_group_ids:

838 + continue

839 if len(gpu_block_ids) >= num_free:

840 break

841 group_manager = self.cpu_coordinator.single_type_managers[g]