diff --git a/mooncake-reshard/pyrightconfig.json b/mooncake-reshard/pyrightconfig.json index 701e4bad8a..ae824d6590 100644 --- a/mooncake-reshard/pyrightconfig.json +++ b/mooncake-reshard/pyrightconfig.json @@ -17,6 +17,7 @@ "python/mooncake/reshard/kv_cache/snapshot.py", "python/mooncake/reshard/kv_cache/snapshot_serde.py", "python/mooncake/reshard/kv_cache/types.py", + "python/mooncake/reshard/kv_cache/identity.py", "python/mooncake/reshard/kv_cache/topology.py", "python/mooncake/reshard/kv_cache/part.py", "python/mooncake/reshard/kv_cache/placement.py", @@ -25,11 +26,14 @@ "python/mooncake/reshard/kv_cache/serde.py", "python/mooncake/reshard/kv_cache/planner.py", "python/mooncake/reshard/kv_cache/plan_serde.py", + "python/mooncake/reshard/kv_cache/_store", "python/mooncake/reshard/weight/_planner", "python/mooncake/reshard/weight/_te", "python/mooncake/reshard/transfer_engine" ], - "extraPaths": ["python"], + "extraPaths": [ + "python" + ], "pythonVersion": "3.10", "typeCheckingMode": "strict", "reportPrivateUsage": "none", diff --git a/mooncake-reshard/python/mooncake/reshard/kv_cache/__init__.py b/mooncake-reshard/python/mooncake/reshard/kv_cache/__init__.py index 43b7876327..57169833e1 100644 --- a/mooncake-reshard/python/mooncake/reshard/kv_cache/__init__.py +++ b/mooncake-reshard/python/mooncake/reshard/kv_cache/__init__.py @@ -1,7 +1,30 @@ """Framework-neutral KV-cache reshard contracts and planning API.""" +from ._store.binding import KVCacheStoreRuntimeBinding +from ._store.manifest import ( + KVCacheStoreFormat, + KVCacheStoreFragment, + KVCacheStoreLayout, + KVCacheStoreManifest, + KVCacheStoreShard, +) +from ._store.planner import ( + KVCacheStoreLoadPlan, + KVCacheStoreLoadRange, + KVCacheStorePlanningLimits, + KVCacheStoreUploadPlan, + plan_kv_cache_store_load, + plan_kv_cache_store_upload, +) +from ._store.serde import ( + kv_cache_store_layout_from_json, + kv_cache_store_layout_to_json, + kv_cache_store_manifest_from_json, + kv_cache_store_manifest_to_json, +) from .completion import KVCacheCompletion, KVCacheTargetReceipt, KVCacheWriterReceipt from .executor import KVCacheTransferEngineExecutor +from .identity import KVCacheModelDomain from .manifest import ( KVCacheBufferBinding, KVCacheComponent, @@ -73,6 +96,7 @@ "KVCacheDescriptor", "KVCacheLayout", "KVCacheLogicalTransferPlan", + "KVCacheModelDomain", "KVCachePlacementManifest", "KVCachePlacementPart", "KVCachePreparedTransferEdge", @@ -86,6 +110,16 @@ "KVCacheRuntimeBuffer", "KVCacheRuntimeTransferPlan", "KVCacheSnapshotDescriptor", + "KVCacheStoreFormat", + "KVCacheStoreFragment", + "KVCacheStoreLayout", + "KVCacheStoreLoadPlan", + "KVCacheStoreLoadRange", + "KVCacheStoreManifest", + "KVCacheStorePlanningLimits", + "KVCacheStoreRuntimeBinding", + "KVCacheStoreShard", + "KVCacheStoreUploadPlan", "KVCacheTargetReceipt", "KVCacheTopology", "KVCacheTopologyParticipant", @@ -110,11 +144,17 @@ "kv_cache_runtime_transfer_to_json", "kv_cache_snapshot_from_json", "kv_cache_snapshot_to_json", + "kv_cache_store_layout_from_json", + "kv_cache_store_layout_to_json", + "kv_cache_store_manifest_from_json", + "kv_cache_store_manifest_to_json", "kv_cache_target_receipt_from_json", "kv_cache_target_receipt_to_json", "kv_cache_writer_receipt_from_json", "kv_cache_writer_receipt_to_json", "placement_fragment_id", + "plan_kv_cache_store_load", + "plan_kv_cache_store_upload", "plan_kv_cache_transfer_to_local_target", "prepare_kv_cache_transfer", "resolve_contiguous_runtime_binding", diff --git a/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/__init__.py b/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/__init__.py new file mode 100644 index 0000000000..ec5348b21c --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/__init__.py @@ -0,0 +1 @@ +"""Internal KV Store contracts and planning; public exports live in kv_cache.""" diff --git a/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/_utils.py b/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/_utils.py new file mode 100644 index 0000000000..b5477edebe --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/_utils.py @@ -0,0 +1,257 @@ +"""Internal helpers for Store KV reshard geometry, batching and metadata.""" + +from __future__ import annotations + +import hashlib +import json +from bisect import bisect_right +from collections.abc import Iterable, Iterator, Mapping +from typing import Generic, NamedTuple, TypeVar, cast + +from ..resolved import KVCacheResolvedRange +from ..types import KVCacheComponent, KVCacheDescriptor + +_T = TypeVar("_T") + + +class HeadIntervalIndex(Generic[_T]): + """Index validated, disjoint intervals with optional exact-interval replicas. + + Candidates within an interval retain caller order. Callers validate source + ownership; the index contains no object keys or runtime addresses. + """ + + def __init__(self, entries: Iterable[tuple[int, int, _T]]) -> None: + grouped: dict[tuple[int, int], list[_T]] = {} + for start, end, value in entries: + grouped.setdefault((start, end), []).append(value) + self._runs = tuple( + (start, end, tuple(grouped[(start, end)])) for start, end in sorted(grouped) + ) + for previous, current in zip(self._runs, self._runs[1:]): + if previous[1] > current[0]: + raise ValueError( + "overlapping source head intervals without exact replicas" + ) + self._starts = tuple(run[0] for run in self._runs) + + def cover(self, start: int, end: int) -> Iterator[tuple[int, int, tuple[_T, ...]]]: + """Yield a complete clipped cover without visiting individual heads.""" + index = max(0, bisect_right(self._starts, start) - 1) + cursor = start + while index < len(self._runs) and self._runs[index][0] < end: + lo, hi, candidates = self._runs[index] + lo, hi = max(start, lo), min(end, hi) + if lo != cursor or hi <= lo: + raise ValueError("source head coverage has a gap") + yield lo, hi, candidates + cursor = hi + index += 1 + if cursor != end: + raise ValueError("source head coverage is incomplete") + + +class KVCacheLogicalShard(NamedTuple): + layer_ids: tuple[int, ...] + head_start: int + head_count: int + replica_ordinal: int = 0 + replica_count: int = 1 + + +def validate_logical_coverage( + descriptor: KVCacheDescriptor, + shards: tuple[KVCacheLogicalShard, ...], + *, + label: str, +) -> None: + """Prove one complete logical copy, allowing only exact declared replicas.""" + layers: dict[int, list[KVCacheLogicalShard]] = { + layer: [] for layer in descriptor.global_layer_ids + } + for shard in shards: + end = shard.head_start + shard.head_count + if end > descriptor.total_kv_heads: + raise ValueError(f"{label} head interval exceeds total_kv_heads") + for layer in shard.layer_ids: + if layer not in layers: + raise ValueError(f"{label} contains an unknown global layer") + layers[layer].append(shard) + for layer, items in layers.items(): + try: + index = HeadIntervalIndex( + (s.head_start, s.head_start + s.head_count, s) for s in items + ) + for _, _, owners in index.cover(0, descriptor.total_kv_heads): + if len(owners) == 1: + if owners[0].replica_count != 1: + raise ValueError("incomplete declared KV-head replica set") + elif {s.replica_count for s in owners} != {len(owners)} or { + s.replica_ordinal for s in owners + } != set(range(len(owners))): + raise ValueError( + "overlapping KV heads must be an exact declared replica set" + ) + except ValueError as error: + raise ValueError(f"{label} layer={layer}: {error}") from error + + +class KVCacheRangeGeometry(NamedTuple): + token_start: int + token_count: int + head_start: int + head_count: int + offset_bytes: int + token_stride_bytes: int + head_stride_bytes: int + + +class KVCacheCopyGeometry(NamedTuple): + source: KVCacheRangeGeometry | KVCacheResolvedRange + target: KVCacheRangeGeometry | KVCacheResolvedRange + token_start: int + token_count: int + head_start: int + head_count: int + token_step: int + head_step: int + head_bytes: int + + @property + def segment_count(self) -> int: + return (self.token_count // self.token_step) * ( + self.head_count // self.head_step + ) + + @property + def nbytes(self) -> int: + return self.token_count * self.head_count * self.head_bytes + + def segments(self) -> Iterator[tuple[int, int, int]]: + """Yield source offset, target offset, length in token/head order.""" + source, target = self.source, self.target + source_base = ( + source.offset_bytes + + (self.token_start - source.token_start) * source.token_stride_bytes + + (self.head_start - source.head_start) * source.head_stride_bytes + ) + target_base = ( + target.offset_bytes + + (self.token_start - target.token_start) * target.token_stride_bytes + + (self.head_start - target.head_start) * target.head_stride_bytes + ) + size = self.token_step * self.head_step * self.head_bytes + for token in range(0, self.token_count, self.token_step): + source_row = source_base + token * source.token_stride_bytes + target_row = target_base + token * target.token_stride_bytes + for head in range(0, self.head_count, self.head_step): + yield ( + source_row + head * source.head_stride_bytes, + target_row + head * target.head_stride_bytes, + size, + ) + + +def plan_range_copy( + source: KVCacheRangeGeometry | KVCacheResolvedRange, + target: KVCacheRangeGeometry | KVCacheResolvedRange, + head_bytes: int, + *, + head_window: tuple[int, int] | None = None, +) -> KVCacheCopyGeometry | None: + """Intersect logical ranges and coalesce only bytes contiguous on both sides.""" + start = max(source.token_start, target.token_start) + end = min( + source.token_start + source.token_count, target.token_start + target.token_count + ) + h0 = max(source.head_start, target.head_start) + h1 = min( + source.head_start + source.head_count, target.head_start + target.head_count + ) + if head_window is not None: + h0, h1 = max(h0, head_window[0]), min(h1, head_window[1]) + if start >= end or h0 >= h1: + return None + head_step = ( + h1 - h0 + if source.head_stride_bytes == target.head_stride_bytes == head_bytes + else 1 + ) + width = head_step * head_bytes + token_step = ( + end - start + if head_step == h1 - h0 + and source.token_stride_bytes == target.token_stride_bytes == width + else 1 + ) + return KVCacheCopyGeometry( + source, + target, + start, + end - start, + h0, + h1 - h0, + token_step, + head_step, + head_bytes, + ) + + +def index_runtime_ranges( + ranges: tuple[KVCacheResolvedRange, ...], +) -> dict[tuple[int, KVCacheComponent], list[KVCacheResolvedRange]]: + """Index each validated runtime range once, preserving its canonical order.""" + result: dict[tuple[int, KVCacheComponent], list[KVCacheResolvedRange]] = {} + for item in ranges: + result.setdefault((item.global_layer_id, item.component), []).append(item) + return result + + +def canonical_json(value: object, *, allow_nan: bool = True) -> str: + """Keep canonical encoding identical across IDs and wire envelopes.""" + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=allow_nan) + + +def canonical_digest(value: object) -> str: + encoded = canonical_json(value).encode() + return hashlib.sha256(encoded).hexdigest() + + +def json_string(value: object, label: str) -> str: + if type(value) is not str or not value: + raise ValueError(f"{label} must be a non-empty string") + return value + + +def json_integer(value: object, label: str, *, minimum: int = 0) -> int: + if type(value) is not int or value < minimum: + raise ValueError(f"{label} must be an integer at least {minimum}") + return value + + +def load_json_object(value: str, label: str) -> Mapping[str, object]: + if type(value) is not str or len(value.encode("utf-8")) > 16 * 1024 * 1024: + raise ValueError("KV-cache JSON must be text within the 16 MiB wire limit") + + def reject_constant(constant: str) -> None: + raise ValueError(f"non-finite JSON number is unsupported: {constant}") + + def reject_duplicate_fields(pairs: list[tuple[str, object]]) -> dict[str, object]: + result: dict[str, object] = {} + for key, item in pairs: + if key in result: + raise ValueError(f"duplicate JSON field: {key}") + result[key] = item + return result + + try: + payload = json.loads( + value, + parse_constant=reject_constant, + object_pairs_hook=reject_duplicate_fields, + ) + except (TypeError, json.JSONDecodeError) as error: + raise ValueError(f"{label} is not valid JSON") from error + if not isinstance(payload, Mapping): + raise ValueError(f"{label} must be a JSON object") # noqa: TRY004 + return cast(Mapping[str, object], payload) diff --git a/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/binding.py b/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/binding.py new file mode 100644 index 0000000000..23cd7ea84b --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/binding.py @@ -0,0 +1,103 @@ +"""Store-specific local binding, independent of the R2R snapshot lifecycle.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from ..placement import KVCachePlacementManifest +from ..resolved import ( + KVCacheRegisteredRegion, + KVCacheResolvedRange, + KVCacheTransferLimits, + _range_key, + _validated_range_spans, +) +from ..types import ( + require_integer, + require_manifest_items, + require_nonempty_string, + require_sha256, +) + + +@dataclass(frozen=True) +class KVCacheStoreRuntimeBinding: + """Operation-scoped local memory for Store I/O; no R2R snapshot identity.""" + + operation_id: str + resource_id: str + placement_id: str + placement_digest: str + instance_id: str + revision: str + participant_id: str + regions: tuple[KVCacheRegisteredRegion, ...] + ranges: tuple[KVCacheResolvedRange, ...] + + def __post_init__(self) -> None: + for name in ( + "operation_id", + "resource_id", + "placement_id", + "instance_id", + "revision", + "participant_id", + ): + require_nonempty_string(getattr(self, name), name) + require_sha256(self.placement_digest, "placement_digest") + regions = require_manifest_items( + self.regions, "registered regions", KVCacheRegisteredRegion + ) + ranges = require_manifest_items( + self.ranges, "resolved ranges", KVCacheResolvedRange + ) + if not regions or not ranges: + raise ValueError("resolved binding must contain regions and ranges") + if len({r.region_id for r in regions}) != len(regions): + raise ValueError("duplicate registered region_id") + if len({r.endpoint for r in regions}) != 1: + raise ValueError("one runtime binding must use a single endpoint") + object.__setattr__( + self, "regions", tuple(sorted(regions, key=lambda r: r.region_id)) + ) + object.__setattr__(self, "ranges", tuple(sorted(ranges, key=_range_key))) + + +def validate_store_binding( + placement: KVCachePlacementManifest, + binding: KVCacheStoreRuntimeBinding, + operation_id: str, + token_start: int, + token_count: int, + limits: KVCacheTransferLimits, +) -> list[tuple[int, int]]: + """Validate Store page coverage, ownership, strides and registered bounds.""" + if not isinstance(binding, KVCacheStoreRuntimeBinding): + raise TypeError("binding must be a KVCacheStoreRuntimeBinding") + if not isinstance(placement, KVCachePlacementManifest): + raise TypeError("placement must be a KVCachePlacementManifest") + if not isinstance(limits, KVCacheTransferLimits): + raise TypeError("limits must be KVCacheTransferLimits") + require_nonempty_string(operation_id, "operation_id") + require_integer(token_start, "token_start") + require_integer(token_count, "token_count", minimum=1) + token_end = require_integer(token_start + token_count, "token_end") + expected = { + "operation_id": operation_id, + "resource_id": placement.resource_id, + "placement_id": placement.placement_id, + "placement_digest": placement.digest, + "revision": placement.revision, + } + for name, value in expected.items(): + if getattr(binding, name) != value: + raise ValueError(f"resolved binding {name} differs") + return _validated_range_spans( + placement, + binding.participant_id, + binding.regions, + binding.ranges, + token_start=token_start, + token_end=token_end, + limits=limits, + ) diff --git a/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/manifest.py b/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/manifest.py new file mode 100644 index 0000000000..5abfd9c5ea --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/manifest.py @@ -0,0 +1,272 @@ +"""Shared, address-free KV Store layout and model metadata; no page snapshots.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +from enum import Enum +from functools import cached_property + +from ..identity import KVCacheModelDomain +from ..types import ( + KVCacheComponent, + KVCacheDescriptor, + require_integer, + require_manifest_items, + require_nonempty_string, + require_sha256, +) +from ._utils import KVCacheLogicalShard, validate_logical_coverage +from ._utils import canonical_digest as _canonical_digest + + +class KVCacheStoreFormat(str, Enum): + """Axis order within each K or V component of a Part; D is contiguous.""" + + PLHD = "PLHD" # [token, layer, head, dim] + LPHD = "LPHD" # [layer, token, head, dim] + HPLD = "HPLD" # [head, token, layer, dim] + + +@dataclass(frozen=True) +class KVCacheStoreShard: + """Logical coverage of one Part. key_suffix is its stable writer label, + retained in the manifest wire format; it is never appended to a page key.""" + + layer_ids: tuple[int, ...] + head_start: int + head_count: int + key_suffix: str + + def __post_init__(self) -> None: + require_nonempty_string(self.key_suffix, "key_suffix") + layers = require_manifest_items(self.layer_ids, "shard layers", int) + if not layers or len(set(layers)) != len(layers): + raise ValueError("shard layers must be nonempty and unique") + for layer in layers: + require_integer(layer, "global_layer_id") + require_integer(self.head_start, "head_start") + require_integer(self.head_count, "head_count", minimum=1) + require_integer(self.head_start + self.head_count, "head_end") + object.__setattr__(self, "layer_ids", tuple(sorted(layers))) + + @property + def logical_shard(self) -> KVCacheLogicalShard: + """Project one durable logical copy without its Store key suffix.""" + return KVCacheLogicalShard(self.layer_ids, self.head_start, self.head_count) + + +@dataclass(frozen=True) +class KVCacheStoreFragment: + """Relative byte mapping for one layer/component and a head interval.""" + + global_layer_id: int + component: KVCacheComponent + head_start: int + head_count: int + part_index: int + part_offset: int + token_stride_bytes: int + head_stride_bytes: int + + def __post_init__(self) -> None: + for name in ("global_layer_id", "head_start", "part_index", "part_offset"): + require_integer(getattr(self, name), name) + for name in ("head_count", "token_stride_bytes", "head_stride_bytes"): + require_integer(getattr(self, name), name, minimum=1) + require_integer(self.head_start + self.head_count, "head_end") + if not isinstance(self.component, KVCacheComponent): + raise ValueError("component must be a KVCacheComponent") # noqa: TRY004 + + +@dataclass(frozen=True) +class KVCacheStoreLayout: + """A fixed-page object template independent of models, pages and workers. + + Persist only shard geometry and axis order. Sizes and strides are derived, + so a manifest cannot give conflicting descriptions of the same bytes. + """ + + descriptor: KVCacheDescriptor + object_format: KVCacheStoreFormat + shards: tuple[KVCacheStoreShard, ...] + + def __post_init__(self) -> None: + if not isinstance(self.descriptor, KVCacheDescriptor): + raise ValueError("descriptor must be a KVCacheDescriptor") # noqa: TRY004 + if not isinstance(self.object_format, KVCacheStoreFormat): + raise ValueError( # noqa: TRY004 + "object_format must be a KVCacheStoreFormat" + ) + shards = require_manifest_items(self.shards, "Store shards", KVCacheStoreShard) + if len({shard.key_suffix for shard in shards}) != len(shards): + raise ValueError("Store shard key suffixes must be unique") + if not shards or len(shards) * 2 > 100_000: + raise ValueError("Store object count limit exceeded or empty layout") + if sum(len(shard.layer_ids) * 2 for shard in shards) > 100_000: + raise ValueError("Store fragment count limit exceeded") + object.__setattr__( + self, + "shards", + tuple(sorted(shards, key=lambda s: (s.layer_ids, s.head_start))), + ) + _validate_layout_coverage(self) + if sum(self.part_sizes) > 1 << 40: + raise ValueError("Store layout byte limit exceeded") + + @property + def token_count(self) -> int: + """Every object contains exactly one full page, never a variable tail.""" + return self.descriptor.page_size + + @property + def part_sizes(self) -> tuple[int, ...]: + """Each Part stores all K bytes followed by all V bytes.""" + d = self.descriptor + return tuple( + require_integer( + len(shard.layer_ids) + * self.token_count + * shard.head_count + * (d.key_head_dim + d.value_head_dim) + * d.itemsize, + "object size", + minimum=1, + ) + for shard in self.shards + ) + + @property + def fragments(self) -> tuple[KVCacheStoreFragment, ...]: + """Derive affine offsets without enumerating tokens or individual bytes.""" + result: list[KVCacheStoreFragment] = [] + d = self.descriptor + for slot, shard in enumerate(self.shards): + layers, heads, tokens = ( + len(shard.layer_ids), + shard.head_count, + self.token_count, + ) + for component_index, component in enumerate(KVCacheComponent): + dim = ( + d.key_head_dim + if component is KVCacheComponent.KEY + else d.value_head_dim + ) + head_bytes = dim * d.itemsize + if self.object_format is KVCacheStoreFormat.PLHD: + layer_stride, token_stride, head_stride = ( + heads * head_bytes, + layers * heads * head_bytes, + head_bytes, + ) + elif self.object_format is KVCacheStoreFormat.LPHD: + layer_stride, token_stride, head_stride = ( + tokens * heads * head_bytes, + heads * head_bytes, + head_bytes, + ) + else: + layer_stride, token_stride, head_stride = ( + head_bytes, + layers * head_bytes, + tokens * layers * head_bytes, + ) + for local_layer, global_layer in enumerate(shard.layer_ids): + result.append( + KVCacheStoreFragment( + global_layer, + component, + shard.head_start, + heads, + slot, + local_layer * layer_stride + + ( + layers * tokens * heads * d.key_head_dim * d.itemsize + if component_index + else 0 + ), + token_stride, + head_stride, + ) + ) + return tuple(result) + + @cached_property + def digest(self) -> str: + return _canonical_digest(asdict(self)) + + @property + def layout_id(self) -> str: + return f"sha256:{self.digest}" + + +def _validate_layout_coverage(layout: KVCacheStoreLayout) -> None: + """Prove each global layer/head occurs once; dense formats then cover bytes.""" + validate_logical_coverage( + layout.descriptor, + tuple(shard.logical_shard for shard in layout.shards), + label="Store layout", + ) + + +@dataclass(frozen=True) +class KVCacheStoreManifest: + """One immutable layout record shared by pages stored in that layout. + + Constructing it does not publish a layout or prove any payload exists. + The page ObjectMetadata binds manifest_key at its first Part write. + """ + + namespace: str + model_id: str + model_revision: str + semantic_fingerprint: str + layout: KVCacheStoreLayout + + def __post_init__(self) -> None: + for name in ("namespace", "model_id", "model_revision", "semantic_fingerprint"): + require_nonempty_string(getattr(self, name), name) + if not isinstance(self.layout, KVCacheStoreLayout): + raise ValueError("layout must be a KVCacheStoreLayout") # noqa: TRY004 + + @cached_property + def model_domain(self) -> str: + """Share discovery across TP/PP geometry and physical axis orders.""" + return KVCacheModelDomain( + self.namespace, + self.model_id, + self.model_revision, + self.semantic_fingerprint, + self.layout.descriptor, + ).digest + + @property + def manifest_key(self) -> str: + return _manifest_key(self.model_domain, self.layout.layout_id) + + @cached_property + def digest(self) -> str: + return _canonical_digest(asdict(self)) + + +def _layout_digest(layout_id: str) -> str: + """Decode a canonical layout ID before using it in object keys.""" + require_nonempty_string(layout_id, "layout_id") + if not layout_id.startswith("sha256:"): + raise ValueError("layout_id must start with sha256:") + return require_sha256(layout_id[7:], "layout digest") + + +def _manifest_key(model_domain: str, layout_id: str) -> str: + """Content-addressed manifest location shared by all pages of a layout.""" + require_sha256(model_domain, "model_domain") + return f"kv/{model_domain}/manifests/{_layout_digest(layout_id)}" + + +__all__ = [ + "KVCacheStoreFormat", + "KVCacheStoreFragment", + "KVCacheStoreLayout", + "KVCacheStoreManifest", + "KVCacheStoreShard", +] diff --git a/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/planner.py b/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/planner.py new file mode 100644 index 0000000000..203e0e4b04 --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/planner.py @@ -0,0 +1,283 @@ +"""Fixed-page Store geometry planning; no snapshots, runtime addresses or I/O.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace + +from ...contracts import ParticipantId +from ..part import KVCachePlacementPart +from ..placement import KVCachePlacementManifest +from ..types import KVCacheComponent, require_integer, require_nonempty_string +from ._utils import HeadIntervalIndex +from ._utils import canonical_digest as _canonical_digest +from .manifest import ( + KVCacheStoreFormat, + KVCacheStoreFragment, + KVCacheStoreLayout, + KVCacheStoreShard, +) + + +@dataclass(frozen=True) +class KVCacheStorePlanningLimits: + """Bound metadata and useful bytes before materializing a transfer plan.""" + + max_parts: int = 100_000 + max_fragments: int = 100_000 + max_bytes: int = 1 << 40 + + def __post_init__(self) -> None: + for name in self.__dataclass_fields__: + require_integer(getattr(self, name), name, minimum=1) + if max(self.max_parts, self.max_fragments) > 100_000: + raise ValueError("planning limits exceed the contract collection limit") + if self.max_bytes > 1 << 40: + raise ValueError("planning byte limit exceeds the 1 TiB contract limit") + + +DEFAULT_STORE_PLANNING_LIMITS = KVCacheStorePlanningLimits() + + +def _selected_parts( + placement: KVCachePlacementManifest, + dp_rank: int | None, + *, + replicas: bool, +) -> tuple[int, tuple[KVCachePlacementPart, ...]]: + """Upload selects one head replica; load fills every declared target replica.""" + if not isinstance(placement, KVCachePlacementManifest): + raise ValueError("placement must be a KVCachePlacementManifest") # noqa: TRY004 + selected = min(placement.dp_ranks) if dp_rank is None else dp_rank + require_integer(selected, "dp_rank") + if selected not in placement.dp_ranks: + raise ValueError("DP rank is absent from placement") + return selected, tuple( + part + for part in placement.parts + if part.rank.dp == selected + and part.layer_ids + and (replicas or part.replica_ordinal == 0) + ) + + +def _check_limits( + layout: KVCacheStoreLayout, limits: KVCacheStorePlanningLimits +) -> None: + if not isinstance(limits, KVCacheStorePlanningLimits): + raise ValueError("limits must be KVCacheStorePlanningLimits") # noqa: TRY004 + if len(layout.shards) > limits.max_parts: + raise ValueError("Store part limit exceeded") + if sum(len(s.layer_ids) * 2 for s in layout.shards) > limits.max_fragments: + raise ValueError("Store fragment limit exceeded") + if sum(layout.part_sizes) > limits.max_bytes: + raise ValueError("Store byte limit exceeded") + + +@dataclass(frozen=True, init=False) +class KVCacheStoreUploadPlan: + """Source-native layout plus one complete writer per Part; no reshard target.""" + + source_placement: KVCachePlacementManifest + source_dp_rank: int + layout: KVCacheStoreLayout + part_writers: tuple[ParticipantId, ...] + + def __init__( + self, + source_placement: KVCachePlacementManifest, + *, + object_format: KVCacheStoreFormat, + source_dp_rank: int | None = None, + limits: KVCacheStorePlanningLimits = DEFAULT_STORE_PLANNING_LIMITS, + ) -> None: + selected, parts = _selected_parts( + source_placement, source_dp_rank, replicas=False + ) + parts = tuple(sorted(parts, key=lambda p: (p.layer_ids, p.head_start))) + layout = KVCacheStoreLayout( + source_placement.descriptor, + object_format, + tuple( + KVCacheStoreShard( + p.layer_ids, + p.head_start, + p.head_count, + f"{p.rank.tp}_{p.rank.pp}" + if source_placement.topology.pp_size > 1 + else str(p.rank.tp), + ) + for p in parts + ), + ) + _check_limits(layout, limits) + object.__setattr__(self, "source_placement", source_placement) + object.__setattr__(self, "source_dp_rank", selected) + object.__setattr__(self, "layout", layout) + object.__setattr__( + self, + "part_writers", + tuple(p.participant_id for p in parts), + ) + + +@dataclass(frozen=True) +class KVCacheStoreLoadRange: + """One target participant's logical range mapped to relative Store bytes.""" + + participant_id: ParticipantId + fragment: KVCacheStoreFragment + + def __post_init__(self) -> None: + require_nonempty_string(self.participant_id, "participant_id") + if not isinstance(self.fragment, KVCacheStoreFragment): + raise ValueError("fragment must be a KVCacheStoreFragment") # noqa: TRY004 + + +@dataclass(frozen=True, init=False) +class KVCacheStoreLoadPlan: + """Derived read geometry for one complete target DP replica; no page or address.""" + + target_placement: KVCachePlacementManifest + layout: KVCacheStoreLayout + target_dp_rank: int + ranges: tuple[KVCacheStoreLoadRange, ...] + + def __init__( + self, + layout: KVCacheStoreLayout, + target_placement: KVCachePlacementManifest, + *, + target_dp_rank: int | None = None, + limits: KVCacheStorePlanningLimits = DEFAULT_STORE_PLANNING_LIMITS, + ) -> None: + if not isinstance(layout, KVCacheStoreLayout): + raise ValueError("layout must be a KVCacheStoreLayout") # noqa: TRY004 + # Source fragments retain their stored strides. Destination strides + # come from the validated runtime binding during range lowering. + selected, parts = _selected_parts( + target_placement, target_dp_rank, replicas=True + ) + if target_placement.descriptor != layout.descriptor: + raise ValueError( + "runtime and Store descriptors differ; conversion is unsupported" + ) + _check_limits(layout, limits) + ranges = _load_ranges(parts, layout, limits) + object.__setattr__(self, "target_placement", target_placement) + object.__setattr__(self, "layout", layout) + object.__setattr__(self, "target_dp_rank", selected) + object.__setattr__(self, "ranges", ranges) + + @property + def digest(self) -> str: + """Identify a reusable layout-to-target template, not a read operation.""" + return _canonical_digest( + { + "placement_id": self.target_placement.placement_id, + "placement_digest": self.target_placement.digest, + "layout_id": self.layout.layout_id, + "target_dp_rank": self.target_dp_rank, + } + ) + + +def _load_ranges( + parts: tuple[KVCachePlacementPart, ...], + layout: KVCacheStoreLayout, + limits: KVCacheStorePlanningLimits, +) -> tuple[KVCacheStoreLoadRange, ...]: + """Join sorted head intervals per layer/component, without a Cartesian scan.""" + if sum(len(p.layer_ids) * 2 for p in parts) > limits.max_fragments: + raise ValueError("runtime fragment limit exceeded") + inventory: dict[tuple[int, KVCacheComponent], list[KVCacheStoreFragment]] = {} + for fragment in layout.fragments: + inventory.setdefault((fragment.global_layer_id, fragment.component), []).append( + fragment + ) + indexes = { + key: HeadIntervalIndex( + (fragment.head_start, fragment.head_start + fragment.head_count, fragment) + for fragment in fragments + ) + for key, fragments in inventory.items() + } + result: list[KVCacheStoreLoadRange] = [] + useful_bytes = 0 + for part in parts: + end = part.head_start + part.head_count + for layer in part.layer_ids: + for component in KVCacheComponent: + key = (layer, component) + for start, stop, candidates in indexes[key].cover(part.head_start, end): + (fragment,) = candidates + if len(result) >= limits.max_fragments: + raise ValueError("transfer range limit exceeded") + dim = ( + layout.descriptor.key_head_dim + if component is KVCacheComponent.KEY + else layout.descriptor.value_head_dim + ) + useful_bytes += ( + (stop - start) + * layout.token_count + * dim + * layout.descriptor.itemsize + ) + if useful_bytes > limits.max_bytes: + raise ValueError("transfer byte limit exceeded") + result.append( + KVCacheStoreLoadRange( + part.participant_id, + replace( + fragment, + head_start=start, + head_count=stop - start, + part_offset=fragment.part_offset + + (start - fragment.head_start) + * fragment.head_stride_bytes, + ), + ) + ) + return tuple(result) + + +def plan_kv_cache_store_upload( + source_placement: KVCachePlacementManifest, + *, + object_format: KVCacheStoreFormat, + source_dp_rank: int | None = None, + limits: KVCacheStorePlanningLimits = DEFAULT_STORE_PLANNING_LIMITS, +) -> KVCacheStoreUploadPlan: + """Build this server's own layout and complete object writers, without Store I/O.""" + return KVCacheStoreUploadPlan( + source_placement, + object_format=object_format, + source_dp_rank=source_dp_rank, + limits=limits, + ) + + +def plan_kv_cache_store_load( + layout: KVCacheStoreLayout, + target_placement: KVCachePlacementManifest, + *, + target_dp_rank: int | None = None, + limits: KVCacheStorePlanningLimits = DEFAULT_STORE_PLANNING_LIMITS, +) -> KVCacheStoreLoadPlan: + """Plan reads from one selected source layout into the target placement.""" + return KVCacheStoreLoadPlan( + layout, + target_placement, + target_dp_rank=target_dp_rank, + limits=limits, + ) + + +__all__ = [ + "KVCacheStoreLoadPlan", + "KVCacheStoreLoadRange", + "KVCacheStorePlanningLimits", + "KVCacheStoreUploadPlan", + "plan_kv_cache_store_load", + "plan_kv_cache_store_upload", +] diff --git a/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/runtime.py b/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/runtime.py new file mode 100644 index 0000000000..743500e8a1 --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/runtime.py @@ -0,0 +1,123 @@ +"""Lower fixed-page Store geometry into checked local memory byte ranges.""" + +from __future__ import annotations + +from collections.abc import Iterator, Sequence +from dataclasses import dataclass, replace + +from ..placement import KVCachePlacementManifest +from ..resolved import ( + KVCacheTransferLimits, + _head_bytes, +) +from ._utils import KVCacheRangeGeometry, index_runtime_ranges, plan_range_copy +from .binding import KVCacheStoreRuntimeBinding, validate_store_binding +from .manifest import KVCacheStoreFragment + + +@dataclass(frozen=True) +class StoreByteRange: + page_index: int + part_index: int + part_offset: int + region_id: str + region_offset: int + nbytes: int + + +def split_store_range_batches( + records: Sequence[StoreByteRange], limits: KVCacheTransferLimits +) -> Iterator[list[tuple[StoreByteRange, int, int]]]: + """Split ranges and fill byte/count budgets for direct and prepared reads.""" + batch: list[tuple[StoreByteRange, int, int]] = [] + used = 0 + for record in records: + offset = 0 + while offset < record.nbytes: + if ( + len(batch) == limits.max_batch_operations + or used == limits.max_batch_bytes + ): + yield batch + batch, used = [], 0 + size = min(record.nbytes - offset, limits.max_batch_bytes - used) + batch.append((record, offset, size)) + used += size + offset += size + if batch: + yield batch + + +def lower_store_ranges( + placement: KVCachePlacementManifest, + binding: KVCacheStoreRuntimeBinding, + fragments: tuple[tuple[KVCacheStoreFragment, ...], ...], + operation_id: str, + limits: KVCacheTransferLimits, +) -> tuple[StoreByteRange, ...]: + """Validate all runtime pages, then join each page's selected Store layout.""" + page_size = placement.descriptor.page_size + validate_store_binding( + placement, binding, operation_id, 0, len(fragments) * page_size, limits + ) + by_layer = index_runtime_ranges(binding.ranges) + records: list[StoreByteRange] = [] + useful_bytes = 0 + for page_index, page_fragments in enumerate(fragments): + page_start = page_index * page_size + for fragment in page_fragments: + source = KVCacheRangeGeometry( + page_start, + page_size, + fragment.head_start, + fragment.head_count, + fragment.part_offset, + fragment.token_stride_bytes, + fragment.head_stride_bytes, + ) + head_bytes = _head_bytes(placement, fragment.component) + for target in by_layer.get( + (fragment.global_layer_id, fragment.component), [] + ): + copy = plan_range_copy(source, target, head_bytes) + if copy is None: + continue + if len(records) + copy.segment_count > limits.max_operations: + raise ValueError("Store physical operation limit exceeded") + useful_bytes += copy.nbytes + if useful_bytes > limits.max_bytes: + raise ValueError("Store transfer byte limit exceeded") + for source_offset, target_offset, nbytes in copy.segments(): + records.append( + StoreByteRange( + page_index, + fragment.part_index, + source_offset, + target.region_id, + target_offset, + nbytes, + ) + ) + records.sort( + key=lambda r: ( + r.page_index, + r.part_index, + r.part_offset, + r.region_id, + r.region_offset, + ) + ) + merged: list[StoreByteRange] = [] + for record in records: + if merged: + last = merged[-1] + if ( + (last.page_index, last.part_index, last.region_id) + == (record.page_index, record.part_index, record.region_id) + and last.part_offset + last.nbytes == record.part_offset + and last.region_offset + last.nbytes == record.region_offset + ): + merged[-1] = replace(last, nbytes=last.nbytes + record.nbytes) + continue + merged.append(record) + return tuple(merged) diff --git a/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/serde.py b/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/serde.py new file mode 100644 index 0000000000..a0f7ac3299 --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/kv_cache/_store/serde.py @@ -0,0 +1,162 @@ +"""Strict JSON for shared Store records; no per-page snapshot serialization.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from dataclasses import asdict +from typing import cast + +from ..serde import JsonValue, _descriptor_from_wire +from ._utils import canonical_json, load_json_object +from ._utils import json_integer as _integer +from ._utils import json_string as _string +from .manifest import ( + KVCacheStoreFormat, + KVCacheStoreLayout, + KVCacheStoreManifest, + KVCacheStoreShard, +) + + +def _fields(value: object, names: set[str], label: str) -> Mapping[str, object]: + if ( + not isinstance(value, Mapping) + or set(cast(Mapping[object, object], value)) != names + ): + raise ValueError(f"{label} fields do not match contract") + return cast(Mapping[str, object], value) + + +def _items(value: object, label: str) -> list[object]: + if type(value) is not list or len(cast(list[object], value)) > 100_000: + raise ValueError(f"{label} must be an array within the collection limit") + return cast(list[object], value) + + +def _load(value: str, fields: set[str], label: str) -> Mapping[str, object]: + try: + payload = _fields(load_json_object(value, label), fields, label) + except RecursionError as error: + raise ValueError("KV Store JSON nesting limit exceeded") from error + return payload + + +def _dump(value: object) -> str: + encoded = canonical_json(value) + if len(encoded.encode("utf-8")) > 16 * 1024 * 1024: + raise ValueError("KV Store JSON exceeds the 16 MiB wire limit") + return encoded + + +def kv_cache_store_layout_to_json(layout: KVCacheStoreLayout) -> str: + """Encode geometry and its canonical identity, with no derived byte fields.""" + if not isinstance(layout, KVCacheStoreLayout): + raise ValueError("layout must be a KVCacheStoreLayout") # noqa: TRY004 + return _dump( + { + **asdict(layout), + "layout_id": layout.layout_id, + "layout_digest": layout.digest, + } + ) + + +def kv_cache_store_layout_from_json(value: str) -> KVCacheStoreLayout: + """Validate exact fields, full shard coverage and canonical digest.""" + payload = _load( + value, + { + "descriptor", + "object_format", + "shards", + "layout_id", + "layout_digest", + }, + "KV Store layout", + ) + shards: list[KVCacheStoreShard] = [] + layer_count = 0 + for item in _items(payload["shards"], "Store shards"): + raw = _fields( + item, {"layer_ids", "head_start", "head_count", "key_suffix"}, "Store shard" + ) + layer_count += len(_items(raw["layer_ids"], "layers")) + if layer_count * 2 > 100_000: + raise ValueError("Store fragment count limit exceeded") + shards.append( + KVCacheStoreShard( + tuple( + _integer(layer, "layer_id") + for layer in _items(raw["layer_ids"], "layers") + ), + _integer(raw["head_start"], "head_start"), + _integer(raw["head_count"], "head_count", minimum=1), + _string(raw["key_suffix"], "key_suffix"), + ) + ) + layout = KVCacheStoreLayout( + _descriptor_from_wire(cast(JsonValue, payload["descriptor"])), + KVCacheStoreFormat(_string(payload["object_format"], "object_format")), + tuple(shards), + ) + if ( + _string(payload["layout_id"], "layout_id"), + _string(payload["layout_digest"], "layout_digest"), + ) != (layout.layout_id, layout.digest): + raise ValueError("stored layout identity does not match content") + return layout + + +def kv_cache_store_manifest_to_json(manifest: KVCacheStoreManifest) -> str: + """Encode one shared model/layout record, not a page or upload attempt.""" + if not isinstance(manifest, KVCacheStoreManifest): + raise ValueError("manifest must be a KVCacheStoreManifest") # noqa: TRY004 + return _dump( + { + "namespace": manifest.namespace, + "model_id": manifest.model_id, + "model_revision": manifest.model_revision, + "semantic_fingerprint": manifest.semantic_fingerprint, + "layout": json.loads(kv_cache_store_layout_to_json(manifest.layout)), + "model_domain": manifest.model_domain, + "manifest_digest": manifest.digest, + } + ) + + +def kv_cache_store_manifest_from_json(value: str) -> KVCacheStoreManifest: + """Decode a self-contained shared manifest; never query Store or page data.""" + payload = _load( + value, + { + "namespace", + "model_id", + "model_revision", + "semantic_fingerprint", + "layout", + "model_domain", + "manifest_digest", + }, + "KV Store manifest", + ) + manifest = KVCacheStoreManifest( + _string(payload["namespace"], "namespace"), + _string(payload["model_id"], "model_id"), + _string(payload["model_revision"], "model_revision"), + _string(payload["semantic_fingerprint"], "semantic_fingerprint"), + kv_cache_store_layout_from_json(_dump(payload["layout"])), + ) + if manifest.model_domain != _string(payload["model_domain"], "model_domain"): + raise ValueError("model domain does not match content") + if manifest.digest != _string(payload["manifest_digest"], "manifest_digest"): + raise ValueError("stored manifest digest does not match content") + return manifest + + +__all__ = [ + "kv_cache_store_layout_from_json", + "kv_cache_store_layout_to_json", + "kv_cache_store_manifest_from_json", + "kv_cache_store_manifest_to_json", +] diff --git a/mooncake-reshard/python/mooncake/reshard/kv_cache/identity.py b/mooncake-reshard/python/mooncake/reshard/kv_cache/identity.py new file mode 100644 index 0000000000..2448d2f915 --- /dev/null +++ b/mooncake-reshard/python/mooncake/reshard/kv_cache/identity.py @@ -0,0 +1,62 @@ +"""Topology-independent semantic identity and logical keys for cached KV.""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import asdict, dataclass +from functools import cached_property + +from .types import KVCacheDescriptor, require_nonempty_string + + +@dataclass(frozen=True) +class KVCacheModelDomain: + """Identify compatible KV content without naming its physical owners. + + The framework supplies the weight revision and a fingerprint of settings + that affect KV values, such as RoPE configuration. Geometry alone cannot + establish model compatibility. TP/PP ranks, buffer addresses and host + axis order do not participate in this identity. + """ + + namespace: str + model_id: str + model_revision: str + semantic_fingerprint: str + descriptor: KVCacheDescriptor + + def __post_init__(self) -> None: + for name in ("namespace", "model_id", "model_revision", "semantic_fingerprint"): + require_nonempty_string(getattr(self, name), name) + if not isinstance(self.descriptor, KVCacheDescriptor): + raise ValueError("descriptor must be a KVCacheDescriptor") # noqa: TRY004 + + @cached_property + def digest(self) -> str: + """Stable semantic-domain digest, compatible with existing cache keys.""" + content = { + "namespace": self.namespace, + "model_id": self.model_id, + "model_revision": self.model_revision, + "semantic_fingerprint": self.semantic_fingerprint, + "descriptor": asdict(self.descriptor), + # Preserve the existing domain encoding. This constant is a + # compatibility marker, never the active host-memory axis order. + "object_format": "PLHD", + } + encoded = json.dumps(content, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest() + + def logical_page_key(self, prefix: str, token_prefix_hash: str) -> str: + """Name one logical page; storage-specific protocol suffixes are external. + + The caller supplies the hash of the complete token prefix, including + its position/context semantics, rather than only the current page. + """ + require_nonempty_string(prefix, "prefix") + require_nonempty_string(token_prefix_hash, "token_prefix_hash") + return f"{prefix}_{token_prefix_hash}_m{self.digest}" + + +__all__ = ["KVCacheModelDomain"] diff --git a/mooncake-reshard/python/mooncake/reshard/kv_cache/resolved.py b/mooncake-reshard/python/mooncake/reshard/kv_cache/resolved.py index ad264dd158..73dd1396da 100644 --- a/mooncake-reshard/python/mooncake/reshard/kv_cache/resolved.py +++ b/mooncake-reshard/python/mooncake/reshard/kv_cache/resolved.py @@ -259,22 +259,39 @@ def _validated_spans( raise ValueError(f"resolved binding {name} differs") if snapshot.resource_id != placement.resource_id: raise ValueError("snapshot resource_id differs from placement") - if len(binding.ranges) > limits.max_ranges: + return _validated_range_spans( + placement, + binding.participant_id, + binding.regions, + binding.ranges, + token_start=snapshot.token_start, + token_end=snapshot.token_end, + limits=limits, + ) + + +def _validated_range_spans( + placement: KVCachePlacementManifest, + participant_id: str, + registered_regions: tuple[KVCacheRegisteredRegion, ...], + ranges: tuple[KVCacheResolvedRange, ...], + *, + token_start: int, + token_end: int, + limits: KVCacheTransferLimits, +) -> list[tuple[int, int]]: + """Validate memory geometry after the caller checks its binding identity.""" + if len(ranges) > limits.max_ranges: raise ValueError("resolved range count limit exceeded") - part = placement.part(ParticipantId(binding.participant_id)) - regions = {r.region_id: r for r in binding.regions} + part = placement.part(ParticipantId(participant_id)) + regions = {r.region_id: r for r in registered_regions} groups: dict[tuple[int, KVCacheComponent], list[KVCacheResolvedRange]] = {} spans: list[tuple[int, int]] = [] - for item in binding.ranges: + for item in ranges: if item.global_layer_id not in part.layer_ids: raise ValueError("resolved range contains an unowned layer") - if ( - not snapshot.token_start - <= item.token_start - < item.token_end - <= snapshot.token_end - ): - raise ValueError("resolved range is outside snapshot token interval") + if not token_start <= item.token_start < item.token_end <= token_end: + raise ValueError("resolved range is outside token interval") if ( not part.head_start <= item.head_start @@ -323,7 +340,7 @@ def _validated_spans( work = 0 for items in groups.values(): boundaries = sorted( - {snapshot.token_start, snapshot.token_end} + {token_start, token_end} | {t for r in items for t in (r.token_start, r.token_end)} ) work += len(boundaries) * len(items) diff --git a/mooncake-reshard/tests/kv_cache/_store_test_utils.py b/mooncake-reshard/tests/kv_cache/_store_test_utils.py new file mode 100644 index 0000000000..5d13ab62fd --- /dev/null +++ b/mooncake-reshard/tests/kv_cache/_store_test_utils.py @@ -0,0 +1,161 @@ +"""Deterministic KV tensors and a minimal byte backend for reshard tests.""" + +import ctypes +import struct + +from mooncake.reshard.kv_cache import ( + KVCacheComponent, + KVCacheRegisteredRegion, + KVCacheResolvedRange, + KVCacheStoreFormat, + KVCacheStoreManifest, + KVCacheStoreRuntimeBinding, + plan_kv_cache_store_upload, +) +from test_kv_cache_reshard import _placement + +FORMATS = tuple(KVCacheStoreFormat) + + +def _layout(tp=1, layers=((0, 1), (2, 3)), fmt=FORMATS[0]): + return plan_kv_cache_store_upload( + _placement("first", layers, tp), object_format=fmt + ).layout + + +def _manifest(layout=None): + return KVCacheStoreManifest( + "cache", "model", "revision", "semantic", layout or _layout() + ) + + +def _offset(fmt, token, layer, head, layers, heads, tokens, head_bytes): + # Independent row-major tensor indexing; do not use production strides. + if fmt is KVCacheStoreFormat.PLHD: + return ((token * layers + layer) * heads + head) * head_bytes + if fmt is KVCacheStoreFormat.LPHD: + return ((layer * tokens + token) * heads + head) * head_bytes + return ((head * tokens + token) * layers + layer) * head_bytes + + +def _reference(layout, layer_ids, head_start, head_count, component): + d = layout.descriptor + dim = d.key_head_dim if component is KVCacheComponent.KEY else d.value_head_dim + result = bytearray(len(layer_ids) * d.page_size * head_count * dim * d.itemsize) + for t in range(d.page_size): + for local_layer, layer in enumerate(layer_ids): + for h in range(head_count): + off = _offset( + layout.object_format, + t, + local_layer, + h, + len(layer_ids), + head_count, + d.page_size, + dim * d.itemsize, + ) + for x in range(dim): + value = layer * 1000 + t * 50 + (head_start + h) * 8 + x + if component is KVCacheComponent.VALUE: + value += 20000 + struct.pack_into("