leavelet
9368c33b23
Two latent CP shared-KV L2 correctness fixes, landed BEFORE L3 (not in an L3 commit). Surfaced by a per-CP8-instance (CP8DP2EP16) scoping review. PREREQ-1 (CP-group scoping). The B1 commit/evict consensus collectives (writing_check ReduceOp.MIN, placement_digest MIN/MAX + its tp_world_size<=1 entry guard, drain_storage_control_queues, the evict/prefetch MINs, the flush barrier) AND cache_controller.prefetch_tp_group all derive from self.tp_group/self.tp_world_size, which was params.tp_cache_group. tp_cache_group equals the CP group ONLY at dp_size=1 (enable_dp_attention False -> tp_cpu_group, and attn_cp_size==tp_size -> _ATTN_CP==_TP) -- the sole reason B1 works today. Under CP8DP2 (DP attention, attn_tp_size=1) tp_cache_group is the size-1 attn-TP group, so every `tp_world_size>1` collective silently no-ops per rank and the placement assert self-disables -> divergent placement -> shared-slab corruption. Fix: for CP hicache (cp_size>1) scope self.tp_group to the CP cpu group (get_attention_cp_group().cpu_group -- already used for the slab-handle broadcast) + self.tp_world_size to its size. A no-op handle change at dp_size=1 (same group object); the intended fix at CP8DP2. The single init-point change propagates to every consensus collective + un-gates prefetch_tp_group + re-enables the placement assert. PREREQ-2 (flush reset). HiRadixCache.reset() cleared the radix tree + host pool but never reset CpSharedL2PageAllocator -> stale free list/ranges/committed after flush_cache (leak; the shared pool was never reclaimed). Added CpSharedL2PageAllocator.reset() (rebuild the per-slab free list all-free, drop ranges + committed, restore the freshly-built placement_digest) called from cache_controller.reset() after the ack queues are cleared. Safe: flush_cache is idle-gated (no in-flight backup/reserve). This is also L3's clear hookpoint. Validation: new test_reset_restores_freshly_constructed_all_free_state + 89/89 pool suite (torch-2.11 container) + import smoke. PREREQ-1 is a no-op at dp_size=1 (live no-regression confirmed on the next prefill restart); CP8DP2 correctness is by construction (CP-group membership verified) pending a 2-machine run. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1569 lines
59 KiB
Python
1569 lines
59 KiB
Python
# Copyright 2023-2024 SGLang Team
|
|
# Licensed under the Apache License, Version 2.0
|
|
"""Ownerless shared physical L2 HiCache metadata and host slab primitives.
|
|
|
|
Feature B adds CPU-level hugetlbfs slab/preflight building blocks only. The
|
|
controller write/load path, cudaHostRegister scheduling, and NUMA placement
|
|
orchestration are intentionally left to later phases.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ctypes
|
|
import hashlib
|
|
import inspect
|
|
import mmap
|
|
import operator
|
|
import os
|
|
import re
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Iterable, Mapping
|
|
|
|
PAYLOAD_TARGET_KV = "target_kv"
|
|
PAYLOAD_DRAFT_KV = "draft_kv"
|
|
PAYLOAD_INDEX_K = "index_k"
|
|
|
|
CP_SHARED_L2_SLAB_BACKEND = "hugetlbfs_2m"
|
|
CP_SHARED_L2_HOST_SLAB_BACKEND = CP_SHARED_L2_SLAB_BACKEND
|
|
CP_SHARED_L2_OS_PAGE_SIZE = 2 * 1024 * 1024
|
|
CP_SHARED_L2_HOST_SLAB_OS_PAGE_SIZE = CP_SHARED_L2_OS_PAGE_SIZE
|
|
CP_SHARED_L2_NUMA_POLICY = "interleave_2m"
|
|
CP_SHARED_L2_HOST_SLAB_NUMA_POLICY = CP_SHARED_L2_NUMA_POLICY
|
|
CP_SHARED_L2_FAILFAST_PREFIX = "[CP_SHARED_L2_FAILFAST]"
|
|
CP_SHARED_L2_NUMA_POLICIES = frozenset((
|
|
"interleave_2m",
|
|
"balanced_local_preferred",
|
|
"strict_local",
|
|
"round_robin",
|
|
))
|
|
|
|
_VALID_PAYLOAD_KINDS = frozenset(
|
|
(PAYLOAD_TARGET_KV, PAYLOAD_DRAFT_KV, PAYLOAD_INDEX_K)
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CpSharedL2SlabInfo:
|
|
payload_kind: str
|
|
slab_id: int
|
|
global_base_page: int
|
|
num_pages: int
|
|
numa_node: int | None = None
|
|
|
|
def __post_init__(self):
|
|
if self.payload_kind not in _VALID_PAYLOAD_KINDS:
|
|
raise ValueError(f"unknown payload_kind: {self.payload_kind!r}")
|
|
if self.slab_id < 0:
|
|
raise ValueError("slab_id must be non-negative")
|
|
if self.global_base_page < 0:
|
|
raise ValueError("global_base_page must be non-negative")
|
|
if self.num_pages <= 0:
|
|
raise ValueError("num_pages must be positive")
|
|
if self.numa_node is not None and self.numa_node < 0:
|
|
raise ValueError("numa_node must be non-negative")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CpSharedL2ObjectRange:
|
|
object_key: str
|
|
payload_kind: str
|
|
slab_id: int
|
|
base_page: int
|
|
num_pages: int
|
|
generation: int
|
|
|
|
def __post_init__(self):
|
|
if self.payload_kind not in _VALID_PAYLOAD_KINDS:
|
|
raise ValueError(f"unknown payload_kind: {self.payload_kind!r}")
|
|
if not self.object_key:
|
|
raise ValueError("object_key must be non-empty")
|
|
if self.slab_id < 0:
|
|
raise ValueError("slab_id must be non-negative")
|
|
if self.base_page < 0:
|
|
raise ValueError("base_page must be non-negative")
|
|
if self.num_pages <= 0:
|
|
raise ValueError("num_pages must be positive")
|
|
if self.generation < 0:
|
|
raise ValueError("generation must be non-negative")
|
|
|
|
|
|
def build_cp_shared_l2_slabs_by_payload(
|
|
pages_per_payload: Mapping[str, int],
|
|
slab_pages_by_payload: Mapping[str, int] | None = None,
|
|
*,
|
|
slab_size_pages: int = 0,
|
|
numa_policy: str = CP_SHARED_L2_NUMA_POLICY,
|
|
) -> dict[str, tuple[CpSharedL2SlabInfo, ...]]:
|
|
"""Build contiguous slab metadata for each shared L2 payload namespace.
|
|
|
|
``slab_size_pages <= 0`` preserves the historical one-physical-slab per
|
|
logical payload behavior. NUMA policy is metadata-only in this phase; slab
|
|
descriptors intentionally keep ``numa_node=None`` until production NUMA
|
|
placement is wired by a later patch.
|
|
"""
|
|
|
|
if numa_policy not in CP_SHARED_L2_NUMA_POLICIES:
|
|
raise ValueError(f"unknown CP shared L2 NUMA policy: {numa_policy!r}")
|
|
if int(slab_size_pages) < 0:
|
|
raise ValueError("slab_size_pages must be non-negative")
|
|
if not pages_per_payload:
|
|
raise ValueError("pages_per_payload must not be empty")
|
|
|
|
result: dict[str, tuple[CpSharedL2SlabInfo, ...]] = {}
|
|
slab_pages_by_payload = dict(slab_pages_by_payload or {})
|
|
for payload_kind, pages_value in pages_per_payload.items():
|
|
if payload_kind not in _VALID_PAYLOAD_KINDS:
|
|
raise ValueError(f"unknown payload_kind: {payload_kind!r}")
|
|
pages = int(pages_value)
|
|
if pages <= 0:
|
|
raise ValueError("pages_per_payload values must be positive")
|
|
payload_slab_pages = int(
|
|
slab_pages_by_payload.get(payload_kind, int(slab_size_pages))
|
|
)
|
|
if payload_slab_pages < 0:
|
|
raise ValueError("slab pages must be non-negative")
|
|
if payload_slab_pages <= 0 or payload_slab_pages >= pages:
|
|
result[payload_kind] = (
|
|
CpSharedL2SlabInfo(
|
|
payload_kind=payload_kind,
|
|
slab_id=0,
|
|
global_base_page=0,
|
|
num_pages=pages,
|
|
numa_node=None,
|
|
),
|
|
)
|
|
continue
|
|
|
|
slabs: list[CpSharedL2SlabInfo] = []
|
|
base_page = 0
|
|
slab_id = 0
|
|
while base_page < pages:
|
|
num_pages = min(payload_slab_pages, pages - base_page)
|
|
slabs.append(
|
|
CpSharedL2SlabInfo(
|
|
payload_kind=payload_kind,
|
|
slab_id=slab_id,
|
|
global_base_page=base_page,
|
|
num_pages=num_pages,
|
|
numa_node=None,
|
|
)
|
|
)
|
|
base_page += num_pages
|
|
slab_id += 1
|
|
result[payload_kind] = tuple(slabs)
|
|
return result
|
|
|
|
|
|
def cp_shared_l2_logical_token_indices(
|
|
object_range: CpSharedL2ObjectRange,
|
|
page_size: int,
|
|
positions: Any,
|
|
):
|
|
"""Translate token positions in an object range to global logical host indices.
|
|
|
|
``object_range.base_page`` is a global logical page id in the shared L2
|
|
namespace. Callers pass positions relative to the start of the ownerless
|
|
object range and receive indices in that global logical host namespace.
|
|
Tensor-like positions must be integral CPU tensors; tensor returns are
|
|
promoted to int64 before adding the global offset.
|
|
"""
|
|
|
|
page_size = int(page_size)
|
|
if page_size <= 0:
|
|
raise ValueError("page_size must be positive")
|
|
object_token_capacity = int(object_range.num_pages) * page_size
|
|
positions = _normalize_cp_shared_l2_token_positions(positions)
|
|
_validate_cp_shared_l2_token_positions(positions, object_token_capacity)
|
|
offset = int(object_range.base_page) * page_size
|
|
if hasattr(positions, "numel"):
|
|
return offset + positions
|
|
if isinstance(positions, range):
|
|
return [offset + int(position) for position in positions]
|
|
if isinstance(positions, list):
|
|
return [offset + int(position) for position in positions]
|
|
if isinstance(positions, tuple):
|
|
return tuple(offset + int(position) for position in positions)
|
|
return offset + int(positions)
|
|
|
|
|
|
def _normalize_cp_shared_l2_token_positions(positions: Any):
|
|
if hasattr(positions, "numel"):
|
|
if bool(getattr(positions, "is_cuda", False)):
|
|
raise ValueError("positions tensor must be on CPU")
|
|
dtype_name = str(getattr(positions, "dtype", "")).lower()
|
|
if (
|
|
not dtype_name
|
|
or "bool" in dtype_name
|
|
or "float" in dtype_name
|
|
or "complex" in dtype_name
|
|
or ("int" not in dtype_name and "uint" not in dtype_name)
|
|
):
|
|
raise TypeError("positions must be integral")
|
|
if not hasattr(positions, "long"):
|
|
raise TypeError("integral tensor positions must support int64 promotion")
|
|
return positions.long()
|
|
if isinstance(positions, range):
|
|
return positions
|
|
if isinstance(positions, list):
|
|
return [
|
|
_coerce_cp_shared_l2_integral_position(position)
|
|
for position in positions
|
|
]
|
|
if isinstance(positions, tuple):
|
|
return tuple(
|
|
_coerce_cp_shared_l2_integral_position(position) for position in positions
|
|
)
|
|
return _coerce_cp_shared_l2_integral_position(positions)
|
|
|
|
|
|
def _coerce_cp_shared_l2_integral_position(position: Any) -> int:
|
|
if isinstance(position, bool):
|
|
raise TypeError("positions must be integral")
|
|
try:
|
|
return operator.index(position)
|
|
except TypeError as exc:
|
|
raise TypeError("positions must be integral") from exc
|
|
|
|
|
|
def _validate_cp_shared_l2_token_positions(positions: Any, object_token_capacity: int):
|
|
if hasattr(positions, "numel"):
|
|
if int(positions.numel()) == 0:
|
|
return
|
|
min_position = int(positions.min().item())
|
|
max_position = int(positions.max().item())
|
|
else:
|
|
if isinstance(positions, range):
|
|
if len(positions) == 0:
|
|
return
|
|
min_position = min(positions)
|
|
max_position = max(positions)
|
|
elif isinstance(positions, (list, tuple)):
|
|
if not positions:
|
|
return
|
|
min_position = min(int(position) for position in positions)
|
|
max_position = max(int(position) for position in positions)
|
|
else:
|
|
min_position = max_position = int(positions)
|
|
if min_position < 0 or max_position >= object_token_capacity:
|
|
raise ValueError(
|
|
"positions must be within object range token capacity: "
|
|
f"min={min_position} max={max_position} "
|
|
f"capacity={object_token_capacity}"
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CpSharedHostSlabHandle:
|
|
name: str
|
|
path: str
|
|
nbytes: int
|
|
shape: tuple[int, ...]
|
|
dtype_name: str
|
|
creator_rank: int
|
|
backend: str = CP_SHARED_L2_SLAB_BACKEND
|
|
hugetlbfs_dir: str = ""
|
|
os_page_size: int = CP_SHARED_L2_OS_PAGE_SIZE
|
|
numa_policy: str = CP_SHARED_L2_NUMA_POLICY
|
|
mapped_nbytes: int = 0
|
|
|
|
def __post_init__(self):
|
|
if not self.name:
|
|
_failfast("host slab name must be non-empty")
|
|
if not self.path:
|
|
_failfast("host slab path must be non-empty")
|
|
if self.nbytes <= 0:
|
|
_failfast("host slab nbytes must be positive")
|
|
if any(dim < 0 for dim in self.shape):
|
|
_failfast("host slab shape dimensions must be non-negative")
|
|
if not self.dtype_name:
|
|
_failfast("host slab dtype_name must be non-empty")
|
|
if self.creator_rank < 0:
|
|
_failfast("host slab creator_rank must be non-negative")
|
|
if not self.hugetlbfs_dir:
|
|
object.__setattr__(self, "hugetlbfs_dir", str(Path(self.path).parent))
|
|
if self.mapped_nbytes <= 0:
|
|
object.__setattr__(
|
|
self,
|
|
"mapped_nbytes",
|
|
round_up_cp_shared_l2_host_slab_bytes(
|
|
self.nbytes, page_size=self.os_page_size
|
|
),
|
|
)
|
|
validate_cp_shared_host_slab_handle_config(self)
|
|
|
|
|
|
|
|
def unlink_cp_shared_host_slab_mappings_all_or_none(mappings) -> bool:
|
|
"""Early-unlink a group of slab mappings with old-path rollback on failure.
|
|
|
|
The source rank uses this for multi-slab groups so an unlink failure on a
|
|
later slab does not leave earlier slabs without their original paths. Backup
|
|
hard links live in the same directory so restore is another metadata-only
|
|
link operation to the same inode.
|
|
"""
|
|
|
|
live_mappings = [mapping for mapping in mappings if not mapping._unlinked]
|
|
if not live_mappings:
|
|
return False
|
|
|
|
for mapping in live_mappings:
|
|
mapping.check_unlink_ready()
|
|
|
|
backups = []
|
|
backup_token = f"{os.getpid()}-{uuid.uuid4().hex}"
|
|
try:
|
|
for idx, mapping in enumerate(live_mappings):
|
|
path = Path(mapping.handle.path)
|
|
backup_path = path.with_name(
|
|
f".cp-shared-l2-unlink-backup-{backup_token}-{idx}"
|
|
)
|
|
os.link(path, backup_path)
|
|
backups.append((mapping, path, backup_path))
|
|
|
|
try:
|
|
for mapping, path, _backup_path in backups:
|
|
os.unlink(path)
|
|
except Exception:
|
|
for _mapping, path, backup_path in backups:
|
|
if not path.exists() and backup_path.exists():
|
|
os.link(backup_path, path)
|
|
raise
|
|
|
|
for mapping, _path, _backup_path in backups:
|
|
mapping._unlinked = True
|
|
mapping._unlink_on_close = False
|
|
return True
|
|
finally:
|
|
for _mapping, _path, backup_path in backups:
|
|
try:
|
|
os.unlink(backup_path)
|
|
except FileNotFoundError:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
@dataclass
|
|
class CpSharedHostSlabMapping:
|
|
handle: CpSharedHostSlabHandle
|
|
mmap: mmap.mmap
|
|
_file: Any
|
|
_unlink_on_close: bool = False
|
|
_unlinked: bool = False
|
|
_closed: bool = False
|
|
|
|
def check_unlink_ready(self) -> None:
|
|
"""Fail before early-unlink if the slab path is not safe to unlink.
|
|
|
|
This is a best-effort preflight for group all-or-none early unlink. It
|
|
cannot make the following POSIX unlink atomic, but it catches expected
|
|
lifecycle/configuration failures before any slab in a group is unlinked.
|
|
"""
|
|
|
|
if self._closed:
|
|
raise RuntimeError(
|
|
"shared host slab mapping is already closed: "
|
|
f"{self.handle.path}"
|
|
)
|
|
if self._unlinked:
|
|
return
|
|
path = Path(self.handle.path)
|
|
if not path.exists():
|
|
raise FileNotFoundError(f"shared host slab path does not exist: {path}")
|
|
parent = path.parent
|
|
if not os.access(parent, os.W_OK):
|
|
raise PermissionError(f"shared host slab directory is not writable: {parent}")
|
|
|
|
def unlink(self) -> bool:
|
|
"""Remove the slab name while keeping the live mmap/fd usable.
|
|
|
|
POSIX unlink only removes the directory entry. Existing mappings and
|
|
open file descriptors remain valid until every process closes them,
|
|
which is the lifecycle we want after all CP ranks have attached and
|
|
registered the shared physical L2 slab.
|
|
"""
|
|
|
|
if self._unlinked:
|
|
return False
|
|
try:
|
|
os.unlink(self.handle.path)
|
|
except FileNotFoundError:
|
|
self._unlinked = True
|
|
self._unlink_on_close = False
|
|
return False
|
|
self._unlinked = True
|
|
self._unlink_on_close = False
|
|
return True
|
|
|
|
def close(self):
|
|
if self._closed:
|
|
return
|
|
try:
|
|
try:
|
|
self.mmap.close()
|
|
finally:
|
|
try:
|
|
self._file.close()
|
|
finally:
|
|
if self._unlink_on_close and not self._unlinked:
|
|
self.unlink()
|
|
finally:
|
|
self._closed = True
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
self.close()
|
|
return False
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CpSharedL2RegistrationRange:
|
|
ptr: int
|
|
nbytes: int
|
|
|
|
def __post_init__(self):
|
|
if self.ptr <= 0:
|
|
_failfast("registration ptr must be a positive address")
|
|
if self.nbytes <= 0:
|
|
_failfast("registration nbytes must be positive")
|
|
|
|
|
|
@dataclass
|
|
class CpSharedL2NodeMetadata:
|
|
logical_len: int
|
|
padded_len: int
|
|
page_size: int
|
|
object_ranges: dict[str, CpSharedL2ObjectRange]
|
|
required_payloads: tuple[str, ...]
|
|
object_key: str = ""
|
|
|
|
def __post_init__(self):
|
|
if self.logical_len < 0:
|
|
raise ValueError("logical_len must be non-negative")
|
|
if self.padded_len < self.logical_len:
|
|
raise ValueError("padded_len must be >= logical_len")
|
|
if self.page_size <= 0:
|
|
raise ValueError("page_size must be positive")
|
|
|
|
for payload_kind in self.required_payloads:
|
|
self._validate_payload_kind(payload_kind)
|
|
if payload_kind not in self.object_ranges:
|
|
raise ValueError(
|
|
f"required payload {payload_kind!r} is missing an object range"
|
|
)
|
|
|
|
for payload_kind, object_range in self.object_ranges.items():
|
|
self._validate_payload_kind(payload_kind)
|
|
if object_range.payload_kind != payload_kind:
|
|
raise ValueError(
|
|
"object_ranges keys must match each CpSharedL2ObjectRange "
|
|
f"payload_kind; got key={payload_kind!r}, "
|
|
f"payload_kind={object_range.payload_kind!r}"
|
|
)
|
|
|
|
@staticmethod
|
|
def _validate_payload_kind(payload_kind: str):
|
|
if payload_kind not in _VALID_PAYLOAD_KINDS:
|
|
raise ValueError(f"unknown payload_kind: {payload_kind!r}")
|
|
|
|
def split(
|
|
self,
|
|
split_len: int,
|
|
*,
|
|
parent_object_key: str | None = None,
|
|
child_object_key: str | None = None,
|
|
) -> tuple["CpSharedL2NodeMetadata", "CpSharedL2NodeMetadata"]:
|
|
split_len = int(split_len)
|
|
if split_len < 0 or split_len > self.logical_len:
|
|
raise ValueError(
|
|
f"split_len must be in [0, {self.logical_len}], got {split_len}"
|
|
)
|
|
if split_len % self.page_size != 0:
|
|
raise ValueError(
|
|
f"split_len ({split_len}) must be a multiple of page_size "
|
|
f"({self.page_size})"
|
|
)
|
|
total_pages = self.padded_len // self.page_size
|
|
split_pages = split_len // self.page_size
|
|
child_pages = total_pages - split_pages
|
|
if split_pages <= 0 or child_pages <= 0:
|
|
raise ValueError(
|
|
"shared L2 metadata split must leave at least one page on each side"
|
|
)
|
|
|
|
parent_key = parent_object_key or self.object_key
|
|
child_key = child_object_key or self.object_key
|
|
|
|
parent_ranges = {}
|
|
child_ranges = {}
|
|
for payload_kind, object_range in self.object_ranges.items():
|
|
if object_range.num_pages != total_pages:
|
|
raise ValueError(
|
|
"shared L2 metadata object range page count must match "
|
|
"padded_len/page_size before split: "
|
|
f"payload={payload_kind!r} range_pages={object_range.num_pages} "
|
|
f"metadata_pages={total_pages}"
|
|
)
|
|
parent_ranges[payload_kind] = CpSharedL2ObjectRange(
|
|
object_key=parent_key,
|
|
payload_kind=object_range.payload_kind,
|
|
slab_id=object_range.slab_id,
|
|
base_page=object_range.base_page,
|
|
num_pages=split_pages,
|
|
generation=object_range.generation,
|
|
)
|
|
child_ranges[payload_kind] = CpSharedL2ObjectRange(
|
|
object_key=child_key,
|
|
payload_kind=object_range.payload_kind,
|
|
slab_id=object_range.slab_id,
|
|
base_page=object_range.base_page + split_pages,
|
|
num_pages=child_pages,
|
|
generation=object_range.generation,
|
|
)
|
|
|
|
return (
|
|
CpSharedL2NodeMetadata(
|
|
logical_len=split_len,
|
|
padded_len=split_len,
|
|
page_size=self.page_size,
|
|
object_ranges=parent_ranges,
|
|
required_payloads=tuple(self.required_payloads),
|
|
object_key=parent_key,
|
|
),
|
|
CpSharedL2NodeMetadata(
|
|
logical_len=self.logical_len - split_len,
|
|
padded_len=self.padded_len - split_len,
|
|
page_size=self.page_size,
|
|
object_ranges=child_ranges,
|
|
required_payloads=tuple(self.required_payloads),
|
|
object_key=child_key,
|
|
),
|
|
)
|
|
|
|
|
|
class CpSharedL2PageAllocator:
|
|
"""Rank-0 logical-object allocator for shared physical L2 page slabs.
|
|
|
|
The allocator models one global page namespace per payload slab. A logical
|
|
object may reserve one range per payload kind, but capacity is charged once
|
|
per reserved shared range, never multiplied by CP size.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
pages_per_payload: Mapping[str, int] | int | None = None,
|
|
slab_ids_by_payload: Mapping[str, int] | None = None,
|
|
slabs_by_payload: Mapping[
|
|
str, Iterable[CpSharedL2SlabInfo | Mapping[str, Any]]
|
|
]
|
|
| None = None,
|
|
):
|
|
pages_by_payload: dict[str, int] | None
|
|
if pages_per_payload is None:
|
|
pages_by_payload = None
|
|
elif isinstance(pages_per_payload, int):
|
|
pages_by_payload = {PAYLOAD_TARGET_KV: int(pages_per_payload)}
|
|
else:
|
|
pages_by_payload = {
|
|
payload: int(pages) for payload, pages in pages_per_payload.items()
|
|
}
|
|
|
|
if slabs_by_payload is None:
|
|
if not pages_by_payload:
|
|
raise ValueError("pages_per_payload must not be empty")
|
|
slabs = self._single_slabs_from_pages(
|
|
pages_by_payload, slab_ids_by_payload=slab_ids_by_payload
|
|
)
|
|
self._capacity_by_payload: dict[str, int] = {}
|
|
self._slab_ids_by_payload: dict[str, int] = {}
|
|
self._slabs_by_payload: dict[str, tuple[CpSharedL2SlabInfo, ...]] = {}
|
|
self._slab_by_payload_id: dict[str, dict[int, CpSharedL2SlabInfo]] = {}
|
|
self._free_by_payload: dict[str, dict[int, list[tuple[int, int]]]] = {}
|
|
if slabs_by_payload is not None:
|
|
slabs = self._normalize_slabs_by_payload(slabs_by_payload)
|
|
if not slabs:
|
|
raise ValueError("slabs_by_payload must not be empty")
|
|
if pages_by_payload is not None:
|
|
self._validate_pages_match_slab_capacity(pages_by_payload, slabs)
|
|
self._install_slabs(slabs)
|
|
|
|
# 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._committed_objects: set[str] = set()
|
|
self._stats = {
|
|
"cp_shared_l2_objects_committed": 0,
|
|
"cp_shared_l2_objects_aborted": 0,
|
|
"cp_shared_l2_objects_evicted": 0,
|
|
}
|
|
|
|
def reset(self) -> None:
|
|
"""Reset to the freshly-constructed all-free state (flush_cache hook).
|
|
|
|
Rebuilds the per-slab free list to fully-free and drops all object ranges +
|
|
the committed set, so flush_cache reclaims the pooled L2 (the allocator
|
|
metadata; the slab bytes are left as-is -- unreferenced once the radix tree
|
|
and ranges are cleared). Rank-uniform: every CP rank runs it at the
|
|
idle-gated flush point, so the post-reset placement_digest is empty and
|
|
identical on every rank. Cumulative stats are preserved.
|
|
"""
|
|
self._generation = 0
|
|
self._ranges_by_object = {}
|
|
self._committed_objects = set()
|
|
self._free_by_payload = {
|
|
payload_kind: {
|
|
slab.slab_id: [(0, slab.num_pages)] for slab in slabs
|
|
}
|
|
for payload_kind, slabs in self._slabs_by_payload.items()
|
|
}
|
|
|
|
def reserve(
|
|
self, object_key: str, payload_kind: str, num_pages: int
|
|
) -> CpSharedL2ObjectRange:
|
|
self._validate_object_key(object_key)
|
|
self._validate_payload_kind(payload_kind)
|
|
if payload_kind not in self._free_by_payload:
|
|
raise ValueError(f"payload {payload_kind!r} has no slab namespace")
|
|
if num_pages <= 0:
|
|
raise ValueError("num_pages must be positive")
|
|
if payload_kind in self._ranges_by_object.get(object_key, {}):
|
|
raise ValueError(
|
|
f"duplicate live reservation for object {object_key!r} "
|
|
f"payload {payload_kind!r}"
|
|
)
|
|
|
|
slab_id, base_page = self._allocate_contiguous(payload_kind, int(num_pages))
|
|
self._generation += 1
|
|
object_range = CpSharedL2ObjectRange(
|
|
object_key=object_key,
|
|
payload_kind=payload_kind,
|
|
slab_id=slab_id,
|
|
base_page=base_page,
|
|
num_pages=int(num_pages),
|
|
generation=self._generation,
|
|
)
|
|
self._ranges_by_object.setdefault(object_key, {})[payload_kind] = object_range
|
|
self._committed_objects.discard(object_key)
|
|
return object_range
|
|
|
|
def abort(self, object_key: str) -> bool:
|
|
dropped = self._drop_object(object_key)
|
|
if dropped:
|
|
self._stats["cp_shared_l2_objects_aborted"] += 1
|
|
return dropped
|
|
|
|
def release(self, object_key: str) -> bool:
|
|
dropped = self._drop_object(object_key)
|
|
if dropped:
|
|
self._stats["cp_shared_l2_objects_evicted"] += 1
|
|
return dropped
|
|
|
|
def mark_object_committed(self, object_key: str) -> bool:
|
|
"""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 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).
|
|
"""
|
|
self._validate_object_key(object_key)
|
|
if object_key not in self._ranges_by_object:
|
|
raise ValueError(f"cannot commit unknown object {object_key!r}")
|
|
if object_key in self._committed_objects:
|
|
return False
|
|
self._committed_objects.add(object_key)
|
|
self._stats["cp_shared_l2_objects_committed"] += 1
|
|
return True
|
|
|
|
def split_committed_object(
|
|
self,
|
|
object_key: str,
|
|
*,
|
|
split_pages_by_payload: Mapping[str, int],
|
|
parent_object_key: str,
|
|
child_object_key: str,
|
|
) -> tuple[dict[str, CpSharedL2ObjectRange], dict[str, CpSharedL2ObjectRange]]:
|
|
"""Split a committed shared-L2 reservation into two live object keys.
|
|
|
|
Radix splits create a new parent node while the old child node keeps the
|
|
suffix. The original shared physical pages must not be returned to the
|
|
free list during that metadata mutation; instead each payload range is
|
|
repartitioned and reattached to the two node object keys so later
|
|
eviction releases exactly the half being evicted.
|
|
"""
|
|
|
|
self._validate_object_key(object_key)
|
|
self._validate_object_key(parent_object_key)
|
|
self._validate_object_key(child_object_key)
|
|
if parent_object_key == child_object_key:
|
|
raise ValueError("parent_object_key and child_object_key must differ")
|
|
if parent_object_key in self._ranges_by_object:
|
|
raise ValueError(f"parent object {parent_object_key!r} already exists")
|
|
if child_object_key != object_key and child_object_key in self._ranges_by_object:
|
|
raise ValueError(f"child object {child_object_key!r} already exists")
|
|
|
|
ranges = self._ranges_by_object.get(object_key)
|
|
if not ranges:
|
|
raise ValueError(f"unknown object {object_key!r}")
|
|
if object_key not in self._committed_objects:
|
|
raise ValueError(f"object {object_key!r} is not committed")
|
|
|
|
parent_ranges: dict[str, CpSharedL2ObjectRange] = {}
|
|
child_ranges: dict[str, CpSharedL2ObjectRange] = {}
|
|
for payload_kind, object_range in ranges.items():
|
|
self._validate_payload_kind(payload_kind)
|
|
if payload_kind not in split_pages_by_payload:
|
|
raise ValueError(f"missing split page count for payload {payload_kind!r}")
|
|
split_pages = int(split_pages_by_payload[payload_kind])
|
|
if split_pages <= 0 or split_pages >= object_range.num_pages:
|
|
raise ValueError(
|
|
"split page count must leave at least one page on each side: "
|
|
f"payload={payload_kind!r} split_pages={split_pages} "
|
|
f"range_pages={object_range.num_pages}"
|
|
)
|
|
parent_ranges[payload_kind] = CpSharedL2ObjectRange(
|
|
object_key=parent_object_key,
|
|
payload_kind=payload_kind,
|
|
slab_id=object_range.slab_id,
|
|
base_page=object_range.base_page,
|
|
num_pages=split_pages,
|
|
generation=object_range.generation,
|
|
)
|
|
child_ranges[payload_kind] = CpSharedL2ObjectRange(
|
|
object_key=child_object_key,
|
|
payload_kind=payload_kind,
|
|
slab_id=object_range.slab_id,
|
|
base_page=object_range.base_page + split_pages,
|
|
num_pages=object_range.num_pages - split_pages,
|
|
generation=object_range.generation,
|
|
)
|
|
|
|
# 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._committed_objects.discard(object_key)
|
|
|
|
for new_key, new_ranges in (
|
|
(parent_object_key, parent_ranges),
|
|
(child_object_key, child_ranges),
|
|
):
|
|
self._ranges_by_object[new_key] = new_ranges
|
|
self._committed_objects.add(new_key)
|
|
|
|
return parent_ranges, child_ranges
|
|
|
|
def get_range(
|
|
self, object_key: str, payload_kind: str
|
|
) -> CpSharedL2ObjectRange | None:
|
|
return self._ranges_by_object.get(object_key, {}).get(payload_kind)
|
|
|
|
def object_ranges(self, object_key: str) -> dict[str, CpSharedL2ObjectRange]:
|
|
return dict(self._ranges_by_object.get(object_key, {}))
|
|
|
|
def is_committed(self, object_key: str) -> bool:
|
|
return object_key in self._committed_objects
|
|
|
|
def used_pages(self, payload_kind: str) -> int:
|
|
self._validate_payload_kind(payload_kind)
|
|
total = self._payload_capacity(payload_kind)
|
|
return total - self.free_pages(payload_kind)
|
|
|
|
def occupancy_by_payload(self) -> dict[str, tuple[int, int]]:
|
|
"""Per-payload-slab ``(used_pages, capacity_pages)`` in a stable payload
|
|
order. For the lane-stats observability probe: the shared pool's REAL fill,
|
|
broken down by payload slab (target_kv/draft_kv/index_k). Derived purely
|
|
from the replicated free list, so identical on every CP rank."""
|
|
return {
|
|
payload: (self.used_pages(payload), self._payload_capacity(payload))
|
|
for payload in sorted(self._free_by_payload)
|
|
}
|
|
|
|
def stats(self) -> dict[str, int]:
|
|
capacity = {
|
|
payload: self._payload_capacity(payload)
|
|
for payload in self._free_by_payload
|
|
}
|
|
used = {
|
|
payload: self.used_pages(payload)
|
|
for payload in self._free_by_payload
|
|
}
|
|
return {
|
|
"cp_shared_l2_pages_capacity": sum(capacity.values()),
|
|
"cp_shared_l2_pages_used": sum(used.values()),
|
|
**self._stats,
|
|
}
|
|
|
|
def free_pages(self, payload_kind: str) -> int:
|
|
self._validate_payload_kind(payload_kind)
|
|
return sum(
|
|
length
|
|
for free_by_slab in self._free_by_payload[payload_kind].values()
|
|
for _, length in free_by_slab
|
|
)
|
|
|
|
def placement_digest(self) -> str:
|
|
"""Deterministic content hash of the full replicated allocator state.
|
|
|
|
The B1 proof-obligation gate (see docs_internal/cp_hicache_2b_pooled_l2_b1_design.md
|
|
sec 2.4, Theorem 1): two allocators that applied the IDENTICAL op sequence
|
|
over identical initial state produce identical digests on every CP rank, so
|
|
the runtime cross-rank assert can prove placement determinism each tick. Any
|
|
per-rank divergence -- a reserve/release fired on a subset of ranks, or with
|
|
a per-rank arg -- changes the digest. Order-independent within each container
|
|
(sorted) so dict/set iteration order cannot perturb it. Hashes the durable
|
|
placement state: the free list, the per-object ranges, and the committed set.
|
|
"""
|
|
h = hashlib.sha256()
|
|
for payload in sorted(self._free_by_payload):
|
|
h.update(b"P" + payload.encode())
|
|
for slab_id in sorted(self._free_by_payload[payload]):
|
|
h.update(b"|s" + str(slab_id).encode() + b":")
|
|
for base, length in sorted(self._free_by_payload[payload][slab_id]):
|
|
h.update(f"{base},{length};".encode())
|
|
for object_key in sorted(self._ranges_by_object):
|
|
ranges = self._ranges_by_object[object_key]
|
|
for payload in sorted(ranges):
|
|
r = ranges[payload]
|
|
h.update(
|
|
f"#{object_key}/{payload}={r.slab_id},{r.base_page},"
|
|
f"{r.num_pages};".encode()
|
|
)
|
|
for object_key in sorted(self._committed_objects):
|
|
h.update(b"@" + object_key.encode() + b";")
|
|
return h.hexdigest()
|
|
|
|
def _allocate_contiguous(self, payload_kind: str, num_pages: int) -> tuple[int, int]:
|
|
for slab in self._slabs_by_payload[payload_kind]:
|
|
free_ranges = self._free_by_payload[payload_kind][slab.slab_id]
|
|
for idx, (local_base, length) in enumerate(free_ranges):
|
|
if length < num_pages:
|
|
continue
|
|
allocated_base = slab.global_base_page + local_base
|
|
if length == num_pages:
|
|
del free_ranges[idx]
|
|
else:
|
|
free_ranges[idx] = (local_base + num_pages, length - num_pages)
|
|
return slab.slab_id, allocated_base
|
|
raise ValueError(
|
|
f"insufficient contiguous shared L2 pages for payload {payload_kind!r}: "
|
|
f"requested {num_pages}, free {self.free_pages(payload_kind)}"
|
|
)
|
|
|
|
def _drop_object(self, object_key: str) -> bool:
|
|
ranges = self._ranges_by_object.pop(object_key, None)
|
|
self._committed_objects.discard(object_key)
|
|
if not ranges:
|
|
return False
|
|
for object_range in ranges.values():
|
|
self._return_range(object_range)
|
|
return True
|
|
|
|
def _return_range(self, object_range: CpSharedL2ObjectRange) -> None:
|
|
slab = self._validate_range_within_known_slab(object_range)
|
|
free_ranges = self._free_by_payload[object_range.payload_kind][
|
|
object_range.slab_id
|
|
]
|
|
free_ranges.append(
|
|
(object_range.base_page - slab.global_base_page, object_range.num_pages)
|
|
)
|
|
free_ranges.sort()
|
|
merged: list[tuple[int, int]] = []
|
|
for base, length in free_ranges:
|
|
if not merged:
|
|
merged.append((base, length))
|
|
continue
|
|
prev_base, prev_length = merged[-1]
|
|
prev_end = prev_base + prev_length
|
|
if base <= prev_end:
|
|
merged[-1] = (prev_base, max(prev_end, base + length) - prev_base)
|
|
else:
|
|
merged.append((base, length))
|
|
self._free_by_payload[object_range.payload_kind][object_range.slab_id] = merged
|
|
|
|
def _payload_capacity(self, payload_kind: str) -> int:
|
|
return self._capacity_by_payload[payload_kind]
|
|
|
|
@classmethod
|
|
def _single_slabs_from_pages(
|
|
cls,
|
|
pages_by_payload: Mapping[str, int],
|
|
*,
|
|
slab_ids_by_payload: Mapping[str, int] | None,
|
|
) -> dict[str, tuple[CpSharedL2SlabInfo, ...]]:
|
|
slabs: dict[str, tuple[CpSharedL2SlabInfo, ...]] = {}
|
|
for payload_kind, pages in pages_by_payload.items():
|
|
cls._validate_payload_kind(payload_kind)
|
|
pages = int(pages)
|
|
if pages <= 0:
|
|
raise ValueError("pages_per_payload values must be positive")
|
|
slab_id = 0
|
|
if slab_ids_by_payload is not None:
|
|
slab_id = int(slab_ids_by_payload[payload_kind])
|
|
slabs[payload_kind] = (
|
|
CpSharedL2SlabInfo(
|
|
payload_kind=payload_kind,
|
|
slab_id=slab_id,
|
|
global_base_page=0,
|
|
num_pages=pages,
|
|
),
|
|
)
|
|
return slabs
|
|
|
|
@classmethod
|
|
def _normalize_slabs_by_payload(
|
|
cls,
|
|
slabs_by_payload: Mapping[
|
|
str, Iterable[CpSharedL2SlabInfo | Mapping[str, Any]]
|
|
],
|
|
) -> dict[str, tuple[CpSharedL2SlabInfo, ...]]:
|
|
normalized: dict[str, tuple[CpSharedL2SlabInfo, ...]] = {}
|
|
for payload_kind, slab_values in slabs_by_payload.items():
|
|
cls._validate_payload_kind(payload_kind)
|
|
slabs: list[CpSharedL2SlabInfo] = []
|
|
for slab_value in slab_values:
|
|
if isinstance(slab_value, CpSharedL2SlabInfo):
|
|
slab = slab_value
|
|
else:
|
|
slab = CpSharedL2SlabInfo(**dict(slab_value))
|
|
if slab.payload_kind != payload_kind:
|
|
raise ValueError(
|
|
f"slab payload_kind {slab.payload_kind!r} does not match "
|
|
f"payload namespace {payload_kind!r}"
|
|
)
|
|
slabs.append(slab)
|
|
if not slabs:
|
|
raise ValueError(f"payload {payload_kind!r} must have at least one slab")
|
|
normalized[payload_kind] = tuple(
|
|
sorted(slabs, key=lambda slab: (slab.global_base_page, slab.slab_id))
|
|
)
|
|
return normalized
|
|
|
|
@staticmethod
|
|
def _validate_pages_match_slab_capacity(
|
|
pages_by_payload: Mapping[str, int],
|
|
slabs_by_payload: Mapping[str, tuple[CpSharedL2SlabInfo, ...]],
|
|
) -> None:
|
|
if set(pages_by_payload) != set(slabs_by_payload):
|
|
raise ValueError("pages_per_payload keys must match slabs_by_payload keys")
|
|
for payload_kind, pages in pages_by_payload.items():
|
|
capacity = sum(slab.num_pages for slab in slabs_by_payload[payload_kind])
|
|
if int(pages) != capacity:
|
|
raise ValueError(
|
|
f"pages_per_payload capacity for {payload_kind!r} ({int(pages)}) "
|
|
f"does not match slabs_by_payload capacity ({capacity})"
|
|
)
|
|
|
|
def _install_slabs(
|
|
self, slabs_by_payload: Mapping[str, tuple[CpSharedL2SlabInfo, ...]]
|
|
) -> None:
|
|
for payload_kind, slabs in slabs_by_payload.items():
|
|
self._validate_payload_kind(payload_kind)
|
|
seen_slab_ids: set[int] = set()
|
|
sorted_by_base = sorted(slabs, key=lambda slab: slab.global_base_page)
|
|
previous_end: int | None = None
|
|
for slab in sorted_by_base:
|
|
if slab.slab_id in seen_slab_ids:
|
|
raise ValueError(
|
|
f"duplicate slab_id {slab.slab_id} for payload {payload_kind!r}"
|
|
)
|
|
seen_slab_ids.add(slab.slab_id)
|
|
if previous_end is not None and slab.global_base_page < previous_end:
|
|
raise ValueError(
|
|
f"overlap in slab page ranges for payload {payload_kind!r}"
|
|
)
|
|
previous_end = slab.global_base_page + slab.num_pages
|
|
|
|
self._slabs_by_payload[payload_kind] = tuple(slabs)
|
|
self._slab_by_payload_id[payload_kind] = {
|
|
slab.slab_id: slab for slab in slabs
|
|
}
|
|
self._free_by_payload[payload_kind] = {
|
|
slab.slab_id: [(0, slab.num_pages)] for slab in slabs
|
|
}
|
|
self._slab_ids_by_payload[payload_kind] = slabs[0].slab_id
|
|
self._capacity_by_payload[payload_kind] = sum(
|
|
slab.num_pages for slab in slabs
|
|
)
|
|
|
|
def _validate_range_within_known_slab(
|
|
self, object_range: CpSharedL2ObjectRange
|
|
) -> CpSharedL2SlabInfo:
|
|
slabs = self._slab_by_payload_id.get(object_range.payload_kind)
|
|
if slabs is None or object_range.slab_id not in slabs:
|
|
raise ValueError(
|
|
f"range for payload {object_range.payload_kind!r} does not reference "
|
|
f"a known slab: slab_id={object_range.slab_id}"
|
|
)
|
|
slab = slabs[object_range.slab_id]
|
|
range_start = object_range.base_page
|
|
range_end = object_range.base_page + object_range.num_pages
|
|
slab_start = slab.global_base_page
|
|
slab_end = slab.global_base_page + slab.num_pages
|
|
if range_start < slab_start or range_end > slab_end:
|
|
raise ValueError(
|
|
f"range for payload {object_range.payload_kind!r} must lie within "
|
|
f"slab {object_range.slab_id}: range=[{range_start}, {range_end}) "
|
|
f"slab=[{slab_start}, {slab_end})"
|
|
)
|
|
return slab
|
|
|
|
@staticmethod
|
|
def _validate_payload_kind(payload_kind: str) -> None:
|
|
if payload_kind not in _VALID_PAYLOAD_KINDS:
|
|
raise ValueError(f"unknown payload_kind: {payload_kind!r}")
|
|
|
|
@staticmethod
|
|
def _validate_object_key(object_key: str) -> None:
|
|
if not object_key:
|
|
raise ValueError("object_key must be non-empty")
|
|
|
|
|
|
def require_cp_shared_l2_same_node(
|
|
pg: Any,
|
|
*,
|
|
source_rank: int = 0,
|
|
same_node_checker: Callable[..., Iterable[bool]] | None = None,
|
|
) -> list[bool]:
|
|
"""Fail fast unless every CP rank is local to ``source_rank``.
|
|
|
|
``same_node_checker`` is injectable for unit tests so importing this module
|
|
does not require torch.distributed initialization.
|
|
"""
|
|
|
|
if same_node_checker is None:
|
|
from sglang.srt.distributed.parallel_state import in_the_same_node_as
|
|
|
|
same_node_checker = in_the_same_node_as
|
|
status = list(same_node_checker(pg, source_rank=source_rank))
|
|
if not status or not all(status):
|
|
raise ValueError(
|
|
f"{CP_SHARED_L2_FAILFAST_PREFIX}[cross_node_cp_group_unsupported] "
|
|
"CP shared physical L2 requires all CP ranks on the same node"
|
|
)
|
|
return status
|
|
|
|
|
|
def _cp_group_local_rank_to_global_rank(
|
|
cp_cpu_group: Any, source_rank: int, *, dist_module: Any
|
|
) -> int:
|
|
"""Resolve CP-group-local source rank to the global rank PyTorch expects."""
|
|
|
|
source_rank = int(source_rank)
|
|
if cp_cpu_group is None:
|
|
return source_rank
|
|
get_global_rank = getattr(dist_module, "get_global_rank", None)
|
|
if get_global_rank is None:
|
|
_failfast(
|
|
"torch.distributed.get_global_rank is required to broadcast a "
|
|
"CP-group-local source rank safely"
|
|
)
|
|
return int(get_global_rank(cp_cpu_group, source_rank))
|
|
|
|
|
|
def _broadcast_object_list_with_cp_group_source(
|
|
object_list: list[Any],
|
|
*,
|
|
cp_cpu_group: Any,
|
|
source_rank: int,
|
|
dist_module: Any,
|
|
) -> None:
|
|
"""Broadcast with CP-local source semantics over a torch distributed group."""
|
|
|
|
global_source_rank = _cp_group_local_rank_to_global_rank(
|
|
cp_cpu_group, source_rank, dist_module=dist_module
|
|
)
|
|
dist_module.broadcast_object_list(
|
|
object_list, src=global_source_rank, group=cp_cpu_group
|
|
)
|
|
|
|
|
|
def broadcast_cp_shared_l2_object(
|
|
value: Any,
|
|
*,
|
|
cp_cpu_group: Any,
|
|
rank: int,
|
|
source_rank: int = 0,
|
|
broadcast_fn: Callable[[list[Any], int, Any], None] | None = None,
|
|
dist_module: Any | None = None,
|
|
) -> Any:
|
|
"""Broadcast one Python object over a CP CPU group behind an injectable seam."""
|
|
|
|
object_list = [value if int(rank) == int(source_rank) else None]
|
|
if broadcast_fn is not None:
|
|
# Preserve the existing injected-test seam: fakes receive the CP-local
|
|
# source rank they were written against. Only the production torch path
|
|
# maps CP-local rank to PyTorch's required global src rank.
|
|
broadcast_fn(object_list, src=source_rank, group=cp_cpu_group)
|
|
return object_list[0]
|
|
|
|
if dist_module is None:
|
|
import torch.distributed as dist
|
|
|
|
dist_module = dist
|
|
_broadcast_object_list_with_cp_group_source(
|
|
object_list,
|
|
cp_cpu_group=cp_cpu_group,
|
|
source_rank=source_rank,
|
|
dist_module=dist_module,
|
|
)
|
|
return object_list[0]
|
|
|
|
|
|
def broadcast_cp_shared_l2_host_slab_handle(
|
|
handle: CpSharedHostSlabHandle | None,
|
|
*,
|
|
cp_cpu_group: Any,
|
|
rank: int,
|
|
source_rank: int = 0,
|
|
broadcast_fn: Callable[[list[Any], int, Any], None] | None = None,
|
|
) -> CpSharedHostSlabHandle:
|
|
return broadcast_cp_shared_l2_object(
|
|
handle,
|
|
cp_cpu_group=cp_cpu_group,
|
|
rank=rank,
|
|
source_rank=source_rank,
|
|
broadcast_fn=broadcast_fn,
|
|
)
|
|
|
|
|
|
def broadcast_cp_shared_l2_object_range(
|
|
object_range: CpSharedL2ObjectRange | None,
|
|
*,
|
|
cp_cpu_group: Any,
|
|
rank: int,
|
|
source_rank: int = 0,
|
|
broadcast_fn: Callable[[list[Any], int, Any], None] | None = None,
|
|
) -> CpSharedL2ObjectRange:
|
|
return broadcast_cp_shared_l2_object(
|
|
object_range,
|
|
cp_cpu_group=cp_cpu_group,
|
|
rank=rank,
|
|
source_rank=source_rank,
|
|
broadcast_fn=broadcast_fn,
|
|
)
|
|
|
|
|
|
def gather_cp_shared_l2_object(
|
|
value: Any,
|
|
*,
|
|
cp_cpu_group: Any,
|
|
gather_fn: Callable[[list[Any], Any, Any], None] | None = None,
|
|
dist_module: Any | None = None,
|
|
) -> list[Any]:
|
|
"""Collect one Python object per CP rank behind an injectable seam."""
|
|
|
|
if gather_fn is not None:
|
|
output: list[Any] = []
|
|
gather_fn(output, value, cp_cpu_group)
|
|
return [item for item in output if item is not None]
|
|
|
|
if cp_cpu_group is None:
|
|
return [value]
|
|
|
|
if dist_module is None:
|
|
import torch.distributed as dist
|
|
|
|
dist_module = dist
|
|
world_size = int(dist_module.get_world_size(group=cp_cpu_group))
|
|
output = [None for _ in range(world_size)]
|
|
dist_module.all_gather_object(output, value, group=cp_cpu_group)
|
|
return [item for item in output if item is not None]
|
|
|
|
|
|
def _failfast(message: str):
|
|
raise ValueError(f"{CP_SHARED_L2_FAILFAST_PREFIX} {message}")
|
|
|
|
|
|
def validate_cp_shared_host_slab_handle_config(
|
|
handle: CpSharedHostSlabHandle | Any,
|
|
) -> None:
|
|
if handle.backend != CP_SHARED_L2_SLAB_BACKEND:
|
|
_failfast(
|
|
f"unsupported host slab backend {handle.backend!r}; "
|
|
f"expected {CP_SHARED_L2_SLAB_BACKEND!r}"
|
|
)
|
|
if handle.os_page_size != CP_SHARED_L2_OS_PAGE_SIZE:
|
|
_failfast(
|
|
f"host slab os_page_size must be 2MiB "
|
|
f"({CP_SHARED_L2_OS_PAGE_SIZE}); got {handle.os_page_size}"
|
|
)
|
|
if handle.numa_policy != CP_SHARED_L2_NUMA_POLICY:
|
|
_failfast(
|
|
f"unsupported host slab numa_policy {handle.numa_policy!r}; "
|
|
f"expected {CP_SHARED_L2_NUMA_POLICY!r}"
|
|
)
|
|
if handle.mapped_nbytes < handle.nbytes:
|
|
_failfast(
|
|
"host slab mapped_nbytes must cover logical nbytes: "
|
|
f"mapped_nbytes={handle.mapped_nbytes}, nbytes={handle.nbytes}"
|
|
)
|
|
if handle.mapped_nbytes % handle.os_page_size != 0:
|
|
_failfast(
|
|
"host slab mapped_nbytes must be aligned to 2MiB huge pages: "
|
|
f"mapped_nbytes={handle.mapped_nbytes}, page_size={handle.os_page_size}"
|
|
)
|
|
|
|
|
|
def round_up_cp_shared_l2_host_slab_bytes(
|
|
nbytes: int, *, page_size: int = CP_SHARED_L2_OS_PAGE_SIZE
|
|
) -> int:
|
|
nbytes = int(nbytes)
|
|
page_size = int(page_size)
|
|
if nbytes <= 0:
|
|
_failfast("host slab nbytes must be positive")
|
|
if page_size <= 0:
|
|
_failfast("host slab page_size must be positive")
|
|
return ((nbytes + page_size - 1) // page_size) * page_size
|
|
|
|
|
|
def validate_hugetlbfs_dir(
|
|
path: str | os.PathLike[str],
|
|
*,
|
|
production: bool,
|
|
mount_checker: Callable[[str], bool | str] | None = None,
|
|
) -> str:
|
|
resolved = str(Path(path))
|
|
if not production:
|
|
return resolved
|
|
|
|
path_obj = Path(resolved)
|
|
if not path_obj.exists():
|
|
_failfast(f"hugetlbfs directory {resolved!r} does not exist")
|
|
if not path_obj.is_dir():
|
|
_failfast(f"hugetlbfs path {resolved!r} is not a directory")
|
|
|
|
if mount_checker is not None:
|
|
mount_evidence = mount_checker(resolved)
|
|
if mount_evidence is False:
|
|
_failfast(f"hugetlbfs directory {resolved!r} is not a mount point")
|
|
if isinstance(mount_evidence, str):
|
|
if mount_evidence != "hugetlbfs":
|
|
_failfast(
|
|
f"hugetlbfs directory {resolved!r} has filesystem "
|
|
f"{mount_evidence!r}; expected 'hugetlbfs'"
|
|
)
|
|
return resolved
|
|
elif not os.path.ismount(resolved):
|
|
_failfast(f"hugetlbfs directory {resolved!r} is not a mount point")
|
|
|
|
fs_type = _find_linux_mount_fs_type(resolved)
|
|
if fs_type != "hugetlbfs":
|
|
if fs_type is None:
|
|
_failfast(
|
|
f"hugetlbfs directory {resolved!r} filesystem type could not be "
|
|
"verified as 'hugetlbfs'"
|
|
)
|
|
_failfast(
|
|
f"hugetlbfs directory {resolved!r} has filesystem {fs_type!r}; "
|
|
"expected 'hugetlbfs'"
|
|
)
|
|
return resolved
|
|
|
|
|
|
def _find_linux_mount_fs_type(path: str) -> str | None:
|
|
try:
|
|
mount_lines = Path("/proc/mounts").read_text().splitlines()
|
|
except OSError:
|
|
return None
|
|
|
|
real_path = os.path.realpath(path)
|
|
best_mount_point = ""
|
|
best_fs_type = None
|
|
for line in mount_lines:
|
|
fields = line.split()
|
|
if len(fields) < 3:
|
|
continue
|
|
mount_point = os.path.realpath(_decode_proc_mount_field(fields[1]))
|
|
if (
|
|
real_path == mount_point
|
|
or real_path.startswith(mount_point.rstrip(os.sep) + os.sep)
|
|
) and len(mount_point) > len(best_mount_point):
|
|
best_mount_point = mount_point
|
|
best_fs_type = fields[2]
|
|
return best_fs_type
|
|
|
|
|
|
def _decode_proc_mount_field(value: str) -> str:
|
|
def replace(match: re.Match[str]) -> str:
|
|
return chr(int(match.group(1), 8))
|
|
|
|
return re.sub(r"\\([0-7]{3})", replace, value)
|
|
|
|
|
|
def validate_effective_2m_page_mapping(
|
|
*,
|
|
smaps_text: str | None = None,
|
|
checker: Callable[[], bool | int] | Callable[[str], bool | int] | None = None,
|
|
path: str | None = None,
|
|
) -> int:
|
|
"""Validate effective mapped page size without requiring hugetlbfs in tests.
|
|
|
|
``checker`` may return True/False or a byte page size. ``smaps_text`` should
|
|
contain Linux smaps-like KernelPageSize/MMUPageSize lines using kB units.
|
|
"""
|
|
|
|
if checker is not None:
|
|
result = _call_page_size_checker(checker, path)
|
|
if result is True:
|
|
return CP_SHARED_L2_OS_PAGE_SIZE
|
|
if result is False:
|
|
_failfast("effective mapping page size is not 2MiB")
|
|
try:
|
|
page_size = int(result)
|
|
except (TypeError, ValueError):
|
|
_failfast(f"page-size checker returned unsupported value {result!r}")
|
|
if page_size != CP_SHARED_L2_OS_PAGE_SIZE:
|
|
_failfast(
|
|
f"effective mapping page size must be 2MiB "
|
|
f"({CP_SHARED_L2_OS_PAGE_SIZE}); got {page_size}"
|
|
)
|
|
return page_size
|
|
|
|
if smaps_text is None:
|
|
_failfast("no page-size evidence supplied for 2MiB validation")
|
|
|
|
page_sizes = _extract_smaps_page_sizes(smaps_text)
|
|
if not page_sizes:
|
|
_failfast("smaps data does not include KernelPageSize/MMUPageSize evidence")
|
|
bad = [size for size in page_sizes if size != CP_SHARED_L2_OS_PAGE_SIZE]
|
|
if bad:
|
|
_failfast(
|
|
f"effective mapping page size must be 2MiB "
|
|
f"({CP_SHARED_L2_OS_PAGE_SIZE}); got {bad[0]}"
|
|
)
|
|
return CP_SHARED_L2_OS_PAGE_SIZE
|
|
|
|
|
|
def validate_effective_2m_page_mapping_for_address(
|
|
*,
|
|
ptr: int,
|
|
nbytes: int | None = None,
|
|
checker: Callable[[], bool | int] | Callable[[str], bool | int] | None = None,
|
|
path: str | None = None,
|
|
) -> int:
|
|
"""Validate the effective page size for a live mapping containing ``ptr``.
|
|
|
|
hugetlbfs files do not have useful smaps evidence until the process has a
|
|
live, faulted/registered mapping. Production callers therefore use this
|
|
helper after ``cudaHostRegister`` succeeds and before the slab path is
|
|
early-unlinked. Tests may still inject ``checker`` to avoid depending on
|
|
Linux ``/proc/self/smaps``.
|
|
"""
|
|
|
|
if checker is not None:
|
|
return validate_effective_2m_page_mapping(checker=checker, path=path)
|
|
return validate_effective_2m_page_mapping(
|
|
smaps_text=_read_smaps_block_for_address(ptr=ptr, nbytes=nbytes)
|
|
)
|
|
|
|
|
|
def _read_smaps_block_for_address(*, ptr: int, nbytes: int | None = None) -> str:
|
|
ptr = int(ptr)
|
|
if ptr <= 0:
|
|
_failfast(f"mapping pointer must be positive for smaps validation; got {ptr}")
|
|
if nbytes is not None and int(nbytes) <= 0:
|
|
_failfast(f"mapping nbytes must be positive for smaps validation; got {nbytes}")
|
|
end_ptr = ptr + int(nbytes) if nbytes is not None else ptr + 1
|
|
|
|
try:
|
|
smaps_text = Path("/proc/self/smaps").read_text()
|
|
except OSError as exc:
|
|
_failfast(f"cannot read /proc/self/smaps for 2MiB validation: {exc}")
|
|
|
|
header_pattern = re.compile(r"^([0-9a-fA-F]+)-([0-9a-fA-F]+)\s")
|
|
current_header: tuple[int, int] | None = None
|
|
current_lines: list[str] = []
|
|
|
|
for line in smaps_text.splitlines():
|
|
header_match = header_pattern.match(line)
|
|
if header_match:
|
|
if current_header is not None:
|
|
start, end = current_header
|
|
if start <= ptr and end_ptr <= end:
|
|
return "\n".join(current_lines)
|
|
current_header = (
|
|
int(header_match.group(1), 16),
|
|
int(header_match.group(2), 16),
|
|
)
|
|
current_lines = [line]
|
|
elif current_header is not None:
|
|
current_lines.append(line)
|
|
|
|
if current_header is not None:
|
|
start, end = current_header
|
|
if start <= ptr and end_ptr <= end:
|
|
return "\n".join(current_lines)
|
|
|
|
_failfast(
|
|
"could not find a /proc/self/smaps mapping covering "
|
|
f"ptr={ptr:#x}, nbytes={nbytes}"
|
|
)
|
|
|
|
|
|
def _call_page_size_checker(
|
|
checker: Callable[[], bool | int] | Callable[[str], bool | int],
|
|
path: str | None,
|
|
) -> bool | int:
|
|
if path is None:
|
|
return checker() # type: ignore[misc]
|
|
|
|
try:
|
|
inspect.signature(checker).bind(path)
|
|
except (TypeError, ValueError):
|
|
return checker() # type: ignore[misc]
|
|
return checker(path) # type: ignore[misc]
|
|
|
|
|
|
def _extract_smaps_page_sizes(smaps_text: str) -> list[int]:
|
|
page_sizes: list[int] = []
|
|
pattern = re.compile(r"^\s*(?:KernelPageSize|MMUPageSize):\s*(\d+)\s*kB\s*$")
|
|
for line in smaps_text.splitlines():
|
|
match = pattern.match(line)
|
|
if match:
|
|
page_sizes.append(int(match.group(1)) * 1024)
|
|
return page_sizes
|
|
|
|
|
|
def create_cp_shared_host_slab(
|
|
*,
|
|
directory: str | os.PathLike[str],
|
|
name: str,
|
|
nbytes: int,
|
|
shape: tuple[int, ...] = (),
|
|
dtype_name: str = "uint8",
|
|
creator_rank: int = 0,
|
|
validate_production: bool = True,
|
|
page_size_checker: Callable[[], bool | int] | Callable[[str], bool | int] | None = None,
|
|
) -> CpSharedHostSlabMapping:
|
|
hugetlbfs_dir = validate_hugetlbfs_dir(directory, production=validate_production)
|
|
if nbytes <= 0:
|
|
_failfast("host slab nbytes must be positive")
|
|
if not name or Path(name).name != name:
|
|
_failfast("host slab name must be a single non-empty file name")
|
|
mapped_nbytes = round_up_cp_shared_l2_host_slab_bytes(nbytes)
|
|
|
|
path = str(Path(hugetlbfs_dir) / name)
|
|
flags = os.O_RDWR | os.O_CREAT | os.O_EXCL
|
|
fd = os.open(path, flags, 0o600)
|
|
file_obj = None
|
|
mapped = None
|
|
try:
|
|
os.ftruncate(fd, mapped_nbytes)
|
|
file_obj = os.fdopen(fd, "r+b", closefd=True)
|
|
fd = -1
|
|
mapped = mmap.mmap(file_obj.fileno(), mapped_nbytes)
|
|
handle = CpSharedHostSlabHandle(
|
|
name=name,
|
|
path=path,
|
|
nbytes=nbytes,
|
|
shape=tuple(shape),
|
|
dtype_name=dtype_name,
|
|
creator_rank=creator_rank,
|
|
hugetlbfs_dir=hugetlbfs_dir,
|
|
mapped_nbytes=mapped_nbytes,
|
|
)
|
|
return CpSharedHostSlabMapping(
|
|
handle=handle,
|
|
mmap=mapped,
|
|
_file=file_obj,
|
|
_unlink_on_close=True,
|
|
)
|
|
except Exception:
|
|
if mapped is not None:
|
|
mapped.close()
|
|
if file_obj is not None:
|
|
file_obj.close()
|
|
elif fd >= 0:
|
|
os.close(fd)
|
|
try:
|
|
os.unlink(path)
|
|
except FileNotFoundError:
|
|
pass
|
|
raise
|
|
|
|
|
|
def attach_cp_shared_host_slab(
|
|
handle: CpSharedHostSlabHandle,
|
|
*,
|
|
validate_production: bool = False,
|
|
page_size_checker: Callable[[], bool | int] | Callable[[str], bool | int] | None = None,
|
|
) -> CpSharedHostSlabMapping:
|
|
validate_cp_shared_host_slab_handle_config(handle)
|
|
if validate_production:
|
|
validate_hugetlbfs_dir(handle.hugetlbfs_dir, production=True)
|
|
if not Path(handle.path).exists():
|
|
_failfast(f"host slab path {handle.path!r} does not exist")
|
|
file_obj = open(handle.path, "r+b")
|
|
try:
|
|
mapped = mmap.mmap(file_obj.fileno(), handle.mapped_nbytes)
|
|
except Exception:
|
|
file_obj.close()
|
|
raise
|
|
return CpSharedHostSlabMapping(handle=handle, mmap=mapped, _file=file_obj)
|
|
|
|
|
|
def expected_cp_shared_l2_registration_range(
|
|
*, ptr: int, nbytes: int
|
|
) -> CpSharedL2RegistrationRange:
|
|
return CpSharedL2RegistrationRange(ptr=int(ptr), nbytes=int(nbytes))
|
|
|
|
|
|
def _known_tensor_like_nbytes(buffer: Any) -> int | None:
|
|
if hasattr(buffer, "numel") and hasattr(buffer, "element_size"):
|
|
return int(buffer.numel()) * int(buffer.element_size())
|
|
if hasattr(buffer, "nbytes"):
|
|
return int(buffer.nbytes)
|
|
return None
|
|
|
|
|
|
def cp_shared_l2_registration_range(
|
|
buffer: Any, nbytes: int | None = None
|
|
) -> CpSharedL2RegistrationRange:
|
|
if hasattr(buffer, "data_ptr"):
|
|
ptr = int(buffer.data_ptr())
|
|
known_nbytes = _known_tensor_like_nbytes(buffer)
|
|
if nbytes is None:
|
|
if known_nbytes is None:
|
|
_failfast("tensor-like buffer requires explicit registration nbytes")
|
|
nbytes = known_nbytes
|
|
else:
|
|
nbytes = int(nbytes)
|
|
if known_nbytes is not None and nbytes > known_nbytes:
|
|
_failfast(
|
|
f"registration nbytes {nbytes} exceeds buffer size "
|
|
f"{known_nbytes}"
|
|
)
|
|
return expected_cp_shared_l2_registration_range(ptr=ptr, nbytes=nbytes)
|
|
|
|
view = memoryview(buffer)
|
|
try:
|
|
size = int(nbytes) if nbytes is not None else view.nbytes
|
|
if size <= 0:
|
|
_failfast("registration nbytes must be positive")
|
|
if size > view.nbytes:
|
|
_failfast(
|
|
f"registration nbytes {size} exceeds buffer size {view.nbytes}"
|
|
)
|
|
if view.readonly:
|
|
_failfast("registration buffer must be writable to expose a stable address")
|
|
c_char_array = ctypes.c_char * view.nbytes
|
|
ptr = ctypes.addressof(c_char_array.from_buffer(view))
|
|
return expected_cp_shared_l2_registration_range(ptr=ptr, nbytes=size)
|
|
finally:
|
|
view.release()
|