[KV Offloading] Support partial-tail prefix reuse with fine-grained prefix matching (#50507)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
Chauncey committed
Aug 5, 2026 at 22:31 UTC
62c5e2162186a369d9fba40b8795d06b9144e8fe
2 files changed
+360
-44
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+114
-7
@@ -6,7 +6,12 @@ from unittest.mock import MagicMock, call
6
import pytest
7
import torch
8
9
+from tests.v1.kv_connector.unit.offloading_connector.test_config import (
10
+ _make_mamba_hybrid_kv_cache_config,
11
+ _make_vllm_config,
12
+)
13
from tests.v1.kv_connector.unit.offloading_connector.utils import (
14
+ MockOffloadingSpec,
15
generate_store_output,
16
to_keys,
17
)
@@ -16,6 +21,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
21
OffloadingConnectorMetadata,
22
OffloadingWorkerMetadata,
23
)
24
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
25
+ build_offloading_config,
26
+)
27
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
28
OffloadingConnectorStats,
29
_ConnectorMetricName,
@@ -25,13 +33,15 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
33
RequestOffloadState,
34
is_store_reachable_swa_chunk,
35
)
28
-from vllm.v1.core.kv_cache_utils import BlockHash
36
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
37
+from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock
38
from vllm.v1.kv_cache_interface import (
39
FullAttentionSpec,
40
KVCacheGroupSpec,
41
SlidingWindowSpec,
42
)
43
from vllm.v1.kv_offload.base import (
44
+ GPULoadStoreSpec,
45
LookupResult,
46
Medium,
47
OffloadingEvent,
@@ -40,12 +50,36 @@ from vllm.v1.kv_offload.base import (
50
ReqContext,
51
RequestOffloadingContext,
52
get_offload_block_hash,
53
+ get_offload_group_idx,
54
make_offload_key,
55
)
56
from vllm.v1.outputs import KVConnectorOutput
57
from vllm.v1.request import RequestStatus
58
59
60
+def _make_partial_tail_scheduler() -> OffloadingConnectorScheduler:
61
+ vllm_config = _make_vllm_config()
62
+ vllm_config.cache_config.prefix_match_unit = 4
63
+ vllm_config.speculative_config = None
64
+ kv_cache_config = _make_mamba_hybrid_kv_cache_config()
65
+ spec = MockOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
66
+ return OffloadingConnectorScheduler(spec, vllm_config, kv_cache_config)
67
+
68
+
69
+def _make_partial_tail_request(
70
+ scheduler: OffloadingConnectorScheduler,
71
+) -> MagicMock:
72
+ request = MagicMock()
73
+ request.request_id = "req"
74
+ request.kv_transfer_params = None
75
+ request.num_prompt_tokens = 30
76
+ request.num_tokens = 30
77
+ request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(7)]
78
+ request.is_finished.return_value = False
79
+ scheduler.on_new_request(request)
80
+ return request
81
+
82
+
83
def _reduce_kv_connector_stats(runner):
84
reduced: dict[str, int | float] = {}
85
for payload in runner.kv_connector_stats:
@@ -59,6 +93,79 @@ def _reduce_kv_connector_stats(runner):
93
return reduced
94
95
96
+def test_partial_tail_store_uses_attention_and_recurrent_cow_sources():
97
+ scheduler = _make_partial_tail_scheduler()
98
+ _make_partial_tail_request(scheduler)
99
+ req_status = scheduler._req_status["req"]
100
+ req_status.group_states[0].block_ids[:] = [11, 12]
101
+ req_status.group_states[1].block_ids[:] = [0, 21]
102
+ scheduler.manager.prepare_store.side_effect = (
103
+ lambda keys, req_context: generate_store_output(keys)
104
+ )
105
+
106
+ output = SimpleNamespace(partial_tail_offloads={"req": [(1, 99, 28)]})
107
+ jobs = scheduler._build_partial_tail_store_jobs(output)
108
+
109
+ assert len(jobs) == 1
110
+ [job_id] = jobs
111
+ src_spec = jobs[job_id].src_spec
112
+ assert isinstance(src_spec, GPULoadStoreSpec)
113
+ assert src_spec.block_ids.tolist() == [12, 99]
114
+ assert src_spec.group_sizes == [1, 1]
115
+ assert src_spec.block_indices == [1, 1]
116
+ assert scheduler._block_id_to_pending_jobs == {
117
+ 12: {job_id},
118
+ 99: {job_id},
119
+ }
120
+
121
+
122
+def test_partial_lookup_returns_exact_boundary_and_group_load_keys():
123
+ scheduler = _make_partial_tail_scheduler()
124
+ request = _make_partial_tail_request(scheduler)
125
+ req_status = scheduler._req_status["req"]
126
+ req_status.num_locally_computed_tokens = 0
127
+ req_status.update_offload_keys()
128
+
129
+ scheduler.manager.lookup.return_value = LookupResult.HIT
130
+ assert scheduler._lookup(req_status) == 28
131
+
132
+ assert req_status.partial_tail_boundary == 28
133
+
134
+ scheduler.update_state_after_alloc(
135
+ request,
136
+ KVCacheBlocks(
137
+ (
138
+ [KVCacheBlock(31), KVCacheBlock(32)],
139
+ [KVCacheBlock(0, is_null=True), KVCacheBlock(41)],
140
+ )
141
+ ),
142
+ num_external_tokens=28,
143
+ )
144
+ [load_job] = scheduler._current_batch_load_jobs.values()
145
+ dst_spec = load_job.dst_spec
146
+ assert isinstance(dst_spec, GPULoadStoreSpec)
147
+ assert dst_spec.block_ids.tolist() == [31, 32, 41]
148
+ assert dst_spec.group_sizes == [2, 1]
149
+ assert dst_spec.block_indices == [0, 1]
150
+ assert req_status.partial_tail_boundary is None
151
+
152
+
153
+def test_partial_lookup_requires_every_cache_group():
154
+ scheduler = _make_partial_tail_scheduler()
155
+ _make_partial_tail_request(scheduler)
156
+ req_status = scheduler._req_status["req"]
157
+ req_status.update_offload_keys()
158
+
159
+ def lookup(key, req_context):
160
+ if get_offload_group_idx(key) == 1 and get_offload_block_hash(key) != b"h3":
161
+ return LookupResult.MISS
162
+ return LookupResult.HIT
163
+
164
+ scheduler.manager.lookup.side_effect = lookup
165
+ assert scheduler._lookup(req_status) == 16
166
+ assert req_status.partial_tail_boundary is None
167
+
168
+
169
def test_scheduler_reports_allocation_failure(request_runner):
170
runner = request_runner(
171
block_size=4,
@@ -2141,7 +2248,7 @@ def test_stale_sliding_window_block_after_prepare_store_failure(
2248
# Now prepare_store succeeds.
2249
# Without the fix, the request would try to offload the stale block_id
2250
# at position 0 (now reused at position 3), causing a duplicate in
2144
- # sliding_window_block_ids and eventually a KeyError.
2251
+ # fenced_block_ids and eventually a KeyError.
2252
runner.manager.prepare_store.side_effect = lambda keys, req_context: (
2253
generate_store_output(keys)
2254
)
@@ -2968,7 +3075,7 @@ class TestEagle:
3075
def test_request_finished_with_pending_stores_populates_fence(request_runner):
3076
"""When a request finishes with in-flight store jobs, the fence index
3077
(_block_id_to_pending_jobs) is correctly populated with the store jobs'
2971
- non_sliding_window_block_ids.
3078
+ deferred_fence_block_ids.
3079
3080
This prevents data corruption when a subsequent request reuses the same
3081
GPU blocks before the store completes.
@@ -3002,7 +3109,7 @@ def test_request_finished_with_pending_stores_populates_fence(request_runner):
3109
)
3110
for js in runner.connector_scheduler._jobs.values():
3111
if js.is_store:
3005
- job_block_ids.update(js.non_sliding_window_block_ids or [])
3112
+ job_block_ids.update(js.deferred_fence_block_ids or [])
3113
3114
# Run 1: create store job, finish request, populate fence.
3115
# With non-blocking drain (#45595), the job stays in-flight.
@@ -3104,7 +3211,7 @@ def test_request_finished_mixed_full_attn_and_sliding_window(
3211
request_runner,
3212
):
3213
"""With both FullAttention and SlidingWindow groups, a single store job
3107
- has both non_sliding_window_block_ids and sliding_window_block_ids.
3214
+ has both deferred_fence_block_ids and fenced_block_ids.
3215
3216
request_finished only registers non-SW blocks in the fence.
3217
SW blocks were already registered at store creation time.
@@ -3160,8 +3267,8 @@ def test_request_finished_mixed_full_attn_and_sliding_window(
3267
)
3268
for js in runner.connector_scheduler._jobs.values():
3269
if js.is_store:
3163
- sw_block_ids.update(js.sliding_window_block_ids or [])
3164
- non_sw_block_ids.update(js.non_sliding_window_block_ids or [])
3270
+ sw_block_ids.update(js.fenced_block_ids or [])
3271
+ non_sw_block_ids.update(js.deferred_fence_block_ids or [])
3272
3273
# Run 1: create store job, finish request, populate fence.
3274
runner.run(
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+246
-37
@@ -74,12 +74,10 @@ class TransferJobStatus:
74
# Offload keys this job covers; passed to manager.complete_*().
75
keys: set[OffloadKey]
76
is_store: bool
77
- # Store src block IDs whose ref_cnt protects them while the request
78
- # runs. Only registered in _block_id_to_pending_jobs on request_finished.
79
- non_sliding_window_block_ids: list[int] | None = None
80
- # Store src block IDs that may be freed before the request finishes.
81
- # Registered in _block_id_to_pending_jobs at store creation time.
82
- sliding_window_block_ids: list[int] | None = None
77
+ # Store source blocks fenced after the request finishes.
78
+ deferred_fence_block_ids: list[int] | None = None
79
+ # Store source blocks fenced when the transfer is created.
80
+ fenced_block_ids: list[int] | None = None
81
82
83
class GroupOffloadConfig(NamedTuple):
@@ -92,6 +90,9 @@ class GroupOffloadConfig(NamedTuple):
90
kv_event_group_spec: OffloadingEventGroupSpec
91
# None below means full attention
92
sliding_window_size_in_chunks: int | None
93
+ # Partial-tail data for this group comes from the scheduler's CoW hand-off
94
+ # rather than the request block table.
95
+ requires_cow_source: bool = False
96
# Number of this group's offloaded chunks per full-attention alignment
97
# segment. Used to skip storing SWA chunks that can never serve a load
98
# hit (e.g. DeepSeek V4 where SWA groups have much smaller block sizes
@@ -162,8 +163,10 @@ def resolve_mamba_align_size(
163
class SchedulerOffloadConfig(NamedTuple):
164
kv_group_configs: tuple[GroupOffloadConfig, ...]
165
blocks_per_chunk: int
166
+ tokens_per_hash: int
167
num_workers: int
168
offload_prompt_only: bool
169
+ supports_partial_tail: bool
170
171
@classmethod
172
def from_spec(
@@ -226,9 +229,14 @@ class SchedulerOffloadConfig(NamedTuple):
229
sorted(eagle_groups),
230
)
231
229
- return cls(
230
- num_workers=vllm_config.parallel_config.world_size,
231
- kv_group_configs=tuple(
232
+ kv_group_configs_list: list[GroupOffloadConfig] = []
233
+ for idx, tokens_per_block in enumerate(spec.tokens_per_block):
234
+ kv_cache_group = kv_cache_config.kv_cache_groups[idx]
235
+ kv_spec = kv_cache_group.kv_cache_spec
236
+ sw = get_sliding_window_size_in_chunks(
237
+ kv_spec, tokens_per_block * spec.blocks_per_chunk
238
+ )
239
+ kv_group_configs_list.append(
240
GroupOffloadConfig(
241
group_idx=idx,
242
tokens_per_block=tokens_per_block,
@@ -237,24 +245,52 @@ class SchedulerOffloadConfig(NamedTuple):
245
(tokens_per_block * spec.blocks_per_chunk)
246
// spec.tokens_per_hash
247
),
240
- sliding_window_size_in_chunks=(
241
- sw := get_sliding_window_size_in_chunks(
242
- kv_cache_config.kv_cache_groups[idx].kv_cache_spec,
243
- tokens_per_block * spec.blocks_per_chunk,
244
- )
245
- ),
248
+ sliding_window_size_in_chunks=sw,
249
alignment_chunk_count=_alignment_chunk_count(
250
tokens_per_block * spec.blocks_per_chunk, sw
251
),
249
- kv_event_group_spec=get_offloading_event_group_spec(
250
- kv_cache_config.kv_cache_groups[idx]
251
- ),
252
+ kv_event_group_spec=get_offloading_event_group_spec(kv_cache_group),
253
is_eagle_group=idx in eagle_groups,
254
+ requires_cow_source=(
255
+ isinstance(kv_spec, MambaSpec)
256
+ and kv_spec.mamba_cache_mode == "align"
257
+ ),
258
)
254
- for idx, tokens_per_block in enumerate(spec.tokens_per_block)
255
- ),
259
+ )
260
+ kv_group_configs = tuple(kv_group_configs_list)
261
+ group_block_sizes = {config.tokens_per_block for config in kv_group_configs}
262
+ has_partial_recurrent_group = any(
263
+ config.requires_cow_source
264
+ and config.tokens_per_block > spec.tokens_per_hash
265
+ for config in kv_group_configs
266
+ )
267
+ # Partial tails currently require one physical block per offload chunk
268
+ # and uniform, non-windowed groups so one boundary identifies every
269
+ # group's source. EAGLE and DCP need additional hand-off semantics.
270
+ supports_partial_tail = (
271
+ spec.blocks_per_chunk == 1
272
+ and len(group_block_sizes) == 1
273
+ and has_partial_recurrent_group
274
+ and not (
275
+ spec.kv_events_config.enable_kv_cache_events
276
+ and spec.kv_events_config.self_describing_kv_events
277
+ )
278
+ and all(
279
+ config.sliding_window_size_in_chunks is None
280
+ or config.requires_cow_source
281
+ for config in kv_group_configs
282
+ )
283
+ and not any(config.is_eagle_group for config in kv_group_configs)
284
+ and vllm_config.parallel_config.decode_context_parallel_size == 1
285
+ )
286
+
287
+ return cls(
288
+ num_workers=vllm_config.parallel_config.world_size,
289
+ kv_group_configs=kv_group_configs,
290
blocks_per_chunk=spec.blocks_per_chunk,
291
+ tokens_per_hash=spec.tokens_per_hash,
292
offload_prompt_only=spec.offload_prompt_only,
293
+ supports_partial_tail=supports_partial_tail,
294
)
295
296
@@ -286,6 +322,9 @@ class RequestOffloadState:
322
# time.monotonic() of this request's first deferred offload lookup;
323
# None once consumed (observed) or while no lookup is pending.
324
deferred_lookup_start_time: float | None = None
325
+ # Fine-grained token boundary selected beyond the last complete offload
326
+ # chunk. It is consumed when the corresponding load is scheduled.
327
+ partial_tail_boundary: int | None = None
328
# True once on_request_finished has been signaled to the manager.
329
finished_signaled: bool = False
330
@@ -479,6 +518,16 @@ class OffloadingConnectorScheduler:
518
self._mamba_align_size: int | None = resolve_mamba_align_size(
519
spec, kv_cache_config
520
)
521
+ self._partial_tail_block_size = (
522
+ self.config.kv_group_configs[0].tokens_per_block
523
+ if self.config.supports_partial_tail
524
+ else 0
525
+ )
526
+ self._cow_source_groups = frozenset(
527
+ config.group_idx
528
+ for config in self.config.kv_group_configs
529
+ if config.requires_cow_source
530
+ )
531
532
self._req_status: dict[ReqId, RequestOffloadState] = {}
533
self._current_batch_load_jobs: dict[int, TransferJob] = {}
@@ -627,8 +676,20 @@ class OffloadingConnectorScheduler:
676
group_state.offload_keys[chunks_to_skip:],
677
req_status.req_context,
678
)
679
+ if req_status.partial_tail_boundary is not None:
680
+ self.manager.touch(
681
+ tuple(
682
+ self._make_boundary_key(
683
+ req_status.req,
684
+ group.group_idx,
685
+ req_status.partial_tail_boundary,
686
+ )
687
+ for group in self.config.kv_group_configs
688
+ ),
689
+ req_status.req_context,
690
+ )
691
631
- def _lookup(self, req_status: RequestOffloadState) -> int | None:
692
+ def _lookup_complete_chunks(self, req_status: RequestOffloadState) -> int | None:
693
"""
694
Find how many tokens beyond num_locally_computed_tokens can be loaded.
695
@@ -801,6 +862,52 @@ class OffloadingConnectorScheduler:
862
863
return num_hit_tokens
864
865
+ def _make_boundary_key(
866
+ self, request: Request, group_idx: int, boundary_tokens: int
867
+ ) -> OffloadKey:
868
+ hash_idx = boundary_tokens // self.config.tokens_per_hash - 1
869
+ return make_offload_key(request.block_hashes[hash_idx], group_idx)
870
+
871
+ def _lookup(self, req_status: RequestOffloadState) -> int | None:
872
+ complete_hit = self._lookup_complete_chunks(req_status)
873
+ req_status.partial_tail_boundary = None
874
+ if complete_hit is None or not self.config.supports_partial_tail:
875
+ return complete_hit
876
+
877
+ local_tokens = req_status.num_locally_computed_tokens
878
+ complete_boundary = local_tokens + complete_hit
879
+ tokens_per_hash = self.config.tokens_per_hash
880
+ block_end = complete_boundary + self._partial_tail_block_size
881
+ max_boundary = round_down(
882
+ min(req_status.req.num_prompt_tokens - 1, block_end - 1), tokens_per_hash
883
+ )
884
+ if max_boundary <= complete_boundary:
885
+ return complete_hit
886
+
887
+ pending = False
888
+ for boundary in range(max_boundary, complete_boundary, -tokens_per_hash):
889
+ boundary_pending = False
890
+ boundary_missed = False
891
+ for group_config in self.config.kv_group_configs:
892
+ key = self._make_boundary_key(
893
+ req_status.req, group_config.group_idx, boundary
894
+ )
895
+ result = self.manager.lookup(key, req_status.req_context)
896
+ if result is LookupResult.MISS:
897
+ boundary_missed = True
898
+ break
899
+ if result in (LookupResult.HIT_PENDING, LookupResult.RETRY):
900
+ boundary_pending = True
901
+
902
+ pending |= boundary_pending
903
+ if not boundary_missed and not boundary_pending:
904
+ req_status.partial_tail_boundary = boundary
905
+ return boundary - local_tokens
906
+
907
+ if pending and complete_hit == 0:
908
+ return None
909
+ return complete_hit
910
+
911
def on_new_request(self, request: Request) -> None:
912
"""Called when a new request is added to the scheduler."""
913
req_context = _create_req_context(request)
@@ -880,6 +987,9 @@ class OffloadingConnectorScheduler:
987
988
num_locally_computed_tokens = req_status.num_locally_computed_tokens
989
num_cached_tokens = num_locally_computed_tokens + num_external_tokens
990
+ partial_tail_boundary = req_status.partial_tail_boundary
991
+ if partial_tail_boundary is not None:
992
+ assert partial_tail_boundary == num_cached_tokens
993
994
keys_to_load: list[OffloadKey] = []
995
dst_block_ids: list[int] = []
@@ -923,12 +1033,19 @@ class OffloadingConnectorScheduler:
1033
)
1034
1035
num_chunks = cdiv(num_cached_tokens, tokens_per_chunk)
926
- assert len(offload_keys) >= num_chunks
1036
if num_pending_gpu_blocks:
1037
start_chunk_idx = (
1038
num_locally_computed_gpu_blocks // self.config.blocks_per_chunk
1039
)
931
- keys_to_load.extend(offload_keys[start_chunk_idx:num_chunks])
1040
+ end_chunk_idx = num_chunks - (partial_tail_boundary is not None)
1041
+ assert len(offload_keys) >= end_chunk_idx
1042
+ keys_to_load.extend(offload_keys[start_chunk_idx:end_chunk_idx])
1043
+ if partial_tail_boundary is not None:
1044
+ keys_to_load.append(
1045
+ self._make_boundary_key(
1046
+ request, group_config.group_idx, partial_tail_boundary
1047
+ )
1048
+ )
1049
1050
dst_block_ids.extend(
1051
block.block_id
@@ -968,6 +1085,7 @@ class OffloadingConnectorScheduler:
1085
1086
if self._chunks_being_loaded is not None:
1087
self._chunks_being_loaded.update(keys_to_load)
1088
+ req_status.partial_tail_boundary = None
1089
1090
def _update_req_states(self, scheduler_output: SchedulerOutput) -> None:
1091
"""
@@ -1020,6 +1138,94 @@ class OffloadingConnectorScheduler:
1138
):
1139
group_state.block_ids[j] = 0
1140
1141
+ def _build_partial_tail_store_jobs(
1142
+ self, scheduler_output: SchedulerOutput
1143
+ ) -> dict[int, TransferJob]:
1144
+ handoffs = scheduler_output.partial_tail_offloads
1145
+ if not self.config.supports_partial_tail or not handoffs:
1146
+ return {}
1147
+
1148
+ store_jobs: dict[int, TransferJob] = {}
1149
+ for req_id, entries in handoffs.items():
1150
+ req_status = self._req_status.get(req_id)
1151
+ assert req_status is not None
1152
+ assert entries
1153
+ boundaries = {boundary for _, _, boundary in entries}
1154
+ assert len(boundaries) == 1
1155
+ boundary = boundaries.pop()
1156
+ req = req_status.req
1157
+ max_boundary = min(
1158
+ req.num_prompt_tokens,
1159
+ req_status.max_offload_tokens or req.num_prompt_tokens,
1160
+ )
1161
+ assert boundary > 0
1162
+ assert boundary % self.config.tokens_per_hash == 0
1163
+ assert boundary <= max_boundary
1164
+
1165
+ cow_blocks = {group_idx: block_id for group_idx, block_id, _ in entries}
1166
+ assert self._cow_source_groups.issubset(cow_blocks)
1167
+
1168
+ assert boundary % self._partial_tail_block_size != 0
1169
+ block_idx = boundary // self._partial_tail_block_size
1170
+ if any(
1171
+ group.group_idx not in self._cow_source_groups
1172
+ and block_idx >= len(req_status.group_states[group.group_idx].block_ids)
1173
+ for group in self.config.kv_group_configs
1174
+ ):
1175
+ continue
1176
+ keys = [
1177
+ self._make_boundary_key(req, group.group_idx, boundary)
1178
+ for group in self.config.kv_group_configs
1179
+ ]
1180
+ block_ids = [
1181
+ cow_blocks[group.group_idx]
1182
+ if group.group_idx in self._cow_source_groups
1183
+ else req_status.group_states[group.group_idx].block_ids[block_idx]
1184
+ for group in self.config.kv_group_configs
1185
+ ]
1186
+ assert all(block_id != 0 for block_id in block_ids)
1187
+
1188
+ store_output = self.manager.prepare_store(keys, req_status.req_context)
1189
+ if store_output is None:
1190
+ self._connector_stats.increase_counter(
1191
+ _ConnectorMetricName.ALLOCATION_FAILURE
1192
+ )
1193
+ continue
1194
+ if not store_output.keys_to_store:
1195
+ continue
1196
+
1197
+ group_by_key = {key: idx for idx, key in enumerate(keys)}
1198
+ accepted_groups = [group_by_key[key] for key in store_output.keys_to_store]
1199
+ group_sizes = [0] * len(self.config.kv_group_configs)
1200
+ block_indices = [0] * len(self.config.kv_group_configs)
1201
+ for group_idx in accepted_groups:
1202
+ group_sizes[group_idx] = 1
1203
+ block_indices[group_idx] = block_idx
1204
+ source_blocks = [block_ids[group_idx] for group_idx in accepted_groups]
1205
+
1206
+ job_id = self._generate_job_id()
1207
+ req_status.transfer_jobs.add(job_id)
1208
+ for block_id in source_blocks:
1209
+ self._block_id_to_pending_jobs.setdefault(block_id, set()).add(job_id)
1210
+ self._jobs[job_id] = TransferJobStatus(
1211
+ req_id=req_id,
1212
+ pending_count=self.config.num_workers,
1213
+ keys=set(store_output.keys_to_store),
1214
+ is_store=True,
1215
+ fenced_block_ids=source_blocks,
1216
+ )
1217
+ store_jobs[job_id] = TransferJob(
1218
+ req_id=req_id,
1219
+ src_spec=GPULoadStoreSpec(
1220
+ source_blocks,
1221
+ group_sizes=group_sizes,
1222
+ block_indices=block_indices,
1223
+ ),
1224
+ dst_spec=store_output.store_spec,
1225
+ )
1226
+
1227
+ return store_jobs
1228
+
1229
def _build_store_jobs(
1230
self,
1231
scheduler_output: SchedulerOutput,
@@ -1119,8 +1325,8 @@ class OffloadingConnectorScheduler:
1325
group_sizes: list[int] = []
1326
block_indices: list[int] = []
1327
src_block_ids: list[int] = []
1122
- sliding_window_block_ids: list[int] = []
1123
- non_sliding_window_block_ids: list[int] = []
1328
+ fenced_block_ids: list[int] = []
1329
+ deferred_fence_block_ids: list[int] = []
1330
for group_config, group_state in zip(
1331
self.config.kv_group_configs, req_status.group_states
1332
):
@@ -1156,9 +1362,9 @@ class OffloadingConnectorScheduler:
1362
src_block_ids.append(block_id)
1363
num_group_blocks += 1
1364
if is_sliding_window:
1159
- sliding_window_block_ids.append(block_id)
1365
+ fenced_block_ids.append(block_id)
1366
else:
1161
- non_sliding_window_block_ids.append(block_id)
1367
+ deferred_fence_block_ids.append(block_id)
1368
1369
group_sizes.append(num_group_blocks)
1370
block_indices.append(start_gpu_block_idx or 0)
@@ -1180,7 +1386,7 @@ class OffloadingConnectorScheduler:
1386
1387
# Watch sliding window blocks as they may get evicted
1388
# before the request finishes
1183
- for bid in sliding_window_block_ids or ():
1389
+ for bid in fenced_block_ids:
1390
self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id)
1391
1392
# the non-sliding window blocks will be watched only
@@ -1190,8 +1396,8 @@ class OffloadingConnectorScheduler:
1396
pending_count=self.config.num_workers,
1397
keys=set(keys_to_store),
1398
is_store=True,
1193
- non_sliding_window_block_ids=non_sliding_window_block_ids,
1194
- sliding_window_block_ids=sliding_window_block_ids or None,
1399
+ deferred_fence_block_ids=deferred_fence_block_ids,
1400
+ fenced_block_ids=fenced_block_ids or None,
1401
)
1402
1403
store_jobs[job_id] = TransferJob(
@@ -1208,7 +1414,7 @@ class OffloadingConnectorScheduler:
1414
1415
if req.is_finished():
1416
# Register non-sliding-window blocks for flush detection.
1211
- for bid in non_sliding_window_block_ids:
1417
+ for bid in deferred_fence_block_ids:
1418
self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id)
1419
if bid in self._current_batch_allocated_block_ids:
1420
self._current_batch_jobs_to_flush.add(job_id)
@@ -1248,9 +1454,11 @@ class OffloadingConnectorScheduler:
1454
for jid in self._block_id_to_pending_jobs[bid]
1455
)
1456
1457
+ partial_store_jobs = self._build_partial_tail_store_jobs(scheduler_output)
1458
+ normal_store_jobs = self._build_store_jobs(scheduler_output)
1459
meta = OffloadingConnectorMetadata(
1460
load_jobs=self._current_batch_load_jobs,
1253
- store_jobs=self._build_store_jobs(scheduler_output),
1461
+ store_jobs=partial_store_jobs | normal_store_jobs,
1462
jobs_to_flush=self._current_batch_jobs_to_flush,
1463
)
1464
@@ -1344,12 +1552,12 @@ class OffloadingConnectorScheduler:
1552
if self._block_id_to_pending_jobs:
1553
# Sliding window blocks are tracked from store creation
1554
# and must be cleaned up unconditionally.
1347
- self._remove_pending_job(job_id, job_status.sliding_window_block_ids)
1555
+ self._remove_pending_job(job_id, job_status.fenced_block_ids)
1556
# Non-sliding-window blocks are only tracked after
1557
# request_finished, so only clean up for finished requests.
1558
if req_status.req.is_finished():
1559
self._remove_pending_job(
1352
- job_id, job_status.non_sliding_window_block_ids
1560
+ job_id, job_status.deferred_fence_block_ids
1561
)
1562
1563
del self._jobs[job_id]
@@ -1404,11 +1612,11 @@ class OffloadingConnectorScheduler:
1612
1613
# Keep req_status alive: _build_store_jobs will process finished_req_ids
1614
# on the next step and handle cleanup after creating store jobs.
1407
- # Register non_sliding_window_block_ids so future block reuse triggers
1408
- # a flush via _block_id_to_pending_jobs.
1615
+ # Register deferred fences so future block reuse triggers a flush via
1616
+ # _block_id_to_pending_jobs.
1617
for job_id in req_status.transfer_jobs:
1618
job_status = self._jobs[job_id]
1411
- for bid in job_status.non_sliding_window_block_ids or ():
1619
+ for bid in job_status.deferred_fence_block_ids or ():
1620
self._block_id_to_pending_jobs.setdefault(bid, set()).add(job_id)
1621
1622
return False, None
@@ -1451,6 +1659,7 @@ class OffloadingConnectorScheduler:
1659
for group_state in status.group_states:
1660
group_state.next_stored_chunk_idx = 0
1661
status.transfer_jobs.clear()
1662
+ status.partial_tail_boundary = None
1663
1664
# Discard jobs and save job_counter to be able to discard worker responses
1665
self._stale_job_threshold = self._job_counter