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

[KV Offload] Emit self-describing events for partial recurrent blocks (#51243)

Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>

Chauncey committed Aug 9, 2026 at 16:28 UTC c4239986427a40389fd79c1a4f6afd48aea6512e

5 files changed +181 -6

docs/features/kv_offloading_usage.md

+1 -1

@@ -76,7 +76,7 @@ vllm serve <model> \

76 | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |

77 | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |

79 -| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |

79 +| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Partial recurrent tails emit the hash-aligned portion from the physical block start through the tail boundary. Other sliding-window/SSM chunks keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |

80 | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |

82 ## Custom Eviction Policies

tests/v1/kv_connector/unit/offloading_connector/test_events.py

+52

@@ -208,6 +208,58 @@ def test_take_events_forwards_locality_to_placeholder_store():

208 assert events[0].locality == "REMOTE"

209

210

211 +def test_partial_tail_event_describes_hash_aligned_physical_block_prefix():

212 + tracker = _tracker()

213 + group_config = _group_config(block_size=16, blocks_per_chunk=1)._replace(

214 + hashes_per_chunk=4

215 + )

216 + req = _request(block_hashes=[_hash(i) for i in range(8)], token_count=32)

217 + key = make_offload_key(req.block_hashes[6], group_config.group_idx)

218 +

219 + tracker.record_partial_store(req, group_config, 28, key)

220 + [event] = tracker.take_events([_stored_event([key])])

221 +

222 + assert isinstance(event, BlockStored)

223 + assert event.block_hashes == [_wire_hash(_hash(i)) for i in range(4, 7)]

224 + assert event.parent_block_hash == _wire_hash(_hash(3))

225 + assert event.token_ids == list(range(17, 29))

226 + assert event.block_size == 4

227 +

228 +

229 +def test_partial_tail_lookup_does_not_overwrite_store_metadata():

230 + tracker = _tracker()

231 + group_config = _group_config()

232 + stored_req = _request(block_hashes=[_hash(0)], token_count=4)

233 + lookup_req = _request(block_hashes=[_hash(0)], token_count=4)

234 + lookup_req.all_token_ids = [9, 9, 9, 9]

235 + key = make_offload_key(stored_req.block_hashes[0], group_config.group_idx)

236 +

237 + tracker.record_partial_store(stored_req, group_config, 4, key)

238 + tracker.record_partial_lookup(lookup_req, group_config, 4, key)

239 + [event] = tracker.take_events([_stored_event([key])])

240 +

241 + assert isinstance(event, BlockStored)

242 + assert event.token_ids == [1, 2, 3, 4]

243 +

244 +

245 +@pytest.mark.parametrize(

246 + "record_method", ["record_partial_store", "record_partial_lookup"]

247 +)

248 +def test_partial_tail_sliding_window_event_uses_placeholder(record_method):

249 + tracker = _tracker()

250 + group_config = _group_config(sliding_window_size_in_chunks=1)

251 + req = _request(block_hashes=[_hash(0)], token_count=4)

252 + key = make_offload_key(req.block_hashes[0], group_config.group_idx)

253 +

254 + getattr(tracker, record_method)(req, group_config, 4, key)

255 + [event] = tracker.take_events([_stored_event([key])])

256 +

257 + assert isinstance(event, BlockStored)

258 + assert event.block_hashes == [_wire_hash(_hash(0))]

259 + assert event.token_ids == []

260 + assert event.block_size == 0

261 +

262 +

263 def test_take_events_forwards_locality_to_remove():

264 tracker = _tracker()

265 req = _request(block_hashes=[_hash(0)], token_count=4)

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

+35 -1

@@ -16,6 +16,7 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (

16 to_keys,

17 )

18 from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID

19 +from vllm.config import KVEventsConfig

20 from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored

21 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (

22 OffloadingConnectorMetadata,

@@ -60,9 +61,12 @@ from vllm.v1.request import RequestStatus

63 def _make_partial_tail_scheduler() -> OffloadingConnectorScheduler:

63 - vllm_config = _make_vllm_config()

64 + vllm_config = _make_vllm_config(extra_config={"self_describing_kv_events": True})

65 vllm_config.cache_config.prefix_match_unit = 4

66 vllm_config.speculative_config = None

67 + vllm_config.kv_events_config = KVEventsConfig(

68 + enable_kv_cache_events=True, publisher="null"

69 + )

70 kv_cache_config = _make_mamba_hybrid_kv_cache_config()

71 spec = MockOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))

72 return OffloadingConnectorScheduler(spec, vllm_config, kv_cache_config)

@@ -77,6 +81,8 @@ def _make_partial_tail_request(

81 request.num_prompt_tokens = 30

82 request.num_tokens = 30

83 request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(7)]

84 + request.all_token_ids = list(range(30))

85 + request.lora_request = None

86 request.is_finished.return_value = False

87 scheduler.on_new_request(request)

88 return request

@@ -119,6 +125,34 @@ def test_partial_tail_store_uses_attention_and_recurrent_cow_sources():

125 12: {job_id},

126 99: {job_id},

127 }

128 + assert scheduler.config.supports_partial_tail

129 +

130 + events = list(

131 + scheduler._events_tracker.take_events(

132 + [

133 + OffloadingEvent(

134 + keys=list(scheduler._jobs[job_id].keys),

135 + medium=Medium.CPU,

136 + removed=False,

137 + )

138 + ]

139 + )

140 + )

141 + assert len(events) == 2

142 + assert all(isinstance(event, BlockStored) for event in events)

143 + assert {event.group_idx for event in events} == {0, 1}

144 + events_by_group = {event.group_idx: event for event in events}

145 + full_attention_event = events_by_group[0]

146 + assert full_attention_event.block_size == 4

147 + assert full_attention_event.token_ids == list(range(16, 28))

148 + assert len(full_attention_event.block_hashes) == 3

149 + assert full_attention_event.parent_block_hash is not None

150 +

151 + recurrent_event = events_by_group[1]

152 + assert recurrent_event.block_size == 0

153 + assert recurrent_event.token_ids == []

154 + assert len(recurrent_event.block_hashes) == 1

155 + assert recurrent_event.parent_block_hash is None

156

157

158 def test_partial_lookup_returns_exact_boundary_and_group_load_keys():

vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py

+79

@@ -145,6 +145,85 @@ class OffloadingEventsTracker:

145 req, group_config, chunk_idx

146 )

147

148 + def record_partial_store(

149 + self,

150 + req: Request,

151 + group_config: "GroupOffloadConfig",

152 + boundary_tokens: int,

153 + offload_key: OffloadKey,

154 + ) -> None:

155 + """Snapshot metadata for a newly stored partial recurrent tail."""

156 + if group_config.sliding_window_size_in_chunks is not None:

157 + return

158 + self._record_partial_tail(req, group_config, boundary_tokens, offload_key)

159 +

160 + def record_partial_lookup(

161 + self,

162 + req: Request,

163 + group_config: "GroupOffloadConfig",

164 + boundary_tokens: int,

165 + offload_key: OffloadKey,

166 + ) -> None:

167 + """Backfill metadata for a partial recurrent tail lookup hit."""

168 + if group_config.sliding_window_size_in_chunks is not None:

169 + return

170 + if offload_key not in self._pending_event_metadata:

171 + self._record_partial_tail(req, group_config, boundary_tokens, offload_key)

172 +

173 + def _record_partial_tail(

174 + self,

175 + req: Request,

176 + group_config: "GroupOffloadConfig",

177 + boundary_tokens: int,

178 + offload_key: OffloadKey,

179 + ) -> None:

180 + """Build metadata for the valid prefix of one physical cache block.

181 +

182 + A partial recurrent tail ends on a hash boundary but before its

183 + physical cache block is full. The event describes only the valid

184 + hashes and tokens, not the unused remainder of that physical block.

185 + """

186 + if not self.self_describing_enabled:

187 + return

188 +

189 + tokens_per_hash = group_config.tokens_per_chunk // group_config.hashes_per_chunk

190 + # Subtract one so the boundary token itself cannot select the next

191 + # physical block when the boundary lies exactly on a block edge.

192 + chunk_start = (

193 + (boundary_tokens - 1) // group_config.tokens_per_chunk

194 + ) * group_config.tokens_per_chunk

195 + first_hash_idx = chunk_start // tokens_per_hash

196 + last_hash_idx = boundary_tokens // tokens_per_hash

197 + assert chunk_start < boundary_tokens

198 + assert boundary_tokens % tokens_per_hash == 0

199 + assert last_hash_idx <= len(req.block_hashes)

200 +

201 + # Unlike a complete chunk, a partial tail contains only the hashes

202 + # between its physical block start and its valid token boundary.

203 + maybe_block_hashes = req.block_hashes[first_hash_idx:last_hash_idx]

204 + block_hashes = tuple(

205 + block_hash for block_hash in maybe_block_hashes if block_hash is not None

206 + )

207 + assert block_hashes and len(block_hashes) == len(maybe_block_hashes)

208 + parent_block_hash = (

209 + req.block_hashes[first_hash_idx - 1] if first_hash_idx > 0 else None

210 + )

211 + assert first_hash_idx == 0 or parent_block_hash is not None

212 +

213 + lora_id = req.lora_request.adapter_id if req.lora_request is not None else None

214 + lora_name = req.lora_request.name if req.lora_request is not None else None

215 + self._pending_event_metadata[offload_key] = _OffloadEventMetadata(

216 + block_hashes=block_hashes,

217 + parent_block_hash=parent_block_hash,

218 + token_ids=tuple(req.all_token_ids[chunk_start:boundary_tokens]),

219 + block_size=tokens_per_hash,

220 + lora_id=lora_id,

221 + lora_name=lora_name,

222 + extra_keys=None,

223 + group_idx=group_config.group_idx,

224 + kv_cache_spec=group_config.kv_event_group_spec,

225 + )

226 +

227 def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]:

228 """Translate raw OffloadingEvents into self-describing KV events.

229

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

+14 -4

@@ -280,10 +280,6 @@ class SchedulerOffloadConfig(NamedTuple):

280 spec.blocks_per_chunk == 1

281 and len(group_block_sizes) == 1

282 and has_partial_recurrent_group

283 - and not (

284 - spec.kv_events_config.enable_kv_cache_events

285 - and spec.kv_events_config.self_describing_kv_events

286 - )

283 and all(

284 config.sliding_window_size_in_chunks is None

285 or config.requires_cow_source

@@ -897,10 +893,12 @@ class OffloadingConnectorScheduler:

893 for boundary in range(max_boundary, complete_boundary, -tokens_per_hash):

894 boundary_pending = False

895 boundary_missed = False

896 + boundary_keys = []

897 for group_config in self.config.kv_group_configs:

898 key = self._make_boundary_key(

899 req_status.req, group_config.group_idx, boundary

900 )

901 + boundary_keys.append(key)

902 result = self.manager.lookup(key, req_status.req_context)

903 if result is LookupResult.MISS:

904 boundary_missed = True

@@ -910,6 +908,12 @@ class OffloadingConnectorScheduler:

908

909 pending |= boundary_pending

910 if not boundary_missed and not boundary_pending:

911 + for group_config, key in zip(

912 + self.config.kv_group_configs, boundary_keys

913 + ):

914 + self._events_tracker.record_partial_lookup(

915 + req_status.req, group_config, boundary, key

916 + )

917 req_status.partial_tail_boundary = boundary

918 return boundary - local_tokens

919

@@ -1203,6 +1207,12 @@ class OffloadingConnectorScheduler:

1207 if not store_output.keys_to_store:

1208 continue

1209

1210 + for group_config, key in zip(self.config.kv_group_configs, keys):

1211 + if key in store_output.keys_to_store:

1212 + self._events_tracker.record_partial_store(

1213 + req, group_config, boundary, key

1214 + )

1215 +

1216 group_by_key = {key: idx for idx, key in enumerate(keys)}

1217 accepted_groups = [group_by_key[key] for key in store_output.keys_to_store]

1218 group_sizes = [0] * len(self.config.kv_group_configs)