2b dead-island GC: collapse pooled-L2 allocator to the B1 commit model
The 2b.1b adopt-then-strip left the entire option-A commit-quorum apparatus
orphaned: B1's writing_check ReduceOp.MIN frontier (mark_object_committed) is
the all-ranks-done consensus, so the per-(rank,layer,payload) gather quorum it
replaced is information-redundant (design sec 2.4 Corollary). This GC removes
the apparatus and collapses the CpSharedL2PageAllocator commit model to exactly
{_ranges_by_object, _committed_objects, _free_by_payload} -- precisely what
placement_digest() hashes -- so the cross-rank digest assert now covers the
whole model and the dead quorum is structurally unrepresentable.
Removed (all proven dead under B1, confirmed by a full tree caller-sweep +
opus adversarial review = SHIP):
- allocator: commit_layer, _has_full_commit, _expected_layers_for_object_payload,
adopt_reserved_range, set_object_required_payloads; fields _commits_by_object,
_expected_ranks, _adopted_ranges, _required_payloads_by_object,
_expected_layers_by_object, _required_payloads, _expected_layers; ctor args
expected_ranks/expected_layers/required_payloads; module fns
broadcast_cp_shared_l2_decision, gather_cp_shared_l2_{preflight,commits,commit}.
reserve/split_committed_object/_drop_object simplified to ranges+committed-bit.
- cache_controller: the 3 dead _commit_cp_shared_l2_* fns, 3 imports, all 5
option-A ctor params (cp_shared_l2_{cpu_group,source_rank,broadcast_fn,
preflight_fn,commit_fn}), the per-object commit-contract builder/setter/call.
- hiradix: allocator construction updated (3 args dropped); live _split_node
caller and _get_cp_shared_l2_rank_and_group@167 intact.
- tests: split tests migrated to mark_object_committed; 4 dead-symbol tests removed.
Lost __init__ payload-has-slab validation is covered by reserve()'s own
ValueError (fail-loud at first write). Default (flag-off) path unchanged: the
allocator is never constructed when enable_cp_shared_physical_l2_hicache=False.
Net -580 LOC. Validated: 87/87 pool suite (syh-dev-new, torch) + 31/31 core
classes locally + edited-module import smoke. Out of scope (flagged separately):
the CpSharedL2NodeMetadata.{required_payloads,committed_payload_layers} vestige.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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}")
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user