diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index ea2c05a13..bd2f170b0 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -49,10 +49,7 @@ from sglang.srt.mem_cache.cp_shared_l2_pool import ( PAYLOAD_TARGET_KV, CpSharedL2NodeMetadata, CpSharedL2PageAllocator, - broadcast_cp_shared_l2_decision, cp_shared_l2_logical_token_indices, - gather_cp_shared_l2_commits, - gather_cp_shared_l2_preflight, ) from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool, NSATokenToKVPool from sglang.srt.mem_cache.page_index_utils import validate_page_aligned_token_indices @@ -393,11 +390,6 @@ class HiCacheController: draft_mem_pool_host: Optional["HostKVCache"] = None, draft_mem_pool_device=None, cp_shared_l2_page_allocator: Optional[CpSharedL2PageAllocator] = None, - cp_shared_l2_cpu_group=None, - cp_shared_l2_source_rank: int = 0, - cp_shared_l2_broadcast_fn=None, - cp_shared_l2_preflight_fn=None, - cp_shared_l2_commit_fn=None, ): self.tp_group = tp_group self.mem_pool_device_allocator = token_to_kv_pool_allocator @@ -409,11 +401,6 @@ class HiCacheController: self.mem_pool_device = mem_pool_device self.cp_shared_kv_layout = cp_shared_kv_layout self.cp_shared_l2_page_allocator = cp_shared_l2_page_allocator - self.cp_shared_l2_cpu_group = cp_shared_l2_cpu_group - self.cp_shared_l2_source_rank = int(cp_shared_l2_source_rank) - self.cp_shared_l2_broadcast_fn = cp_shared_l2_broadcast_fn - self.cp_shared_l2_preflight_fn = cp_shared_l2_preflight_fn - self.cp_shared_l2_commit_fn = cp_shared_l2_commit_fn self._cp_shared_l2_stats = { "cp_shared_l2_load_hits": 0, "cp_shared_l2_load_misses": 0, @@ -1323,46 +1310,6 @@ class HiCacheController: start_layer = int(getattr(self.mem_pool_host, "start_layer", 0) or 0) return tuple(range(start_layer, start_layer + index_layer_num)) - def _cp_shared_l2_expected_layers_by_payload( - self, required_payloads: List[str] - ) -> dict[str, object]: - expected_layers_by_payload = {} - if PAYLOAD_TARGET_KV in required_payloads: - expected_layers_by_payload[PAYLOAD_TARGET_KV] = range(self.layer_num) - if PAYLOAD_DRAFT_KV in required_payloads: - if self.draft_mem_pool_device is None: - raise RuntimeError( - "CP shared physical L2 draft payload requires draft KV device pool" - ) - draft_layer_num = int(getattr(self.draft_mem_pool_device, "layer_num", 0)) - if draft_layer_num <= 0: - raise RuntimeError( - "CP shared physical L2 draft payload requires positive draft layer count" - ) - expected_layers_by_payload[PAYLOAD_DRAFT_KV] = range(draft_layer_num) - if PAYLOAD_INDEX_K in required_payloads: - index_layer_ids = self._cp_shared_l2_index_layer_ids() - if not index_layer_ids: - raise RuntimeError( - "CP shared physical L2 index_k payload requires active NSA index layers" - ) - expected_layers_by_payload[PAYLOAD_INDEX_K] = index_layer_ids - return expected_layers_by_payload - - def _set_cp_shared_l2_object_commit_contract( - self, allocator, object_key: str, required_payloads: List[str] - ) -> None: - set_required = getattr(allocator, "set_object_required_payloads", None) - if set_required is None: - return - set_required( - object_key, - required_payloads, - expected_layers_by_payload=self._cp_shared_l2_expected_layers_by_payload( - required_payloads - ), - ) - def _cp_shared_l2_index_payload_active(self) -> bool: return bool(self._cp_shared_l2_index_layer_ids()) @@ -1428,13 +1375,9 @@ class HiCacheController: required_host_slots=padded_len, reason="shared_l2_capacity", ) - # MF3: reserve() only sets the allocator-wide default payload set; pin - # the actual per-object contract (required payloads + expected layers) - # so split_committed_object reads correct state. - self._set_cp_shared_l2_object_commit_contract( - allocator, object_key, required_payloads - ) - + # B1 (Tier-2): no per-object commit contract -- commit consensus is the + # writing_check MIN frontier (mark_object_committed), and the allocator's + # entire state is the placement_digest-hashed {ranges, committed, free}. object_range = object_ranges[PAYLOAD_TARGET_KV] owned_mask = layout.owned_by_this_rank(padded_device_indices) owned_positions = owned_mask.nonzero(as_tuple=True)[0].cpu() @@ -1490,34 +1433,6 @@ class HiCacheController: allocator.abort(object_key) raise - def _commit_cp_shared_l2_all_required_layers( - self, reservation: HiCacheWriteReservation - ) -> None: - metadata = reservation.metadata - required_payloads = tuple( - getattr(metadata, "required_payloads", (PAYLOAD_TARGET_KV,)) - ) - for layer_id in range(self.layer_num): - if PAYLOAD_TARGET_KV in required_payloads: - self._commit_cp_shared_l2_layer_if_needed( - reservation, PAYLOAD_TARGET_KV, layer_id - ) - if PAYLOAD_DRAFT_KV in required_payloads: - if self.draft_mem_pool_device is None: - raise RuntimeError( - "CP shared physical L2 draft commit requires draft KV device pool" - ) - draft_layer_num = int(getattr(self.draft_mem_pool_device, "layer_num", 0)) - for layer_id in range(draft_layer_num): - self._commit_cp_shared_l2_layer_if_needed( - reservation, PAYLOAD_DRAFT_KV, layer_id - ) - if PAYLOAD_INDEX_K in required_payloads: - for layer_id in self._cp_shared_l2_index_layer_ids(): - self._commit_cp_shared_l2_layer_if_needed( - reservation, PAYLOAD_INDEX_K, layer_id - ) - def submit_write_cp_all_layer(self, reservation: HiCacheWriteReservation) -> None: logger.warning( "[CP_HICACHE_FALLBACK][submit_write_cp_all_layer] " @@ -1726,75 +1641,6 @@ class HiCacheController: reservation.draft_host_indices is not None, ) - def _commit_cp_shared_l2_layer_if_needed( - self, reservation: HiCacheWriteReservation, payload: str, layer_id: int - ) -> bool: - return self._commit_cp_shared_l2_layers_batched( - [reservation], payload, layer_id - ) - - def _commit_cp_shared_l2_layers_batched( - self, - reservations: List[HiCacheWriteReservation], - payload: str, - layer_id: int, - ) -> bool: - allocator = self.cp_shared_l2_page_allocator - if allocator is None: - return False - reservations_by_object_key: Dict[str, HiCacheWriteReservation] = {} - local_commits = [] - for reservation in reservations: - metadata = reservation.metadata - object_key = getattr(metadata, "object_key", None) - if object_key is None: - continue - reservations_by_object_key[object_key] = reservation - local_commits.append( - ( - object_key, - payload, - int(layer_id), - int(self.cp_shared_kv_layout.cp_rank), - ) - ) - if not local_commits: - return False - local_commits_tuple = tuple(local_commits) - if self.cp_shared_l2_commit_fn is not None: - gathered_commits = self.cp_shared_l2_commit_fn( - local_commits_tuple, - cp_cpu_group=self.cp_shared_l2_cpu_group, - rank=int(self.cp_shared_kv_layout.cp_rank), - ) - else: - gathered_commits = gather_cp_shared_l2_commits( - local_commits_tuple, - cp_cpu_group=self.cp_shared_l2_cpu_group, - ) - committed = False - for commit_object_key, commit_payload, commit_layer_id, commit_rank in ( - gathered_commits - ): - if commit_payload != payload: - continue - if allocator.get_range(commit_object_key, commit_payload) is None: - continue - committed = allocator.commit_layer( - commit_object_key, - commit_payload, - int(commit_layer_id), - int(commit_rank), - ) or committed - for object_key, reservation in reservations_by_object_key.items(): - metadata = reservation.metadata - getattr(metadata, "committed_payload_layers", {}).setdefault( - payload, set() - ).add(layer_id) - if allocator.is_committed(object_key): - metadata.committed = True - return committed - def _abort_cp_shared_l2_states( self, states: List[HiCacheLayerWriteState], *, fence_write_stream: bool = False ) -> None: diff --git a/python/sglang/srt/mem_cache/cp_shared_l2_pool.py b/python/sglang/srt/mem_cache/cp_shared_l2_pool.py index 592907990..236965507 100644 --- a/python/sglang/srt/mem_cache/cp_shared_l2_pool.py +++ b/python/sglang/srt/mem_cache/cp_shared_l2_pool.py @@ -584,9 +584,6 @@ class CpSharedL2PageAllocator: str, Iterable[CpSharedL2SlabInfo | Mapping[str, Any]] ] | None = None, - expected_ranks: Iterable[int] = (0,), - expected_layers: Iterable[int] = (0,), - required_payloads: Iterable[str] = (PAYLOAD_TARGET_KV,), ): pages_by_payload: dict[str, int] | None if pages_per_payload is None: @@ -617,33 +614,16 @@ class CpSharedL2PageAllocator: self._validate_pages_match_slab_capacity(pages_by_payload, slabs) self._install_slabs(slabs) - self._expected_ranks = frozenset(int(rank) for rank in expected_ranks) - self._expected_layers = frozenset(int(layer) for layer in expected_layers) - self._required_payloads = frozenset(required_payloads) - if not self._expected_ranks: - raise ValueError("expected_ranks must not be empty") - if not self._expected_layers: - raise ValueError("expected_layers must not be empty") - if not self._required_payloads: - raise ValueError("required_payloads must not be empty") - if any(rank < 0 for rank in self._expected_ranks): - raise ValueError("expected_ranks must be non-negative") - if any(layer < 0 for layer in self._expected_layers): - raise ValueError("expected_layers must be non-negative") - for payload_kind in self._required_payloads: - self._validate_payload_kind(payload_kind) - if payload_kind not in self._free_by_payload: - raise ValueError( - f"required payload {payload_kind!r} has no slab namespace" - ) - + # B1 commit model (Tier-2 collapse, design sec 9.12): the allocator carries + # ONLY the durable placement state the placement_digest() hashes -- the + # per-object ranges and the committed set (free list lives in + # _free_by_payload, installed above). The option-A per-(rank,layer,payload) + # commit quorum (commit_layer/_has_full_commit/_commits_by_object/ + # _expected_ranks) and the per-object contract are GONE: commit consensus is + # the writing_check ReduceOp.MIN frontier driving mark_object_committed. self._generation = 0 self._ranges_by_object: dict[str, dict[str, CpSharedL2ObjectRange]] = {} - self._adopted_ranges: set[tuple[str, str]] = set() - self._commits_by_object: dict[str, dict[str, dict[int, set[int]]]] = {} self._committed_objects: set[str] = set() - self._required_payloads_by_object: dict[str, frozenset[str]] = {} - self._expected_layers_by_object: dict[str, dict[str, frozenset[int]]] = {} self._stats = { "cp_shared_l2_objects_committed": 0, "cp_shared_l2_objects_aborted": 0, @@ -676,122 +656,9 @@ class CpSharedL2PageAllocator: generation=self._generation, ) self._ranges_by_object.setdefault(object_key, {})[payload_kind] = object_range - self._adopted_ranges.discard((object_key, payload_kind)) - self._commits_by_object.setdefault(object_key, {}) - self._required_payloads_by_object.setdefault( - object_key, frozenset(self._required_payloads) - ) self._committed_objects.discard(object_key) return object_range - def adopt_reserved_range(self, object_range: CpSharedL2ObjectRange) -> None: - """Register a broadcast reservation on non-source ranks without charging free pages. - - Rank 0 owns allocation from the free list. Other CP ranks receive the - same object range by broadcast and need a local commit table entry so - commit_layer can validate expected ranks/layers without performing a - second allocation. Calling this on rank 0 after reserve is idempotent. - """ - - self._validate_object_key(object_range.object_key) - self._validate_payload_kind(object_range.payload_kind) - self._validate_range_within_known_slab(object_range) - existing = self._ranges_by_object.setdefault(object_range.object_key, {}).get( - object_range.payload_kind - ) - if existing is not None: - if existing != object_range: - raise ValueError( - f"conflicting shared L2 range for object {object_range.object_key!r} " - f"payload {object_range.payload_kind!r}" - ) - return - self._ranges_by_object[object_range.object_key][ - object_range.payload_kind - ] = object_range - self._adopted_ranges.add((object_range.object_key, object_range.payload_kind)) - self._commits_by_object.setdefault(object_range.object_key, {}) - self._required_payloads_by_object.setdefault( - object_range.object_key, frozenset(self._required_payloads) - ) - self._committed_objects.discard(object_range.object_key) - - def set_object_required_payloads( - self, - object_key: str, - payloads: Iterable[str], - *, - expected_layers_by_payload: Mapping[str, Iterable[int]] | None = None, - ) -> None: - self._validate_object_key(object_key) - required = frozenset(payloads) - if not required: - raise ValueError("required payloads must not be empty") - for payload_kind in required: - self._validate_payload_kind(payload_kind) - if payload_kind not in self._free_by_payload: - raise ValueError( - f"required payload {payload_kind!r} has no slab namespace" - ) - self._required_payloads_by_object[object_key] = required - if expected_layers_by_payload is None: - self._expected_layers_by_object.pop(object_key, None) - return - layer_map: dict[str, frozenset[int]] = {} - for payload_kind, layers in expected_layers_by_payload.items(): - self._validate_payload_kind(payload_kind) - if payload_kind not in required: - raise ValueError( - f"expected layers provided for non-required payload {payload_kind!r}" - ) - expected_layers = frozenset(int(layer) for layer in layers) - if not expected_layers: - raise ValueError( - f"expected layers for payload {payload_kind!r} must not be empty" - ) - if any(layer < 0 for layer in expected_layers): - raise ValueError("expected layer ids must be non-negative") - layer_map[payload_kind] = expected_layers - for payload_kind in required: - layer_map.setdefault(payload_kind, self._expected_layers) - self._expected_layers_by_object[object_key] = layer_map - - def commit_layer( - self, object_key: str, payload: str, layer_id: int, rank: int - ) -> bool: - self._validate_object_key(object_key) - self._validate_payload_kind(payload) - layer_id = int(layer_id) - rank = int(rank) - if object_key not in self._ranges_by_object: - raise ValueError(f"unknown object {object_key!r}") - if payload not in self._ranges_by_object[object_key]: - raise ValueError( - f"object {object_key!r} has no reservation for payload {payload!r}" - ) - required_payloads = self._required_payloads_by_object.get( - object_key, self._required_payloads - ) - if payload not in required_payloads: - raise ValueError(f"payload {payload!r} is not required for commit") - expected_layers = self._expected_layers_for_object_payload(object_key, payload) - if layer_id not in expected_layers: - raise ValueError(f"unexpected layer_id {layer_id}") - if rank not in self._expected_ranks: - raise ValueError(f"unexpected rank {rank}") - - payload_commits = self._commits_by_object.setdefault(object_key, {}).setdefault( - payload, {} - ) - payload_commits.setdefault(layer_id, set()).add(rank) - if self._has_full_commit(object_key): - was_committed = object_key in self._committed_objects - self._committed_objects.add(object_key) - if not was_committed: - self._stats["cp_shared_l2_objects_committed"] += 1 - return True - return False - def abort(self, object_key: str) -> bool: dropped = self._drop_object(object_key) if dropped: @@ -805,16 +672,14 @@ class CpSharedL2PageAllocator: return dropped def mark_object_committed(self, object_key: str) -> bool: - """B1 commit: mark an object committed directly, bypassing the per-(rank, - layer, payload) commit_layer quorum. + """B1 commit: mark an object committed directly. Under the B1 collective-free allocator the all-ranks-done consensus is the writing_check ReduceOp.MIN frontier (the same barrier that releases the - node's write lock), NOT a per-layer all-gather. When that frontier commits a - node, every CP rank calls this with the same object_key, so the committed set - transitions rank-uniformly (and feeds placement_digest). Does not populate - _commits_by_object -- the per-layer quorum (_has_full_commit) is unused under - B1; is_committed / release / split_committed_object read _committed_objects, + node's write lock), NOT a per-layer all-gather quorum. When that frontier + commits a node, every CP rank calls this with the same object_key, so the + committed set transitions rank-uniformly (and feeds placement_digest). + is_committed / release / split_committed_object read _committed_objects, which this maintains. Returns True on the first commit, False if already committed. Raises on an unknown object (fail-loud: the node's reservation must outlive its write lock). @@ -861,12 +726,6 @@ class CpSharedL2PageAllocator: if object_key not in self._committed_objects: raise ValueError(f"object {object_key!r} is not committed") - required_payloads = self._required_payloads_by_object.get( - object_key, self._required_payloads - ) - expected_layers = self._expected_layers_by_object.get(object_key) - commits = self._commits_by_object.get(object_key, {}) - parent_ranges: dict[str, CpSharedL2ObjectRange] = {} child_ranges: dict[str, CpSharedL2ObjectRange] = {} for payload_kind, object_range in ranges.items(): @@ -897,18 +756,9 @@ class CpSharedL2PageAllocator: generation=object_range.generation, ) - old_adopted_payloads = { - payload_kind - for payload_kind in ranges - if (object_key, payload_kind) in self._adopted_ranges - } - for payload_kind in old_adopted_payloads: - self._adopted_ranges.discard((object_key, payload_kind)) - + # B1 (Tier-2): split carries ONLY the durable placement (ranges) + the + # committed bit. No commit-quorum / contract / adopted-range tables exist. self._ranges_by_object.pop(object_key, None) - self._commits_by_object.pop(object_key, None) - self._required_payloads_by_object.pop(object_key, None) - self._expected_layers_by_object.pop(object_key, None) self._committed_objects.discard(object_key) for new_key, new_ranges in ( @@ -916,23 +766,7 @@ class CpSharedL2PageAllocator: (child_object_key, child_ranges), ): self._ranges_by_object[new_key] = new_ranges - self._commits_by_object[new_key] = { - payload_kind: { - int(layer_id): set(ranks) - for layer_id, ranks in layer_commits.items() - } - for payload_kind, layer_commits in commits.items() - } - self._required_payloads_by_object[new_key] = frozenset(required_payloads) - if expected_layers is not None: - self._expected_layers_by_object[new_key] = { - payload_kind: frozenset(layers) - for payload_kind, layers in expected_layers.items() - } self._committed_objects.add(new_key) - for payload_kind in old_adopted_payloads: - if payload_kind in new_ranges: - self._adopted_ranges.add((new_key, payload_kind)) return parent_ranges, child_ranges @@ -1025,17 +859,10 @@ class CpSharedL2PageAllocator: def _drop_object(self, object_key: str) -> bool: ranges = self._ranges_by_object.pop(object_key, None) - self._commits_by_object.pop(object_key, None) - self._required_payloads_by_object.pop(object_key, None) - self._expected_layers_by_object.pop(object_key, None) self._committed_objects.discard(object_key) if not ranges: return False for object_range in ranges.values(): - adopted_key = (object_range.object_key, object_range.payload_kind) - if adopted_key in self._adopted_ranges: - self._adopted_ranges.discard(adopted_key) - continue self._return_range(object_range) return True @@ -1061,30 +888,6 @@ class CpSharedL2PageAllocator: merged.append((base, length)) self._free_by_payload[object_range.payload_kind][object_range.slab_id] = merged - def _has_full_commit(self, object_key: str) -> bool: - payload_commits = self._commits_by_object.get(object_key, {}) - object_ranges = self._ranges_by_object.get(object_key, {}) - required_payloads = self._required_payloads_by_object.get( - object_key, self._required_payloads - ) - for payload_kind in required_payloads: - if payload_kind not in object_ranges: - return False - layer_commits = payload_commits.get(payload_kind, {}) - for layer_id in self._expected_layers_for_object_payload( - object_key, payload_kind - ): - if layer_commits.get(layer_id, set()) != self._expected_ranks: - return False - return True - - def _expected_layers_for_object_payload( - self, object_key: str, payload_kind: str - ) -> frozenset[int]: - return self._expected_layers_by_object.get(object_key, {}).get( - payload_kind, self._expected_layers - ) - def _payload_capacity(self, payload_kind: str) -> int: return self._capacity_by_payload[payload_kind] @@ -1348,23 +1151,6 @@ def broadcast_cp_shared_l2_object_range( ) -def broadcast_cp_shared_l2_decision( - decision: Any, - *, - cp_cpu_group: Any, - rank: int, - source_rank: int = 0, - broadcast_fn: Callable[[list[Any], int, Any], None] | None = None, -) -> Any: - return broadcast_cp_shared_l2_object( - decision, - cp_cpu_group=cp_cpu_group, - rank=rank, - source_rank=source_rank, - broadcast_fn=broadcast_fn, - ) - - def gather_cp_shared_l2_object( value: Any, *, @@ -1392,66 +1178,6 @@ def gather_cp_shared_l2_object( return [item for item in output if item is not None] -def gather_cp_shared_l2_preflight( - blocked: bool, - *, - cp_cpu_group: Any, - gather_fn: Callable[[list[Any], Any, Any], None] | None = None, - dist_module: Any | None = None, -) -> list[bool]: - """Collect shared-L2 reserve preflight blocked flags from all CP ranks.""" - - return [ - bool(item) - for item in gather_cp_shared_l2_object( - bool(blocked), - cp_cpu_group=cp_cpu_group, - gather_fn=gather_fn, - dist_module=dist_module, - ) - ] - - -def gather_cp_shared_l2_commits( - commit_infos: tuple[tuple[str, str, int, int], ...], - *, - cp_cpu_group: Any, - gather_fn: Callable[[list[Any], Any, Any], None] | None = None, - dist_module: Any | None = None, -) -> list[tuple[str, str, int, int]]: - """Collect per-rank batches of shared-L2 layer commit facts.""" - - gathered_batches = gather_cp_shared_l2_object( - tuple(commit_infos), - cp_cpu_group=cp_cpu_group, - gather_fn=gather_fn, - dist_module=dist_module, - ) - commits: list[tuple[str, str, int, int]] = [] - for batch in gathered_batches: - for item in batch: - commits.append(tuple(item)) - return commits - - -def gather_cp_shared_l2_commit( - commit_info: tuple[str, str, int, int], - *, - cp_cpu_group: Any, - rank: int, - gather_fn: Callable[[list[Any], Any, Any], None] | None = None, - dist_module: Any | None = None, -) -> list[tuple[str, str, int, int]]: - """Collect one shared-L2 layer commit fact per rank.""" - - return gather_cp_shared_l2_commits( - (commit_info,), - cp_cpu_group=cp_cpu_group, - gather_fn=gather_fn, - dist_module=dist_module, - ) - - def _failfast(message: str): raise ValueError(f"{CP_SHARED_L2_FAILFAST_PREFIX} {message}") diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index d9732e64d..60dcd7b7e 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1039,7 +1039,6 @@ class HiRadixCache(RadixCache): self.load_cache_event = threading.Event() cp_shared_kv_layout = None cp_shared_l2_page_allocator = None - cp_shared_l2_cpu_group = None if self._uses_cp_hicache: cp_shared_kv_layout = CpSharedKVLayout( page_size=self.page_size, @@ -1048,13 +1047,9 @@ class HiRadixCache(RadixCache): ) if getattr(server_args, "enable_cp_shared_physical_l2_hicache", False): from sglang.srt.mem_cache.cp_shared_l2_pool import ( - PAYLOAD_TARGET_KV, CpSharedL2PageAllocator, ) - _cp_rank, cp_size, cp_shared_l2_cpu_group = _get_cp_shared_l2_rank_and_group( - params.token_to_kv_pool_allocator - ) if cp_shared_l2_pages_per_payload_planned is not None: pages_per_payload = dict(cp_shared_l2_pages_per_payload_planned) else: @@ -1064,12 +1059,12 @@ class HiRadixCache(RadixCache): draft_token_to_kv_pool=draft_token_to_kv_pool, draft_host_token_capacity=draft_host_token_capacity, ) + # B1 (Tier-2): the allocator carries only placement state; the + # option-A ctor args (expected_ranks/expected_layers/required_payloads + # feeding the dead commit quorum) are gone. cp_shared_l2_page_allocator = CpSharedL2PageAllocator( pages_per_payload=pages_per_payload, slabs_by_payload=cp_shared_l2_slabs_by_payload, - expected_ranks=range(cp_size), - expected_layers=range(self.kv_cache.layer_num), - required_payloads=(PAYLOAD_TARGET_KV,), ) self.cache_controller = HiCacheController( params.token_to_kv_pool_allocator, @@ -1088,7 +1083,6 @@ class HiRadixCache(RadixCache): enable_storage_metrics=self.enable_storage_metrics, cp_shared_kv_layout=cp_shared_kv_layout, cp_shared_l2_page_allocator=cp_shared_l2_page_allocator, - cp_shared_l2_cpu_group=cp_shared_l2_cpu_group, ) if draft_token_to_kv_pool is not None: self.attach_draft_kv_pool( diff --git a/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py b/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py index 931ac60c9..9e17c0cba 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py +++ b/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py @@ -963,20 +963,10 @@ class TestCpSharedHostSlabPrimitives(unittest.TestCase): class TestCpSharedL2PageAllocator(unittest.TestCase): - def make_allocator( - self, - *, - pages=8, - ranks=(0, 1, 2, 3), - layers=(0, 1), - payloads=(PAYLOAD_TARGET_KV,), - ): + def make_allocator(self, *, pages=8): return _cp_shared_l2_pool.CpSharedL2PageAllocator( pages_per_payload={PAYLOAD_TARGET_KV: pages, PAYLOAD_DRAFT_KV: pages}, slab_ids_by_payload={PAYLOAD_TARGET_KV: 10, PAYLOAD_DRAFT_KV: 11}, - expected_ranks=ranks, - expected_layers=layers, - required_payloads=payloads, ) @@ -1058,9 +1048,8 @@ class TestCpSharedL2PageAllocator(unittest.TestCase): def test_split_committed_object_preserves_live_capacity_accounting(self): allocator = self.make_allocator(pages=8) original = allocator.reserve("node-12", PAYLOAD_TARGET_KV, 4) - for layer in (0, 1): - for rank in (0, 1, 2, 3): - allocator.commit_layer("node-12", PAYLOAD_TARGET_KV, layer, rank) + # B1: commit rides the writing_check MIN frontier -> mark_object_committed. + allocator.mark_object_committed("node-12") self.assertTrue(allocator.is_committed("node-12")) parent, child = allocator.split_committed_object( @@ -1234,11 +1223,9 @@ class TestCpSharedL2PageAllocator(unittest.TestCase): } ] }, - expected_ranks=(0,), - expected_layers=(0,), ) original = allocator.reserve("node-global", PAYLOAD_TARGET_KV, 4) - self.assertTrue(allocator.commit_layer("node-global", PAYLOAD_TARGET_KV, 0, 0)) + allocator.mark_object_committed("node-global") parent, child = allocator.split_committed_object( "node-global", @@ -1324,112 +1311,9 @@ class TestCpSharedL2PageAllocator(unittest.TestCase): }, ) - def test_adopt_reserved_range_validates_known_slab_and_namespace(self): - allocator = _cp_shared_l2_pool.CpSharedL2PageAllocator( - slabs_by_payload={ - PAYLOAD_TARGET_KV: [ - { - "payload_kind": PAYLOAD_TARGET_KV, - "slab_id": 3, - "global_base_page": 50, - "num_pages": 4, - } - ] - } - ) - - allocator.adopt_reserved_range( - CpSharedL2ObjectRange( - object_key="adopted", - payload_kind=PAYLOAD_TARGET_KV, - slab_id=3, - base_page=51, - num_pages=2, - generation=9, - ) - ) - self.assertEqual(allocator.free_pages(PAYLOAD_TARGET_KV), 4) - with self.assertRaisesRegex(ValueError, "known slab"): - allocator.adopt_reserved_range( - CpSharedL2ObjectRange( - object_key="bad-slab", - payload_kind=PAYLOAD_TARGET_KV, - slab_id=4, - base_page=50, - num_pages=1, - generation=9, - ) - ) - with self.assertRaisesRegex(ValueError, "within slab"): - allocator.adopt_reserved_range( - CpSharedL2ObjectRange( - object_key="bad-range", - payload_kind=PAYLOAD_TARGET_KV, - slab_id=3, - base_page=49, - num_pages=1, - generation=9, - ) - ) - def test_object_range_still_has_no_numa_node_field(self): self.assertNotIn("numa_node", CpSharedL2ObjectRange.__dataclass_fields__) - def test_commit_requires_all_expected_ranks_layers_and_payloads_and_is_idempotent(self): - allocator = self.make_allocator( - pages=10, - ranks=(0, 1), - layers=(0, 1), - payloads=(PAYLOAD_TARGET_KV, PAYLOAD_DRAFT_KV), - ) - allocator.reserve("obj", PAYLOAD_TARGET_KV, 2) - allocator.reserve("obj", PAYLOAD_DRAFT_KV, 1) - - self.assertFalse(allocator.commit_layer("obj", PAYLOAD_TARGET_KV, 0, 0)) - self.assertFalse(allocator.is_committed("obj")) - self.assertFalse(allocator.commit_layer("obj", PAYLOAD_TARGET_KV, 0, 0)) - self.assertFalse(allocator.commit_layer("obj", PAYLOAD_TARGET_KV, 0, 1)) - self.assertFalse(allocator.commit_layer("obj", PAYLOAD_TARGET_KV, 1, 0)) - self.assertFalse(allocator.commit_layer("obj", PAYLOAD_TARGET_KV, 1, 1)) - self.assertFalse(allocator.commit_layer("obj", PAYLOAD_DRAFT_KV, 0, 0)) - self.assertFalse(allocator.commit_layer("obj", PAYLOAD_DRAFT_KV, 0, 1)) - self.assertFalse(allocator.commit_layer("obj", PAYLOAD_DRAFT_KV, 1, 0)) - - self.assertTrue(allocator.commit_layer("obj", PAYLOAD_DRAFT_KV, 1, 1)) - self.assertTrue(allocator.is_committed("obj")) - self.assertTrue(allocator.commit_layer("obj", PAYLOAD_DRAFT_KV, 1, 1)) - - def test_object_commit_contract_supports_payload_specific_layer_sets(self): - allocator = self.make_allocator( - pages=10, - ranks=(0,), - layers=(0, 1), - payloads=(PAYLOAD_TARGET_KV,), - ) - allocator.reserve("obj-draft-small", PAYLOAD_TARGET_KV, 2) - allocator.reserve("obj-draft-small", PAYLOAD_DRAFT_KV, 2) - allocator.set_object_required_payloads( - "obj-draft-small", - (PAYLOAD_TARGET_KV, PAYLOAD_DRAFT_KV), - expected_layers_by_payload={ - PAYLOAD_TARGET_KV: range(2), - PAYLOAD_DRAFT_KV: range(1), - }, - ) - - self.assertFalse( - allocator.commit_layer("obj-draft-small", PAYLOAD_TARGET_KV, 0, 0) - ) - self.assertFalse( - allocator.commit_layer("obj-draft-small", PAYLOAD_TARGET_KV, 1, 0) - ) - self.assertTrue( - allocator.commit_layer("obj-draft-small", PAYLOAD_DRAFT_KV, 0, 0) - ) - self.assertTrue(allocator.is_committed("obj-draft-small")) - with self.assertRaisesRegex(ValueError, "unexpected layer_id"): - allocator.commit_layer("obj-draft-small", PAYLOAD_DRAFT_KV, 1, 0) - def test_same_node_helper_accepts_local_and_failfasts_for_remote_rank(self): self.assertEqual( _cp_shared_l2_pool.require_cp_shared_l2_same_node( @@ -1486,27 +1370,6 @@ class TestCpSharedL2PageAllocator(unittest.TestCase): self.assertEqual(calls[0], ("get_global_rank", "cp-group", 0)) self.assertEqual(calls[1][0:3], ("broadcast", 5, "cp-group")) - def test_broadcast_helpers_use_injected_function_without_distributed_init(self): - calls = [] - - def fake_broadcast(object_list, src, group): - calls.append((src, group, list(object_list))) - if object_list[0] is None: - object_list[0] = {"decision": "commit", "object_key": "obj"} - - decision_rank0 = _cp_shared_l2_pool.broadcast_cp_shared_l2_decision( - {"decision": "commit", "object_key": "obj"}, - cp_cpu_group="group", - rank=0, - broadcast_fn=fake_broadcast, - ) - decision_rank1 = _cp_shared_l2_pool.broadcast_cp_shared_l2_decision( - None, cp_cpu_group="group", rank=1, broadcast_fn=fake_broadcast - ) - - self.assertEqual(decision_rank0["decision"], "commit") - self.assertEqual(decision_rank1["object_key"], "obj") - self.assertEqual(len(calls), 2) @@ -2716,9 +2579,6 @@ class TestCpSharedL2PlacementDigest(unittest.TestCase): return _cp_shared_l2_pool.CpSharedL2PageAllocator( pages_per_payload={PAYLOAD_TARGET_KV: pages, PAYLOAD_DRAFT_KV: pages}, slab_ids_by_payload={PAYLOAD_TARGET_KV: 10, PAYLOAD_DRAFT_KV: 11}, - expected_ranks=(0, 1, 2, 3), - expected_layers=(0, 1), - required_payloads=(PAYLOAD_TARGET_KV,), ) def _apply(self, allocator, ops): @@ -2793,9 +2653,6 @@ class TestCpSharedL2MarkObjectCommitted(unittest.TestCase): return _cp_shared_l2_pool.CpSharedL2PageAllocator( pages_per_payload={PAYLOAD_TARGET_KV: pages}, slab_ids_by_payload={PAYLOAD_TARGET_KV: 10}, - expected_ranks=(0, 1, 2, 3), - expected_layers=(0, 1), - required_payloads=(PAYLOAD_TARGET_KV,), ) def test_mark_commits_without_per_layer_quorum(self): @@ -2836,8 +2693,8 @@ class TestCpSharedL2MarkObjectCommitted(unittest.TestCase): self.assertEqual(a.placement_digest(), b.placement_digest()) def test_split_of_marked_committed_object_is_coherent(self): - # MF4: split a B1-committed object (no per-layer commit facts) -> children - # are committed, pages preserved, no error from the empty _commits_by_object. + # B1: split a marked-committed object -> children are committed, pages + # preserved (the allocator carries only ranges + the committed bit). a = self.make_allocator() a.reserve("child", PAYLOAD_TARGET_KV, 6) a.mark_object_committed("child") @@ -2877,9 +2734,6 @@ class TestCpSharedL2EightRankReserveDeterminism(unittest.TestCase): _cp_shared_l2_pool.CpSharedL2PageAllocator( pages_per_payload={PAYLOAD_TARGET_KV: pages, PAYLOAD_DRAFT_KV: pages}, slab_ids_by_payload={PAYLOAD_TARGET_KV: 10, PAYLOAD_DRAFT_KV: 11}, - expected_ranks=range(self.CP), - expected_layers=(0, 1), - required_payloads=(PAYLOAD_TARGET_KV,), ) for _ in range(self.CP) ]