diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index ea5f78d4e7aa..7a7f7e5ffcd0 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -780,6 +780,12 @@ Please consult the documentation below and [server_args.py](https://github.com/s `False` bool flag (set to enable) + + `--python-thread-switch-interval` + Optional positive, finite Python thread switch interval in seconds for each spawned scheduler process. Affects all Python threads in that process. Unset preserves interpreter defaults; CPU affinity and thread priority are unchanged. + `None` + Type: float + `--batch-notify-size` Number of streaming notifications to batch before yielding to the event loop. Reduces asyncio wakeup overhead under high concurrency. diff --git a/python/sglang/srt/arg_groups/field_order.py b/python/sglang/srt/arg_groups/field_order.py index cd77b9ebb305..bbe67ac38d89 100644 --- a/python/sglang/srt/arg_groups/field_order.py +++ b/python/sglang/srt/arg_groups/field_order.py @@ -521,5 +521,18 @@ "enable_decoder_swa_bounded_replay", "enable_response_store", "disaggregation_decode_host_receive_threshold", + # Freeze the previously appended fields before adding another namespace leaf. + "qsa_indexer_dtype", + "disaggregation_decode_allocation_policy", + "speculative_boundary_reduction", + "speculative_use_block_verification", + "speculative_draft_sink_size", + "disable_aiter_allreduce_fusion_in_prefill", + "disable_aiter_allreduce_fusion_in_decode", + "dsa_triton_union", + "enable_dsa_fused_indexer", + "enable_cp_tp_group_sharing", + "grpc_response_timeout_secs", + "python_thread_switch_interval", ) # fmt: on diff --git a/python/sglang/srt/arg_groups/fields/device.py b/python/sglang/srt/arg_groups/fields/device.py index 614860cc9d79..9f709881a17b 100644 --- a/python/sglang/srt/arg_groups/fields/device.py +++ b/python/sglang/srt/arg_groups/fields/device.py @@ -9,7 +9,6 @@ ) import msgspec - from sglang.srt.arg_groups.arg_utils import A, Derived @@ -72,6 +71,10 @@ class Device(msgspec.Struct): "Set soft watchdog timeout in seconds. If a forward batch takes longer than this, the server will dump information for debugging.", ] = None sleep_on_idle: A[bool, "Reduce CPU usage when sglang is idle."] = False + python_thread_switch_interval: A[ + Optional[float], + "Optional positive, finite Python thread switch interval in seconds for each scheduler process. Affects all Python threads in that process; unset preserves interpreter defaults. Does not change CPU affinity or thread priority.", + ] = None use_ray: A[ bool, "Use Ray actors for scheduler process management.", diff --git a/python/sglang/srt/arg_groups/fields/memory.py b/python/sglang/srt/arg_groups/fields/memory.py index c79a33346a85..af98341b88f2 100644 --- a/python/sglang/srt/arg_groups/fields/memory.py +++ b/python/sglang/srt/arg_groups/fields/memory.py @@ -145,11 +145,12 @@ class Memory(msgspec.Struct): hicache_storage_backend: A[ Optional[str], Arg( - help="The storage backend for hierarchical KV cache. Built-in backends: file, mooncake, npu_memcache, hf3fs, nixl, aibrix, tensorcast, seaweedfs. For dynamic backend, use --hicache-storage-backend-extra-config to specify: backend_name (custom name), module_path (Python module path), class_name (backend class name).", + help="The storage backend for hierarchical KV cache. Built-in backends: file, mooncake, kvcr, npu_memcache, hf3fs, nixl, aibrix, tensorcast, seaweedfs. For dynamic backend, use --hicache-storage-backend-extra-config to specify: backend_name (custom name), module_path (Python module path), class_name (backend class name).", choices=[ "file", "sim", "mooncake", + "kvcr", "npu_memcache", "hf3fs", "nixl", diff --git a/python/sglang/srt/arg_groups/hicache_hook.py b/python/sglang/srt/arg_groups/hicache_hook.py index 901fa2b14d70..8ad77b730a78 100644 --- a/python/sglang/srt/arg_groups/hicache_hook.py +++ b/python/sglang/srt/arg_groups/hicache_hook.py @@ -179,7 +179,7 @@ def resolve_layout_io_compatibility(server_args: Any): def resolve_storage_layout_compatibility(server_args: Any): cfg = resolving_view(server_args) if ( - cfg.hicache_storage_backend not in ("mooncake", "npu_memcache") + cfg.hicache_storage_backend not in ("mooncake", "npu_memcache", "kvcr") or cfg.hicache_mem_layout != "layer_first" ): return diff --git a/python/sglang/srt/arg_groups/pipeline.py b/python/sglang/srt/arg_groups/pipeline.py index 0ff314930f04..f3fabe5b3ac5 100644 --- a/python/sglang/srt/arg_groups/pipeline.py +++ b/python/sglang/srt/arg_groups/pipeline.py @@ -79,11 +79,13 @@ def run_resolution_pipeline(server_args: Any) -> None: default_unset_prefill_decode_interval, validate_experimental_sgl_marlin, validate_prefill_decode_interval, + validate_python_thread_scheduling, validate_response_store, validate_sampling_mask_max_tokens, ) run_hook(validate_prefill_decode_interval, server_args) + run_hook(validate_python_thread_scheduling, server_args) run_hook(validate_response_store, server_args) run_hook(validate_sampling_mask_max_tokens, server_args) diff --git a/python/sglang/srt/arg_groups/validation_hook.py b/python/sglang/srt/arg_groups/validation_hook.py index b5785fa8e6de..11649f02a2a8 100644 --- a/python/sglang/srt/arg_groups/validation_hook.py +++ b/python/sglang/srt/arg_groups/validation_hook.py @@ -25,6 +25,16 @@ logger = logging.getLogger(__name__) +def validate_python_thread_scheduling(server_args: Any) -> None: + from sglang.srt.managers.scheduler_components.thread_scheduling import ( + validate_python_thread_switch_interval, + ) + + validate_python_thread_switch_interval( + resolving_view(server_args).python_thread_switch_interval + ) + + def validate_response_store(server_args: Any) -> None: cfg = resolving_view(server_args) if cfg.enable_response_store and cfg.disaggregation_mode != "null": diff --git a/python/sglang/srt/disaggregation/decode_hicache_mixin.py b/python/sglang/srt/disaggregation/decode_hicache_mixin.py index 713737b5529f..7b8103b5f015 100644 --- a/python/sglang/srt/disaggregation/decode_hicache_mixin.py +++ b/python/sglang/srt/disaggregation/decode_hicache_mixin.py @@ -95,6 +95,7 @@ def _build_decode_prefix_match( prefix_keys, extra_key=req.extra_key, cache_salt=req.cache_salt, + kv_hints=req.kv_hints, ) return DecodePrefixMatch( @@ -139,6 +140,7 @@ def _start_hicache_prefetch( prefix_keys, extra_key=req.extra_key, cache_salt=req.cache_salt, + kv_hints=req.kv_hints, ) prefix_match.prefetch_registered = self.tree_cache.has_ongoing_prefetch( req.cache_request_handle diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 8a27dbf38ba7..3e7cae4c0641 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -22,11 +22,11 @@ from typing import TYPE_CHECKING, Callable, List, NamedTuple, Optional import torch - from sglang.srt.mem_cache.hicache_storage import ( STORAGE_BATCH_SIZE, HiCacheStorageConfig, HiCacheStorageExtraInfo, + PoolHitPolicy, PoolName, PoolTransfer, count_pool_hits, @@ -234,6 +234,7 @@ def __init__( last_hash: Optional[str] = None, hash_value: Optional[List[str]] = None, prefix_keys: Optional[List[str]] = None, + kv_hints=None, ): self.host_indices = host_indices self.token_ids = token_ids @@ -241,6 +242,7 @@ def __init__( self.completed_tokens = 0 self.hash_value = hash_value if hash_value is not None else [] self.prefix_keys = prefix_keys + self.kv_hints = kv_hints # Full queried page-hash chain, set by _storage_hit_query before # hash_value is truncated to the hit boundary; the tail is the # absence signal that invalidates buffer-mode existence beliefs. @@ -562,13 +564,6 @@ def attach_storage_backend( self.storage_config = self._generate_storage_config( model_name, storage_backend_extra_config ) - # for MLA models, only one rank needs to backup the KV cache - self.backup_skip = ( - self.storage_config.is_mla_model - # todo: load balancing - and self.storage_config.tp_rank != 0 - ) - # Use storage backend factory for dynamic backend creation from sglang.srt.mem_cache.storage import StorageBackendFactory @@ -577,6 +572,13 @@ def attach_storage_backend( storage_backend, self.storage_config, self.storage_host_pool ) self.storage_backend.register_mem_pool_host(self.storage_host_pool) + # Rank-local stores need every MLA writer. A backend that explicitly + # provides shared replicated storage can use the TP0-only path. + self.backup_skip = ( + self.storage_config.is_mla_model + and self.storage_config.tp_rank != 0 + and self.storage_backend.requires_rank_local_backup is not True + ) self.enable_storage = True # todo: threshold policy for prefetching @@ -608,6 +610,7 @@ def attach_storage_backend( in [ "hf3fs", "mooncake", + "kvcr", "npu_memcache", "eic", "nixl", @@ -1094,6 +1097,7 @@ def _generic_page_get( def _page_transfer(self, operation: PrefetchOperation) -> int: # Transfer batch by batch + batch_pages = self.storage_backend.prefetch_batch_pages prefix_keys = operation.prefix_keys kv_derived_transfers = [ transfer @@ -1102,20 +1106,23 @@ def _page_transfer(self, operation: PrefetchOperation) -> int: ] all_success = True completed_pages = 0 - for i in range(0, len(operation.hash_value), STORAGE_BATCH_SIZE): + for i in range(0, len(operation.hash_value), batch_pages): # When an error is occurred, we should keep looping and produce the same number of # PrefetchAck as other ranks do, because prefetch_sync_thread (i.e. consumer of # prefetch_sync_queue) perform reduce on the results. This is so tricky. if all_success and operation.is_terminated(): all_success = False if all_success: - batch_hashes = operation.hash_value[i : i + STORAGE_BATCH_SIZE] + batch_hashes = operation.hash_value[i : i + batch_pages] batch_host_indices = operation.host_indices[ i * self.page_size : (i + len(batch_hashes)) * self.page_size ] # Get one batch token, and update the completed_tokens if succeed - extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys) + extra_info = HiCacheStorageExtraInfo( + prefix_keys=prefix_keys, + extra_info={"kv_hints": operation.kv_hints}, + ) try: hit_pages = self._page_transfer_kv_batch( @@ -1162,6 +1169,55 @@ def _page_transfer_kv_batch( Here, "batch" means a single unit of L3 read, not a "batch" in model forward. """ + if ( + kv_derived_transfers + and getattr(self.storage_backend, "supports_combined_page_reads", False) + is True + and all( + transfer.hit_policy == PoolHitPolicy.ALL_PAGES + for transfer in kv_derived_transfers + ) + ): + transfers = [ + PoolTransfer( + name=name, host_indices=batch_host_indices, keys=batch_hashes + ) + for name in [PoolName.KV, *(t.name for t in kv_derived_transfers)] + ] + required_names = {str(transfer.name) for transfer in transfers} + if len(required_names) != len(transfers): + logger.error( + "Duplicate combined storage pools for request %s", + operation.request_id, + ) + return 0 + results = self.storage_backend.batch_get_v2( + transfers, extra_info=extra_info + ) + # Missing or truncated sidecar results must not silently drop a + # required pool from the prefix intersection. A malformed response + # is a miss; keep producing the normal number of cross-rank acks. + if not isinstance(results, dict) or any( + name not in results + or not isinstance(results[name], list) + or len(results[name]) != len(batch_hashes) + or any(type(value) is not bool for value in results[name]) + for name in required_names + ): + logger.error( + "Malformed combined storage result for request %s", + operation.request_id, + ) + return 0 + hits = count_pool_hits({name: results[name] for name in required_names}) + completed = min(hits.values()) + if completed < len(batch_hashes): + logger.warning( + "Prefetch operation %s failed to retrieve page %s.", + operation.request_id, + batch_hashes[completed], + ) + return completed # Read from KV pool. kv_hits = self.page_get_func( operation, batch_hashes, batch_host_indices, extra_info @@ -1258,9 +1314,12 @@ def _storage_hit_query(self, operation) -> tuple[list[str], int]: ) operation.all_hash_values = page_hashes - for start in range(0, len(page_hashes), STORAGE_BATCH_SIZE): - batch_hashes = page_hashes[start : start + STORAGE_BATCH_SIZE] - extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys) + batch_pages = self.storage_backend.query_batch_pages + for start in range(0, len(page_hashes), batch_pages): + batch_hashes = page_hashes[start : start + batch_pages] + extra_info = HiCacheStorageExtraInfo( + prefix_keys=prefix_keys, extra_info={"kv_hints": operation.kv_hints} + ) hit_page_num = self.storage_backend.batch_exists(batch_hashes, extra_info) hash_value.extend(batch_hashes[:hit_page_num]) storage_query_count += hit_page_num * self.page_size diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 1a3f60b8576e..1f3ca79441a6 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -276,6 +276,9 @@ from sglang.srt.managers.scheduler_components.request_receiver import ( SchedulerRequestReceiver, ) +from sglang.srt.managers.scheduler_components.thread_scheduling import ( + configure_python_thread_scheduling, +) from sglang.srt.managers.scheduler_components.weight_updater import ( SchedulerWeightUpdaterManager, ) @@ -3185,6 +3188,7 @@ def _prefetch_kvcache(self, req: Req, storage_hit_end: Optional[int] = None): extra_key=req.extra_key, cache_salt=req.cache_salt, storage_hit_end=storage_hit_end, + kv_hints=req.kv_hints, ) def _process_storage_prefetch_retries(self): @@ -5786,6 +5790,13 @@ def maybe_sleep_on_idle(self): def handle_freeze_gc(self, recv_req: FreezeGCReq): """Handle freeze_gc request: freeze scheduler's GC and forward to detokenizer.""" freeze_gc("Scheduler") + logger.info( + "GC freeze completed: pid=%d tp=%d cp=%d pp=%d", + os.getpid(), + get_parallel().tp_rank, + get_parallel().attn_cp_rank, + get_parallel().pp_rank, + ) self.ipc_channels.send_to_detokenizer.send_output(recv_req, recv_req) return None @@ -5919,6 +5930,7 @@ def configure_scheduler_process( # Configure the logger configure_logger(server_args, prefix=prefix) suppress_other_loggers() + configure_python_thread_scheduling(get_device().python_thread_switch_interval) # Set cpu affinity to this gpu process if envs.SGLANG_SET_CPU_AFFINITY.get(): diff --git a/python/sglang/srt/managers/scheduler_components/thread_scheduling.py b/python/sglang/srt/managers/scheduler_components/thread_scheduling.py new file mode 100644 index 000000000000..16aabfe530c0 --- /dev/null +++ b/python/sglang/srt/managers/scheduler_components/thread_scheduling.py @@ -0,0 +1,39 @@ +"""Explicit, process-local Python thread scheduling for scheduler workers.""" + +import logging +import math +import os +import sys +from typing import Optional + +logger = logging.getLogger(__name__) + + +def validate_python_thread_switch_interval(value: Optional[float]) -> None: + if value is not None and ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or value <= 0 + ): + raise ValueError("--python-thread-switch-interval must be positive and finite") + + +def configure_python_thread_scheduling(interval: Optional[float]) -> None: + """Leave interpreter defaults untouched unless explicitly requested. + + Call in the spawned scheduler process, after runtime-context publication; + setting a parent interpreter does not configure spawned children. This is + process-wide, not a thread priority or a CPU-affinity control. + """ + validate_python_thread_switch_interval(interval) + if interval is None: + return + sys.setswitchinterval(interval) + logger.info( + "Python scheduler thread switch interval pid=%d requested_seconds=%.9g " + "actual_seconds=%.9g", + os.getpid(), + interval, + sys.getswitchinterval(), + ) diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index fd795d55d6b9..9bbcf1237555 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -11,7 +11,6 @@ from typing import TYPE_CHECKING, Any, List, Optional, Set import torch - from sglang.srt.environ import envs if TYPE_CHECKING: @@ -173,6 +172,30 @@ class HiCacheStorage(ABC): It abstracts the underlying storage mechanism, allowing different implementations to be used. """ + # An opt-in backend may read primary KV and KV-derived ALL_PAGES sidecars + # together. It must return a same-length per-entry result for EVERY pool; + # completion must protect all destinations until the native IO is terminal. + # Trailing/checkpoint pools retain their separate resume-boundary path. + supports_combined_page_reads = False + + # Opt-in writes use one terminal completion for primary KV and KV-derived + # ALL_PAGES sidecars. A backup ACK must never claim success for KV alone. + supports_combined_page_writes = False + + # Shared-store MLA backends can persist primary KV/indexer on TP0 only. + # A rank-local store must override this to retain every shard's writer. + requires_rank_local_backup = False + + # Read batching only; writes retain STORAGE_BATCH_SIZE. All + # ranks participating in prefetch ACK reduction must use the same value, + # including after a failed batch, to preserve their collective sequence. + prefetch_batch_pages = STORAGE_BATCH_SIZE + + # Ordinary KV prefix lookups can amortize request-scoped metadata parsing + # independently of payload reads. Backends keep the historical default; + # hybrid pool queries retain their own full-chain boundary contract. + query_batch_pages = STORAGE_BATCH_SIZE + # todo, the page size of storage backend does not have to be the same as the same as host memory pool def register_mem_pool_host(self, mem_pool_host: HostKVCache): self.mem_pool_host = mem_pool_host diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index 6698a40bce03..f9555396e400 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -28,6 +28,7 @@ ) from sglang.srt.mem_cache.base_prefix_cache import CacheRequestHandle from sglang.srt.mem_cache.hicache_storage import ( + STORAGE_BATCH_SIZE, HiCacheStorageExtraInfo, PoolHitPolicy, PoolName, @@ -89,6 +90,7 @@ class PPPrefetchTicket: matched_prefix_tokens: List[int] pool_specs: tuple[PPPrefetchPoolSpec, ...] storage_hit_count: int = 0 + kv_hints: Any = None class PPPrefetchDecision(Enum): @@ -136,8 +138,11 @@ def __init__( hash_value: Optional[List[str]] = None, prefix_keys: Optional[List[str]] = None, pool_transfers: Optional[list[PoolTransfer]] = None, + kv_hints=None, ): - super().__init__(host_indices, token_ids, last_hash, hash_value, prefix_keys) + super().__init__( + host_indices, token_ids, last_hash, hash_value, prefix_keys, kv_hints + ) self.pool_transfers = pool_transfers self.pool_storage_result = PoolTransferResult.empty() @@ -151,6 +156,7 @@ def __init__( prefix_keys: Optional[List[str]] = None, pool_transfers: Optional[list[PoolTransfer]] = None, assume_stored: bool = False, + kv_hints=None, ): self.handle = handle self.request_id = handle.rid @@ -167,6 +173,7 @@ def __init__( last_hash, prefix_keys=prefix_keys, pool_transfers=pool_transfers, + kv_hints=kv_hints, ) self.pool_transfers_done = not bool(pool_transfers) # The hit query's verdict per pool, as the chain prefix length (in @@ -1006,6 +1013,7 @@ def prefetch( prefix_keys: Optional[List[str]] = None, extra_pools: Optional[list[PoolTransfer]] = None, assume_stored: bool = False, + kv_hints=None, ) -> PrefetchOperation: operation = PrefetchOperation( handle, @@ -1014,6 +1022,7 @@ def prefetch( prefix_keys=prefix_keys, pool_transfers=extra_pools, assume_stored=assume_stored, + kv_hints=kv_hints, ) self.prefetch_queue.put(operation) return operation @@ -1041,6 +1050,7 @@ def submit_prefetch( matched_prefix_tokens: Optional[List[int]], pool_transfers: Optional[list[PoolTransfer]], assume_stored: bool = False, + kv_hints=None, ) -> PrefetchSubmission: if self.pp_prefetch_command_group is None: return PrefetchSubmission( @@ -1051,6 +1061,7 @@ def submit_prefetch( prefix_keys, extra_pools=pool_transfers, assume_stored=assume_stored, + kv_hints=kv_hints, ) ) @@ -1059,6 +1070,7 @@ def submit_prefetch( rid = handle.rid ticket = PPPrefetchTicket( + kv_hints=kv_hints, handle=handle, prefetch_key=RadixKey( array("q", prefetch_key.token_ids), @@ -1080,6 +1092,7 @@ def submit_prefetch( last_hash, prefix_keys=ticket.prefix_keys, pool_transfers=pool_transfers, + kv_hints=kv_hints, ) storage_hit_count = len(ticket.prefetch_key.token_ids) @@ -1279,6 +1292,7 @@ def pp_prefetch_command_thread_func(self) -> None: ticket.handle, ticket.prefetch_key, ticket.last_hash, + kv_hints=ticket.kv_hints, prefix_keys=ticket.prefix_keys, # Keep sidecar ACKs aligned even if allocation fails. pool_transfers=[ @@ -1365,7 +1379,7 @@ def _storage_hit_query( extra_info = HiCacheStorageExtraInfo( prefix_keys=operation.prefix_keys.copy() if operation.prefix_keys else None, - extra_info={"pp_rank": pp_rank} if pp_rank is not None else None, + extra_info={"kv_hints": operation.kv_hints, "pp_rank": pp_rank}, ) if operation.pool_transfers: hit_result = self.storage_backend.batch_exists_v2( @@ -1378,6 +1392,9 @@ def _storage_hit_query( ) kv_hit_pages = hit_result.kv_hit_pages + operation.pool_storage_result.restorable_prefix_pages = ( + hit_result.restorable_prefix_pages + ) operation.pool_storage_result.update_kv_hit_pages(kv_hit_pages) # Each pool's own boundary (the KV pool's where the backend reports # it, else the folded cut); the beliefs heal pool by pool. A pool the @@ -1396,13 +1413,20 @@ def _storage_hit_query( ) def _sync_prefetch_hit_query(self, operation, storage_hit_count: int) -> int: - """One MIN collective for the folded hit count and every pool's own - boundary, so per-pool beliefs stay rank-identical. A pool omitted - by any rank heals from the folded cut without learning presence.""" + """Reduce pool verdicts, then intersect explicit legal resume boundaries. + + The packed presence flags make the bitmap collective rank-identical. + Mixed explicit/legacy replies fail closed; a scalar maximum cannot + establish that another rank has a checkpoint at that boundary. + """ pools = list(PoolName) local = operation.query_pool_hit_pages + boundaries = operation.pool_storage_result.restorable_prefix_pages + explicit = boundaries is not None packed = torch.tensor( - [storage_hit_count] + [local.get(pool, _NO_POOL_VERDICT) for pool in pools], + [storage_hit_count] + + [local.get(pool, _NO_POOL_VERDICT) for pool in pools] + + [int(explicit), int(not explicit)], dtype=torch.int, ) self._all_reduce( @@ -1413,10 +1437,28 @@ def _sync_prefetch_hit_query(self, operation, storage_hit_count: int) -> int: reduced = packed.tolist() operation.query_pool_hit_pages = { pool: pages - for pool, pages in zip(pools, reduced[1:], strict=True) + for pool, pages in zip(pools, reduced[1 : 1 + len(pools)], strict=True) if pages != _NO_POOL_VERDICT } - return int(reduced[0]) + if reduced[-1]: + return int(reduced[0]) + if not reduced[-2]: + operation.pool_storage_result.restorable_prefix_pages = [] + return 0 + valid = set(boundaries) + pages = len(operation.all_hash_values) + if not pages: + operation.pool_storage_result.restorable_prefix_pages = [] + return 0 + mask = torch.tensor([int(page in valid) for page in range(1, pages + 1)]) + self._all_reduce( + mask, + torch.distributed.ReduceOp.MIN, + self.prefetch_hits_sync_groups, + ) + common = [page for page, present in enumerate(mask.tolist(), 1) if present] + operation.pool_storage_result.restorable_prefix_pages = common + return max(common, default=0) * self.page_size def _move_pool_indices( self, host_pool, host_indices, device_indices, *, write_back_jit: bool @@ -1506,6 +1548,7 @@ def _page_transfer_sidecar( sidecar_hashes, transfers_nonkv, ): + extra_info.extra_info = {"kv_hints": operation.kv_hints} results.update( self.storage_backend.batch_get_v2( transfers, extra_info=extra_info @@ -1563,6 +1606,21 @@ def _page_backup_with_stable_layout(self, operation): if backup_transfers: self._resolve_sidecar_kv_derived_pool_transfers(operation) self._resolve_sidecar_nonkv_derived_pool_transfers(operation) + if ( + not self.backup_skip + and self.mem_pool_host.kv_buffer is not None + and getattr( + self.storage_backend, "supports_combined_page_writes", False + ) + is True + and all( + transfer.indices_from_pool == PoolName.KV + and transfer.hit_policy == PoolHitPolicy.ALL_PAGES + for transfer in backup_transfers + ) + ): + self._page_backup_combined(operation, backup_transfers) + return results = {} for extra_info, transfers in _trailing_chain_groups( operation.prefix_keys, operation.hash_value, backup_transfers @@ -1596,6 +1654,60 @@ def _page_backup_with_stable_layout(self, operation): len(operation.hash_value) * self.page_size if sidecar_ok else 0 ) + def _page_backup_combined(self, operation, sidecars): + """Persist aligned KV/sidecar batches before acknowledging any pages. + + The caller holds the host layout lease and the queued operation keeps + every source slot protected. Trailing or independently indexed pools + deliberately retain their existing backup path. + """ + names = [PoolName.KV, *(transfer.name for transfer in sidecars)] + if len(set(names)) != len(names): + logger.error("Duplicate combined storage backup pools") + return + accumulated = {str(name): [] for name in names} + prefix_keys = operation.prefix_keys + for start in range(0, len(operation.hash_value), STORAGE_BATCH_SIZE): + keys = operation.hash_value[start : start + STORAGE_BATCH_SIZE] + indices = operation.host_indices[ + start * self.page_size : (start + len(keys)) * self.page_size + ] + transfers = [ + PoolTransfer(name=PoolName.KV, keys=keys, host_indices=indices), + *( + replace(transfer, keys=keys, host_indices=indices) + for transfer in sidecars + ), + ] + results = self.storage_backend.batch_set_v2( + transfers, + extra_info=HiCacheStorageExtraInfo(prefix_keys=prefix_keys), + ) + if not isinstance(results, dict) or any( + name not in results + or not isinstance(results[name], list) + or len(results[name]) != len(keys) + or any(type(value) is not bool for value in results[name]) + for name in accumulated + ): + logger.error("Malformed combined storage backup result") + break + for name in accumulated: + accumulated[name].extend(results[name]) + if not all(all(results[name]) for name in accumulated): + logger.warning("Combined storage backup failed for %d pages", len(keys)) + break + operation.completed_tokens += len(keys) * self.page_size + if prefix_keys is not None: + prefix_keys = prefix_keys + keys + operation.pool_storage_result.update_extra_pool_hit_pages( + { + name: hits + for name, hits in count_pool_hits(accumulated).items() + if name != str(PoolName.KV) + } + ) + def should_backup(self, transfer: PoolTransfer) -> bool: if not self.backup_skip: return True diff --git a/python/sglang/srt/mem_cache/pool_host/common.py b/python/sglang/srt/mem_cache/pool_host/common.py index 5728dbde58ad..804da6453b80 100644 --- a/python/sglang/srt/mem_cache/pool_host/common.py +++ b/python/sglang/srt/mem_cache/pool_host/common.py @@ -7,7 +7,6 @@ from functools import lru_cache import torch - from sglang.srt.environ import envs from sglang.srt.mem_cache.storage.mmap import alloc_mmap from sglang.srt.runtime_context import get_memory @@ -87,6 +86,10 @@ def get_allocator_from_storage(allocator_type): "Fallback to use default allocator." ) return HostTensorAllocator() + elif allocator_type == "kvcr": + from sglang.srt.mem_cache.storage.kvcr.numa_pool import allocator_for_config + + return allocator_for_config(get_memory().hicache_storage_backend_extra_config) elif allocator_type == "mori": try: from sglang.srt.mem_cache.storage.umbp.umbp_host_allocator import ( @@ -187,6 +190,12 @@ def _cuda_host_register( # cudaHostUnregister to receive each base pointer, not just the tensor's # original base once after several independent registrations. setattr(buffer, _CUDA_HOST_REGISTERED_RANGES_ATTR, registered_ranges) + if hasattr(buffer, "_sglang_owned_numa_host_proof"): + from sglang.srt.mem_cache.storage.kvcr.numa_pool import ( + audit_registered_host, + ) + + audit_registered_host(buffer, registered_ranges) except Exception: remaining_ranges = _cuda_host_unregister_ranges( cudart, registered_ranges, operation="registration rollback" diff --git a/python/sglang/srt/mem_cache/storage/backend_factory.py b/python/sglang/srt/mem_cache/storage/backend_factory.py index 7537f2877070..62375c0da622 100644 --- a/python/sglang/srt/mem_cache/storage/backend_factory.py +++ b/python/sglang/srt/mem_cache/storage/backend_factory.py @@ -162,7 +162,7 @@ def _create_builtin_backend( return backend_class(storage_config) elif backend_name == "nixl": return backend_class(storage_config) - elif backend_name == "mooncake": + elif backend_name in ("mooncake", "kvcr"): backend = backend_class(storage_config, mem_pool_host) return backend elif backend_name == "npu_memcache": @@ -221,6 +221,10 @@ def _create_builtin_backend( "MooncakeStore", ) +StorageBackendFactory.register_backend( + "kvcr", "sglang.srt.mem_cache.storage.kvcr.kvcr_store", "KVCRStore" +) + StorageBackendFactory.register_backend( "npu_memcache", "sglang.srt.mem_cache.storage.npu_memcache.npu_memcache_store", diff --git a/python/sglang/srt/mem_cache/storage/kvcr/README.md b/python/sglang/srt/mem_cache/storage/kvcr/README.md new file mode 100644 index 000000000000..b365100455df --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/kvcr/README.md @@ -0,0 +1,107 @@ +# KVCR HiCacheStorage backend + +This optional backend uses KVCR's indexed-region public API and NIXL/UCX +transport. Install a compatible KVCR build separately; selecting another +HiCache backend does not import KVCR. + +## Ownership and data path + +SGLang owns GPU memory, HiCache host pages, scheduling, and host-to-GPU copies. +KVCR owns its deposited DRAM cache: + +- Backup: GPU -> HiCache host pages -> KVCR-owned DRAM. +- Restore: local or peer KVCR DRAM -> locked HiCache host pages -> GPU. + +This is not a DirectLinker or direct-HBM restore. Peer delivery does not require +an intermediate allocation in the target's KVCR DRAM cache. Unpinned SGLang +host pages are never offered as remote sources. + +## Configuration + +Enable hierarchical caching, choose `--hicache-storage-backend kvcr`, and use +`--hicache-mem-layout page_first` or `page_first_direct`. Supply inline JSON +through `--hicache-storage-backend-extra-config`, for example: + +```json +{ + "control_advertise_host": "", + "control_port": 32100, + "local_dram_bytes": 17179869184, + "combine_page_reads": true, + "combine_page_writes": true, + "prefetch_batch_pages": 512, + "align_storage_prefix": true +} +``` + +This is an illustrative configuration, not a model-specific capacity +recommendation. Account separately for GPU cache, HiCache host staging, and +KVCR DRAM. The default rank-local mode allocates KVCR storage on each rank. +Control ports must accommodate the rank offsets and be reachable by peers. +The control endpoint's TCP scheme does not select the NIXL payload transport. + +Combined reads/writes are opt-in and apply only to primary KV plus +KV-derived `ALL_PAGES` sidecars. Independent and trailing/checkpoint pools +retain their own transfer and resume-boundary handling. Read batching defaults +to 128 pages; configure the same value on ranks participating in the same +prefetch collective. Query and backup batches remain 128 pages. + +Prefix-recency alignment additionally requires +`--hicache-storage-pass-prefix-keys`. It adjusts eviction preference; it does +not reserve source blocks or guarantee residency. + +An optional positive integer `startup_timeout_seconds` requires the companion +KVCR startup budget change. Omitting it preserves KVCR's normal startup budget. +Transfer and abandonment deadlines are separate from this initialization setting. + +## Remote hints and replicated MLA + +Cross-replica discovery requires request-scoped source hints, such as those +provided by the companion in-process Dynamo SGLang integration. A query is +advisory; terminal delivery determines which pages are usable. Missing hints +do not authorize fetching from an unrelated replica. + +`share_replicated_mla: true` is an explicit assertion that KV and indexer bytes +are identical across attention TP ranks. It stores one copy on TP0 while every +rank restores into its own host buffers. It requires physical MLA KV, CP=1, +PP=1, and no heterogeneous TP head splitting, and rejects pools other than KV +and indexer. Do not enable it for sharded hybrid state or SWA pools. + +The adapter preserves valid resume-boundary sets for trailing/checkpoint +pools. The hybrid controller intersects these sets across ranks instead of +assuming that the minimum scalar hit count is a valid checkpoint. + +## Failure handling and diagnostics + +Successful or missed operations must reach terminal native completion before +HiCache can reuse their buffers. If native transfer state is uncertain, the +worker exits because HiCache cannot quarantine potentially in-flight pages. +Startup/progress/shutdown failures likewise must not release registered +memory while native access may remain possible. + +`diagnostic_telemetry` and `verify_transfer_bytes` are disabled by default. +They are validation tools, not free performance instrumentation. Diagnostic +logging can include storage-key and request identities; retain those logs +according to the deployment's data-handling policy. + +## Validation + +CPU adapter/controller contracts: + +```bash +python test/registered/unit/mem_cache/test_kvcr_storage.py -v +python test/registered/unit/mem_cache/test_buffer_mode_sidecar.py -v +python test/registered/unit/mem_cache/test_pp_prefetch_ticket.py -v +``` + +Real NIXL host-byte round trips, including multi-pool and shared-MLA cases: + +```bash +python test/manual/test_kvcr_storage_native.py -v +``` + +The native gate requires Linux, compatible KVCR/NIXL/UCX, and adequate shared +memory; provision `/dev/shm` explicitly in containers. It verifies host bytes, +not model accuracy or inference performance. End-to-end validation was run +with GLM5.3, two independently complete TP4 replicas, and in-process Dynamo. +That does not establish arbitrary hybrid-model or multi-node-engine support. diff --git a/python/sglang/srt/mem_cache/storage/kvcr/__init__.py b/python/sglang/srt/mem_cache/storage/kvcr/__init__.py new file mode 100644 index 000000000000..c9bb6077b376 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/kvcr/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""KVCR-backed HiCache L3 storage (optional dependency).""" diff --git a/python/sglang/srt/mem_cache/storage/kvcr/delivery.py b/python/sglang/srt/mem_cache/storage/kvcr/delivery.py new file mode 100644 index 000000000000..8e2802d148ce --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/kvcr/delivery.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Bound public delivery calls without changing the adapter's block mapping.""" + +from __future__ import annotations + +import logging +import os +import time +from collections import deque +from dataclasses import dataclass +from itertools import islice, pairwise + +logger = logging.getLogger(__name__) + + +def validate_destination_spans(blocks, layouts): + """Reject aliases across the entire call before any destination is writable.""" + components = { + part.name: part for layout in layouts.values() for part in layout.components + } + spans = [] + for refs in blocks.values(): + for ref in refs: + part = components[ref.label] + start = part.address + ref.element_index * part.stride + spans.append((start, start + part.size)) + spans.sort() + for previous, current in pairwise(spans): + if current[0] < previous[1]: + raise ValueError("Overlapping KVCR delivery destination spans") + + +def block_chunks(blocks): + """Keep every component of each pool-qualified block in the same chunk.""" + entries = iter(blocks.items()) + while chunk := dict(islice(entries, 512)): + yield chunk + + +@dataclass +class Delivery: + blocks: dict + parent_operation: int | None = None + operation: int | None = None + prepared_epoch: float | None = None + finished_epoch: float | None = None + prepared: float | None = None + submitted: float | None = None + finished: float | None = None + result: object = None + + +def run_window(store, deliveries, request_id, seen_handles): + """Return raw completions only after ALL writes in this hint are terminal. + + The caller MUST hold the active hint scope. Submission itself may enqueue + before raising, and even bookkeeping can lose a returned handle. There is + no safe Python unwind in either case. Do not attempt an unbounded drain of + known handles: it cannot establish quiescence of an unknown operation. + """ + try: + pending, completed = deque(), [] + remaining = iter(deliveries) + exhausted = False + while pending or not exhausted: + while len(pending) < 4 and not exhausted: + delivery = next(remaining, None) + if delivery is None: + exhausted = True + break + delivery.prepared_epoch = ( + time.time() if store._diagnostic_telemetry else None + ) + delivery.prepared = ( + time.perf_counter() if store._diagnostic_telemetry else None + ) + handle = store._runner.deliver(delivery.blocks, request_id=request_id) + if type(handle) is not int or handle <= 0 or handle in seen_handles: + raise RuntimeError("Invalid or duplicate KVCR delivery handle") + seen_handles.add(handle) + delivery.operation = handle + delivery.submitted = ( + time.perf_counter() if store._diagnostic_telemetry else None + ) + pending.append(delivery) + if pending: + delivery = pending[0] + # The existing shared completion mailbox retains reverse-order + # completions and completions belonging to concurrent callers. + delivery.result = store._wait(delivery.operation) + delivery.finished_epoch = ( + time.time() if store._diagnostic_telemetry else None + ) + delivery.finished = ( + time.perf_counter() if store._diagnostic_telemetry else None + ) + completed.append(delivery) + pending.popleft() + return completed + except BaseException: # noqa: BLE001 - even cancellation must not release live DMA slots + # This must run INSIDE the hint. A logging failure must also terminate + # without discarding hints or returning HiCache's still-exposed slots. + try: + logger.critical( + "Unsafe KVCR delivery pipeline; terminating worker", exc_info=True + ) + finally: + os._exit(1) + raise RuntimeError("KVCR fatal exit unexpectedly returned") + + +def requested_results(deliveries): + """Interpret only this handle's requested identities, after phase completion.""" + return { + key: delivery.result[key] + for delivery in deliveries + for key in delivery.blocks + if key in delivery.result + } diff --git a/python/sglang/srt/mem_cache/storage/kvcr/diagnostics.py b/python/sglang/srt/mem_cache/storage/kvcr/diagnostics.py new file mode 100644 index 000000000000..7c5234246d08 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/kvcr/diagnostics.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Bounded interval collector for KVCR's optional native telemetry seam. + +No per-operation samples or page keys are retained. KVCR calls this collector +from its public main-side API, serialized by KVCRStore's IO lock. These are +diagnostic measurements, not HiCache's logical-page or bandwidth metrics. +""" + +from __future__ import annotations + +import bisect +import json + + +class DiagnosticStats: + # Native durations are seconds. Keep a fixed histogram, including overflow, + # rather than accumulating an unbounded list while the server is idle. + BOUNDS = (0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1, 5, 30) + + def __init__(self, max_series=128): + self._max_series = max_series + self._series = {} + self._dropped = 0 + + def _entry(self, kind, name, labels): + key = (kind, name, tuple(labels)) + if key not in self._series: + if len(self._series) >= self._max_series: + self._dropped += 1 + return None + self._series[key] = {} + return self._series[key] + + def increase_counter(self, name, value, labelvalues=()): + entry = self._entry("counter", name, labelvalues) + if entry is not None: + entry["value"] = entry.get("value", 0) + value + + def set_gauge(self, name, value, labelvalues=()): + entry = self._entry("gauge", name, labelvalues) + if entry is not None: + entry["value"] = value + + def observe_histogram(self, name, value, labelvalues=()): + entry = self._entry("histogram", name, labelvalues) + if entry is None: + return + entry["count"] = entry.get("count", 0) + 1 + entry["sum"] = entry.get("sum", 0) + value + entry["min"] = min(entry.get("min", value), value) + entry["max"] = max(entry.get("max", value), value) + bucket = bisect.bisect_left(self.BOUNDS, value) + field = f"bucket_{bucket}" + entry[field] = entry.get(field, 0) + 1 + + def reduce(self): + """Drain one interval; JSON tuple keys preserve metric/label identity. + + Histogram buckets are non-cumulative and indexed into BOUNDS; index + len(BOUNDS) is overflow. Consumers can report bounds, not exact p90s. + """ + result = { + json.dumps((kind, name, labels, field), separators=(",", ":")): value + for (kind, name, labels), entry in self._series.items() + for field, value in entry.items() + } + if self._dropped: + result["dropped_observations"] = self._dropped + self._series.clear() + self._dropped = 0 + return result + + def is_empty(self): + return not self._series and not self._dropped diff --git a/python/sglang/srt/mem_cache/storage/kvcr/hints.py b/python/sglang/srt/mem_cache/storage/kvcr/hints.py new file mode 100644 index 000000000000..87c52817553b --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/kvcr/hints.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Request-scoped hints; no hint is ever shared between unrelated requests.""" + +from __future__ import annotations + +from urllib.parse import urlsplit + +import msgspec + + +def rank_offset(config) -> int: + return config.tp_rank + config.tp_size * ( + config.attn_cp_rank + config.attn_cp_size * config.pp_rank + ) + + +def rank_endpoint(endpoint: str, offset: int) -> str | None: + """Select a neighboring endpoint, rejecting malformed control addresses.""" + try: + source = urlsplit(endpoint) + port = source.port + if ( + source.scheme != "tcp" + or not source.hostname + or port is None + or not 0 < port + offset <= 65535 + or source.path + or source.query + or source.fragment + or source.username is not None + ): + return None + except ValueError: + return None + host = source.hostname + if ":" in host: + host = f"[{host}]" + return f"tcp://{host}:{port + offset}" + + +def peer_hint(extra_info, offset: int) -> dict | None: + """Validate the supported action and select this rank's source endpoint. + + Dynamo advertises the TP/CP/PP-zero endpoint of the chosen DP replica. + Neighbor ports represent the other shards in that same replica. + """ + if extra_info is None or not extra_info.extra_info: + return None + envelope = extra_info.extra_info.get("kv_hints") + if isinstance(envelope, msgspec.Struct): + envelope = msgspec.to_builtins(envelope) + if not isinstance(envelope, dict) or envelope.get("protocol_version") != "0.1": + return None + actions = envelope.get("actions") + if not isinstance(actions, list): + return None + for action in actions: + if not isinstance(action, dict) or ( + action.get("action_type") != "kv.fetch" + or action.get("action_version") != "1.0" + ): + continue + payload = action.get("payload") + if not isinstance(payload, dict) or payload.get("mode", "copy") != "copy": + continue + if "no_retain" in payload: + continue + endpoint = payload.get("source_control_endpoint") + hashes = payload.get("block_hashes") + if not isinstance(endpoint, str) or not isinstance(hashes, list) or not hashes: + continue + if any(type(h) is not int or not 0 <= h < 1 << 64 for h in hashes): + continue + selected_endpoint = rank_endpoint(endpoint, offset) + if selected_endpoint is None: + continue + return { + "protocol_version": "0.1", + "message_id": envelope.get("message_id", ""), + "actions": [ + { + "action_id": action.get("action_id", ""), + "action_type": "kv.fetch", + "action_version": "1.0", + "payload": { + "source_control_endpoint": selected_endpoint, + "block_hashes": hashes, + "mode": "copy", + }, + } + ], + } + return None diff --git a/python/sglang/srt/mem_cache/storage/kvcr/kvcr_store.py b/python/sglang/srt/mem_cache/storage/kvcr/kvcr_store.py new file mode 100644 index 000000000000..f31dfd998975 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/kvcr/kvcr_store.py @@ -0,0 +1,1096 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Host-staged KVCR L3: SGLang HBM -> HiCache host -> KVCR DRAM. + +Reads deliver directly into the host slots owned/locked by HiCache. KVCR +serves only its own deposited DRAM, never unpinned framework host pages. +""" + +from __future__ import annotations + +import ctypes +import hashlib +import json +import logging +import mmap +import os +import threading +import time +import uuid +from contextlib import ExitStack, contextmanager +from queue import Empty, SimpleQueue + +from sglang.srt.mem_cache.hicache_storage import ( + STORAGE_BATCH_SIZE, + HiCacheStorage, + PoolHitPolicy, + PoolName, + PoolTransfer, + PoolTransferResult, +) +from sglang.srt.mem_cache.storage.kvcr.delivery import ( + Delivery, + block_chunks, + requested_results, + run_window, + validate_destination_spans, +) +from sglang.srt.mem_cache.storage.kvcr.hints import ( + peer_hint, + rank_endpoint, + rank_offset, +) +from sglang.srt.mem_cache.storage.kvcr.layout import ( + StorageKeyAdapter, + block_key, + describe_pool, + page_slots, + resume_boundaries, + storage_namespace, +) + +logger = logging.getLogger(__name__) + + +def _allocate_dram_pool(length): + """Own process-private pages and permit THP before NIXL registers them. + + mmap(-1, length) defaults to shared /dev/zero memory on Linux. Such pages + cannot use THP when the host's shmem policy is "never", even if anonymous + THP is available. No other process maps this pool; private anonymous memory + preserves ownership while avoiding that accidental restriction. + """ + if not hasattr(mmap, "MAP_PRIVATE") or not hasattr(mmap, "MAP_ANONYMOUS"): + return mmap.mmap(-1, length) + allocation = mmap.mmap(-1, length, flags=mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS) + if hasattr(allocation, "madvise") and hasattr(mmap, "MADV_HUGEPAGE"): + try: + allocation.madvise(mmap.MADV_HUGEPAGE) + except OSError: + # Advice is optional: ordinary pages must remain fully functional + # on kernels that do not implement or permit transparent huge pages. + logger.debug("KVCR DRAM huge-page advice unavailable", exc_info=True) + return allocation + + +# Pools HiCache restores as TRAILING_PAGES windows (SWA, its DeepSeek V4 +# compressor states, Mamba): a page need not have state in them. +_TRAILING_POOLS = frozenset( + str(name) + for name in ( + PoolName.SWA, + PoolName.DRAFT_SWA, + PoolName.MAMBA, + PoolName.DEEPSEEK_V4_C4_STATE, + PoolName.DEEPSEEK_V4_C4_INDEXER_STATE, + PoolName.DEEPSEEK_V4_C128_STATE, + ) +) + + +class _NoFrameworkPinning: + """Answer misses when a peer asks for a non-KVCR-owned source.""" + + def __init__(self): + self._results = SimpleQueue() + self._next_id = 0 + + def request_pin(self, keys): + self._next_id += 1 + self._results.put((self._next_id, None)) + return self._next_id + + def poll_pin_results(self): + results = [] + while True: + try: + results.append(self._results.get_nowait()) + except Empty: + return results + + def release_pin(self, handle): + return False + + +class KVCRStore(HiCacheStorage): + """Optional NIXL/UCX L3 backend using KVCR's indexed-region public API.""" + + def __init__(self, storage_config, mem_pool_host): + self.config = storage_config + self.mem_pool_host = mem_pool_host + self.registered_pools = {} + self._layouts = {} + self._runner = None + self._maps = [] + self._lock = threading.RLock() + self._completion = threading.Condition(self._lock) + self._completed = {} + self._poll_thread = None + self._stop_poll = threading.Event() + self._closed = False + self._pinning = _NoFrameworkPinning() + self._extra = dict(storage_config.extra_config or {}) + from sglang.srt.mem_cache.storage.kvcr.numa_pool import node_option + + self._local_dram_numa_node = node_option( + self._extra.get("local_dram_numa_node"), "local_dram_numa_node" + ) + self._delivery_pipeline_window = self._extra.get("delivery_pipeline_window", 1) + if type( + self._delivery_pipeline_window + ) is not int or self._delivery_pipeline_window not in (1, 4): + raise ValueError("delivery_pipeline_window must be integer 1 or 4") + if "startup_timeout_seconds" in self._extra and ( + type(self._extra["startup_timeout_seconds"]) is not int + or self._extra["startup_timeout_seconds"] <= 0 + ): + raise ValueError("startup_timeout_seconds must be a positive integer") + if ( + "coalesce_remote_dram_descriptors" in self._extra + and type(self._extra["coalesce_remote_dram_descriptors"]) is not bool + ): + raise ValueError("coalesce_remote_dram_descriptors must be a boolean") + if "nixl_ucx_split_batch_size" in self._extra: + split = self._extra["nixl_ucx_split_batch_size"] + if type(split) is not int or not 1 <= split <= 2**32 - 1: + raise ValueError( + "nixl_ucx_split_batch_size must be a positive uint32 integer" + ) + self.prefetch_batch_pages = self._extra.get( + "prefetch_batch_pages", STORAGE_BATCH_SIZE + ) + if type(self.prefetch_batch_pages) is not int or self.prefetch_batch_pages <= 0: + raise ValueError("prefetch_batch_pages must be a positive integer") + self.query_batch_pages = self._extra.get( + "query_batch_pages", STORAGE_BATCH_SIZE + ) + if ( + type(self.query_batch_pages) is not int + or not 1 <= self.query_batch_pages <= 2048 + ): + raise ValueError("query_batch_pages must be an integer from 1 to 2048") + if self.query_batch_pages != STORAGE_BATCH_SIZE and storage_config.is_mla_model: + raise ValueError( + "query_batch_pages opt-in requires a non-MLA KV-only layout" + ) + self._metrics = None + if storage_config.enable_storage_metrics: + from sglang.srt.observability.metrics_collector import StorageMetrics + + self._metrics = StorageMetrics() + self._verify_transfer_bytes = self._extra.get("verify_transfer_bytes", False) + if type(self._verify_transfer_bytes) is not bool: + raise ValueError("verify_transfer_bytes must be a boolean") + self._align_storage_prefix = self._extra.get("align_storage_prefix", False) + if type(self._align_storage_prefix) is not bool: + raise ValueError("align_storage_prefix must be a boolean") + self._diagnostic_telemetry = self._extra.get("diagnostic_telemetry", False) + if type(self._diagnostic_telemetry) is not bool: + raise ValueError("diagnostic_telemetry must be a boolean") + self.supports_combined_page_reads = self._extra.get("combine_page_reads", False) + if type(self.supports_combined_page_reads) is not bool: + raise ValueError("combine_page_reads must be a boolean") + self.supports_combined_page_writes = self._extra.get( + "combine_page_writes", False + ) + if type(self.supports_combined_page_writes) is not bool: + raise ValueError("combine_page_writes must be a boolean") + self._diagnostic_epoch = time.monotonic() + self._offset = rank_offset(storage_config) + self._share_replicated_mla = self._extra.get("share_replicated_mla", False) + if type(self._share_replicated_mla) is not bool: + raise ValueError("share_replicated_mla must be a boolean") + if self._share_replicated_mla and ( + not storage_config.is_mla_model + or storage_config.pp_size != 1 + or storage_config.attn_cp_size != 1 + or storage_config.should_split_heads + or mem_pool_host.kv_buffer is None + ): + raise ValueError( + "Shared MLA storage requires replicated physical KV, CP=1 and PP=1" + ) + # Shared mode is an explicit claim that KV/indexer bytes are identical + # across attention TP ranks. Never apply it to sharded hybrid state. + self.requires_rank_local_backup = not self._share_replicated_mla + self._source_offset = 0 if self._share_replicated_mla else self._offset + self._agent_name = f"sglang-kvcr-{uuid.uuid4().hex}" + if storage_config.should_split_heads: + raise ValueError("KVCR does not support heterogeneous TP head splitting") + if storage_config.pp_size != 1: + raise ValueError("KVCR HiCacheStorage currently requires PP=1") + self.register_mem_pool_host(mem_pool_host) + if self.query_batch_pages != STORAGE_BATCH_SIZE and set(self._layouts) != { + "kv" + }: + raise ValueError( + "query_batch_pages opt-in requires a non-MLA KV-only layout" + ) + + def register_mem_pool_host(self, mem_pool_host): + self.mem_pool_host = mem_pool_host + # Hybrid logical anchors have no payload. The real pools arrive via v2. + if mem_pool_host.kv_buffer is not None: + self.register_mem_host_pool_v2(mem_pool_host, PoolName.KV) + + def register_mem_host_pool_v2(self, host_pool, host_pool_name): + name = str(host_pool_name) + with self._lock: + if self.query_batch_pages != STORAGE_BATCH_SIZE and name != str( + PoolName.KV + ): + raise ValueError( + "query_batch_pages opt-in requires a non-MLA KV-only layout" + ) + if self._share_replicated_mla and name not in ("kv", "indexer"): + raise ValueError( + "Shared MLA storage only supports replicated KV/indexer pools" + ) + if name in self.registered_pools: + if self.registered_pools[name] is not host_pool: + raise ValueError(f"Host pool {name} was registered twice") + return + if name == str(PoolName.KV) and host_pool.kv_buffer is None: + return + if self._runner is not None: + raise ValueError( + "KVCR host pools must be registered before the first IO" + ) + self._layouts[name] = describe_pool(name, host_pool) + self.registered_pools[name] = host_pool + + def _start(self): + if self._closed: + raise RuntimeError("KVCR storage backend is closed") + if self._runner is not None: + return + from kvcr import KVCR, KVCRBindings + from kvcr.config import KVCRBackendConfigs, KVCRConfig, LocalDramOptions + from kvcr.control_channels import ZmqPeerControlChannel + from kvcr.types import KVCRStartupError, RegionDescriptor + from sglang.srt.mem_cache.storage.kvcr.diagnostics import DiagnosticStats + + if not self._layouts: + raise ValueError("KVCR has no registered physical host pools") + self._namespace = storage_namespace( + self.config, self._layouts, shared_mla=self._share_replicated_mla + ) + total = sum( + component.size * component.count + for layout in self._layouts.values() + for component in layout.components + ) + budget = self._extra.get("local_dram_bytes", 8 << 30) + pages_by_pool = self._extra.get("local_dram_pool_pages", {}) + if ( + type(budget) is not int + or budget <= 0 + or not isinstance(pages_by_pool, dict) + ): + raise ValueError("Invalid KVCR local DRAM budget") + base_port = self._extra.get("control_port", 32100) + if type(base_port) is not int or not 0 < base_port <= 65535: + raise ValueError("KVCR control_port must be a valid integer port") + nixl_base_port = self._extra.get("nixl_listen_port", base_port + 1024) + if type(nixl_base_port) is not int or not 0 < nixl_base_port <= 65535: + raise ValueError("KVCR nixl_listen_port must be a valid integer port") + dp_stride = self.config.tp_size * self.config.attn_cp_size * self.config.pp_size + offset = self.config.dp_rank * dp_stride + self._offset + port, nixl_port = base_port + offset, nixl_base_port + offset + if not 0 < port <= 65535 or not 0 < nixl_port <= 65535 or port == nixl_port: + raise ValueError("KVCR control/NIXL port range is invalid or overlaps") + layouts, regions, plans = [], [], [] + for name, layout in self._layouts.items(): + pages = pages_by_pool.get( + name, budget * layout.components[0].count // total + ) + if type(pages) is not int or pages < 1: + raise ValueError(f"KVCR DRAM budget cannot hold one page of {name}") + for component in layout.components: + length = pages * component.size + if not self._share_replicated_mla or self.config.tp_rank == 0: + plans.append((component.name, length)) + layouts.append((component.name, component.size)) + regions.append( + RegionDescriptor( + addr=component.address, + size=component.size, + stride=component.stride, + count=component.count, + label=component.name, + ) + ) + host = self._extra.get("control_advertise_host", "127.0.0.1") + channel = ZmqPeerControlChannel( + bind_host=self._extra.get("control_host", "0.0.0.0"), + bind_port=port, + advertise_host=host, + ) + self._control_endpoint = channel.endpoint + self._owner_endpoint = rank_endpoint(channel.endpoint, -self._offset) + if self._owner_endpoint is None: + channel.close() + raise ValueError("Cannot identify the replica's TP0 control endpoint") + allocations, mappings = [], [] + with ExitStack() as resources: + resources.callback(channel.close) + for name, length in plans: + allocation = _allocate_dram_pool(length) + resources.callback(allocation.close) + if self._local_dram_numa_node is not None: + from sglang.srt.mem_cache.storage.kvcr.numa_pool import ( + bind_new_private_pool, + ) + + bind_new_private_pool(allocation, self._local_dram_numa_node) + mappings.append(allocation) + address = ctypes.addressof(ctypes.c_char.from_buffer(allocation)) + allocations.append((name, address, length)) + try: + # Do not pass the new option when unset: ordinary deployments + # remain compatible with KVCR versions predating this field. + startup_options = ( + {"startup_timeout_seconds": self._extra["startup_timeout_seconds"]} + if "startup_timeout_seconds" in self._extra + else {} + ) + # Leave the default off and omit the new keyword for older + # KVCR versions unless this experimental path is opted into. + if self._extra.get("coalesce_remote_dram_descriptors", False): + startup_options["coalesce_remote_dram_descriptors"] = True + if "nixl_ucx_split_batch_size" in self._extra: + startup_options["nixl_ucx_split_batch_size"] = self._extra[ + "nixl_ucx_split_batch_size" + ] + self._runner = KVCR( + config=KVCRConfig( + nixl_agent_name=self._agent_name, + nixl_listen_port=nixl_port, + pool_layouts=layouts, + enable_telemetry=self._diagnostic_telemetry, + operation_timeout_ms=self._extra.get( + "operation_timeout_ms", 30000 + ), + abandon_timeout_ms=self._extra.get("abandon_timeout_ms", 60000), + **startup_options, + ), + bindings=KVCRBindings( + request_pin=self._pinning.request_pin, + poll_pin_results=self._pinning.poll_pin_results, + release_pin=self._pinning.release_pin, + framework_control=channel, + key_adapter=StorageKeyAdapter(), + on_resilience_event=self._resilience_event, + stats_factory=( + DiagnosticStats if self._diagnostic_telemetry else None + ), + inventory_sink=( + self._diagnostic_inventory + if self._diagnostic_telemetry + else None + ), + ), + backend_configs=KVCRBackendConfigs( + framework_regions=regions, + local_dram=( + LocalDramOptions(pools=allocations) if allocations else None + ), + ), + ) + except KVCRStartupError: + # Native work still owns these mappings. They cannot be reclaimed. + resources.pop_all() + logger.critical( + "Nonquiescent KVCR startup; terminating worker", exc_info=True + ) + os._exit(1) + except Exception: + # A dead prefetch thread otherwise strands wait_complete while + # unrelated HTTP health checks continue to report success. + logger.critical( + "KVCR startup failed; terminating worker", exc_info=True + ) + os._exit(1) + self._maps = mappings + resources.pop_all() + # Peer source operations also need the public main-side poll while this + # worker is idle. Polling only from a local read/write strands source + # pin/refusal and lease-release callbacks until another request arrives. + self._poll_thread = threading.Thread( + target=self._poll_loop, name="kvcr-storage-completions", daemon=True + ) + self._poll_thread.start() + logger.info( + "KVCR L3 ready: endpoint=%s physical_pools=%s bytes=%d namespace=%s", + channel.endpoint, + layouts, + sum(size for _, _, size in allocations), + self._namespace, + ) + + @staticmethod + def _resilience_event(error): + from kvcr.types import TransferError + + if isinstance(error, TransferError) and error.state == "uncertain": + # HiCache cannot quarantine arbitrary source/destination pages yet. + # Returning a miss would let it reuse memory still exposed to NIXL. + logger.critical("Unsafe KVCR transfer; terminating worker: %s", error) + os._exit(1) + logger.warning("KVCR resilience event: %s", error) + + def _operation_hint(self, extra_info, keys=None): + hint = ( + peer_hint(extra_info, self._source_offset) + if self._extra.get("enable_remote_hint", True) + else None + ) + if hint is not None: + return hint + if not self._share_replicated_mla or self.config.tp_rank == 0 or not keys: + return None + # Readers without an external hint use their own replica's TP0. Their + # query is advisory; TP0's real local query plus HiCache's MIN reduction + # bounds the hit length. Final delivery results are also rank-reduced. + return { + "protocol_version": "0.1", + "message_id": "replica-local", + "actions": [ + { + "action_id": "replica-local", + "action_type": "kv.fetch", + "action_version": "1.0", + "payload": { + "source_control_endpoint": self._owner_endpoint, + "block_hashes": list( + dict.fromkeys(int(key[:16], 16) for key in keys) + ), + "mode": "copy", + }, + } + ], + } + + @contextmanager + def _hint(self, extra_info, keys=None): + request_id = None + hint = self._operation_hint(extra_info, keys) + if hint is not None: + request_id = uuid.uuid4().hex + self._runner.submit_hint(hint, request_id=request_id) + try: + yield request_id + finally: + if request_id is not None: + self._runner.discard_hint(request_id) + + def _wait(self, handle): + # There can be concurrent backup and prefetch threads. One caller may + # drain another caller's completion; keep it until the owner consumes it. + with self._completion: + while handle not in self._completed: + self._completed.update(self._runner.poll_completed()) + if handle not in self._completed: + self._completion.wait(timeout=0.001) + result = self._completed.pop(handle) + self._completion.notify_all() + return result + + def _poll_loop(self): + while not self._stop_poll.is_set(): + try: + with self._completion: + self._completed.update(self._runner.poll_completed()) + self._completion.notify_all() + if self._diagnostic_telemetry: + self._drain_diagnostics() + except Exception: + # Native progress failure cannot safely be treated as a cache + # miss: HiCache could recycle a buffer still exposed to NIXL. + logger.critical( + "KVCR L3 progress failed; terminating worker", exc_info=True + ) + os._exit(1) + self._stop_poll.wait(0.001) + + def _diagnostic_inventory(self, event): + # Called on the public main side. Full storage keys permit exact offline + # joins across peers; u64 router hints alone cannot distinguish pools. + # Published removal proves absence then, not that every later miss was + # an eviction. A later store event supersedes the removal. + logger.info( + "KVCR L3 inventory %s", + json.dumps( + { + "epoch": time.time(), + "agent": self._agent_name, + "endpoint": self._control_endpoint, + "namespace": self._namespace, + "tier": event.tier.value, + "removed": event.removed, + "keys": [key.decode() for key in event.keys], + }, + separators=(",", ":"), + ), + ) + + def _drain_diagnostics(self, *, force=False): + now = time.monotonic() + if not force and now - self._diagnostic_epoch < 5: + return + stats = self._runner.get_stats() + if stats is not None: + logger.info( + "KVCR L3 native stats %s", + json.dumps( + { + "epoch": time.time(), + "agent": self._agent_name, + "endpoint": self._control_endpoint, + "interval_seconds": now - self._diagnostic_epoch, + "histogram_bounds_seconds": stats.BOUNDS, + "metrics": stats.reduce(), + }, + separators=(",", ":"), + ), + ) + self._diagnostic_epoch = now + + def _transfer_digests(self, blocks, results): + """Diagnostic-only hashes while HiCache still protects the host pages. + + Use after native completion, before returning the pages to the + controller. This reads CPU buffers only, and covers every component. + It does not prove the subsequent host-to-GPU copy or model numerics. + """ + components = { + component.name: component + for layout in self._layouts.values() + for component in layout.components + } + digests = {} + for key, refs in blocks.items(): + if key not in results or not results[key].success: + continue + digest = hashlib.blake2b(digest_size=32) + for ref in refs: + component = components[ref.label] + digest.update(ref.label.encode()) + digest.update( + ctypes.string_at( + component.address + ref.element_index * component.stride, + component.size, + ) + ) + digests[key.decode()] = digest.hexdigest() + return digests + + def _io(self, transfers, *, is_set, extra_info): + from kvcr.types import MemoryRef + + if is_set and self._share_replicated_mla and self.config.tp_rank != 0: + raise RuntimeError("Replicated MLA L3 backup must run on TP0") + diagnostic = self._diagnostic_telemetry + entered = time.perf_counter() if diagnostic else None + entered_epoch = time.time() if diagnostic else None + with self._lock: + acquired = time.perf_counter() if diagnostic else None + self._start() + blocks, groups = {}, {} + for transfer in transfers: + name = str(transfer.name) + keys = transfer.keys or [] + if name == str(PoolName.KV) and name not in self._layouts: + groups[name] = [(None, True) for _ in keys] + continue + layout = self._layouts[name] + slots = page_slots(layout, keys, transfer.host_indices) + groups.setdefault(name, []) + for key, slot in zip(keys, slots, strict=True): + identity = block_key(key, self._namespace, name) + refs = [ + MemoryRef( + end_point_name=self._agent_name, + label=component.name, + element_index=slot, + ) + for component in layout.components + ] + if identity in blocks: + raise ValueError("Duplicate pool/page in one KVCR operation") + blocks[identity] = refs + groups[name].append((identity, False)) + if not blocks: + return {name: [True] * len(group) for name, group in groups.items()} + if not is_set and self._delivery_pipeline_window == 4: + validate_destination_spans(blocks, self._layouts) + described = time.perf_counter() if diagnostic else None + started = time.perf_counter() if self._metrics is not None else None + sequence = None + if is_set and self._align_storage_prefix: + # HiCache owns a different LRU. A reused parent may never be + # read from KVCR before its new tail is backed up. Refresh it + # BEFORE deposit can allocate/evict, and again after completion + # to give the new pages the same recency and tail-first order. + # This changes eviction preference, not residency or leases: + # align_sequence ignores absent/non-ready native blocks. + # Enable hicache_storage_pass_prefix_keys to include parents. + prefix = getattr(extra_info, "prefix_keys", None) or [] + sequence = [ + block_key(key, self._namespace, str(transfer.name)) + for transfer in transfers + if str(transfer.name) in self._layouts + for key in [*prefix, *(transfer.keys or [])] + ] + self._runner.align_sequence(sequence, use_current_time=True) + aligned = time.perf_counter() if diagnostic else None + with self._hint(extra_info, keys=None if is_set else blocks) as request_id: + prepared = time.perf_counter() if diagnostic else None + if not is_set and self._delivery_pipeline_window == 4: + result = self._deliver_pipeline( + blocks, + transfers, + extra_info, + request_id, + entered=entered, + entered_epoch=entered_epoch, + acquired=acquired, + described=described, + aligned=aligned, + prepared=prepared, + ) + else: + handle = ( + self._runner.deposit(blocks) + if is_set + else self._runner.deliver(blocks, request_id=request_id) + ) + submitted = time.perf_counter() if diagnostic else None + # _wait releases the condition's recursive lock while sleeping. + result = self._wait(handle) + finished = time.perf_counter() if diagnostic else None + if diagnostic: + logger.info( + "KVCR L3 operation %s", + json.dumps( + { + "epoch": time.time(), + "start_epoch": entered_epoch, + "agent": self._agent_name, + "endpoint": self._control_endpoint, + "operation": handle, + "stage": "primary", + "direction": "deposit" if is_set else "deliver", + "source": ( + self._operation_hint(extra_info, blocks)[ + "actions" + ][0]["payload"]["source_control_endpoint"] + if request_id is not None + else None + ), + "lock_wait_ms": (acquired - entered) * 1000, + "prepare_ms": (prepared - acquired) * 1000, + "descriptor_ms": (described - acquired) * 1000, + "align_before_ms": (aligned - described) * 1000, + "hint_ms": (prepared - aligned) * 1000, + "submit_ms": (submitted - prepared) * 1000, + "wait_ms": (finished - submitted) * 1000, + "requested_blocks": len(blocks), + "requested_spans": sum(map(len, blocks.values())), + "pools": [str(t.name) for t in transfers], + "failed_keys": [ + key.decode() + for key in blocks + if key not in result or not result[key].success + ], + }, + separators=(",", ":"), + ), + ) + if ( + not is_set + and self._share_replicated_mla + and self.config.tp_rank != 0 + and request_id is not None + ): + failed = { + key: refs + for key, refs in blocks.items() + if key not in result or not result[key].success + } + if ( + failed + and self._operation_hint(extra_info, blocks)["actions"][0][ + "payload" + ]["source_control_endpoint"] + != self._owner_endpoint + ): + # An external hint must not hide this replica's own + # TP0-resident pages. Retry only terminal misses, under + # a fresh scope, while HiCache still owns every slot. + # Successful external pages are never written again; + # uncertain native transfers still terminate the worker. + fallback_epoch = time.time() if diagnostic else None + fallback_started = ( + time.perf_counter() if diagnostic else None + ) + with self._hint(None, failed) as owner_request_id: + fallback_prepared = ( + time.perf_counter() if diagnostic else None + ) + fallback_handle = self._runner.deliver( + failed, request_id=owner_request_id + ) + fallback_submitted = ( + time.perf_counter() if diagnostic else None + ) + recovered = self._wait(fallback_handle) + fallback_finished = ( + time.perf_counter() if diagnostic else None + ) + if diagnostic: + logger.info( + "KVCR L3 operation %s", + json.dumps( + { + "epoch": time.time(), + "start_epoch": fallback_epoch, + "agent": self._agent_name, + "endpoint": self._control_endpoint, + "operation": fallback_handle, + "parent_operation": handle, + "stage": "replica-local-fallback", + "direction": "deliver", + "source": self._owner_endpoint, + "lock_wait_ms": 0, + "prepare_ms": ( + fallback_prepared - fallback_started + ) + * 1000, + "submit_ms": ( + fallback_submitted - fallback_prepared + ) + * 1000, + "wait_ms": ( + fallback_finished - fallback_submitted + ) + * 1000, + "requested_blocks": len(failed), + "requested_spans": sum( + map(len, failed.values()) + ), + "pools": [str(t.name) for t in transfers], + "failed_keys": [ + key.decode() + for key in failed + if key not in recovered + or not recovered[key].success + ], + }, + separators=(",", ":"), + ), + ) + result = { + **result, + **{ + key: recovered[key] + for key in failed + if key in recovered + }, + } + if sequence is not None: + self._runner.align_sequence(sequence, use_current_time=True) + if self._verify_transfer_bytes: + logger.info( + "KVCR L3 byte verification %s", + json.dumps( + { + "direction": "deposit" if is_set else "deliver", + "namespace": self._namespace, + "digests": self._transfer_digests(blocks, result), + }, + sort_keys=True, + ), + ) + if started is not None: + self._record_metrics( + transfers, + blocks, + result, + elapsed=time.perf_counter() - started, + is_set=is_set, + ) + return { + name: [ + logical or (key in result and result[key].success) + for key, logical in group + ] + for name, group in groups.items() + } + + def _deliver_pipeline( + self, + blocks, + transfers, + extra_info, + request_id, + *, + entered, + entered_epoch, + acquired, + described, + aligned, + prepared, + ): + """Run inside the primary hint; retain original per-handle diagnostics.""" + seen_handles = set() + primary = run_window( + self, + (Delivery(chunk) for chunk in block_chunks(blocks)), + request_id, + seen_handles, + ) + # No result objects or diagnostic callbacks are touched while any + # primary write is outstanding. In particular, a miss cannot stop drain. + result = requested_results(primary) + hint = self._operation_hint(extra_info, blocks) if request_id else None + source = ( + hint["actions"][0]["payload"]["source_control_endpoint"] if hint else None + ) + fallback = [] + if ( + self._share_replicated_mla + and self.config.tp_rank != 0 + and request_id is not None + and source != self._owner_endpoint + ): + # Retain the primary chunk as the fallback's parent even if a + # subset of that chunk failed. Never resubmit successful entries. + for delivery in primary: + failed = { + key: refs + for key, refs in delivery.blocks.items() + if key not in result or not result[key].success + } + if failed: + fallback.append(Delivery(failed, delivery.operation)) + if fallback: + failed_blocks = { + key: refs for item in fallback for key, refs in item.blocks.items() + } + fallback_started = ( + time.perf_counter() if self._diagnostic_telemetry else None + ) + with self._hint(None, failed_blocks) as owner_request_id: + fallback_prepared = ( + time.perf_counter() if self._diagnostic_telemetry else None + ) + fallback = run_window( + self, fallback, owner_request_id, seen_handles + ) + result.update(requested_results(fallback)) + if self._diagnostic_telemetry: + # Both phases are now terminal, including all misses. Logging, + # JSON encoding and user-defined result properties may all throw. + for chunk_index, delivery in enumerate([*primary, *fallback]): + is_primary = delivery.parent_operation is None + first_primary = delivery is primary[0] + first_fallback = bool(fallback) and delivery is fallback[0] + fields = { + "epoch": delivery.finished_epoch, + "start_epoch": delivery.prepared_epoch, + "logical_start_epoch": entered_epoch, + "delivery_pipeline_window": 4, + "pipeline_chunk_index": chunk_index, + "agent": self._agent_name, + "endpoint": self._control_endpoint, + "operation": delivery.operation, + "stage": "primary" if is_primary else "replica-local-fallback", + "direction": "deliver", + "source": source if is_primary else self._owner_endpoint, + "lock_wait_ms": (acquired - entered) * 1000 if first_primary else 0, + "prepare_ms": ( + prepared - acquired + if first_primary + else fallback_prepared - fallback_started + if first_fallback + else 0 + ) + * 1000, + "submit_ms": (delivery.submitted - delivery.prepared) * 1000, + "wait_ms": (delivery.finished - delivery.submitted) * 1000, + "requested_blocks": len(delivery.blocks), + "requested_spans": sum(map(len, delivery.blocks.values())), + "pools": [str(t.name) for t in transfers], + "failed_keys": [ + key.decode() + for key in delivery.blocks + if key not in delivery.result + or not delivery.result[key].success + ], + } + if is_primary: + fields.update( + descriptor_ms=(described - acquired) * 1000 + if first_primary + else 0, + align_before_ms=(aligned - described) * 1000 + if first_primary + else 0, + hint_ms=(prepared - aligned) * 1000 if first_primary else 0, + ) + else: + fields["parent_operation"] = delivery.parent_operation + logger.info( + "KVCR L3 operation %s", json.dumps(fields, separators=(",", ":")) + ) + return result + + def _record_metrics(self, transfers, blocks, results, *, elapsed, is_set): + """Count completed logical pages once, but every transferred byte. + + A primary page with a failed requested sidecar is not a usable page. + Its successfully moved primary bytes still consumed transport bandwidth. + HiCache separately accounts for admission and resume-boundary validity. + Called with the IO lock held, after terminal native completion. + """ + components = { + component.name: component + for layout in self._layouts.values() + for component in layout.components + } + page_results, transferred_bytes = {}, 0 + for identity, refs in blocks.items(): + successful = identity in results and results[identity].success + key = identity.split(b"#", 1)[0].decode() + page_results.setdefault(key, []).append(successful) + if successful: + transferred_bytes += sum(components[ref.label].size for ref in refs) + primary_keys = { + key + for transfer in transfers + if str(transfer.name) == str(PoolName.KV) + for key in (transfer.keys or []) + } + pages = sum(all(page_results.get(key, [False])) for key in primary_keys) + bandwidth = transferred_bytes / max(elapsed, 1e-9) / 1e9 + if is_set: + self._metrics.backup_pgs.append(pages) + self._metrics.backup_bandwidth.append(bandwidth) + else: + self._metrics.prefetch_pgs.append(pages) + self._metrics.prefetch_bandwidth.append(bandwidth) + + def get_stats(self): + """Atomically hand HiCache its standard interval snapshot.""" + with self._lock: + if self._metrics is None: + return None + from sglang.srt.observability.metrics_collector import StorageMetrics + + snapshot, self._metrics = self._metrics, StorageMetrics() + return snapshot + + def batch_get_v1(self, keys, host_indices, extra_info=None): + return self.batch_get_v2( + [PoolTransfer(PoolName.KV, host_indices=host_indices, keys=keys)], + extra_info, + )[str(PoolName.KV)] + + def batch_set_v1(self, keys, host_indices, extra_info=None): + return self.batch_set_v2( + [PoolTransfer(PoolName.KV, host_indices=host_indices, keys=keys)], + extra_info, + )[str(PoolName.KV)] + + def batch_get_v2(self, transfers, extra_info=None): + return self._io(transfers, is_set=False, extra_info=extra_info) + + def batch_set_v2(self, transfers, extra_info=None): + return self._io(transfers, is_set=True, extra_info=extra_info) + + def batch_exists_v2(self, keys, pool_transfers=None, extra_info=None): + from kvcr.types import QueryStatus + + required = list(pool_transfers or []) + if str(PoolName.KV) in self._layouts: + required = [PoolTransfer(PoolName.KV)] + required + elif not required: + # A logical KV anchor has no payload to prove a page. Answer a + # bare KV query from the pools that hold every stored page. + required = [ + PoolTransfer(name) + for name in self._layouts + if name not in _TRAILING_POOLS + ] + if not required or not keys: + # An empty intersection is not evidence of stored data. In + # particular, trailing-only layouts cannot prove a bare prefix. + return PoolTransferResult(0, {}, []) + with self._lock: + self._start() + identities = [ + block_key(key, self._namespace, str(transfer.name)) + for transfer in required + for key in keys + ] + with self._hint(extra_info, identities) as request_id: + statuses = self._runner.query(identities, request_id=request_id) + restorable, hits = set(range(1, len(keys) + 1)), {} + for index, transfer in enumerate(required): + status = statuses[index * len(keys) : (index + 1) * len(keys)] + exists = [s in (QueryStatus.HIT, QueryStatus.FETCHABLE) for s, _ in status] + trailing = None + if transfer.hit_policy == PoolHitPolicy.TRAILING_PAGES: + trailing = max(1, len(transfer.keys or [])) + elif transfer.hit_policy != PoolHitPolicy.ALL_PAGES: + raise ValueError(f"Unsupported pool hit policy: {transfer.hit_policy}") + valid = resume_boundaries(exists, trailing=trailing) + hits[str(transfer.name)] = max(valid, default=0) + restorable.intersection_update(valid) + valid = sorted(restorable) + return PoolTransferResult(max(valid, default=0), hits, valid) + + def batch_exists(self, keys, extra_info=None): + return self.batch_exists_v2(keys, extra_info=extra_info).kv_hit_pages + + def exists(self, key): + return bool(self.batch_exists([key])) + + # Tensor-valued v0 IO cannot identify registered, lifetime-protected pages. + # The factory/controller explicitly select the indexed v1/v2 path instead. + def get(self, key, target_location=None, target_sizes=None): + raise NotImplementedError("KVCR requires indexed HiCache host pages") + + def batch_get(self, keys, target_locations=None, target_sizes=None): + raise NotImplementedError("KVCR requires indexed HiCache host pages") + + def set(self, key, value=None, target_location=None, target_sizes=None): + raise NotImplementedError("KVCR requires indexed HiCache host pages") + + def batch_set(self, keys, values=None, target_locations=None, target_sizes=None): + raise NotImplementedError("KVCR requires indexed HiCache host pages") + + def close(self): + with self._lock: + if self._closed: + return + self._closed = True + self._stop_poll.set() + # Do not hold the IO lock while joining its completion-poll owner. + if self._poll_thread is not None: + self._poll_thread.join() + with self._lock: + if self._runner is not None: + try: + if self._diagnostic_telemetry: + self._drain_diagnostics(force=True) + self._runner.close() + except Exception: + logger.critical( + "KVCR L3 close failed; retaining buffers until process exit", + exc_info=True, + ) + os._exit(1) + for allocation in self._maps: + allocation.close() + self._maps.clear() diff --git a/python/sglang/srt/mem_cache/storage/kvcr/layout.py b/python/sglang/srt/mem_cache/storage/kvcr/layout.py new file mode 100644 index 000000000000..afccf65aa4ea --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/kvcr/layout.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Translate stable host-page layouts into indexed KVCR registrations.""" + +from __future__ import annotations + +import hashlib +import json + +import msgspec +import torch + + +class PageComponent(msgspec.Struct, frozen=True): + name: str + address: int + size: int + stride: int + count: int + + +class PoolLayout(msgspec.Struct, frozen=True): + name: str + page_size: int + components: tuple[PageComponent, ...] + format: str + + @property + def bytes_per_page(self) -> int: + return sum(component.size for component in self.components) + + +def describe_pool(name: str, pool) -> PoolLayout: + """Derive the layout from the same accessor that performs storage IO. + + Each component may have a different element size (e.g. temporal and conv + state), so it gets its own KVCR allocation pool. A logical page groups all + its components under one block key; a successful read must carry them all. + """ + count = pool.size // pool.page_size + if pool.layout not in ("page_first", "page_first_direct") or count < 2: + raise ValueError("KVCR requires at least two host slots in a page-first layout") + pointers, sizes = pool.get_page_buffer_meta(torch.arange(2 * pool.page_size)) + if not pointers or len(pointers) % 2 or len(pointers) != len(sizes): + raise ValueError(f"Invalid host-page metadata for {name}") + parts = len(pointers) // 2 + components = [] + for part in range(parts): + size = sizes[part] + stride = pointers[parts + part] - pointers[part] + if size <= 0 or sizes[parts + part] != size or stride < size: + raise ValueError(f"Nonuniform host-page component in {name}") + components.append( + PageComponent(f"{name}_{part}", pointers[part], size, stride, count) + ) + return PoolLayout( + name, pool.page_size, tuple(components), f"{pool.layout}:{pool.dtype}" + ) + + +def page_slots(layout: PoolLayout, keys: list[str], indices: torch.Tensor) -> list[int]: + """Validate whole, contiguous physical pages before exposing memory to NIXL.""" + if indices.device.type != "cpu" or indices.ndim != 1: + raise ValueError("KVCR storage indices must be a one-dimensional CPU tensor") + if indices.dtype not in (torch.int32, torch.int64): + raise ValueError("KVCR storage indices must be integers") + if len(indices) != len(keys) * layout.page_size: + raise ValueError("Keys and host-page indices have different lengths") + pages = indices.reshape(-1, layout.page_size) + starts = pages[:, 0] + if len(keys) and ( + torch.any(starts < 0) + or torch.any(starts % layout.page_size != 0) + or torch.any(starts // layout.page_size >= layout.components[0].count) + or not torch.equal( + pages, starts[:, None] + torch.arange(layout.page_size)[None, :] + ) + ): + raise ValueError("KVCR received unaligned, noncontiguous or out-of-range pages") + return (starts // layout.page_size).tolist() + + +def storage_namespace( + config, layouts: dict[str, PoolLayout], *, shared_mla=False +) -> str: + """Model, byte format and shard identity, but not replica/DP identity.""" + identity = [ + "sglang-kvcr-shared-mla-v1" if shared_mla else "sglang-kvcr-v1", + config.model_name, + config.tp_size, + 0 if shared_mla else config.tp_rank, + config.pp_size, + config.pp_rank, + config.attn_cp_size, + config.attn_cp_rank, + [ + [ + name, + layout.page_size, + layout.format, + [part.size for part in layout.components], + ] + for name, layout in sorted(layouts.items()) + ], + ] + return hashlib.sha256(json.dumps(identity).encode()).hexdigest() + + +def block_key(page_key: str, namespace: str, pool: str) -> bytes: + # Keep the complete hash for storage; only hint membership uses its u64 prefix. + if len(page_key) != 64: + raise ValueError("KVCR expects a full SHA256 storage-page hash") + bytes.fromhex(page_key) + return f"{page_key}#{namespace}#{pool}".encode() + + +class StorageKeyAdapter: + def encode(self, framework_key: object) -> bytes: + if not isinstance(framework_key, bytes): + raise TypeError("KVCR storage keys must be bytes") + return framework_key + + def decode(self, key: bytes) -> int: + # Inverts SGLang's hash_str_to_int64, as the unsigned router hash. + return int(key[:16], 16) + + +def resume_boundaries(page_exists: list[bool], *, trailing: int | None) -> list[int]: + """A missing checkpoint creates holes, not a monotonically shrinking prefix.""" + if trailing is None: + boundary = next( + (i for i, hit in enumerate(page_exists) if not hit), len(page_exists) + ) + return list(range(1, boundary + 1)) + if trailing < 1: + raise ValueError("Trailing windows must contain at least one page") + return [ + end + for end in range(1, len(page_exists) + 1) + if all(page_exists[max(0, end - trailing) : end]) + ] diff --git a/python/sglang/srt/mem_cache/storage/kvcr/numa_pool.py b/python/sglang/srt/mem_cache/storage/kvcr/numa_pool.py new file mode 100644 index 000000000000..b6c1983ac24b --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/kvcr/numa_pool.py @@ -0,0 +1,284 @@ +"""Opt-in ownership-scoped NUMA allocation for HiCache and KVCR DRAM. + +Only NEW mappings are bound. Never change a task policy or migrate old pages. +The ordinary allocator remains the default when host_pool_numa_node is absent. +""" + +import ctypes +import ctypes.util +import json +import logging +import math +import mmap +import os +import re +from pathlib import Path + +logger = logging.getLogger(__name__) +HOST_PROOF_ATTR = "_sglang_owned_numa_host_proof" +MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23) + + +def node_option(value, name): + if value is not None and (type(value) is not int or not 0 <= value < 65536): + raise ValueError(f"{name} must be a nonnegative integer NUMA node or null") + return value + + +def load_config(raw): + if raw is None or raw == "": + return {} + if not isinstance(raw, str): + raise ValueError( + "KVCR host allocator extra config must be a JSON string or @file" + ) + if raw.startswith("@"): + path = Path(raw[1:]) + if not raw[1:]: + raise ValueError("KVCR host allocator config path is empty") + if path.suffix.lower() == ".json": + value = json.loads(path.read_text()) + elif path.suffix.lower() == ".toml": + import tomllib + + with path.open("rb") as stream: + value = tomllib.load(stream) + elif path.suffix.lower() in (".yaml", ".yml"): + import yaml + + value = yaml.safe_load(path.read_text()) + else: + raise ValueError("Unsupported KVCR host allocator config file extension") + else: + value = json.loads(raw) + if not isinstance(value, dict): + raise ValueError("KVCR host allocator extra config must contain an object") + return value + + +def prove_range(maps, numa, smaps, address, length, node, shared): + """A fully resident uniform covering VMA proves this exact owned subrange.""" + if address <= 0 or length <= 0 or address % mmap.PAGESIZE: + raise RuntimeError("Invalid owned NUMA mapping range") + rows = [] + for line in maps.splitlines(): + fields = line.split(None, 5) + lo, hi = (int(v, 16) for v in fields[0].split("-")) + if lo <= address and address + length <= hi: + rows.append((lo, hi, fields, line)) + if len(rows) != 1: + raise RuntimeError("Owned NUMA range has no unique complete covering VMA") + lo, hi, fields, line = rows[0] + if fields[1] != ("rw-s" if shared else "rw-p"): + raise RuntimeError("Owned NUMA mapping sharing mode differs") + if shared: + if ( + len(fields) != 6 + or fields[5] != "/dev/zero (deleted)" + or int(fields[4]) <= 0 + ): + raise RuntimeError("Owned host mapping is not new shared anonymous memory") + elif len(fields) != 5 or fields[4] != "0": + raise RuntimeError("Owned local DRAM mapping is not private anonymous memory") + policy = [row for row in numa.splitlines() if int(row.split()[0], 16) == lo] + if len(policy) != 1 or policy[0].split()[1] != f"bind:{node}": + raise RuntimeError( + "Requested strict owned-mapping NUMA policy was not installed" + ) + counts = {int(k): int(v) for k, v in re.findall(r"\bN(\d+)=(\d+)", policy[0])} + if ( + counts.get(node) != (hi - lo) // mmap.PAGESIZE + or sum(counts.values()) != (hi - lo) // mmap.PAGESIZE + ): + raise RuntimeError("Owned NUMA mapping is incomplete or has wrong-node pages") + header = re.search(r"(?m)^" + re.escape(fields[0]) + r" .*$", smaps) + if header is None: + raise RuntimeError("Owned NUMA mapping has no smaps evidence") + segment = smaps[header.end() :] + following = re.search(r"(?m)^[0-9a-f]+-[0-9a-f]+ ", segment) + if following: + segment = segment[: following.start()] + backing = {} + for key in ("Size", "Rss", "AnonHugePages", "KernelPageSize", "ShmemPmdMapped"): + match = re.search(r"(?m)^" + key + r":\s*(\d+) kB", segment) + if match is None: + raise RuntimeError("Incomplete owned NUMA page-backing evidence") + backing[key] = int(match[1]) + if backing["Size"] != (hi - lo) // 1024 or backing["Rss"] != backing["Size"]: + raise RuntimeError("Owned NUMA mapping is not fully resident") + return dict( + address=address, + length=length, + end=address + length, + node=node, + covering_start=lo, + covering_end=hi, + maps_line=line, + numa_line=policy[0], + node_pages=counts, + smaps=backing, + shared=shared, + all_covering_pages_resident_on_requested_node=True, + ) + + +def snapshot_range(address, length, node, shared): + return prove_range( + *( + Path("/proc/self/" + name).read_text() + for name in ("maps", "numa_maps", "smaps") + ), + address, + length, + node, + shared, + ) + + +def bind_new_mapping_before_use(allocation, node, *, shared): + """Call only for a newly created owned, unfaulted mapping; flags=0.""" + node = node_option(node, "owned_mapping_numa_node") + if node is None: + raise ValueError("An explicit NUMA node is required for strict binding") + node_path = Path(f"/sys/devices/system/node/node{node}") + if not node_path.is_dir(): + raise RuntimeError(f"Requested NUMA node {node} is not present") + address = ctypes.addressof(ctypes.c_char.from_buffer(allocation)) + length = len(allocation) + if address % mmap.PAGESIZE or length <= 0 or length % mmap.PAGESIZE: + raise ValueError("Owned NUMA mapping must be page-aligned") + name = ctypes.util.find_library("numa") + if name is None: + raise RuntimeError("Strict owned-mapping NUMA placement requires libnuma") + library = ctypes.CDLL(name, use_errno=True) + library.mbind.argtypes = [ + ctypes.c_void_p, + ctypes.c_ulong, + ctypes.c_int, + ctypes.POINTER(ctypes.c_ulong), + ctypes.c_ulong, + ctypes.c_uint, + ] + library.mbind.restype = ctypes.c_long + word_bits = ctypes.sizeof(ctypes.c_ulong) * 8 + mask = (ctypes.c_ulong * (node // word_bits + 1))() + mask[node // word_bits] = 1 << (node % word_bits) + ctypes.set_errno(0) + if library.mbind(address, length, 2, mask, len(mask) * word_bits + 1, 0) != 0: + error = ctypes.get_errno() + raise RuntimeError(f"Owned-mapping mbind failed for node {node}: errno {error}") + # A late bind is not a substitute: prefault only AFTER installing the policy. + allocation.madvise(MADV_POPULATE_WRITE) + return snapshot_range(address, length, node, shared) + + +def allocate_bound_host(dims, dtype, device, node): + import torch + from sglang.srt.environ import envs + + if device != "cpu": + raise ValueError("KVCR NUMA host allocator only accepts CPU tensors") + if envs.SGLANG_HUGEPAGE_SIZE.get(): + raise ValueError( + "Explicit hugeTLB host allocation is not qualified with KVCR NUMA placement" + ) + elements = math.prod(dims) + if elements <= 0: + raise ValueError("KVCR NUMA host allocation must contain at least one element") + count = elements * torch.empty([], dtype=dtype).element_size() + length = math.ceil(count / mmap.PAGESIZE) * mmap.PAGESIZE + allocation = mmap.mmap( + -1, + length, + flags=mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, + prot=mmap.PROT_READ | mmap.PROT_WRITE, + ) + try: + proof = bind_new_mapping_before_use(allocation, node, shared=True) + tensor = torch.frombuffer(allocation, dtype=dtype, count=elements).reshape(dims) + except BaseException: + allocation.close() + raise + setattr(tensor, HOST_PROOF_ATTR, proof) + return tensor + + +def allocator_for_config(raw): + from sglang.srt.mem_cache.pool_host.common import HostTensorAllocator + + config = load_config(raw) + node = node_option(config.get("host_pool_numa_node"), "host_pool_numa_node") + if node is None: + return HostTensorAllocator() + + class BoundHostTensorAllocator(HostTensorAllocator): + def allocate(self, dims, dtype, device): + self.dtype, self.dims = dtype, dims + return allocate_bound_host(dims, dtype, device, node) + + return BoundHostTensorAllocator() + + +def audit_registered_host(tensor, registrations): + """Join actual cudaHostRegister intervals to the owned allocation proof.""" + proof = getattr(tensor, HOST_PROOF_ATTR) + size = tensor.numel() * tensor.element_size() + address = tensor.data_ptr() + if address != proof["address"] or size > proof["length"]: + raise RuntimeError( + "CUDA-registered tensor does not match its owned NUMA allocation" + ) + cursor = address + for pointer, length in registrations: + if pointer != cursor or length <= 0: + raise RuntimeError( + "CUDA host registrations do not completely cover the owned tensor" + ) + cursor += length + if cursor != address + size: + raise RuntimeError( + "CUDA host registrations have incomplete or excessive owned-tensor coverage" + ) + fresh = snapshot_range(address, proof["length"], proof["node"], True) + stable = ( + "address", + "length", + "end", + "node", + "covering_start", + "covering_end", + "maps_line", + "node_pages", + "shared", + ) + if any(fresh[key] != proof[key] for key in stable): + raise RuntimeError("Owned host placement changed during CUDA registration") + logger.info( + "KVCR owned host NUMA registration %s", + json.dumps( + dict( + pid=os.getpid(), + tensor_bytes=size, + registration_ranges=registrations, + placement=fresh, + task_policy_changed=False, + pages_migrated=False, + ) + ), + ) + + +def bind_new_private_pool(allocation, node): + proof = bind_new_mapping_before_use(allocation, node, shared=False) + logger.info( + "KVCR owned local DRAM NUMA placement %s", + json.dumps( + dict( + pid=os.getpid(), + placement=proof, + task_policy_changed=False, + pages_migrated=False, + ) + ), + ) + return proof diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmcache_unified_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmcache_unified_radix_cache.py index c502e2e2ec4d..4918420542f7 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmcache_unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/lmcache/lmcache_unified_radix_cache.py @@ -185,6 +185,7 @@ def prefetch_from_storage( extra_key: Optional[str] = None, cache_salt: Optional[str] = None, storage_hit_end: Optional[int] = None, + kv_hints=None, ) -> None: del last_hash, prefix_keys, storage_hit_end req_id = request.rid diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 37a06b11c154..d0423af5b311 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -1974,6 +1974,7 @@ def query_storage_hit_length( prefix_keys: Optional[list[str]] = None, extra_key: Optional[str] = None, cache_salt: Optional[str] = None, + kv_hints=None, ) -> int: """Probe L3 with the request namespace.""" if ( @@ -1997,6 +1998,7 @@ def query_storage_hit_length( prefetch_key, last_hash, prefix_keys, + kv_hints=kv_hints, ) _, storage_hit_count = self.cache_controller._storage_hit_query(operation) storage_hit_count_tensor = torch.tensor(storage_hit_count, dtype=torch.int) @@ -2019,6 +2021,7 @@ def prefetch_from_storage( extra_key: Optional[str] = None, cache_salt: Optional[str] = None, storage_hit_end: Optional[int] = None, + kv_hints=None, ) -> Optional[bool]: if not self.enable_storage or self.cache_controller is None: return @@ -2131,6 +2134,7 @@ def prefetch_from_storage( matched_prefix_tokens, aux_xfers or None, assume_stored=assume_stored, + kv_hints=kv_hints, ) operation = submission.operation if operation is None: diff --git a/test/manual/test_kvcr_storage_native.py b/test/manual/test_kvcr_storage_native.py new file mode 100644 index 000000000000..aa5b47c0d501 --- /dev/null +++ b/test/manual/test_kvcr_storage_native.py @@ -0,0 +1,653 @@ +"""Manual real-NIXL gate for the HiCacheStorage backend (no model required). + +Requires Linux, KVCR, NIXL/UCX and Torch. This intentionally does not substitute +the transport: two independent KVCR agents own their DRAM and exchange bytes. +Run: python test/manual/test_kvcr_storage_native.py -v +""" + +import hashlib +import itertools +import json +import socket +import unittest +from contextlib import ExitStack +from copy import deepcopy +from queue import Queue +from types import SimpleNamespace + +import torch + +from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation +from sglang.srt.mem_cache.hicache_storage import ( + HiCacheStorageConfig, + HiCacheStorageExtraInfo, + PoolName, + PoolTransfer, +) +from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + HybridCacheController, + StorageOperation, +) +from sglang.srt.mem_cache.storage.kvcr.kvcr_store import KVCRStore +from sglang.test.test_utils import CustomTestCase + + +class HostPages: + """Stable page-first CPU buffers with unequal-sized page components.""" + + def __init__(self, sizes, page_size=2, pages=16): + self.page_size, self.size = page_size, pages * page_size + self.layout, self.dtype = "page_first", torch.uint8 + self.buffers = [torch.zeros((pages, size), dtype=self.dtype) for size in sizes] + self.kv_buffer = self.buffers[0] + + def get_page_buffer_meta(self, indices): + pointers, sizes = [], [] + for start in indices.tolist()[:: self.page_size]: + for buffer in self.buffers: + pointers.append(buffer[start // self.page_size].data_ptr()) + sizes.append(buffer.shape[1]) + return pointers, sizes + + +# Port zero selects the outbound ephemeral range. NIXL/UCX can consume a +# previously probed port while initializing, before the control listener binds. +# This manual Linux fixture uses unique non-ephemeral ports in its own process; +# bind probes are still advisory, not a general cross-process reservation API. +_TEST_PORTS = itertools.count(20000) + + +def free_port(): + for _ in range(5000): + port = next(_TEST_PORTS) + if port >= 25000: + break + try: + with socket.socket() as sock: + sock.bind(("0.0.0.0", port)) + return port + except OSError: + continue + raise RuntimeError("No available non-ephemeral native-fixture control port") + + +def free_port_pair(): + for _ in range(100): + port = free_port() + # Consume the neighbor as well so separately selected TP pairs cannot + # repeatedly overlap before their listeners have been constructed. + next(_TEST_PORTS) + try: + with socket.socket() as first, socket.socket() as second: + first.bind(("0.0.0.0", port)) + second.bind(("0.0.0.0", port + 1)) + return port + except OSError: + continue + raise RuntimeError("No available neighboring control ports") + + +class TestKVCRNativeStorage(CustomTestCase): + def test_large_prefetch_batches_restore_every_component_through_real_nixl(self): + """Combined backups and three reads restore every real-NIXL component. + + The target starts cold; real independent agents and the real controller + must populate unequal KV/indexer components and produce the exact ACK + prefix sequence. This is not a GPU inference/performance gate. + """ + pages, page_size = 1100, 2 + keys = [hashlib.sha256(f"large-{i}".encode()).hexdigest() for i in range(pages)] + ports = [free_port(), free_port()] + while ports[0] == ports[1]: + ports[1] = free_port() + with ExitStack() as resources: + stores, pools = [], [] + for port in ports: + host = { + PoolName.KV: HostPages((8192, 4096), pages=pages), + PoolName.INDEXER: HostPages((2048,), pages=pages), + } + config = HiCacheStorageConfig( + tp_rank=0, + tp_size=1, + pp_rank=0, + pp_size=1, + attn_cp_rank=0, + attn_cp_size=1, + is_mla_model=True, + enable_storage_metrics=False, + is_page_first_layout=True, + model_name="glm53-large-read-byte-gate", + extra_config={ + "control_port": port, + "control_advertise_host": "127.0.0.1", + "local_dram_pool_pages": {"kv": pages, "indexer": pages}, + "combine_page_reads": True, + "combine_page_writes": True, + "prefetch_batch_pages": 512, + }, + ) + store = KVCRStore(config, host[PoolName.KV]) + store.register_mem_host_pool_v2( + host[PoolName.INDEXER], PoolName.INDEXER + ) + resources.callback(store.close) + stores.append(store) + pools.append(host) + source, target = stores + for pool_index, pool in enumerate(pools[0].values()): + for part, buffer in enumerate(pool.buffers): + pattern = ( + torch.arange(pages).reshape(-1, 1) * 31 + + torch.arange(buffer.shape[1]).reshape(1, -1) * 17 + + pool_index * 43 + + part * 11 + ) % 251 + buffer.copy_(pattern.to(torch.uint8)) + indices = torch.arange(pages * page_size) + backup = HybridCacheController.__new__(HybridCacheController) + backup.storage_backend, backup.page_size = source, page_size + backup.backup_skip = False + backup.mem_pool_host = SimpleNamespace( + kv_buffer=pools[0][PoolName.KV].kv_buffer, layout_lease=ExitStack + ) + stored = StorageOperation( + indices, + list(range(pages * page_size)), + hash_value=keys, + pool_transfers=[ + PoolTransfer(PoolName.INDEXER, indices_from_pool=PoolName.KV) + ], + ) + backup._page_backup(stored) + self.assertEqual(stored.completed_tokens, pages * page_size) + self.assertEqual( + stored.pool_storage_result.extra_pool_hit_pages, {"indexer": pages} + ) + self.assertEqual( + target.batch_exists_v2( + keys, [PoolTransfer(PoolName.INDEXER)] + ).kv_hit_pages, + 0, + ) + c = HiCacheController.__new__(HiCacheController) + c.storage_backend, c.page_size = target, page_size + c.page_get_func = c._page_get_zero_copy + c.prefetch_sync_queue = Queue() + op = PrefetchOperation("native-large-read", list(range(pages * page_size))) + op.host_indices, op.hash_value = indices, keys + op.pool_transfers = [ + PoolTransfer(PoolName.INDEXER, indices_from_pool=PoolName.KV) + ] + op.kv_hints = { + "protocol_version": "0.1", + "actions": [ + { + "action_type": "kv.fetch", + "action_version": "1.0", + "payload": { + "source_control_endpoint": f"tcp://127.0.0.1:{ports[0]}", + "block_hashes": [int(key[:16], 16) for key in keys], + }, + } + ], + } + self.assertEqual(c._page_transfer(op), pages) + acks = [ + c.prefetch_sync_queue.get_nowait().completed_tokens for _ in range(3) + ] + self.assertEqual(acks, [1024, 2048, 2200]) + self.assertTrue(c.prefetch_sync_queue.empty()) + for name in pools[0]: + for original, restored in zip( + pools[0][name].buffers, pools[1][name].buffers, strict=True + ): + self.assertTrue(torch.equal(original, restored)) + + def test_shared_mla_real_reader_restores_from_replica_owner_and_remote_owner(self): + """No secondary DRAM copy; local owner truth bounds advisory cold hits. + + Restore every primary/indexer byte into rank-one host slots through + real NIXL, first from its own TP0 and then an independent replica's + TP0, including a read mixing pages owned by the two replicas. An + external hint must not hide pages still resident on the reader's own + TP0. The secondary has no KVCR-owned local DRAM allocation. + """ + keys = [hashlib.sha256(f"shared-{i}".encode()).hexdigest() for i in range(4)] + remote_keys = [ + hashlib.sha256(f"shared-remote-{i}".encode()).hexdigest() for i in range(4) + ] + bases = [free_port_pair() for _ in range(2)] + while abs(bases[0] - bases[1]) < 2: + bases[1] = free_port_pair() + nixl = [free_port_pair() for _ in range(2)] + while ( + any(abs(a - b) < 2 for a in nixl for b in bases) + or abs(nixl[0] - nixl[1]) < 2 + ): + nixl = [free_port_pair() for _ in range(2)] + with ExitStack() as resources: + stores, pools = [], [] + for replica, rank in ((0, 0), (1, 0), (1, 1)): + host = { + PoolName.KV: HostPages((8192, 4096)), + PoolName.INDEXER: HostPages((2048,)), + } + config = HiCacheStorageConfig( + tp_rank=rank, + tp_size=2, + pp_rank=0, + pp_size=1, + attn_cp_rank=0, + attn_cp_size=1, + is_mla_model=True, + enable_storage_metrics=False, + is_page_first_layout=True, + model_name="glm53-replicated-native", + extra_config={ + "control_port": bases[replica], + "nixl_listen_port": nixl[replica], + "control_advertise_host": "127.0.0.1", + "local_dram_pool_pages": {"kv": 16, "indexer": 16}, + "share_replicated_mla": True, + }, + ) + store = KVCRStore(config, host[PoolName.KV]) + store.register_mem_host_pool_v2( + host[PoolName.INDEXER], PoolName.INDEXER + ) + resources.callback(store.close) + stores.append(store) + pools.append(host) + source, owner, reader = stores + # Real query statuses: TP0's local truth makes the cross-rank MIN + # zero despite the secondary's advisory own-replica coverage. + owner_hits = owner.batch_exists_v2( + keys, [PoolTransfer(PoolName.INDEXER)] + ).kv_hit_pages + reader_hits = reader.batch_exists_v2( + keys, [PoolTransfer(PoolName.INDEXER)] + ).kv_hit_pages + self.assertEqual(owner_hits, 0) + self.assertEqual(reader_hits, 4) + self.assertEqual(min(owner_hits, reader_hits), 0) + self.assertEqual(reader._maps, []) + + def transfers(indices, page_keys=keys): + return [ + PoolTransfer(name, keys=page_keys, host_indices=indices) + for name in pools[0] + ] + + for replica in (0, 1): + for pool in pools[replica].values(): + for part, buffer in enumerate(pool.buffers): + for page in range(4): + buffer[page].copy_( + ( + ( + torch.arange(buffer.shape[1]) * 17 + + replica * 31 + + part + + page + ) + % 251 + ).to(torch.uint8) + ) + result = stores[replica].batch_set_v2( + transfers(torch.arange(8), remote_keys if replica == 0 else keys) + ) + self.assertTrue(all(all(rows) for rows in result.values())) + self.assertEqual(owner._namespace, reader._namespace) + self.assertEqual(source._namespace, reader._namespace) + self.assertEqual( + owner.batch_exists_v2( + keys, [PoolTransfer(PoolName.INDEXER)] + ).kv_hit_pages, + 4, + ) + external = HiCacheStorageExtraInfo( + extra_info={ + "kv_hints": { + "protocol_version": "0.1", + "actions": [ + { + "action_type": "kv.fetch", + "action_version": "1.0", + "payload": { + "source_control_endpoint": f"tcp://127.0.0.1:{bases[0]}", + "block_hashes": [ + int(key[:16], 16) for key in remote_keys + ], + }, + } + ], + } + } + ) + for replica, info, indices, page_keys in ( + (1, None, torch.arange(8, 16), keys), + (0, external, torch.arange(16, 24), remote_keys), + ): + restored = reader.batch_get_v2(transfers(indices, page_keys), info) + self.assertTrue(all(all(rows) for rows in restored.values()), restored) + slot = int(indices[0]) // 2 + for name, host in pools[replica].items(): + for expected, actual in zip( + host.buffers, pools[2][name].buffers, strict=True + ): + self.assertTrue( + torch.equal(expected[:4], actual[slot : slot + 4]) + ) + self.assertEqual(reader._maps, []) + + # Router coverage is advisory, not a promise that every hinted + # page remains on that source. TP0 can combine its local pages + # with externally owned pages; its secondary must do so too. + mixed_keys = [keys[0], remote_keys[1], keys[2], remote_keys[3]] + mixed_info = deepcopy(external) + mixed_info.extra_info["kv_hints"]["actions"][0]["payload"][ + "block_hashes" + ] = [int(key[:16], 16) for key in mixed_keys] + for store, destination in ((owner, pools[1]), (reader, pools[2])): + with self.subTest(reader_rank=store.config.tp_rank): + self.assertEqual( + store.batch_exists_v2( + mixed_keys, [PoolTransfer(PoolName.INDEXER)], mixed_info + ).kv_hit_pages, + 4, + ) + restored = store.batch_get_v2( + transfers(torch.arange(24, 32), mixed_keys), mixed_info + ) + self.assertEqual( + restored, + {str(name): [True] * 4 for name in pools[0]}, + ) + for name, host in destination.items(): + for part, actual in enumerate(host.buffers): + for page, replica in enumerate((1, 0, 1, 0)): + self.assertTrue( + torch.equal( + pools[replica][name].buffers[part][page], + actual[12 + page], + ), + (name, part, page, replica), + ) + self.assertEqual(reader._maps, []) + + def test_prefix_alignment_preserves_bounded_native_lru(self): + """Real native allocation must evict an unrelated page, not a hot parent. + + Run an unaligned control and an aligned candidate at identical three- + page capacity for both KV and indexer. Query and restored bytes, not + a mocked callback, prove the effect of pre-allocation alignment. + """ + keys = [hashlib.sha256(f"lru-{i}".encode()).hexdigest() for i in range(4)] + for aligned in (False, True): + with self.subTest(aligned=aligned), ExitStack() as resources: + pools = { + PoolName.KV: HostPages((8192, 4096)), + PoolName.INDEXER: HostPages((2048,)), + } + control_port = free_port() + nixl_port = free_port() + while nixl_port == control_port: + nixl_port = free_port() + config = HiCacheStorageConfig( + tp_rank=0, + tp_size=1, + pp_rank=0, + pp_size=1, + attn_cp_rank=0, + attn_cp_size=1, + is_mla_model=True, + enable_storage_metrics=False, + is_page_first_layout=True, + model_name="glm53-native-bounded-lru", + extra_config={ + "control_port": control_port, + "nixl_listen_port": nixl_port, + "local_dram_pool_pages": {"kv": 3, "indexer": 3}, + "align_storage_prefix": aligned, + }, + ) + store = KVCRStore(config, pools[PoolName.KV]) + store.register_mem_host_pool_v2( + pools[PoolName.INDEXER], PoolName.INDEXER + ) + resources.callback(store.close) + for pool in pools.values(): + for part, buffer in enumerate(pool.buffers): + for slot in range(4): + buffer[slot].copy_( + ( + (torch.arange(buffer.shape[1]) * 37 + slot + part) + % 251 + ).to(torch.uint8) + ) + + def transfers(page_keys, indices, names=tuple(pools)): + return [ + PoolTransfer(name, keys=page_keys, host_indices=indices) + for name in names + ] + + for page_keys, indices, info in ( + (keys[:2], torch.arange(4), None), + (keys[2:3], torch.arange(4, 6), None), + ( + keys[3:], + torch.arange(6, 8), + HiCacheStorageExtraInfo(prefix_keys=keys[:2]), + ), + ): + result = store.batch_set_v2(transfers(page_keys, indices), info) + self.assertTrue( + all(all(values) for values in result.values()), result + ) + chain = keys[:2] + keys[3:] + self.assertEqual( + store.batch_exists_v2( + chain, [PoolTransfer(PoolName.INDEXER)] + ).kv_hit_pages, + 3 if aligned else 0, + ) + if aligned: + result = store.batch_get_v2(transfers(chain, torch.arange(8, 14))) + self.assertTrue( + all(all(values) for values in result.values()), result + ) + self.assertEqual( + store.batch_exists_v2( + keys[2:3], [PoolTransfer(PoolName.INDEXER)] + ).kv_hit_pages, + 0, + ) + for pool in pools.values(): + for buffer in pool.buffers: + self.assertTrue(torch.equal(buffer[[0, 1, 3]], buffer[4:7])) + + def test_remote_restore_all_components_and_missing_pool(self): + """Remote success must populate every byte; missing indexer stays invalid.""" + self._round_trip(((8192, 4096), (2048,)), page_size=2) + + def test_glm53_page_sized_remote_restore(self): + """Exercise the observed GLM5.3 MLA/indexer physical page byte sizes. + + This is a storage transport gate, not an inference or GPU-copy test. + Nonuniform payloads expose within-page truncation and offset errors. + """ + self._round_trip(((2875392,), (658944,)), page_size=64) + + def test_native_diagnostics_record_remote_bytes_and_inventory(self): + """The flag alone is insufficient: native stats require a wired factory. + + Validate a real peer restore, then inspect the native interval handed + back during close. Logical adapter bandwidth cannot satisfy this gate. + """ + from kvcr.api import DURATION_METRIC, TRANSFER_BYTES_METRIC + + with self.assertLogs( + "sglang.srt.mem_cache.storage.kvcr.kvcr_store", level="INFO" + ) as captured: + self._round_trip(((8192, 4096), (2048,)), page_size=2, diagnostic=True) + snapshots = [ + json.loads(line.split("KVCR L3 native stats ", 1)[1])["metrics"] + for line in captured.output + if "KVCR L3 native stats " in line + ] + source_bytes = json.dumps( + ("counter", TRANSFER_BYTES_METRIC, ("source_write",), "value"), + separators=(",", ":"), + ) + remote_duration = json.dumps( + ("histogram", DURATION_METRIC, ("remote_deliver", "success"), "count"), + separators=(",", ":"), + ) + self.assertGreater( + sum(snapshot.get(source_bytes, 0) for snapshot in snapshots), 0 + ) + self.assertGreater( + sum(snapshot.get(remote_duration, 0) for snapshot in snapshots), 0 + ) + inventory = [ + json.loads(line.split("KVCR L3 inventory ", 1)[1]) + for line in captured.output + if "KVCR L3 inventory " in line + ] + self.assertTrue( + any(not event["removed"] and event["keys"] for event in inventory) + ) + + def _round_trip(self, sizes, page_size, diagnostic=False): + keys = [hashlib.sha256(f"page-{i}".encode()).hexdigest() for i in range(4)] + port = free_port() + target_port = free_port() + while target_port == port: + target_port = free_port() + source_pools = { + PoolName.KV: HostPages(sizes[0], page_size), + PoolName.INDEXER: HostPages(sizes[1], page_size), + } + target_pools = { + PoolName.KV: HostPages(sizes[0], page_size), + PoolName.INDEXER: HostPages(sizes[1], page_size), + } + with ExitStack() as resources: + stores = [] + for control_port, pools in ( + (port, source_pools), + (target_port, target_pools), + ): + config = HiCacheStorageConfig( + tp_rank=0, + tp_size=1, + pp_rank=0, + pp_size=1, + attn_cp_rank=0, + attn_cp_size=1, + is_mla_model=True, + enable_storage_metrics=False, + is_page_first_layout=True, + model_name="glm53-native-byte-gate", + extra_config={ + "control_port": control_port, + "control_advertise_host": "127.0.0.1", + "local_dram_pool_pages": {"kv": 16, "indexer": 16}, + "diagnostic_telemetry": diagnostic, + }, + ) + store = KVCRStore(config, pools[PoolName.KV]) + store.register_mem_host_pool_v2( + pools[PoolName.INDEXER], PoolName.INDEXER + ) + resources.callback(store.close) + stores.append(store) + source, target = stores + for pool_index, pool in enumerate(source_pools.values()): + for part_index, buffer in enumerate(pool.buffers): + for slot in range(4): + pattern = torch.arange(buffer.shape[1], dtype=torch.int64) + pattern = ( + pattern * 37 + 11 + 31 * pool_index + 7 * part_index + slot + ) % 251 + buffer[slot].copy_(pattern.to(torch.uint8)) + written = source.batch_set_v2( + [ + PoolTransfer( + name, keys=keys, host_indices=torch.arange(4 * page_size) + ) + for name in source_pools + ] + ) + self.assertTrue(all(all(results) for results in written.values())) + transfers = [ + PoolTransfer( + name, + keys=keys, + host_indices=torch.arange(8 * page_size, 12 * page_size), + ) + for name in target_pools + ] + # An empty target cannot silently manufacture a local hit. + self.assertEqual( + target.batch_exists_v2( + keys, [PoolTransfer(PoolName.INDEXER)] + ).kv_hit_pages, + 0, + ) + hint = HiCacheStorageExtraInfo( + extra_info={ + "kv_hints": { + "protocol_version": "0.1", + "message_id": "native-restore", + "actions": [ + { + "action_id": "fetch", + "action_type": "kv.fetch", + "action_version": "1.0", + "payload": { + "source_control_endpoint": f"tcp://127.0.0.1:{port}", + "block_hashes": [int(key[:16], 16) for key in keys], + }, + } + ], + } + } + ) + restored = target.batch_get_v2(transfers, hint) + self.assertTrue( + all(all(results) for results in restored.values()), restored + ) + for name, source_pool in source_pools.items(): + for source_buffer, target_buffer in zip( + source_pool.buffers, target_pools[name].buffers, strict=True + ): + self.assertTrue(torch.equal(source_buffer[:4], target_buffer[8:12])) + missing_key = hashlib.sha256(b"missing-indexer").hexdigest() + source.batch_set_v1( + [missing_key], torch.arange(4 * page_size, 5 * page_size) + ) + hint.extra_info["kv_hints"]["actions"][0]["payload"]["block_hashes"] = [ + int(missing_key[:16], 16) + ] + missing = target.batch_get_v2( + [ + PoolTransfer( + name, + keys=[missing_key], + host_indices=torch.arange(12 * page_size, 13 * page_size), + ) + for name in target_pools + ], + hint, + ) + self.assertEqual(missing["indexer"], [False]) + # KVCR may refuse the whole remote operation when a required key + # is absent; no caller may infer a valid combined page from it. + self.assertFalse(all(all(results) for results in missing.values())) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/disaggregation/test_decode_hicache_tree_core.py b/test/registered/unit/disaggregation/test_decode_hicache_tree_core.py index 1a55433b62e0..c715bd3591cb 100644 --- a/test/registered/unit/disaggregation/test_decode_hicache_tree_core.py +++ b/test/registered/unit/disaggregation/test_decode_hicache_tree_core.py @@ -23,7 +23,11 @@ def test_storage_probe_and_prefetch_use_node_handles(self): ongoing_prefetch = {} def register_prefetch(req_id, *_args, **_kwargs): - ongoing_prefetch[req_id] = object() + # Only the router-selected source can satisfy this request. + if _kwargs.get("kv_hints") == hint: + ongoing_prefetch[req_id] = object() + + hint = {"protocol_version": "0.1", "message_id": "decode-peer"} tree_cache = SimpleNamespace( hicache_storage_pass_prefix_keys=True, @@ -33,7 +37,11 @@ def register_prefetch(req_id, *_args, **_kwargs): is_root=Mock(return_value=False), get_last_hash_value=Mock(return_value="h2"), get_prefix_hash_values=Mock(return_value=["h0", "h1"]), - query_storage_hit_length=Mock(return_value=2), + query_storage_hit_length=Mock( + side_effect=lambda *_args, **kwargs: ( + 2 if kwargs.get("kv_hints") == hint else 0 + ) + ), prefetch_from_storage=Mock(side_effect=register_prefetch), prefix_device_indices=Mock(return_value=torch.tensor([10, 11])), ) @@ -47,6 +55,7 @@ def register_prefetch(req_id, *_args, **_kwargs): origin_input_ids=[0, 1, 2, 3, 4, 5, 6, 7], extra_key="model", cache_salt="tenant-a", + kv_hints=hint, ) result = SimpleNamespace( device_prefix_len=2, @@ -67,6 +76,7 @@ def register_prefetch(req_id, *_args, **_kwargs): ["h0", "h1"], extra_key="model", cache_salt="tenant-a", + kv_hints=hint, ) DecodeHiCachePreallocMixin._start_hicache_prefetch(harness, req, prefix_match) @@ -80,6 +90,7 @@ def register_prefetch(req_id, *_args, **_kwargs): ["h0", "h1"], extra_key="model", cache_salt="tenant-a", + kv_hints=hint, ) def test_stale_prefetch_anchor_degrades_to_l2(self): diff --git a/test/registered/unit/managers/scheduler_components/test_thread_scheduling.py b/test/registered/unit/managers/scheduler_components/test_thread_scheduling.py new file mode 100644 index 000000000000..a57c742a78d2 --- /dev/null +++ b/test/registered/unit/managers/scheduler_components/test_thread_scheduling.py @@ -0,0 +1,47 @@ +"""Explicit scheduler scheduling leaves defaults and CPU placement alone.""" + +import sys +import unittest +from unittest.mock import patch + +from sglang.srt.managers.scheduler_components.thread_scheduling import ( + configure_python_thread_scheduling, + validate_python_thread_switch_interval, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +class TestPythonThreadScheduling(CustomTestCase): + def test_default_does_not_modify_the_interpreter(self): + with patch("sys.setswitchinterval") as setter: + configure_python_thread_scheduling(None) + setter.assert_not_called() + + def test_explicit_setting_applies_and_records_actual_process_value(self): + before = sys.getswitchinterval() + self.addCleanup(sys.setswitchinterval, before) + with self.assertLogs( + "sglang.srt.managers.scheduler_components.thread_scheduling", level="INFO" + ) as logs: + configure_python_thread_scheduling(0.001) + self.assertAlmostEqual(sys.getswitchinterval(), 0.001, places=8) + self.assertIn("requested_seconds=0.001 actual_seconds=0.001", logs.output[0]) + + def test_invalid_values_rejected_before_mutation(self): + for value in (0, -1, float("nan"), float("inf"), -float("inf"), True, "0.001"): + with self.subTest(value=value), patch("sys.setswitchinterval") as setter: + with self.assertRaisesRegex(ValueError, "positive and finite"): + configure_python_thread_scheduling(value) + setter.assert_not_called() + + def test_finite_positive_intervals_validate(self): + for value in (None, 0.0001, 0.001, 0.005, 1): + with self.subTest(value=value): + validate_python_thread_switch_interval(value) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py b/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py index e25988c02e82..dcb2313393d0 100644 --- a/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py +++ b/test/registered/unit/mem_cache/test_buffer_mode_sidecar.py @@ -44,6 +44,7 @@ ) from sglang.srt.mem_cache.utils import get_storage_hash_str from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=8, suite="base-a-test-cpu") @@ -548,5 +549,64 @@ def test_completed_prefetch_keeps_dsv4_full_and_swa_sidecars_for_h2d(self): self.assertEqual(cache.prefetch_loaded_storage_start_by_reqid[req_id], 0) +class TestResumeBoundarySync(CustomTestCase): + def test_disjoint_checkpoint_maxima_choose_common_boundary(self): + """min(4, 3)=3 is invalid on rank zero; the common checkpoint is 1.""" + for local_boundaries, peer_boundaries in (([1, 4], [1, 3]), ([1, 3], [1, 4])): + with self.subTest(local=local_boundaries): + operation = PrefetchOperation( + CacheRequestHandle("holes", 0), list(range(8)) + ) + operation.all_hash_values = [f"page-{i}" for i in range(4)] + operation.query_pool_hit_pages = {PoolName.KV: 4} + operation.pool_storage_result.restorable_prefix_pages = local_boundaries + calls = [] + + def reduce_with_peer(values, reduce_op, groups): + calls.append(values.numel()) + if values.numel() == 4: + peer = torch.tensor( + [int(i in peer_boundaries) for i in range(1, 5)] + ) + values.copy_(torch.minimum(values, peer)) + else: + values[0] = min(values[0], max(peer_boundaries) * 2) + + controller = SimpleNamespace( + page_size=2, + prefetch_hits_sync_groups=[], + _all_reduce=reduce_with_peer, + ) + hit = HybridCacheController._sync_prefetch_hit_query( + controller, operation, max(local_boundaries) * 2 + ) + self.assertEqual(hit, 2) + self.assertEqual( + operation.pool_storage_result.restorable_prefix_pages, [1] + ) + self.assertEqual(calls, [len(PoolName) + 3, 4]) + + def test_mixed_legacy_reply_never_enters_a_rank_only_collective(self): + """An explicit/absent-set mix must reject the hit, not hang one rank.""" + operation = PrefetchOperation(CacheRequestHandle("mixed", 0), list(range(8))) + operation.all_hash_values = [f"page-{i}" for i in range(4)] + operation.query_pool_hit_pages = {PoolName.KV: 4} + operation.pool_storage_result.restorable_prefix_pages = [1, 4] + calls = [] + + def reduce_with_legacy(values, *_): + calls.append(values.numel()) + values[-2:] = 0 # collective MIN of (1,0) and (0,1) + + controller = SimpleNamespace( + page_size=2, prefetch_hits_sync_groups=[], _all_reduce=reduce_with_legacy + ) + self.assertEqual( + HybridCacheController._sync_prefetch_hit_query(controller, operation, 8), 0 + ) + self.assertEqual(operation.pool_storage_result.restorable_prefix_pages, []) + self.assertEqual(calls, [len(PoolName) + 3]) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_kvcr_storage.py b/test/registered/unit/mem_cache/test_kvcr_storage.py new file mode 100644 index 000000000000..029779824fe7 --- /dev/null +++ b/test/registered/unit/mem_cache/test_kvcr_storage.py @@ -0,0 +1,2603 @@ +"""CPU contracts for indexed HiCache IO; native NIXL is a separate integration gate.""" + +import ctypes +import json +import mmap +import os +import threading +import types +import unittest +from contextlib import ExitStack +from enum import Enum +from queue import Queue +from types import SimpleNamespace +from unittest.mock import patch + +import msgspec +import torch +from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation +from sglang.srt.mem_cache.hicache_storage import ( + HiCacheStorageConfig, + HiCacheStorageExtraInfo, + PoolHitPolicy, + PoolName, + PoolTransfer, +) +from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + HybridCacheController, + StorageOperation, +) +from sglang.srt.mem_cache.storage.kvcr.delivery import ( + Delivery, + block_chunks, + requested_results, + run_window, + validate_destination_spans, +) +from sglang.srt.mem_cache.storage.kvcr.diagnostics import DiagnosticStats +from sglang.srt.mem_cache.storage.kvcr.hints import peer_hint +from sglang.srt.mem_cache.storage.kvcr.kvcr_store import KVCRStore, _allocate_dram_pool +from sglang.srt.mem_cache.storage.kvcr.layout import ( + block_key, + describe_pool, + page_slots, + resume_boundaries, + storage_namespace, +) +from sglang.srt.mem_cache.storage.kvcr.numa_pool import ( + HOST_PROOF_ATTR, + audit_registered_host, + load_config, + node_option, + prove_range, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class _HostPool: + def __init__(self, sizes=(8, 3), page_size=2, pages=8): + self.page_size, self.size = page_size, pages * page_size + self.layout, self.dtype = "page_first", torch.uint8 + self.buffers = [torch.zeros((pages, size), dtype=self.dtype) for size in sizes] + self.kv_buffer = self.buffers[0] + + def get_page_buffer_meta(self, indices): + pointers, sizes = [], [] + for start in indices.tolist()[:: self.page_size]: + for buffer in self.buffers: + slot = start // self.page_size + pointers.append(buffer[slot].data_ptr()) + sizes.append(buffer.shape[1]) + return pointers, sizes + + +class _MemoryRef(msgspec.Struct, frozen=True, kw_only=True): + end_point_name: str + label: str + element_index: int + + +class _QueryStatus(Enum): + HIT = "hit" + FETCHABLE = "fetchable" + MISS = "miss" + + +class _CopyTransport: + """External API substitute that copies the addresses selected by real adapter logic.""" + + def __init__(self, layouts): + self.components = { + part.name: part for p in layouts.values() for part in p.components + } + self.storage, self.completed, self.hints = {}, {}, {} + self.next_handle = 0 + + def deposit(self, blocks): + for key, refs in blocks.items(): + self.storage[key] = { + ref.label: ctypes.string_at( + self.components[ref.label].address + + ref.element_index * self.components[ref.label].stride, + self.components[ref.label].size, + ) + for ref in refs + } + return self._complete({key: SimpleNamespace(success=True) for key in blocks}) + + def deliver(self, blocks, request_id=None): + results = {} + for key, refs in blocks.items(): + data = self.storage.get(key) + ok = data is not None and all(ref.label in data for ref in refs) + if ok: + for ref in refs: + part = self.components[ref.label] + ctypes.memmove( + part.address + ref.element_index * part.stride, + data[ref.label], + part.size, + ) + results[key] = SimpleNamespace(success=ok) + return self._complete(results) + + def _complete(self, result): + self.next_handle += 1 + self.completed[self.next_handle] = result + return self.next_handle + + def poll_completed(self): + result = list(self.completed.items()) + self.completed.clear() + return result + + def query(self, keys, request_id=None): + return [ + (_QueryStatus.HIT if k in self.storage else _QueryStatus.MISS, None) + for k in keys + ] + + def submit_hint(self, hint, request_id=None): + self.hints[request_id] = hint + + def discard_hint(self, request_id): + del self.hints[request_id] + + +class _DeferredCopyTransport(_CopyTransport): + """Queue public operations; perform real byte writes only at completion poll.""" + + def __init__(self, layouts, peers=None): + super().__init__(layouts) + self.pending, self.reads, self.writes, self.events = {}, [], [], [] + self.peers, self.max_pending = peers, 0 + self.result_hook = None + self._copying = False + + def deliver(self, blocks, request_id=None): + self.next_handle += 1 + handle = self.next_handle + hint = self.hints.get(request_id) + source = ( + hint["actions"][0]["payload"]["source_control_endpoint"] if hint else None + ) + self.pending[handle] = (dict(blocks), request_id, source) + self.reads.append((handle, dict(blocks), request_id, source)) + self.events.append(("submit", handle)) + self.max_pending = max(self.max_pending, len(self.pending)) + return handle + + def _complete(self, result): + return result if self._copying else super()._complete(result) + + def poll_completed(self): + completed = super().poll_completed() + # Reverse order deliberately exercises the adapter's shared mailbox. + for handle in reversed(list(self.pending)): + blocks, request_id, source = self.pending[handle] + if request_id is not None and request_id not in self.hints: + raise AssertionError("Hint discarded before native writes completed") + prior = self.storage + try: + if self.peers is not None: + self.storage = self.peers.get(source, {}) + self._copying = True + result = super().deliver(blocks, request_id) + finally: + self._copying = False + self.storage = prior + for key, refs in blocks.items(): + if result[key].success: + self.writes.append((source, key, tuple(refs))) + if self.result_hook is not None: + result = self.result_hook(handle, blocks, result) + completed.append((handle, result)) + del self.pending[handle] + self.events.append(("complete", handle)) + return completed + + def discard_hint(self, request_id): + if any(item[1] == request_id for item in self.pending.values()): + raise AssertionError("Discarding a hint with pending writes") + self.events.append(("discard", request_id)) + super().discard_hint(request_id) + + +class _BoundedCopyTransport(_CopyTransport): + """External bounded LRU with the native API's tail-first alignment contract.""" + + def __init__(self, layouts, capacity): + super().__init__(layouts) + self.capacity, self.recency, self.clock = capacity, {}, 0 + + def deposit(self, blocks): + self.clock += 1 + for key in blocks: + if key in self.storage: + continue + pool = key.rsplit(b"#", 1)[1] + occupants = [ + item for item in self.storage if item.rsplit(b"#", 1)[1] == pool + ] + if len(occupants) >= self.capacity: + victim = min(occupants, key=self.recency.__getitem__) + del self.storage[victim] + del self.recency[victim] + self.recency[key] = (self.clock, 0) + return super().deposit(blocks) + + def align_sequence(self, ordered_keys, use_current_time=False): + self.clock += 1 + for position, key in enumerate(dict.fromkeys(ordered_keys)): + if key in self.storage: + self.recency[key] = (self.clock, -position) + + +class _PeerCopyTransport(_CopyTransport): + """Route actual CPU byte copies by the public request-scoped source hint.""" + + def __init__(self, layouts, peers): + super().__init__(layouts) + self.peers, self.reads = peers, [] + + def deliver(self, blocks, request_id=None): + source = self.hints[request_id]["actions"][0]["payload"][ + "source_control_endpoint" + ] + self.reads.append((source, tuple(blocks))) + prior = self.storage + try: + self.storage = self.peers.get(source, {}) + return super().deliver(blocks, request_id) + finally: + self.storage = prior + + +class TestKVCRStorage(CustomTestCase): + def setUp(self): + self.config = HiCacheStorageConfig( + tp_rank=0, + tp_size=2, + pp_rank=0, + pp_size=1, + attn_cp_rank=0, + attn_cp_size=1, + is_mla_model=False, + enable_storage_metrics=False, + is_page_first_layout=True, + model_name="glm-test", + ) + self.pool = _HostPool() + self.backend = KVCRStore(self.config, self.pool) + self.backend._namespace = storage_namespace(self.config, self.backend._layouts) + self.transport = self.backend._runner = _CopyTransport(self.backend._layouts) + self.keys = [f"{i + 1:064x}" for i in range(4)] + types_module = types.ModuleType("kvcr.types") + types_module.MemoryRef, types_module.QueryStatus = _MemoryRef, _QueryStatus + self.stub = patch.dict( + "sys.modules", + {"kvcr": types.ModuleType("kvcr"), "kvcr.types": types_module}, + ) + self.stub.start() + self.addCleanup(self.stub.stop) + + @unittest.skipUnless(hasattr(os, "fork"), "Private-map isolation requires fork") + def test_dram_pool_is_private_and_owns_writable_bytes(self): + """Reverting to mmap's shared default leaks a child's writes to the owner.""" + with _allocate_dram_pool(8192) as allocation: + allocation[0] = 17 + child = os.fork() + if child == 0: + allocation[0] = 29 + os._exit(0) + _, status = os.waitpid(child, 0) + self.assertEqual(status, 0) + self.assertEqual(allocation[0], 17) + address = ctypes.addressof(ctypes.c_char.from_buffer(allocation)) + ctypes.memset(address + 4096, 43, 4096) + self.assertEqual(allocation[4096:], bytes([43]) * 4096) + self.assertTrue(allocation.closed) + + def test_dram_pool_remains_usable_when_hugepage_advice_is_rejected(self): + """THP is an allocation hint, never a new storage-readiness requirement.""" + + class AdviceRejected(mmap.mmap): + def madvise(self, *_args): + raise OSError("THP is not supported") + + with ( + patch( + "sglang.srt.mem_cache.storage.kvcr.kvcr_store.mmap.mmap", + AdviceRejected, + ), + _allocate_dram_pool(8192) as allocation, + ): + allocation[:4] = b"kvcr" + address = ctypes.addressof(ctypes.c_char.from_buffer(allocation)) + self.assertEqual(ctypes.string_at(address, 4), b"kvcr") + self.assertTrue(allocation.closed) + + def test_whole_page_components_round_trip_to_different_slots(self): + """A restore must move every unequal-sized component, not only primary KV.""" + for part, buffer in enumerate(self.pool.buffers): + buffer[0].fill_(17 + part) + buffer[1].fill_(29 + part) + self.assertEqual( + self.backend.batch_set_v1(self.keys[:2], torch.arange(4)), [True, True] + ) + self.assertEqual( + self.backend.batch_get_v1(self.keys[:2], torch.arange(8, 12)), [True, True] + ) + for buffer in self.pool.buffers: + self.assertTrue(torch.equal(buffer[:2], buffer[4:6])) + # Missing one component must not produce a valid page. + identity = block_key(self.keys[0], self.backend._namespace, "kv") + self.transport.storage[identity].pop("kv_1") + self.assertEqual( + self.backend.batch_get_v1(self.keys[:1], torch.arange(12, 14)), [False] + ) + + def test_invalid_startup_budget_is_rejected_before_registration(self): + """An invalid optional deadline must not reach native registration.""" + for budget in (0, -1, True, 1.5, "30000", None): + with self.subTest(budget=budget): + self.config.extra_config = {"startup_timeout_seconds": budget} + with self.assertRaisesRegex(ValueError, "startup_timeout_seconds"): + KVCRStore(self.config, self.pool) + + def test_invalid_descriptor_coalescing_is_rejected_before_registration(self): + """A truthy string/integer must not silently enable a copy-order change.""" + for value in (0, 1, None, "true", []): + with self.subTest(value=value): + self.config.extra_config = {"coalesce_remote_dram_descriptors": value} + with self.assertRaisesRegex( + ValueError, "coalesce_remote_dram_descriptors" + ): + KVCRStore(self.config, self.pool) + + def test_invalid_ucx_split_is_rejected_before_registration(self): + """Only native descriptor counts, not booleans or coercible values, bind.""" + for value in (0, -1, True, False, 1.5, "32", None, 2**32): + with self.subTest(value=value): + self.config.extra_config = {"nixl_ucx_split_batch_size": value} + with self.assertRaisesRegex(ValueError, "nixl_ucx_split_batch_size"): + KVCRStore(self.config, self.pool) + + def test_logical_anchor_queries_intersect_physical_all_page_pools(self): + """A DSv4 logical anchor proves nothing without its physical sidecars.""" + backend = KVCRStore(self.config, SimpleNamespace(kv_buffer=None)) + for name in ( + PoolName.DEEPSEEK_V4_C4, + PoolName.DEEPSEEK_V4_C4_INDEXER, + PoolName.SWA, + PoolName.DEEPSEEK_V4_C4_STATE, + ): + backend.register_mem_host_pool_v2(_HostPool(), name) + backend._namespace = storage_namespace(self.config, backend._layouts) + transport = backend._runner = _CopyTransport(backend._layouts) + self.assertEqual(backend.batch_exists(self.keys), 0) + + for name, count in ( + (PoolName.DEEPSEEK_V4_C4, 4), + (PoolName.DEEPSEEK_V4_C4_INDEXER, 2), + ): + for key in self.keys[:count]: + transport.storage[block_key(key, backend._namespace, str(name))] = {} + # Missing SWA/state pages do not constrain the bare all-pages probe. + result = backend.batch_exists_v2(self.keys) + self.assertEqual(result.kv_hit_pages, 2) + self.assertEqual(result.restorable_prefix_pages, [1, 2]) + + def test_bare_logical_anchor_without_all_page_evidence_is_a_miss(self): + """Empty or trailing-only registrations cannot prove a full prefix.""" + for pools in ((), (PoolName.SWA, PoolName.MAMBA)): + with self.subTest(pools=pools): + backend = KVCRStore(self.config, SimpleNamespace(kv_buffer=None)) + for name in pools: + backend.register_mem_host_pool_v2(_HostPool(), name) + backend._namespace = storage_namespace(self.config, backend._layouts) + backend._runner = _CopyTransport(backend._layouts) + self.assertEqual(backend.batch_exists(self.keys), 0) + self.assertEqual(backend.batch_exists([]), 0) + + def test_constructor_failure_is_fatal_and_retains_unsafe_mappings(self): + """A startup error cannot leave wait_complete behind a dead IO thread. + + A quiescent failure may unwind mappings, but a nonquiescent failure + must retain them until process exit. The unset optional deadline must + also work with older KVCR configuration constructors. + """ + + class StartupError(RuntimeError): + pass + + original_mmap = mmap.mmap + for error_type in (RuntimeError, StartupError): + for budget, coalesce, split in ( + (None, None, None), + (180, None, None), + (None, False, None), + (None, True, None), + (180, True, None), + (None, None, 32), + (180, True, 32), + (None, True, 1), + ): + with self.subTest( + error_type=error_type, budget=budget, coalesce=coalesce, split=split + ): + self.config.extra_config = {"local_dram_bytes": 128} + if budget is not None: + self.config.extra_config["startup_timeout_seconds"] = budget + if coalesce is not None: + self.config.extra_config["coalesce_remote_dram_descriptors"] = ( + coalesce + ) + if split is not None: + self.config.extra_config["nixl_ucx_split_batch_size"] = split + backend = KVCRStore(self.config, self.pool) + captured, mappings = [], [] + + def construct(captured=captured, error_type=error_type, **kwargs): + captured.append(vars(kwargs["config"])) + raise error_type("registration failed") + + def allocate(*args, mappings=mappings, **kwargs): + allocation = original_mmap(*args, **kwargs) + mappings.append(allocation) + return allocation + + modules = { + "kvcr": SimpleNamespace( + KVCR=construct, KVCRBindings=SimpleNamespace + ), + "kvcr.config": SimpleNamespace( + KVCRBackendConfigs=SimpleNamespace, + KVCRConfig=SimpleNamespace, + LocalDramOptions=SimpleNamespace, + ), + "kvcr.control_channels": SimpleNamespace( + ZmqPeerControlChannel=lambda **kwargs: SimpleNamespace( + endpoint="tcp://127.0.0.1:32100", close=lambda: None + ) + ), + "kvcr.types": SimpleNamespace( + KVCRStartupError=StartupError, + RegionDescriptor=SimpleNamespace, + ), + } + try: + with ( + patch.dict("sys.modules", modules), + patch( + "sglang.srt.mem_cache.storage.kvcr.kvcr_store.mmap.mmap", + side_effect=allocate, + ), + patch( + "sglang.srt.mem_cache.storage.kvcr.kvcr_store.os._exit", + side_effect=SystemExit(1), + ), + self.assertLogs( + "sglang.srt.mem_cache.storage.kvcr.kvcr_store", + level="CRITICAL", + ), + self.assertRaises(SystemExit) as fatal, + ): + backend._start() + self.assertEqual(fatal.exception.code, 1) + self.assertIsNone(backend._runner) + self.assertIsNone(backend._poll_thread) + self.assertTrue(mappings) + self.assertEqual( + [allocation.closed for allocation in mappings], + [error_type is RuntimeError] * len(mappings), + ) + if budget is None: + self.assertNotIn("startup_timeout_seconds", captured[0]) + else: + self.assertEqual( + captured[0]["startup_timeout_seconds"], budget + ) + if coalesce is True: + self.assertIs( + captured[0]["coalesce_remote_dram_descriptors"], True + ) + else: + self.assertNotIn( + "coalesce_remote_dram_descriptors", captured[0] + ) + if split is None: + self.assertNotIn("nixl_ucx_split_batch_size", captured[0]) + else: + self.assertEqual( + captured[0]["nixl_ucx_split_batch_size"], split + ) + finally: + for allocation in mappings: + allocation.close() + + def test_byte_verification_covers_every_successful_page_component(self): + """Diagnostic digests follow physical bytes, not slot IDs or only primary KV.""" + self.backend._verify_transfer_bytes = True + self.pool.buffers[0][0].fill_(17) + self.pool.buffers[1][0].fill_(29) + with self.assertLogs( + "sglang.srt.mem_cache.storage.kvcr.kvcr_store", level="INFO" + ) as captured: + self.backend.batch_set_v1(self.keys[:1], torch.arange(2)) + self.backend.batch_get_v1(self.keys[:1], torch.arange(8, 10)) + # Changing only the unequal-sized side component changes the digest. + self.pool.buffers[1][4, 1] ^= 1 + self.backend.batch_set_v1(self.keys[:1], torch.arange(8, 10)) + missing = self.backend.batch_get_v1(self.keys[1:2], torch.arange(12, 14)) + records = [ + json.loads(line.split("KVCR L3 byte verification ", 1)[1]) + for line in captured.output + ] + self.assertEqual(records[0]["digests"], records[1]["digests"]) + self.assertNotEqual(records[1]["digests"], records[2]["digests"]) + self.assertEqual(records[3]["digests"], {}) + self.assertEqual(missing, [False]) + + def test_metrics_count_complete_logical_pages_and_reset_intervals(self): + """A failed sidecar excludes its page, not successfully moved KV bytes.""" + self.config.enable_storage_metrics = True + indexer = _HostPool(sizes=(5,)) + backend = KVCRStore(self.config, self.pool) + backend.register_mem_host_pool_v2(indexer, PoolName.INDEXER) + backend._namespace = storage_namespace(self.config, backend._layouts) + transport = backend._runner = _CopyTransport(backend._layouts) + transfers = [ + PoolTransfer(name, keys=self.keys[:2], host_indices=torch.arange(4)) + for name in (PoolName.KV, PoolName.INDEXER) + ] + with patch( + "sglang.srt.mem_cache.storage.kvcr.kvcr_store.time.perf_counter", + side_effect=[10.0, 12.0, 20.0, 24.0], + ): + backend.batch_set_v2(transfers) + # Two primary pages exist, but only one has its required indexer. + transport.storage.pop( + block_key(self.keys[1], backend._namespace, "indexer") + ) + restored = backend.batch_get_v2( + [ + PoolTransfer( + name, keys=self.keys[:3], host_indices=torch.arange(8, 14) + ) + for name in (PoolName.KV, PoolName.INDEXER) + ] + ) + self.assertEqual(restored["kv"], [True, True, False]) + self.assertEqual(restored["indexer"], [True, False, False]) + snapshot = backend.get_stats() + self.assertEqual(snapshot.backup_pgs, [2]) + self.assertEqual(snapshot.prefetch_pgs, [1]) + self.assertAlmostEqual(snapshot.backup_bandwidth[0], 32 / 2 / 1e9) + self.assertAlmostEqual(snapshot.prefetch_bandwidth[0], 27 / 4 / 1e9) + self.assertEqual(backend.get_stats().prefetch_pgs, []) + self.assertEqual(backend.get_stats().backup_pgs, []) + + def test_prefix_alignment_preserves_parents_before_tail_allocation(self): + """Refreshing after deposit is too late: allocation already evicted a parent. + + A hot prefix exists in both KV and indexer while an unrelated backup + fills each three-page L3 pool. Backing up its next page must preserve + every parent in every pool, not only the primary KV. The unaligned + control loses the root; alignment evicts the unrelated page instead. + Native bounded-LRU behavior is also covered by the manual NIXL gate. + """ + for aligned in (False, True): + with self.subTest(aligned=aligned): + config = SimpleNamespace(**vars(self.config)) + config.extra_config = {"align_storage_prefix": aligned} + backend = KVCRStore(config, self.pool) + indexer = _HostPool(sizes=(5,)) + backend.register_mem_host_pool_v2(indexer, PoolName.INDEXER) + backend._namespace = storage_namespace(config, backend._layouts) + backend._runner = _BoundedCopyTransport(backend._layouts, 3) + for pool in (self.pool, indexer): + for part, buffer in enumerate(pool.buffers): + for slot in range(4): + buffer[slot].fill_(11 + slot + part) + + def transfers(keys, indices): + return [ + PoolTransfer(name, keys=keys, host_indices=indices) + for name in (PoolName.KV, PoolName.INDEXER) + ] + + backend.batch_set_v2(transfers(self.keys[:2], torch.arange(4))) + # Unrelated newest page competes with the reused prefix. + backend.batch_set_v2(transfers(self.keys[2:3], torch.arange(4, 6))) + backend.batch_set_v2( + transfers(self.keys[3:], torch.arange(6, 8)), + HiCacheStorageExtraInfo(prefix_keys=self.keys[:2]), + ) + chain = self.keys[:2] + self.keys[3:] + hit = backend.batch_exists_v2(chain, [PoolTransfer(PoolName.INDEXER)]) + self.assertEqual(hit.kv_hit_pages, 3 if aligned else 0) + if aligned: + restored = backend.batch_get_v2( + transfers(chain, torch.arange(8, 14)) + ) + self.assertTrue(all(all(values) for values in restored.values())) + self.assertEqual( + backend.batch_exists_v2( + self.keys[2:3], [PoolTransfer(PoolName.INDEXER)] + ).kv_hit_pages, + 0, + ) + for pool in (self.pool, indexer): + for buffer in pool.buffers: + self.assertTrue(torch.equal(buffer[[0, 1, 3]], buffer[4:7])) + + def test_alignment_config_rejects_truthy_non_booleans(self): + for value in ("false", 1, None): + with ( + self.subTest(value=value), + self.assertRaisesRegex( + ValueError, "align_storage_prefix must be a boolean" + ), + ): + config = SimpleNamespace(**vars(self.config)) + config.extra_config = {"align_storage_prefix": value} + KVCRStore(config, self.pool) + + def test_diagnostic_intervals_preserve_labels_and_bound_sample_memory(self): + """Unpolled telemetry must not grow with operations or merge distinct scopes.""" + stats = DiagnosticStats(max_series=3) + for _ in range(10000): + stats.observe_histogram("duration", 0.001, ("source_write", "success")) + stats.observe_histogram("duration", 0.02, ("source_write", "success")) + stats.increase_counter("bytes", 512, ("source_write",)) + stats.increase_counter("bytes", 128, ("source_write",)) + stats.increase_counter("bytes", 256, ("local_fill",)) + # An unexpected series storm drops observations instead of growing + # without bound or failing an otherwise valid KV transfer. + for index in range(100): + stats.set_gauge("unexpected", index, (str(index),)) + snapshot = stats.reduce() + self.assertEqual(snapshot["dropped_observations"], 100) + result = { + tuple(json.loads(k)[:2]) + (tuple(json.loads(k)[2]), json.loads(k)[3]): v + for k, v in snapshot.items() + if k != "dropped_observations" + } + hist = ("histogram", "duration", ("source_write", "success")) + self.assertEqual(result[(*hist, "count")], 10001) + self.assertAlmostEqual(result[(*hist, "sum")], 10.02) + self.assertEqual(result[(*hist, "bucket_2")], 10000) + self.assertEqual(result[(*hist, "bucket_5")], 1) + self.assertEqual(result[("counter", "bytes", ("source_write",), "value")], 640) + self.assertEqual(result[("counter", "bytes", ("local_fill",), "value")], 256) + self.assertLessEqual(len(result), 9) + self.assertTrue(stats.is_empty()) + self.assertEqual(stats.reduce(), {}) + + def test_diagnostic_io_reports_terminal_miss_without_changing_restoration(self): + """Diagnostics must not turn a failed physical page into a successful hit.""" + self.backend._diagnostic_telemetry = True + self.backend._control_endpoint = "tcp://source:32000" + self.pool.buffers[0][0].fill_(43) + self.pool.buffers[1][0].fill_(57) + with self.assertLogs( + "sglang.srt.mem_cache.storage.kvcr.kvcr_store", level="INFO" + ) as captured: + self.backend.batch_set_v1(self.keys[:1], torch.arange(2)) + restored = self.backend.batch_get_v1(self.keys[:2], torch.arange(8, 12)) + records = [ + json.loads(line.split("KVCR L3 operation ", 1)[1]) + for line in captured.output + ] + self.assertEqual(restored, [True, False]) + for buffer in self.pool.buffers: + self.assertTrue(torch.equal(buffer[0], buffer[4])) + self.assertEqual(records[0]["failed_keys"], []) + self.assertEqual( + records[1]["failed_keys"], + [block_key(self.keys[1], self.backend._namespace, "kv").decode()], + ) + self.assertEqual(records[1]["requested_blocks"], 2) + self.assertEqual(records[1]["requested_spans"], 4) + self.assertGreaterEqual(records[1]["wait_ms"], 0) + for record in records: + self.assertAlmostEqual( + record["prepare_ms"], + sum( + record[field] + for field in ("descriptor_ms", "align_before_ms", "hint_ms") + ), + ) + + def test_diagnostic_config_rejects_truthy_non_booleans(self): + for value in ("false", 1, None): + with ( + self.subTest(value=value), + self.assertRaisesRegex( + ValueError, "diagnostic_telemetry must be a boolean" + ), + ): + config = SimpleNamespace(**vars(self.config)) + config.extra_config = {"diagnostic_telemetry": value} + KVCRStore(config, self.pool) + + def _controller_with_indexer( + self, *, combined=True, pages=4, batch_pages=128, pipeline_window=1 + ): + config = SimpleNamespace(**vars(self.config)) + config.extra_config = { + "combine_page_reads": combined, + "prefetch_batch_pages": batch_pages, + "delivery_pipeline_window": pipeline_window, + } + pools = { + PoolName.KV: _HostPool(pages=pages * 2), + PoolName.INDEXER: _HostPool(sizes=(5,), pages=pages * 2), + } + backend = KVCRStore(config, pools[PoolName.KV]) + backend.register_mem_host_pool_v2(pools[PoolName.INDEXER], PoolName.INDEXER) + backend._namespace = storage_namespace(config, backend._layouts) + transport = backend._runner = _CopyTransport(backend._layouts) + controller = HiCacheController.__new__(HiCacheController) + controller.storage_backend = backend + controller.page_size = 2 + controller.page_get_func = controller._page_get_zero_copy + controller.prefetch_sync_queue = Queue() + operation = PrefetchOperation("combined-page-read", list(range(pages * 2))) + operation.host_indices = torch.arange(pages * 2, pages * 4) + operation.hash_value = ( + self.keys if pages == 4 else [f"{page:064x}" for page in range(pages)] + ) + operation.pool_transfers = [ + PoolTransfer(PoolName.INDEXER, indices_from_pool=PoolName.KV) + ] + return controller, operation, pools, transport + + def test_combined_controller_reads_restore_all_components_in_one_completion(self): + """Coalescing changes submissions, not KV/indexer bytes or valid prefix length.""" + for combined in (False, True): + with self.subTest(combined=combined): + c, operation, pools, transport = self._controller_with_indexer( + combined=combined + ) + for pool in pools.values(): + for part, buffer in enumerate(pool.buffers): + for slot in range(4): + buffer[slot].fill_(13 + slot + part) + c.storage_backend.batch_set_v2( + [ + PoolTransfer(name, keys=self.keys, host_indices=torch.arange(8)) + for name in pools + ] + ) + before = transport.next_handle + hits = c._page_transfer_kv_batch( + operation, + self.keys, + operation.host_indices, + None, + operation.pool_transfers, + ) + self.assertEqual(hits, 4) + self.assertEqual(transport.next_handle - before, 1 if combined else 2) + for pool in pools.values(): + for buffer in pool.buffers: + self.assertTrue(torch.equal(buffer[:4], buffer[4:8])) + + def test_combined_partial_pool_miss_keeps_every_batch_ack(self): + """A sidecar hole clamps the prefix, but later skipped batches still get acks. + + Returning KV's full hit count admits incomplete indexer data. Returning + early from the outer transfer loop strands ranks waiting for later acks. + """ + c, operation, pools, transport = self._controller_with_indexer(batch_pages=2) + for pool in pools.values(): + for buffer in pool.buffers: + buffer[:4].fill_(63) + c.storage_backend.batch_set_v2( + [ + PoolTransfer(name, keys=self.keys, host_indices=torch.arange(8)) + for name in pools + ] + ) + transport.storage.pop( + block_key(self.keys[1], c.storage_backend._namespace, "indexer") + ) + with self.assertLogs(level="WARNING"): + completed = c._page_transfer(operation) + self.assertEqual(completed, 1) + acks = [c.prefetch_sync_queue.get_nowait() for _ in range(2)] + self.assertEqual([ack.completed_tokens for ack in acks], [2, 2]) + self.assertTrue(c.prefetch_sync_queue.empty()) + for pool in pools.values(): + for buffer in pool.buffers: + self.assertTrue(torch.equal(buffer[0], buffer[4])) + self.assertEqual(int(buffer[6].sum()), 0) + + def test_prefetch_batch_size_preserves_bytes_and_ack_prefix_after_sidecar_hole( + self, + ): + """Larger batches reduce completions without admitting incomplete pages. + + A hole inside either the first or middle batch must stop the usable + prefix and leave later batches untouched, while emitting the same ACK + count as a successful peer. Default and opt-in paths restore identical + bytes, including a short final batch. + """ + pages = 1100 + cases = ( + (128, None, 9, [256, 512, 768, 1024, 1280, 1536, 1792, 2048, 2200]), + (512, None, 3, [1024, 2048, 2200]), + (512, 600, 2, [1024, 1200, 1200]), + (1024, None, 2, [2048, 2200]), + (1024, 600, 1, [1200, 1200]), + ) + for batch_pages, missing, submissions, expected_acks in cases: + with self.subTest(batch_pages=batch_pages, missing=missing): + c, op, pools, transport = self._controller_with_indexer( + pages=pages, batch_pages=batch_pages + ) + for pool in pools.values(): + for part, buffer in enumerate(pool.buffers): + values = ( + torch.arange(pages * buffer.shape[1]).reshape( + pages, buffer.shape[1] + ) + + part * 17 + ) % 251 + buffer[:pages].copy_(values.to(torch.uint8)) + c.storage_backend.batch_set_v2( + [ + PoolTransfer( + name, + keys=op.hash_value, + host_indices=torch.arange(pages * 2), + ) + for name in pools + ] + ) + if missing is not None: + transport.storage.pop( + block_key( + op.hash_value[missing], + c.storage_backend._namespace, + "indexer", + ) + ) + before = transport.next_handle + with ( + self.assertLogs(level="WARNING") + if missing is not None + else ExitStack() + ): + restored = c._page_transfer(op) + self.assertEqual(restored, pages if missing is None else missing) + self.assertEqual(transport.next_handle - before, submissions) + acks = [] + while not c.prefetch_sync_queue.empty(): + acks.append(c.prefetch_sync_queue.get_nowait().completed_tokens) + self.assertEqual(acks, expected_acks) + for pool in pools.values(): + for buffer in pool.buffers: + self.assertTrue( + torch.equal( + buffer[:restored], buffer[pages : pages + restored] + ) + ) + if missing is not None: + self.assertEqual(int(buffer[pages + 1024 :].sum()), 0) + if missing is not None: + self.assertEqual( + int(pools[PoolName.INDEXER].buffers[0][pages + missing].sum()), + 0, + ) + + def test_prefetch_batch_config_rejects_values_that_break_collective_sequence(self): + for value in (0, -1, True, False, 1.5, "512", None): + with ( + self.subTest(value=value), + self.assertRaisesRegex( + ValueError, "prefetch_batch_pages must be a positive integer" + ), + ): + config = SimpleNamespace(**vars(self.config)) + config.extra_config = {"prefetch_batch_pages": value} + KVCRStore(config, self.pool) + + def test_query_batch_default_and_optin_do_not_change_payload_batching(self): + self.assertEqual(self.backend.query_batch_pages, 128) + for pages in (1, 128, 2048): + with self.subTest(pages=pages): + config = SimpleNamespace(**vars(self.config)) + config.extra_config = {"query_batch_pages": pages} + store = KVCRStore(config, self.pool) + self.addCleanup(store.close) + self.assertEqual(store.query_batch_pages, pages) + self.assertEqual(store.prefetch_batch_pages, 128) + self.assertEqual(store._delivery_pipeline_window, 1) + self.assertEqual(set(store._layouts), {"kv"}) + + def test_query_batch_config_rejects_non_integer_or_unbounded_values(self): + for value in (0, -1, True, False, 1.5, "2048", None, 2049): + with ( + self.subTest(value=value), + self.assertRaisesRegex( + ValueError, "query_batch_pages must be an integer from 1 to 2048" + ), + ): + config = SimpleNamespace(**vars(self.config)) + config.extra_config = {"query_batch_pages": value} + KVCRStore(config, self.pool) + + def test_query_batch_optin_rejects_MLA_and_logical_anchor(self): + for mla, pool in ((True, self.pool), (False, SimpleNamespace(kv_buffer=None))): + with ( + self.subTest(mla=mla), + self.assertRaisesRegex( + ValueError, + "query_batch_pages opt-in requires a non-MLA KV-only layout", + ), + ): + config = SimpleNamespace(**vars(self.config)) + config.is_mla_model = mla + config.extra_config = {"query_batch_pages": 2048} + KVCRStore(config, pool) + + def test_query_batch_optin_rejects_sidecars_before_mutating_layouts(self): + config = SimpleNamespace(**vars(self.config)) + config.extra_config = {"query_batch_pages": 2048} + store = KVCRStore(config, self.pool) + self.addCleanup(store.close) + for name in (PoolName.INDEXER, PoolName.MAMBA): + with ( + self.subTest(name=name), + self.assertRaisesRegex( + ValueError, + "query_batch_pages opt-in requires a non-MLA KV-only layout", + ), + ): + store.register_mem_host_pool_v2(_HostPool(), name) + self.assertEqual(set(store._layouts), {"kv"}) + self.assertEqual(set(store.registered_pools), {"kv"}) + + def test_query_batch_controller_preserves_exact_prefix_and_hint_scope(self): + pages, page_size = 4097, 2 + hashes = [f"{i + 1:064x}" for i in range(pages)] + for batch in (128, 2048): + for missing in (0, 1, 127, 128, 129, 2047, 2048, 2049, pages - 1, pages): + for prefix in (None, [], ["prior-prefix"]): + with self.subTest(batch=batch, missing=missing, prefix=prefix): + calls = [] + hint = {"request": "not-a-global-cache-key"} + operation = PrefetchOperation( + "query-prefix", + list(range(pages * page_size)), + last_hash="previous-hash", + prefix_keys=prefix, + ) + operation.kv_hints = hint + + def exists(keys, extra): + start = len(calls) * batch + self.assertEqual(keys, hashes[start : start + batch]) + self.assertIs(extra.extra_info["kv_hints"], hint) + expected = ( + None if prefix is None else prefix + hashes[:start] + ) + self.assertEqual(extra.prefix_keys, expected) + calls.append(list(keys)) + return min(len(keys), max(0, missing - start)) + + controller = HiCacheController.__new__(HiCacheController) + controller.page_size = page_size + controller.storage_backend = SimpleNamespace( + query_batch_pages=batch, batch_exists=exists + ) + with patch( + "sglang.srt.managers.cache_controller.get_storage_hash_str", + return_value=hashes, + ) as codec: + result = controller._storage_hit_query(operation) + self.assertEqual( + result, (hashes[:missing], missing * page_size) + ) + self.assertEqual(operation.all_hash_values, hashes) + self.assertEqual(operation.prefix_keys, prefix) + self.assertEqual( + len(calls), + min(missing // batch + 1, (pages + batch - 1) // batch), + ) + codec.assert_called_once_with( + operation.token_ids, "previous-hash", page_size=page_size + ) + + def test_query_batch_controller_empty_chain_has_no_backend_call(self): + controller = HiCacheController.__new__(HiCacheController) + controller.page_size = 2 + + def unexpected(*args): + self.fail("Empty input must not query storage") + + controller.storage_backend = SimpleNamespace( + query_batch_pages=2048, batch_exists=unexpected + ) + operation = PrefetchOperation("empty-query", []) + with patch( + "sglang.srt.managers.cache_controller.get_storage_hash_str", return_value=[] + ): + self.assertEqual(controller._storage_hit_query(operation), ([], 0)) + self.assertEqual(operation.all_hash_values, []) + + def test_combined_result_rejects_missing_or_truncated_sidecar(self): + """Dropping a pool/result entry must not masquerade as an all-pool hit.""" + c, operation, _, _ = self._controller_with_indexer() + malformed = [ + {"kv": [True] * 4}, + {"kv": [True] * 4, "indexer": [True]}, + {"kv": [True] * 4, "indexer": None}, + {"kv": [True] * 4, "indexer": [1] * 4}, + ] + for result in malformed: + with ( + self.subTest(result=result), + patch.object(c.storage_backend, "batch_get_v2", return_value=result), + self.assertLogs(level="ERROR"), + ): + self.assertEqual( + c._page_transfer_kv_batch( + operation, + self.keys, + operation.host_indices, + None, + operation.pool_transfers, + ), + 0, + ) + + def test_combined_read_config_rejects_truthy_non_booleans(self): + for value in ("false", 1, None): + with ( + self.subTest(value=value), + self.assertRaisesRegex( + ValueError, "combine_page_reads must be a boolean" + ), + ): + config = SimpleNamespace(**vars(self.config)) + config.extra_config = {"combine_page_reads": value} + KVCRStore(config, self.pool) + + def _backup_with_indexer(self, *, combined=True, pages=130): + base, read, pools, transport = self._controller_with_indexer(pages=pages) + controller = HybridCacheController.__new__(HybridCacheController) + controller.__dict__.update(base.__dict__) + controller.mem_pool_host = SimpleNamespace( + kv_buffer=pools[PoolName.KV].kv_buffer, layout_lease=ExitStack + ) + controller.backup_skip = False + controller.page_set_func = controller._page_set_zero_copy + controller.storage_backend.supports_combined_page_writes = combined + operation = StorageOperation( + torch.arange(pages * 2), + list(range(pages * 2)), + hash_value=read.hash_value, + prefix_keys=["f" * 64], + pool_transfers=[ + PoolTransfer(PoolName.INDEXER, indices_from_pool=PoolName.KV) + ], + ) + for pool in pools.values(): + for part, buffer in enumerate(pool.buffers): + values = ( + torch.arange(pages * buffer.shape[1]).reshape(pages, -1) + part * 17 + ) % 251 + buffer[:pages].copy_(values.to(torch.uint8)) + return controller, operation, pools, transport + + def test_combined_backup_preserves_bytes_and_short_final_batch(self): + """One completion covers KV plus indexer, with the same restorable bytes.""" + for combined in (False, True): + with self.subTest(combined=combined): + c, operation, pools, transport = self._backup_with_indexer( + combined=combined + ) + c._page_backup(operation) + self.assertEqual(operation.completed_tokens, 260) + self.assertEqual( + operation.pool_storage_result.extra_pool_hit_pages, {"indexer": 130} + ) + self.assertEqual(transport.next_handle, 2 if combined else 3) + restored = c.storage_backend.batch_get_v2( + [ + PoolTransfer( + name, + keys=operation.hash_value, + host_indices=torch.arange(260, 520), + ) + for name in pools + ] + ) + self.assertTrue(all(all(values) for values in restored.values())) + for pool in pools.values(): + for buffer in pool.buffers: + self.assertTrue(torch.equal(buffer[:130], buffer[130:])) + + def test_combined_backup_never_acknowledges_incomplete_sidecar_batch(self): + """Malformed or failed sidecars stop ACK growth and leave later batches idle.""" + malformed = ( + {"kv": [True] * 128}, + {"kv": [True] * 128, "indexer": [True]}, + {"kv": [True] * 128, "indexer": [1] * 128}, + {"kv": [True] * 128, "indexer": [False] + [True] * 127}, + ) + for result in malformed: + with self.subTest(result=result): + c, operation, _, transport = self._backup_with_indexer(pages=270) + real_set = c.storage_backend.batch_set_v2 + batches = [] + + def fail_second( + transfers, + extra_info=None, + batches=batches, + result=result, + real_set=real_set, + ): + batches.append(list(extra_info.prefix_keys)) + if len(batches) == 2: + return result + return real_set(transfers, extra_info) + + with ( + patch.object(c.storage_backend, "batch_set_v2", fail_second), + self.assertLogs(level="WARNING"), + ): + c._page_backup(operation) + self.assertEqual(operation.completed_tokens, 256) + self.assertEqual(len(batches), 2) + self.assertEqual(batches[0], ["f" * 64]) + self.assertEqual(batches[1], ["f" * 64] + operation.hash_value[:128]) + self.assertEqual(transport.next_handle, 1) + self.assertNotIn( + block_key( + operation.hash_value[256], c.storage_backend._namespace, "kv" + ), + transport.storage, + ) + + def test_combined_backup_does_not_reinterpret_trailing_sidecar_indices(self): + """A separately indexed trailing checkpoint must retain its own source slot.""" + c, operation, pools, transport = self._backup_with_indexer() + operation.pool_transfers = [ + PoolTransfer( + PoolName.INDEXER, + keys=operation.hash_value[-1:], + host_indices=torch.arange(258, 260), + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + ] + c._page_backup(operation) + self.assertEqual(operation.completed_tokens, 260) + self.assertEqual(transport.next_handle, 3) + self.assertEqual( + operation.pool_storage_result.extra_pool_hit_pages, {"indexer": 1} + ) + result = c.storage_backend.batch_get_v2( + [ + PoolTransfer( + PoolName.INDEXER, + keys=operation.hash_value[-1:], + host_indices=torch.arange(260, 262), + ) + ] + ) + self.assertEqual(result, {"indexer": [True]}) + for buffer in pools[PoolName.INDEXER].buffers: + self.assertTrue(torch.equal(buffer[129], buffer[130])) + + def test_combined_write_config_rejects_truthy_non_booleans(self): + for value in ("false", 1, None): + with ( + self.subTest(value=value), + self.assertRaisesRegex( + ValueError, "combine_page_writes must be a boolean" + ), + ): + config = SimpleNamespace(**vars(self.config)) + config.extra_config = {"combine_page_writes": value} + KVCRStore(config, self.pool) + + def _shared_reader(self, rank, **changes): + config = SimpleNamespace( + **{ + **vars(self.config), + "tp_rank": rank, + "is_mla_model": True, + "extra_config": {"share_replicated_mla": True}, + **changes, + } + ) + store = KVCRStore(config, _HostPool()) + store._namespace = storage_namespace(config, store._layouts, shared_mla=True) + store._runner = _CopyTransport(store._layouts) + store._owner_endpoint = "tcp://127.0.0.1:32100" + return store + + def test_shared_mla_namespace_is_rank_invariant_but_not_legacy_compatible(self): + """Replicated readers share identity; ordinary shard keys remain isolated.""" + owner, reader = self._shared_reader(0), self._shared_reader(1) + self.assertEqual(owner._namespace, reader._namespace) + self.assertNotEqual( + reader._namespace, storage_namespace(reader.config, reader._layouts) + ) + changed = self._shared_reader(1, model_name="different-model") + self.assertNotEqual(reader._namespace, changed._namespace) + config = SimpleNamespace(**vars(owner.config)) + config.tp_rank = 1 + self.assertNotEqual( + storage_namespace(owner.config, owner._layouts), + storage_namespace(config, owner._layouts), + ) + + def test_shared_reader_selects_tp0_and_retains_replica_local_fallback(self): + """Missing/disabled external hints must still allow own-replica L3 reads.""" + reader = self._shared_reader(1) + key = block_key(self.keys[0], reader._namespace, "kv") + extra = HiCacheStorageExtraInfo( + extra_info={ + "kv_hints": { + "protocol_version": "0.1", + "actions": [ + { + "action_type": "kv.fetch", + "action_version": "1.0", + "payload": { + "source_control_endpoint": "tcp://127.0.0.1:33000", + "block_hashes": [0], + }, + } + ], + } + } + ) + external = reader._operation_hint(extra, [key]) + self.assertEqual( + external["actions"][0]["payload"]["source_control_endpoint"], + "tcp://127.0.0.1:33000", + ) + for disabled in (False, True): + reader._extra["enable_remote_hint"] = not disabled + own = reader._operation_hint(extra if disabled else None, [key]) + self.assertEqual( + own["actions"][0]["payload"]["source_control_endpoint"], + reader._owner_endpoint, + ) + self.assertEqual( + own["actions"][0]["payload"]["block_hashes"], [int(key[:16], 16)] + ) + owner = self._shared_reader(0) + self.assertIsNone(owner._operation_hint(None, [key])) + + def test_shared_mla_refuses_sharded_topologies_and_auxiliary_state(self): + """Sharing a namespace must never alias KDA/SWA or context/head shards.""" + for changes in ( + {"is_mla_model": False}, + {"pp_size": 2}, + {"attn_cp_size": 2}, + {"should_split_heads": True}, + ): + with ( + self.subTest(changes=changes), + self.assertRaisesRegex(ValueError, "Shared MLA storage requires"), + ): + self._shared_reader(0, **changes) + reader = self._shared_reader(1) + with self.assertRaisesRegex(ValueError, "only supports replicated KV/indexer"): + reader.register_mem_host_pool_v2(_HostPool(), PoolName.MAMBA) + for value in ("false", 1, None): + with ( + self.subTest(value=value), + self.assertRaisesRegex( + ValueError, "share_replicated_mla must be a boolean" + ), + ): + self._shared_reader(0, extra_config={"share_replicated_mla": value}) + + def test_shared_reader_recovers_only_failed_external_pages_from_local_owner(self): + """An external hint cannot hide local-owner pages or turn real misses into hits. + + Both sources complete before the adapter reports success. Already + restored external pages are not submitted again, genuinely absent + pages stay misses, and diagnostics preserve each operation's result. + """ + reader = self._shared_reader(1) + reader._runner = None + reader.register_mem_host_pool_v2(_HostPool((5,)), PoolName.INDEXER) + reader._namespace = storage_namespace( + reader.config, reader._layouts, shared_mla=True + ) + reader._diagnostic_telemetry = True + reader._control_endpoint = "tcp://127.0.0.1:32101" + external_endpoint = "tcp://127.0.0.1:33000" + peers = {reader._owner_endpoint: {}, external_endpoint: {}} + for name, layout in reader._layouts.items(): + for page, source in enumerate( + (reader._owner_endpoint, external_endpoint, external_endpoint) + ): + peers[source][block_key(self.keys[page], reader._namespace, name)] = { + part.name: bytes([17 + page * 13 + component]) * part.size + for component, part in enumerate(layout.components) + } + transport = reader._runner = _PeerCopyTransport(reader._layouts, peers) + reader._metrics = SimpleNamespace(prefetch_pgs=[], prefetch_bandwidth=[]) + for pool in reader.registered_pools.values(): + for buffer in pool.buffers: + buffer.fill_(163) + extra = HiCacheStorageExtraInfo( + extra_info={ + "kv_hints": { + "protocol_version": "0.1", + "actions": [ + { + "action_type": "kv.fetch", + "action_version": "1.0", + "payload": { + "source_control_endpoint": external_endpoint, + "block_hashes": [ + int(key[:16], 16) for key in self.keys + ], + }, + } + ], + } + } + ) + with self.assertLogs( + "sglang.srt.mem_cache.storage.kvcr.kvcr_store", level="INFO" + ) as captured: + restored = reader.batch_get_v2( + [ + PoolTransfer(name, keys=self.keys, host_indices=torch.arange(8, 16)) + for name in reader._layouts + ], + extra, + ) + self.assertEqual( + restored, {name: [True, True, True, False] for name in reader._layouts} + ) + for pool in reader.registered_pools.values(): + for component, buffer in enumerate(pool.buffers): + for page in range(3): + self.assertTrue( + torch.all(buffer[4 + page] == 17 + page * 13 + component) + ) + self.assertTrue(torch.all(buffer[7] == 163)) + failed = tuple( + block_key(key, reader._namespace, name) + for name in reader._layouts + for key in (self.keys[0], self.keys[3]) + ) + self.assertEqual(len(transport.reads), 2) + self.assertEqual(transport.reads[1], (reader._owner_endpoint, failed)) + self.assertEqual(transport.hints, {}) + self.assertEqual(reader._metrics.prefetch_pgs, [3]) + records = [ + json.loads(line.split("KVCR L3 operation ", 1)[1]) + for line in captured.output + ] + self.assertEqual( + [record["stage"] for record in records], + ["primary", "replica-local-fallback"], + ) + self.assertEqual(records[0]["failed_keys"], [key.decode() for key in failed]) + self.assertEqual(records[1]["requested_blocks"], 4) + self.assertEqual( + records[1]["failed_keys"], + [ + block_key(self.keys[3], reader._namespace, name).decode() + for name in reader._layouts + ], + ) + # Already-selected local TP0 is never retried against itself. + with self.assertLogs( + "sglang.srt.mem_cache.storage.kvcr.kvcr_store", level="INFO" + ): + self.assertEqual( + reader.batch_get_v1(self.keys, torch.arange(8, 16)), + [True, False, False, False], + ) + self.assertEqual(len(transport.reads), 3) + self.assertEqual(transport.hints, {}) + + def test_shared_secondary_cannot_create_an_independent_backup(self): + """A misrouted backup fails before submission, never creating replica-local shards.""" + reader = self._shared_reader(1) + with self.assertRaisesRegex(RuntimeError, "must run on TP0"): + reader.batch_set_v1(self.keys[:1], torch.arange(2)) + self.assertEqual(reader._runner.storage, {}) + self.assertEqual(reader._runner.next_handle, 0) + + def _pipeline_fixture(self, pages, *, window=4, logical=False, sidecars=()): + config = SimpleNamespace(**vars(self.config)) + config.extra_config = ( + {} if window is None else {"delivery_pipeline_window": window} + ) + primary = ( + SimpleNamespace(kv_buffer=None) + if logical + else _HostPool(pages=2 * pages + 8) + ) + store = KVCRStore(config, primary) + for name, sizes, page_size in sidecars: + store.register_mem_host_pool_v2( + _HostPool(sizes, page_size, pages=2 * pages + 8), name + ) + store._namespace = storage_namespace(config, store._layouts) + transport = store._runner = _DeferredCopyTransport(store._layouts) + keys = [f"{i + 1:064x}" for i in range(pages)] + for layout in store._layouts.values(): + for component, part in enumerate(layout.components): + for slot in range(pages): + ctypes.memset( + part.address + slot * part.stride, + (slot + component * 17) % 251 + 1, + part.size, + ) + return store, transport, keys + + def _pipeline_transfers(self, store, keys, slots): + return [ + PoolTransfer( + name, + keys=keys, + host_indices=torch.tensor( + [ + slot * layout.page_size + i + for slot in slots + for i in range(layout.page_size) + ] + ), + ) + for name, layout in store._layouts.items() + ] + + def test_pipeline_config_is_strict_before_host_registration(self): + for value in (None, True, False, 0, -1, 2, 3, 5, 1.0, 4.0, "4", [], {}): + with self.subTest(value=value): + config = SimpleNamespace(**vars(self.config)) + config.extra_config = {"delivery_pipeline_window": value} + with patch.object(KVCRStore, "register_mem_pool_host") as register: + with self.assertRaisesRegex(ValueError, "delivery_pipeline_window"): + KVCRStore(config, self.pool) + register.assert_not_called() + + def test_pipeline_default_shape_and_unsplit_backup(self): + """Absent/1 keeps one large public operation; opting in only splits reads.""" + pages = 1812 + for window in (None, 1, 4): + with self.subTest(window=window): + store, transport, keys = self._pipeline_fixture( + pages, window=window, sidecars=((PoolName.INDEXER, (5, 2), 1),) + ) + store._diagnostic_telemetry = True + store._control_endpoint = "tcp://target:32100" + with self.assertLogs( + "sglang.srt.mem_cache.storage.kvcr.kvcr_store", level="INFO" + ) as captured: + saved = store.batch_set_v2( + self._pipeline_transfers(store, keys, range(pages)) + ) + self.assertTrue(all(all(hits) for hits in saved.values())) + self.assertEqual(transport.next_handle, 1) + restored = store.batch_get_v2( + self._pipeline_transfers(store, keys, range(pages, 2 * pages)) + ) + self.assertTrue(all(all(hits) for hits in restored.values())) + expected = [512] * 7 + [40] if window == 4 else [2 * pages] + self.assertEqual([len(read[1]) for read in transport.reads], expected) + self.assertEqual(store.prefetch_batch_pages, 128) + records = [ + json.loads(line.split("KVCR L3 operation ", 1)[1]) + for line in captured.output + ] + self.assertEqual(len(records), 1 + len(expected)) + self.assertEqual(records[0]["direction"], "deposit") + self.assertEqual(records[0]["requested_blocks"], 2 * pages) + if window == 4: + reads = records[1:] + self.assertEqual( + [r["pipeline_chunk_index"] for r in reads], list(range(8)) + ) + self.assertEqual(len({r["logical_start_epoch"] for r in reads}), 1) + for index, record in enumerate(reads): + self.assertEqual(record["delivery_pipeline_window"], 4) + self.assertLessEqual( + record["logical_start_epoch"], record["start_epoch"] + ) + self.assertLessEqual(record["start_epoch"], record["epoch"]) + if index: + self.assertTrue( + all( + record[field] == 0 + for field in ( + "prepare_ms", + "lock_wait_ms", + "descriptor_ms", + "align_before_ms", + "hint_ms", + ) + ) + ) + self.assertAlmostEqual( + reads[0]["prepare_ms"], + sum( + reads[0][field] + for field in ("descriptor_ms", "align_before_ms", "hint_ms") + ), + ) + for pool in store.registered_pools.values(): + for buffer in pool.buffers: + self.assertTrue( + torch.equal(buffer[:pages], buffer[pages : 2 * pages]) + ) + + def test_pipeline_boundaries_reverse_completions_and_arbitrary_slots(self): + for pages in (1, 511, 512, 513, 1812, 2048, 2601): + with self.subTest(pages=pages): + store, transport, keys = self._pipeline_fixture(pages) + store.batch_set_v2(self._pipeline_transfers(store, keys, range(pages))) + transport.completed[999999] = {b"other-caller": True} + destination = list(reversed(range(pages, 2 * pages))) + restored = store.batch_get_v2( + self._pipeline_transfers(store, keys, destination) + ) + self.assertEqual(restored, {"kv": [True] * pages}) + chunks = (pages + 511) // 512 + self.assertEqual(len(transport.reads), chunks) + self.assertEqual(transport.max_pending, min(4, chunks)) + first_complete = next( + i + for i, event in enumerate(transport.events) + if event[0] == "complete" + ) + self.assertEqual(first_complete, min(4, chunks)) + self.assertEqual(store._wait(999999), {b"other-caller": True}) + for part in store._layouts["kv"].components: + for source, target in enumerate(destination): + self.assertEqual( + ctypes.string_at( + part.address + source * part.stride, part.size + ), + ctypes.string_at( + part.address + target * part.stride, part.size + ), + ) + self.assertEqual(transport.pending, {}) + + def test_pipeline_misses_ignore_other_chunks_keys_and_keep_one_metric_and_digest( + self, + ): + pages = 1100 + store, transport, keys = self._pipeline_fixture( + pages, sidecars=((PoolName.INDEXER, (5,), 2),) + ) + store.batch_set_v2(self._pipeline_transfers(store, keys, range(pages))) + missing = block_key(keys[600], store._namespace, "indexer") + omitted = block_key(keys[10], store._namespace, "kv") + transport.storage.pop(missing) + transport.storage.pop(omitted) + + def corrupt_result(_handle, blocks, result): + result.pop(omitted, None) + # Both are real identities from this logical call, but not this handle. + for key in (missing, omitted): + if key not in blocks: + result[key] = SimpleNamespace(success=True) + result[b"unsolicited"] = SimpleNamespace(success=True) + return result + + transport.result_hook = corrupt_result + store._metrics = SimpleNamespace(prefetch_pgs=[], prefetch_bandwidth=[]) + store._verify_transfer_bytes = store._diagnostic_telemetry = True + store._control_endpoint = "tcp://target:32100" + real_metrics = store._record_metrics + with patch.object(store, "_record_metrics", wraps=real_metrics) as metrics: + with self.assertLogs( + "sglang.srt.mem_cache.storage.kvcr.kvcr_store", level="INFO" + ) as captured: + restored = store.batch_get_v2( + self._pipeline_transfers(store, keys, range(pages, 2 * pages)) + ) + metrics.assert_called_once() + elapsed = metrics.call_args.kwargs["elapsed"] + self.assertEqual([i for i, hit in enumerate(restored["kv"]) if not hit], [10]) + self.assertEqual( + [i for i, hit in enumerate(restored["indexer"]) if not hit], [600] + ) + self.assertEqual(store._metrics.prefetch_pgs, [pages - 2]) + self.assertAlmostEqual( + store._metrics.prefetch_bandwidth[0], (16 * pages - 16) / elapsed / 1e9 + ) + operation_records = [ + json.loads(line.split("KVCR L3 operation ", 1)[1]) + for line in captured.output + if "KVCR L3 operation " in line + ] + self.assertEqual( + [r["requested_blocks"] for r in operation_records], + [512, 512, 512, 512, 152], + ) + self.assertEqual( + {k for r in operation_records for k in r["failed_keys"]}, + {missing.decode(), omitted.decode()}, + ) + digests = json.loads( + next( + line.split("KVCR L3 byte verification ", 1)[1] + for line in captured.output + if "KVCR L3 byte verification " in line + ) + )["digests"] + self.assertEqual(len(digests), 2 * pages - 2) + self.assertNotIn(missing.decode(), digests) + self.assertNotIn(omitted.decode(), digests) + self.assertEqual(transport.pending, {}) + for name, pool in store.registered_pools.items(): + for buffer in pool.buffers: + absent = 10 if name == "kv" else 600 + self.assertTrue( + torch.equal(buffer[:absent], buffer[pages : pages + absent]) + ) + self.assertTrue( + torch.equal( + buffer[absent + 1 : pages], + buffer[pages + absent + 1 : 2 * pages], + ) + ) + self.assertEqual(int(buffer[pages + absent].sum()), 0) + + def test_pipeline_logical_anchor_and_independent_trailing_checkpoint_indices(self): + pages = 600 + store, transport, keys = self._pipeline_fixture( + pages, + logical=True, + sidecars=( + (PoolName.DEEPSEEK_V4_C4, (7, 3), 2), + (PoolName.MAMBA, (5, 2), 1), + ), + ) + all_pages = PoolName.DEEPSEEK_V4_C4 + transfers = [ + PoolTransfer(PoolName.KV, keys=keys, host_indices=None), + PoolTransfer(all_pages, keys=keys, host_indices=torch.arange(2 * pages)), + PoolTransfer( + PoolName.MAMBA, + keys=[keys[9], keys[-1]], + host_indices=torch.tensor([3, 17]), + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ), + ] + store.batch_set_v2(transfers) + restored = store.batch_get_v2( + [ + PoolTransfer(PoolName.KV, keys=keys), + PoolTransfer( + all_pages, + keys=keys, + host_indices=torch.arange(2 * pages, 4 * pages), + ), + PoolTransfer( + PoolName.MAMBA, + keys=[keys[9], keys[-1]], + host_indices=torch.tensor([pages + 7, pages + 31]), + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ), + ] + ) + self.assertEqual( + restored, + { + "kv": [True] * pages, + str(all_pages): [True] * pages, + "mamba": [True, True], + }, + ) + self.assertEqual([len(read[1]) for read in transport.reads], [512, 90]) + for buffer in store.registered_pools[str(all_pages)].buffers: + self.assertTrue(torch.equal(buffer[:pages], buffer[pages : 2 * pages])) + for part in store._layouts["mamba"].components: + for source, target in ((3, pages + 7), (17, pages + 31)): + self.assertEqual( + ctypes.string_at(part.address + source * part.stride, part.size), + ctypes.string_at(part.address + target * part.stride, part.size), + ) + boundary = store.batch_exists_v2( + keys, + [ + PoolTransfer( + PoolName.MAMBA, + keys=[keys[-1]], + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ), + PoolTransfer(all_pages), + ], + ) + self.assertEqual(boundary.restorable_prefix_pages, [10, pages]) + + def test_pipeline_duplicate_identity_and_destination_alias_preflight(self): + for alias in ( + "identity", + "slot", + "cross-pool", + "cross-component", + "partial-component", + ): + with self.subTest(alias=alias): + pages = 513 + store, transport, keys = self._pipeline_fixture(pages) + if alias == "identity": + keys[-1] = keys[0] + if alias in ("cross-pool", "cross-component", "partial-component"): + # Re-register real overlapping tensor views before first IO. + store._runner = None + pool = store.registered_pools["kv"] + if alias == "cross-pool": + store.register_mem_host_pool_v2(pool, PoolName.INDEXER) + else: + pool.buffers[1] = ( + pool.buffers[0][:, :3] + if alias == "cross-component" + else pool.buffers[0][:, 6:8] + ) + store._layouts["kv"] = describe_pool("kv", pool) + store._runner = transport = _DeferredCopyTransport(store._layouts) + slots = list(range(pages, 2 * pages)) + if alias == "slot": + slots[-1] = slots[0] + before = [ + ctypes.string_at(p.address, (p.count - 1) * p.stride + p.size) + for layout in store._layouts.values() + for p in layout.components + ] + with self.assertRaisesRegex(ValueError, "Duplicate|Overlapping"): + store.batch_get_v2(self._pipeline_transfers(store, keys, slots)) + self.assertEqual(transport.reads, []) + self.assertEqual(transport.hints, {}) + after = [ + ctypes.string_at(p.address, (p.count - 1) * p.stride + p.size) + for layout in store._layouts.values() + for p in layout.components + ] + self.assertEqual(before, after) + + def test_pipeline_controller_partial_prefix_keeps_logical_ack_count(self): + pages = 1812 + for missing in (None, 600): + with self.subTest(missing=missing): + controller, operation, pools, transport = self._controller_with_indexer( + pages=pages, batch_pages=2048, pipeline_window=4 + ) + for pool in pools.values(): + for part, buffer in enumerate(pool.buffers): + buffer[:pages].fill_(31 + part) + controller.storage_backend.batch_set_v2( + [ + PoolTransfer( + name, + keys=operation.hash_value, + host_indices=torch.arange(2 * pages), + ) + for name in pools + ] + ) + if missing is not None: + transport.storage.pop( + block_key( + operation.hash_value[missing], + controller.storage_backend._namespace, + "indexer", + ) + ) + with ( + self.assertLogs(level="WARNING") + if missing is not None + else ExitStack() + ): + completed = controller._page_transfer(operation) + prefix = pages if missing is None else missing + self.assertEqual(completed, prefix) + self.assertEqual( + controller.prefetch_sync_queue.get_nowait().completed_tokens, + prefix * 2, + ) + self.assertTrue(controller.prefetch_sync_queue.empty()) + self.assertEqual( + transport.next_handle, 9 + ) # one deposit, eight read chunks + for pool in pools.values(): + for buffer in pool.buffers: + self.assertTrue( + torch.equal(buffer[:prefix], buffer[pages : pages + prefix]) + ) + + def _pipeline_mla_fixture(self, pages=2200): + store, _, keys = self._pipeline_fixture( + pages, sidecars=((PoolName.INDEXER, (5,), 1),) + ) + store._share_replicated_mla = True + store.config.is_mla_model, store.config.tp_rank = True, 1 + store._source_offset = 0 + store._owner_endpoint = "tcp://owner:32100" + store._control_endpoint = "tcp://target:32101" + store._namespace = storage_namespace( + store.config, store._layouts, shared_mla=True + ) + external = "tcp://external:33000" + peers = {external: {}, store._owner_endpoint: {}} + for name, layout in store._layouts.items(): + for page, key in enumerate(keys): + identity = block_key(key, store._namespace, name) + # Owner also holds different bytes for external successes, so + # accidentally retrying them is observable in the final bytes. + if page % 3 != 2: + peers[store._owner_endpoint][identity] = { + part.name: bytes([73 + component]) * part.size + for component, part in enumerate(layout.components) + } + if page % 3 == 1: + peers[external][identity] = { + part.name: bytes([29 + component]) * part.size + for component, part in enumerate(layout.components) + } + transport = store._runner = _DeferredCopyTransport(store._layouts, peers) + extra = HiCacheStorageExtraInfo( + extra_info={ + "kv_hints": { + "protocol_version": "0.1", + "actions": [ + { + "action_type": "kv.fetch", + "action_version": "1.0", + "payload": { + "source_control_endpoint": external, + "block_hashes": [0], + }, + } + ], + } + } + ) + return store, transport, keys, extra + + def test_pipeline_mla_retries_only_failed_chunks_after_primary_terminal(self): + pages = 2200 + store, transport, keys, extra = self._pipeline_mla_fixture(pages) + store._metrics = SimpleNamespace(prefetch_pgs=[], prefetch_bandwidth=[]) + store._diagnostic_telemetry = True + with self.assertLogs( + "sglang.srt.mem_cache.storage.kvcr.kvcr_store", level="INFO" + ) as captured: + restored = store.batch_get_v2( + self._pipeline_transfers(store, keys, range(pages, 2 * pages)), extra + ) + expected = [page % 3 != 2 for page in range(pages)] + self.assertEqual(restored, {"kv": expected, "indexer": expected}) + self.assertEqual(store._metrics.prefetch_pgs, [sum(expected)]) + self.assertEqual(transport.max_pending, 4) + self.assertFalse(transport.pending) + self.assertFalse(transport.hints) + external_reads = [r for r in transport.reads if r[3] == "tcp://external:33000"] + local_reads = [r for r in transport.reads if r[3] == store._owner_endpoint] + self.assertEqual(len(external_reads), 9) + self.assertEqual(len(local_reads), 9) + self.assertEqual(len({r[2] for r in external_reads + local_reads}), 2) + event_index = {event: i for i, event in enumerate(transport.events)} + self.assertLess( + max(event_index[("complete", r[0])] for r in external_reads), + min(event_index[("submit", r[0])] for r in local_reads), + ) + for source, identity, _ in transport.writes: + page = int(identity.split(b"#", 1)[0], 16) - 1 + self.assertEqual( + source, + "tcp://external:33000" if page % 3 == 1 else store._owner_endpoint, + ) + self.assertEqual(len(transport.writes), 2 * sum(expected)) + for layout in store._layouts.values(): + for component, part in enumerate(layout.components): + for page in range(pages): + value = ( + 0 + if page % 3 == 2 + else (29 if page % 3 == 1 else 73) + component + ) + self.assertEqual( + ctypes.string_at( + part.address + (pages + page) * part.stride, part.size + ), + bytes([value]) * part.size, + ) + records = [ + json.loads(line.split("KVCR L3 operation ", 1)[1]) + for line in captured.output + ] + primary = {r["operation"]: r for r in records if r["stage"] == "primary"} + self.assertEqual([r["pipeline_chunk_index"] for r in records], list(range(18))) + self.assertEqual(len({r["logical_start_epoch"] for r in records}), 1) + for record in records: + self.assertLessEqual(record["logical_start_epoch"], record["start_epoch"]) + self.assertLessEqual(record["start_epoch"], record["epoch"]) + self.assertGreaterEqual(record["submit_ms"], 0) + self.assertGreaterEqual(record["wait_ms"], 0) + if record["stage"] != "primary": + parent = primary[record["parent_operation"]] + self.assertEqual(record["requested_blocks"], len(parent["failed_keys"])) + self.assertTrue( + set(record["failed_keys"]) <= set(parent["failed_keys"]) + ) + self.assertEqual(record["source"], store._owner_endpoint) + before = len(transport.reads) + self.assertEqual( + store.batch_get_v1(keys, torch.arange(2 * pages, 4 * pages)), + [page % 3 != 2 for page in range(pages)], + ) + self.assertEqual(len(transport.reads) - before, 5) # no fallback to own owner + self.assertTrue( + all(r[3] == store._owner_endpoint for r in transport.reads[before:]) + ) + + def test_pipeline_fatal_entry_precedes_primary_or_fallback_hint_discard(self): + """A Python mock unwinds later; only callback ENTRY proves ordering here.""" + for phase in ("primary", "fallback"): + for failure in ( + "unknown-submit", + "none-handle", + "duplicate-handle", + "wait", + "bookkeeping", + ): + with self.subTest(phase=phase, failure=failure): + pages = 1100 + store, transport, keys, extra = self._pipeline_mla_fixture(pages) + real_deliver, real_wait = transport.deliver, store._wait + events, selected = [], [] + + def is_selected( + request_id, transport=transport, store=store, phase=phase + ): + source = transport.hints[request_id]["actions"][0]["payload"][ + "source_control_endpoint" + ] + return (source == store._owner_endpoint) == ( + phase == "fallback" + ) + + def submit( + blocks, + request_id=None, + real_deliver=real_deliver, + is_selected=is_selected, + selected=selected, + failure=failure, + ): + handle = real_deliver(blocks, request_id) + if is_selected(request_id): + selected.append(handle) + if len(selected) == 2: + if failure == "unknown-submit": + raise KeyboardInterrupt( + "enqueued without returning handle" + ) + if failure == "none-handle": + return None + if failure == "duplicate-handle": + return selected[0] + return handle + + def wait( + handle, selected=selected, failure=failure, real_wait=real_wait + ): + if selected and handle == selected[0] and failure == "wait": + # One completion may already exist before another poll fails. + raise GeneratorExit("lost poll state") + return real_wait(handle) + + def clock(failure=failure, selected=selected): + if failure == "bookkeeping" and len(selected) == 2: + raise SystemExit("post-submit clock failed") + return 1.0 + + def fatal(code, events=events, transport=transport): + events.append("fatal-entry") + self.assertEqual(code, 1) + self.assertTrue(transport.pending) + self.assertTrue(transport.hints) + self.assertNotIn("discard", events) + self.assertNotIn("ack", events) + raise SystemExit(code) + + def discard(request_id, events=events, transport=transport): + # Simulate the unwind caused by a Python-raised exit mock. + events.append("discard") + transport.hints.pop(request_id) + + store._diagnostic_telemetry = failure == "bookkeeping" + with ( + patch.object(transport, "deliver", submit), + patch.object(store, "_wait", wait), + patch.object(transport, "discard_hint", discard), + patch( + "sglang.srt.mem_cache.storage.kvcr.delivery.time.perf_counter", + clock, + ), + patch( + "sglang.srt.mem_cache.storage.kvcr.delivery.logger.critical", + side_effect=RuntimeError("logging failed"), + ), + patch( + "sglang.srt.mem_cache.storage.kvcr.delivery.os._exit", fatal + ), + self.assertRaises(SystemExit), + ): + store.batch_get_v2( + self._pipeline_transfers( + store, keys, range(pages, 2 * pages) + ), + extra, + ) + events.append("ack") + self.assertEqual(events[0], "fatal-entry") + + def test_pipeline_result_and_diagnostic_exceptions_only_after_all_writes_terminal( + self, + ): + for failure in ("result", "diagnostic"): + with self.subTest(failure=failure): + store, transport, keys = self._pipeline_fixture(2601) + store.batch_set_v2(self._pipeline_transfers(store, keys, range(2601))) + + class ExplodingResult: + def __init__(inner, transport): + inner.transport = transport + + @property + def success(inner): + self.assertFalse(inner.transport.pending) + raise ValueError("terminal result failure") + + if failure == "result": + + def replace(_handle, _blocks, result, transport=transport): + return dict.fromkeys(result, ExplodingResult(transport)) + + transport.result_hook = replace + store._diagnostic_telemetry = failure == "diagnostic" + store._control_endpoint = "tcp://target:32100" + + def log(*_args, transport=transport, **_kwargs): + self.assertFalse(transport.pending) + raise ValueError("terminal diagnostic failure") + + with ( + patch( + "sglang.srt.mem_cache.storage.kvcr.kvcr_store.logger.info", log + ), + patch( + "sglang.srt.mem_cache.storage.kvcr.delivery.os._exit" + ) as fatal, + self.assertRaisesRegex(ValueError, "terminal"), + ): + store.batch_get_v2( + self._pipeline_transfers(store, keys, range(2601, 5202)) + ) + fatal.assert_not_called() + self.assertEqual(len(transport.reads), 6) + self.assertFalse(transport.pending) + + def test_sparse_checkpoints_intersect_prefix_sets(self): + """KV max=4 and checkpoint max=3 do not imply that boundary 2 is valid.""" + state = _HostPool(sizes=(5,), page_size=1) + # Sidecar-specific pools (notably DSA indexer) expose page accessors, + # not the KV anchor's kv_buffer attribute. + del state.kv_buffer + self.backend._runner = None + self.backend.register_mem_host_pool_v2(state, PoolName.MAMBA) + self.backend._namespace = storage_namespace(self.config, self.backend._layouts) + self.backend._runner = self.transport = _CopyTransport(self.backend._layouts) + self.backend.batch_set_v1(self.keys, torch.arange(8)) + self.backend.batch_set_v2( + [ + PoolTransfer( + PoolName.MAMBA, + keys=[self.keys[0], self.keys[2]], + host_indices=torch.tensor([1, 3]), + ) + ] + ) + result = self.backend.batch_exists_v2( + self.keys, + [ + PoolTransfer( + PoolName.MAMBA, + keys=[self.keys[-1]], + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + ], + ) + self.assertEqual(result.restorable_prefix_pages, [1, 3]) + self.assertEqual(result.kv_hit_pages, 3) + + def test_invalid_physical_pages_are_rejected_before_io(self): + layout = describe_pool("kv", self.pool) + for indices in ([0, 3], [1, 2], [-2, -1], [16, 17], [0]): + with self.subTest(indices=indices), self.assertRaises(ValueError): + page_slots(layout, self.keys[:1], torch.tensor(indices)) + self.assertFalse(self.transport.storage) + + def test_completion_drain_preserves_another_callers_result(self): + """One poll drain must not drop a concurrent backup/prefetch completion.""" + self.transport.completed = {11: {"one": True}, 12: {"two": True}} + self.assertEqual(self.backend._wait(12), {"two": True}) + self.assertEqual(self.backend._wait(11), {"one": True}) + + def test_hint_is_rank_aligned_and_discarded(self): + info = HiCacheStorageExtraInfo( + extra_info={ + "kv_hints": { + "protocol_version": "0.1", + "message_id": "r", + "actions": [ + { + "action_id": "a", + "action_type": "kv.fetch", + "action_version": "1.0", + "payload": { + "source_control_endpoint": "tcp://source:32000", + "block_hashes": [(1 << 63) + 1], + }, + } + ], + } + } + ) + self.assertEqual( + peer_hint(info, 3)["actions"][0]["payload"]["source_control_endpoint"], + "tcp://source:32003", + ) + with self.backend._hint(info) as request_id: + self.assertEqual(len(self.transport.hints), 1) + self.assertIn(request_id, self.transport.hints) + self.assertEqual(self.transport.hints, {}) + info.extra_info["kv_hints"]["actions"][0]["payload"]["block_hashes"] = [True] + self.assertIsNone(peer_hint(info, 0)) + + def test_namespace_separates_model_format_and_shard_not_replica(self): + namespace = storage_namespace(self.config, self.backend._layouts) + self.config.dp_rank = 1 + self.assertEqual( + storage_namespace(self.config, self.backend._layouts), namespace + ) + for attribute, value in (("model_name", "other-model"), ("tp_rank", 1)): + with self.subTest(attribute=attribute): + config = SimpleNamespace(**vars(self.config)) + setattr(config, attribute, value) + self.assertNotEqual( + storage_namespace(config, self.backend._layouts), namespace + ) + + def test_trailing_window_holes_do_not_become_contiguous_prefixes(self): + self.assertEqual(resume_boundaries([False, True, True, False], trailing=2), [3]) + self.assertEqual(resume_boundaries([True, False, True], trailing=None), [1]) + + +class TestKVCRDeliveryWindow(CustomTestCase): + """Dependency-free boundary cases also runnable without loading Torch. + + The CPU suite imports normally. A minimal local gate can load these exact + fixture/class definitions and the adapter class without its import graph. + """ + + def _fixture(self, pages=2601, *, diagnostic=False): + buffers = [(ctypes.c_ubyte * (pages * size))() for size in (5, 3)] + parts = tuple( + SimpleNamespace( + name=f"kv_{i}", + address=ctypes.addressof(buffer), + size=size, + stride=size, + count=pages, + ) + for i, (buffer, size) in enumerate(zip(buffers, (5, 3))) + ) + layouts = {"kv": SimpleNamespace(components=parts)} + store = KVCRStore.__new__(KVCRStore) + store._diagnostic_telemetry = diagnostic + store._completion = threading.Condition(threading.RLock()) + store._completed = {} + transport = store._runner = _DeferredCopyTransport(layouts) + blocks = { + f"page-{page}".encode(): [ + SimpleNamespace(label=part.name, element_index=page) for part in parts + ] + for page in range(pages) + } + for page, key in enumerate(blocks): + transport.storage[key] = { + part.name: bytes([(page + i * 17) % 251 + 1]) * part.size + for i, part in enumerate(parts) + } + return store, transport, blocks, layouts, buffers + + def test_window_uses_actual_bytes_for_boundaries_short_tail_and_reverse_results( + self, + ): + for count in (1, 511, 512, 513, 1812, 2048, 2601): + with self.subTest(count=count): + store, transport, blocks, layouts, buffers = self._fixture(count) + transport.completed[999999] = {b"other": True} + validate_destination_spans(blocks, layouts) + completed = run_window( + store, + (Delivery(chunk) for chunk in block_chunks(blocks)), + None, + set(), + ) + self.assertEqual(len(requested_results(completed)), count) + self.assertEqual(transport.max_pending, min(4, (count + 511) // 512)) + self.assertEqual( + [len(d.blocks) for d in completed], + [512] * (count // 512) + ([count % 512] if count % 512 else []), + ) + for i, buffer in enumerate(buffers): + size = layouts["kv"].components[i].size + expected = b"".join( + bytes([(page + i * 17) % 251 + 1]) * size + for page in range(count) + ) + self.assertEqual(bytes(buffer), expected) + self.assertEqual(store._wait(999999), {b"other": True}) + self.assertFalse(transport.pending) + + def test_overlap_preflight_covers_aliases_across_chunk_boundaries(self): + for alias in ("slot", "component", "partial", "pool"): + with self.subTest(alias=alias): + _store, transport, blocks, layouts, buffers = self._fixture(513) + refs = blocks[b"page-512"] + first = layouts["kv"].components[0] + if alias == "slot": + refs[0].element_index = 0 + else: + part = SimpleNamespace(**vars(first)) + part.name = "sidecar" + part.address += 1 if alias == "partial" else 0 + part.size = 2 + layouts["sidecar"] = SimpleNamespace(components=(part,)) + new_ref = SimpleNamespace(label=part.name, element_index=0) + if alias == "pool": + blocks[b"other-pool"] = [new_ref] + else: + refs.append(new_ref) + with self.assertRaisesRegex(ValueError, "Overlapping"): + validate_destination_spans(blocks, layouts) + self.assertEqual(transport.reads, []) + self.assertTrue(all(not any(buffer) for buffer in buffers)) + + def test_unrequested_successes_never_override_missing_or_false_own_results(self): + store, transport, blocks, _, buffers = self._fixture(1100) + transport.storage.pop(b"page-10") + transport.storage.pop(b"page-600") + + def result_hook(_handle, chunk, result): + result.pop(b"page-10", None) + for key in (b"page-10", b"page-600"): + if key not in chunk: + result[key] = SimpleNamespace(success=True) + return result + + transport.result_hook = result_hook + completed = run_window( + store, (Delivery(c) for c in block_chunks(blocks)), None, set() + ) + result = requested_results(completed) + self.assertNotIn(b"page-10", result) + self.assertFalse(result[b"page-600"].success) + for buffer, size in zip(buffers, (5, 3)): + self.assertEqual(bytes(buffer[10 * size : 11 * size]), bytes(size)) + self.assertEqual(bytes(buffer[600 * size : 601 * size]), bytes(size)) + self.assertTrue(any(buffer[601 * size : 602 * size])) + self.assertFalse(transport.pending) + + def _assert_fatal(self, store, transport, blocks, action): + """Assert ENTRY ordering, not production finally suppression by Python mocks.""" + hint = { + "actions": [{"payload": {"source_control_endpoint": "tcp://source:32100"}}] + } + events = [] + + def fatal(code): + events.append("fatal") + self.assertEqual(code, 1) + self.assertTrue(transport.pending) + self.assertTrue(transport.hints) + self.assertNotIn("discard", events) + self.assertNotIn("reuse", events) + raise SystemExit(code) + + def discard(request_id): + events.append("discard") + transport.hints.pop(request_id) + + with ( + patch.object(store, "_operation_hint", return_value=hint), + patch.object(transport, "discard_hint", discard), + patch( + "sglang.srt.mem_cache.storage.kvcr.delivery.logger.critical", + side_effect=RuntimeError("logger failed"), + ), + patch("sglang.srt.mem_cache.storage.kvcr.delivery.os._exit", fatal), + self.assertRaises(SystemExit), + ): + with store._hint(None, blocks) as request_id: + action(request_id) + events.append("reuse") + self.assertEqual(events[0], "fatal") + + def test_fatal_entry_for_unknown_submit_and_invalid_or_duplicate_handles(self): + for invalid in (None, True, 0, -1, "7", [], "duplicate", "unknown"): + with self.subTest(invalid=invalid): + store, transport, blocks, _, buffers = self._fixture(1100) + real_deliver = transport.deliver + handles = [] + + def submit( + chunk, + request_id=None, + real_deliver=real_deliver, + handles=handles, + invalid=invalid, + ): + handle = real_deliver(chunk, request_id) + handles.append(handle) + if len(handles) == 2: + if invalid == "unknown": + raise KeyboardInterrupt("enqueue then raise") + return handles[0] if invalid == "duplicate" else invalid + return handle + + with patch.object(transport, "deliver", submit): + self._assert_fatal( + store, + transport, + blocks, + lambda request, store=store, blocks=blocks: run_window( + store, + (Delivery(c) for c in block_chunks(blocks)), + request, + set(), + ), + ) + self.assertTrue(all(not any(buffer) for buffer in buffers)) + + def test_fatal_entry_for_late_bookkeeping_and_iterator_failures(self): + class BrokenSet(set): + def add(self, _handle): + raise MemoryError("lost handle while adding") + + for failure in ("seen", "clock", "iterator"): + with self.subTest(failure=failure): + store, transport, blocks, _, buffers = self._fixture( + 1100, diagnostic=True + ) + + def deliveries(blocks=blocks, failure=failure): + for i, chunk in enumerate(block_chunks(blocks)): + if i == 2 and failure == "iterator": + raise GeneratorExit("late chunk preparation failed") + yield Delivery(chunk) + + def clock(failure=failure, transport=transport): + if failure == "clock" and transport.pending: + raise SystemExit("post-submit timestamp failed") + return 1.0 + + with patch( + "sglang.srt.mem_cache.storage.kvcr.delivery.time.perf_counter", + clock, + ): + self._assert_fatal( + store, + transport, + blocks, + lambda request, store=store, failure=failure: run_window( + store, + deliveries(), + request, + BrokenSet() if failure == "seen" else set(), + ), + ) + self.assertTrue(all(not any(buffer) for buffer in buffers)) + + def test_fatal_entry_after_partial_completion_then_poll_failure(self): + store, transport, blocks, _, buffers = self._fixture(2601) + real_poll = transport.poll_completed + polls = [] + + def poll(): + polls.append(True) + if len(polls) > 1: + raise KeyboardInterrupt("poll failed with writes still pending") + # Complete just the last handle; _wait must retain it while polling + # for its requested first handle. Other writes remain outstanding. + last = next(reversed(transport.pending)) + pending = transport.pending + transport.pending = {last: pending.pop(last)} + try: + return real_poll() + finally: + transport.pending = pending + + with patch.object(transport, "poll_completed", poll): + self._assert_fatal( + store, + transport, + blocks, + lambda request: run_window( + store, (Delivery(c) for c in block_chunks(blocks)), request, set() + ), + ) + self.assertEqual(len(transport.writes), 512) + self.assertEqual(bytes(buffers[0][:5]), bytes(5)) + self.assertTrue(any(buffers[0][1536 * 5 : 2048 * 5])) + + def test_handle_reuse_after_completed_window_is_fatal(self): + store, transport, blocks, _, _ = self._fixture(2601) + real_deliver = transport.deliver + handles = [] + + def submit(chunk, request_id=None): + handle = real_deliver(chunk, request_id) + handles.append(handle) + return handles[0] if len(handles) == 5 else handle + + with patch.object(transport, "deliver", submit): + self._assert_fatal( + store, + transport, + blocks, + lambda request: run_window( + store, (Delivery(c) for c in block_chunks(blocks)), request, set() + ), + ) + self.assertEqual(len(transport.writes), 2048) + + def test_native_timestamps_bracket_submit_and_completion_only_when_enabled(self): + for diagnostic in (False, True): + with self.subTest(diagnostic=diagnostic): + store, transport, blocks, _, _ = self._fixture( + 1812, diagnostic=diagnostic + ) + clock_events = [] + ticks = iter(range(100, 200)) + + def epoch(ticks=ticks, clock_events=clock_events, transport=transport): + value = next(ticks) + clock_events.append( + (value, len(transport.reads), len(transport.writes)) + ) + return value + + with patch( + "sglang.srt.mem_cache.storage.kvcr.delivery.time.time", epoch + ): + completed = run_window( + store, (Delivery(c) for c in block_chunks(blocks)), None, set() + ) + if diagnostic: + self.assertEqual(len(clock_events), 8) + self.assertEqual( + [event[1] for event in clock_events[:4]], [0, 1, 2, 3] + ) + self.assertTrue(all(event[2] == 1812 for event in clock_events[4:])) + self.assertTrue( + all(d.prepared_epoch < d.finished_epoch for d in completed) + ) + else: + self.assertEqual(clock_events, []) + self.assertTrue( + all( + d.prepared_epoch is None and d.finished_epoch is None + for d in completed + ) + ) + + +class TestKVCROwnedNUMAProof(CustomTestCase): + """Reject placement or registration claims that do not cover owned bytes.""" + + def _mapping(self, *, shared=True): + page = mmap.PAGESIZE + start, end = 16 * page, 18 * page + span = f"{start:x}-{end:x}" + tail = ( + "rw-s 00000000 00:01 42 /dev/zero (deleted)" + if shared + else ("rw-p 00000000 00:00 0") + ) + maps = f"{span} {tail}\n" + numa = f"{start:x} bind:1 dirty=2 N1=2\n" + kib = 2 * page // 1024 + smaps = maps + ( + f"Size: {kib} kB\nRss: {kib} kB\nAnonHugePages: 0 kB\n" + f"KernelPageSize: {page // 1024} kB\nShmemPmdMapped: 0 kB\n" + ) + return maps, numa, smaps, start, 2 * page + + def test_node_config_rejects_boolean_coercion_and_non_objects(self): + for value in (True, False, -1, 65536, 1.0, "1"): + with self.subTest(value=value), self.assertRaises(ValueError): + node_option(value, "host_pool_numa_node") + for raw in ("[]", "1", "null", [], "@"): + with self.subTest(raw=raw), self.assertRaises(ValueError): + load_config(raw) + self.assertEqual(load_config(None), {}) + self.assertEqual( + load_config('{"host_pool_numa_node": 0}'), {"host_pool_numa_node": 0} + ) + + def test_fully_resident_owned_mapping_and_covered_subrange(self): + for shared in (True, False): + with self.subTest(shared=shared): + maps, numa, smaps, address, length = self._mapping(shared=shared) + proof = prove_range(maps, numa, smaps, address, length, 1, shared) + self.assertEqual(proof["node_pages"], {1: 2}) + self.assertTrue(proof["all_covering_pages_resident_on_requested_node"]) + subrange = prove_range( + maps, numa, smaps, address, length // 2, 1, shared + ) + self.assertEqual(subrange["covering_end"], address + length) + self.assertEqual(subrange["length"], length // 2) + + def test_wrong_node_incomplete_residency_and_wrong_sharing_fail_closed(self): + maps, numa, smaps, address, length = self._mapping() + cases = [ + (maps, numa.replace("N1=2", "N0=2"), smaps, True), + (maps, numa.replace("N1=2", "N1=1"), smaps, True), + (maps, numa.replace("bind:1", "default"), smaps, True), + (maps, numa, smaps.replace(f"Rss: {length // 1024} kB", "Rss: 0 kB"), True), + (maps, numa, smaps, False), + ] + for m, n, s, shared in cases: + with self.subTest(numa=n, shared=shared), self.assertRaises(RuntimeError): + prove_range(m, n, s, address, length, 1, shared) + with self.assertRaises(RuntimeError): + prove_range(maps, numa, smaps, address, length + mmap.PAGESIZE, 1, True) + + def test_cuda_registration_requires_exact_contiguous_owned_coverage(self): + maps, numa, smaps, address, length = self._mapping() + proof = prove_range(maps, numa, smaps, address, length, 1, True) + tensor = SimpleNamespace( + data_ptr=lambda: address, numel=lambda: length, element_size=lambda: 1 + ) + setattr(tensor, HOST_PROOF_ATTR, proof) + target = "sglang.srt.mem_cache.storage.kvcr.numa_pool.snapshot_range" + with patch(target, return_value=proof): + audit_registered_host(tensor, [(address, length)]) + for ranges in ( + [], + [(address, length - 1)], + [(address + 1, length)], + [(address, length + 1)], + [(address, length), (address, 1)], + ): + with self.subTest(ranges=ranges), self.assertRaises(RuntimeError): + audit_registered_host(tensor, ranges) + with patch(target, return_value=dict(proof, node_pages={0: 2})): + with self.assertRaisesRegex(RuntimeError, "changed during CUDA"): + audit_registered_host(tensor, [(address, length)]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_pp_prefetch_ticket.py b/test/registered/unit/mem_cache/test_pp_prefetch_ticket.py index 4bd1669d6361..dc2bfc3a8b7e 100644 --- a/test/registered/unit/mem_cache/test_pp_prefetch_ticket.py +++ b/test/registered/unit/mem_cache/test_pp_prefetch_ticket.py @@ -6,6 +6,7 @@ from array import array from contextlib import nullcontext from queue import Empty, Queue +from types import SimpleNamespace from unittest.mock import Mock, call, patch import torch @@ -30,11 +31,32 @@ from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache from sglang.srt.mem_cache.utils import get_storage_hash_str from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=3, suite="base-a-test-cpu") -class TestPPPrefetchTicket(unittest.TestCase): +class TestPPPrefetchTicket(CustomTestCase): + def test_decode_storage_query_delivers_hint_to_controller(self): + hint = {"protocol_version": "0.1", "message_id": "decode-query"} + cache = UnifiedRadixCache.__new__(UnifiedRadixCache) + cache.tree_core = Mock(enable_storage=True, is_eagle=False) + cache.page_size = cache.prefetch_threshold = 4 + + def query(operation): + # Observable hit length changes only if the request envelope + # survives UnifiedRadixCache -> PrefetchOperation -> controller. + return [], 8 if operation.kv_hints == hint else 0 + + cache.cache_controller = SimpleNamespace( + prefetch_rate_limited=lambda: False, _storage_hit_query=query + ) + cache._all_reduce_attn_groups = Mock() + self.assertEqual( + cache.query_storage_hit_length(0, list(range(8)), kv_hints=hint), 8 + ) + self.assertEqual(cache.query_storage_hit_length(0, list(range(8))), 0) + def setUp(self): self.c = c = HybridCacheController.__new__(HybridCacheController) c.page_size = c.prefetch_threshold = 4 @@ -247,9 +269,9 @@ def test_allocation_error_preserves_ack_sequence_and_next_ticket(self): c.mem_pool_host.free.assert_called_once_with(kv, pool=PoolName.KV) c.page_get_func = Mock(return_value=1) c.storage_backend = Mock() + c.storage_backend.prefetch_batch_pages = 1 c.storage_backend.batch_get_v2.return_value = {"swa": [True]} with ( - patch("sglang.srt.managers.cache_controller.STORAGE_BATCH_SIZE", 1), patch.object( c.storage_stop_event, "is_set", side_effect=[False, False, True] ), diff --git a/test/registered/unit/server_args/test_python_thread_switch_interval.py b/test/registered/unit/server_args/test_python_thread_switch_interval.py new file mode 100644 index 000000000000..5408104da359 --- /dev/null +++ b/test/registered/unit/server_args/test_python_thread_switch_interval.py @@ -0,0 +1,109 @@ +"""Thread scheduling is an explicit seed propagated to the scheduler child.""" + +import argparse +import hashlib +import json +import multiprocessing +import os +import sys +import unittest + +from sglang.srt.managers.scheduler_components.thread_scheduling import ( + configure_python_thread_scheduling, +) +from sglang.srt.runtime_context import get_device, publish, reset_context +from sglang.srt.server_args import ServerArgs +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=12, suite="base-a-test-cpu") + + +def _scheduler_child(seed, connection): + try: + publish(seed, role="scheduler") + configure_python_thread_scheduling(get_device().python_thread_switch_interval) + connection.send((os.getpid(), sys.getswitchinterval())) + finally: + connection.close() + + +class TestPythonThreadSwitchIntervalSeed(CustomTestCase): + def setUp(self): + reset_context() + self.addCleanup(reset_context) + + def test_unset_remains_none(self): + seed = ServerArgs(model_path="dummy") + publish(seed, role="test") + self.assertIsNone(seed.python_thread_switch_interval) + self.assertIsNone(get_device().python_thread_switch_interval) + + def test_existing_positional_field_order_is_preserved(self): + fields = ServerArgs.__struct_fields__ + index = fields.index("sleep_on_idle") + self.assertEqual(fields[index + 1], "use_ray") + self.assertEqual(fields[-1], "python_thread_switch_interval") + legacy = json.dumps(list(fields[:-1]), separators=(",", ":")).encode() + self.assertEqual(len(fields[:-1]), 527) + self.assertEqual( + hashlib.sha256(legacy).hexdigest(), + "7ee1f836a5af82275fc8fb37e9dd74903e956d8aef6ac786a51053a2bc5cfae9", + ) + + def test_cli_and_published_namespace_preserve_explicit_value(self): + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + seed = ServerArgs.from_cli_args( + parser.parse_args( + ["--model", "dummy", "--python-thread-switch-interval", "0.001"] + ) + ) + publish(seed, role="test") + self.assertEqual(seed.python_thread_switch_interval, 0.001) + self.assertEqual(get_device().python_thread_switch_interval, 0.001) + + def test_argument_validation_rejects_invalid_intervals(self): + for value in (0.0, -0.001, float("nan"), float("inf")): + with ( + self.subTest(value=value), + self.assertRaisesRegex(ValueError, "positive and finite"), + ): + publish( + ServerArgs(model_path="dummy", python_thread_switch_interval=value), + role="test", + ) + reset_context() + + def test_spawned_child_uses_seed_not_parent_interpreter(self): + old = sys.getswitchinterval() + self.addCleanup(sys.setswitchinterval, old) + sys.setswitchinterval(0.003) + context = multiprocessing.get_context("spawn") + receiver, sender = context.Pipe(duplex=False) + process = context.Process( + target=_scheduler_child, + args=( + ServerArgs(model_path="dummy", python_thread_switch_interval=0.001), + sender, + ), + ) + process.start() + sender.close() + try: + self.assertTrue(receiver.poll(45), "scheduler child did not report") + pid, interval = receiver.recv() + process.join(10) + self.assertEqual(process.exitcode, 0) + self.assertNotEqual(pid, os.getpid()) + self.assertAlmostEqual(interval, 0.001, places=8) + self.assertAlmostEqual(sys.getswitchinterval(), 0.003, places=8) + finally: + receiver.close() + if process.is_alive(): + process.terminate() + process.join(10) + + +if __name__ == "__main__": + unittest.main()