diff --git a/python/sglang/srt/mem_cache/cp_shared_l2_pool.py b/python/sglang/srt/mem_cache/cp_shared_l2_pool.py new file mode 100644 index 000000000..b0119b624 --- /dev/null +++ b/python/sglang/srt/mem_cache/cp_shared_l2_pool.py @@ -0,0 +1,1779 @@ +# 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 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, ...] + committed_payload_layers: dict[str, set[int]] + committed: bool = False + 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}" + ) + + for payload_kind, layers in self.committed_payload_layers.items(): + self._validate_payload_kind(payload_kind) + if not isinstance(layers, set): + raise ValueError("committed_payload_layers values must be sets") + if any(layer < 0 for layer in layers): + raise ValueError("committed payload layer ids must be non-negative") + + @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, + ) + + committed_layers = { + payload_kind: set(layers) + for payload_kind, layers in self.committed_payload_layers.items() + } + 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), + committed_payload_layers={ + payload_kind: set(layers) + for payload_kind, layers in committed_layers.items() + }, + committed=self.committed, + 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), + committed_payload_layers={ + payload_kind: set(layers) + for payload_kind, layers in committed_layers.items() + }, + committed=self.committed, + 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, + 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: + 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) + + 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" + ) + + 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, + "cp_shared_l2_objects_evicted": 0, + } + + 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._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: + 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 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") + + 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(): + 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, + ) + + 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)) + + 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 ( + (parent_object_key, parent_ranges), + (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 + + 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 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 _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._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 + + 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 _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] + + @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 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, + *, + 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 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}") + + +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() diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index fd5cafba0..6ebc3d878 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -7,7 +7,7 @@ import weakref from collections import defaultdict from dataclasses import dataclass from functools import lru_cache, wraps -from typing import Optional, Tuple +from typing import Any, Optional, Sequence, Tuple import psutil import torch @@ -76,6 +76,15 @@ class PreparedLoadDescriptor: tai_index_h2d_descriptor: Optional[object] = None +@dataclass(frozen=True) +class HostTransferSegment: + """One physical host tensor segment for a logical host-index transfer.""" + + tensor: Any + host_indices: torch.Tensor + device_indices: torch.Tensor + + @lru_cache(maxsize=1) def _load_tai_transfer_kv_per_layer_mla_lf_pf(): try: @@ -228,6 +237,662 @@ class HostTensorAllocator(abc.ABC): return tensor +class SharedHostTensorAllocator(HostTensorAllocator): + """Host tensor allocator backed by a CP shared physical L2 host slab. + + One allocator instance owns exactly one slab mapping. Callers that need + multiple independently shared buffers (for example NSA main KV plus index + payload) must pass one allocator per buffer so each payload has distinct + handle/registration metadata. + """ + + def __init__( + self, + *, + directory: Optional[str] = None, + name: Optional[str] = None, + creator_rank: int = 0, + handle=None, + validate_production: bool = True, + create_fn=None, + attach_fn=None, + page_size_checker=None, + rank: int = 0, + source_rank: int = 0, + cp_cpu_group=None, + broadcast_handle_fn=None, + registration_status_gather_fn=None, + unlink_result_broadcast_fn=None, + ): + super().__init__() + if handle is None and (not directory or not name): + raise ValueError( + "SharedHostTensorAllocator requires either a slab handle or " + "both directory and name" + ) + if handle is not None and (directory is not None or name is not None): + raise ValueError( + "SharedHostTensorAllocator attach mode must not also specify " + "directory/name" + ) + self.directory = directory + self.name = name + self.creator_rank = int(creator_rank) + self.validate_production = bool(validate_production) + self._create_fn = create_fn + self._attach_fn = attach_fn + self._page_size_checker = page_size_checker + self.rank = int(rank) + self.source_rank = int(source_rank) + self.cp_cpu_group = cp_cpu_group + self._broadcast_handle_fn = broadcast_handle_fn + self._registration_status_gather_fn = registration_status_gather_fn + self._unlink_result_broadcast_fn = unlink_result_broadcast_fn + self.mapping = None + self.slab_handle = handle + self.registration_range = None + self.tensor = None + self._cuda_host_register_finalized = False + + @property + def handle(self): + return self.slab_handle + + def allocate(self, dims: tuple, dtype: torch.dtype, device: str) -> torch.Tensor: + if str(device) != "cpu": + raise ValueError( + "SharedHostTensorAllocator only supports CPU tensor views; " + f"got device={device!r}" + ) + if self.tensor is not None: + raise RuntimeError( + "SharedHostTensorAllocator instances are single-use; create a " + "separate allocator for each shared slab payload" + ) + + from sglang.srt.mem_cache.cp_shared_l2_pool import ( + attach_cp_shared_host_slab, + cp_shared_l2_registration_range, + create_cp_shared_host_slab, + ) + + dims = tuple(int(dim) for dim in dims) + numel = 1 + for dim in dims: + if dim < 0: + raise ValueError(f"shared host tensor dims must be non-negative: {dims}") + numel *= dim + element_size = torch.empty((), dtype=dtype).element_size() + nbytes = int(numel) * int(element_size) + if nbytes <= 0: + raise ValueError("shared host tensor allocation must be non-empty") + dtype_name = str(dtype).removeprefix("torch.") + + if self.slab_handle is None: + if self.cp_cpu_group is not None and self.rank != self.source_rank: + self.slab_handle = self._broadcast_handle(None) + else: + create_fn = self._create_fn or create_cp_shared_host_slab + self.mapping = create_fn( + directory=self.directory, + name=self.name, + nbytes=nbytes, + shape=dims, + dtype_name=dtype_name, + creator_rank=self.creator_rank, + validate_production=self.validate_production, + page_size_checker=self._page_size_checker, + ) + self.slab_handle = self.mapping.handle + if self.cp_cpu_group is not None: + self._broadcast_handle(self.slab_handle) + if self.mapping is None: + if tuple(self.slab_handle.shape) != dims: + raise ValueError( + "shared host slab shape mismatch: " + f"handle={tuple(self.slab_handle.shape)} requested={dims}" + ) + if self.slab_handle.dtype_name != dtype_name: + raise ValueError( + "shared host slab dtype mismatch: " + f"handle={self.slab_handle.dtype_name!r} requested={dtype_name!r}" + ) + if int(self.slab_handle.nbytes) != nbytes: + raise ValueError( + "shared host slab byte-size mismatch: " + f"handle={self.slab_handle.nbytes} requested={nbytes}" + ) + attach_fn = self._attach_fn or attach_cp_shared_host_slab + self.mapping = attach_fn( + self.slab_handle, + validate_production=self.validate_production, + page_size_checker=self._page_size_checker, + ) + + tensor = torch.frombuffer( + self.mapping.mmap, dtype=dtype, count=numel + ).reshape(dims) + self.dtype = dtype + self.dims = dims + self.tensor = tensor + self.registration_range = cp_shared_l2_registration_range(tensor, nbytes=nbytes) + return tensor + + def _broadcast_handle(self, handle): + from sglang.srt.mem_cache.cp_shared_l2_pool import ( + broadcast_cp_shared_l2_host_slab_handle, + ) + + broadcast_fn = ( + self._broadcast_handle_fn or broadcast_cp_shared_l2_host_slab_handle + ) + return broadcast_fn( + handle, + cp_cpu_group=self.cp_cpu_group, + rank=self.rank, + source_rank=self.source_rank, + ) + + def complete_cuda_host_register( + self, + *, + ptr: int, + nbytes: int, + success: bool = True, + error_message: Optional[str] = None, + ) -> None: + """Finalize shared-slab CUDA registration and early-unlink the creator path. + + All CP ranks call this after attempting ``cudaHostRegister`` for the + shared slab. The source rank unlinks the hugetlbfs path only after the + all-rank status gather confirms that every rank has attached and + registered successfully. Unlinking the path does not invalidate this + allocator's mmap/tensor view. + """ + + if self._cuda_host_register_finalized: + return + status = self._cuda_host_register_status( + ptr=ptr, nbytes=nbytes, success=success, error_message=error_message + ) + statuses = self._gather_cuda_host_register_status(status) + self._raise_cuda_host_register_failures(statuses) + self._early_unlink_after_cuda_host_register() + + def _cuda_host_register_status( + self, + *, + ptr: int, + nbytes: int, + success: bool = True, + error_message: Optional[str] = None, + ): + if self.slab_handle is None: + raise RuntimeError( + "SharedHostTensorAllocator cannot finalize CUDA registration " + "before allocate() creates or attaches a slab handle." + ) + + local_success = bool(success) + local_error = None if local_success else ( + error_message or "cudaHostRegister failed" + ) + if local_success and self.validate_production: + try: + from sglang.srt.mem_cache.cp_shared_l2_pool import ( + validate_effective_2m_page_mapping_for_address, + ) + + validate_effective_2m_page_mapping_for_address( + ptr=ptr, + nbytes=nbytes, + checker=self._page_size_checker, + path=self.slab_handle.path, + ) + except Exception as exc: + local_success = False + local_error = ( + "2MiB effective page-size validation failed after " + f"cudaHostRegister: {exc}" + ) + + return { + "rank": int(self.rank), + "success": local_success, + "error": local_error, + "ptr": int(ptr), + "nbytes": int(nbytes), + "path": self.slab_handle.path, + } + + @staticmethod + def _raise_cuda_host_register_failures(statuses) -> None: + failures = [item for item in statuses if not bool(item.get("success"))] + if failures: + raise RuntimeError( + "CP shared physical L2 cudaHostRegister failed before early unlink: " + + "; ".join( + f"rank {item.get('rank')}: {item.get('error') or 'unknown error'}" + for item in failures + ) + ) + + def _early_unlink_after_cuda_host_register(self) -> None: + if self._cuda_host_register_finalized: + return + unlink_result = None + if self.rank == self.source_rank: + try: + unlinked = False + if self.mapping is not None: + unlinked = bool(self.mapping.unlink()) + unlink_result = { + "success": True, + "rank": int(self.rank), + "path": self.slab_handle.path, + "unlinked": unlinked, + } + except Exception as exc: + unlink_result = { + "success": False, + "rank": int(self.rank), + "path": self.slab_handle.path, + "error": str(exc), + } + + unlink_result = self._broadcast_cuda_host_register_unlink_result(unlink_result) + if unlink_result is not None and not bool(unlink_result.get("success")): + raise RuntimeError( + "CP shared physical L2 early unlink failed on source rank " + f"{unlink_result.get('rank')}: {unlink_result.get('error')}" + ) + self._cuda_host_register_finalized = True + + def _gather_cuda_host_register_status(self, status): + if self._registration_status_gather_fn is not None: + output = [] + self._registration_status_gather_fn(output, status, self.cp_cpu_group) + return [item for item in output if item is not None] + if self.cp_cpu_group is None: + return [status] + from sglang.srt.mem_cache.cp_shared_l2_pool import gather_cp_shared_l2_object + + return gather_cp_shared_l2_object(status, cp_cpu_group=self.cp_cpu_group) + + def _broadcast_cuda_host_register_unlink_result(self, unlink_result): + if self.cp_cpu_group is None: + return unlink_result + from sglang.srt.mem_cache.cp_shared_l2_pool import broadcast_cp_shared_l2_object + + return broadcast_cp_shared_l2_object( + unlink_result, + cp_cpu_group=self.cp_cpu_group, + rank=self.rank, + source_rank=self.source_rank, + broadcast_fn=self._unlink_result_broadcast_fn, + ) + + def release_tensor_view_for_close(self): + """Drop this allocator's tensor reference before explicitly closing mmap. + + Callers must also release any external tensor views obtained from + ``allocate`` before calling ``close``. Accessing such views after this + method and close is invalid. + """ + + self.tensor = None + + def close(self): + if self.tensor is not None: + raise RuntimeError( + "SharedHostTensorAllocator.close() refuses to close an mmap " + "while its torch.frombuffer tensor view is still retained. " + "Release external tensor references, then call " + "release_tensor_view_for_close() before close()." + ) + mapping = self.mapping + self.mapping = None + if mapping is not None: + mapping.close() + + +@dataclass(frozen=True) +class SharedHostTensorSlabView: + """Tensor view and lifecycle owner for one shared host slab in a group.""" + + slab_info: Any + tensor: torch.Tensor + allocator: SharedHostTensorAllocator + + @property + def handle(self): + return self.allocator.handle + + @property + def registration_range(self): + return self.allocator.registration_range + + +class SharedHostTensorGroup: + """Collection of independently mapped CP shared-L2 host slab tensors.""" + + def __init__(self, slab_views: Sequence[SharedHostTensorSlabView]): + slab_views = tuple(slab_views) + if not slab_views: + raise ValueError("SharedHostTensorGroup requires at least one slab view") + view_by_slab_id = {int(view.slab_info.slab_id): view for view in slab_views} + if len(view_by_slab_id) != len(slab_views): + raise ValueError("SharedHostTensorGroup requires unique slab ids") + self._slab_views = slab_views + self._view_by_slab_id = view_by_slab_id + self._slab_infos = tuple(view.slab_info for view in slab_views) + self._allocators = tuple(view.allocator for view in slab_views) + self._released = False + self._closed = False + + def _ensure_tensor_views_live(self) -> None: + if self._closed: + raise RuntimeError( + "SharedHostTensorGroup tensor views are unavailable after close()." + ) + if self._released: + raise RuntimeError( + "SharedHostTensorGroup tensor views were released for close; " + "do not access tensors after release_tensor_view_for_close()." + ) + + @property + def slab_views(self): + self._ensure_tensor_views_live() + return self._slab_views + + @property + def slab_infos(self): + return self._slab_infos + + @property + def tensors(self): + self._ensure_tensor_views_live() + return tuple(view.tensor for view in self._slab_views) + + @property + def handles(self): + return tuple(allocator.handle for allocator in self._allocators) + + @property + def registration_ranges(self): + return tuple(allocator.registration_range for allocator in self._allocators) + + def tensor_for_slab(self, slab_id: int) -> torch.Tensor: + self._ensure_tensor_views_live() + try: + return self._view_by_slab_id[int(slab_id)].tensor + except KeyError as exc: + raise KeyError(f"unknown CP shared L2 slab_id: {slab_id}") from exc + + def complete_cuda_host_register(self, registration_results=None) -> None: + """Finalize registration/early-unlink for every slab in this group. + + The group validates/gathers registration status for all slabs first, then + unlinks paths. This prevents a later slab failure from partially + early-unlinking the group before the failure is known. + """ + + self._ensure_tensor_views_live() + result_by_slab_id = None + if registration_results is not None: + result_by_slab_id = { + int(result["slab_id"]): result for result in registration_results + } + gathered_statuses = [] + for view in self._slab_views: + tensor = view.tensor + allocator = view.allocator + if allocator._cuda_host_register_finalized: + continue + nbytes = int(tensor.numel()) * int(tensor.element_size()) + ptr = tensor.data_ptr() + success = True + error_message = None + if result_by_slab_id is not None: + result = result_by_slab_id[int(view.slab_info.slab_id)] + ptr = int(result["ptr"]) + nbytes = int(result["nbytes"]) + success = bool(result["success"]) + error_message = result.get("error") + status = allocator._cuda_host_register_status( + ptr=ptr, + nbytes=nbytes, + success=success, + error_message=error_message, + ) + gathered_statuses.append(allocator._gather_cuda_host_register_status(status)) + + for statuses in gathered_statuses: + SharedHostTensorAllocator._raise_cuda_host_register_failures(statuses) + + self._early_unlink_all_or_none_after_cuda_host_register() + + def _early_unlink_all_or_none_after_cuda_host_register(self) -> None: + pending_allocators = [ + allocator + for allocator in self._allocators + if not allocator._cuda_host_register_finalized + ] + if not pending_allocators: + return + + unlink_result = None + source_allocator = pending_allocators[0] + if source_allocator.rank == source_allocator.source_rank: + try: + from sglang.srt.mem_cache.cp_shared_l2_pool import ( + unlink_cp_shared_host_slab_mappings_all_or_none, + ) + + mappings = [ + allocator.mapping + for allocator in pending_allocators + if allocator.mapping is not None + ] + unlinked = unlink_cp_shared_host_slab_mappings_all_or_none(mappings) + unlink_result = { + "success": True, + "rank": int(source_allocator.rank), + "paths": [ + allocator.slab_handle.path for allocator in pending_allocators + ], + "unlinked": bool(unlinked), + } + except Exception as exc: + unlink_result = { + "success": False, + "rank": int(source_allocator.rank), + "paths": [ + allocator.slab_handle.path for allocator in pending_allocators + ], + "error": str(exc), + } + + unlink_result = source_allocator._broadcast_cuda_host_register_unlink_result( + unlink_result + ) + if unlink_result is not None and not bool(unlink_result.get("success")): + raise RuntimeError( + "CP shared physical L2 group early unlink failed on source rank " + f"{unlink_result.get('rank')}: {unlink_result.get('error')}" + ) + for allocator in pending_allocators: + allocator._cuda_host_register_finalized = True + + def release_tensor_view_for_close(self): + if self._released: + return + for allocator in self._allocators: + allocator.release_tensor_view_for_close() + self._slab_views = None + self._view_by_slab_id = None + self._released = True + + def close(self): + if self._closed: + return + if not self._released: + raise RuntimeError( + "SharedHostTensorGroup.close() refuses to close mmap-backed slabs " + "while group tensor views are still retained. Call " + "release_tensor_view_for_close() before close()." + ) + for allocator in self._allocators: + allocator.close() + self._closed = True + + +class SharedHostTensorGroupAllocator: + """Allocator for logical tensors split across multiple shared host slabs.""" + + def __init__( + self, + *, + slabs: Sequence[Any], + directory: Optional[str] = None, + name: Optional[str] = None, + creator_rank: int = 0, + validate_production: bool = True, + create_fn=None, + attach_fn=None, + page_size_checker=None, + rank: int = 0, + source_rank: int = 0, + cp_cpu_group=None, + broadcast_handle_fn=None, + registration_status_gather_fn=None, + unlink_result_broadcast_fn=None, + ): + if not slabs: + raise ValueError("SharedHostTensorGroupAllocator requires non-empty slabs") + if not directory or not name: + raise ValueError( + "SharedHostTensorGroupAllocator requires both directory and base name" + ) + self.slabs = tuple(slabs) + seen_slab_ids = set() + for slab in self.slabs: + slab_id = int(slab.slab_id) + if slab_id in seen_slab_ids: + raise ValueError(f"duplicate CP shared L2 slab_id: {slab_id}") + seen_slab_ids.add(slab_id) + self.allocators = tuple( + SharedHostTensorAllocator( + directory=directory, + name=f"{name}-slab{int(slab.slab_id)}", + creator_rank=creator_rank, + validate_production=validate_production, + create_fn=create_fn, + attach_fn=attach_fn, + page_size_checker=page_size_checker, + rank=rank, + source_rank=source_rank, + cp_cpu_group=cp_cpu_group, + broadcast_handle_fn=broadcast_handle_fn, + registration_status_gather_fn=registration_status_gather_fn, + unlink_result_broadcast_fn=unlink_result_broadcast_fn, + ) + for slab in self.slabs + ) + + @property + def slab_infos(self): + return self.slabs + + def allocate_group( + self, + dims: tuple, + dtype: torch.dtype, + device: str, + *, + page_dim: int, + pin_memory: bool = False, + ) -> SharedHostTensorGroup: + if str(device) != "cpu": + raise ValueError( + "SharedHostTensorGroupAllocator only supports CPU tensor views; " + f"got device={device!r}" + ) + dims = tuple(int(dim) for dim in dims) + if not dims: + raise ValueError("shared host tensor group dims must be non-empty") + if page_dim < 0: + page_dim += len(dims) + if page_dim < 0 or page_dim >= len(dims): + raise ValueError( + f"page_dim must be a valid dimension for dims={dims}; got {page_dim}" + ) + for dim in dims: + if dim < 0: + raise ValueError(f"shared host tensor dims must be non-negative: {dims}") + total_pages = sum(int(slab.num_pages) for slab in self.slabs) + if total_pages != dims[page_dim]: + raise ValueError( + "shared host tensor group page_dim size must equal total slab pages: " + f"dims[{page_dim}]={dims[page_dim]} total_slab_pages={total_pages}" + ) + + slab_views = [] + for slab, allocator in zip(self.slabs, self.allocators): + slab_dims = list(dims) + slab_dims[page_dim] = int(slab.num_pages) + tensor = allocator.allocate(tuple(slab_dims), dtype=dtype, device=device) + slab_views.append( + SharedHostTensorSlabView( + slab_info=slab, tensor=tensor, allocator=allocator + ) + ) + group = SharedHostTensorGroup(slab_views) + if pin_memory: + self._register_group_tensors(group) + return group + + @staticmethod + def _register_group_tensors(group: SharedHostTensorGroup) -> None: + registration_results = [] + for view in group.slab_views: + tensor = view.tensor + ptr = tensor.data_ptr() + nbytes = int(tensor.numel()) * int(tensor.element_size()) + registered = False + register_error = None + try: + _check_torch_cudart( + torch.cuda.cudart().cudaHostRegister(ptr, nbytes, 0), + "cudaHostRegister", + ) + registered = True + except Exception as exc: + register_error = exc + registration_results.append( + { + "slab_id": int(view.slab_info.slab_id), + "ptr": ptr, + "nbytes": nbytes, + "success": registered, + "error": str(register_error) if register_error else None, + } + ) + + try: + group.complete_cuda_host_register(registration_results=registration_results) + except Exception: + for result in registration_results: + if result["success"]: + _cuda_host_unregister(int(result["ptr"])) + raise + + for view, result in zip(group.slab_views, registration_results): + if result["success"]: + weakref.finalize(view.tensor, _cuda_host_unregister, int(result["ptr"])) + + def get_allocator_from_storage(allocator_type): if allocator_type == "mooncake": try: @@ -261,12 +926,37 @@ def alloc_with_host_register( buffer = allocator.allocate(dims, dtype=dtype, device=device) if pin_memory: ptr = buffer.data_ptr() - _check_torch_cudart( - torch.cuda.cudart().cudaHostRegister( - ptr, buffer.numel() * buffer.element_size(), 0 - ), - "cudaHostRegister", - ) + nbytes = buffer.numel() * buffer.element_size() + registered = False + register_error = None + try: + _check_torch_cudart( + torch.cuda.cudart().cudaHostRegister(ptr, nbytes, 0), + "cudaHostRegister", + ) + registered = True + except Exception as exc: + register_error = exc + + finalize_error = None + finalize_register = getattr(allocator, "complete_cuda_host_register", None) + if callable(finalize_register): + try: + finalize_register( + ptr=ptr, + nbytes=nbytes, + success=registered, + error_message=str(register_error) if register_error else None, + ) + except Exception as exc: + finalize_error = exc + + if register_error is not None: + raise register_error + if finalize_error is not None: + if registered: + _cuda_host_unregister(ptr) + raise finalize_error weakref.finalize(buffer, _cuda_host_unregister, ptr) return buffer @@ -327,13 +1017,16 @@ class HostKVCache(abc.ABC): device: str, allocator_type: str = "default", host_token_capacity: Optional[int] = None, + host_tensor_allocator: Optional[HostTensorAllocator] = None, ): self.device_pool = device_pool self.page_size = page_size self.layout = layout self.pin_memory = pin_memory self.device = device - self.allocator = get_allocator_from_storage(allocator_type) + self.allocator = host_tensor_allocator or get_allocator_from_storage( + allocator_type + ) self.dtype = device_pool.store_dtype self.size_per_token = self.get_size_per_token() @@ -507,11 +1200,14 @@ class HostKVCache(abc.ABC): host_indices: torch.Tensor, device_indices: torch.Tensor, page_size: int, + host_buffer: Optional[Any] = None, ) -> Optional[object]: if descriptor.io_backend != "direct": return None if descriptor.layout not in ("page_first_direct", "layer_page_first"): return None + if self._is_group_backed_buffer(self.kv_buffer if host_buffer is None else host_buffer): + return None try: prepare = _load_tai_prepare_h2d_page_descriptor() except RuntimeError as exc: @@ -538,6 +1234,133 @@ class HostKVCache(abc.ABC): f"page_size={page_size} tokens={int(host_indices.numel())}" ) from exc + @staticmethod + def _is_group_backed_buffer(buffer: Any) -> bool: + return isinstance(buffer, SharedHostTensorGroup) + + def _alloc_host_tensor( + self, + dims: tuple, + *, + page_dim: Optional[int] = None, + dtype: Optional[torch.dtype] = None, + allocator: Optional[HostTensorAllocator] = None, + ): + dtype = self.dtype if dtype is None else dtype + allocator = self.allocator if allocator is None else allocator + if isinstance(allocator, SharedHostTensorGroupAllocator): + if self.layout not in ("page_first_direct", "layer_page_first"): + raise RuntimeError( + "[CP_HICACHE_FAILFAST][shared_host_group_unsupported_layout] " + f"SharedHostTensorGroupAllocator requires a page-first direct " + f"layout; got layout={self.layout!r}" + ) + if page_dim is None: + raise RuntimeError( + "[CP_HICACHE_FAILFAST][shared_host_group_missing_page_dim] " + "SharedHostTensorGroupAllocator allocation requires page_dim" + ) + return allocator.allocate_group( + dims, + dtype=dtype, + device=self.device, + page_dim=page_dim, + pin_memory=self.pin_memory, + ) + + alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] + return alloc_func( + dims, + dtype=dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=allocator, + ) + + def _host_transfer_segments( + self, + host_buffer: Any, + host_indices: torch.Tensor, + device_indices: torch.Tensor, + *, + page_size: int, + ) -> list[HostTransferSegment]: + if not isinstance(host_buffer, SharedHostTensorGroup): + return [HostTransferSegment(host_buffer, host_indices, device_indices)] + + page_size = int(page_size) + if host_indices.numel() == 0: + return [] + if not host_indices.is_cuda: + validate_page_aligned_token_indices( + host_indices, page_size, "host_indices" + ) + if not device_indices.is_cuda: + validate_page_aligned_token_indices( + device_indices, page_size, "device_indices" + ) + + host_pages = ( + host_indices.reshape(-1, page_size)[:, 0] // page_size + ).contiguous() + slab_views = sorted( + host_buffer.slab_views, + key=lambda view: int(view.slab_info.global_base_page), + ) + + def find_view(global_page: int) -> SharedHostTensorSlabView: + for view in slab_views: + base = int(view.slab_info.global_base_page) + end = base + int(view.slab_info.num_pages) + if base <= global_page < end: + return view + slab_ranges = [ + (int(v.slab_info.global_base_page), int(v.slab_info.num_pages)) + for v in slab_views + ] + raise RuntimeError( + "[CP_HICACHE_FAILFAST][shared_host_group_page_uncovered] " + f"global_page={global_page} page_size={page_size} is not covered " + f"by group slabs={slab_ranges}" + ) + + segments = [] + current_view = None + current_host_chunks = [] + current_device_chunks = [] + for page_offset, global_page in enumerate(host_pages.detach().cpu().tolist()): + view = find_view(int(global_page)) + token_start = page_offset * page_size + token_end = token_start + page_size + base_token = int(view.slab_info.global_base_page) * page_size + local_host_chunk = ( + host_indices[token_start:token_end] - base_token + ).contiguous() + device_chunk = device_indices[token_start:token_end].contiguous() + if current_view is not None and view is not current_view: + segments.append( + HostTransferSegment( + current_view.tensor, + torch.cat(current_host_chunks).contiguous(), + torch.cat(current_device_chunks).contiguous(), + ) + ) + current_host_chunks = [] + current_device_chunks = [] + current_view = view + current_host_chunks.append(local_host_chunk) + current_device_chunks.append(device_chunk) + + if current_view is not None: + segments.append( + HostTransferSegment( + current_view.tensor, + torch.cat(current_host_chunks).contiguous(), + torch.cat(current_device_chunks).contiguous(), + ) + ) + return segments + def _log_missing_tai_prepared_h2d_once(self, message: str) -> None: if getattr(self, "_missing_tai_prepared_h2d_warning_emitted", False): return @@ -999,6 +1822,7 @@ class MHATokenToKVPoolHost(HostKVCache): device: str = "cpu", allocator_type: str = "default", host_token_capacity: Optional[int] = None, + host_tensor_allocator: Optional[HostTensorAllocator] = None, ): super().__init__( device_pool, @@ -1010,24 +1834,26 @@ class MHATokenToKVPoolHost(HostKVCache): device, allocator_type, host_token_capacity=host_token_capacity, + host_tensor_allocator=host_tensor_allocator, ) self.element_dim = self.device_pool.head_num * self.device_pool.head_dim self.can_use_jit = _is_cuda and can_use_hicache_jit_kernel( element_size=self.element_dim * self.dtype.itemsize ) - self.k_data_refs = [self.k_buffer[i] for i in range(self.layer_num)] - self.v_data_refs = [self.v_buffer[i] for i in range(self.layer_num)] - self.k_data_ptrs = torch.tensor( - [x.data_ptr() for x in self.k_data_refs], - dtype=torch.uint64, - device=self.device_pool.device, - ) - self.v_data_ptrs = torch.tensor( - [x.data_ptr() for x in self.v_data_refs], - dtype=torch.uint64, - device=self.device_pool.device, - ) + if not self._is_group_backed_buffer(self.kv_buffer): + self.k_data_refs = [self.k_buffer[i] for i in range(self.layer_num)] + self.v_data_refs = [self.v_buffer[i] for i in range(self.layer_num)] + self.k_data_ptrs = torch.tensor( + [x.data_ptr() for x in self.k_data_refs], + dtype=torch.uint64, + device=self.device_pool.device, + ) + self.v_data_ptrs = torch.tensor( + [x.data_ptr() for x in self.v_data_refs], + dtype=torch.uint64, + device=self.device_pool.device, + ) def get_size_per_token(self): self.head_num = self.device_pool.head_num @@ -1076,15 +1902,12 @@ class MHATokenToKVPoolHost(HostKVCache): self.token_stride_size = self.head_num * self.head_dim * self.dtype.itemsize self.layout_dim = self.token_stride_size * self.layer_num - alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] - buffer = alloc_func( - dims, - dtype=self.dtype, - device=self.device, - pin_memory=self.pin_memory, - allocator=self.allocator, - ) - return buffer + page_dim = None + if self.layout == "page_first_direct": + page_dim = 1 + elif self.layout == "layer_page_first": + page_dim = 2 + return self._alloc_host_tensor(dims, page_dim=page_dim) @property def k_buffer(self): @@ -1165,29 +1988,43 @@ class MHATokenToKVPoolHost(HostKVCache): page_size=self.page_size, ) elif self.layout == "page_first_direct": - _load_tai_transfer_kv_per_layer_direct_pf_lf()( - src_ptrs=[self.k_buffer, self.v_buffer], - dst_ptrs=[ - device_pool.k_buffer[layer_id], - device_pool.v_buffer[layer_id], - ], - src_indices=host_indices, - dst_indices=device_indices, - layer_id=layer_id, + tai_transfer = _load_tai_transfer_kv_per_layer_direct_pf_lf() + for segment in self._host_transfer_segments( + self.kv_buffer, + host_indices, + device_indices, page_size=self.page_size, - ) + ): + tai_transfer( + src_ptrs=[segment.tensor[0], segment.tensor[1]], + dst_ptrs=[ + device_pool.k_buffer[layer_id], + device_pool.v_buffer[layer_id], + ], + src_indices=segment.host_indices, + dst_indices=segment.device_indices, + layer_id=layer_id, + page_size=self.page_size, + ) elif self.layout == "layer_page_first": - _load_tai_transfer_kv_per_layer_direct_lpf_lf()( - src_ptrs=[self.k_buffer, self.v_buffer], - dst_ptrs=[ - device_pool.k_buffer[layer_id], - device_pool.v_buffer[layer_id], - ], - src_indices=host_indices, - dst_indices=device_indices, - layer_id=layer_id, + tai_transfer = _load_tai_transfer_kv_per_layer_direct_lpf_lf() + for segment in self._host_transfer_segments( + self.kv_buffer, + host_indices, + device_indices, page_size=self.page_size, - ) + ): + tai_transfer( + src_ptrs=[segment.tensor[0], segment.tensor[1]], + dst_ptrs=[ + device_pool.k_buffer[layer_id], + device_pool.v_buffer[layer_id], + ], + src_indices=segment.host_indices, + dst_indices=segment.device_indices, + layer_id=layer_id, + page_size=self.page_size, + ) else: raise ValueError(f"Unsupported layout: {self.layout}") elif io_backend == "kernel_ascend": @@ -1281,31 +2118,43 @@ class MHATokenToKVPoolHost(HostKVCache): # per-layer direct op, which owns the version-specific ABI. tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_pf() for layer_id in range(self.layer_num): - tai_transfer( - src_ptrs=[ - device_pool.k_buffer[layer_id], - device_pool.v_buffer[layer_id], - ], - dst_ptrs=[self.k_buffer, self.v_buffer], - src_indices=device_indices, - dst_indices=host_indices, - layer_id=layer_id, + for segment in self._host_transfer_segments( + self.kv_buffer, + host_indices, + device_indices, page_size=self.page_size, - ) + ): + tai_transfer( + src_ptrs=[ + device_pool.k_buffer[layer_id], + device_pool.v_buffer[layer_id], + ], + dst_ptrs=[segment.tensor[0], segment.tensor[1]], + src_indices=segment.device_indices, + dst_indices=segment.host_indices, + layer_id=layer_id, + page_size=self.page_size, + ) elif self.layout == "layer_page_first": tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_lpf() for layer_id in range(self.layer_num): - tai_transfer( - src_ptrs=[ - device_pool.k_buffer[layer_id], - device_pool.v_buffer[layer_id], - ], - dst_ptrs=[self.k_buffer, self.v_buffer], - src_indices=device_indices, - dst_indices=host_indices, - layer_id=layer_id, + for segment in self._host_transfer_segments( + self.kv_buffer, + host_indices, + device_indices, page_size=self.page_size, - ) + ): + tai_transfer( + src_ptrs=[ + device_pool.k_buffer[layer_id], + device_pool.v_buffer[layer_id], + ], + dst_ptrs=[segment.tensor[0], segment.tensor[1]], + src_indices=segment.device_indices, + dst_indices=segment.host_indices, + layer_id=layer_id, + page_size=self.page_size, + ) else: raise ValueError(f"Unsupported layout: {self.layout}") elif io_backend == "kernel_ascend": @@ -1365,29 +2214,43 @@ class MHATokenToKVPoolHost(HostKVCache): page_size=self.page_size, ) elif self.layout == "page_first_direct": - _load_tai_transfer_kv_per_layer_direct_lf_pf()( - src_ptrs=[ - device_pool.k_buffer[layer_id], - device_pool.v_buffer[layer_id], - ], - dst_ptrs=[self.k_buffer, self.v_buffer], - src_indices=device_indices, - dst_indices=host_indices, - layer_id=layer_id, + tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_pf() + for segment in self._host_transfer_segments( + self.kv_buffer, + host_indices, + device_indices, page_size=self.page_size, - ) + ): + tai_transfer( + src_ptrs=[ + device_pool.k_buffer[layer_id], + device_pool.v_buffer[layer_id], + ], + dst_ptrs=[segment.tensor[0], segment.tensor[1]], + src_indices=segment.device_indices, + dst_indices=segment.host_indices, + layer_id=layer_id, + page_size=self.page_size, + ) elif self.layout == "layer_page_first": - _load_tai_transfer_kv_per_layer_direct_lf_lpf()( - src_ptrs=[ - device_pool.k_buffer[layer_id], - device_pool.v_buffer[layer_id], - ], - dst_ptrs=[self.k_buffer, self.v_buffer], - src_indices=device_indices, - dst_indices=host_indices, - layer_id=layer_id, + tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_lpf() + for segment in self._host_transfer_segments( + self.kv_buffer, + host_indices, + device_indices, page_size=self.page_size, - ) + ): + tai_transfer( + src_ptrs=[ + device_pool.k_buffer[layer_id], + device_pool.v_buffer[layer_id], + ], + dst_ptrs=[segment.tensor[0], segment.tensor[1]], + src_indices=segment.device_indices, + dst_indices=segment.host_indices, + layer_id=layer_id, + page_size=self.page_size, + ) else: raise ValueError(f"Unsupported layout: {self.layout}") elif io_backend == "kernel_ascend": @@ -1602,6 +2465,7 @@ class MLATokenToKVPoolHost(HostKVCache): allocator_type: str = "default", override_kv_cache_dim: Optional[int] = None, host_token_capacity: Optional[int] = None, + host_tensor_allocator: Optional[HostTensorAllocator] = None, ): self.override_kv_cache_dim = override_kv_cache_dim super().__init__( @@ -1614,13 +2478,15 @@ class MLATokenToKVPoolHost(HostKVCache): device, allocator_type, host_token_capacity=host_token_capacity, + host_tensor_allocator=host_tensor_allocator, ) - self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)] - self.data_ptrs = torch.tensor( - [x.data_ptr() for x in self.data_refs], - dtype=torch.uint64, - device=self.device_pool.device, - ) + if not self._is_group_backed_buffer(self.kv_buffer): + self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)] + self.data_ptrs = torch.tensor( + [x.data_ptr() for x in self.data_refs], + dtype=torch.uint64, + device=self.device_pool.device, + ) def get_size_per_token(self): self.kv_lora_rank = self.device_pool.kv_lora_rank @@ -1706,15 +2572,12 @@ class MLATokenToKVPoolHost(HostKVCache): self.token_stride_size = self.kv_cache_dim * self.dtype.itemsize self.layout_dim = self.token_stride_size * self.layer_num - alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] - buffer = alloc_func( - dims, - dtype=self.dtype, - device=self.device, - pin_memory=self.pin_memory, - allocator=self.allocator, - ) - return buffer + page_dim = None + if self.layout == "page_first_direct": + page_dim = 0 + elif self.layout == "layer_page_first": + page_dim = 1 + return self._alloc_host_tensor(dims, page_dim=page_dim) def load_to_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend @@ -1762,14 +2625,21 @@ class MLATokenToKVPoolHost(HostKVCache): layer_id=layer_id, ) return - _load_tai_transfer_kv_per_layer_direct_pf_lf()( - src_ptrs=[self.kv_buffer], - dst_ptrs=[device_pool.kv_buffer[layer_id]], - src_indices=host_indices, - dst_indices=device_indices, - layer_id=layer_id, + tai_transfer = _load_tai_transfer_kv_per_layer_direct_pf_lf() + for segment in self._host_transfer_segments( + self.kv_buffer, + host_indices, + device_indices, page_size=self.page_size, - ) + ): + tai_transfer( + src_ptrs=[segment.tensor], + dst_ptrs=[device_pool.kv_buffer[layer_id]], + src_indices=segment.host_indices, + dst_indices=segment.device_indices, + layer_id=layer_id, + page_size=self.page_size, + ) elif self.layout == "layer_page_first": if descriptor is not None and descriptor.tai_h2d_descriptor is not None: _load_tai_submit_h2d_layer()( @@ -1779,14 +2649,21 @@ class MLATokenToKVPoolHost(HostKVCache): layer_id=layer_id, ) return - _load_tai_transfer_kv_per_layer_direct_lpf_lf()( - src_ptrs=[self.kv_buffer], - dst_ptrs=[device_pool.kv_buffer[layer_id]], - src_indices=host_indices, - dst_indices=device_indices, - layer_id=layer_id, + tai_transfer = _load_tai_transfer_kv_per_layer_direct_lpf_lf() + for segment in self._host_transfer_segments( + self.kv_buffer, + host_indices, + device_indices, page_size=self.page_size, - ) + ): + tai_transfer( + src_ptrs=[segment.tensor], + dst_ptrs=[device_pool.kv_buffer[layer_id]], + src_indices=segment.host_indices, + dst_indices=segment.device_indices, + layer_id=layer_id, + page_size=self.page_size, + ) else: raise ValueError(f"Unsupported layout: {self.layout}") elif io_backend == "kernel_ascend": @@ -1851,25 +2728,37 @@ class MLATokenToKVPoolHost(HostKVCache): # per-layer direct op, which owns the version-specific ABI. tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_pf() for layer_id in range(self.layer_num): - tai_transfer( - src_ptrs=[device_pool.kv_buffer[layer_id]], - dst_ptrs=[self.kv_buffer], - src_indices=device_indices, - dst_indices=host_indices, - layer_id=layer_id, + for segment in self._host_transfer_segments( + self.kv_buffer, + host_indices, + device_indices, page_size=self.page_size, - ) + ): + tai_transfer( + src_ptrs=[device_pool.kv_buffer[layer_id]], + dst_ptrs=[segment.tensor], + src_indices=segment.device_indices, + dst_indices=segment.host_indices, + layer_id=layer_id, + page_size=self.page_size, + ) elif self.layout == "layer_page_first": tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_lpf() for layer_id in range(self.layer_num): - tai_transfer( - src_ptrs=[device_pool.kv_buffer[layer_id]], - dst_ptrs=[self.kv_buffer], - src_indices=device_indices, - dst_indices=host_indices, - layer_id=layer_id, + for segment in self._host_transfer_segments( + self.kv_buffer, + host_indices, + device_indices, page_size=self.page_size, - ) + ): + tai_transfer( + src_ptrs=[device_pool.kv_buffer[layer_id]], + dst_ptrs=[segment.tensor], + src_indices=segment.device_indices, + dst_indices=segment.host_indices, + layer_id=layer_id, + page_size=self.page_size, + ) else: raise ValueError(f"Unsupported layout: {self.layout}") elif io_backend == "kernel_ascend": @@ -1930,23 +2819,37 @@ class MLATokenToKVPoolHost(HostKVCache): page_size=self.page_size, ) elif self.layout == "page_first_direct": - _load_tai_transfer_kv_per_layer_direct_lf_pf()( - src_ptrs=[device_pool.kv_buffer[layer_id]], - dst_ptrs=[self.kv_buffer], - src_indices=device_indices, - dst_indices=host_indices, - layer_id=layer_id, + tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_pf() + for segment in self._host_transfer_segments( + self.kv_buffer, + host_indices, + device_indices, page_size=self.page_size, - ) + ): + tai_transfer( + src_ptrs=[device_pool.kv_buffer[layer_id]], + dst_ptrs=[segment.tensor], + src_indices=segment.device_indices, + dst_indices=segment.host_indices, + layer_id=layer_id, + page_size=self.page_size, + ) elif self.layout == "layer_page_first": - _load_tai_transfer_kv_per_layer_direct_lf_lpf()( - src_ptrs=[device_pool.kv_buffer[layer_id]], - dst_ptrs=[self.kv_buffer], - src_indices=device_indices, - dst_indices=host_indices, - layer_id=layer_id, + tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_lpf() + for segment in self._host_transfer_segments( + self.kv_buffer, + host_indices, + device_indices, page_size=self.page_size, - ) + ): + tai_transfer( + src_ptrs=[device_pool.kv_buffer[layer_id]], + dst_ptrs=[segment.tensor], + src_indices=segment.device_indices, + dst_indices=segment.host_indices, + layer_id=layer_id, + page_size=self.page_size, + ) else: raise ValueError(f"Unsupported layout: {self.layout}") elif io_backend == "kernel_ascend": @@ -2092,8 +2995,11 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): device: str = "cpu", allocator_type: str = "default", host_token_capacity: Optional[int] = None, + host_tensor_allocator: Optional[HostTensorAllocator] = None, + index_host_tensor_allocator: Optional[HostTensorAllocator] = None, ): # Initialize indexer metadata before HostKVCache.__init__ calls get_size_per_token. + self.index_host_tensor_allocator = index_host_tensor_allocator self.index_head_dim = device_pool.index_head_dim self.indexer_quant_block_size = device_pool.quant_block_size self.indexer_dtype = NSATokenToKVPool.index_k_with_scale_buffer_dtype @@ -2124,6 +3030,7 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): allocator_type, override_kv_cache_dim=device_pool.kv_cache_dim, host_token_capacity=host_token_capacity, + host_tensor_allocator=host_tensor_allocator, ) self.indexer_page_stride_size = ( self.indexer_size_per_token * self.page_size * self.indexer_dtype.itemsize @@ -2148,6 +3055,16 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): def _init_indexer_buffers(self): alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] + index_allocator = self.index_host_tensor_allocator or self.allocator + if isinstance(index_allocator, SharedHostTensorGroupAllocator) and self.layout not in ( + "page_first_direct", + "layer_page_first", + ): + raise RuntimeError( + "[CP_HICACHE_FAILFAST][shared_host_group_unsupported_layout] " + "SharedHostTensorGroupAllocator index buffer requires " + "page_first_direct or layer_page_first layout" + ) self.index_k_device_ptrs = torch.tensor( [x.data_ptr() for x in self.device_pool.index_k_with_scale_buffer], dtype=torch.uint64, @@ -2160,7 +3077,7 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): dtype=self.indexer_dtype, device=self.device, pin_memory=self.pin_memory, - allocator=self.allocator, + allocator=index_allocator, ) for _ in range(self.index_active_layer_num) ] @@ -2174,40 +3091,72 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): device=self.device_pool.device, ) elif self.layout in ["page_first", "page_first_direct"]: - self.index_k_with_scale_buffer = alloc_func( - ( - self.indexer_page_num, - self.index_active_layer_num, - 1, - self.indexer_page_stride_size, - ), - dtype=self.indexer_dtype, - device=self.device, - pin_memory=self.pin_memory, - allocator=self.allocator, - ) + if isinstance(index_allocator, SharedHostTensorGroupAllocator): + if self.layout != "page_first_direct": + raise RuntimeError( + "[CP_HICACHE_FAILFAST][shared_host_group_unsupported_layout] " + "SharedHostTensorGroupAllocator index buffer requires " + "page_first_direct or layer_page_first layout" + ) + self.index_k_with_scale_buffer = self._alloc_host_tensor( + ( + self.indexer_page_num, + self.index_active_layer_num, + 1, + self.indexer_page_stride_size, + ), + page_dim=0, + dtype=self.indexer_dtype, + allocator=index_allocator, + ) + else: + self.index_k_with_scale_buffer = alloc_func( + ( + self.indexer_page_num, + self.index_active_layer_num, + 1, + self.indexer_page_stride_size, + ), + dtype=self.indexer_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=index_allocator, + ) elif self.layout == "layer_page_first": - self.index_k_with_scale_buffer = alloc_func( - ( - self.index_active_layer_num, - self.indexer_page_num, - 1, - self.indexer_page_stride_size, - ), - dtype=self.indexer_dtype, - device=self.device, - pin_memory=self.pin_memory, - allocator=self.allocator, - ) - self.index_k_data_refs = [ - self.index_k_with_scale_buffer[i] - for i in range(self.index_active_layer_num) - ] - self.index_k_data_ptrs = torch.tensor( - [x.data_ptr() for x in self.index_k_data_refs], - dtype=torch.uint64, - device=self.device_pool.device, - ) + if isinstance(index_allocator, SharedHostTensorGroupAllocator): + self.index_k_with_scale_buffer = self._alloc_host_tensor( + ( + self.index_active_layer_num, + self.indexer_page_num, + 1, + self.indexer_page_stride_size, + ), + page_dim=1, + dtype=self.indexer_dtype, + allocator=index_allocator, + ) + else: + self.index_k_with_scale_buffer = alloc_func( + ( + self.index_active_layer_num, + self.indexer_page_num, + 1, + self.indexer_page_stride_size, + ), + dtype=self.indexer_dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=index_allocator, + ) + self.index_k_data_refs = [ + self.index_k_with_scale_buffer[i] + for i in range(self.index_active_layer_num) + ] + self.index_k_data_ptrs = torch.tensor( + [x.data_ptr() for x in self.index_k_data_refs], + dtype=torch.uint64, + device=self.device_pool.device, + ) else: raise ValueError(f"Unsupported layout: {self.layout}") @@ -2295,6 +3244,7 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): host_indices=host_page_indices, device_indices=device_page_indices, page_size=1, + host_buffer=self.index_k_with_scale_buffer, ) self._active_load_indexer_page_indices = ( descriptor.index_host_page_indices, @@ -2377,16 +3327,23 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): layer_id=host_layer_slot, ) return - _load_tai_transfer_kv_per_layer_direct_pf_lf()( - src_ptrs=[self.index_k_with_scale_buffer], - dst_ptrs=[ - device_pool.index_k_with_scale_buffer[device_layer_slot] - ], - src_indices=host_page_indices, - dst_indices=device_page_indices, - layer_id=host_layer_slot, + tai_transfer = _load_tai_transfer_kv_per_layer_direct_pf_lf() + for segment in self._host_transfer_segments( + self.index_k_with_scale_buffer, + host_page_indices, + device_page_indices, page_size=1, - ) + ): + tai_transfer( + src_ptrs=[segment.tensor], + dst_ptrs=[ + device_pool.index_k_with_scale_buffer[device_layer_slot] + ], + src_indices=segment.host_indices, + dst_indices=segment.device_indices, + layer_id=host_layer_slot, + page_size=1, + ) elif self.layout == "layer_page_first": if ( descriptor is not None @@ -2399,16 +3356,23 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): layer_id=host_layer_slot, ) return - _load_tai_transfer_kv_per_layer_direct_lpf_lf()( - src_ptrs=[self.index_k_with_scale_buffer], - dst_ptrs=[ - device_pool.index_k_with_scale_buffer[device_layer_slot] - ], - src_indices=host_page_indices, - dst_indices=device_page_indices, - layer_id=host_layer_slot, + tai_transfer = _load_tai_transfer_kv_per_layer_direct_lpf_lf() + for segment in self._host_transfer_segments( + self.index_k_with_scale_buffer, + host_page_indices, + device_page_indices, page_size=1, - ) + ): + tai_transfer( + src_ptrs=[segment.tensor], + dst_ptrs=[ + device_pool.index_k_with_scale_buffer[device_layer_slot] + ], + src_indices=segment.host_indices, + dst_indices=segment.device_indices, + layer_id=host_layer_slot, + page_size=1, + ) else: raise ValueError(f"Unsupported layout: {self.layout}") else: @@ -2509,16 +3473,24 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): device_pool, layer_id ) host_layer_slot = self._host_index_layer_slot(layer_id) - tai_transfer( - src_ptrs=[ - device_pool.index_k_with_scale_buffer[device_layer_slot] - ], - dst_ptrs=[self.index_k_with_scale_buffer], - src_indices=device_page_indices, - dst_indices=host_page_indices, - layer_id=host_layer_slot, + for segment in self._host_transfer_segments( + self.index_k_with_scale_buffer, + host_page_indices, + device_page_indices, page_size=1, - ) + ): + tai_transfer( + src_ptrs=[ + device_pool.index_k_with_scale_buffer[ + device_layer_slot + ] + ], + dst_ptrs=[segment.tensor], + src_indices=segment.device_indices, + dst_indices=segment.host_indices, + layer_id=host_layer_slot, + page_size=1, + ) elif self.layout == "layer_page_first": tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_lpf() for layer_id in active_layer_ids: @@ -2526,16 +3498,24 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): device_pool, layer_id ) host_layer_slot = self._host_index_layer_slot(layer_id) - tai_transfer( - src_ptrs=[ - device_pool.index_k_with_scale_buffer[device_layer_slot] - ], - dst_ptrs=[self.index_k_with_scale_buffer], - src_indices=device_page_indices, - dst_indices=host_page_indices, - layer_id=host_layer_slot, + for segment in self._host_transfer_segments( + self.index_k_with_scale_buffer, + host_page_indices, + device_page_indices, page_size=1, - ) + ): + tai_transfer( + src_ptrs=[ + device_pool.index_k_with_scale_buffer[ + device_layer_slot + ] + ], + dst_ptrs=[segment.tensor], + src_indices=segment.device_indices, + dst_indices=segment.host_indices, + layer_id=host_layer_slot, + page_size=1, + ) else: raise ValueError(f"Unsupported layout: {self.layout}") else: @@ -2585,27 +3565,41 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): page_size=1, ) elif self.layout == "page_first_direct": - _load_tai_transfer_kv_per_layer_direct_lf_pf()( - src_ptrs=[ - device_pool.index_k_with_scale_buffer[device_layer_slot] - ], - dst_ptrs=[self.index_k_with_scale_buffer], - src_indices=device_page_indices, - dst_indices=host_page_indices, - layer_id=host_layer_slot, + tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_pf() + for segment in self._host_transfer_segments( + self.index_k_with_scale_buffer, + host_page_indices, + device_page_indices, page_size=1, - ) + ): + tai_transfer( + src_ptrs=[ + device_pool.index_k_with_scale_buffer[device_layer_slot] + ], + dst_ptrs=[segment.tensor], + src_indices=segment.device_indices, + dst_indices=segment.host_indices, + layer_id=host_layer_slot, + page_size=1, + ) elif self.layout == "layer_page_first": - _load_tai_transfer_kv_per_layer_direct_lf_lpf()( - src_ptrs=[ - device_pool.index_k_with_scale_buffer[device_layer_slot] - ], - dst_ptrs=[self.index_k_with_scale_buffer], - src_indices=device_page_indices, - dst_indices=host_page_indices, - layer_id=host_layer_slot, + tai_transfer = _load_tai_transfer_kv_per_layer_direct_lf_lpf() + for segment in self._host_transfer_segments( + self.index_k_with_scale_buffer, + host_page_indices, + device_page_indices, page_size=1, - ) + ): + tai_transfer( + src_ptrs=[ + device_pool.index_k_with_scale_buffer[device_layer_slot] + ], + dst_ptrs=[segment.tensor], + src_indices=segment.device_indices, + dst_indices=segment.host_indices, + layer_id=host_layer_slot, + page_size=1, + ) else: raise ValueError(f"Unsupported layout: {self.layout}") else: @@ -2625,6 +3619,51 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): self._load_indexer_to_device_per_layer( device_pool, host_indices, device_indices, layer_id, io_backend ) + + def load_kv_to_device_per_layer( + self, + device_pool, + host_indices, + device_indices, + layer_id, + io_backend, + ): + """Load only target KV, leaving NSA index_k transfer to an explicit call.""" + + super().load_to_device_per_layer( + device_pool, host_indices, device_indices, layer_id, io_backend + ) + + def load_indexer_to_device_per_layer( + self, + device_pool, + host_indices, + device_indices, + layer_id, + io_backend, + ): + """Load only NSA index_k pages. + + CP shared physical L2 can reserve target_kv and index_k from different + shared slab namespaces. The public combined load path derives index + page indices from the target host range, so the controller uses this + split entry point when index_k has its own object range. + """ + + descriptor = getattr(self, "_active_load_descriptor", None) + active_indexer_page_indices = getattr( + self, "_active_load_indexer_page_indices", None + ) + self._active_load_descriptor = None + self._active_load_indexer_page_indices = None + try: + self._load_indexer_to_device_per_layer( + device_pool, host_indices, device_indices, layer_id, io_backend + ) + finally: + self._active_load_descriptor = descriptor + self._active_load_indexer_page_indices = active_indexer_page_indices + def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend ): @@ -2635,6 +3674,24 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): device_pool, host_indices, device_indices, io_backend ) + def backup_kv_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend + ): + """Back up only target KV, leaving NSA index_k to an explicit call.""" + + super().backup_from_device_all_layer( + device_pool, host_indices, device_indices, io_backend + ) + + def backup_indexer_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend + ): + """Back up only NSA index_k pages into a separately reserved range.""" + + self._backup_indexer_from_device_all_layer( + device_pool, host_indices, device_indices, io_backend + ) + def backup_from_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend ): @@ -2644,3 +3701,21 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): self._backup_indexer_from_device_per_layer( device_pool, host_indices, device_indices, layer_id, io_backend ) + + def backup_kv_from_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ): + """Back up only target KV, leaving NSA index_k to an explicit call.""" + + super().backup_from_device_per_layer( + device_pool, host_indices, device_indices, layer_id, io_backend + ) + + def backup_indexer_from_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ): + """Back up only NSA index_k pages into a separately reserved range.""" + + self._backup_indexer_from_device_per_layer( + device_pool, host_indices, device_indices, layer_id, io_backend + ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 68709d69f..87523e5d4 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -593,6 +593,10 @@ class ServerArgs: hicache_storage_backend: Optional[str] = None hicache_storage_prefetch_policy: str = "best_effort" hicache_storage_backend_extra_config: Optional[str] = None + enable_cp_shared_physical_l2_hicache: bool = False + cp_shared_l2_hugetlbfs_dir: str = "/mnt/huge_2m" + cp_shared_l2_slab_size_gb: int = 0 + cp_shared_l2_numa_policy: str = "interleave_2m" # Hierarchical sparse attention enable_hisparse: bool = False @@ -788,7 +792,9 @@ class ServerArgs: # hf_config through get_model_config(). self._handle_nsa_index_model_override_args() - # Validate CP shared KV constraints early (before dummy-model short-circuit). + # Validate CP shared physical L2 constraints before the broader CP shared KV + # checks so users get the most actionable error for the experimental flag. + self._handle_cp_shared_physical_l2_hicache_validation() self._handle_cp_shared_kv_validation() # Normalize load balancing defaults early (before dummy-model short-circuit). @@ -1080,6 +1086,51 @@ class ServerArgs: "layer_page_first is supported only with the direct IO backend." ) + def _handle_cp_shared_physical_l2_hicache_validation(self): + if self.cp_shared_l2_slab_size_gb < 0: + raise ValueError("cp_shared_l2_slab_size_gb must be non-negative") + if not self.enable_cp_shared_physical_l2_hicache: + return + + valid_numa_policies = { + "interleave_2m", + "balanced_local_preferred", + "strict_local", + "round_robin", + } + if self.cp_shared_l2_numa_policy not in valid_numa_policies: + raise ValueError( + "cp_shared_l2_numa_policy must be one of " + f"{sorted(valid_numa_policies)}; got {self.cp_shared_l2_numa_policy!r}" + ) + + if not self.enable_nsa_prefill_cp_shared_kv: + raise ValueError( + "--enable-cp-shared-physical-l2-hicache requires " + "--enable-nsa-prefill-cp-shared-kv." + ) + if not self.enable_hierarchical_cache: + raise ValueError( + "--enable-cp-shared-physical-l2-hicache requires " + "--enable-hierarchical-cache." + ) + if self.hicache_io_backend != "direct": + raise ValueError( + "--enable-cp-shared-physical-l2-hicache requires " + "--hicache-io-backend direct." + ) + if self.hicache_mem_layout != "layer_page_first": + raise ValueError( + "--enable-cp-shared-physical-l2-hicache requires " + "--hicache-mem-layout layer_page_first." + ) + if self.hicache_storage_backend is not None: + raise ValueError( + "--enable-cp-shared-physical-l2-hicache does not support " + "--hicache-storage-backend; layer_page_first page-buffer " + "metadata is not a storage contract." + ) + def _handle_deprecated_args(self): # Handle deprecated tool call parsers deprecated_tool_call_parsers = {"qwen25": "qwen", "glm45": "glm"} @@ -5501,6 +5552,46 @@ class ServerArgs: default=ServerArgs.hicache_storage_backend_extra_config, help="A dictionary in JSON string format, or a string starting with a leading '@' and a config file in JSON/YAML/TOML format, containing extra configuration for the storage backend.", ) + parser.add_argument( + "--enable-cp-shared-physical-l2-hicache", + action="store_true", + default=ServerArgs.enable_cp_shared_physical_l2_hicache, + help=( + "Enable experimental ownerless CP shared physical L2 HiCache " + "metadata. Requires CP shared KV HiCache with direct " + "layer_page_first host layout and no storage backend." + ), + ) + parser.add_argument( + "--cp-shared-l2-hugetlbfs-dir", + type=str, + default=ServerArgs.cp_shared_l2_hugetlbfs_dir, + help=( + "hugetlbfs directory reserved for CP shared physical L2 " + "HiCache slabs. The directory is validated by later " + "production integration; this flag does not mount or configure " + "hugetlbfs." + ), + ) + parser.add_argument( + "--cp-shared-l2-slab-size-gb", + type=int, + default=ServerArgs.cp_shared_l2_slab_size_gb, + help=( + "Physical host tensor slab size in decimal GB for CP shared " + "physical L2 HiCache. 0 preserves the legacy single-slab path." + ), + ) + parser.add_argument( + "--cp-shared-l2-numa-policy", + type=str, + default=ServerArgs.cp_shared_l2_numa_policy, + help=( + "Metadata-only NUMA policy for CP shared physical L2 slabs. " + "Allowed values: interleave_2m, balanced_local_preferred, " + "strict_local, round_robin." + ), + ) # Hierarchical sparse attention parser.add_argument( diff --git a/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py b/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py new file mode 100644 index 000000000..dadd3c692 --- /dev/null +++ b/test/registered/unit/mem_cache/test_cp_shared_l2_pool.py @@ -0,0 +1,2711 @@ +import argparse +import dataclasses +import importlib +import importlib.util +import json +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +_REPO_ROOT = Path(__file__).resolve().parents[4] + + +def _load_cp_shared_l2_pool_module(): + module_name = "_test_cp_shared_l2_pool_module" + module_path = _REPO_ROOT / "python" / "sglang" / "srt" / "mem_cache" / "cp_shared_l2_pool.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + with patch.dict(sys.modules, {module_name: module}): + spec.loader.exec_module(module) + return module + + +_cp_shared_l2_pool = _load_cp_shared_l2_pool_module() +CpSharedL2NodeMetadata = _cp_shared_l2_pool.CpSharedL2NodeMetadata +CpSharedL2ObjectRange = _cp_shared_l2_pool.CpSharedL2ObjectRange +CpSharedL2SlabInfo = _cp_shared_l2_pool.CpSharedL2SlabInfo +build_cp_shared_l2_slabs_by_payload = _cp_shared_l2_pool.build_cp_shared_l2_slabs_by_payload +cp_shared_l2_logical_token_indices = ( + _cp_shared_l2_pool.cp_shared_l2_logical_token_indices +) +PAYLOAD_DRAFT_KV = _cp_shared_l2_pool.PAYLOAD_DRAFT_KV +PAYLOAD_INDEX_K = _cp_shared_l2_pool.PAYLOAD_INDEX_K +PAYLOAD_TARGET_KV = _cp_shared_l2_pool.PAYLOAD_TARGET_KV + + +def _server_args_import_stubs(): + sglang_pkg = types.ModuleType("sglang") + sglang_pkg.__path__ = [str(_REPO_ROOT / "python" / "sglang")] + sglang_srt_pkg = types.ModuleType("sglang.srt") + sglang_srt_pkg.__path__ = [str(_REPO_ROOT / "python" / "sglang" / "srt")] + + safetensors_torch = types.ModuleType("safetensors.torch") + safetensors_torch.load = lambda *args, **kwargs: None + safetensors_torch.save = lambda *args, **kwargs: None + + orjson_stub = types.ModuleType("orjson") + orjson_stub.loads = lambda value, *args, **kwargs: json.loads( + value.decode() if isinstance(value, bytes) else value + ) + orjson_stub.dumps = lambda value, *args, **kwargs: json.dumps(value).encode() + + psutil_stub = types.ModuleType("psutil") + psutil_stub.NoSuchProcess = RuntimeError + psutil_stub.virtual_memory = lambda: types.SimpleNamespace(available=0, total=0) + psutil_stub.cpu_count = lambda logical=True: 1 + psutil_stub.Process = lambda *args, **kwargs: types.SimpleNamespace( + pid=0, + children=lambda recursive=True: [], + cpu_affinity=lambda *a, **k: [0], + ) + + triton_stub = types.ModuleType("triton") + triton_stub.__version__ = "0.0.0" + + pil_stub = types.ModuleType("PIL") + image_stub = types.ModuleType("PIL.Image") + image_stub.Image = object + pil_stub.Image = image_stub + + starlette_routing_stub = types.ModuleType("starlette.routing") + starlette_routing_stub.Mount = object + + connector_stub = types.ModuleType("sglang.srt.connector") + + class _ConnectorType: + INSTANCE = "instance" + + connector_stub.ConnectorType = _ConnectorType + + chunk_delta_h_stub = types.ModuleType("sglang.srt.layers.attention.fla.chunk_delta_h") + chunk_delta_h_stub.CHUNK_SIZE = 128 + + lora_registry_stub = types.ModuleType("sglang.srt.lora.lora_registry") + lora_registry_stub.LoRARef = type("LoRARef", (), {}) + + reasoning_parser_stub = types.ModuleType("sglang.srt.parser.reasoning_parser") + reasoning_parser_stub.ReasoningParser = type( + "ReasoningParser", (), {"DetectorMap": {"deepseek-r1": object}} + ) + + environ_stub = types.ModuleType("sglang.srt.environ") + # ServerArgs.__post_init__ -> _handle_cp_shared_kv_validation reads this env + # (Phase-1 fail-fast); stub it to its default (0 = no waiting timeout) so the + # flag-validation tests below construct ServerArgs without touching real env. + environ_stub.envs = types.SimpleNamespace( + SGLANG_REQ_WAITING_TIMEOUT=types.SimpleNamespace(get=lambda *a, **k: 0), + ) + + def _human_readable_int(value): + return int(value) + + _human_readable_int.__doc__ = "test stub" + + common_stub = types.ModuleType("sglang.srt.utils.common") + for name in ( + "cpu_has_amx_support", + "get_bool_env_var", + "is_blackwell_supported", + "is_cpu", + "is_cuda", + "is_flashinfer_available", + "is_hip", + "is_hopper_with_cuda_12_3", + "is_mps", + "is_musa", + "is_no_spec_infer_or_topk_one", + "is_npu", + "is_remote_url", + "is_sm90_supported", + "is_sm100_supported", + "is_sm120_supported", + "is_triton_kernels_available", + "is_xpu", + "xpu_has_xmx_support", + ): + setattr(common_stub, name, lambda *args, **kwargs: False) + common_stub.LORA_TARGET_ALL_MODULES = "all" + common_stub.SUPPORTED_LORA_TARGET_MODULES = [] + common_stub.get_device = lambda *args, **kwargs: "cuda" + common_stub.get_device_memory_capacity = lambda *args, **kwargs: 0 + common_stub.get_device_name = lambda *args, **kwargs: "" + common_stub.get_device_sm = lambda *args, **kwargs: 0 + common_stub.get_int_env_var = lambda *args, **kwargs: 0 + common_stub.get_quantization_config = lambda *args, **kwargs: None + common_stub.human_readable_int = _human_readable_int + common_stub.json_list_type = lambda value: [] if value is None else value.split(",") + common_stub.nullable_str = ( + lambda value: None if value in (None, "None", "none", "") else value + ) + common_stub.parse_connector_type = lambda value: None + common_stub.torch_release = (0, 0) + + hf_utils_stub = types.ModuleType("sglang.srt.utils.hf_transformers_utils") + hf_utils_stub.check_gguf_file = lambda *args, **kwargs: False + + network_stub = types.ModuleType("sglang.srt.utils.network") + network_stub.NetworkAddress = lambda host, port: types.SimpleNamespace( + to_url=lambda scheme: f"{scheme}://{host}:{port}" + ) + network_stub.get_free_port = lambda *args, **kwargs: 0 + network_stub.wait_port_available = lambda *args, **kwargs: None + + sglang_utils_stub = types.ModuleType("sglang.utils") + sglang_utils_stub.is_in_ci = lambda: False + sglang_utils_stub.convert_json_schema_to_str = lambda schema: str(schema) + + function_call_parser_stub = types.ModuleType( + "sglang.srt.function_call.function_call_parser" + ) + function_call_parser_stub.FunctionCallParser = type( + "FunctionCallParser", + (), + {"ToolCallParserEnum": {"qwen": object, "qwen25": object}}, + ) + + return { + "sglang": sglang_pkg, + "sglang.srt": sglang_srt_pkg, + "safetensors": types.ModuleType("safetensors"), + "safetensors.torch": safetensors_torch, + "orjson": orjson_stub, + "psutil": psutil_stub, + "triton": triton_stub, + "PIL": pil_stub, + "PIL.Image": image_stub, + "starlette": types.ModuleType("starlette"), + "starlette.routing": starlette_routing_stub, + "sglang.srt.connector": connector_stub, + "sglang.srt.layers.attention.fla.chunk_delta_h": chunk_delta_h_stub, + "sglang.srt.lora.lora_registry": lora_registry_stub, + "sglang.srt.parser.reasoning_parser": reasoning_parser_stub, + "sglang.srt.environ": environ_stub, + "sglang.srt.utils.common": common_stub, + "sglang.srt.utils.hf_transformers_utils": hf_utils_stub, + "sglang.srt.utils.network": network_stub, + "sglang.utils": sglang_utils_stub, + "sglang.srt.function_call.function_call_parser": function_call_parser_stub, + } + + +class ServerArgsTestCase(unittest.TestCase): + def setUp(self): + self._module_patcher = patch.dict(sys.modules, _server_args_import_stubs()) + self._module_patcher.start() + sys.modules.pop("sglang.srt.server_args", None) + self.ServerArgs = importlib.import_module("sglang.srt.server_args").ServerArgs + + def tearDown(self): + sys.modules.pop("sglang.srt.server_args", None) + self._module_patcher.stop() + + def make_server_args(self, **overrides): + values = { + "model_path": "dummy", + "enable_cp_shared_physical_l2_hicache": True, + "enable_hierarchical_cache": True, + "enable_nsa_prefill_cp_shared_kv": True, + "hicache_io_backend": "direct", + "hicache_mem_layout": "layer_page_first", + "hicache_storage_backend": None, + } + values.update(overrides) + return self.ServerArgs(**values) + + +class TestCpSharedPhysicalL2ServerArgs(ServerArgsTestCase): + def test_defaults_disable_shared_physical_l2_and_use_standard_hugetlbfs_dir(self): + fields = {field.name: field for field in dataclasses.fields(self.ServerArgs)} + + self.assertFalse(fields["enable_cp_shared_physical_l2_hicache"].default) + self.assertEqual(fields["cp_shared_l2_hugetlbfs_dir"].default, "/mnt/huge_2m") + self.assertEqual(fields["cp_shared_l2_slab_size_gb"].default, 0) + self.assertEqual(fields["cp_shared_l2_numa_policy"].default, "interleave_2m") + + def test_cli_parser_accepts_shared_physical_l2_options(self): + parser = argparse.ArgumentParser() + self.ServerArgs.add_cli_args(parser) + + raw_args = parser.parse_args( + [ + "--model-path", + "dummy", + "--enable-hierarchical-cache", + "--enable-nsa-prefill-cp-shared-kv", + "--hicache-io-backend", + "direct", + "--hicache-mem-layout", + "layer_page_first", + "--enable-cp-shared-physical-l2-hicache", + "--cp-shared-l2-hugetlbfs-dir", + "/dev/hugepages", + "--cp-shared-l2-slab-size-gb", + "2", + "--cp-shared-l2-numa-policy", + "round_robin", + ] + ) + args = self.ServerArgs.from_cli_args(raw_args) + + self.assertTrue(args.enable_cp_shared_physical_l2_hicache) + self.assertEqual(args.cp_shared_l2_hugetlbfs_dir, "/dev/hugepages") + self.assertEqual(args.cp_shared_l2_slab_size_gb, 2) + self.assertEqual(args.cp_shared_l2_numa_policy, "round_robin") + + def test_post_init_accepts_cp_shared_kv_direct_layer_page_first_without_storage(self): + args = self.make_server_args() + + self.assertTrue(args.enable_cp_shared_physical_l2_hicache) + + def test_post_init_rejects_non_hierarchical_cache(self): + with self.assertRaisesRegex(ValueError, r"--enable-hierarchical-cache"): + self.make_server_args(enable_hierarchical_cache=False) + + def test_post_init_rejects_non_direct_backend(self): + with self.assertRaisesRegex(ValueError, r"--hicache-io-backend direct"): + self.make_server_args(hicache_io_backend="kernel") + + def test_post_init_rejects_non_layer_page_first_layout(self): + with self.assertRaisesRegex(ValueError, r"--hicache-mem-layout layer_page_first"): + self.make_server_args(hicache_mem_layout="layer_first") + + def test_post_init_rejects_storage_backend_with_physical_l2_error(self): + with self.assertRaisesRegex( + ValueError, + r"--enable-cp-shared-physical-l2-hicache.*--hicache-storage-backend", + ): + self.make_server_args(hicache_storage_backend="file") + + def test_post_init_rejects_without_cp_shared_kv(self): + with self.assertRaisesRegex(ValueError, r"--enable-nsa-prefill-cp-shared-kv"): + self.make_server_args(enable_nsa_prefill_cp_shared_kv=False) + + def test_post_init_rejects_negative_slab_size_even_when_feature_disabled(self): + with self.assertRaisesRegex(ValueError, "cp_shared_l2_slab_size_gb.*non-negative"): + self.make_server_args( + enable_cp_shared_physical_l2_hicache=False, + cp_shared_l2_slab_size_gb=-1, + ) + + def test_post_init_rejects_unknown_numa_policy_when_feature_enabled(self): + with self.assertRaisesRegex(ValueError, "cp_shared_l2_numa_policy"): + self.make_server_args(cp_shared_l2_numa_policy="bind_0") + + def test_post_init_accepts_allowed_numa_policies_when_feature_enabled(self): + for policy in ( + "interleave_2m", + "balanced_local_preferred", + "strict_local", + "round_robin", + ): + with self.subTest(policy=policy): + args = self.make_server_args(cp_shared_l2_numa_policy=policy) + self.assertEqual(args.cp_shared_l2_numa_policy, policy) + + +class TestCpSharedL2Metadata(unittest.TestCase): + def test_logical_token_indices_treat_base_page_as_global_page_id(self): + import torch + + object_range = CpSharedL2ObjectRange( + object_key="node-6/target-kv", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=2, + base_page=7, + num_pages=3, + generation=4, + ) + positions = torch.tensor([0, 5, 127, 191], dtype=torch.int64) + + host_indices = cp_shared_l2_logical_token_indices( + object_range, page_size=64, positions=positions + ) + + self.assertTrue( + torch.equal(host_indices, torch.tensor([448, 453, 575, 639])) + ) + + def test_logical_token_indices_validates_page_size_and_position_bounds(self): + import torch + + object_range = CpSharedL2ObjectRange( + object_key="node-6/target-kv", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=2, + base_page=7, + num_pages=3, + generation=4, + ) + + with self.assertRaisesRegex(ValueError, "page_size must be positive"): + cp_shared_l2_logical_token_indices( + object_range, page_size=0, positions=torch.tensor([0]) + ) + with self.assertRaisesRegex(ValueError, "within object range"): + cp_shared_l2_logical_token_indices( + object_range, page_size=64, positions=torch.tensor([-1]) + ) + with self.assertRaisesRegex(ValueError, "within object range"): + cp_shared_l2_logical_token_indices( + object_range, page_size=64, positions=torch.tensor([192]) + ) + + def test_logical_token_indices_supports_integral_non_tensor_positions(self): + object_range = CpSharedL2ObjectRange( + object_key="node-6/target-kv", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=2, + base_page=7, + num_pages=3, + generation=4, + ) + + self.assertEqual( + cp_shared_l2_logical_token_indices( + object_range, page_size=64, positions=5 + ), + 453, + ) + self.assertEqual( + cp_shared_l2_logical_token_indices( + object_range, page_size=64, positions=[0, 5] + ), + [448, 453], + ) + self.assertEqual( + cp_shared_l2_logical_token_indices( + object_range, page_size=64, positions=(0, 5) + ), + (448, 453), + ) + self.assertEqual( + cp_shared_l2_logical_token_indices( + object_range, page_size=64, positions=range(2) + ), + [448, 449], + ) + + def test_logical_token_indices_promotes_int32_tensor_to_int64_before_offset(self): + import torch + + object_range = CpSharedL2ObjectRange( + object_key="node-large/target-kv", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=2, + base_page=2**31 // 64, + num_pages=2, + generation=4, + ) + + host_indices = cp_shared_l2_logical_token_indices( + object_range, + page_size=64, + positions=torch.tensor([0, 1, 127], dtype=torch.int32), + ) + + self.assertEqual(host_indices.dtype, torch.int64) + self.assertTrue( + torch.equal( + host_indices, + torch.tensor([2**31, 2**31 + 1, 2**31 + 127], dtype=torch.int64), + ) + ) + + def test_logical_token_indices_rejects_non_integral_positions(self): + import torch + + object_range = CpSharedL2ObjectRange( + object_key="node-6/target-kv", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=2, + base_page=7, + num_pages=3, + generation=4, + ) + + for positions in (1.5, [0, 1.5], (0, 1.5), torch.tensor([0.0])): + with self.subTest(positions=positions): + with self.assertRaisesRegex(TypeError, "integral"): + cp_shared_l2_logical_token_indices( + object_range, page_size=64, positions=positions + ) + + def test_logical_token_indices_rejects_cuda_tensor_like_positions(self): + class CudaTensorLike: + is_cuda = True + + def numel(self): + raise AssertionError("CUDA-like tensor should be rejected first") + + object_range = CpSharedL2ObjectRange( + object_key="node-6/target-kv", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=2, + base_page=7, + num_pages=3, + generation=4, + ) + + with self.assertRaisesRegex(ValueError, "CPU"): + cp_shared_l2_logical_token_indices( + object_range, page_size=64, positions=CudaTensorLike() + ) + + def test_logical_token_indices_empty_inputs_preserve_container_shape(self): + import torch + + object_range = CpSharedL2ObjectRange( + object_key="node-6/target-kv", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=2, + base_page=7, + num_pages=3, + generation=4, + ) + + empty_tensor = cp_shared_l2_logical_token_indices( + object_range, page_size=64, positions=torch.tensor([], dtype=torch.int32) + ) + + self.assertEqual(empty_tensor.dtype, torch.int64) + self.assertEqual(empty_tensor.numel(), 0) + self.assertEqual( + cp_shared_l2_logical_token_indices( + object_range, page_size=64, positions=[] + ), + [], + ) + + def test_metadata_uses_payload_specific_object_ranges_without_numa_node(self): + ranges = { + PAYLOAD_TARGET_KV: CpSharedL2ObjectRange( + object_key="node-7/target-kv", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=3, + base_page=10, + num_pages=4, + generation=2, + ), + PAYLOAD_DRAFT_KV: CpSharedL2ObjectRange( + object_key="node-7/draft-kv", + payload_kind=PAYLOAD_DRAFT_KV, + slab_id=3, + base_page=14, + num_pages=2, + generation=2, + ), + PAYLOAD_INDEX_K: CpSharedL2ObjectRange( + object_key="node-7/index-k", + payload_kind=PAYLOAD_INDEX_K, + slab_id=4, + base_page=0, + num_pages=1, + generation=2, + ), + } + + metadata = CpSharedL2NodeMetadata( + logical_len=257, + padded_len=320, + page_size=64, + object_ranges=ranges, + required_payloads=(PAYLOAD_TARGET_KV, PAYLOAD_INDEX_K), + committed_payload_layers={PAYLOAD_TARGET_KV: {0, 1}, PAYLOAD_INDEX_K: {0}}, + ) + + self.assertFalse(metadata.committed) + self.assertIs( + metadata.object_ranges[PAYLOAD_TARGET_KV], ranges[PAYLOAD_TARGET_KV] + ) + self.assertEqual( + metadata.object_ranges[PAYLOAD_DRAFT_KV].payload_kind, PAYLOAD_DRAFT_KV + ) + self.assertNotIn("numa_node", CpSharedL2ObjectRange.__dataclass_fields__) + + def test_object_range_is_frozen_stable_descriptor(self): + object_range = CpSharedL2ObjectRange( + object_key="node-8/target-kv", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=2, + base_page=1, + num_pages=1, + generation=1, + ) + + with self.assertRaises(dataclasses.FrozenInstanceError): + object_range.base_page = 2 + + def test_object_range_rejects_payload_key_mismatch(self): + with self.assertRaisesRegex(ValueError, "payload_kind"): + CpSharedL2NodeMetadata( + logical_len=64, + padded_len=64, + page_size=64, + object_ranges={ + PAYLOAD_TARGET_KV: CpSharedL2ObjectRange( + object_key="node-9/draft-kv", + payload_kind=PAYLOAD_DRAFT_KV, + slab_id=1, + base_page=0, + num_pages=1, + generation=1, + ) + }, + required_payloads=(PAYLOAD_TARGET_KV,), + committed_payload_layers={}, + ) + + def test_shared_metadata_split_partitions_payload_ranges_by_page(self): + metadata = CpSharedL2NodeMetadata( + logical_len=192, + padded_len=256, + page_size=64, + object_ranges={ + PAYLOAD_TARGET_KV: CpSharedL2ObjectRange( + object_key="node-10", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=2, + base_page=10, + num_pages=4, + generation=7, + ), + PAYLOAD_DRAFT_KV: CpSharedL2ObjectRange( + object_key="node-10", + payload_kind=PAYLOAD_DRAFT_KV, + slab_id=3, + base_page=40, + num_pages=4, + generation=8, + ), + }, + required_payloads=(PAYLOAD_TARGET_KV, PAYLOAD_DRAFT_KV), + committed_payload_layers={ + PAYLOAD_TARGET_KV: {0, 1}, + PAYLOAD_DRAFT_KV: {0}, + }, + committed=True, + object_key="node-10", + ) + + parent, child = metadata.split(128) + + self.assertEqual(parent.logical_len, 128) + self.assertEqual(parent.padded_len, 128) + self.assertEqual(child.logical_len, 64) + self.assertEqual(child.padded_len, 128) + self.assertTrue(parent.committed) + self.assertTrue(child.committed) + self.assertEqual(parent.required_payloads, metadata.required_payloads) + self.assertEqual(child.required_payloads, metadata.required_payloads) + self.assertEqual( + parent.object_ranges[PAYLOAD_TARGET_KV], + CpSharedL2ObjectRange( + object_key="node-10", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=2, + base_page=10, + num_pages=2, + generation=7, + ), + ) + self.assertEqual( + child.object_ranges[PAYLOAD_TARGET_KV], + CpSharedL2ObjectRange( + object_key="node-10", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=2, + base_page=12, + num_pages=2, + generation=7, + ), + ) + self.assertEqual( + parent.object_ranges[PAYLOAD_DRAFT_KV], + CpSharedL2ObjectRange( + object_key="node-10", + payload_kind=PAYLOAD_DRAFT_KV, + slab_id=3, + base_page=40, + num_pages=2, + generation=8, + ), + ) + self.assertEqual( + child.object_ranges[PAYLOAD_DRAFT_KV], + CpSharedL2ObjectRange( + object_key="node-10", + payload_kind=PAYLOAD_DRAFT_KV, + slab_id=3, + base_page=42, + num_pages=2, + generation=8, + ), + ) + + def test_shared_metadata_split_rejects_non_page_aligned_split(self): + metadata = CpSharedL2NodeMetadata( + logical_len=96, + padded_len=128, + page_size=64, + object_ranges={ + PAYLOAD_TARGET_KV: CpSharedL2ObjectRange( + object_key="node-11", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=1, + base_page=0, + num_pages=2, + generation=1, + ) + }, + required_payloads=(PAYLOAD_TARGET_KV,), + committed_payload_layers={PAYLOAD_TARGET_KV: {0}}, + object_key="node-11", + ) + + with self.assertRaisesRegex(ValueError, "multiple of page_size"): + metadata.split(32) + + +class TestCpSharedHostSlabPrimitives(unittest.TestCase): + def test_create_attach_visibility_and_lifecycle_owner_unlinks_only_on_creator_close(self): + import tempfile + + with tempfile.TemporaryDirectory() as tmpdir: + creator = _cp_shared_l2_pool.create_cp_shared_host_slab( + directory=tmpdir, + name="slab0", + nbytes=4096, + shape=(1024,), + dtype_name="uint32", + creator_rank=3, + validate_production=False, + ) + self.assertEqual(creator.handle.name, "slab0") + self.assertEqual(creator.handle.nbytes, 4096) + self.assertEqual(creator.handle.mapped_nbytes, 2 * 1024 * 1024) + self.assertEqual(creator.handle.shape, (1024,)) + self.assertEqual(creator.handle.dtype_name, "uint32") + self.assertEqual(creator.handle.creator_rank, 3) + self.assertEqual(creator.handle.backend, _cp_shared_l2_pool.CP_SHARED_L2_SLAB_BACKEND) + self.assertEqual(creator.handle.hugetlbfs_dir, str(Path(tmpdir))) + self.assertEqual(creator.handle.os_page_size, 2 * 1024 * 1024) + self.assertEqual(creator.handle.numa_policy, "interleave_2m") + self.assertTrue(Path(creator.handle.path).exists()) + + attached = _cp_shared_l2_pool.attach_cp_shared_host_slab(creator.handle) + creator.mmap[8:12] = b"l2ok" + self.assertEqual(attached.mmap[8:12], b"l2ok") + + attached.close() + self.assertTrue(Path(creator.handle.path).exists()) + creator.close() + self.assertFalse(Path(creator.handle.path).exists()) + + def test_create_rounds_mapping_length_to_2m_hugepage_boundary(self): + import tempfile + + with tempfile.TemporaryDirectory() as tmpdir: + logical_nbytes = 4097 + mapped_nbytes = 2 * 1024 * 1024 + creator = _cp_shared_l2_pool.create_cp_shared_host_slab( + directory=tmpdir, + name="unaligned-slab", + nbytes=logical_nbytes, + shape=(logical_nbytes,), + dtype_name="uint8", + validate_production=False, + ) + try: + self.assertEqual(creator.handle.nbytes, logical_nbytes) + self.assertEqual(creator.handle.mapped_nbytes, mapped_nbytes) + self.assertEqual(len(creator.mmap), mapped_nbytes) + self.assertEqual(Path(creator.handle.path).stat().st_size, mapped_nbytes) + + attached = _cp_shared_l2_pool.attach_cp_shared_host_slab(creator.handle) + try: + self.assertEqual(len(attached.mmap), mapped_nbytes) + creator.mmap[logical_nbytes : logical_nbytes + 4] = b"pad!" + self.assertEqual( + attached.mmap[logical_nbytes : logical_nbytes + 4], + b"pad!", + ) + finally: + attached.close() + finally: + creator.close() + + def test_attach_close_does_not_unlink_when_creator_still_open(self): + import tempfile + + with tempfile.TemporaryDirectory() as tmpdir: + creator = _cp_shared_l2_pool.create_cp_shared_host_slab( + directory=tmpdir, name="slab1", nbytes=128, validate_production=False + ) + attached = _cp_shared_l2_pool.attach_cp_shared_host_slab(creator.handle) + + attached.close() + + self.assertTrue(Path(creator.handle.path).exists()) + creator.close() + + def test_validation_failfasts_for_missing_path_wrong_config_and_non_2m_pages(self): + with self.assertRaisesRegex(ValueError, r"\[CP_SHARED_L2_FAILFAST\].*does not exist"): + _cp_shared_l2_pool.validate_hugetlbfs_dir("/definitely/missing/cp-l2", production=True) + + with self.assertRaisesRegex(ValueError, r"\[CP_SHARED_L2_FAILFAST\].*mount point"): + _cp_shared_l2_pool.validate_hugetlbfs_dir( + ".", production=True, mount_checker=lambda path: False + ) + + with self.assertRaisesRegex(ValueError, r"\[CP_SHARED_L2_FAILFAST\].*hugetlbfs"): + _cp_shared_l2_pool.validate_hugetlbfs_dir( + ".", production=True, mount_checker=lambda path: "tmpfs" + ) + + with self.assertRaisesRegex(ValueError, r"\[CP_SHARED_L2_FAILFAST\].*backend"): + _cp_shared_l2_pool.validate_cp_shared_host_slab_handle_config( + _cp_shared_l2_pool.CpSharedHostSlabHandle( + name="bad", + path="/tmp/bad", + nbytes=1, + shape=(1,), + dtype_name="uint8", + creator_rank=0, + backend="tmpfs", + hugetlbfs_dir="/tmp", + os_page_size=2 * 1024 * 1024, + numa_policy="interleave_2m", + ) + ) + + with self.assertRaisesRegex(ValueError, r"\[CP_SHARED_L2_FAILFAST\].*2MiB"): + _cp_shared_l2_pool.CpSharedHostSlabHandle( + name="bad-page", + path="/tmp/bad-page", + nbytes=1, + shape=(1,), + dtype_name="uint8", + creator_rank=0, + backend="hugetlbfs_2m", + hugetlbfs_dir="/tmp", + os_page_size=4096, + numa_policy="interleave_2m", + ) + + with self.assertRaisesRegex(ValueError, r"\[CP_SHARED_L2_FAILFAST\].*numa_policy"): + _cp_shared_l2_pool.CpSharedHostSlabHandle( + name="bad-policy", + path="/tmp/bad-policy", + nbytes=1, + shape=(1,), + dtype_name="uint8", + creator_rank=0, + backend="hugetlbfs_2m", + hugetlbfs_dir="/tmp", + os_page_size=2 * 1024 * 1024, + numa_policy="bind_0", + ) + + smaps_4k = "7f-80 rw-s 00000000 00:00 0 /tmp/slab\nKernelPageSize: 4 kB\nMMUPageSize: 4 kB\n" + with self.assertRaisesRegex(ValueError, r"\[CP_SHARED_L2_FAILFAST\].*2MiB"): + _cp_shared_l2_pool.validate_effective_2m_page_mapping(smaps_text=smaps_4k) + + def test_hugetlbfs_validation_accepts_injected_hugetlbfs_mount_proof(self): + import tempfile + + with tempfile.TemporaryDirectory() as tmpdir: + self.assertEqual( + _cp_shared_l2_pool.validate_hugetlbfs_dir( + tmpdir, + production=True, + mount_checker=lambda path: "hugetlbfs", + ), + str(Path(tmpdir)), + ) + + def test_validation_accepts_injected_checker_or_smaps_with_2m_pages(self): + self.assertEqual( + _cp_shared_l2_pool.validate_effective_2m_page_mapping(checker=lambda: True), + 2 * 1024 * 1024, + ) + smaps_2m = "KernelPageSize: 2048 kB\nMMUPageSize: 2048 kB\n" + self.assertEqual( + _cp_shared_l2_pool.validate_effective_2m_page_mapping(smaps_text=smaps_2m), + 2 * 1024 * 1024, + ) + + def test_create_defers_effective_page_size_validation_until_mapping_is_live(self): + import tempfile + + calls = [] + + def checker(path=None): + calls.append(path) + return True + + with tempfile.TemporaryDirectory() as tmpdir: + with patch.object( + _cp_shared_l2_pool, + "_find_linux_mount_fs_type", + return_value="hugetlbfs", + ), patch.object(_cp_shared_l2_pool.os.path, "ismount", return_value=True): + mapping = _cp_shared_l2_pool.create_cp_shared_host_slab( + directory=tmpdir, + name="deferred-page-size-check", + nbytes=128, + validate_production=True, + page_size_checker=checker, + ) + try: + self.assertEqual(calls, []) + self.assertTrue(Path(mapping.handle.path).exists()) + finally: + mapping.close() + + + def test_page_size_checker_internal_type_error_propagates_without_retry(self): + calls = [] + + def checker(path=None): + calls.append(path) + raise TypeError("internal checker bug") + + with self.assertRaisesRegex(TypeError, "internal checker bug"): + _cp_shared_l2_pool.validate_effective_2m_page_mapping( + checker=checker, path="/tmp/slab" + ) + self.assertEqual(calls, ["/tmp/slab"]) + + def test_creator_close_marks_closed_when_unlink_raises_and_second_close_is_noop(self): + import tempfile + + with tempfile.TemporaryDirectory() as tmpdir: + creator = _cp_shared_l2_pool.create_cp_shared_host_slab( + directory=tmpdir, name="unlink-fail", nbytes=128, validate_production=False + ) + unlink_calls = [] + + def failing_unlink(path): + unlink_calls.append(path) + raise OSError("unlink denied") + + with patch.object(_cp_shared_l2_pool.os, "unlink", side_effect=failing_unlink): + with self.assertRaisesRegex(OSError, "unlink denied"): + creator.close() + self.assertTrue(creator._closed) + creator.close() + + self.assertEqual(unlink_calls, [creator.handle.path]) + Path(creator.handle.path).unlink() + + def test_creator_can_unlink_name_while_mapping_remains_usable(self): + import tempfile + + with tempfile.TemporaryDirectory() as tmpdir: + creator = _cp_shared_l2_pool.create_cp_shared_host_slab( + directory=tmpdir, + name="early-unlink", + nbytes=128, + validate_production=False, + ) + attached = _cp_shared_l2_pool.attach_cp_shared_host_slab(creator.handle) + + self.assertTrue(Path(creator.handle.path).exists()) + + self.assertTrue(creator.unlink()) + self.assertFalse(Path(creator.handle.path).exists()) + + creator.mmap[0:4] = b"live" + self.assertEqual(attached.mmap[0:4], b"live") + + attached.mmap[8:12] = b"rank" + self.assertEqual(creator.mmap[8:12], b"rank") + + self.assertFalse(creator.unlink()) + attached.close() + creator.close() + + def test_registration_contract_helper_uses_explicit_pointer_and_buffer_size(self): + import mmap + + explicit = _cp_shared_l2_pool.expected_cp_shared_l2_registration_range( + ptr=0x1000, nbytes=8192 + ) + self.assertEqual((explicit.ptr, explicit.nbytes), (0x1000, 8192)) + + mapping = mmap.mmap(-1, 64) + try: + reg = _cp_shared_l2_pool.cp_shared_l2_registration_range(mapping) + self.assertIsInstance(reg.ptr, int) + self.assertGreater(reg.ptr, 0) + self.assertEqual(reg.nbytes, 64) + finally: + mapping.close() + + def test_registration_helper_rejects_explicit_nbytes_beyond_tensor_backing_size(self): + class FakeTensor: + def data_ptr(self): + return 0x2000 + + def numel(self): + return 4 + + def element_size(self): + return 2 + + with self.assertRaisesRegex(ValueError, r"\[CP_SHARED_L2_FAILFAST\].*exceeds"): + _cp_shared_l2_pool.cp_shared_l2_registration_range( + FakeTensor(), nbytes=9 + ) + + + +class TestCpSharedL2PageAllocator(unittest.TestCase): + def make_allocator( + self, + *, + pages=8, + ranks=(0, 1, 2, 3), + layers=(0, 1), + payloads=(PAYLOAD_TARGET_KV,), + ): + 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, + ) + + + def test_build_slabs_by_payload_splits_contiguous_global_ranges(self): + slabs = build_cp_shared_l2_slabs_by_payload( + {PAYLOAD_TARGET_KV: 10}, + slab_size_pages=4, + numa_policy="round_robin", + ) + + self.assertEqual( + slabs[PAYLOAD_TARGET_KV], + ( + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, 0, 0, 4), + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, 1, 4, 4), + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, 2, 8, 2), + ), + ) + + def test_build_slabs_by_payload_defaults_to_single_slab(self): + slabs = build_cp_shared_l2_slabs_by_payload( + {PAYLOAD_TARGET_KV: 10, PAYLOAD_INDEX_K: 3}, + slab_size_pages=0, + ) + + self.assertEqual(len(slabs[PAYLOAD_TARGET_KV]), 1) + self.assertEqual((slabs[PAYLOAD_TARGET_KV][0].global_base_page, slabs[PAYLOAD_TARGET_KV][0].num_pages), (0, 10)) + self.assertEqual(len(slabs[PAYLOAD_INDEX_K]), 1) + self.assertEqual((slabs[PAYLOAD_INDEX_K][0].global_base_page, slabs[PAYLOAD_INDEX_K][0].num_pages), (0, 3)) + + def test_rank0_reserve_broadcasts_one_shared_range_to_nonzero_ranks(self): + allocator = self.make_allocator(pages=16) + rank0_range = allocator.reserve("obj-a", PAYLOAD_TARGET_KV, 3) + broadcasts = [] + + def fake_broadcast(object_list, src, group): + broadcasts.append((src, group, list(object_list))) + object_list[0] = rank0_range + + received = [ + _cp_shared_l2_pool.broadcast_cp_shared_l2_object_range( + rank0_range if rank == 0 else None, + cp_cpu_group="fake-cpu-group", + rank=rank, + source_rank=0, + broadcast_fn=fake_broadcast, + ) + for rank in range(4) + ] + + self.assertEqual(allocator.used_pages(PAYLOAD_TARGET_KV), 3) + self.assertEqual(allocator.free_pages(PAYLOAD_TARGET_KV), 13) + self.assertEqual(len({r for r in received}), 1) + for got in received: + self.assertEqual(got.payload_kind, PAYLOAD_TARGET_KV) + self.assertEqual(got.slab_id, 10) + self.assertEqual(got.base_page, 0) + self.assertEqual(got.generation, rank0_range.generation) + self.assertEqual(len(broadcasts), 4) + self.assertTrue( + all(src == 0 and group == "fake-cpu-group" for src, group, _ in broadcasts) + ) + + def test_release_and_abort_return_capacity_once_without_double_free(self): + allocator = self.make_allocator(pages=5) + allocator.reserve("release-me", PAYLOAD_TARGET_KV, 2) + + self.assertTrue(allocator.release("release-me")) + self.assertEqual(allocator.free_pages(PAYLOAD_TARGET_KV), 5) + self.assertFalse(allocator.release("release-me")) + self.assertFalse(allocator.abort("release-me")) + self.assertEqual(allocator.free_pages(PAYLOAD_TARGET_KV), 5) + + allocator.reserve("abort-me", PAYLOAD_TARGET_KV, 3) + self.assertTrue(allocator.abort("abort-me")) + self.assertFalse(allocator.abort("abort-me")) + self.assertEqual(allocator.free_pages(PAYLOAD_TARGET_KV), 5) + + 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) + self.assertTrue(allocator.is_committed("node-12")) + + parent, child = allocator.split_committed_object( + "node-12", + split_pages_by_payload={PAYLOAD_TARGET_KV: 2}, + parent_object_key="node-99", + child_object_key="node-12", + ) + + self.assertFalse(allocator.is_committed("node-12:old")) + self.assertTrue(allocator.is_committed("node-99")) + self.assertTrue(allocator.is_committed("node-12")) + self.assertEqual( + parent[PAYLOAD_TARGET_KV], + CpSharedL2ObjectRange( + object_key="node-99", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=original.slab_id, + base_page=0, + num_pages=2, + generation=original.generation, + ), + ) + self.assertEqual( + child[PAYLOAD_TARGET_KV], + CpSharedL2ObjectRange( + object_key="node-12", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=original.slab_id, + base_page=2, + num_pages=2, + generation=original.generation, + ), + ) + self.assertEqual(allocator.free_pages(PAYLOAD_TARGET_KV), 4) + + self.assertTrue(allocator.release("node-99")) + self.assertEqual(allocator.free_pages(PAYLOAD_TARGET_KV), 6) + self.assertTrue(allocator.release("node-12")) + self.assertEqual(allocator.free_pages(PAYLOAD_TARGET_KV), 8) + + def test_duplicate_live_object_payload_reservations_are_rejected(self): + allocator = self.make_allocator(pages=5) + allocator.reserve("dup", PAYLOAD_TARGET_KV, 1) + + with self.assertRaisesRegex(ValueError, r"duplicate.*dup.*target_kv"): + allocator.reserve("dup", PAYLOAD_TARGET_KV, 1) + + self.assertEqual(allocator.used_pages(PAYLOAD_TARGET_KV), 1) + + def test_contiguous_free_list_merges_adjacent_ranges_and_reuse_increments_generation(self): + allocator = self.make_allocator(pages=6) + a = allocator.reserve("a", PAYLOAD_TARGET_KV, 2) + b = allocator.reserve("b", PAYLOAD_TARGET_KV, 2) + c = allocator.reserve("c", PAYLOAD_TARGET_KV, 2) + self.assertEqual((a.base_page, b.base_page, c.base_page), (0, 2, 4)) + + self.assertTrue(allocator.release("b")) + self.assertTrue(allocator.release("a")) + merged = allocator.reserve("merged", PAYLOAD_TARGET_KV, 4) + + self.assertEqual(merged.base_page, 0) + self.assertEqual(merged.num_pages, 4) + self.assertGreater(merged.generation, c.generation) + self.assertEqual(allocator.free_pages(PAYLOAD_TARGET_KV), 0) + + def test_multi_slab_reserve_uses_global_page_ids_and_slab_id(self): + allocator = _cp_shared_l2_pool.CpSharedL2PageAllocator( + slabs_by_payload={ + PAYLOAD_TARGET_KV: [ + CpSharedL2SlabInfo( + payload_kind=PAYLOAD_TARGET_KV, + slab_id=4, + global_base_page=100, + num_pages=3, + numa_node=0, + ), + { + "payload_kind": PAYLOAD_TARGET_KV, + "slab_id": 5, + "global_base_page": 200, + "num_pages": 4, + }, + ] + } + ) + + first = allocator.reserve("multi-a", PAYLOAD_TARGET_KV, 2) + second = allocator.reserve("multi-b", PAYLOAD_TARGET_KV, 4) + + self.assertEqual( + (first.slab_id, first.base_page, first.num_pages), (4, 100, 2) + ) + self.assertEqual( + (second.slab_id, second.base_page, second.num_pages), (5, 200, 4) + ) + self.assertEqual(allocator.free_pages(PAYLOAD_TARGET_KV), 1) + self.assertEqual(allocator.used_pages(PAYLOAD_TARGET_KV), 6) + self.assertEqual( + allocator.stats()["cp_shared_l2_pages_capacity"], + 7, + ) + + def test_multi_slab_reservation_never_crosses_slabs(self): + allocator = _cp_shared_l2_pool.CpSharedL2PageAllocator( + slabs_by_payload={ + PAYLOAD_TARGET_KV: [ + { + "payload_kind": PAYLOAD_TARGET_KV, + "slab_id": 0, + "global_base_page": 0, + "num_pages": 3, + }, + { + "payload_kind": PAYLOAD_TARGET_KV, + "slab_id": 1, + "global_base_page": 100, + "num_pages": 3, + }, + ] + } + ) + + with self.assertRaisesRegex(ValueError, "insufficient contiguous"): + allocator.reserve("too-large", PAYLOAD_TARGET_KV, 4) + + self.assertEqual(allocator.free_pages(PAYLOAD_TARGET_KV), 6) + + def test_multi_slab_release_and_abort_return_pages_to_original_slab_only(self): + allocator = _cp_shared_l2_pool.CpSharedL2PageAllocator( + slabs_by_payload={ + PAYLOAD_TARGET_KV: [ + { + "payload_kind": PAYLOAD_TARGET_KV, + "slab_id": 0, + "global_base_page": 0, + "num_pages": 2, + }, + { + "payload_kind": PAYLOAD_TARGET_KV, + "slab_id": 1, + "global_base_page": 10, + "num_pages": 2, + }, + ] + } + ) + first = allocator.reserve("first", PAYLOAD_TARGET_KV, 2) + second = allocator.reserve("second", PAYLOAD_TARGET_KV, 2) + self.assertEqual((first.slab_id, second.slab_id), (0, 1)) + + self.assertTrue(allocator.release("first")) + reused_first = allocator.reserve("reuse-first", PAYLOAD_TARGET_KV, 2) + self.assertEqual((reused_first.slab_id, reused_first.base_page), (0, 0)) + + self.assertTrue(allocator.abort("second")) + reused_second = allocator.reserve("reuse-second", PAYLOAD_TARGET_KV, 2) + self.assertEqual((reused_second.slab_id, reused_second.base_page), (1, 10)) + + def test_split_committed_object_on_nonzero_base_slab_preserves_global_pages_and_slab_id( + self, + ): + allocator = _cp_shared_l2_pool.CpSharedL2PageAllocator( + slabs_by_payload={ + PAYLOAD_TARGET_KV: [ + { + "payload_kind": PAYLOAD_TARGET_KV, + "slab_id": 7, + "global_base_page": 500, + "num_pages": 6, + } + ] + }, + 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)) + + parent, child = allocator.split_committed_object( + "node-global", + split_pages_by_payload={PAYLOAD_TARGET_KV: 3}, + parent_object_key="node-global-parent", + child_object_key="node-global-child", + ) + + self.assertEqual( + parent[PAYLOAD_TARGET_KV], + CpSharedL2ObjectRange( + object_key="node-global-parent", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=7, + base_page=500, + num_pages=3, + generation=original.generation, + ), + ) + self.assertEqual( + child[PAYLOAD_TARGET_KV], + CpSharedL2ObjectRange( + object_key="node-global-child", + payload_kind=PAYLOAD_TARGET_KV, + slab_id=7, + base_page=503, + num_pages=1, + generation=original.generation, + ), + ) + + def test_multi_slab_config_rejects_duplicate_ids_overlaps_and_capacity_mismatch(self): + with self.assertRaisesRegex(ValueError, "duplicate slab_id"): + _cp_shared_l2_pool.CpSharedL2PageAllocator( + slabs_by_payload={ + PAYLOAD_TARGET_KV: [ + { + "payload_kind": PAYLOAD_TARGET_KV, + "slab_id": 0, + "global_base_page": 0, + "num_pages": 1, + }, + { + "payload_kind": PAYLOAD_TARGET_KV, + "slab_id": 0, + "global_base_page": 10, + "num_pages": 1, + }, + ] + } + ) + with self.assertRaisesRegex(ValueError, "overlap"): + _cp_shared_l2_pool.CpSharedL2PageAllocator( + slabs_by_payload={ + PAYLOAD_TARGET_KV: [ + { + "payload_kind": PAYLOAD_TARGET_KV, + "slab_id": 0, + "global_base_page": 0, + "num_pages": 3, + }, + { + "payload_kind": PAYLOAD_TARGET_KV, + "slab_id": 1, + "global_base_page": 2, + "num_pages": 3, + }, + ] + } + ) + with self.assertRaisesRegex(ValueError, "capacity"): + _cp_shared_l2_pool.CpSharedL2PageAllocator( + pages_per_payload={PAYLOAD_TARGET_KV: 3}, + slabs_by_payload={ + PAYLOAD_TARGET_KV: [ + { + "payload_kind": PAYLOAD_TARGET_KV, + "slab_id": 0, + "global_base_page": 0, + "num_pages": 4, + } + ] + }, + ) + + 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( + "pg", + same_node_checker=lambda pg, source_rank=0: [ + True, + True, + True, + True, + ], + ), + [True, True, True, True], + ) + + with self.assertRaisesRegex( + ValueError, + r"\[CP_SHARED_L2_FAILFAST\]\[cross_node_cp_group_unsupported\]", + ): + _cp_shared_l2_pool.require_cp_shared_l2_same_node( + "pg", + same_node_checker=lambda pg, source_rank=0: [ + True, + False, + True, + True, + ], + ) + + def test_default_broadcast_maps_cp_local_source_rank_to_global_rank(self): + calls = [] + + class FakeDist: + @staticmethod + def get_global_rank(group, group_rank): + calls.append(("get_global_rank", group, group_rank)) + self.assertEqual(group, "cp-group") + self.assertEqual(group_rank, 0) + return 5 + + @staticmethod + def broadcast_object_list(object_list, *, src, group): + calls.append(("broadcast", src, group, list(object_list))) + object_list[0] = {"handle": "from-global-rank-5"} + + got = _cp_shared_l2_pool.broadcast_cp_shared_l2_object( + {"handle": "from-local-rank-0"}, + cp_cpu_group="cp-group", + rank=0, + source_rank=0, + dist_module=FakeDist, + ) + + self.assertEqual(got["handle"], "from-global-rank-5") + 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) + + + + +# Patch D: shared host tensor allocator construction tests. These tests load +# memory_pool_host behind lightweight stubs so they do not require CUDA, +# torch.distributed, psutil, or sgl-kernel. +def _load_memory_pool_host_module_for_patch_d(): + import torch + + module_name = "_test_memory_pool_host_patch_d" + module_path = _REPO_ROOT / "python" / "sglang" / "srt" / "mem_cache" / "memory_pool_host.py" + + psutil_stub = types.ModuleType("psutil") + psutil_stub.virtual_memory = lambda: types.SimpleNamespace(available=10**15) + + hicache_jit_stub = types.ModuleType("sglang.jit_kernel.hicache") + hicache_jit_stub.can_use_hicache_jit_kernel = lambda *a, **k: False + hicache_jit_stub.transfer_hicache_all_layer = lambda *a, **k: None + hicache_jit_stub.transfer_hicache_one_layer = lambda *a, **k: None + + memory_pool_stub = types.ModuleType("sglang.srt.mem_cache.memory_pool") + + class KVCache: + pass + + class MHATokenToKVPool(KVCache): + pass + + class MLATokenToKVPool(KVCache): + pass + + class NSATokenToKVPool(MLATokenToKVPool): + index_k_with_scale_buffer_dtype = torch.uint8 + + memory_pool_stub.KVCache = KVCache + memory_pool_stub.MHATokenToKVPool = MHATokenToKVPool + memory_pool_stub.MLATokenToKVPool = MLATokenToKVPool + memory_pool_stub.NSATokenToKVPool = NSATokenToKVPool + + page_index_stub = types.ModuleType("sglang.srt.mem_cache.page_index_utils") + page_index_stub.validate_page_aligned_token_indices = lambda *a, **k: None + + utils_stub = types.ModuleType("sglang.srt.utils") + for name in ("is_cuda", "is_mps", "is_npu", "is_xpu"): + setattr(utils_stub, name, lambda: False) + + kvcacheio_stub = types.ModuleType("sgl_kernel.kvcacheio") + for name in ( + "transfer_kv_all_layer", + "transfer_kv_all_layer_direct_lf_pf", + "transfer_kv_all_layer_lf_pf", + "transfer_kv_all_layer_lf_ph", + "transfer_kv_all_layer_mla", + "transfer_kv_all_layer_mla_lf_pf", + "transfer_kv_direct", + "transfer_kv_per_layer", + "transfer_kv_per_layer_mla", + "transfer_kv_per_layer_mla_pf_lf", + "transfer_kv_per_layer_pf_lf", + "transfer_kv_per_layer_ph_lf", + ): + setattr(kvcacheio_stub, name, lambda *a, **k: None) + + stubs = { + "psutil": psutil_stub, + "sglang.jit_kernel.hicache": hicache_jit_stub, + "sglang.srt.mem_cache.memory_pool": memory_pool_stub, + "sglang.srt.mem_cache.page_index_utils": page_index_stub, + "sglang.srt.utils": utils_stub, + "sgl_kernel": types.ModuleType("sgl_kernel"), + "sgl_kernel.kvcacheio": kvcacheio_stub, + } + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + with patch.dict(sys.modules, {module_name: module, **stubs}): + spec.loader.exec_module(module) + module._test_memory_pool_stub = memory_pool_stub + return module + + +class TestPatchDSharedHostTensorAllocator(unittest.TestCase): + def test_shared_allocator_close_refuses_live_tensor_view_until_released(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + with tempfile.TemporaryDirectory() as tmpdir: + allocator = mph.SharedHostTensorAllocator( + directory=tmpdir, + name="close-live-view", + creator_rank=0, + validate_production=False, + ) + tensor = allocator.allocate((4,), dtype=torch.int32, device="cpu") + tensor[0] = 9 + + with self.assertRaisesRegex(RuntimeError, "release_tensor_view_for_close"): + allocator.close() + + allocator.release_tensor_view_for_close() + del tensor + allocator.close() + + def test_shared_allocator_tensor_view_attach_sees_writes_and_exposes_registration(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + with tempfile.TemporaryDirectory() as tmpdir: + creator = mph.SharedHostTensorAllocator( + directory=tmpdir, + name="shared-tensor", + creator_rank=0, + validate_production=False, + ) + tensor = creator.allocate((2, 4), dtype=torch.int32, device="cpu") + tensor.fill_(0) + tensor[1, 2] = 12345 + + attached = mph.SharedHostTensorAllocator( + handle=creator.slab_handle, + validate_production=False, + ) + attached_tensor = attached.allocate((2, 4), dtype=torch.int32, device="cpu") + + self.assertEqual(attached_tensor[1, 2].item(), 12345) + self.assertEqual(creator.slab_handle.shape, (2, 4)) + self.assertEqual(creator.slab_handle.dtype_name, "int32") + self.assertEqual(creator.registration_range.ptr, tensor.data_ptr()) + self.assertEqual(creator.registration_range.nbytes, tensor.numel() * tensor.element_size()) + self.assertEqual(attached.registration_range.nbytes, creator.registration_range.nbytes) + + attached.release_tensor_view_for_close() + creator.release_tensor_view_for_close() + del attached_tensor + del tensor + attached.close() + creator.close() + + def test_mla_layer_page_first_shared_allocator_preserves_normal_shape_and_refs(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + device_pool = types.SimpleNamespace( + store_dtype=torch.float16, + size=8, + start_layer=0, + end_layer=3, + device="cpu", + kv_lora_rank=16, + qk_rope_head_dim=4, + layer_num=3, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + shared_allocator = mph.SharedHostTensorAllocator( + directory=tmpdir, + name="mla-target", + creator_rank=0, + validate_production=False, + ) + shared_pool = mph.MLATokenToKVPoolHost( + device_pool=device_pool, + host_to_device_ratio=2.0, + host_size=0, + page_size=4, + layout="layer_page_first", + pin_memory=False, + device="cpu", + host_token_capacity=16, + host_tensor_allocator=shared_allocator, + ) + normal_pool = mph.MLATokenToKVPoolHost( + device_pool=device_pool, + host_to_device_ratio=2.0, + host_size=0, + page_size=4, + layout="layer_page_first", + pin_memory=False, + device="cpu", + host_token_capacity=16, + ) + + self.assertEqual(tuple(shared_pool.kv_buffer.shape), tuple(normal_pool.kv_buffer.shape)) + self.assertEqual(shared_pool.layout, normal_pool.layout) + self.assertEqual(shared_pool.page_size, normal_pool.page_size) + self.assertEqual(shared_pool.page_num, normal_pool.page_num) + self.assertEqual(shared_pool.size, normal_pool.size) + self.assertEqual(len(shared_pool.data_refs), len(normal_pool.data_refs)) + self.assertEqual(tuple(shared_pool.data_refs[0].shape), tuple(normal_pool.data_refs[0].shape)) + self.assertIs(shared_pool.allocator, shared_allocator) + self.assertEqual(shared_allocator.slab_handle.shape, tuple(shared_pool.kv_buffer.shape)) + shared_allocator.release_tensor_view_for_close() + del shared_pool + del normal_pool + shared_allocator.close() + + def test_nsa_shared_allocator_uses_separate_main_and_index_slabs(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + indexer_dtype = mph._test_memory_pool_stub.NSATokenToKVPool.index_k_with_scale_buffer_dtype + device_pool = types.SimpleNamespace( + store_dtype=torch.float16, + size=8, + start_layer=0, + end_layer=3, + device="cpu", + kv_lora_rank=16, + qk_rope_head_dim=4, + kv_cache_dim=24, + layer_num=3, + index_head_dim=16, + quant_block_size=8, + index_k_with_scale_buffer=[torch.empty((5, 96), dtype=indexer_dtype) for _ in range(3)], + ) + + with tempfile.TemporaryDirectory() as tmpdir: + main_allocator = mph.SharedHostTensorAllocator( + directory=tmpdir, + name="nsa-target", + creator_rank=0, + validate_production=False, + ) + index_allocator = mph.SharedHostTensorAllocator( + directory=tmpdir, + name="nsa-index", + creator_rank=0, + validate_production=False, + ) + host_pool = mph.NSATokenToKVPoolHost( + device_pool=device_pool, + host_to_device_ratio=2.0, + host_size=0, + page_size=4, + layout="layer_page_first", + pin_memory=False, + device="cpu", + host_token_capacity=16, + host_tensor_allocator=main_allocator, + index_host_tensor_allocator=index_allocator, + ) + + self.assertEqual(tuple(host_pool.kv_buffer.shape), (3, 4, 4, 1, 24)) + self.assertEqual( + tuple(host_pool.index_k_with_scale_buffer.shape), + (3, host_pool.indexer_page_num, 1, host_pool.indexer_page_stride_size), + ) + self.assertNotEqual(main_allocator.slab_handle.path, index_allocator.slab_handle.path) + self.assertEqual(main_allocator.slab_handle.shape, tuple(host_pool.kv_buffer.shape)) + self.assertEqual(index_allocator.slab_handle.shape, tuple(host_pool.index_k_with_scale_buffer.shape)) + main_allocator.release_tensor_view_for_close() + index_allocator.release_tensor_view_for_close() + del host_pool + main_allocator.close() + index_allocator.close() + + def test_broadcasting_shared_allocator_creates_once_and_nonzero_rank_attaches_handle(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + broadcast_state = {} + calls = [] + + def fake_broadcast(handle, *, cp_cpu_group, rank, source_rank): + calls.append((rank, source_rank, cp_cpu_group, handle is not None)) + if rank == source_rank: + broadcast_state["handle"] = handle + return handle + return broadcast_state["handle"] + + with tempfile.TemporaryDirectory() as tmpdir: + rank0 = mph.SharedHostTensorAllocator( + directory=tmpdir, + name="broadcast-target", + creator_rank=0, + validate_production=False, + rank=0, + source_rank=0, + cp_cpu_group="fake-cp-cpu", + broadcast_handle_fn=fake_broadcast, + ) + t0 = rank0.allocate((4,), dtype=torch.int16, device="cpu") + t0[2] = 77 + rank1 = mph.SharedHostTensorAllocator( + directory=tmpdir, + name="ignored-on-attach", + creator_rank=0, + validate_production=False, + rank=1, + source_rank=0, + cp_cpu_group="fake-cp-cpu", + broadcast_handle_fn=fake_broadcast, + ) + t1 = rank1.allocate((4,), dtype=torch.int16, device="cpu") + + self.assertEqual(t1[2].item(), 77) + self.assertEqual(calls, [(0, 0, "fake-cp-cpu", True), (1, 0, "fake-cp-cpu", False)]) + self.assertEqual(rank0.slab_handle.path, rank1.slab_handle.path) + rank1.release_tensor_view_for_close() + rank0.release_tensor_view_for_close() + del t1 + del t0 + rank1.close() + rank0.close() + + def test_shared_allocator_early_unlinks_after_all_ranks_register_and_stays_live(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + broadcast_state = {} + unlink_results = {} + calls = [] + + def fake_broadcast(handle, *, cp_cpu_group, rank, source_rank): + if rank == source_rank: + broadcast_state["handle"] = handle + return handle + return broadcast_state["handle"] + + def fake_gather(output, value, group): + calls.append(("gather", value["rank"], value["success"], group)) + output.extend( + [ + {"rank": 0, "success": True, "error": None}, + {"rank": 1, "success": True, "error": None}, + ] + ) + + def fake_unlink_broadcast(object_list, src, group): + calls.append(("broadcast-unlink", src, group, list(object_list))) + if object_list[0] is not None: + unlink_results["value"] = object_list[0] + else: + object_list[0] = unlink_results["value"] + + with tempfile.TemporaryDirectory() as tmpdir: + rank0 = mph.SharedHostTensorAllocator( + directory=tmpdir, + name="early-unlink-target", + creator_rank=0, + validate_production=False, + rank=0, + source_rank=0, + cp_cpu_group="fake-cp-cpu", + broadcast_handle_fn=fake_broadcast, + registration_status_gather_fn=fake_gather, + unlink_result_broadcast_fn=fake_unlink_broadcast, + ) + t0 = rank0.allocate((4,), dtype=torch.int16, device="cpu") + rank1 = mph.SharedHostTensorAllocator( + directory=tmpdir, + name="ignored-on-attach", + creator_rank=0, + validate_production=False, + rank=1, + source_rank=0, + cp_cpu_group="fake-cp-cpu", + broadcast_handle_fn=fake_broadcast, + registration_status_gather_fn=fake_gather, + unlink_result_broadcast_fn=fake_unlink_broadcast, + ) + t1 = rank1.allocate((4,), dtype=torch.int16, device="cpu") + path = Path(rank0.slab_handle.path) + self.assertTrue(path.exists()) + + rank0.complete_cuda_host_register(ptr=t0.data_ptr(), nbytes=t0.numel() * t0.element_size()) + rank1.complete_cuda_host_register(ptr=t1.data_ptr(), nbytes=t1.numel() * t1.element_size()) + + self.assertFalse(path.exists()) + t0[1] = 314 + self.assertEqual(t1[1].item(), 314) + t1[2] = 271 + self.assertEqual(t0[2].item(), 271) + self.assertEqual( + [call[0] for call in calls], + ["gather", "broadcast-unlink", "gather", "broadcast-unlink"], + ) + + rank1.release_tensor_view_for_close() + rank0.release_tensor_view_for_close() + del t1 + del t0 + rank1.close() + rank0.close() + + def test_shared_allocator_production_validation_runs_before_early_unlink(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + page_checks = [] + + def create_without_real_hugetlbfs(**kwargs): + return _cp_shared_l2_pool.create_cp_shared_host_slab( + directory=kwargs["directory"], + name=kwargs["name"], + nbytes=kwargs["nbytes"], + shape=kwargs["shape"], + dtype_name=kwargs["dtype_name"], + creator_rank=kwargs["creator_rank"], + validate_production=False, + ) + + def checker(path=None): + page_checks.append((path, Path(path).exists())) + return True + + with tempfile.TemporaryDirectory() as tmpdir: + allocator = mph.SharedHostTensorAllocator( + directory=tmpdir, + name="validate-before-unlink", + creator_rank=0, + validate_production=True, + create_fn=create_without_real_hugetlbfs, + page_size_checker=checker, + rank=0, + source_rank=0, + ) + tensor = allocator.allocate((4,), dtype=torch.int16, device="cpu") + path = Path(allocator.slab_handle.path) + self.assertTrue(path.exists()) + + allocator.complete_cuda_host_register( + ptr=tensor.data_ptr(), + nbytes=tensor.numel() * tensor.element_size(), + ) + + self.assertEqual(page_checks, [(str(path), True)]) + self.assertFalse(path.exists()) + allocator.release_tensor_view_for_close() + del tensor + allocator.close() + + def test_shared_allocator_production_validation_failure_prevents_unlink(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + + def create_without_real_hugetlbfs(**kwargs): + return _cp_shared_l2_pool.create_cp_shared_host_slab( + directory=kwargs["directory"], + name=kwargs["name"], + nbytes=kwargs["nbytes"], + shape=kwargs["shape"], + dtype_name=kwargs["dtype_name"], + creator_rank=kwargs["creator_rank"], + validate_production=False, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + allocator = mph.SharedHostTensorAllocator( + directory=tmpdir, + name="validate-fail-keeps-path", + creator_rank=0, + validate_production=True, + create_fn=create_without_real_hugetlbfs, + page_size_checker=lambda path=None: False, + rank=0, + source_rank=0, + ) + tensor = allocator.allocate((4,), dtype=torch.int16, device="cpu") + path = Path(allocator.slab_handle.path) + + with self.assertRaisesRegex(RuntimeError, "2MiB effective page-size"): + allocator.complete_cuda_host_register( + ptr=tensor.data_ptr(), + nbytes=tensor.numel() * tensor.element_size(), + ) + + self.assertTrue(path.exists()) + allocator.release_tensor_view_for_close() + del tensor + allocator.close() + + def test_shared_allocator_does_not_unlink_when_any_rank_register_fails(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + broadcast_state = {} + + def fake_broadcast(handle, *, cp_cpu_group, rank, source_rank): + if rank == source_rank: + broadcast_state["handle"] = handle + return handle + return broadcast_state["handle"] + + def fake_gather(output, value, group): + output.extend( + [ + {"rank": 0, "success": True, "error": None}, + {"rank": 1, "success": False, "error": "cudaHostRegister failed"}, + ] + ) + + with tempfile.TemporaryDirectory() as tmpdir: + rank0 = mph.SharedHostTensorAllocator( + directory=tmpdir, + name="early-unlink-fail", + creator_rank=0, + validate_production=False, + rank=0, + source_rank=0, + cp_cpu_group="fake-cp-cpu", + broadcast_handle_fn=fake_broadcast, + registration_status_gather_fn=fake_gather, + ) + t0 = rank0.allocate((4,), dtype=torch.int16, device="cpu") + path = Path(rank0.slab_handle.path) + + with self.assertRaisesRegex(RuntimeError, "rank 1.*cudaHostRegister failed"): + rank0.complete_cuda_host_register( + ptr=t0.data_ptr(), + nbytes=t0.numel() * t0.element_size(), + ) + + self.assertTrue(path.exists()) + rank0.release_tensor_view_for_close() + del t0 + rank0.close() + + def test_alloc_with_host_register_notifies_shared_allocator_before_return(self): + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + calls = [] + + class FakeCudart: + def cudaHostRegister(self, ptr, nbytes, flags): + calls.append(("register", ptr, nbytes, flags)) + return 0 + + def cudaHostUnregister(self, ptr): + calls.append(("unregister", ptr)) + return 0 + + def cudaGetErrorString(self, err): + return b"ok" + + class RecordingAllocator(mph.HostTensorAllocator): + def __init__(self): + super().__init__() + self.finalized = [] + + def complete_cuda_host_register( + self, *, ptr, nbytes, success=True, error_message=None + ): + self.finalized.append( + { + "ptr": ptr, + "nbytes": nbytes, + "success": success, + "error": error_message, + } + ) + + fake_cudart = FakeCudart() + allocator = RecordingAllocator() + with patch.object(mph.torch.cuda, "cudart", return_value=fake_cudart): + tensor = mph.alloc_with_host_register( + (4,), torch.int16, "cpu", True, allocator + ) + + self.assertEqual(calls[0], ("register", tensor.data_ptr(), 8, 0)) + self.assertEqual( + allocator.finalized, + [ + { + "ptr": tensor.data_ptr(), + "nbytes": 8, + "success": True, + "error": None, + } + ], + ) + + def test_group_allocator_splits_logical_shape_and_uses_deterministic_slab_names(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + slabs = ( + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=2, global_base_page=0, num_pages=3), + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=5, global_base_page=3, num_pages=2), + ) + + with tempfile.TemporaryDirectory() as tmpdir: + allocator = mph.SharedHostTensorGroupAllocator( + slabs=slabs, + directory=tmpdir, + name="target-kv", + creator_rank=0, + validate_production=False, + ) + group = allocator.allocate_group( + (4, 5, 2), dtype=torch.int16, device="cpu", page_dim=1 + ) + + self.assertIsInstance(group, mph.SharedHostTensorGroup) + self.assertEqual([tuple(t.shape) for t in group.tensors], [(4, 3, 2), (4, 2, 2)]) + self.assertEqual(group.slab_infos, slabs) + self.assertIs(group.tensor_for_slab(2), group.tensors[0]) + self.assertIs(group.tensor_for_slab(5), group.tensors[1]) + self.assertEqual( + [Path(handle.path).name for handle in group.handles], + ["target-kv-slab2", "target-kv-slab5"], + ) + self.assertNotEqual(group.handles[0].path, group.handles[1].path) + self.assertEqual( + [rng.nbytes for rng in group.registration_ranges], + [group.tensors[0].numel() * group.tensors[0].element_size(), group.tensors[1].numel() * group.tensors[1].element_size()], + ) + + group.release_tensor_view_for_close() + group.close() + + def test_group_close_requires_release_and_release_invalidates_tensor_access(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + slabs = ( + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=0, global_base_page=0, num_pages=2), + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=1, global_base_page=2, num_pages=2), + ) + + with tempfile.TemporaryDirectory() as tmpdir: + allocator = mph.SharedHostTensorGroupAllocator( + slabs=slabs, + directory=tmpdir, + name="group-release", + creator_rank=0, + validate_production=False, + ) + group = allocator.allocate_group( + (4,), dtype=torch.int16, device="cpu", page_dim=0 + ) + self.assertEqual(len(group.slab_views), 2) + self.assertEqual(len(group.tensors), 2) + self.assertEqual(tuple(group.tensor_for_slab(0).shape), (2,)) + + with self.assertRaisesRegex(RuntimeError, "release_tensor_view_for_close"): + group.close() + + group.release_tensor_view_for_close() + + with self.assertRaisesRegex(RuntimeError, "released for close"): + _ = group.slab_views + with self.assertRaisesRegex(RuntimeError, "released for close"): + _ = group.tensors + with self.assertRaisesRegex(RuntimeError, "released for close"): + group.tensor_for_slab(0) + + # Metadata and allocator-owned handles remain inspectable after release. + self.assertEqual(group.slab_infos, slabs) + self.assertEqual(len(group.handles), 2) + group.close() + group.close() + + with self.assertRaisesRegex(RuntimeError, "after close"): + _ = group.tensors + + def test_group_release_drops_allocator_tensor_refs_before_close(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + slabs = ( + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=0, global_base_page=0, num_pages=1), + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=1, global_base_page=1, num_pages=1), + ) + + with tempfile.TemporaryDirectory() as tmpdir: + allocator = mph.SharedHostTensorGroupAllocator( + slabs=slabs, + directory=tmpdir, + name="group-release-refs", + creator_rank=0, + validate_production=False, + ) + group = allocator.allocate_group( + (2,), dtype=torch.int16, device="cpu", page_dim=0 + ) + group.tensors[0][0] = 11 + + group.release_tensor_view_for_close() + + self.assertTrue(all(slab_allocator.tensor is None for slab_allocator in allocator.allocators)) + with self.assertRaisesRegex(RuntimeError, "released for close"): + group.tensor_for_slab(1) + group.close() + + def test_group_allocator_fake_broadcast_attaches_all_slabs_and_shares_writes(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + slabs = ( + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=0, global_base_page=0, num_pages=2), + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=1, global_base_page=2, num_pages=3), + ) + broadcast_state = {"handles": [], "next_attach": 0} + calls = [] + + def fake_broadcast(handle, *, cp_cpu_group, rank, source_rank): + calls.append((rank, handle is not None)) + if rank == source_rank: + broadcast_state["handles"].append(handle) + return handle + handle = broadcast_state["handles"][broadcast_state["next_attach"]] + broadcast_state["next_attach"] += 1 + return handle + + with tempfile.TemporaryDirectory() as tmpdir: + rank0_allocator = mph.SharedHostTensorGroupAllocator( + slabs=slabs, + directory=tmpdir, + name="broadcast-group", + creator_rank=0, + validate_production=False, + rank=0, + source_rank=0, + cp_cpu_group="fake-cp-cpu", + broadcast_handle_fn=fake_broadcast, + ) + rank0_group = rank0_allocator.allocate_group( + (5,), dtype=torch.int32, device="cpu", page_dim=0 + ) + rank0_group.tensor_for_slab(0)[1] = 101 + rank0_group.tensor_for_slab(1)[2] = 202 + + rank1_allocator = mph.SharedHostTensorGroupAllocator( + slabs=slabs, + directory=tmpdir, + name="ignored-on-attach", + creator_rank=0, + validate_production=False, + rank=1, + source_rank=0, + cp_cpu_group="fake-cp-cpu", + broadcast_handle_fn=fake_broadcast, + ) + rank1_group = rank1_allocator.allocate_group( + (5,), dtype=torch.int32, device="cpu", page_dim=0 + ) + + self.assertEqual(rank1_group.tensor_for_slab(0)[1].item(), 101) + self.assertEqual(rank1_group.tensor_for_slab(1)[2].item(), 202) + rank1_group.tensor_for_slab(0)[0] = 303 + rank1_group.tensor_for_slab(1)[1] = 404 + self.assertEqual(rank0_group.tensor_for_slab(0)[0].item(), 303) + self.assertEqual(rank0_group.tensor_for_slab(1)[1].item(), 404) + self.assertEqual(calls, [(0, True), (0, True), (1, False), (1, False)]) + self.assertEqual( + [r0.path for r0 in rank0_group.handles], + [r1.path for r1 in rank1_group.handles], + ) + + rank1_group.release_tensor_view_for_close() + rank0_group.release_tensor_view_for_close() + del rank1_group + del rank0_group + rank1_allocator.allocators[0].close() + rank1_allocator.allocators[1].close() + rank0_allocator.allocators[0].close() + rank0_allocator.allocators[1].close() + + def test_group_complete_cuda_host_register_early_unlinks_all_slabs_after_all_rank_success(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + slabs = ( + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=0, global_base_page=0, num_pages=2), + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=1, global_base_page=2, num_pages=2), + ) + broadcast_state = {"handles": [], "next_attach": 0} + unlink_results = [] + calls = [] + + def fake_broadcast(handle, *, cp_cpu_group, rank, source_rank): + if rank == source_rank: + broadcast_state["handles"].append(handle) + return handle + handle = broadcast_state["handles"][broadcast_state["next_attach"]] + broadcast_state["next_attach"] += 1 + return handle + + def fake_gather(output, value, group): + calls.append(("gather", value["rank"], Path(value["path"]).name, value["success"])) + output.extend( + [ + {"rank": 0, "success": True, "error": None}, + {"rank": 1, "success": True, "error": None}, + ] + ) + + def fake_unlink_broadcast(object_list, src, group): + calls.append(("broadcast-unlink", src, group, object_list[0])) + if object_list[0] is not None: + unlink_results.append(object_list[0]) + else: + object_list[0] = unlink_results.pop(0) + + with tempfile.TemporaryDirectory() as tmpdir: + rank0_allocator = mph.SharedHostTensorGroupAllocator( + slabs=slabs, + directory=tmpdir, + name="early-unlink-group", + creator_rank=0, + validate_production=False, + rank=0, + source_rank=0, + cp_cpu_group="fake-cp-cpu", + broadcast_handle_fn=fake_broadcast, + registration_status_gather_fn=fake_gather, + unlink_result_broadcast_fn=fake_unlink_broadcast, + ) + rank0_group = rank0_allocator.allocate_group( + (4,), dtype=torch.int16, device="cpu", page_dim=0 + ) + rank1_allocator = mph.SharedHostTensorGroupAllocator( + slabs=slabs, + directory=tmpdir, + name="ignored-on-attach", + creator_rank=0, + validate_production=False, + rank=1, + source_rank=0, + cp_cpu_group="fake-cp-cpu", + broadcast_handle_fn=fake_broadcast, + registration_status_gather_fn=fake_gather, + unlink_result_broadcast_fn=fake_unlink_broadcast, + ) + rank1_group = rank1_allocator.allocate_group( + (4,), dtype=torch.int16, device="cpu", page_dim=0 + ) + paths = [Path(handle.path) for handle in rank0_group.handles] + self.assertTrue(all(path.exists() for path in paths)) + + rank0_group.complete_cuda_host_register() + rank1_group.complete_cuda_host_register() + + self.assertTrue(all(not path.exists() for path in paths)) + rank0_group.tensor_for_slab(0)[1] = 314 + rank0_group.tensor_for_slab(1)[0] = 271 + self.assertEqual(rank1_group.tensor_for_slab(0)[1].item(), 314) + self.assertEqual(rank1_group.tensor_for_slab(1)[0].item(), 271) + self.assertEqual( + [call[0] for call in calls], + ["gather", "gather", "broadcast-unlink", "gather", "gather", "broadcast-unlink"], + ) + + rank1_group.release_tensor_view_for_close() + rank0_group.release_tensor_view_for_close() + del rank1_group + del rank0_group + rank1_allocator.allocators[0].close() + rank1_allocator.allocators[1].close() + rank0_allocator.allocators[0].close() + rank0_allocator.allocators[1].close() + + def test_group_failed_registration_keeps_slab_paths_and_raises(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + slabs = ( + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=0, global_base_page=0, num_pages=2), + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=1, global_base_page=2, num_pages=2), + ) + + def create_without_real_hugetlbfs(**kwargs): + return _cp_shared_l2_pool.create_cp_shared_host_slab( + directory=kwargs["directory"], + name=kwargs["name"], + nbytes=kwargs["nbytes"], + shape=kwargs["shape"], + dtype_name=kwargs["dtype_name"], + creator_rank=kwargs["creator_rank"], + validate_production=False, + ) + + def checker(path=None): + return not str(path).endswith("slab1") + + with tempfile.TemporaryDirectory() as tmpdir: + allocator = mph.SharedHostTensorGroupAllocator( + slabs=slabs, + directory=tmpdir, + name="validate-group", + creator_rank=0, + validate_production=True, + create_fn=create_without_real_hugetlbfs, + page_size_checker=checker, + rank=0, + source_rank=0, + ) + group = allocator.allocate_group( + (4,), dtype=torch.int16, device="cpu", page_dim=0 + ) + paths = [Path(handle.path) for handle in group.handles] + + with self.assertRaisesRegex(RuntimeError, "2MiB effective page-size"): + group.complete_cuda_host_register() + + self.assertTrue(all(path.exists() for path in paths)) + group.release_tensor_view_for_close() + del group + allocator.allocators[0].close() + allocator.allocators[1].close() + + def test_group_pin_memory_validation_failure_does_not_unlink_prior_slab(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + slabs = ( + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=0, global_base_page=0, num_pages=2), + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=1, global_base_page=2, num_pages=2), + ) + calls = [] + + class FakeCudart: + def cudaHostRegister(self, ptr, nbytes, flags): + calls.append(("register", ptr, nbytes, flags)) + return 0 + + def cudaHostUnregister(self, ptr): + calls.append(("unregister", ptr)) + return 0 + + def cudaGetErrorString(self, err): + return b"ok" + + def create_without_real_hugetlbfs(**kwargs): + return _cp_shared_l2_pool.create_cp_shared_host_slab( + directory=kwargs["directory"], + name=kwargs["name"], + nbytes=kwargs["nbytes"], + shape=kwargs["shape"], + dtype_name=kwargs["dtype_name"], + creator_rank=kwargs["creator_rank"], + validate_production=False, + ) + + def checker(path=None): + return not str(path).endswith("slab1") + + with tempfile.TemporaryDirectory() as tmpdir: + allocator = mph.SharedHostTensorGroupAllocator( + slabs=slabs, + directory=tmpdir, + name="pin-fail-group", + creator_rank=0, + validate_production=True, + create_fn=create_without_real_hugetlbfs, + page_size_checker=checker, + rank=0, + source_rank=0, + ) + + with patch.object(mph.torch.cuda, "cudart", return_value=FakeCudart()): + with self.assertRaisesRegex(RuntimeError, "2MiB effective page-size"): + allocator.allocate_group( + (4,), + dtype=torch.int16, + device="cpu", + page_dim=0, + pin_memory=True, + ) + + paths = [Path(slab_allocator.slab_handle.path) for slab_allocator in allocator.allocators] + self.assertTrue(all(path.exists() for path in paths)) + self.assertEqual([call[0] for call in calls], ["register", "register", "unregister", "unregister"]) + + for slab_allocator in allocator.allocators: + slab_allocator.release_tensor_view_for_close() + slab_allocator.close() + + def test_group_pin_memory_actual_unlink_failure_restores_original_paths(self): + import importlib + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + cp_pool_module = importlib.import_module("sglang.srt.mem_cache.cp_shared_l2_pool") + slabs = ( + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=0, global_base_page=0, num_pages=2), + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=1, global_base_page=2, num_pages=2), + ) + calls = [] + + class FakeCudart: + def cudaHostRegister(self, ptr, nbytes, flags): + calls.append(("register", ptr, nbytes, flags)) + return 0 + + def cudaHostUnregister(self, ptr): + calls.append(("unregister", ptr)) + return 0 + + def cudaGetErrorString(self, err): + return b"ok" + + real_unlink = cp_pool_module.os.unlink + + def fail_second_original_unlink(path): + path = Path(path) + if path.name == "actual-unlink-fail-group-slab1": + raise RuntimeError("injected actual unlink failure") + return real_unlink(path) + + with tempfile.TemporaryDirectory() as tmpdir: + allocator = mph.SharedHostTensorGroupAllocator( + slabs=slabs, + directory=tmpdir, + name="actual-unlink-fail-group", + creator_rank=0, + validate_production=False, + rank=0, + source_rank=0, + ) + + with patch.object(mph.torch.cuda, "cudart", return_value=FakeCudart()): + with patch.object(cp_pool_module.os, "unlink", side_effect=fail_second_original_unlink): + with self.assertRaisesRegex(RuntimeError, "group early unlink failed"): + allocator.allocate_group( + (4,), + dtype=torch.int16, + device="cpu", + page_dim=0, + pin_memory=True, + ) + + paths = [ + Path(slab_allocator.slab_handle.path) + for slab_allocator in allocator.allocators + ] + self.assertTrue(all(path.exists() for path in paths)) + self.assertEqual( + [call[0] for call in calls], + ["register", "register", "unregister", "unregister"], + ) + self.assertFalse( + any( + slab_allocator._cuda_host_register_finalized + for slab_allocator in allocator.allocators + ) + ) + + for slab_allocator in allocator.allocators: + slab_allocator.release_tensor_view_for_close() + slab_allocator.close() + + def test_group_pin_memory_unlink_preflight_failure_does_not_partial_unlink(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + slabs = ( + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=0, global_base_page=0, num_pages=2), + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=1, global_base_page=2, num_pages=2), + ) + calls = [] + + class FakeCudart: + def cudaHostRegister(self, ptr, nbytes, flags): + calls.append(("register", ptr, nbytes, flags)) + return 0 + + def cudaHostUnregister(self, ptr): + calls.append(("unregister", ptr)) + return 0 + + def cudaGetErrorString(self, err): + return b"ok" + + def create_with_second_slab_preflight_failure(**kwargs): + mapping = _cp_shared_l2_pool.create_cp_shared_host_slab( + directory=kwargs["directory"], + name=kwargs["name"], + nbytes=kwargs["nbytes"], + shape=kwargs["shape"], + dtype_name=kwargs["dtype_name"], + creator_rank=kwargs["creator_rank"], + validate_production=False, + ) + if kwargs["name"].endswith("slab1"): + def fail_preflight(): + raise RuntimeError("injected unlink readiness failure") + + mapping.check_unlink_ready = fail_preflight + return mapping + + with tempfile.TemporaryDirectory() as tmpdir: + allocator = mph.SharedHostTensorGroupAllocator( + slabs=slabs, + directory=tmpdir, + name="preflight-fail-group", + creator_rank=0, + validate_production=False, + create_fn=create_with_second_slab_preflight_failure, + rank=0, + source_rank=0, + ) + + with patch.object(mph.torch.cuda, "cudart", return_value=FakeCudart()): + with self.assertRaisesRegex(RuntimeError, "group early unlink failed"): + allocator.allocate_group( + (4,), + dtype=torch.int16, + device="cpu", + page_dim=0, + pin_memory=True, + ) + + paths = [ + Path(slab_allocator.slab_handle.path) + for slab_allocator in allocator.allocators + ] + self.assertTrue(all(path.exists() for path in paths)) + self.assertEqual( + [call[0] for call in calls], + ["register", "register", "unregister", "unregister"], + ) + self.assertFalse( + any( + slab_allocator._cuda_host_register_finalized + for slab_allocator in allocator.allocators + ) + ) + + for slab_allocator in allocator.allocators: + slab_allocator.release_tensor_view_for_close() + slab_allocator.close() + + def test_group_allocator_rejects_shape_page_dim_mismatch_and_invalid_inputs(self): + import tempfile + import torch + + mph = _load_memory_pool_host_module_for_patch_d() + slabs = ( + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=0, global_base_page=0, num_pages=2), + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=1, global_base_page=2, num_pages=3), + ) + + with tempfile.TemporaryDirectory() as tmpdir: + allocator = mph.SharedHostTensorGroupAllocator( + slabs=slabs, + directory=tmpdir, + name="bad-shape", + creator_rank=0, + validate_production=False, + ) + with self.assertRaisesRegex(ValueError, "total slab pages"): + allocator.allocate_group((4,), dtype=torch.int16, device="cpu", page_dim=0) + with self.assertRaisesRegex(ValueError, "page_dim"): + allocator.allocate_group((5,), dtype=torch.int16, device="cpu", page_dim=1) + with self.assertRaisesRegex(ValueError, "CPU tensor"): + allocator.allocate_group((5,), dtype=torch.int16, device="cuda", page_dim=0) + + duplicate_slabs = ( + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=0, global_base_page=0, num_pages=2), + CpSharedL2SlabInfo(PAYLOAD_TARGET_KV, slab_id=0, global_base_page=2, num_pages=3), + ) + with self.assertRaisesRegex(ValueError, "duplicate.*slab_id"): + mph.SharedHostTensorGroupAllocator( + slabs=duplicate_slabs, + directory=tmpdir, + name="duplicate", + creator_rank=0, + validate_production=False, + ) + with self.assertRaisesRegex(ValueError, "non-empty slabs"): + mph.SharedHostTensorGroupAllocator( + slabs=(), + directory=tmpdir, + name="empty", + creator_rank=0, + validate_production=False, + ) + + def test_transfer_api_signatures_remain_unchanged(self): + import inspect + + mph = _load_memory_pool_host_module_for_patch_d() + expected = ["self", "device_pool", "host_indices", "device_indices", "layer_id", "io_backend"] + for cls in (mph.MHATokenToKVPoolHost, mph.MLATokenToKVPoolHost, mph.NSATokenToKVPoolHost): + self.assertEqual(list(inspect.signature(cls.backup_from_device_per_layer).parameters), expected) + self.assertEqual(list(inspect.signature(cls.load_to_device_per_layer).parameters), expected) + + +if __name__ == "__main__": + unittest.main()