From 9e779c32b285e00577a7829a75192069d12685dc Mon Sep 17 00:00:00 2001 From: FujitsuPolycom <87842395+FujitsuPolycom@users.noreply.github.com> Date: Sun, 30 Aug 2026 00:18:14 -0500 Subject: [PATCH] Consume proven recurrent replay-boundary blocks SparkCache opts into vLLM recurrent-boundary hand-offs and accepts them only when request, group, token boundary, recurrent topology, and non-null physical block agree. Every aligned recurrent group must have exactly one proven block before a store plan is created. Absent, incomplete, contradictory, or preemption-stale metadata cancels publication. SparkCache never scans arithmetic, running, or speculative slots for a substitute. GPU-free coverage exercises the 6,912-token boundary with seven DFlash verification slots through SparkContextCacheConnector and ManifestStore, including malformed coverage and request-lifetime cleanup. The exact GLM vLLM lease contract accepts and requires the recurrent-boundary runtime postimages and SchedulerOutput interface. Cache namespace impact: none. CacheIdentity values, digest salts, chunk geometry, manifest schemas, and page-delta wire bytes are unchanged. Deployable SparkCache source SHA-256: 01cc59bf2c45af60f02813b771a484e65829e4f464eb27aa1756ca7067c82c9f. Validation: python -m pytest sparkcache -q (752 passed, 7 skipped); python -m pytest deploy -q (108 passed, 1 skipped); python -m ruff check .; source-attestation tests (2 passed). One concurrent validation run exceeded an existing sub-millisecond timing assertion; the isolated test and the following full SparkCache run passed. --- deploy/deepseek_v4/tp4_profile.json | 2 +- deploy/glm52_35bpw/profile.json | 2 +- .../source-receipt.json | 2 +- sparkcache/README.md | 8 +- ...st_glm53_b12x_kda_adaptive_mtp_contract.py | 27 +- ...-contract-glm53-b12x-kda-adaptive-mtp.json | 37 ++- sparkcache/spark_context_cache_connector.py | 162 +++++++++- sparkcache/test_defect_regressions.py | 296 ++++++++++++++++++ 8 files changed, 510 insertions(+), 26 deletions(-) diff --git a/deploy/deepseek_v4/tp4_profile.json b/deploy/deepseek_v4/tp4_profile.json index 35f820f..00266ad 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": "bc7cae86732c869ee8b2205d48ac5be6f580ee8b77a3e4ffd4c69dcd4f1bfae5" + "source_sha256": "01cc59bf2c45af60f02813b771a484e65829e4f464eb27aa1756ca7067c82c9f" }, "model": { "repository": "deepseek-ai/DeepSeek-V4-Flash-0731", diff --git a/deploy/glm52_35bpw/profile.json b/deploy/glm52_35bpw/profile.json index 312be95..438c8b9 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": "bc7cae86732c869ee8b2205d48ac5be6f580ee8b77a3e4ffd4c69dcd4f1bfae5" + "source_sha256": "01cc59bf2c45af60f02813b771a484e65829e4f464eb27aa1756ca7067c82c9f" }, "model": { "repository": "brandonmusic/GLM-5.2-EXL3-TR3v4-3.5bpw-MTP78", diff --git a/patches/vllm-glm53-b12x-kda-adaptive-mtp/source-receipt.json b/patches/vllm-glm53-b12x-kda-adaptive-mtp/source-receipt.json index 2296e78..b32aa50 100644 --- a/patches/vllm-glm53-b12x-kda-adaptive-mtp/source-receipt.json +++ b/patches/vllm-glm53-b12x-kda-adaptive-mtp/source-receipt.json @@ -62,7 +62,7 @@ ], "contract": { "path": "sparkcache/runtime_patches/vllm-kv-block-lease-contract-glm53-b12x-kda-adaptive-mtp.json", - "sha256": "6defde9551cbb586fd09bb2d3020495531b6573397875a767eaae1dbad126024" + "sha256": "45d7a92b38b836a4f829f02df85e339cfeea860e1080e4663a8340af6c125125" }, "result": "All four patches apply in order to the LF source tree, and the eleven-file SparkCache contract verifies the resulting source bytes, including the live-tensor B12X KDA implementation." } diff --git a/sparkcache/README.md b/sparkcache/README.md index b48f193..39919e6 100644 --- a/sparkcache/README.md +++ b/sparkcache/README.md @@ -225,7 +225,13 @@ 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. The + 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 `sparkcache-page-delta-manifest/v1` schema embeds its authenticated base graph, allowing capacity maintenance to retain shared objects after predecessor roots are removed. Restore reconstructs the verified full diff --git a/sparkcache/runtime_patches/test_glm53_b12x_kda_adaptive_mtp_contract.py b/sparkcache/runtime_patches/test_glm53_b12x_kda_adaptive_mtp_contract.py index bceaaf6..a7547e7 100644 --- a/sparkcache/runtime_patches/test_glm53_b12x_kda_adaptive_mtp_contract.py +++ b/sparkcache/runtime_patches/test_glm53_b12x_kda_adaptive_mtp_contract.py @@ -19,6 +19,13 @@ CONTAINERFILE = ROOT / "deploy/glm53_flash/Containerfile.b12x-kda-adaptive-mtp" VLLM_COMMIT = "0b67266a0f37d6146a8403fb8482403c62f412d5" SOURCE_ROLE = "source_built_glm53_b12x_kda_adaptive_mtp" +RECURRENT_BOUNDARY_ROLE = "recurrent_boundary_contract" +RECURRENT_BOUNDARY_FILES = { + "vllm/v1/core/kv_cache_manager.py", + "vllm/v1/core/sched/output.py", + "vllm/v1/core/sched/scheduler.py", + "vllm/v1/core/single_type_kv_cache_manager.py", +} KDA_PATH = "vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py" E105_KDA_SHA256 = ( "a879af0081f69ba8288ef909e1d69b5bbb85bdff7e5aa0d3c11ad892bfea8410" @@ -71,12 +78,24 @@ def test_glm53_b12x_kda_adaptive_mtp_contract_attests_the_complete_sparkcache_vl "vllm/v1/kv_cache_interface.py", KDA_PATH, } - assert all( - set(record["accepted_sha256"]) == {SOURCE_ROLE} - for record in contract["files"] - ) + for record in contract["files"]: + expected_roles = {SOURCE_ROLE} + if record["path"] in RECURRENT_BOUNDARY_FILES: + expected_roles.add(RECURRENT_BOUNDARY_ROLE) + assert set(record["accepted_sha256"]) == expected_roles assert all(record["required_symbols"] for record in contract["files"]) + by_path = {record["path"]: record for record in contract["files"]} + assert by_path["vllm/v1/core/sched/output.py"]["accepted_sha256"][ + RECURRENT_BOUNDARY_ROLE + ] == "9911b3f9d21815a185285852b5a6176e5484e1ab0ff5c30f7caaa68ea0fab543" + assert "SchedulerOutput.recurrent_boundary_blocks" in by_path[ + "vllm/v1/core/sched/output.py" + ]["required_symbols"] + assert "KVCacheManager.take_recurrent_boundary_blocks" in by_path[ + "vllm/v1/core/kv_cache_manager.py" + ]["required_symbols"] + def test_glm53_b12x_kda_contract_rejects_the_e105_kda_source( monkeypatch: pytest.MonkeyPatch, diff --git a/sparkcache/runtime_patches/vllm-kv-block-lease-contract-glm53-b12x-kda-adaptive-mtp.json b/sparkcache/runtime_patches/vllm-kv-block-lease-contract-glm53-b12x-kda-adaptive-mtp.json index 02e8007..9ac1dd6 100644 --- a/sparkcache/runtime_patches/vllm-kv-block-lease-contract-glm53-b12x-kda-adaptive-mtp.json +++ b/sparkcache/runtime_patches/vllm-kv-block-lease-contract-glm53-b12x-kda-adaptive-mtp.json @@ -24,9 +24,10 @@ ] }, { - "path": "vllm/v1/core/sched/scheduler.py", - "accepted_sha256": { - "source_built_glm53_b12x_kda_adaptive_mtp": "6d397c97f31e67a75efc01b5ddd89fa58db425de14fa43965ef2d6146b6b9bdb" + "path": "vllm/v1/core/sched/scheduler.py", + "accepted_sha256": { + "source_built_glm53_b12x_kda_adaptive_mtp": "6d397c97f31e67a75efc01b5ddd89fa58db425de14fa43965ef2d6146b6b9bdb", + "recurrent_boundary_contract": "260f36ce8fabf70c193b20009ea465eea7b1b6c8e9fb72f2307a01ba8fcf7b2a" }, "required_symbols": [ "Scheduler.schedule", @@ -39,9 +40,10 @@ ] }, { - "path": "vllm/v1/core/kv_cache_manager.py", - "accepted_sha256": { - "source_built_glm53_b12x_kda_adaptive_mtp": "ee03dc9ce2b720c0be6e9f572d23580ba96eff68fe3406250557e83071654af0" + "path": "vllm/v1/core/kv_cache_manager.py", + "accepted_sha256": { + "source_built_glm53_b12x_kda_adaptive_mtp": "ee03dc9ce2b720c0be6e9f572d23580ba96eff68fe3406250557e83071654af0", + "recurrent_boundary_contract": "c5b83d382c96b2bf8c466a993ed77123a14a971e2661797128533319388d0b5f" }, "required_symbols": [ "KVCacheManager.get_block_ids", @@ -51,17 +53,20 @@ "KVCacheManager.expire_shared_prefix_leases", "KVCacheManager.discard_shared_prefix_lease", "KVCacheManager.evict_shared_prefix_leases_until_free", - "KVCacheManager.take_kv_cache_block_copies" + "KVCacheManager.take_kv_cache_block_copies", + "KVCacheManager.take_recurrent_boundary_blocks" ] }, { - "path": "vllm/v1/core/sched/output.py", - "accepted_sha256": { - "source_built_glm53_b12x_kda_adaptive_mtp": "65235eba652e5a3ccee18bf3cbfeac9bf4da8fb9c61e961580f612cfb7e593bc" + "path": "vllm/v1/core/sched/output.py", + "accepted_sha256": { + "source_built_glm53_b12x_kda_adaptive_mtp": "65235eba652e5a3ccee18bf3cbfeac9bf4da8fb9c61e961580f612cfb7e593bc", + "recurrent_boundary_contract": "9911b3f9d21815a185285852b5a6176e5484e1ab0ff5c30f7caaa68ea0fab543" }, "required_symbols": [ "SchedulerOutput.preempted_req_ids", - "SchedulerOutput.kv_cache_block_copies" + "SchedulerOutput.kv_cache_block_copies", + "SchedulerOutput.recurrent_boundary_blocks" ] }, { @@ -96,9 +101,10 @@ ] }, { - "path": "vllm/v1/core/single_type_kv_cache_manager.py", - "accepted_sha256": { - "source_built_glm53_b12x_kda_adaptive_mtp": "e4b1c5c38b63b708fd55aa40a9ab0d008b266d006a63dcfcef55890ac1371cb8" + "path": "vllm/v1/core/single_type_kv_cache_manager.py", + "accepted_sha256": { + "source_built_glm53_b12x_kda_adaptive_mtp": "e4b1c5c38b63b708fd55aa40a9ab0d008b266d006a63dcfcef55890ac1371cb8", + "recurrent_boundary_contract": "f67a1850a7e0288baaa6d42e7ec55b22b09c156720767e23acaabedcae333c8a" }, "required_symbols": [ "SingleTypeKVCacheManager.add_local_computed_blocks", @@ -107,7 +113,8 @@ "SingleTypeKVCacheManager.take_pending_cow_copies", "SingleTypeKVCacheManager.pop_blocks_for_free", "MambaManager.remove_skipped_blocks", - "MambaManager.allocate_new_blocks" + "MambaManager.allocate_new_blocks", + "SingleTypeKVCacheManager.take_pending_aligned_recurrent_boundaries" ] }, { diff --git a/sparkcache/spark_context_cache_connector.py b/sparkcache/spark_context_cache_connector.py index 09375bf..22b98fd 100644 --- a/sparkcache/spark_context_cache_connector.py +++ b/sparkcache/spark_context_cache_connector.py @@ -194,6 +194,11 @@ class _ReqPlan: # stored prefix. Empty fields select ordinary full-snapshot publication. base_context_digest: str = "" base_span_tokens: int = 0 + # vLLM may retain a recurrent replay-boundary page outside the request's + # arithmetic block-table slot after later forward work advances the live + # state. Each pair is an exact (group index, physical block id) override + # proven by vLLM for this plan's span_tokens boundary. + recurrent_boundary_blocks: tuple[tuple[int, int], ...] = () # Authenticated row-prefix roots whose descriptors must match the leading # descriptors of this plan. Workers validate these roots before the # leader's blocks may back a shorter shared-prefix lease. @@ -533,6 +538,11 @@ def is_empty(self) -> bool: class SparkContextCacheConnector(KVConnectorBase_V1, SupportsHMA): """Store/restore each rank's DCP shard on rank-local NVMe.""" + # Exact-vLLM runtimes use this opt-in before pinning and exporting aligned + # recurrent replay-boundary blocks. wait_for_save synchronously detaches + # every referenced page before request/preemption cleanup may release it. + supports_recurrent_boundary_blocks = True + configure_streaming_snapshot_runtime = staticmethod( configure_streaming_snapshot_runtime ) @@ -743,6 +753,9 @@ def __init__( # checkpoint/test adapters that seed that tuple remain compatible. self._store_token_ids: dict[str, tuple[int, ...]] = {} self._store_bases: dict[str, tuple[str, int]] = {} + self._store_recurrent_boundaries: dict[ + str, tuple[tuple[int, int], ...] + ] = {} self.counters: dict[str, int] = { "store_committed": 0, "store_failed": 0, @@ -751,6 +764,7 @@ def __init__( "store_skipped_present": 0, "store_skipped_quorum": 0, "page_delta_compactions": 0, + "recurrent_boundary_metadata_rejected": 0, "prefix_alias_publication_attempted": 0, "prefix_alias_publication_failed": 0, "prefix_aliases_published": 0, @@ -1002,11 +1016,33 @@ def _select_group_blocks_for_span( self, groups: tuple[tuple[int, ...], ...], span_tokens: int, + *, + recurrent_boundary_blocks: Sequence[tuple[int, int]] = (), ) -> tuple[tuple[int, ...], ...]: if len(groups) != len(self._group_topology): raise HybridCodecError("request block tables disagree with page groups") + overrides: dict[int, int] = {} + for entry in recurrent_boundary_blocks: + if ( + not isinstance(entry, (list, tuple)) + or len(entry) != 2 + or any(type(value) is not int for value in entry) + ): + raise HybridCodecError("recurrent boundary override is malformed") + group_index, block_id = entry + if ( + not 0 <= group_index < len(self._group_topology) + or group_index in overrides + or block_id <= 0 + or self._group_topology[group_index]["reuse_policy"] + != "recurrent_align" + ): + raise HybridCodecError("recurrent boundary override is incompatible") + overrides[group_index] = block_id trimmed = [] - for group, topology in zip(groups, self._group_topology): + for group_index, (group, topology) in enumerate( + zip(groups, self._group_topology) + ): block_size = int(topology["block_size"]) required = (span_tokens + block_size - 1) // block_size if len(group) < required: @@ -1039,7 +1075,10 @@ def _select_group_blocks_for_span( # SparkCache manifest names one exact full span; its constituent # chunks are never matched as independent prefixes, so only the # reuse window at that span's final boundary is required. - chosen = group[required - selected : required] + if policy == "recurrent_align" and group_index in overrides: + chosen = (overrides[group_index],) + else: + chosen = group[required - selected : required] if any(block <= 0 for block in chosen): raise HybridCodecError( "selected page window contains vLLM's null block" @@ -1569,6 +1608,76 @@ def _append_streaming_snapshot_offer( ) ) + def _validated_recurrent_boundary_blocks( + self, + scheduler_output: "SchedulerOutput", + request_id: str, + boundary_tokens: 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``. + """ + + def reject(reason: str) -> None: + self.counters["recurrent_boundary_metadata_rejected"] += 1 + logger.warning( + "spark-context-cache: recurrent boundary metadata rejected" + " request=%s boundary=%d: %s", + request_id, + boundary_tokens, + reason, + ) + return None + + required_groups = { + group_index + for group_index, topology in enumerate(self._group_topology) + if topology["reuse_policy"] == "recurrent_align" + } + 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") + 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") + if not isinstance(entries, (list, tuple)): + return reject("request value is not a sequence") + + overrides: list[tuple[int, int]] = [] + seen_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") + group_index, block_id, entry_boundary = entry + if any(type(value) is not int for value in entry): + 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) + overrides.append((group_index, block_id)) + if seen_groups != required_groups: + return reject("entries do not cover every aligned recurrent group") + return tuple(sorted(overrides)) + def build_connector_meta( self, scheduler_output: "SchedulerOutput" ) -> KVConnectorMetadata: @@ -1577,6 +1686,11 @@ def build_connector_meta( sorted(getattr(scheduler_output, "preempted_req_ids", None) or ()) ) ) + # vLLM releases request-lifetime recurrent boundary pins on preemption. + # Keep token/table accumulation for a possible resume, but require a + # fresh hash-proven hand-off before the resumed request may publish. + for request_id in meta.preempted_request_ids: + self._store_recurrent_boundaries.pop(request_id, None) for request_id, ( digest, span, @@ -1627,6 +1741,15 @@ 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( @@ -1648,6 +1771,10 @@ 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 already >= span: meta.plans.append( _ReqPlan( @@ -1660,6 +1787,9 @@ 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: @@ -1673,6 +1803,10 @@ 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 @@ -1682,6 +1816,19 @@ 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)) + recurrent_boundary_blocks = self._validated_recurrent_boundary_blocks( + scheduler_output, + req_id, + span, + ) + if recurrent_boundary_blocks is None: + 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) + continue + if recurrent_boundary_blocks: + self._store_recurrent_boundaries[req_id] = recurrent_boundary_blocks new_block_ids = cached.new_block_ids[index] appended = ( [ @@ -1715,6 +1862,7 @@ def build_connector_meta( 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) if self._has_full_quorum(digest): self.counters["store_skipped_quorum"] += 1 continue @@ -1737,6 +1885,10 @@ def build_connector_meta( 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 @@ -1752,6 +1904,7 @@ 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: @@ -3491,6 +3644,7 @@ def request_finished( self._store_progress.pop(request_id, None) self._store_token_ids.pop(request_id, None) self._store_bases.pop(request_id, None) + self._store_recurrent_boundaries.pop(request_id, None) runtime = self._streaming_runtime if runtime is None: return False, None @@ -3747,7 +3901,9 @@ def _snapshot_hybrid_store(self, plan: _ReqPlan) -> _HybridStoreSnapshot: if layout is None: raise RuntimeError("block-page layout was not registered") groups = self._select_group_blocks_for_span( - plan.group_block_ids, plan.span_tokens + plan.group_block_ids, + plan.span_tokens, + recurrent_boundary_blocks=plan.recurrent_boundary_blocks, ) if len(groups) != len(layout.groups): raise HybridCodecError("request block tables disagree with page groups") diff --git a/sparkcache/test_defect_regressions.py b/sparkcache/test_defect_regressions.py index e2f68b7..f6c72e2 100644 --- a/sparkcache/test_defect_regressions.py +++ b/sparkcache/test_defect_regressions.py @@ -1023,6 +1023,302 @@ class MambaSpec: ) +class DefectD17RecurrentBoundaryMetadataTests(unittest.TestCase): + """D-17: vLLM identifies off-table recurrent replay boundaries.""" + + BOUNDARY = 6912 + PROMPT_TOKENS = 6992 + BOUNDARY_BLOCK = 42 + + @staticmethod + def _config() -> types.SimpleNamespace: + class FullAttentionSpec: + block_size = 2304 + storage_block_size = 2304 + page_size_bytes = 64 + + class MambaSpec: + block_size = 2304 + storage_block_size = 2304 + page_size_bytes = 64 + mamba_cache_mode = "align" + tokens_per_state = 2304 + num_speculative_blocks = 7 + num_prefill_checkpoint_blocks = 0 + + return types.SimpleNamespace( + kv_cache_groups=( + types.SimpleNamespace( + kv_cache_spec=FullAttentionSpec(), + is_eagle_group=False, + layer_names=("full",), + ), + types.SimpleNamespace( + kv_cache_spec=MambaSpec(), + is_eagle_group=False, + layer_names=("recurrent",), + ), + ) + ) + + @classmethod + def _tables(cls) -> tuple[tuple[int, ...], ...]: + # The recurrent replay-boundary slot is null after vLLM advances the + # running state. Entries 71..78 are the later running state and seven + # DFlash verification slots, so none can substitute for block 42. + return ((11, 12, 13, 14), (0, 0, 0, 71, 72, 73, 74, 75, 76, 77, 78)) + + @classmethod + def _scheduler_output( + cls, + recurrent_boundary_blocks: object = None, + ) -> types.SimpleNamespace: + request_id = "dflash-recurrent-boundary" + output = types.SimpleNamespace( + scheduled_new_reqs=[ + types.SimpleNamespace( + req_id=request_id, + prompt_token_ids=list(range(cls.PROMPT_TOKENS)), + block_ids=cls._tables(), + num_computed_tokens=0, + ) + ], + scheduled_cached_reqs=types.SimpleNamespace( + req_ids=[], + resumed_req_ids=set(), + num_computed_tokens=[], + new_block_ids=[], + ), + num_scheduled_tokens={request_id: cls.PROMPT_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) + config = self._config() + scheduler = _make_connector( + root, + 0, + block_size=256, + role=KVConnectorRole.SCHEDULER, + override_worker_rank=False, + tp=1, + dcp=1, + 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) + ] + } + ) + + self.assertTrue(scheduler.supports_recurrent_boundary_blocks) + metadata = scheduler.build_connector_meta(output) + + self.assertEqual(len(metadata.plans), 1) + plan = metadata.plans[0] + self.assertEqual(plan.span_tokens, self.BOUNDARY) + self.assertEqual(plan.recurrent_boundary_blocks, ((1, 42),)) + worker = _make_connector( + root, + 0, + block_size=256, + tp=1, + dcp=1, + kv_cache_config=config, + extra_config={"spark_cache_model_profile": "glm53-flash-hybrid"}, + ) + pools = { + name: ( + torch.arange(128 * 64, dtype=torch.int32) + .add(offset) + .remainder(251) + .to(torch.uint8) + .reshape(128, 1, 64) + ) + for name, offset in (("full", 0), ("recurrent", 29)) + } + worker.register_kv_caches(pools) + expected_full = pools["full"][[11, 12, 13]].clone() + expected_recurrent = pools["recurrent"][[self.BOUNDARY_BLOCK]].clone() + + worker._store_one(plan) + + lookup = worker._store.lookup(worker._identity(0), plan.digest) + self.assertTrue(lookup.is_hit, lookup.reason) + destination = ( + (90, 91, 92), + (0, 0, 93, 94, 95, 96, 97, 98, 99, 100, 101), + ) + pools["full"][[90, 91, 92]].zero_() + pools["recurrent"][[93]].zero_() + self.assertTrue( + worker._load_one( + _ReqPlan( + "dflash-recurrent-restore", + plan.digest, + self.BOUNDARY, + destination[0], + False, + block_ids_by_group=destination, + ) + ) + ) + self.assertTrue(torch.equal(pools["full"][[90, 91, 92]], expected_full)) + self.assertTrue( + torch.equal(pools["recurrent"][[93]], expected_recurrent) + ) + + def test_missing_or_wrong_request_metadata_skips_publication(self) -> None: + for boundary_metadata in ( + None, + {"another-request": [(1, self.BOUNDARY_BLOCK, self.BOUNDARY)]}, + ): + with self.subTest(boundary_metadata=boundary_metadata): + 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(boundary_metadata) + full, recurrent = output.scheduled_new_reqs[0].block_ids + 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) + self.assertEqual(metadata.plans, []) + self.assertEqual( + connector.counters[ + "recurrent_boundary_metadata_rejected" + ], + 1, + ) + + def test_contradictory_boundary_metadata_skips_publication(self) -> None: + invalid_entries = ( + [], + [(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)], + [(0, self.BOUNDARY_BLOCK, self.BOUNDARY)], + ) + for entries in invalid_entries: + with self.subTest(entries=entries): + 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" + }, + ) + metadata = connector.build_connector_meta( + self._scheduler_output( + {"dflash-recurrent-boundary": entries} + ) + ) + self.assertEqual(metadata.plans, []) + self.assertEqual( + connector.counters[ + "recurrent_boundary_metadata_rejected" + ], + 1, + ) + + def test_partial_recurrent_group_coverage_skips_publication(self) -> None: + config = self._config() + second_recurrent = types.SimpleNamespace( + kv_cache_spec=config.kv_cache_groups[1].kv_cache_spec, + is_eagle_group=False, + 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.scheduled_new_reqs[0].block_ids = ( + *self._tables(), + (0, 0, 0, 81, 82, 83, 84, 85, 86, 87, 88), + ) + 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=config, + extra_config={"spark_cache_model_profile": "glm53-flash-hybrid"}, + ) + + metadata = connector.build_connector_meta(output) + + self.assertEqual(metadata.plans, []) + self.assertEqual( + connector.counters["recurrent_boundary_metadata_rejected"], + 1, + ) + + def test_preemption_discards_request_lifetime_boundary_block(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._store_recurrent_boundaries[request_id] = ( + (1, self.BOUNDARY_BLOCK), + ) + output = types.SimpleNamespace( + scheduled_new_reqs=[], + scheduled_cached_reqs=types.SimpleNamespace( + req_ids=[], + resumed_req_ids=set(), + num_computed_tokens=[], + new_block_ids=[], + ), + num_scheduled_tokens={}, + preempted_req_ids={request_id}, + ) + + metadata = connector.build_connector_meta(output) + + self.assertEqual(metadata.preempted_request_ids, (request_id,)) + self.assertNotIn(request_id, connector._store_recurrent_boundaries) + + class DigestNamespaceTests(unittest.TestCase): """D-4: context digests are identical across roles and physical ranks."""