diff --git a/deploy/deepseek_v4/tp4_profile.json b/deploy/deepseek_v4/tp4_profile.json index c586573..cda756a 100644 --- a/deploy/deepseek_v4/tp4_profile.json +++ b/deploy/deepseek_v4/tp4_profile.json @@ -4,7 +4,7 @@ "cache_model_profile": "deepseek-v4-fp8-hma", "published_runtime_base": "ghcr.io/fujitsupolycom/gb10-vllm-serving@sha256:6fc26fdad81a18f0fff67ce0a05f6d90165625ea2e1cac8a6f39bfb462017028", "sparkcache": { - "source_sha256": "83853050f790b18af95d424fec837abeb1a9a33f0538b5e4b97c16fb9c681781" + "source_sha256": "788686e858ba4af01f535e95122c7650f412fddc40cd221a0924f4ce2b32ff98" }, "model": { "repository": "deepseek-ai/DeepSeek-V4-Flash-0731", diff --git a/deploy/glm52_35bpw/profile.json b/deploy/glm52_35bpw/profile.json index 547ad8b..d8e66c8 100644 --- a/deploy/glm52_35bpw/profile.json +++ b/deploy/glm52_35bpw/profile.json @@ -5,7 +5,7 @@ "published_runtime_base": "ghcr.io/fujitsupolycom/gb10-vllm-serving@sha256:6fc26fdad81a18f0fff67ce0a05f6d90165625ea2e1cac8a6f39bfb462017028", "base_image_requirement": "exact GLM-5.2 3.5-bpw R7 image recorded by the source container inspection", "sparkcache": { - "source_sha256": "83853050f790b18af95d424fec837abeb1a9a33f0538b5e4b97c16fb9c681781" + "source_sha256": "788686e858ba4af01f535e95122c7650f412fddc40cd221a0924f4ce2b32ff98" }, "model": { "repository": "brandonmusic/GLM-5.2-EXL3-TR3v4-3.5bpw-MTP78", diff --git a/sparkcache/README.md b/sparkcache/README.md index 3998788..e763e6f 100644 --- a/sparkcache/README.md +++ b/sparkcache/README.md @@ -225,13 +225,20 @@ when placement completes and intentionally excludes that bookkeeping. `sparkcache-hybrid-page-delta/v1` codec reuses only byte-identical page prefixes and binds the base snapshot, layout, block counts, and semantic token boundaries. A boundary inside an HMA page replaces that complete page - while retaining earlier byte-identical pages. For an aligned recurrent group, - vLLM may retain the replay-boundary page outside the advancing request block - table. Its `SchedulerOutput.recurrent_boundary_blocks` hand-off names the - pinned physical block by request, group, and token boundary. SparkCache uses - that block only after all three identities and the recurrent topology match; - missing or contradictory metadata skips publication rather than scanning - later running or speculative state. The + while retaining earlier byte-identical pages. At an exact recurrent-page + boundary, vLLM may retain the replay-boundary page outside the advancing + request block table. Its `SchedulerOutput.recurrent_boundary_blocks` hand-off + names the pinned physical block by request, group, and token boundary. + SparkCache defers a new recurrent request until a later cached scheduler step, + when the preceding forward's hand-off can be observed. It latches one matching + entry for every recurrent group, including a partial-tail CoW target when the + boundary lies inside a recurrent page. Valid entries for an earlier checkpoint + are ignored while the request advances; outputs with no target-boundary entry + leave publication pending. Incomplete, future, contradictory, or changed + target evidence cancels it. A store is emitted only after every recurrent + group has a proven pinned block at the exact publication boundary. SparkCache + never substitutes an accumulated request-table ID because vLLM may have + replaced that source block while producing the durable CoW target. The `sparkcache-page-delta-manifest/v2` schema embeds its authenticated base graph and groups delta bytes into immutable objects of at most 64 MiB. A 1,575,821,491-byte delta therefore uses at most 24 physical delta objects diff --git a/sparkcache/spark_context_cache_connector.py b/sparkcache/spark_context_cache_connector.py index 22b98fd..4803044 100644 --- a/sparkcache/spark_context_cache_connector.py +++ b/sparkcache/spark_context_cache_connector.py @@ -570,6 +570,11 @@ def __init__( self._storage_mode = config.storage_mode self._publication_schema = config.publication_schema self._group_topology = config.group_topology + self._recurrent_group_indexes = frozenset( + group_index + for group_index, topology in enumerate(self._group_topology) + if topology["reuse_policy"] == "recurrent_align" + ) self._chunk_tokens = config.chunk_tokens self._root = config.root self._store = ManifestStore(self._root) @@ -1613,15 +1618,21 @@ def _validated_recurrent_boundary_blocks( scheduler_output: "SchedulerOutput", request_id: str, boundary_tokens: int, + *, + latched: tuple[tuple[int, int], ...] = (), ) -> tuple[tuple[int, int], ...] | None: """Validate vLLM's exact recurrent replay-boundary block hand-off. - An empty tuple is valid only when the registered topology has no - aligned recurrent group. None means the metadata is absent, - incomplete, or contradictory, so publication must be skipped. - SparkCache never derives a replacement from another non-null table - entry because later entries can hold running or speculative state - beyond ``boundary_tokens``. + vLLM may expose an earlier aligned checkpoint while the request is + still advancing toward this store boundary, followed by a partial-tail + CoW target on a later scheduler output. Valid older entries and absent + per-request metadata therefore preserve ``latched`` and keep the store + pending. None means supplied metadata is incomplete, malformed, ahead + of the plan, or conflicts with an earlier same-boundary latch, so this + publication attempt must be poisoned. SparkCache never derives a + replacement from another request-table entry because it can name + overwritten running or speculative state instead of vLLM's pinned CoW + target. """ def reject(reason: str) -> None: @@ -1635,26 +1646,24 @@ def reject(reason: str) -> None: ) return None - required_groups = { - group_index - for group_index, topology in enumerate(self._group_topology) - if topology["reuse_policy"] == "recurrent_align" - } + required_groups = self._recurrent_group_indexes if not required_groups: return () raw = getattr(scheduler_output, "recurrent_boundary_blocks", None) if raw is None: - return reject("vLLM supplied no recurrent boundary mapping") + return latched if not isinstance(raw, Mapping): return reject("top-level value is not a mapping") entries = raw.get(request_id) if entries is None: - return reject("request has no recurrent boundary entries") + return latched if not isinstance(entries, (list, tuple)): return reject("request value is not a sequence") + if not entries: + return reject("request has no recurrent boundary entries") overrides: list[tuple[int, int]] = [] - seen_groups: set[int] = set() + seen_target_groups: set[int] = set() for entry in entries: if not isinstance(entry, (list, tuple)) or len(entry) != 3: return reject("entry is not a group, block, boundary triple") @@ -1663,20 +1672,31 @@ def reject(reason: str) -> None: return reject("entry values are not integers") if not 0 <= group_index < len(self._group_topology): return reject("group index is outside the registered topology") - if group_index in seen_groups: - return reject("multiple blocks claim the same recurrent group") topology = self._group_topology[group_index] if topology["reuse_policy"] != "recurrent_align": return reject("group is not an aligned recurrent cache") if block_id <= 0: return reject("physical block is vLLM's null block") - if entry_boundary != boundary_tokens: - return reject("entry boundary differs from the store plan") - seen_groups.add(group_index) + if entry_boundary < boundary_tokens: + continue + if entry_boundary > boundary_tokens: + return reject( + "entry boundary is ahead of the store plan" + f" observed={entry_boundary} target={boundary_tokens}" + f" group={group_index} block={block_id}" + ) + if group_index in seen_target_groups: + return reject("multiple blocks claim the same recurrent group") + seen_target_groups.add(group_index) overrides.append((group_index, block_id)) - if seen_groups != required_groups: + if not overrides: + return latched + if seen_target_groups != required_groups: return reject("entries do not cover every aligned recurrent group") - return tuple(sorted(overrides)) + validated = tuple(sorted(overrides)) + if latched and validated != latched: + return reject("entries conflict with the latched recurrent boundary") + return validated def build_connector_meta( self, scheduler_output: "SchedulerOutput" @@ -1741,15 +1761,6 @@ def build_connector_meta( if self._has_full_quorum(digest): self.counters["store_skipped_quorum"] += 1 continue - recurrent_boundary_blocks = ( - self._validated_recurrent_boundary_blocks( - scheduler_output, - req_id, - span, - ) - ) - if recurrent_boundary_blocks is None: - continue already = new_req.num_computed_tokens + scheduled if self._streaming_snapshots_enabled: self._append_streaming_snapshot_offer( @@ -1771,10 +1782,34 @@ def build_connector_meta( [list(group) for group in group_blocks], ) self._store_token_ids[req_id] = exact_token_ids - if recurrent_boundary_blocks: - self._store_recurrent_boundaries[req_id] = ( - recurrent_boundary_blocks - ) + elif self._recurrent_group_indexes: + recurrent_boundary_blocks = ( + self._validated_recurrent_boundary_blocks( + scheduler_output, + req_id, + span, + ) + ) + if recurrent_boundary_blocks is None: + continue + # Full-page proof and partial-tail CoW hand-offs can arrive + # after the prefill which began this store. Retain the + # complete request table and any early proof until a later + # cached step has both finished the span and proven every + # recurrent group. + self._store_progress[req_id] = ( + digest, + span, + already, + [list(group) for group in group_blocks], + ) + self._store_token_ids[req_id] = exact_token_ids + if recurrent_boundary_blocks: + self._store_recurrent_boundaries[req_id] = ( + recurrent_boundary_blocks + ) + if base_digest: + self._store_bases[req_id] = (base_digest, base_span) elif already >= span: meta.plans.append( _ReqPlan( @@ -1787,9 +1822,6 @@ def build_connector_meta( token_ids=exact_token_ids, base_context_digest=base_digest, base_span_tokens=base_span, - recurrent_boundary_blocks=( - recurrent_boundary_blocks - ), ) ) else: @@ -1803,10 +1835,6 @@ def build_connector_meta( [list(group) for group in group_blocks], ) self._store_token_ids[req_id] = exact_token_ids - if recurrent_boundary_blocks: - self._store_recurrent_boundaries[req_id] = ( - recurrent_boundary_blocks - ) if base_digest: self._store_bases[req_id] = (base_digest, base_span) cached = scheduler_output.scheduled_cached_reqs @@ -1816,10 +1844,18 @@ def build_connector_meta( digest, span, done, blocks_by_group = self._store_progress[req_id] exact_token_ids = self._store_token_ids.get(req_id, ()) base_digest, base_span = self._store_bases.get(req_id, ("", 0)) + if self._has_full_quorum(digest): + del self._store_progress[req_id] + self._store_token_ids.pop(req_id, None) + self._store_bases.pop(req_id, None) + self._store_recurrent_boundaries.pop(req_id, None) + self.counters["store_skipped_quorum"] += 1 + continue recurrent_boundary_blocks = self._validated_recurrent_boundary_blocks( scheduler_output, req_id, span, + latched=self._store_recurrent_boundaries.get(req_id, ()), ) if recurrent_boundary_blocks is None: del self._store_progress[req_id] @@ -1829,30 +1865,31 @@ def build_connector_meta( continue if recurrent_boundary_blocks: self._store_recurrent_boundaries[req_id] = recurrent_boundary_blocks - new_block_ids = cached.new_block_ids[index] - appended = ( - [ - list(group) - for group in self._normalize_group_blocks( - new_block_ids, - allow_empty_groups=True, - ) - ] - if new_block_ids is not None - else [[] for _ in blocks_by_group] - ) - if len(appended) != len(blocks_by_group): - raise RuntimeError( - "spark-context-cache: KV-cache group count changed while" - " accumulating a store" + if done < span or req_id in cached.resumed_req_ids: + new_block_ids = cached.new_block_ids[index] + appended = ( + [ + list(group) + for group in self._normalize_group_blocks( + new_block_ids, + allow_empty_groups=True, + ) + ] + if new_block_ids is not None + else [[] for _ in blocks_by_group] ) - if req_id in cached.resumed_req_ids: - blocks_by_group = appended - else: - blocks_by_group = [ - existing + added - for existing, added in zip(blocks_by_group, appended) - ] + if len(appended) != len(blocks_by_group): + raise RuntimeError( + "spark-context-cache: KV-cache group count changed while" + " accumulating a store" + ) + if req_id in cached.resumed_req_ids: + blocks_by_group = appended + else: + blocks_by_group = [ + existing + added + for existing, added in zip(blocks_by_group, appended) + ] blocks = blocks_by_group[0] done = cached.num_computed_tokens[index] + ( scheduler_output.num_scheduled_tokens.get(req_id, 0) @@ -1882,16 +1919,18 @@ def build_connector_meta( block_ids=blocks, ) elif done >= span: + if self._recurrent_group_indexes and not recurrent_boundary_blocks: + self._store_progress[req_id] = ( + digest, + span, + done, + blocks_by_group, + ) + continue del self._store_progress[req_id] self._store_token_ids.pop(req_id, None) self._store_bases.pop(req_id, None) - recurrent_boundary_blocks = self._store_recurrent_boundaries.pop( - req_id, - (), - ) - if self._has_full_quorum(digest): - self.counters["store_skipped_quorum"] += 1 - continue + self._store_recurrent_boundaries.pop(req_id, None) normalized = tuple(tuple(group) for group in blocks_by_group) meta.plans.append( _ReqPlan( diff --git a/sparkcache/test_defect_regressions.py b/sparkcache/test_defect_regressions.py index 70e4bd2..c7d3a1a 100644 --- a/sparkcache/test_defect_regressions.py +++ b/sparkcache/test_defect_regressions.py @@ -1025,7 +1025,10 @@ class DefectD17RecurrentBoundaryMetadataTests(unittest.TestCase): BOUNDARY = 6912 PROMPT_TOKENS = 6992 + NONALIGNED_BOUNDARY = 8192 + NONALIGNED_PROMPT_TOKENS = 8256 BOUNDARY_BLOCK = 42 + COW_BLOCK = 142 @staticmethod def _config() -> types.SimpleNamespace: @@ -1093,6 +1096,37 @@ def _scheduler_output( output.recurrent_boundary_blocks = recurrent_boundary_blocks return output + @classmethod + def _cached_scheduler_output( + cls, + *, + num_computed_tokens: int, + recurrent_boundary_blocks: object = None, + group_count: int = 2, + num_scheduled_tokens: int = 1, + resumed: bool = False, + new_block_ids: object = None, + ) -> types.SimpleNamespace: + request_id = "dflash-recurrent-boundary" + output = types.SimpleNamespace( + scheduled_new_reqs=[], + scheduled_cached_reqs=types.SimpleNamespace( + req_ids=[request_id], + resumed_req_ids={request_id} if resumed else set(), + num_computed_tokens=[num_computed_tokens], + new_block_ids=[ + new_block_ids + if new_block_ids is not None + else tuple(() for _ in range(group_count)) + ], + ), + num_scheduled_tokens={request_id: num_scheduled_tokens}, + preempted_req_ids=set(), + ) + if recurrent_boundary_blocks is not None: + output.recurrent_boundary_blocks = recurrent_boundary_blocks + return output + def test_explicit_boundary_block_round_trips_through_manifest_store(self) -> None: with tempfile.TemporaryDirectory() as directory: root = Path(directory) @@ -1108,12 +1142,24 @@ def test_explicit_boundary_block_round_trips_through_manifest_store(self) -> Non kv_cache_config=config, extra_config={"spark_cache_model_profile": "glm53-flash-hybrid"}, ) - output = self._scheduler_output( - { - "dflash-recurrent-boundary": [ - (1, self.BOUNDARY_BLOCK, self.BOUNDARY) - ] - } + boundary_metadata = { + "dflash-recurrent-boundary": [ + (1, self.BOUNDARY_BLOCK, self.BOUNDARY) + ] + } + first_metadata = scheduler.build_connector_meta( + self._scheduler_output(boundary_metadata) + ) + self.assertEqual(first_metadata.plans, []) + self.assertIn("dflash-recurrent-boundary", scheduler._store_progress) + self.assertEqual( + scheduler._store_recurrent_boundaries[ + "dflash-recurrent-boundary" + ], + ((1, self.BOUNDARY_BLOCK),), + ) + output = self._cached_scheduler_output( + num_computed_tokens=self.PROMPT_TOKENS, ) self.assertTrue(scheduler.supports_recurrent_boundary_blocks) @@ -1173,7 +1219,164 @@ def test_explicit_boundary_block_round_trips_through_manifest_store(self) -> Non torch.equal(pools["recurrent"][[93]], expected_recurrent) ) - def test_missing_or_wrong_request_metadata_skips_publication(self) -> None: + def test_nonaligned_boundary_waits_for_partial_tail_cow_mapping(self) -> None: + with tempfile.TemporaryDirectory() as directory: + connector = _make_connector( + Path(directory), + 0, + block_size=256, + role=KVConnectorRole.SCHEDULER, + override_worker_rank=False, + tp=1, + dcp=1, + kv_cache_config=self._config(), + extra_config={"spark_cache_model_profile": "glm53-flash-hybrid"}, + ) + output = self._scheduler_output() + request = output.scheduled_new_reqs[0] + request.prompt_token_ids = list(range(self.NONALIGNED_PROMPT_TOKENS)) + output.num_scheduled_tokens[request.req_id] = ( + self.NONALIGNED_PROMPT_TOKENS + ) + output.recurrent_boundary_blocks = { + request.req_id: [(1, self.BOUNDARY_BLOCK, self.BOUNDARY)] + } + + first_metadata = connector.build_connector_meta(output) + self.assertEqual(first_metadata.plans, []) + self.assertIn(request.req_id, connector._store_progress) + self.assertNotIn(request.req_id, connector._store_recurrent_boundaries) + self.assertEqual( + connector.counters["recurrent_boundary_metadata_rejected"], + 0, + ) + pending = connector.build_connector_meta( + self._cached_scheduler_output( + num_computed_tokens=self.NONALIGNED_PROMPT_TOKENS, + ) + ) + self.assertEqual(pending.plans, []) + self.assertIn(request.req_id, connector._store_progress) + metadata = connector.build_connector_meta( + self._cached_scheduler_output( + num_computed_tokens=self.NONALIGNED_PROMPT_TOKENS + 1, + recurrent_boundary_blocks={ + request.req_id: [ + (1, self.COW_BLOCK, self.NONALIGNED_BOUNDARY) + ] + }, + ) + ) + + self.assertEqual(len(metadata.plans), 1) + plan = metadata.plans[0] + self.assertEqual(plan.span_tokens, self.NONALIGNED_BOUNDARY) + self.assertEqual(plan.recurrent_boundary_blocks, ((1, self.COW_BLOCK),)) + self.assertEqual( + connector._select_group_blocks_for_span( + plan.block_ids_by_group, + plan.span_tokens, + recurrent_boundary_blocks=plan.recurrent_boundary_blocks, + ), + ((11, 12, 13, 14), (self.COW_BLOCK,)), + ) + self.assertEqual( + connector.counters["recurrent_boundary_metadata_rejected"], + 0, + ) + + def test_conflicting_mapping_poisons_latched_publication(self) -> None: + output = self._scheduler_output( + { + "dflash-recurrent-boundary": [ + (1, self.BOUNDARY_BLOCK, self.BOUNDARY) + ] + } + ) + with tempfile.TemporaryDirectory() as directory: + connector = _make_connector( + Path(directory), + 0, + block_size=256, + role=KVConnectorRole.SCHEDULER, + override_worker_rank=False, + tp=1, + dcp=1, + kv_cache_config=self._config(), + extra_config={"spark_cache_model_profile": "glm53-flash-hybrid"}, + ) + + first_metadata = connector.build_connector_meta(output) + self.assertEqual(first_metadata.plans, []) + metadata = connector.build_connector_meta( + self._cached_scheduler_output( + num_computed_tokens=self.NONALIGNED_PROMPT_TOKENS, + recurrent_boundary_blocks={ + "dflash-recurrent-boundary": [ + (1, self.BOUNDARY_BLOCK + 1, self.BOUNDARY) + ] + }, + ) + ) + + self.assertEqual(metadata.plans, []) + self.assertEqual( + connector.counters["recurrent_boundary_metadata_rejected"], + 1, + ) + + def test_chunked_prefill_validates_only_at_publication_step(self) -> None: + with tempfile.TemporaryDirectory() as directory: + connector = _make_connector( + Path(directory), + 0, + block_size=256, + role=KVConnectorRole.SCHEDULER, + override_worker_rank=False, + tp=1, + dcp=1, + kv_cache_config=self._config(), + extra_config={"spark_cache_model_profile": "glm53-flash-hybrid"}, + ) + first = self._scheduler_output() + first.num_scheduled_tokens["dflash-recurrent-boundary"] = 2304 + + self.assertEqual(connector.build_connector_meta(first).plans, []) + middle = self._cached_scheduler_output( + num_computed_tokens=2304, + num_scheduled_tokens=2304, + recurrent_boundary_blocks={ + "dflash-recurrent-boundary": [ + (1, self.BOUNDARY_BLOCK - 1, 2304) + ] + }, + ) + self.assertEqual(connector.build_connector_meta(middle).plans, []) + self.assertIn("dflash-recurrent-boundary", connector._store_progress) + self.assertEqual( + connector.counters["recurrent_boundary_metadata_rejected"], + 0, + ) + final = self._cached_scheduler_output( + num_computed_tokens=4608, + num_scheduled_tokens=2304, + recurrent_boundary_blocks={ + "dflash-recurrent-boundary": [ + (1, self.BOUNDARY_BLOCK, self.BOUNDARY) + ] + }, + ) + + metadata = connector.build_connector_meta(final) + + self.assertEqual(len(metadata.plans), 1) + self.assertEqual( + metadata.plans[0].recurrent_boundary_blocks, + ((1, self.BOUNDARY_BLOCK),), + ) + self.assertNotIn("dflash-recurrent-boundary", connector._store_progress) + + def test_missing_or_wrong_request_metadata_keeps_publication_pending(self) -> None: for boundary_metadata in ( None, {"another-request": [(1, self.BOUNDARY_BLOCK, self.BOUNDARY)]}, @@ -1198,19 +1401,42 @@ def test_missing_or_wrong_request_metadata_skips_publication(self) -> None: recurrent = list(recurrent) recurrent[2] = 69 # stale or recycled, not boundary-proven output.scheduled_new_reqs[0].block_ids = (full, tuple(recurrent)) - metadata = connector.build_connector_meta(output) + first_metadata = connector.build_connector_meta(output) + self.assertEqual(first_metadata.plans, []) + metadata = connector.build_connector_meta( + self._cached_scheduler_output( + num_computed_tokens=self.PROMPT_TOKENS, + recurrent_boundary_blocks=boundary_metadata, + ) + ) self.assertEqual(metadata.plans, []) self.assertEqual( connector.counters[ "recurrent_boundary_metadata_rejected" ], - 1, + 0, + ) + self.assertIn( + "dflash-recurrent-boundary", connector._store_progress + ) + connector.request_finished( + types.SimpleNamespace( + request_id="dflash-recurrent-boundary" + ), + [], + ) + self.assertNotIn( + "dflash-recurrent-boundary", connector._store_progress + ) + self.assertNotIn( + "dflash-recurrent-boundary", + connector._store_recurrent_boundaries, ) def test_contradictory_boundary_metadata_skips_publication(self) -> None: invalid_entries = ( [], - [(1, self.BOUNDARY_BLOCK, self.BOUNDARY - 256)], + [(1, self.BOUNDARY_BLOCK, self.BOUNDARY + 256)], [(2, self.BOUNDARY_BLOCK, self.BOUNDARY)], [(1, self.BOUNDARY_BLOCK, self.BOUNDARY), (1, 43, self.BOUNDARY)], [(1, 0, self.BOUNDARY)], @@ -1232,9 +1458,16 @@ def test_contradictory_boundary_metadata_skips_publication(self) -> None: "spark_cache_model_profile": "glm53-flash-hybrid" }, ) + first_metadata = connector.build_connector_meta( + self._scheduler_output() + ) + self.assertEqual(first_metadata.plans, []) metadata = connector.build_connector_meta( - self._scheduler_output( - {"dflash-recurrent-boundary": entries} + self._cached_scheduler_output( + num_computed_tokens=self.PROMPT_TOKENS, + recurrent_boundary_blocks={ + "dflash-recurrent-boundary": entries + }, ) ) self.assertEqual(metadata.plans, []) @@ -1245,6 +1478,52 @@ def test_contradictory_boundary_metadata_skips_publication(self) -> None: 1, ) + def test_future_boundary_rejection_logs_observed_mapping_identity(self) -> None: + output = self._scheduler_output() + request = output.scheduled_new_reqs[0] + request.prompt_token_ids = list(range(self.NONALIGNED_PROMPT_TOKENS)) + output.num_scheduled_tokens[request.req_id] = self.NONALIGNED_PROMPT_TOKENS + with tempfile.TemporaryDirectory() as directory: + connector = _make_connector( + Path(directory), + 0, + block_size=256, + role=KVConnectorRole.SCHEDULER, + override_worker_rank=False, + tp=1, + dcp=1, + kv_cache_config=self._config(), + extra_config={"spark_cache_model_profile": "glm53-flash-hybrid"}, + ) + connector.build_connector_meta(output) + + with mock.patch.object( + connector_module.logger, "warning" + ) as warning: + metadata = connector.build_connector_meta( + self._cached_scheduler_output( + num_computed_tokens=self.NONALIGNED_PROMPT_TOKENS, + recurrent_boundary_blocks={ + request.req_id: [ + ( + 1, + self.COW_BLOCK, + self.NONALIGNED_BOUNDARY + 256, + ) + ] + }, + ) + ) + + self.assertEqual(metadata.plans, []) + warning.assert_called_once() + log_args = warning.call_args.args + self.assertIn( + "entry boundary is ahead of the store plan" + " observed=8448 target=8192 group=1 block=142", + log_args[0] % log_args[1:], + ) + def test_partial_recurrent_group_coverage_skips_publication(self) -> None: config = self._config() second_recurrent = types.SimpleNamespace( @@ -1253,9 +1532,7 @@ def test_partial_recurrent_group_coverage_skips_publication(self) -> None: layer_names=("recurrent-2",), ) config.kv_cache_groups = (*config.kv_cache_groups, second_recurrent) - output = self._scheduler_output( - {"dflash-recurrent-boundary": [(1, 42, self.BOUNDARY)]} - ) + output = self._scheduler_output() output.scheduled_new_reqs[0].block_ids = ( *self._tables(), (0, 0, 0, 81, 82, 83, 84, 85, 86, 87, 88), @@ -1273,7 +1550,17 @@ def test_partial_recurrent_group_coverage_skips_publication(self) -> None: extra_config={"spark_cache_model_profile": "glm53-flash-hybrid"}, ) - metadata = connector.build_connector_meta(output) + first_metadata = connector.build_connector_meta(output) + self.assertEqual(first_metadata.plans, []) + metadata = connector.build_connector_meta( + self._cached_scheduler_output( + num_computed_tokens=self.PROMPT_TOKENS, + recurrent_boundary_blocks={ + "dflash-recurrent-boundary": [(1, 42, self.BOUNDARY)] + }, + group_count=3, + ) + ) self.assertEqual(metadata.plans, []) self.assertEqual( @@ -1295,9 +1582,16 @@ def test_preemption_discards_request_lifetime_boundary_block(self) -> None: extra_config={"spark_cache_model_profile": "glm53-flash-hybrid"}, ) request_id = "dflash-recurrent-boundary" - connector._store_recurrent_boundaries[request_id] = ( - (1, self.BOUNDARY_BLOCK), + connector.build_connector_meta( + self._scheduler_output( + { + request_id: [ + (1, self.BOUNDARY_BLOCK, self.BOUNDARY) + ] + } + ) ) + self.assertIn(request_id, connector._store_recurrent_boundaries) output = types.SimpleNamespace( scheduled_new_reqs=[], scheduled_cached_reqs=types.SimpleNamespace( @@ -1314,6 +1608,66 @@ def test_preemption_discards_request_lifetime_boundary_block(self) -> None: self.assertEqual(metadata.preempted_request_ids, (request_id,)) self.assertNotIn(request_id, connector._store_recurrent_boundaries) + self.assertIn(request_id, connector._store_progress) + + resumed = connector.build_connector_meta( + self._cached_scheduler_output( + num_computed_tokens=self.PROMPT_TOKENS, + resumed=True, + new_block_ids=self._tables(), + ) + ) + self.assertEqual(resumed.plans, []) + self.assertNotIn(request_id, connector._store_recurrent_boundaries) + published = connector.build_connector_meta( + self._cached_scheduler_output( + num_computed_tokens=self.PROMPT_TOKENS + 1, + recurrent_boundary_blocks={ + request_id: [ + (1, self.BOUNDARY_BLOCK + 10, self.BOUNDARY) + ] + }, + ) + ) + self.assertEqual(len(published.plans), 1) + self.assertEqual( + published.plans[0].recurrent_boundary_blocks, + ((1, self.BOUNDARY_BLOCK + 10),), + ) + + def test_quorum_retires_pending_store_without_boundary_proof(self) -> None: + with tempfile.TemporaryDirectory() as directory: + connector = _make_connector( + Path(directory), + 0, + block_size=256, + role=KVConnectorRole.SCHEDULER, + override_worker_rank=False, + tp=1, + dcp=1, + kv_cache_config=self._config(), + extra_config={"spark_cache_model_profile": "glm53-flash-hybrid"}, + ) + request_id = "dflash-recurrent-boundary" + connector.build_connector_meta(self._scheduler_output()) + digest = connector._store_progress[request_id][0] + connector._quorum[digest] = {0} + + metadata = connector.build_connector_meta( + self._cached_scheduler_output( + num_computed_tokens=self.PROMPT_TOKENS, + recurrent_boundary_blocks="malformed-but-unneeded", + ) + ) + + self.assertEqual(metadata.plans, []) + self.assertNotIn(request_id, connector._store_progress) + self.assertNotIn(request_id, connector._store_recurrent_boundaries) + self.assertEqual(connector.counters["store_skipped_quorum"], 1) + self.assertEqual( + connector.counters["recurrent_boundary_metadata_rejected"], + 0, + ) class DigestNamespaceTests(unittest.TestCase):