diff --git a/python/sglang/srt/arg_groups/fields/memory.py b/python/sglang/srt/arg_groups/fields/memory.py index be8d204e6..c0bd80d6e 100644 --- a/python/sglang/srt/arg_groups/fields/memory.py +++ b/python/sglang/srt/arg_groups/fields/memory.py @@ -153,7 +153,7 @@ 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. 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, npu_memcache, hf3fs, nixl, aibrix, tensorcast. 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", @@ -167,6 +167,7 @@ class Memory(msgspec.Struct): "simm", "mori", "shm", + "tensorcast", ], ), ] = None diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index c28184141..c652c1b09 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -601,6 +601,7 @@ class HiCacheController: "nixl", "simm", "mori", + "tensorcast", ] ) or ( self.storage_backend_type == "dynamic" diff --git a/python/sglang/srt/mem_cache/pool_host/common.py b/python/sglang/srt/mem_cache/pool_host/common.py index 183e76d20..5728dbde5 100644 --- a/python/sglang/srt/mem_cache/pool_host/common.py +++ b/python/sglang/srt/mem_cache/pool_host/common.py @@ -103,6 +103,18 @@ def get_allocator_from_storage(allocator_type): return HostTensorAllocator() elif allocator_type == "shm": return ShmHostTensorAllocator() + elif allocator_type == "tensorcast": + try: + from sglang.srt.mem_cache.storage.tensorcast_store.host_allocator import ( + get_tensorcast_host_allocator_from_runtime, + ) + + return get_tensorcast_host_allocator_from_runtime() + except ImportError: + logger.warning( + "TensorCast's tensor allocator requires tensorcast >= 0.1.1. Please install TensorCast by 'pip install tensorcast' or build from source by following https://tensorcast.ai/development/build-from-source/. Fallback to use default allocator" + ) + return HostTensorAllocator() else: return HostTensorAllocator() diff --git a/python/sglang/srt/mem_cache/storage/backend_factory.py b/python/sglang/srt/mem_cache/storage/backend_factory.py index 42c92a3f0..c55ae3980 100644 --- a/python/sglang/srt/mem_cache/storage/backend_factory.py +++ b/python/sglang/srt/mem_cache/storage/backend_factory.py @@ -192,6 +192,8 @@ class StorageBackendFactory: return backend_class(storage_config, mem_pool_host) elif backend_name == "shm": return backend_class(storage_config, mem_pool_host) + elif backend_name == "tensorcast": + return backend_class(storage_config) else: raise ValueError(f"Unknown built-in backend: {backend_name}") @@ -258,3 +260,9 @@ StorageBackendFactory.register_backend( "sglang.srt.mem_cache.storage.shm", "HiCacheShm", ) + +StorageBackendFactory.register_backend( + "tensorcast", + "sglang.srt.mem_cache.storage.tensorcast_store.tensorcast_store", + "TensorcastStore", +) diff --git a/python/sglang/srt/mem_cache/storage/tensorcast_store/README.md b/python/sglang/srt/mem_cache/storage/tensorcast_store/README.md new file mode 100644 index 000000000..25ff3758f --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/tensorcast_store/README.md @@ -0,0 +1,368 @@ +# TensorCast as an L3 KV Cache + +This document describes how to use TensorCast as the L3 storage backend for +SGLang HiCache. The initial integration targets the Unified Radix Cache FULL +pool and uses TensorCast's public process-scoped +`RegionBackedArtifactSession`; SGLang does not construct daemon requests, +region layouts, or canonical TensorCast artifact IDs. + +Related documentation: + +- [TensorCast project](https://tensorcast.ai) +- [TensorCast repository](https://github.com/tensorcast-ai/tensorcast) +- [HiCache system design](https://docs.sglang.io/advanced_features/hicache_design.html) + +## About TensorCast + +TensorCast manages model state, KV caches, checkpoints, and other tensor state +as distributed artifacts. It separates cluster-wide discovery and routing from +host-local memory and data transfer: + +- One **Global Store** manages artifact metadata, replica state, and routing. +- One **StoreDaemon** on each serving host owns local memory regions and moves + artifact bytes locally or between hosts. +- Each SGLang rank attaches one process-scoped Session to its node-local + StoreDaemon. The Session owns that rank's regions, RPC health, and transfer + protocol details. + +SGLang HiCache continues to decide when pages are published and prefetched. +TensorCast stores and moves the K/V fragments that make up those pages. + +### Transfer modes + +Both modes use the same artifact identity and HiCache operation path. + +1. **Allocator-backed direct mode** is the default and recommended mode. + - The TensorCast Session allocates the CPU tensors used to back the standard + SGLang HostPool. + - L2/L3 operations submit spans in those tensors directly, without an + SGLang-side staging copy. + - Separate HostPool allocations, including asymmetric MHA K and V tensors, + become separate Session-managed regions automatically. +2. **Scratch mode** is the compatibility mode. + - The standard SGLang allocator owns the HostPool tensors. + - The Session lazily creates one fixed-capacity get arena and one + fixed-capacity put arena and copies between them and the HostPool. + - Operators must size `scratch.capacity_bytes` for the largest storage + batch. Registration fails at startup if the configured capacity is too + small. + +## Requirements and Initial Scope + +The initial integration requires: + +- Linux and SGLang's CUDA backend; +- Unified Radix Cache with the primary FULL KV pool; +- `--hicache-host-memory-mode cache`; +- either `page_first` with the `kernel` I/O backend or + `page_first_direct` with the `direct` I/O backend; +- a Global Store and node-local StoreDaemon started and ready before SGLang; +- `engine.cpu_shared_memory.enabled: true` in the StoreDaemon config; and +- a daemon-visible SGLang owner PID. Containers must share a PID namespace + or provide an equivalent PID visibility arrangement. + +The FULL path supports standard MHA, asymmetric MHA with separate K/V HostPool +tensors, and MLA. Currently TensorCast L3 does not support `buffer_only`, non-CUDA +platforms, `layer_first`, `page_head`, split-head layouts, DCP/attention CP, +packed MTP draft pools, or storage-v2 sidecar transfers such as Mamba and SWA. +There is no automatic fallback from allocator mode to scratch mode. + +## Install TensorCast + +Install an ABI-compatible TensorCast SDK and daemon: + +```bash +pip install tensorcast +``` + +If the published wheel's Torch or CUDA build does not match the SGLang +environment, build TensorCast from source instead. See the +[TensorCast build guide](https://github.com/tensorcast-ai/tensorcast/blob/main/docs/development/build-from-source.md). + +TensorCast is an optional SGLang dependency. Importing SGLang does not import +TensorCast unless this storage backend is selected. + +## Deployment + +The recommended deployment is operator-managed services plus SDK attachment: + +- The operator starts one Global Store for the TensorCast cluster. +- The operator starts one StoreDaemon on every SGLang host and waits for it to + become ready. +- SGLang ranks attach to the local daemon during HostPool construction. +- SGLang never starts, restarts, supervises, or stops TensorCast services. + +### Single-host deployment + +Run the following commands from the SGLang repository root. The checked-in +files under `configs/` are starter configurations and must be sized and secured +for the target deployment. + +**Step 1: Prepare the environment** + +```bash +export TC_CONFIG_DIR="$PWD/python/sglang/srt/mem_cache/storage/tensorcast_store/configs" +export TC_GLOBAL_SESSION=sglang-tensorcast-global +export TC_DAEMON_SESSION=sglang-tensorcast-daemon + +# Large KV publishes can retain more than 1,024 artifact memfds. +ulimit -n 65535 +``` + +**Step 2: Start the Global Store** + +```bash +tensorcast-cli global start \ + --config "${TC_CONFIG_DIR}/global_store_config.yaml" \ + --gs-session "${TC_GLOBAL_SESSION}" +``` + +The checked-in config listens on port `50051`. + +**Step 3: Start the StoreDaemon** + +```bash +tensorcast-cli daemon start \ + --config "${TC_CONFIG_DIR}/store_daemon_config.yaml" \ + --global-store-mode connect \ + --global-store-address 127.0.0.1:50051 \ + --session "${TC_DAEMON_SESSION}" +``` + +The checked-in daemon config listens for SDK RPCs on port `50052` and uses +port `65090` for P2P transfers. + +Its capability-token secret is fixed test material for local validation only. +A production copy must replace +`capability_tokens.active.secret` with independently generated secret material. +The Global Store and every daemon in a deployment must also use a consistent, +deployment-specific cluster identity. + +**Step 4: Verify service readiness** + +```bash +tensorcast-cli global status --gs-session "${TC_GLOBAL_SESSION}" +tensorcast-cli daemon status --session "${TC_DAEMON_SESSION}" +``` + +Do not start SGLang until both commands succeed. Session attachment additionally +checks that the daemon is ready, CPU shared memory is enabled, its local handle +service is reachable, and the endpoint is node-local. + +**Step 5: Start SGLang in allocator mode** + +Allocator mode is selected when `transfer_mode` is omitted: + +```bash +python -m sglang.launch_server \ + --model-path \ + --enable-hierarchical-cache \ + --hicache-host-memory-mode cache \ + --hicache-ratio 1 \ + --hicache-mem-layout page_first_direct \ + --hicache-io-backend direct \ + --hicache-write-policy write_through \ + --hicache-storage-backend tensorcast \ + --hicache-storage-backend-extra-config '{ + "tensorcast": { + "daemon_address": "127.0.0.1:50052" + } + }' +``` + +`page_first` plus `kernel` is also supported. Transfer mode controls L2/L3 +movement; it does not silently select or change the HostPool layout or L1/L2 +I/O backend. + +**Scratch-mode alternative** + +```bash +python -m sglang.launch_server \ + --model-path \ + --enable-hierarchical-cache \ + --hicache-host-memory-mode cache \ + --hicache-ratio 1 \ + --hicache-mem-layout page_first \ + --hicache-io-backend kernel \ + --hicache-write-policy write_through \ + --hicache-storage-backend tensorcast \ + --hicache-storage-backend-extra-config '{ + "tensorcast": { + "daemon_address": "127.0.0.1:50052", + "transfer_mode": "scratch", + "scratch": { + "capacity_bytes": 4294967296 + } + } + }' +``` + +The 4 GiB value is only an example. Use the sizing rule below for the selected +model, rank topology, dtype, page size, and HostPool capacity. + +**Step 6: Shut down in ownership order** + +First terminate SGLang normally and allow its storage workers to join. Then +stop the StoreDaemon before the Global Store: + +```bash +tensorcast-cli daemon stop --session "${TC_DAEMON_SESSION}" +tensorcast-cli global stop --gs-session "${TC_GLOBAL_SESSION}" +``` + +Do not stop the daemon while attached SGLang ranks are still running. The +StoreDaemon reclaims a rank's process-pinned stable region backing after that +rank process exits. + +### Multi-host and multi-instance deployment + +One Global Store may serve many hosts. Each host runs a StoreDaemon, and all +SGLang ranks on that host attach to its node-local endpoint: + +```text + Global Store + (one per cluster) + | + +---------------+---------------+ + | | + StoreDaemon A <----- P2P -----> StoreDaemon B + / | \ / \ + rank 0 rank 1 rank N rank 0 rank N +``` + +Multiple ranks and multiple SGLang instances may attach to the same local +daemon. Every rank owns a distinct process Session and one or more distinct +regions. SGLang appends `rank{world_rank}of{world_size}` to the configured +Session and region prefixes, while TensorCast makes concrete region names +PID-unique. + +Instances reuse artifacts only when their artifact identity inputs agree, +including namespace, model ID, model version, FULL layout, dtype, page size, +and TP/PP rank topology. Each instance's rank 0 consumes the artifact shard +published by rank 0 of a compatible instance; ranks do not consume one +another's shards. + +For cross-host traffic, configure every daemon to register with the same Global +Store and ensure its advertised and P2P addresses are reachable. Enable and +tune RDMA in the daemon config only when the host fabric and drivers support +it; otherwise configure the supported TCP transport. + +## Configuration + +The full extra-config value is a JSON object, or an `@path` reference to a +JSON, YAML, or TOML file. TensorCast-specific fields must be nested under +`tensorcast`: + +```json +{ + "prefetch_threshold": 256, + "prefetch_timeout_base": 1.0, + "prefetch_timeout_per_ki_token": 0.25, + "hicache_storage_pass_prefix_keys": false, + "tensorcast": { + "daemon_address": "127.0.0.1:50052", + "namespace": "default", + "transfer_mode": "allocator", + "model_id": null, + "model_version": "unversioned", + "session_name_prefix": "sglang", + "region_name_prefix": "sglang_tensorcast", + "exists_timeout_s": 30.0, + "transfer_timeout_s": null, + "scratch": { + "capacity_bytes": 16777216 + } + } +} +``` + +### TensorCast fields + +| Field | Default | Meaning | +|---|---|---| +| `daemon_address` | required | StoreDaemon RPC endpoint passed unchanged to TensorCast, for example `127.0.0.1:50052`. | +| `namespace` | `default` | Byte-artifact namespace. Compatible publishers and consumers must use the same value. | +| `transfer_mode` | `allocator` | Selects `allocator` or `scratch` for the lifetime of the rank process. | +| `model_id` | SGLang storage model name | Optional stable model identity override. | +| `model_version` | `unversioned` | Immutable model/config revision. Production deployments should set a release- or content-specific value to prevent stale KV reuse after a model change. | +| `session_name_prefix` | `sglang` | Diagnostic Session prefix; SGLang appends the world-rank label. | +| `region_name_prefix` | `sglang_tensorcast` | Diagnostic region prefix; SGLang appends the world-rank label and TensorCast makes concrete names PID-unique. | +| `exists_timeout_s` | `30.0` | Positive metadata-only exists deadline in seconds. | +| `transfer_timeout_s` | `null` | Optional positive transfer deadline passed unchanged to TensorCast. Region get/put uses zero transparent retries. | +| `scratch.capacity_bytes` | `16777216` | Positive capacity in bytes for each scratch-direction arena. Ignored by allocator mode. | + +### Generic HiCache fields + +These optional fields remain at the top level of the extra config: + +| Field | Default | Meaning | +|---|---|---| +| `prefetch_threshold` | `256` | Minimum prefix length in tokens before storage prefetch is attempted. | +| `prefetch_timeout_base` | `1.0` | Fixed portion of the Unified prefetch timeout in seconds. | +| `prefetch_timeout_per_ki_token` | `0.25` | Additional timeout in seconds per 1,024 tokens. | +| `hicache_storage_pass_prefix_keys` | `false` | Pass prefix keys to storage backends; TensorCast FULL v1 does not use them. | + +### Scratch capacity + +The exact minimum is known only after SGLang constructs and registers the +HostPool: + +```text +maximum_batch_pages = min(128, host_pool.page_num) +page_bytes = sum(bytes of all registered fragments in one logical page) +minimum_capacity_bytes = maximum_batch_pages * page_bytes +``` + +For MHA, `page_bytes` includes both K and V. For MLA, it contains the single +combined KV fragment. The configured value applies independently to the lazy +get and put arenas, so budget up to `2 * scratch.capacity_bytes` of additional +host memory after both directions have been used. + +The 16 MiB default is syntactically valid but may be too small for a real +model. An insufficient value fails rank startup with the configured and +required byte counts plus the calculated page geometry. Runtime scratch growth +is not supported. + +### Supported layout matrix + +| Host-memory mode | Transfer mode | Host layout | I/O backend | Result | +|---|---|---|---|---| +| `cache` | `allocator` | `page_first` | `kernel` | Supported; allocator is the default transfer mode. | +| `cache` | `allocator` | `page_first_direct` | `direct` | Supported and recommended for direct L1/L2 I/O. | +| `cache` | `scratch` | `page_first` | `kernel` | Supported with one copy at the L2/L3 boundary. | +| `cache` | `scratch` | `page_first_direct` | `direct` | Supported with one copy at the L2/L3 boundary. | +| `buffer_only` | either | any | any | Rejected. | +| `cache` | either | other layouts | any | Rejected initially. | + +## Runtime and Failure Semantics + +Session attachment is an early startup operation. A daemon readiness, +configuration, or allocator failure aborts rank startup; SGLang does not fall +back to another transfer mode. + +During a supported FULL-v1 exists/get/put operation, a fatal Session or adapter +exception permanently disables TensorCast L3 for that rank. Subsequent calls +report no storage hits or false page results, while the SGLang worker continues +running and may recompute the missing prefix locally. Other ranks retain their own +rank-local Session health. + +Transfer mode, daemon endpoint, and Session options cannot be changed at +runtime. A failed or terminated Session cannot reattach in the same rank +process. Restart the rank to establish a new Session. + +`TensorcastStore.close()` terminates Session admission but does not directly +release process-pinned allocator or scratch regions. Their mappings remain +valid for the HostPool lifetime, and the daemon reclaims stable backing only +after the SGLang rank exits. + +## Troubleshooting + +- **Attach fails before model startup:** verify Global Store and StoreDaemon + status, `cpu_shared_memory.enabled`, the configured daemon address, local + handle socket access, and PID visibility. +- **`Too many open files` from `memfd_create`:** raise the StoreDaemon's + inherited `RLIMIT_NOFILE`, for example with `ulimit -n 65535`, and restart + the failed daemon and rank processes. +- **Scratch capacity is insufficient:** use the exact required byte count in + the startup error and increase `scratch.capacity_bytes`; remember that get + and put may allocate two arenas of that size. diff --git a/python/sglang/srt/mem_cache/storage/tensorcast_store/configs/global_store_config.yaml b/python/sglang/srt/mem_cache/storage/tensorcast_store/configs/global_store_config.yaml new file mode 100644 index 000000000..9c42a507b --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/tensorcast_store/configs/global_store_config.yaml @@ -0,0 +1,93 @@ +# Config metadata for auditing and cluster safety. +meta: + schema_version: "v1" + description: "Example config file for TensorCast Global Store" + cluster_token: "" + +database: + db_file: null + +server: + listen: + host: 0.0.0.0 + port: 50051 + advertise: + host: "" + port: 0 + max_workers: 10 + grpc: + max_concurrent_streams: 0 + keepalive_time: 120s + keepalive_timeout: 0s + max_connection_idle: 0s + max_connection_age: 0s + tcp_nodelay: false + so_reuseport: false + tls: + enabled: false + cert_file: "" + key_file: "" + client_ca_file: "" + metrics_port: 18000 + +worker_policy: + heartbeat_timeout: 30s + cleanup_interval: 60s + default_heartbeat_interval: 5s + memory_tiers: + snapshot_retention: 600s + snapshot_max_rows: 200 + publish_interval: 5s + key_mapping: + alias_cache_ttl: 1s + transport_scheduler: + mode: TRANSPORT_SCHEDULER_MODE_GROUP_DISPATCH + source_balance_weights: + replica_load_weight: 1.0 + worker_load_weight: 1.0 + recent_assignment_penalty_weight: 1.0 + diffusion_bonus_weight: 1.0 + group_dispatch: + fairness_floor_ratio: 0.25 + completion_bias_weight: 1.0 + starvation_aging_threshold: 5s + queue_scan_limit: 128 + dispatch_batch_limit: 16 + +limits: + digest_writes: + max_leaf_writes_per_request: 16384 + max_proof_digests_per_request: 16384 + max_total_digests_per_request: 32768 + max_digest_bytes_per_request: 2097152 + operation_leases: + default_ttl: 30s + max_ttl: 300s + operation_writes: + min_status_update_interval: 1s + retention: + operations_ttl: 86400s + assembly_proof_commitments_ttl: 86400s + piece_proof_digests_ttl: 86400s + +observability: + logging: + level: INFO + file: "" + otel_context_enabled: false + sink_file: "" + vlog_level: 0 + otel: + enabled: false + exporter_otlp_endpoint: http://127.0.0.1:4317 + exporter_protocol: grpc + service_name: tensorcast-global-store + sampler: parentbased_traceidratio + sampler_arg: "1.0" + otel_cxx: + sdk_disabled: false + exporter_insecure: false + exporter_headers: {} + console_exporter: false + tracing: + chrome_trace_dir: "" diff --git a/python/sglang/srt/mem_cache/storage/tensorcast_store/configs/store_daemon_config.yaml b/python/sglang/srt/mem_cache/storage/tensorcast_store/configs/store_daemon_config.yaml new file mode 100644 index 000000000..ffe1e18b7 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/tensorcast_store/configs/store_daemon_config.yaml @@ -0,0 +1,389 @@ +# Config metadata for auditing and cluster safety. +meta: + # Schema version label for auditing; keep empty unless you enforce schema checks. + schema_version: "" + # Free-form description for operators and change tracking. + description: "Example config file for TensorCast Store Daemon" + # Cluster identity token to prevent split-brain when talking to Global Store. + cluster_token: "" + +# gRPC server settings. +server: + # Control-plane listen socket for the gRPC API. + listen: + # Bind host for gRPC; use 0.0.0.0 to expose all NICs or a specific interface to fence access. + host: 0.0.0.0 + # gRPC listen port; this is the port advertised to Global Store regardless of advertise.port. + port: 50052 + # Advertised address for Global Store registration; set when auto-detect might pick a non-routable IP. + advertise: + # Routable host that clients and Global Store should dial; avoid loopback and link-local. + # If empty or non-routable, the daemon resolves in order: listen.host (routable), + # route IP toward Global Store, then default interface IP. + host: null # automatically resolve + # Dedicated P2P listen socket; use a separate NIC/port to isolate data-plane traffic. + p2p_listen: + # Bind host for P2P transfers; often set to the RDMA or high-bandwidth interface. + host: 0.0.0.0 + # P2P port; must be non-zero when high_availability.enabled is true. + port: 65090 + # Optional default root for relative disk_path values. + # When empty, disk materialization accepts absolute disk_path only. + storage_path: "" + # StoreEngine I/O worker threads; tune for disk/P2P parallelism vs CPU contention. + num_threads: 16 + # gRPC server channel args and TLS. + grpc: + # Max concurrent HTTP/2 streams; 0 uses gRPC default. + max_concurrent_streams: 0 + # Keepalive ping interval; set 0s to disable periodic pings. + keepalive_time: 30s + # Keepalive timeout before terminating a dead connection. + keepalive_timeout: 10s + # Close idle connections after this duration; 0s disables. + max_connection_idle: 10m + # Hard connection max age; 0s disables. + max_connection_age: 0s + # Disable Nagle when true; set true for low-latency RPCs. + tcp_nodelay: false + # Enable SO_REUSEPORT for listener load balancing. + so_reuseport: false + tls: + # Enable TLS for the daemon gRPC server. + enabled: false + # PEM-encoded server certificate. + cert_file: "" + # PEM-encoded private key for cert_file. + key_file: "" + # Optional client CA bundle for mTLS; empty disables client cert checks. + client_ca_file: "" + +# StoreEngine and data-plane memory/transfer tuning. +engine: + # UMA chunk granularity; larger chunks reduce metadata but make eviction/leases coarser. + artifact_chunk_bytes: 256MB + # Default depth for StreamingPinnedBuffer instances created by StoreEngine. + streaming_buffer_chunks: 8 + # Enable external target verification for MaterializeIntoTarget (off by default). + enable_external_target_verification: false + # Enable memfd-backed UMA CPU allocations for zero-copy CPU tensor materialization. + cpu_shared_memory: + enabled: true + # Memory tiers define stable/preemptible behavior; tune for DRAM pressure and HA telemetry. + memory_tiers: + # When false, CPU chunks stay resident (no preemptible marking); use for stability or debugging. + enable_preemptible: false + # Stable budget in bytes; must fit in host DRAM minus total pinned pools (sum(pinned_memory.classes[].pool_bytes)) or startup fails. + # Sized for the single-GPU local validation host. This leaves room for the + # SGLang worker and its HiCache HostPool while retaining ample L3 capacity. + stable_bytes: 32GB + # Cap for preemptible pool; ignored when enable_preemptible=false (0 means no preemptible budget). + preemptible_limit_bytes: 0 + # Reclaim target ratio for preemptible pressure; lower values reclaim more aggressively. + preemptible_low_watermark_ratio: 0.4 + +# Promotion policy for P2P routing. +promotion: + # Promotion controls whether the daemon will "export" newly materialized + # replicas to become P2P-routable sources in the Global Store (i.e., publish + # `export_state=EXPORTABLE` + transport metadata). + # + # IMPORTANT: Promotion is only attempted when the *client request* asks for + # it via `export_policy` (SDK `GetArtifactOptions.export_policy`): + # - `never` (default): do not attempt promotion for this materialize. + # - `auto`: allow policy-driven promotion (if the daemon policy permits). + # - `force`: request an immediate promotion attempt (still gated by daemon policy). + # + # Policy values: + # - PROMOTION_POLICY_NEVER: + # Never promote; replicas remain presence-only and are not P2P targets. + # - PROMOTION_POLICY_ON_MATERIALIZE: + # Promote after any successful materialize when `export_policy` is `auto` or `force`. + # - PROMOTION_POLICY_ON_HOTNESS: + # Promote only when `export_policy=force` (conservative default; treat "force" as + # the caller's "this is hot" signal; `auto` does not promote). + # - PROMOTION_POLICY_ON_POLICY: + # Currently behaves like ON_MATERIALIZE (promote on `auto`/`force`); reserved for + # future policy-based gating beyond the request signal. + policy: PROMOTION_POLICY_ON_MATERIALIZE + # Require verification metadata before exporting. + require_verified: false + # Bounded drain timeout before forced demotion. + demotion_drain_timeout: 30s + +# Unified pinned memory configuration for all daemon-side pinned usage. +pinned_memory: + # Phase 1 fixed-allocation: total pinned bytes is derived as sum(classes[].pool_bytes). + # Startup fail-fast: before allocating pinned pools, the daemon checks current available memory and + # exits early if insufficient. Required bytes = pinned_total + engine.memory_tiers.stable_bytes + headroom, + # where headroom = min(10% * (pinned_total + stable_bytes), 10GiB). + # Deadline for pinned slice acquisition under pressure. + allocation_timeout: 30s + classes: + - name: engine + slice_bytes: 256MB + # 1 GPU * streaming_buffer_chunks(8) = 8 slices required at startup. + # 8 * 256MB = 2GB minimum engine pinned pool. + pool_bytes: 2GB + - name: comm_gpu + slice_bytes: 16MB + # 3GB covers the configured 16 buffers/flow, 8 TCP connections, and + # expected_gpu_channels=8 startup sizing on the local validation host. + pool_bytes: 3GB + rdma_preregister: true + - name: comm_cpu + slice_bytes: 16MB + # 16 buffers/flow require at least 16 slices (256MB). + pool_bytes: 512MB + rdma_preregister: true + +# Post-seal policies for assembly -> mi2 alias resolution. +post_seal: + # When true, migrate selected assembly views under the sealed mi2_id. + migrate_views: false + # When true, only migrate views that require transforms (e.g., transpose). + migrate_transpose_only: false + # When true, allow safe reuse of CGID-scoped view replicas after sealing. + reuse_views_if_safe: false + # When true, retire CGID-scoped piece replicas after sealing. + retire_pieces: false + +# Capability token keys (daemon-issued capability envelope). +capability_tokens: + active: + # Token version id; rotate by incrementing version and moving old key to previous. + version: 1 + # Secret bytes (base64-encoded in YAML/JSON). This is a dummy non-production key as an example; production deployments + # must replace it with independently generated secret material. + secret: "c2dsYW5nLXRlbnNvcmNhc3QtbG9jYWwtdGVzdC1rZXk=" + previous: [] + +# Retention handle settings (control-plane only; requires capability_tokens). +retention_handles: + # When true, daemon issues retention handles for local stable DRAM. + enabled: false + # Default TTL when caller does not specify ttl_ms. + default_ttl: 10m + # Maximum allowed TTL for any handle (issuer clamp). + max_ttl: 24h + +# Capability directory publishing (Global Store capability flags). +capability_directory: + enabled: true + +# Global Store registration and heartbeat coordination. +high_availability: + # When true, the daemon registers with Global Store and emits heartbeats/sync; requires p2p_listen.port. + enabled: true + # Global Store endpoints; keep the first entry as the primary registration target. + global_store_endpoints: + - + # Host for Global Store gRPC; must be reachable from the daemon. + host: 127.0.0.1 + # Port for Global Store gRPC. + port: 50051 + # Heartbeat cadence; keep comfortably below Global Store worker_policy.heartbeat_timeout. + heartbeat_interval: 10s + # Chunk inventory sync cadence; set 0s to disable sync thread. + periodic_sync_interval: 30s + # Retry attempts for HA registration/requests; currently uses GlobalStoreClient defaults. + max_retries: 3 + # Backoff between registration retries; currently not wired in the daemon. + registration_retry_delay: 500ms +# Background cleanup, eviction, and TTL scheduling. +lifecycle: + # GPU memory fraction that triggers eviction when periodic eviction is enabled; lower values evict earlier. + gpu_memory_limit_fraction: 0.75 + # Enable background eviction loop; disable if you want strictly demand-driven eviction. + enable_periodic_eviction: false + # Eviction poll cadence; smaller improves responsiveness but increases background work. + eviction_loop_interval: 1s + # Legacy alias for eviction_loop_interval; 0s keeps eviction_loop_interval in effect. + eviction_check_interval: 0s + # PID liveness polling interval used when pidfd is unavailable; shorter detects dead clients faster. + proc_check_interval: 5s + # Sweep interval for session TTL cleanup; keep below sessions_ttl for timely reclamation. + sessions_sweep_interval: 10s + # Sweep interval for transport lock expiry; keep below locks_ttl to avoid stale locks. + locks_sweep_interval: 10s + # Sweep interval for verification tracker cleanup/timeouts; smaller reduces tail latency on failures. + verification_sweep_interval: 500ms + # Session TTL; must exceed client keepalive cadence to avoid premature eviction. + sessions_ttl: 60s + # Transport lock TTL; long enough for worst-case transfers but short enough to recover from leaks. + locks_ttl: 120s + # Local-only Unix domain socket for handle FD handoff and lease release. + handle_leases: + # Optional override. Leave empty to let the daemon auto-select + # /local_handle.sock for same-pod/local SDKs. + # daemon_state_dir defaults to: + # $TENSORCAST_HOME/hosts//sessions//session + # (or ~/.tensorcast/hosts//sessions//session) + # Auto-discovery relies on TENSORCAST_INSTANCE (CLI/SDK-launched daemon). + # If TENSORCAST_INSTANCE is not set, the daemon falls back to: + # $TENSORCAST_HOME/hosts//runtime/daemons//local_handle.sock + # If the selected socket path exceeds AF_UNIX limits, the daemon falls back to: + # $TENSORCAST_HOME/uds/lh-.sock + # Set explicitly when daemon and client SDK run in different pods and need a shared path. + # Socket path must have a daemon-owned parent directory (not world-writable); create the directory beforehand. + local_handle_socket_path: "" + # Handle lease TTL as a crash-safety fallback; live clients release leases when tensors are freed. + # Set to 0s to disable TTL (leases rely on explicit release and PID-exit cleanup). + ttl: 0s + # Best-effort guardrail: limit lease-bearing handle mints per second (0 => unlimited). + max_mints_per_second: 0 + +# Communicator transport configuration (TCP/RDMA). +communicator: + # Enable RDMA transport; requires RDMA-capable NICs and correct kernel/driver setup. + enable_rdma: True + # Example tuning block (defaults shown below); adjust per workload. + stager: + # Stage CPU tensors for RDMA; disable only if you rely on direct GPU RDMA paths. + stage_cpu_for_rdma: true + # Buffers per flow; more buffers increase overlap but require a larger pool. + buffers_per_flow: 16 + # Cap concurrent in-flight segments; 0 inherits buffers_per_flow. + max_window_segments: 0 + # Cap concurrent GPU channels so pool sizing is predictable (8 for 8-GPU defaults; 0 = auto). + expected_gpu_channels: 8 + rdma: + # Outstanding WRs per QP; increase for bandwidth, decrease to reduce CPU/latency. + outstanding_wr: 64 + # ACK TTL; increase on high-latency fabrics to avoid premature cleanup. + ack_ttl_ms: 30000 + # QP tuning; adjust only if you understand fabric QoS/timeouts. + traffic_class: 186 + qp_timeout: 20 + qp_retry: 7 + # Multi-QP configuration for high-throughput scenarios + qp_count: 1 # Number of QPs per transport (1-16, default: 1) + bonding_balance: false # Enable LAG port balancing (default: false) + # Reuse preregistered stable local MRs for read_plan target windows when available. + # Disable only for A/B or debugging request-scoped MR registration behavior. + enable_stable_local_mr_reuse: true + # Slot-aligned MR reuse chunk size for stable local backing. + stable_local_mr_reuse_chunk_slots: 32 + # Async all-rail chunk prewarm. Unset or 0 disables; values >0 enable one + # prewarm job per visible rail on a bounded communicator-local worker pool. + stable_local_mr_reuse_prewarm_workers: 1 + transport: + # TCP connections per peer; more lanes help bandwidth but add CPU overhead. + tcp_conn_count: 8 + # Connect timeout; increase when handshakes are slow or firewalled. + connect_timeout_sec: 10 + # TOS/DSCP value; keep 0 unless fabric QoS is configured. + tcp_tos: 0 + # Allow SO_REUSEPORT on listeners; disable if OS load balancing is undesirable. + so_reuseport: true + affinity: + # Reserved for future CPU pinning policies. + enable: false + simple_numa: + # NIC/GPU locality mapping for multi-socket nodes. + enable: false + # Example topology mapping: + # nodes: + # - id: 0 + # nics: ["mlx5_0", "mlx5_1"] + # gpus: [0, 1] + # is_default: true + nodes: [] + +# Logging and tracing configuration (best-effort). +observability: + # Logging output and verbosity controls. + logging: + # Minimum severity emitted; higher levels reduce noise at the cost of debugging detail. + level: INFO + # Optional plain-text logfile; empty logs only to stderr. + file: "" + # When true, include trace/span IDs in logs (if tracing is enabled). + otel_context_enabled: false + # Optional enriched sink file with trace/span IDs; leave empty to disable. + sink_file: "" + # VLOG verbosity level; 0 disables VLOG. + vlog_level: 0 + # OpenTelemetry export settings. + otel: + # Enable OTLP export; when false, other OTel settings are ignored. + enabled: false + # Collector endpoint; must be reachable from the daemon host. + exporter_otlp_endpoint: http://127.0.0.1:4317 + # Transport protocol for OTLP; aliases like "grpc" or "http/protobuf" are normalized. + exporter_protocol: grpc + # Stable logical service name for grouping traces/metrics across nodes. + service_name: tensorcast-store-daemon + # Sampler name for Python OTel SDK; C++ currently ignores this field. + sampler: parentbased_traceidratio + # Sampler argument (ratio for traceidratio); keep 1.0 for full sampling. + sampler_arg: "1.0" + otel_cxx: + # Disable the C++ SDK when embedding in C++ processes. + sdk_disabled: false + # Force insecure OTLP gRPC; set true when using non-TLS collectors. + exporter_insecure: false + # Static OTLP headers (e.g., auth tokens). + exporter_headers: {} + # Emit spans to stdout for debugging. + console_exporter: false + tracing: + # Chrome trace output directory; empty disables. + chrome_trace_dir: "" + +# Compatibility toggles for legacy behavior. +compatibility: + # Require explicit disk_path for confirm RPCs; reserved for legacy clients. + confirm_requires_disk_path: false + # Status to return on verification timeouts; set OK or DEADLINE when wired. + verification_timeout_status: VERIFICATION_TIMEOUT_STATUS_UNSPECIFIED + # Evict replicas on dead PID detection; reserved until wired. + evict_on_dead_pid: false + # Auto-register disk loads; reserved until wired. + auto_register_disk_loads: false + # Force full-digest verification on load; reserved until wired. + force_full_digest_on_load: false + +# Debug-only toggles; keep disabled in production unless explicitly needed. +debug: + # CUDA debug toggles used for tests and local runs. + cuda: + # Same-process CUDA IPC fallback for tests; breaks real IPC semantics across processes. + enable_same_process_ipc_fallback: false + +# Byte-artifact routing and payload transport defaults. Primarily used for distributed KV store +byte_artifact_routing: + # Number of routing shards used for authority fanout. + # Recommended: 2N shards for N TensorCast daemons in the cluster + shard_count: 8 + # How long route resolution can be reused before forcing a refresh. + route_staleness_budget: 5s + # Lease TTL for published byte-artifact routes. + lease_ttl: 30s + # Keepalive cadence for refreshing active route leases. + keepalive_interval: 10s + # Worker-directory cache freshness budget used before refetching active workers. + worker_directory_staleness_budget: 30s + payload_transport: + # Per-transport chunk size limit used when segmenting payloads. + max_chunk_bytes: 1048576 + # Batch transport protocol revision. + # v2 uses P2P communicator, v1 uses gRPC developed at early ages. Use v2 + batch_transport_protocol_version: 2 + # Use communicator/materialize source export paths when available. + communicator_source_enabled: true + # Allow host-memory export for direct-write-capable CPU paths. + host_memory_export_enabled: true + # 0 means no payload-size cap + max_batch_payload_bytes: 0 + # 0 means no per-batch item-count cap + max_batch_items: 0 + source_publish_prereg: + # Pre-register publish-time exports so later batch_get can reuse them directly. + enabled: true + # Keep preregistered exports alive long enough for the delayed consumer phase. + ttl: 60s + # Cap live preregistered export entries to avoid unbounded retention. + max_live_entries: 4096 + # Cap total bytes retained by live preregistered exports. + max_live_bytes: 17179869184 diff --git a/python/sglang/srt/mem_cache/storage/tensorcast_store/host_allocator.py b/python/sglang/srt/mem_cache/storage/tensorcast_store/host_allocator.py new file mode 100644 index 000000000..57a5e0633 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/tensorcast_store/host_allocator.py @@ -0,0 +1,476 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to SGLang project + +"""Configuration and host-allocation boundary for TensorCast HiCache.""" + +from __future__ import annotations + +import json +import sys +import threading +from collections.abc import Mapping +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import TYPE_CHECKING, Annotated, TypeAlias, cast + +import torch +from pydantic import BaseModel, ConfigDict, Field, field_validator + +from sglang.srt.mem_cache.pool_host.common import HostTensorAllocator + +if TYPE_CHECKING: + from tensorcast.api.store import ( + RegionBackedArtifactSession, + RegionBackedArtifactSessionOptions, + ) + + +_SUPPORTED_LAYOUT_IO_PAIRS = frozenset( + { + ("page_first", "kernel"), + ("page_first_direct", "direct"), + } +) + + +class TensorcastTransferMode(str, Enum): + """SGLang's supported TensorCast host-transfer modes. + allocator (default): TensorCast daemon allocates and owns a shared memory slab + with direct RDMA + scratch: uses SGlang-owned common HostTensorAllocator. TensorCast uses a scratch memory slab as staging buffer to copy KV pages into it. More compatible with performance overhead. + """ + + ALLOCATOR = "allocator" + SCRATCH = "scratch" + + +class _FrozenConfig(BaseModel): + model_config = ConfigDict( + extra="forbid", + frozen=True, + validate_default=True, + ) + + +class TensorcastScratchConfig(_FrozenConfig): + """Fixed per-direction scratch-arena configuration.""" + + capacity_bytes: Annotated[int, Field(strict=True, gt=0)] = 16 * 1024 * 1024 + + +class TensorcastConfig(_FrozenConfig): + """Validated TensorCast-specific HiCache configuration.""" + + daemon_address: Annotated[str, Field(strict=True)] + namespace: Annotated[str, Field(strict=True)] = "default" + transfer_mode: TensorcastTransferMode = TensorcastTransferMode.ALLOCATOR + model_id: Annotated[str, Field(strict=True)] | None = None + model_version: Annotated[str, Field(strict=True)] = "unversioned" + session_name_prefix: Annotated[str, Field(strict=True)] = "sglang" + region_name_prefix: Annotated[str, Field(strict=True)] = "sglang_tensorcast" + exists_timeout_s: Annotated[ + float, + Field(gt=0.0, allow_inf_nan=False), + ] = 30.0 + transfer_timeout_s: ( + Annotated[ + float, + Field(gt=0.0, allow_inf_nan=False), + ] + | None + ) = None + scratch: TensorcastScratchConfig = TensorcastScratchConfig() + + @field_validator( + "namespace", + "model_version", + "session_name_prefix", + "region_name_prefix", + ) + @classmethod + def validate_required_string(cls, value: str) -> str: + return _strip_non_empty(value) + + @field_validator("model_id") + @classmethod + def validate_optional_model_id(cls, value: str | None) -> str | None: + if value is None: + return None + return _strip_non_empty(value) + + +TensorcastConfigSource: TypeAlias = TensorcastConfig | str | Mapping[str, object] + + +class TensorcastSessionRegistryError(RuntimeError): + """Raised when process-local TensorCast Session admission is invalid.""" + + +@dataclass(slots=True) +class _TensorcastSessionRegistryState: + session: RegionBackedArtifactSession | None = None + session_options: RegionBackedArtifactSessionOptions | None = None + early_attach_completed: bool = False + active_store_owner: object | None = None + terminal: bool = False + + +_PROCESS_SESSION_REGISTRY_LOCK = threading.Lock() +_PROCESS_SESSION_REGISTRY = _TensorcastSessionRegistryState() + + +def normalize_tensorcast_config(source: TensorcastConfigSource) -> TensorcastConfig: + """Normalize raw or controller-parsed HiCache extra configuration.""" + + if isinstance(source, TensorcastConfig): + return source + + raw = _load_extra_config(source) if isinstance(source, str) else source + if "tensorcast" not in raw: + raise ValueError( + "TensorCast HiCache configuration requires a 'tensorcast' object" + ) + + return TensorcastConfig.model_validate(raw["tensorcast"]) + + +def format_tensorcast_rank_label(world_rank: int, world_size: int) -> str: + """Return the stable diagnostic label for one SGLang rank.""" + + if world_size <= 0: + raise ValueError(f"world_size must be positive, got {world_size}") + if world_rank < 0 or world_rank >= world_size: + raise ValueError( + "world_rank must be in [0, world_size), got " + f"world_rank={world_rank}, world_size={world_size}" + ) + return f"rank{world_rank}of{world_size}" + + +def build_tensorcast_session_options( + config: TensorcastConfig, + *, + world_rank: int | None = None, + world_size: int | None = None, +) -> RegionBackedArtifactSessionOptions: + """Build public TensorCast Session options without exposing SDK internals.""" + + try: + from tensorcast.api.store import ( + AllocatorTransferOptions, + RegionBackedArtifactSessionOptions, + ScratchTransferOptions, + ) + except ImportError as exc: + raise ImportError( + "The TensorCast HiCache backend requires the optional 'tensorcast' " + "Python package. Install TensorCast before selecting " + "--hicache-storage-backend tensorcast." + ) from exc + + if (world_rank is None) != (world_size is None): + raise ValueError("world_rank and world_size must be provided together") + if world_rank is None: + from sglang.srt.runtime_context import get_parallel + + parallel = get_parallel() + world_rank = parallel.world_rank + world_size = parallel.world_size + + rank_label = format_tensorcast_rank_label(world_rank, cast(int, world_size)) + if config.transfer_mode == TensorcastTransferMode.ALLOCATOR: + transfer = AllocatorTransferOptions() + else: + transfer = ScratchTransferOptions( + capacity_bytes=config.scratch.capacity_bytes, + ) + + return RegionBackedArtifactSessionOptions( + daemon_address=config.daemon_address, + session_name=f"{config.session_name_prefix}-{rank_label}", + transfer=transfer, + transfer_timeout_s=config.transfer_timeout_s, + exists_timeout_s=config.exists_timeout_s, + region_name_prefix=f"{config.region_name_prefix}-{rank_label}", + ) + + +def _attach_process_session( + options: RegionBackedArtifactSessionOptions, +) -> RegionBackedArtifactSession: + from tensorcast.api.store import RegionBackedArtifactSession + + return RegionBackedArtifactSession.attach(options) + + +def attach_early_process_session( + options: RegionBackedArtifactSessionOptions, +) -> RegionBackedArtifactSession: + """Attach once before HostPool allocation and retain the process Session.""" + + with _PROCESS_SESSION_REGISTRY_LOCK: + state = _PROCESS_SESSION_REGISTRY + if state.terminal: + raise TensorcastSessionRegistryError( + "the TensorCast process Session is terminal; restart the rank process" + ) + if state.early_attach_completed: + if options != state.session_options: + raise TensorcastSessionRegistryError( + "TensorCast early attach conflicts with the process Session options" + ) + if state.session is None: + raise RuntimeError( + "TensorCast Session registry is inconsistent after early attach" + ) + return state.session + + try: + session = _attach_process_session(options) + except BaseException: + state.terminal = True + raise + + state.session = session + state.session_options = options + state.early_attach_completed = True + return session + + +def claim_tensorcast_store_session( + options: RegionBackedArtifactSessionOptions, + *, + owner: object, +) -> RegionBackedArtifactSession: + """Give one Store ownership of the already attached process Session.""" + + if owner is None: + raise ValueError("TensorCast Store owner must not be None") + with _PROCESS_SESSION_REGISTRY_LOCK: + state = _PROCESS_SESSION_REGISTRY + if state.terminal: + raise TensorcastSessionRegistryError( + "the TensorCast process Session is terminal; restart the rank process" + ) + if not state.early_attach_completed or state.session is None: + raise TensorcastSessionRegistryError( + "TensorCast Store cannot attach for the first time; the HostPool " + "allocator must complete early Session attach" + ) + if options != state.session_options: + raise TensorcastSessionRegistryError( + "TensorCast Store options conflict with the early process Session" + ) + if state.active_store_owner is None: + state.active_store_owner = owner + elif state.active_store_owner is not owner: + raise TensorcastSessionRegistryError( + "another TensorCast Store already owns the process Session" + ) + return state.session + + +def terminate_tensorcast_store_session(*, owner: object) -> bool: + """Make Store admission terminal and terminate the Session exactly once.""" + + with _PROCESS_SESSION_REGISTRY_LOCK: + state = _PROCESS_SESSION_REGISTRY + if state.active_store_owner is not owner: + raise TensorcastSessionRegistryError( + "only the active TensorCast Store owner may terminate the process Session" + ) + if state.terminal: + return False + if state.session is None: + raise RuntimeError( + "TensorCast Session registry has an owner without an attached Session" + ) + state.terminal = True + session = state.session + + session.terminate_process_session() + return True + + +class TensorcastHostTensorAllocator(HostTensorAllocator): + """Delegate exact HostPool tensor allocations to one TensorCast Session.""" + + def __init__(self, session: RegionBackedArtifactSession) -> None: + super().__init__() + self._session = session + self._allocation_sequence = 0 + + @property + def session(self) -> RegionBackedArtifactSession: + return self._session + + def allocate( + self, + dims: tuple[int, ...], + dtype: torch.dtype, + device: str, + ) -> torch.Tensor: + if device != "cpu": + raise ValueError( + f"TensorCast host allocation requires CPU memory, got device={device!r}" + ) + self.dims = dims + self.dtype = dtype + self._allocation_sequence += 1 + return self._session.allocate_host_tensor( + dims, + dtype, + name=f"host-pool-{self._allocation_sequence}", + ) + + +def create_tensorcast_host_allocator( + source: TensorcastConfigSource, + *, + host_memory_mode: str, + host_layout: str, + io_backend: str, + platform_name: str, + is_cuda_backend: bool, + world_rank: int | None = None, + world_size: int | None = None, +) -> HostTensorAllocator: + """Validate, attach, and select the configured TensorCast host allocator.""" + + config = normalize_tensorcast_config(source) + validate_tensorcast_startup_configuration( + host_memory_mode=host_memory_mode, + host_layout=host_layout, + io_backend=io_backend, + platform_name=platform_name, + is_cuda_backend=is_cuda_backend, + ) + options = build_tensorcast_session_options( + config, + world_rank=world_rank, + world_size=world_size, + ) + session = attach_early_process_session(options) + if config.transfer_mode == TensorcastTransferMode.SCRATCH: + return HostTensorAllocator() + return TensorcastHostTensorAllocator(session) + + +def get_tensorcast_host_allocator_from_runtime() -> HostTensorAllocator: + """Select the TensorCast allocator from the published SGLang config.""" + + from sglang.srt.platforms import current_platform + from sglang.srt.runtime_context import get_memory + + memory = get_memory() + source = memory.hicache_storage_backend_extra_config + if source is None: + raise ValueError( + "--hicache-storage-backend tensorcast requires " + "--hicache-storage-backend-extra-config with tensorcast.daemon_address" + ) + return create_tensorcast_host_allocator( + source, + host_memory_mode=memory.hicache_host_memory_mode, + host_layout=memory.hicache_mem_layout, + io_backend=memory.hicache_io_backend, + platform_name=sys.platform, + is_cuda_backend=current_platform.is_cuda(), + ) + + +def resolve_tensorcast_model_id( + config: TensorcastConfig, + storage_model_name: str | None, +) -> str: + """Resolve the artifact model identity at later Store registration.""" + + if config.model_id is not None: + return config.model_id + if storage_model_name is None: + raise ValueError( + "TensorCast artifact identity requires tensorcast.model_id or an " + "SGLang storage model name" + ) + try: + return _strip_non_empty(storage_model_name) + except ValueError as exc: + raise ValueError( + "TensorCast artifact identity requires tensorcast.model_id or a " + "non-empty SGLang storage model name" + ) from exc + + +def validate_tensorcast_startup_configuration( + *, + host_memory_mode: str, + host_layout: str, + io_backend: str, + platform_name: str, + is_cuda_backend: bool, +) -> None: + """Reject initial-scope runtime combinations before Session attachment.""" + + if platform_name != "linux": + raise ValueError( + "TensorCast HiCache initially supports Linux only, got " + f"platform={platform_name!r}" + ) + if not is_cuda_backend: + raise ValueError("TensorCast HiCache initially supports the CUDA backend only") + if host_memory_mode != "cache": + raise ValueError( + "TensorCast HiCache requires --hicache-host-memory-mode=cache, got " + f"{host_memory_mode!r}" + ) + if (host_layout, io_backend) not in _SUPPORTED_LAYOUT_IO_PAIRS: + raise ValueError( + "TensorCast HiCache requires (page_first, kernel) or " + "(page_first_direct, direct), got " + f"layout={host_layout!r}, io_backend={io_backend!r}" + ) + + +def _strip_non_empty(value: str) -> str: + stripped = value.strip() + if not stripped: + raise ValueError("value must not be empty") + return stripped + + +def _load_extra_config(source: str) -> Mapping[str, object]: + if not source.startswith("@"): + return _require_mapping(json.loads(source), source_name="inline JSON") + + path_text = source[1:] + if not path_text: + raise ValueError("TensorCast HiCache config path must not be empty") + path = Path(path_text) + suffix = path.suffix.lower() + if suffix == ".json": + with path.open(encoding="utf-8") as config_file: + parsed = json.load(config_file) + elif suffix == ".toml": + import tomllib + + with path.open("rb") as config_file: + parsed = tomllib.load(config_file) + elif suffix in {".yaml", ".yml"}: + import yaml + + with path.open(encoding="utf-8") as config_file: + parsed = yaml.safe_load(config_file) + else: + raise ValueError( + f"Unsupported TensorCast HiCache config file extension {suffix!r}" + ) + return _require_mapping(parsed, source_name=str(path)) + + +def _require_mapping(value: object, *, source_name: str) -> Mapping[str, object]: + if not isinstance(value, Mapping): + raise ValueError( + f"TensorCast HiCache config from {source_name} must be an object" + ) + return cast(Mapping[str, object], value) diff --git a/python/sglang/srt/mem_cache/storage/tensorcast_store/tensorcast_store.py b/python/sglang/srt/mem_cache/storage/tensorcast_store/tensorcast_store.py new file mode 100644 index 000000000..c1487b5a6 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/tensorcast_store/tensorcast_store.py @@ -0,0 +1,860 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to SGLang project + +"""TensorCast-backed synchronous L3 storage adapter.""" + +from __future__ import annotations + +import logging +import threading +from dataclasses import dataclass +from enum import Enum +from hashlib import sha256 +from typing import Any, TypeVar + +import torch +from tensorcast.api.store import ( + ByteArtifactKeyspace, + ByteArtifactSpec, + HostMemorySpan, + RegionArtifactInputError, + RegionArtifactTransfer, + RegionBackedArtifactSession, + RegionSessionFailedError, + RegionSessionTerminatedError, +) + +from sglang.srt.mem_cache.hicache_storage import ( + STORAGE_BATCH_SIZE, + HiCacheStorage, + HiCacheStorageConfig, + HiCacheStorageExtraInfo, + PoolTransfer, + PoolTransferResult, +) +from sglang.srt.mem_cache.pool_host import HostKVCache +from sglang.srt.mem_cache.pool_host.mha import ( + AsymmetricMHATokenToKVPoolHost, + MHATokenToKVPoolHost, +) +from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost +from sglang.srt.mem_cache.storage.tensorcast_store.host_allocator import ( + TensorcastConfig, + TensorcastHostTensorAllocator, + TensorcastTransferMode, + build_tensorcast_session_options, + claim_tensorcast_store_session, + normalize_tensorcast_config, + resolve_tensorcast_model_id, + terminate_tensorcast_store_session, +) + +logger = logging.getLogger(__name__) + +# Schema used for TensorCast byte-artifact key namespace +ARTIFACT_LAYOUT_SCHEMA_VERSION = "full-fragment-v1" +_ARTIFACT_LAYOUT_SCHEMA_TOKEN = "ff1" +_SUPPORTED_FULL_LAYOUTS = frozenset({"page_first", "page_first_direct"}) +_FULL_LAYOUT_TOKENS = { + "page_first": "pf", + "page_first_direct": "pfd", +} +_DTYPE_LAYOUT_TOKENS = { + "bfloat16": "bf16", + "float16": "f16", + "float32": "f32", + "float64": "f64", + "float8_e4m3fn": "f8e4m3fn", + "float8_e4m3fnuz": "f8e4m3fnuz", + "float8_e5m2": "f8e5m2", + "uint8": "u8", +} +_ENGINE_KEY_DOMAIN = b"sglang-hicache-engine-key-v1\0" +_ResultT = TypeVar("_ResultT") + + +class FragmentComponent(str, Enum): + """Logical contiguous components in one FULL page.""" + + K = "k" + V = "v" + KV = "kv" + # TODO: support SWA/Mamba in V2 + + +class _PoolFamily(str, Enum): + MHA = "mha" + MLA = "mla" + # TODO: support SWA/Mamba in V2 + + +@dataclass(frozen=True, slots=True) +class FragmentSchema: + """Immutable byte contract for one component of every logical page.""" + + component: FragmentComponent + byte_length: int + + +@dataclass(frozen=True, slots=True) +class PageFragment: + """One caller-owned identity and optional host range in a page plan.""" + + logical_page_index: int + component: FragmentComponent + engine_key: bytes + host_address: int | None + byte_length: int + owner: torch.Tensor | None + + +@dataclass(frozen=True, slots=True) +class _RegisteredPool: + pool: HostKVCache + family: _PoolFamily + fragment_schema: tuple[FragmentSchema, ...] + roots: tuple[torch.Tensor, ...] + keyspace: ByteArtifactKeyspace + rank_suffix: str + layout_id: str + + +class TensorcastStore(HiCacheStorage): + """Adapt SGLang KV pages to TensorCast byte-artifact Session calls.""" + + def __init__(self, storage_config: HiCacheStorageConfig) -> None: + source = storage_config.extra_config + if source is None: + raise ValueError( + "TensorCast HiCache requires storage backend extra configuration" + ) + tensorcast_config = normalize_tensorcast_config(source) + session_options = build_tensorcast_session_options(tensorcast_config) + + self._storage_config = storage_config + self._tensorcast_config = tensorcast_config + self._session_options = session_options + self._registered: _RegisteredPool | None = None + self._availability_lock = threading.Lock() + self._disabled = False + self._failure_logged = False + + # This must remain the final fallible/stateful constructor step so a local + # validation failure cannot strand ownership in the process registry. + self._session = claim_tensorcast_store_session(session_options, owner=self) + + @property + def session(self) -> RegionBackedArtifactSession: + return self._session + + def register_mem_pool_host(self, mem_pool_host: HostKVCache) -> None: + registered = self._registered + if registered is not None: + if registered.pool is mem_pool_host: + return + raise ValueError( + "TensorCast Store is already registered with a different HostPool" + ) + + candidate = _build_registered_pool( + mem_pool_host, + storage_config=self._storage_config, + tensorcast_config=self._tensorcast_config, + ) + _validate_transfer_mode_registration( + candidate, + config=self._tensorcast_config, + session=self._session, + ) + + super().register_mem_pool_host(mem_pool_host) + self._registered = candidate + + def exists(self, key: str) -> bool: + return self.batch_exists([key]) == 1 + + def batch_exists( + self, + keys: list[str], + extra_info: HiCacheStorageExtraInfo | None = None, + ) -> int: + if not keys: + return 0 + if not self._adapter_is_available(): + return 0 + try: + specs = self._build_artifact_specs(keys) + result = self._session.batch_exists(specs) + registered = self._require_registered() + hit_pages = _leading_complete_page_count( + result.existence_mask, + page_count=len(keys), + fragments_per_page=len(registered.fragment_schema), + ) + except Exception as exc: + self._disable_after_runtime_exception("batch_exists", exc) + return 0 + return self._publish_if_available(hit_pages, fallback=0) + + def _build_page_fragments(self, keys: list[str]) -> tuple[PageFragment, ...]: + registered = self._require_registered() + return _expand_page_fragments(registered, keys) + + def _build_artifact_specs(self, keys: list[str]) -> tuple[ByteArtifactSpec, ...]: + registered = self._require_registered() + return _expand_artifact_specs(registered, keys) + + def _require_registered(self) -> _RegisteredPool: + if self._registered is None: + raise RuntimeError("TensorCast Store has no registered HostPool") + return self._registered + + def get( + self, + key: str, + target_location: Any | None = None, + target_sizes: Any | None = None, + ) -> torch.Tensor | None: + raise NotImplementedError("TensorCast does not support value-oriented get()") + + def batch_get( + self, + keys: list[str], + target_locations: Any | None = None, + target_sizes: Any | None = None, + ) -> list[torch.Tensor | None] | int: + raise NotImplementedError( + "TensorCast does not support value-oriented batch_get()" + ) + + def set( + self, + key: str, + value: Any | None = None, + target_location: Any | None = None, + target_sizes: Any | None = None, + ) -> bool: + raise NotImplementedError("TensorCast does not support value-oriented set()") + + def batch_set( + self, + keys: list[str], + values: Any | None = None, + target_locations: Any | None = None, + target_sizes: Any | None = None, + ) -> bool: + raise NotImplementedError( + "TensorCast does not support value-oriented batch_set()" + ) + + def batch_get_v1( + self, + keys: list[str], + host_indices: torch.Tensor, + extra_info: HiCacheStorageExtraInfo | None = None, + ) -> list[bool]: + unavailable = [False] * len(keys) + if not keys: + return [] + if not self._adapter_is_available(): + return unavailable + try: + transfers = self._build_artifact_transfers(keys, host_indices) + result = self._session.batch_get_into(transfers) + registered = self._require_registered() + page_mask = _fold_fragment_mask( + result.success_mask, + page_count=len(keys), + fragments_per_page=len(registered.fragment_schema), + ) + result_mask = list(_normalize_leading_page_mask(page_mask)) + except Exception as exc: + self._disable_after_runtime_exception("batch_get_v1", exc) + return unavailable + return self._publish_if_available(result_mask, fallback=unavailable) + + def batch_set_v1( + self, + keys: list[str], + host_indices: torch.Tensor, + extra_info: HiCacheStorageExtraInfo | None = None, + ) -> list[bool]: + unavailable = [False] * len(keys) + if not keys: + return [] + if not self._adapter_is_available(): + return unavailable + try: + transfers = self._build_artifact_transfers(keys, host_indices) + result = self._session.batch_put_from(transfers) + registered = self._require_registered() + result_mask = list( + _fold_fragment_mask( + result.success_mask, + page_count=len(keys), + fragments_per_page=len(registered.fragment_schema), + ) + ) + except Exception as exc: + self._disable_after_runtime_exception("batch_set_v1", exc) + return unavailable + return self._publish_if_available(result_mask, fallback=unavailable) + + def _adapter_is_available(self) -> bool: + with self._availability_lock: + return not self._disabled + + def _publish_if_available( + self, + result: _ResultT, + *, + fallback: _ResultT, + ) -> _ResultT: + with self._availability_lock: + if self._disabled: + return fallback + return result + + def _disable_after_runtime_exception( + self, + operation: str, + error: Exception, + ) -> None: + with self._availability_lock: + if self._disabled: + return + self._disabled = True + should_log = not self._failure_logged + self._failure_logged = True + + if not should_log: + return + if isinstance(error, RegionSessionFailedError): + failure = error.failure + logger.error( + "TensorCast L3 unavailable: category=session_failed " + "exception_type=RegionSessionFailedError adapter_operation=%s " + "failure_code=%s session_operation=%s operation_id=%s message=%s", + operation, + failure.code.value, + failure.operation_kind.value, + failure.operation_id, + failure.message, + ) + return + if isinstance(error, RegionArtifactInputError): + logger.error( + "TensorCast L3 unavailable: category=adapter_input_failure " + "exception_type=RegionArtifactInputError adapter_operation=%s " + "message=%s", + operation, + error, + ) + return + if isinstance(error, RegionSessionTerminatedError): + logger.error( + "TensorCast L3 unavailable: category=unexpected_terminated " + "exception_type=RegionSessionTerminatedError adapter_operation=%s " + "message=%s", + operation, + error, + ) + return + logger.exception( + "TensorCast L3 unavailable: category=adapter_exception " + "exception_type=%s adapter_operation=%s message=%s", + type(error).__name__, + operation, + error, + ) + + def close(self) -> None: + with self._availability_lock: + self._disabled = True + terminate_tensorcast_store_session(owner=self) + + def _build_artifact_transfers( + self, + keys: list[str], + host_indices: torch.Tensor, + ) -> tuple[RegionArtifactTransfer, ...]: + registered = self._require_registered() + return _expand_artifact_transfers(registered, keys, host_indices) + + def batch_exists_v2( + self, + keys: list[str], + pool_transfers: list[PoolTransfer] | None = None, + extra_info: HiCacheStorageExtraInfo | None = None, + ) -> PoolTransferResult: + raise NotImplementedError("TensorCast initially supports FULL v1 only") + + def batch_get_v2( + self, + transfers: list[PoolTransfer], + extra_info: HiCacheStorageExtraInfo | None = None, + ) -> dict[str, list[bool]]: + raise NotImplementedError("TensorCast initially supports FULL v1 only") + + def batch_set_v2( + self, + transfers: list[PoolTransfer], + extra_info: HiCacheStorageExtraInfo | None = None, + ) -> dict[str, list[bool]]: + raise NotImplementedError("TensorCast initially supports FULL v1 only") + + def clear(self) -> None: + raise NotImplementedError("TensorCast has no namespace-wide clear operation") + + +def _build_registered_pool( + mem_pool_host: HostKVCache, + *, + storage_config: HiCacheStorageConfig, + tensorcast_config: TensorcastConfig, +) -> _RegisteredPool: + family, components, roots = _select_pool_adapter(mem_pool_host) + _validate_storage_config( + mem_pool_host, + family=family, + storage_config=storage_config, + ) + fragment_schema = _sample_fragment_schema(mem_pool_host, components=components) + _validate_roots(roots, pool_type=type(mem_pool_host).__name__) + + effective_model_id = resolve_tensorcast_model_id( + tensorcast_config, + storage_config.model_name, + ) + rank_suffix = _build_rank_suffix(storage_config, family=family) + layout_id = _build_layout_id(mem_pool_host, family=family) + keyspace = ByteArtifactKeyspace( + namespace=tensorcast_config.namespace, + engine="sglang", + model_id=effective_model_id, + model_version=tensorcast_config.model_version, + layout_id=layout_id, + ) + return _RegisteredPool( + pool=mem_pool_host, + family=family, + fragment_schema=fragment_schema, + roots=roots, + keyspace=keyspace, + rank_suffix=rank_suffix, + layout_id=layout_id, + ) + + +def _select_pool_adapter( + mem_pool_host: HostKVCache, +) -> tuple[_PoolFamily, tuple[FragmentComponent, ...], tuple[torch.Tensor, ...]]: + pool_type = type(mem_pool_host) + if pool_type is AsymmetricMHATokenToKVPoolHost: + asymmetric_pool = mem_pool_host + return ( + _PoolFamily.MHA, + (FragmentComponent.K, FragmentComponent.V), + (asymmetric_pool.k_buffer, asymmetric_pool.v_buffer), + ) + if pool_type is MHATokenToKVPoolHost: + mha_pool = mem_pool_host + return ( + _PoolFamily.MHA, + (FragmentComponent.K, FragmentComponent.V), + (mha_pool.kv_buffer, mha_pool.kv_buffer), + ) + if pool_type is MLATokenToKVPoolHost: + mla_pool = mem_pool_host + return ( + _PoolFamily.MLA, + (FragmentComponent.KV,), + (mla_pool.kv_buffer,), + ) + raise NotImplementedError( + "TensorCast FULL v1 does not support HostPool type " + f"{pool_type.__module__}.{pool_type.__qualname__}" + ) + + +def _validate_storage_config( + mem_pool_host: HostKVCache, + *, + family: _PoolFamily, + storage_config: HiCacheStorageConfig, +) -> None: + """Check that the storage configuration matches the implemented data path. + + The current checks cover only FULL v1. Extend this validation as support for + additional pool families and storage API versions is implemented. + """ + if mem_pool_host.layout not in _SUPPORTED_FULL_LAYOUTS: + raise NotImplementedError( + "TensorCast FULL v1 supports only page_first and page_first_direct, " + f"got layout={mem_pool_host.layout!r}" + ) + if mem_pool_host.mtp_draft_device_pools: + raise NotImplementedError( + "TensorCast FULL v1 does not support packed MTP draft pools" + ) + if mem_pool_host.dcp_size != 1 or mem_pool_host.dcp_rank != 0: + raise NotImplementedError( + "TensorCast FULL v1 does not support DCP-aware HostPool layouts" + ) + if storage_config.attn_cp_size != 1 or storage_config.attn_cp_rank != 0: + raise NotImplementedError( + "TensorCast FULL v1 requires attention context parallel size 1" + ) + if storage_config.tp_lcm_size is not None: + raise NotImplementedError("TensorCast FULL v1 does not support TP-LCM") + if storage_config.should_split_heads: + raise NotImplementedError("TensorCast FULL v1 does not support split heads") + expected_mla = family is _PoolFamily.MLA + if storage_config.is_mla_model is not expected_mla: + raise ValueError( + "TensorCast HostPool family conflicts with is_mla_model: " + f"pool_family={family.value}, is_mla_model={storage_config.is_mla_model}" + ) + if mem_pool_host.page_size <= 0 or mem_pool_host.page_num <= 0: + raise ValueError( + "TensorCast FULL HostPool requires positive page_size and page_num" + ) + + +def _sample_fragment_schema( + mem_pool_host: HostKVCache, + *, + components: tuple[FragmentComponent, ...], +) -> tuple[FragmentSchema, ...]: + sample_indices = torch.arange(mem_pool_host.page_size, dtype=torch.int64) + pointers, byte_lengths = mem_pool_host.get_page_buffer_meta(sample_indices) + expected_count = len(components) + if len(pointers) != expected_count or len(byte_lengths) != expected_count: + raise ValueError( + "TensorCast FULL registration metadata must contain exactly one page: " + f"expected_components={expected_count}, pointers={len(pointers)}, " + f"byte_lengths={len(byte_lengths)}" + ) + if any(type(pointer) is not int or pointer <= 0 for pointer in pointers): + raise ValueError( + "TensorCast FULL registration metadata contains a non-positive pointer" + ) + if any( + type(byte_length) is not int or byte_length <= 0 for byte_length in byte_lengths + ): + raise ValueError( + "TensorCast FULL registration metadata contains an invalid byte length" + ) + if type(mem_pool_host) is MHATokenToKVPoolHost and ( + byte_lengths[0] != byte_lengths[1] + ): + raise ValueError( + "TensorCast standard MHA registration requires equal K/V byte lengths" + ) + return tuple( + FragmentSchema(component=component, byte_length=byte_length) + for component, byte_length in zip(components, byte_lengths, strict=True) + ) + + +def _validate_roots(roots: tuple[torch.Tensor, ...], *, pool_type: str) -> None: + for root in roots: + if type(root) is not torch.Tensor: + raise TypeError( + f"TensorCast {pool_type} backing roots must be torch.Tensor objects" + ) + if root.device.type != "cpu" or not root.is_contiguous(): + raise ValueError( + f"TensorCast {pool_type} backing roots must be contiguous CPU tensors" + ) + + +def _validate_transfer_mode_registration( + registered: _RegisteredPool, + *, + config: TensorcastConfig, + session: RegionBackedArtifactSession, +) -> None: + if config.transfer_mode is TensorcastTransferMode.ALLOCATOR: + allocator = registered.pool.allocator + if type(allocator) is not TensorcastHostTensorAllocator: + raise ValueError( + "TensorCast allocator mode requires TensorcastHostTensorAllocator " + "backing for the registered HostPool" + ) + if allocator.session is not session: + raise ValueError( + "TensorCast HostPool allocator and Store must use the same Session" + ) + return + + maximum_logical_batch_pages = min(STORAGE_BATCH_SIZE, registered.pool.page_num) + page_bytes = sum(fragment.byte_length for fragment in registered.fragment_schema) + required_capacity_bytes = maximum_logical_batch_pages * page_bytes + configured_capacity_bytes = config.scratch.capacity_bytes + if configured_capacity_bytes < required_capacity_bytes: + component_lengths = tuple( + fragment.byte_length for fragment in registered.fragment_schema + ) + raise ValueError( + "TensorCast scratch capacity is insufficient: " + f"configured_capacity_bytes={configured_capacity_bytes}, " + f"required_capacity_bytes={required_capacity_bytes}, " + f"maximum_logical_batch_pages={maximum_logical_batch_pages}, " + f"page_bytes={page_bytes}, " + f"pool_type={type(registered.pool).__name__}, " + f"layout={registered.pool.layout}, " + f"component_byte_lengths={component_lengths}" + ) + + +def _build_layout_id(mem_pool_host: HostKVCache, *, family: _PoolFamily) -> str: + dtype_name = str(mem_pool_host.dtype) + if dtype_name.startswith("torch."): + dtype_name = dtype_name.removeprefix("torch.") + dtype_token = _DTYPE_LAYOUT_TOKENS.get(dtype_name, dtype_name) + layout_token = _FULL_LAYOUT_TOKENS[mem_pool_host.layout] + return ( + f"{_ARTIFACT_LAYOUT_SCHEMA_TOKEN}_{layout_token}_{dtype_token}_" + f"p{mem_pool_host.page_size}_{family.value}" + ) + + +def _build_rank_suffix( + storage_config: HiCacheStorageConfig, + *, + family: _PoolFamily, +) -> str: + if family is _PoolFamily.MLA: + return f"pp{storage_config.pp_rank}of{storage_config.pp_size}" + tp_suffix = f"tp{storage_config.tp_rank}of{storage_config.tp_size}" + if storage_config.pp_size == 1: + return tp_suffix + return f"{tp_suffix}_pp{storage_config.pp_rank}of{storage_config.pp_size}" + + +def _build_engine_key( + rank_suffix: str, + logical_key: str, + component: FragmentComponent, +) -> bytes: + if type(logical_key) is not str: + raise TypeError( + "TensorCast logical page keys must be exact str values, got " + f"{type(logical_key).__name__}" + ) + component_suffix = ( + b"k" if component is FragmentComponent.KV else component.value.encode() + ) + rank_bytes = rank_suffix.encode("ascii") + logical_key_bytes = logical_key.encode("utf-8") + payload = b"".join( + ( + _ENGINE_KEY_DOMAIN, + len(rank_bytes).to_bytes(4, byteorder="big"), + rank_bytes, + len(logical_key_bytes).to_bytes(8, byteorder="big"), + logical_key_bytes, + component_suffix, + ) + ) + return sha256(payload).digest() + + +def _expand_page_fragments( + registered: _RegisteredPool, + keys: list[str], +) -> tuple[PageFragment, ...]: + fragments = tuple( + PageFragment( + logical_page_index=page_index, + component=schema.component, + engine_key=_build_engine_key( + registered.rank_suffix, + logical_key, + schema.component, + ), + host_address=None, + byte_length=schema.byte_length, + owner=None, + ) + for page_index, logical_key in enumerate(keys) + for schema in registered.fragment_schema + ) + engine_keys = tuple(fragment.engine_key for fragment in fragments) + if len(set(engine_keys)) != len(engine_keys): + raise ValueError( + "TensorCast logical keys must produce unique fragment identities " + "within one batch" + ) + return fragments + + +def _expand_artifact_specs( + registered: _RegisteredPool, + keys: list[str], +) -> tuple[ByteArtifactSpec, ...]: + return tuple( + ByteArtifactSpec( + keyspace=registered.keyspace, + engine_key=fragment.engine_key, + byte_length=fragment.byte_length, + ) + for fragment in _expand_page_fragments(registered, keys) + ) + + +def _expand_transfer_fragments( + registered: _RegisteredPool, + keys: list[str], + host_indices: torch.Tensor, +) -> tuple[PageFragment, ...]: + if type(host_indices) is not torch.Tensor: + raise TypeError( + "TensorCast FULL host_indices must be a torch.Tensor, got " + f"{type(host_indices).__name__}" + ) + if host_indices.ndim != 1: + raise ValueError( + "TensorCast FULL host_indices must be one-dimensional, got " + f"shape={tuple(host_indices.shape)}" + ) + expected_index_count = len(keys) * registered.pool.page_size + if len(host_indices) != expected_index_count: + raise ValueError( + "TensorCast FULL logical-key/host-index count mismatch: " + f"logical_pages={len(keys)}, page_size={registered.pool.page_size}, " + f"expected_host_indices={expected_index_count}, " + f"actual_host_indices={len(host_indices)}" + ) + if not keys: + return () + + pointers, byte_lengths = registered.pool.get_page_buffer_meta(host_indices) + fragments = _expand_page_fragments(registered, keys) + expected_fragment_count = len(fragments) + if ( + len(pointers) != expected_fragment_count + or len(byte_lengths) != expected_fragment_count + ): + raise ValueError( + "TensorCast FULL runtime metadata count mismatch: " + f"expected_fragments={expected_fragment_count}, " + f"pointers={len(pointers)}, byte_lengths={len(byte_lengths)}" + ) + + fragments_per_page = len(registered.fragment_schema) + transfer_fragments: list[PageFragment] = [] + for fragment_index, (fragment, pointer, byte_length) in enumerate( + zip(fragments, pointers, byte_lengths, strict=True) + ): + schema_index = fragment_index % fragments_per_page + schema = registered.fragment_schema[schema_index] + if type(pointer) is not int or pointer <= 0: + raise ValueError( + "TensorCast FULL runtime metadata contains a non-positive pointer: " + f"fragment_index={fragment_index}, pointer={pointer!r}" + ) + if type(byte_length) is not int or byte_length <= 0: + raise ValueError( + "TensorCast FULL runtime metadata contains an invalid byte length: " + f"fragment_index={fragment_index}, byte_length={byte_length!r}" + ) + if byte_length != schema.byte_length: + raise ValueError( + "TensorCast FULL runtime byte length differs from the registered " + f"schema: fragment_index={fragment_index}, " + f"component={schema.component.value}, expected={schema.byte_length}, " + f"actual={byte_length}" + ) + transfer_fragments.append( + PageFragment( + logical_page_index=fragment.logical_page_index, + component=fragment.component, + engine_key=fragment.engine_key, + host_address=pointer, + byte_length=byte_length, + owner=registered.roots[schema_index], + ) + ) + return tuple(transfer_fragments) + + +def _expand_artifact_transfers( + registered: _RegisteredPool, + keys: list[str], + host_indices: torch.Tensor, +) -> tuple[RegionArtifactTransfer, ...]: + transfers: list[RegionArtifactTransfer] = [] + for fragment in _expand_transfer_fragments(registered, keys, host_indices): + if fragment.host_address is None or fragment.owner is None: + raise RuntimeError("TensorCast transfer fragment has no owned host range") + offset_bytes = fragment.host_address - int(fragment.owner.data_ptr()) + span = HostMemorySpan.from_tensor( + fragment.owner, + offset_bytes=offset_bytes, + byte_length=fragment.byte_length, + ) + transfers.append( + RegionArtifactTransfer( + artifact=ByteArtifactSpec( + keyspace=registered.keyspace, + engine_key=fragment.engine_key, + byte_length=fragment.byte_length, + ), + span=span, + ) + ) + return tuple(transfers) + + +def _fold_fragment_mask( + fragment_mask: tuple[bool, ...], + *, + page_count: int, + fragments_per_page: int, +) -> tuple[bool, ...]: + if fragments_per_page <= 0: + raise ValueError("TensorCast fragments_per_page must be positive") + expected_count = page_count * fragments_per_page + if len(fragment_mask) != expected_count: + raise ValueError( + "TensorCast transfer result length does not match the logical page plan: " + f"expected={expected_count}, actual={len(fragment_mask)}" + ) + return tuple( + all( + fragment_mask[ + page_index * fragments_per_page : (page_index + 1) * fragments_per_page + ] + ) + for page_index in range(page_count) + ) + + +def _normalize_leading_page_mask(page_mask: tuple[bool, ...]) -> tuple[bool, ...]: + prefix_complete = True + normalized: list[bool] = [] + for page_complete in page_mask: + prefix_complete = prefix_complete and page_complete + normalized.append(prefix_complete) + return tuple(normalized) + + +def _leading_complete_page_count( + fragment_mask: tuple[bool, ...], + *, + page_count: int, + fragments_per_page: int, +) -> int: + page_mask = _fold_fragment_mask( + fragment_mask, + page_count=page_count, + fragments_per_page=fragments_per_page, + ) + prefix = 0 + for page_complete in page_mask: + if not page_complete: + break + prefix += 1 + return prefix diff --git a/python/sglang/srt/mem_cache/storage/tensorcast_store/test_tensorcast_host_allocator.py b/python/sglang/srt/mem_cache/storage/tensorcast_store/test_tensorcast_host_allocator.py new file mode 100644 index 000000000..29a508be8 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/tensorcast_store/test_tensorcast_host_allocator.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to SGLang project + +from __future__ import annotations + +import subprocess +import sys +import textwrap + + +def _run_fresh_python(source: str) -> subprocess.CompletedProcess[str]: + """Run a scenario with pristine import and process-lifecycle state. + + TensorCast SDK modules may already be cached by pytest collection, and the + adapter's process-scoped Session registry is intentionally not resettable. + A newly executed interpreter preserves those production invariants while + keeping import-boundary and lifecycle tests independent of test order. + """ + return subprocess.run( + [sys.executable, "-c", textwrap.dedent(source)], + check=False, + capture_output=True, + text=True, + ) + + +# Attach, claim, fail, or terminate the process-scoped Session registry. +# Each complete lifecycle should belong to a fresh interpreter. +def test_allocator_delegates_exact_tensors_in_subprocess() -> None: + result = _run_fresh_python(""" + import torch + + from sglang.srt.mem_cache.storage.tensorcast_store import host_allocator + + class FakeSession: + def __init__(self): + self.calls = [] + self.results = [ + torch.empty((2, 3), dtype=torch.float32), + torch.empty((5,), dtype=torch.uint8), + torch.empty((1,), dtype=torch.int64), + ] + + def allocate_host_tensor(self, shape, dtype, *, name): + self.calls.append((shape, dtype, name)) + return self.results[len(self.calls) - 1] + + def terminate_process_session(self): + raise AssertionError("allocator terminated the Session") + + session = FakeSession() + attach_calls = [] + + def fake_attach(options): + attach_calls.append(options) + return session + + host_allocator._attach_process_session = fake_attach + source = {"tensorcast": {"daemon_address": "127.0.0.1:8073"}} + first_allocator = host_allocator.create_tensorcast_host_allocator( + source, + host_memory_mode="cache", + host_layout="page_first", + io_backend="kernel", + platform_name="linux", + is_cuda_backend=True, + world_rank=0, + world_size=1, + ) + second_allocator = host_allocator.create_tensorcast_host_allocator( + source, + host_memory_mode="cache", + host_layout="page_first", + io_backend="kernel", + platform_name="linux", + is_cuda_backend=True, + world_rank=0, + world_size=1, + ) + + assert isinstance( + first_allocator, host_allocator.TensorcastHostTensorAllocator + ) + assert isinstance( + second_allocator, host_allocator.TensorcastHostTensorAllocator + ) + assert first_allocator.session is session + assert second_allocator.session is session + assert len(attach_calls) == 1 + + first = first_allocator.allocate((2, 3), torch.float32, "cpu") + second = first_allocator.allocate((5,), torch.uint8, "cpu") + third = second_allocator.allocate((1,), torch.int64, "cpu") + assert first is session.results[0] + assert second is session.results[1] + assert third is session.results[2] + assert session.calls == [ + ((2, 3), torch.float32, "host-pool-1"), + ((5,), torch.uint8, "host-pool-2"), + ((1,), torch.int64, "host-pool-1"), + ] + assert first_allocator.dims == (5,) + assert first_allocator.dtype is torch.uint8 + + try: + first_allocator.allocate((1,), torch.float32, "cuda") + except ValueError as exc: + assert "requires CPU memory" in str(exc) + else: + raise AssertionError("non-CPU allocation was accepted") + assert len(session.calls) == 3 + """) + assert result.returncode == 0, result.stderr + + +def test_registry_equal_conflicting_and_store_claims_in_subprocess() -> None: + result = _run_fresh_python(""" + from sglang.srt.mem_cache.storage.tensorcast_store import host_allocator + + class FakeSession: + def __init__(self): + self.terminate_calls = 0 + + def terminate_process_session(self): + self.terminate_calls += 1 + + session = FakeSession() + attach_calls = [] + + def fake_attach(options): + attach_calls.append(options) + return session + + host_allocator._attach_process_session = fake_attach + first_config = host_allocator.TensorcastConfig( + daemon_address="127.0.0.1:8073" + ) + first_options = host_allocator.build_tensorcast_session_options( + first_config, world_rank=0, world_size=1 + ) + conflicting_options = host_allocator.build_tensorcast_session_options( + host_allocator.TensorcastConfig( + daemon_address="127.0.0.1:8073", exists_timeout_s=31.0 + ), + world_rank=0, + world_size=1, + ) + + assert host_allocator.attach_early_process_session(first_options) is session + assert host_allocator.attach_early_process_session(first_options) is session + assert len(attach_calls) == 1 + try: + host_allocator.attach_early_process_session(conflicting_options) + except host_allocator.TensorcastSessionRegistryError as exc: + assert "conflicts" in str(exc) + else: + raise AssertionError("conflicting early attach was accepted") + + owner = object() + other_owner = object() + assert ( + host_allocator.claim_tensorcast_store_session( + first_options, owner=owner + ) + is session + ) + assert ( + host_allocator.claim_tensorcast_store_session( + first_options, owner=owner + ) + is session + ) + try: + host_allocator.claim_tensorcast_store_session( + first_options, owner=other_owner + ) + except host_allocator.TensorcastSessionRegistryError as exc: + assert "already owns" in str(exc) + else: + raise AssertionError("a second Store owner was accepted") + + assert host_allocator.terminate_tensorcast_store_session(owner=owner) is True + assert host_allocator.terminate_tensorcast_store_session(owner=owner) is False + assert session.terminate_calls == 1 + for operation in ( + lambda: host_allocator.attach_early_process_session(first_options), + lambda: host_allocator.claim_tensorcast_store_session( + first_options, owner=owner + ), + ): + try: + operation() + except host_allocator.TensorcastSessionRegistryError as exc: + assert "terminal" in str(exc) + else: + raise AssertionError("terminal registry admitted an operation") + """) + assert result.returncode == 0, result.stderr + + +def test_store_close_lifecycle_is_terminal_and_retains_allocator_roots_in_subprocess() -> ( + None +): + result = _run_fresh_python(""" + from types import SimpleNamespace + + import torch + + import sglang.srt.runtime_context as runtime_context + from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig + from sglang.srt.mem_cache.storage.tensorcast_store import host_allocator + from sglang.srt.mem_cache.storage.tensorcast_store.tensorcast_store import ( + TensorcastStore, + ) + + runtime_context.get_parallel = lambda: SimpleNamespace( + world_rank=0, + world_size=1, + ) + + class FakeSession: + def __init__(self): + self.terminate_calls = 0 + + def allocate_host_tensor(self, shape, dtype, *, name): + del name + return torch.empty(shape, dtype=dtype) + + def terminate_process_session(self): + self.terminate_calls += 1 + + session = FakeSession() + host_allocator._attach_process_session = lambda options: session + source = { + "tensorcast": { + "daemon_address": "127.0.0.1:8073", + "model_id": "close-lifecycle-model", + } + } + allocator = host_allocator.create_tensorcast_host_allocator( + source, + host_memory_mode="cache", + host_layout="page_first_direct", + io_backend="direct", + platform_name="linux", + is_cuda_backend=True, + world_rank=0, + world_size=1, + ) + root = allocator.allocate((8,), torch.float32, "cpu") + root.fill_(7) + 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=False, + enable_storage_metrics=False, + is_page_first_layout=True, + model_name="unused", + extra_config=source, + ) + store = TensorcastStore(config) + store.close() + store.close() + + assert session.terminate_calls == 1 + assert store.batch_exists(["after-close"]) == 0 + indices = torch.tensor([0, 1], dtype=torch.int64) + assert store.batch_get_v1(["after-close"], indices) == [False] + assert store.batch_set_v1(["after-close"], indices) == [False] + assert torch.equal(root, torch.full((8,), 7.0)) + root.add_(1) + assert torch.equal(root, torch.full((8,), 8.0)) + + try: + TensorcastStore(config) + except host_allocator.TensorcastSessionRegistryError as exc: + assert "terminal" in str(exc) + else: + raise AssertionError("terminal registry admitted a second Store") + + try: + host_allocator.create_tensorcast_host_allocator( + source, + host_memory_mode="cache", + host_layout="page_first_direct", + io_backend="direct", + platform_name="linux", + is_cuda_backend=True, + world_rank=0, + world_size=1, + ) + except host_allocator.TensorcastSessionRegistryError as exc: + assert "terminal" in str(exc) + else: + raise AssertionError("terminal registry admitted allocator reclaim") + """) + assert result.returncode == 0, result.stderr diff --git a/python/sglang/srt/mem_cache/storage/tensorcast_store/test_tensorcast_store.py b/python/sglang/srt/mem_cache/storage/tensorcast_store/test_tensorcast_store.py new file mode 100644 index 000000000..e6e8cc9b8 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/tensorcast_store/test_tensorcast_store.py @@ -0,0 +1,1263 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to SGLang project + +from __future__ import annotations + +import logging +import threading +from collections.abc import Sequence +from datetime import datetime, timezone +from types import SimpleNamespace +from typing import cast + +import pytest +import torch +from tensorcast.api.store import ( + RegionArtifactExistsResult, + RegionArtifactInputError, + RegionArtifactTransferResult, + RegionSessionFailedError, + RegionSessionFailure, + RegionSessionFailureCode, + RegionSessionOperationKind, + RegionSessionTerminatedError, +) + +from sglang.srt.managers.cache_controller import HiCacheController +from sglang.srt.mem_cache.hicache_storage import ( + STORAGE_BATCH_SIZE, + HiCacheStorage, + HiCacheStorageConfig, + PoolName, +) +from sglang.srt.mem_cache.pool_host import ( + HostKVCache, + HostPoolGroup, + HostTensorAllocator, + PoolEntry, +) +from sglang.srt.mem_cache.pool_host.mha import ( + AsymmetricMHATokenToKVPoolHost, + MHATokenToKOnlyPoolHost, + MHATokenToKVPoolHost, +) +from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost +from sglang.srt.mem_cache.storage.tensorcast_store.host_allocator import ( + TensorcastConfig, + TensorcastHostTensorAllocator, + normalize_tensorcast_config, +) +from sglang.srt.mem_cache.storage.tensorcast_store.tensorcast_store import ( + ARTIFACT_LAYOUT_SCHEMA_VERSION, + FragmentComponent, + TensorcastStore, + _build_registered_pool, + _expand_artifact_specs, + _expand_transfer_fragments, + _fold_fragment_mask, + _leading_complete_page_count, + _normalize_leading_page_mask, + _validate_transfer_mode_registration, +) + + +class _FakeSession: + def __init__( + self, + existence_mask: Sequence[bool] = (), + *, + get_mask: Sequence[bool] = (), + put_mask: Sequence[bool] = (), + exists_error: Exception | None = None, + get_error: Exception | None = None, + put_error: Exception | None = None, + ) -> None: + self.existence_mask = tuple(existence_mask) + self.get_mask = tuple(get_mask) + self.put_mask = tuple(put_mask) + self.exists_calls: list[tuple[object, ...]] = [] + self.get_calls: list[tuple[object, ...]] = [] + self.put_calls: list[tuple[object, ...]] = [] + self.exists_error = exists_error + self.get_error = get_error + self.put_error = put_error + self.terminate_calls = 0 + self.allocation_calls: list[ + tuple[tuple[int, ...], torch.dtype, str, torch.Tensor] + ] = [] + + def allocate_host_tensor( + self, + shape: tuple[int, ...], + dtype: torch.dtype, + *, + name: str, + ) -> torch.Tensor: + tensor = torch.empty(shape, dtype=dtype) + self.allocation_calls.append((shape, dtype, name, tensor)) + return tensor + + def batch_exists(self, specs: Sequence[object]) -> RegionArtifactExistsResult: + self.exists_calls.append(tuple(specs)) + if self.exists_error is not None: + raise self.exists_error + return RegionArtifactExistsResult( + existence_mask=self.existence_mask, + rpc_elapsed_s=0.0, + ) + + def batch_get_into( + self, transfers: Sequence[object] + ) -> RegionArtifactTransferResult: + snapshot = tuple(transfers) + self.get_calls.append(snapshot) + if self.get_error is not None: + raise self.get_error + return RegionArtifactTransferResult( + success_mask=self.get_mask, + operation_id="fake-get" if snapshot else None, + pack_elapsed_s=0.0, + copy_elapsed_s=0.0, + rpc_elapsed_s=0.0, + ) + + def batch_put_from( + self, transfers: Sequence[object] + ) -> RegionArtifactTransferResult: + snapshot = tuple(transfers) + self.put_calls.append(snapshot) + if self.put_error is not None: + raise self.put_error + return RegionArtifactTransferResult( + success_mask=self.put_mask, + operation_id="fake-put" if snapshot else None, + pack_elapsed_s=0.0, + copy_elapsed_s=0.0, + rpc_elapsed_s=0.0, + ) + + def terminate_process_session(self) -> None: + self.terminate_calls += 1 + + +def _wrapped_config( + *, + transfer_mode: str = "scratch", + scratch_capacity: int = 4096, + **overrides: object, +) -> dict[str, object]: + tensorcast: dict[str, object] = { + "daemon_address": "127.0.0.1:8073", + "transfer_mode": transfer_mode, + "scratch": {"capacity_bytes": scratch_capacity}, + "model_id": "model-a", + } + tensorcast.update(overrides) + return {"tensorcast": tensorcast} + + +def _storage_config( + *, + family: str = "mha", + extra_config: dict[str, object] | None = None, + **overrides: object, +) -> HiCacheStorageConfig: + values: dict[str, object] = { + "tp_rank": 0, + "tp_size": 1, + "pp_rank": 0, + "pp_size": 1, + "attn_cp_rank": 0, + "attn_cp_size": 1, + "is_mla_model": family == "mla", + "enable_storage_metrics": False, + "is_page_first_layout": True, + "model_name": "fallback-model", + "tp_lcm_size": None, + "should_split_heads": False, + "extra_config": extra_config or _wrapped_config(), + } + values.update(overrides) + return HiCacheStorageConfig(**values) # type: ignore[arg-type] + + +def _make_pool( + family: str, + layout: str, + *, + page_num: int = 3, + page_size: int = 2, + dtype: torch.dtype = torch.float32, + allocator: HostTensorAllocator | None = None, + allocate_with_allocator: bool = False, +) -> HostKVCache: + size = page_num * page_size + layer_num = 1 + head_num = 1 + + def allocate_root(shape: tuple[int, ...]) -> torch.Tensor: + if not allocate_with_allocator: + return torch.empty(shape, dtype=dtype) + if allocator is None: + raise ValueError("allocator-backed test pool requires an allocator") + return allocator.allocate(shape, dtype, "cpu") + + if family == "mha": + pool = object.__new__(MHATokenToKVPoolHost) + head_dim = 2 + if layout == "page_first_direct": + root = allocate_root( + (2, page_num, layer_num, page_size, head_num, head_dim) + ) + else: + root = allocate_root((2, size, layer_num, head_num, head_dim)) + pool.kv_buffer = root + pool.head_dim = head_dim + elif family == "asymmetric_mha": + pool = object.__new__(AsymmetricMHATokenToKVPoolHost) + head_dim = 2 + v_head_dim = 3 + if layout == "page_first_direct": + k_root = allocate_root((page_num, layer_num, page_size, head_num, head_dim)) + v_root = allocate_root( + (page_num, layer_num, page_size, head_num, v_head_dim) + ) + else: + k_root = allocate_root((size, layer_num, head_num, head_dim)) + v_root = allocate_root((size, layer_num, head_num, v_head_dim)) + pool.kv_buffer = (k_root, v_root) + pool.head_dim = head_dim + pool.v_head_dim = v_head_dim + elif family == "mla": + pool = object.__new__(MLATokenToKVPoolHost) + kv_cache_dim = 3 + if layout == "page_first_direct": + root = allocate_root((page_num, layer_num, page_size, 1, kv_cache_dim)) + else: + root = allocate_root((size, layer_num, 1, kv_cache_dim)) + pool.kv_buffer = root + pool.kv_cache_dim = kv_cache_dim + else: + raise ValueError(f"unknown test pool family {family}") + + pool.layout = layout + pool.page_num = page_num + pool.page_size = page_size + pool.size = size + pool.layer_num = layer_num + pool.head_num = head_num + pool.dtype = dtype + pool.mtp_draft_device_pools = () + pool.dcp_size = 1 + pool.dcp_rank = 0 + pool.allocator = allocator or HostTensorAllocator() + return cast(HostKVCache, pool) + + +def _candidate( + family: str, + layout: str, + *, + pool: HostKVCache | None = None, + storage_overrides: dict[str, object] | None = None, + config: TensorcastConfig | None = None, +): + normalized_family = "mla" if family == "mla" else "mha" + storage_config = _storage_config( + family=normalized_family, + **(storage_overrides or {}), + ) + return _build_registered_pool( + pool or _make_pool(family, layout), + storage_config=storage_config, + tensorcast_config=config + or normalize_tensorcast_config(storage_config.extra_config), + ) + + +@pytest.mark.parametrize("layout", ["page_first", "page_first_direct"]) +@pytest.mark.parametrize( + ("family", "components", "lengths"), + [ + ("mha", (FragmentComponent.K, FragmentComponent.V), (16, 16)), + ( + "asymmetric_mha", + (FragmentComponent.K, FragmentComponent.V), + (16, 24), + ), + ("mla", (FragmentComponent.KV,), (24,)), + ], +) +def test_fragment_registration_matrix_and_root_retention( + family: str, + components: tuple[FragmentComponent, ...], + lengths: tuple[int, ...], + layout: str, +) -> None: + pool = _make_pool(family, layout) + registered = _candidate(family, layout, pool=pool) + + assert tuple(item.component for item in registered.fragment_schema) == components + assert tuple(item.byte_length for item in registered.fragment_schema) == lengths + if family == "mha": + assert registered.roots[0] is pool.kv_buffer + assert registered.roots[1] is pool.kv_buffer + elif family == "asymmetric_mha": + assert registered.roots == (pool.k_buffer, pool.v_buffer) + else: + assert registered.roots == (pool.kv_buffer,) + + +def test_fragment_registration_dispatch_is_exact_and_specific() -> None: + asymmetric = _candidate("asymmetric_mha", "page_first") + assert tuple(fragment.byte_length for fragment in asymmetric.fragment_schema) == ( + 16, + 24, + ) + + class UnknownMhaPool(MHATokenToKVPoolHost): + pass + + unknown = object.__new__(UnknownMhaPool) + with pytest.raises(NotImplementedError, match="does not support HostPool type"): + _candidate("mha", "page_first", pool=cast(HostKVCache, unknown)) + + k_only = object.__new__(MHATokenToKOnlyPoolHost) + with pytest.raises(NotImplementedError, match="does not support HostPool type"): + _candidate("mha", "page_first", pool=cast(HostKVCache, k_only)) + + +@pytest.mark.parametrize( + ("pool_changes", "storage_changes", "message"), + [ + ({"layout": "layer_first"}, {}, "supports only"), + ({"layout": "page_head"}, {}, "supports only"), + ({"mtp_draft_device_pools": (object(),)}, {}, "MTP"), + ({"dcp_size": 2}, {}, "DCP"), + ({}, {"attn_cp_size": 2}, "context parallel"), + ({}, {"tp_lcm_size": 2}, "TP-LCM"), + ({}, {"should_split_heads": True}, "split heads"), + ], +) +def test_scope_rejects_unsupported_full_semantics( + pool_changes: dict[str, object], + storage_changes: dict[str, object], + message: str, +) -> None: + pool = _make_pool("mha", "page_first") + for name, value in pool_changes.items(): + setattr(pool, name, value) + with pytest.raises(NotImplementedError, match=message): + _candidate( + "mha", + "page_first", + pool=pool, + storage_overrides=storage_changes, + ) + + +def test_scope_rejects_family_mismatch_and_invalid_page_geometry() -> None: + pool = _make_pool("mha", "page_first") + with pytest.raises(ValueError, match="conflicts with is_mla_model"): + _candidate( + "mha", + "page_first", + pool=pool, + storage_overrides={"is_mla_model": True}, + ) + + pool.page_num = 0 + with pytest.raises(ValueError, match="positive page_size and page_num"): + _candidate("mha", "page_first", pool=pool) + + +def test_fragment_registration_rejects_invalid_sample_metadata() -> None: + pool = _make_pool("mha", "page_first") + pool.get_page_buffer_meta = lambda indices: ([1], [8]) + with pytest.raises(ValueError, match="exactly one page"): + _candidate("mha", "page_first", pool=pool) + + pool.get_page_buffer_meta = lambda indices: ([1, 2], [8, 9]) + with pytest.raises(ValueError, match="equal K/V"): + _candidate("mha", "page_first", pool=pool) + + +def test_registration_is_idempotent_only_for_same_pool() -> None: + session = _FakeSession() + source = _wrapped_config() + store = _store_without_registry(source, session=session) + pool = _make_pool("mha", "page_first") + + store.register_mem_pool_host(pool) + registered = store._registered + store.register_mem_pool_host(pool) + assert store._registered is registered + assert store.mem_pool_host is pool + + with pytest.raises(ValueError, match="different HostPool"): + store.register_mem_pool_host(_make_pool("mha", "page_first")) + + +def test_registration_constructor_claim_is_final_local_step( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import sglang.srt.runtime_context as runtime_context + from sglang.srt.mem_cache.storage.tensorcast_store import tensorcast_store + + session = _FakeSession() + claim_calls: list[object] = [] + + monkeypatch.setattr( + runtime_context, + "get_parallel", + lambda: SimpleNamespace(world_rank=0, world_size=1), + ) + + def claim(options: object, *, owner: TensorcastStore) -> _FakeSession: + assert owner._storage_config.model_name == "fallback-model" + assert owner._registered is None + assert owner._session_options is options + claim_calls.append(owner) + return session + + monkeypatch.setattr(tensorcast_store, "claim_tensorcast_store_session", claim) + store = TensorcastStore(_storage_config()) + + assert store.session is session + assert claim_calls == [store] + + +def test_identity_inputs_are_canonical_and_transfer_mode_independent() -> None: + source = _wrapped_config( + namespace=" Tenant ", + model_id=" Model/A ", + model_version=" Release-1 ", + ) + config = normalize_tensorcast_config(source) + storage_overrides = { + "tp_rank": 2, + "tp_size": 4, + "pp_rank": 1, + "pp_size": 2, + } + registered = _candidate( + "mha", + "page_first_direct", + storage_overrides=storage_overrides, + config=config, + ) + logical_key = "0123456789abcdef" * 4 + specs = _expand_artifact_specs(registered, [logical_key]) + + assert registered.layout_id == "ff1_pfd_f32_p2_mha" + assert registered.keyspace.namespace == "Tenant" + assert registered.keyspace.model_id == "Model/A" + assert registered.keyspace.model_version == "Release-1" + assert registered.keyspace.engine == "sglang" + assert specs[0].engine_key.hex() == ( + "0fb433b34c33e3a7ae1cfdf0c0b5017a58695d3fe63cde10e9d388572381aae1" + ) + assert specs[1].engine_key.hex() == ( + "23e9888f2289481bd631f48f7518cba486e75bc06203c077403d4797ffb37adc" + ) + assert len(specs[0].engine_key) == len(specs[1].engine_key) == 32 + assert tuple(spec.byte_length for spec in specs) == (16, 16) + assert ARTIFACT_LAYOUT_SCHEMA_VERSION == "full-fragment-v1" + + scratch = _candidate("mha", "page_first", config=config) + allocator_config = normalize_tensorcast_config( + _wrapped_config( + transfer_mode="allocator", + namespace=" Tenant ", + model_id=" Model/A ", + model_version=" Release-1 ", + ) + ) + allocator = _candidate("mha", "page_first", config=allocator_config) + assert _expand_artifact_specs(scratch, ["same-key"]) == _expand_artifact_specs( + allocator, ["same-key"] + ) + + +@pytest.mark.parametrize( + ("dtype", "dtype_token"), + [ + (torch.bfloat16, "bf16"), + (torch.float16, "f16"), + (torch.float32, "f32"), + (torch.float64, "f64"), + (torch.uint8, "u8"), + ], +) +def test_identity_layout_uses_compact_dtype_tokens( + dtype: torch.dtype, + dtype_token: str, +) -> None: + pool = _make_pool("mha", "page_first", dtype=dtype) + registered = _candidate("mha", "page_first", pool=pool) + + assert registered.layout_id == f"ff1_pf_{dtype_token}_p2_mha" + + +def test_identity_mla_uses_pp_scope_and_historical_k_suffix() -> None: + registered = _candidate( + "mla", + "page_first", + storage_overrides={"tp_rank": 3, "tp_size": 8, "pp_rank": 1, "pp_size": 2}, + ) + specs = _expand_artifact_specs(registered, ["opaque"]) + + assert len(specs) == 1 + assert specs[0].engine_key.hex() == ( + "5549b303b603995e8fe9873e7934566218131a9a3f547982cdd809cd6ca0ef27" + ) + assert registered.layout_id == "ff1_pf_f32_p2_mla" + + +def test_identity_rejects_duplicate_logical_keys_and_non_string_keys() -> None: + registered = _candidate("mha", "page_first") + with pytest.raises(ValueError, match="unique fragment identities"): + _expand_artifact_specs(registered, ["duplicate", "duplicate"]) + with pytest.raises(TypeError, match="exact str"): + _expand_artifact_specs(registered, [cast(str, b"not-text")]) + + +@pytest.mark.parametrize( + ("mask", "expected"), + [ + ((True, True, True, True, True, True), 3), + ((False, True, True, True, True, True), 0), + ((True, False, True, True, True, True), 0), + ((True, True, False, True, True, True), 1), + ((True, True, True, True, True, False), 2), + ], +) +def test_exists_folds_fragments_to_leading_page_prefix( + mask: tuple[bool, ...], expected: int +) -> None: + session = _FakeSession(mask) + store = _registered_store(session=session) + + assert store.batch_exists(["a", "b", "c"]) == expected + assert len(session.exists_calls) == 1 + assert len(session.exists_calls[0]) == 6 + + +def test_exists_single_item_and_empty_call_behavior() -> None: + session = _FakeSession((True, True)) + store = _registered_store(session=session) + + assert store.exists("one") is True + assert len(session.exists_calls) == 1 + session.existence_mask = (False, True) + assert store.exists("one") is False + assert len(session.exists_calls) == 2 + assert store.batch_exists([]) == 0 + assert len(session.exists_calls) == 2 + + +def test_exists_rejects_result_length_mismatch() -> None: + with pytest.raises(ValueError, match="result length"): + _leading_complete_page_count((True,), page_count=1, fragments_per_page=2) + + +def test_scratch_capacity_accepts_exact_and_rejects_one_byte_short() -> None: + session = _FakeSession() + pool = _make_pool("asymmetric_mha", "page_first", page_num=3) + exact = 3 * (16 + 24) + exact_config = normalize_tensorcast_config(_wrapped_config(scratch_capacity=exact)) + registered = _candidate( + "asymmetric_mha", + "page_first", + pool=pool, + config=exact_config, + ) + _validate_transfer_mode_registration( + registered, + config=exact_config, + session=cast(object, session), + ) + + short_config = normalize_tensorcast_config( + _wrapped_config(scratch_capacity=exact - 1) + ) + with pytest.raises(ValueError) as error_info: + _validate_transfer_mode_registration( + registered, + config=short_config, + session=cast(object, session), + ) + message = str(error_info.value) + for expected_diagnostic in ( + f"configured_capacity_bytes={exact - 1}", + f"required_capacity_bytes={exact}", + "maximum_logical_batch_pages=3", + "page_bytes=40", + "pool_type=AsymmetricMHATokenToKVPoolHost", + "layout=page_first", + "component_byte_lengths=(16, 24)", + ): + assert expected_diagnostic in message + + +def test_scratch_capacity_caps_page_count_at_storage_batch_size() -> None: + session = _FakeSession() + pool = _make_pool("mla", "page_first", page_num=STORAGE_BATCH_SIZE + 1) + required = STORAGE_BATCH_SIZE * 24 + config = normalize_tensorcast_config(_wrapped_config(scratch_capacity=required)) + registered = _candidate("mla", "page_first", pool=pool, config=config) + + _validate_transfer_mode_registration( + registered, + config=config, + session=cast(object, session), + ) + + +@pytest.mark.parametrize( + ("family", "layout", "fragments_per_page"), + [ + ("mha", "page_first_direct", 2), + ("mla", "page_first", 1), + ], +) +def test_span_batch_get_v1_and_batch_set_v1_construct_ordered_transfers( + family: str, + layout: str, + fragments_per_page: int, +) -> None: + fragment_count = 2 * fragments_per_page + session = _FakeSession( + get_mask=(True,) * fragment_count, + put_mask=(True,) * fragment_count, + ) + pool = _make_pool(family, layout) + source = _wrapped_config() + store = _store_without_registry(source, session=session, family=family) + store.register_mem_pool_host(pool) + keys = ["page-a", "page-b"] + host_indices = torch.tensor([2, 3, 4, 5], dtype=torch.int64) + expected_pointers, expected_lengths = pool.get_page_buffer_meta(host_indices) + + assert store.batch_set_v1(keys, host_indices) == [True, True] + assert store.batch_get_v1(keys, host_indices) == [True, True] + assert len(session.put_calls) == len(session.get_calls) == 1 + + for call in (session.put_calls[0], session.get_calls[0]): + assert len(call) == fragment_count + assert [transfer.span.address for transfer in call] == expected_pointers + assert [transfer.span.byte_length for transfer in call] == expected_lengths + assert [transfer.artifact.byte_length for transfer in call] == expected_lengths + + registered = store._require_registered() + fragments = _expand_transfer_fragments(registered, keys, host_indices) + assert [fragment.host_address for fragment in fragments] == expected_pointers + for fragment_index, fragment in enumerate(fragments): + schema_index = fragment_index % fragments_per_page + root = registered.roots[schema_index] + assert fragment.owner is root + assert fragment.host_address is not None + assert fragment.host_address - root.data_ptr() >= 0 + + +def test_span_asymmetric_mha_retains_distinct_component_roots_and_lengths() -> None: + session = _FakeSession(put_mask=(True, True)) + pool = _make_pool("asymmetric_mha", "page_first_direct") + store = _store_without_registry(_wrapped_config(), session=session) + store.register_mem_pool_host(pool) + host_indices = torch.tensor([2, 3], dtype=torch.int64) + + assert store.batch_set_v1(["page"], host_indices) == [True] + fragments = _expand_transfer_fragments( + store._require_registered(), ["page"], host_indices + ) + assert tuple(fragment.owner for fragment in fragments) == ( + pool.k_buffer, + pool.v_buffer, + ) + assert tuple(fragment.byte_length for fragment in fragments) == (16, 24) + assert tuple(transfer.span.address for transfer in session.put_calls[0]) == tuple( + fragment.host_address for fragment in fragments + ) + + +def test_folding_get_normalizes_tail_after_first_incomplete_page() -> None: + session = _FakeSession( + get_mask=(True, True, False, True, True, True), + put_mask=(True, True, False, True, True, True), + ) + store = _registered_store(session=session) + keys = ["a", "b", "c"] + host_indices = torch.arange(6, dtype=torch.int64) + + assert store.batch_set_v1(keys, host_indices) == [True, False, True] + assert store.batch_get_v1(keys, host_indices) == [True, False, False] + assert _fold_fragment_mask( + session.get_mask, + page_count=3, + fragments_per_page=2, + ) == (True, False, True) + assert _normalize_leading_page_mask((True, False, True)) == ( + True, + False, + False, + ) + + +def test_batch_get_v1_and_batch_set_v1_empty_calls_skip_session_rpc() -> None: + session = _FakeSession() + store = _registered_store(session=session) + empty_indices = torch.empty(0, dtype=torch.int64) + + assert store.batch_get_v1([], empty_indices) == [] + assert store.batch_set_v1([], empty_indices) == [] + assert session.get_calls == [] + assert session.put_calls == [] + + +def test_failure_runtime_metadata_validation_disables_before_session_rpc() -> None: + def make_case() -> tuple[_FakeSession, HostKVCache, TensorcastStore]: + session = _FakeSession(put_mask=(True, True)) + pool = _make_pool("mha", "page_first") + store = _store_without_registry(_wrapped_config(), session=session) + store.register_mem_pool_host(pool) + return session, pool, store + + session, _, store = make_case() + assert store.batch_set_v1(["page"], torch.tensor([0], dtype=torch.int64)) == [False] + assert store._disabled is True + assert session.put_calls == [] + + session, _, store = make_case() + assert store.batch_set_v1(["page"], torch.tensor([[0, 1]], dtype=torch.int64)) == [ + False + ] + assert store._disabled is True + assert session.put_calls == [] + + session, pool, store = make_case() + pool.get_page_buffer_meta = lambda indices: ([pool.kv_buffer.data_ptr()], [16]) + assert store.batch_set_v1(["page"], torch.tensor([0, 1], dtype=torch.int64)) == [ + False + ] + assert store._disabled is True + assert session.put_calls == [] + + session, pool, store = make_case() + pool.get_page_buffer_meta = lambda indices: ( + [pool.kv_buffer.data_ptr(), pool.v_buffer.data_ptr()], + [16, 15], + ) + assert store.batch_set_v1(["page"], torch.tensor([0, 1], dtype=torch.int64)) == [ + False + ] + assert store._disabled is True + assert session.put_calls == [] + + +def test_scratch_and_allocator_build_identical_transfer_artifacts() -> None: + keys = ["same-page"] + host_indices = torch.tensor([0, 1], dtype=torch.int64) + scratch = _candidate("mha", "page_first") + allocator_config = normalize_tensorcast_config( + _wrapped_config(transfer_mode="allocator") + ) + allocator = _candidate("mha", "page_first", config=allocator_config) + + scratch_fragments = _expand_transfer_fragments(scratch, keys, host_indices) + allocator_fragments = _expand_transfer_fragments(allocator, keys, host_indices) + assert tuple(fragment.engine_key for fragment in scratch_fragments) == tuple( + fragment.engine_key for fragment in allocator_fragments + ) + assert tuple(fragment.byte_length for fragment in scratch_fragments) == tuple( + fragment.byte_length for fragment in allocator_fragments + ) + + +@pytest.mark.parametrize( + ("family", "allocation_count", "fragments_per_page"), + [ + ("mha", 1, 2), + ("asymmetric_mha", 2, 2), + ("mla", 1, 1), + ], +) +def test_allocator_direct_no_copy_multi_region_transfer( + family: str, + allocation_count: int, + fragments_per_page: int, +) -> None: + fragment_count = 2 * fragments_per_page + session = _FakeSession( + get_mask=(True,) * fragment_count, + put_mask=(True,) * fragment_count, + ) + allocator = TensorcastHostTensorAllocator(cast(object, session)) + pool = _make_pool( + family, + "page_first_direct", + allocator=allocator, + allocate_with_allocator=True, + ) + allocated_roots = tuple(call[3] for call in session.allocation_calls) + assert len(allocated_roots) == allocation_count + + source = _wrapped_config(transfer_mode="allocator") + store = _store_without_registry( + source, + session=session, + family="mla" if family == "mla" else "mha", + ) + store.register_mem_pool_host(pool) + registered = store._require_registered() + expected_roots = ( + (allocated_roots[0], allocated_roots[0]) if family == "mha" else allocated_roots + ) + assert len(registered.roots) == len(expected_roots) + assert all( + registered_root is expected_root + for registered_root, expected_root in zip( + registered.roots, expected_roots, strict=True + ) + ) + + keys = ["allocator-page-0", "allocator-page-1"] + host_indices = torch.tensor([0, 1, 2, 3], dtype=torch.int64) + allocation_calls_before_transfer = tuple(session.allocation_calls) + + assert store.batch_set_v1(keys, host_indices) == [True, True] + assert store.batch_get_v1(keys, host_indices) == [True, True] + assert tuple(session.allocation_calls) == allocation_calls_before_transfer + assert len(session.put_calls) == len(session.get_calls) == 1 + + fragments = _expand_transfer_fragments(registered, keys, host_indices) + expected_owners = tuple( + registered.roots[index % fragments_per_page] for index in range(fragment_count) + ) + assert all( + fragment.owner is expected_owner + for fragment, expected_owner in zip(fragments, expected_owners, strict=True) + ) + for call in (session.put_calls[0], session.get_calls[0]): + assert tuple(transfer.span.address for transfer in call) == tuple( + fragment.host_address for fragment in fragments + ) + assert tuple(transfer.span.byte_length for transfer in call) == tuple( + fragment.byte_length for fragment in fragments + ) + + if family == "asymmetric_mha": + assert allocated_roots[0] is pool.k_buffer + assert allocated_roots[1] is pool.v_buffer + assert {fragment.owner.data_ptr() for fragment in fragments} == { + pool.k_buffer.data_ptr(), + pool.v_buffer.data_ptr(), + } + + +def _public_runtime_error( + error_kind: str, + operation: str, +) -> Exception: + if error_kind == "input": + return RegionArtifactInputError("invalid runtime artifact input") + if error_kind == "terminated": + return RegionSessionTerminatedError("unexpected terminal Session") + if error_kind != "failed": + raise ValueError(f"unknown public error kind {error_kind!r}") + operation_kind = { + "exists": RegionSessionOperationKind.EXISTS, + "get": RegionSessionOperationKind.GET_INTO, + "put": RegionSessionOperationKind.PUT_FROM, + }[operation] + return RegionSessionFailedError( + RegionSessionFailure( + code=RegionSessionFailureCode.TRANSPORT, + message="daemon transport failed", + operation_kind=operation_kind, + operation_id=f"{operation}-operation", + occurred_at=datetime.now(timezone.utc), + ) + ) + + +def _session_with_runtime_error( + operation: str, + error: Exception, +) -> _FakeSession: + if operation == "exists": + return _FakeSession(exists_error=error) + if operation == "get": + return _FakeSession(get_error=error) + if operation == "put": + return _FakeSession(put_error=error) + raise ValueError(f"unknown runtime operation {operation!r}") + + +def _invoke_runtime_operation(store: TensorcastStore, operation: str) -> object: + if operation == "exists": + return store.batch_exists(["page"]) + host_indices = torch.tensor([0, 1], dtype=torch.int64) + if operation == "get": + return store.batch_get_v1(["page"], host_indices) + if operation == "put": + return store.batch_set_v1(["page"], host_indices) + raise ValueError(f"unknown runtime operation {operation!r}") + + +@pytest.mark.parametrize("operation", ["exists", "get", "put"]) +@pytest.mark.parametrize("error_kind", ["input", "failed", "terminated"]) +def test_failure_public_exceptions_disable_once_and_skip_future_session_calls( + operation: str, + error_kind: str, + caplog: pytest.LogCaptureFixture, +) -> None: + session = _session_with_runtime_error( + operation, + _public_runtime_error(error_kind, operation), + ) + store = _registered_store(session=session) + caplog.set_level(logging.ERROR, logger=TensorcastStore.__module__) + + expected = 0 if operation == "exists" else [False] + assert _invoke_runtime_operation(store, operation) == expected + assert store._disabled is True + transition_records = [ + record + for record in caplog.records + if record.message.startswith("TensorCast L3 unavailable:") + ] + assert len(transition_records) == 1 + expected_category = { + "input": "adapter_input_failure", + "failed": "session_failed", + "terminated": "unexpected_terminated", + }[error_kind] + assert f"category={expected_category}" in transition_records[0].message + + assert store.batch_exists(["later"]) == 0 + later_indices = torch.tensor([0, 1], dtype=torch.int64) + assert store.batch_get_v1(["later"], later_indices) == [False] + assert store.batch_set_v1(["later"], later_indices) == [False] + assert ( + len(session.exists_calls) + len(session.get_calls) + len(session.put_calls) == 1 + ) + assert ( + len( + [ + record + for record in caplog.records + if record.message.startswith("TensorCast L3 unavailable:") + ] + ) + == 1 + ) + + +def test_failure_unexpected_planner_span_session_and_folding_errors_are_contained( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + caplog.set_level(logging.ERROR, logger=TensorcastStore.__module__) + + planner_store = _registered_store(session=_FakeSession()) + monkeypatch.setattr( + planner_store, + "_build_artifact_specs", + lambda keys: (_ for _ in ()).throw(RuntimeError("planner defect")), + ) + assert planner_store.batch_exists(["page"]) == 0 + + span_store = _registered_store(session=_FakeSession()) + + def fail_span( + cls: type[object], + tensor: torch.Tensor, + *, + offset_bytes: int, + byte_length: int, + ) -> object: + del cls, tensor, offset_bytes, byte_length + raise RuntimeError("span construction defect") + + from sglang.srt.mem_cache.storage.tensorcast_store import tensorcast_store + + monkeypatch.setattr( + tensorcast_store.HostMemorySpan, + "from_tensor", + classmethod(fail_span), + ) + indices = torch.tensor([0, 1], dtype=torch.int64) + assert span_store.batch_get_v1(["page"], indices) == [False] + monkeypatch.undo() + + session_store = _registered_store( + session=_FakeSession(put_error=RuntimeError("Session boundary defect")) + ) + assert session_store.batch_set_v1(["page"], indices) == [False] + + folding_store = _registered_store(session=_FakeSession(get_mask=(True,))) + assert folding_store.batch_get_v1(["page"], indices) == [False] + + transition_records = [ + record + for record in caplog.records + if record.message.startswith("TensorCast L3 unavailable:") + ] + assert len(transition_records) == 4 + assert all( + "category=adapter_exception" in record.message for record in transition_records + ) + + +def test_failure_ordinary_item_misses_do_not_disable_adapter() -> None: + session = _FakeSession( + existence_mask=(False, False), + get_mask=(False, False), + put_mask=(False, False), + ) + store = _registered_store(session=session) + indices = torch.tensor([0, 1], dtype=torch.int64) + + assert store.batch_exists(["page"]) == 0 + assert store.batch_get_v1(["page"], indices) == [False] + assert store.batch_set_v1(["page"], indices) == [False] + assert store._disabled is False + + session.existence_mask = (True, True) + session.get_mask = (True, True) + session.put_mask = (True, True) + assert store.batch_exists(["page"]) == 1 + assert store.batch_get_v1(["page"], indices) == [True] + assert store.batch_set_v1(["page"], indices) == [True] + + +def test_concurrent_failure_publication_gate_suppresses_racing_success() -> None: + class BlockingSession(_FakeSession): + def __init__(self) -> None: + super().__init__() + self.exists_started = threading.Event() + self.release_exists = threading.Event() + + def batch_exists( + self, + specs: Sequence[object], + ) -> RegionArtifactExistsResult: + self.exists_calls.append(tuple(specs)) + self.exists_started.set() + if not self.release_exists.wait(timeout=5): + raise RuntimeError("test did not release blocked exists") + return RegionArtifactExistsResult( + existence_mask=(True, True), + rpc_elapsed_s=0.0, + ) + + def batch_put_from( + self, + transfers: Sequence[object], + ) -> RegionArtifactTransferResult: + self.put_calls.append(tuple(transfers)) + raise RegionArtifactInputError("disable while exists is in flight") + + session = BlockingSession() + store = _registered_store(session=session) + published: list[int] = [] + exists_thread = threading.Thread( + target=lambda: published.append(store.batch_exists(["page"])) + ) + exists_thread.start() + assert session.exists_started.wait(timeout=5) + + indices = torch.tensor([0, 1], dtype=torch.int64) + assert store.batch_set_v1(["page"], indices) == [False] + session.release_exists.set() + exists_thread.join(timeout=5) + + assert not exists_thread.is_alive() + assert published == [0] + assert len(session.exists_calls) == len(session.put_calls) == 1 + + +def test_worker_callbacks_survive_get_and_backup_failures() -> None: + indices = torch.tensor([0, 1], dtype=torch.int64) + operation = SimpleNamespace(request_id="worker-operation") + + get_store = _registered_store( + session=_FakeSession(get_error=RuntimeError("prefetch failure")) + ) + get_controller = object.__new__(HiCacheController) + get_controller.storage_backend = get_store + assert ( + get_controller._page_get_zero_copy( + operation, + ["page"], + indices, + ) + == 0 + ) + + put_store = _registered_store( + session=_FakeSession(put_error=RuntimeError("backup failure")) + ) + put_controller = object.__new__(HiCacheController) + put_controller.storage_backend = put_store + assert put_controller._page_set_zero_copy(["page"], indices) is False + + +def test_controller_tensorcast_factory_and_host_pool_group_use_v1_anchor( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from sglang.srt.mem_cache.storage.backend_factory import StorageBackendFactory + + registry_entry = StorageBackendFactory._registry["tensorcast"] + assert registry_entry["module_path"] == ( + "sglang.srt.mem_cache.storage.tensorcast_store.tensorcast_store" + ) + assert registry_entry["class_name"] == "TensorcastStore" + + class FakeStorageBackend: + def __init__(self) -> None: + self.registered: list[HostKVCache] = [] + + def register_mem_pool_host(self, pool: HostKVCache) -> None: + self.registered.append(pool) + + backend = FakeStorageBackend() + + def create_backend( + cls: type[StorageBackendFactory], + backend_name: str, + storage_config: object, + mem_pool_host: HostKVCache, + **kwargs: object, + ) -> FakeStorageBackend: + del cls, storage_config, mem_pool_host, kwargs + assert backend_name == "tensorcast" + return backend + + monkeypatch.setattr( + StorageBackendFactory, + "create_backend", + classmethod(create_backend), + ) + + anchor = _make_pool("mha", "page_first") + anchor.can_use_write_back_jit = False + anchor.device = "cpu" + group = HostPoolGroup( + [ + PoolEntry( + name=PoolName.KV, + host_pool=anchor, + device_pool=object(), + layer_mapper=lambda layer_id: layer_id, + is_primary_index_anchor=True, + ) + ] + ) + controller = object.__new__(HiCacheController) + controller.enable_storage = False + controller.storage_stop_event = threading.Event() + controller.storage_host_pool = group.anchor_entry.host_pool + controller.mem_pool_host = group + controller.host_memory_mode = "cache" + controller.page_size = anchor.page_size + controller._stop_storage_threads = lambda: None + controller._start_storage_threads = lambda: None + controller._create_sync_groups = lambda: [] + controller._generate_storage_config = lambda model_name, extra_config: ( + SimpleNamespace(is_mla_model=False, tp_rank=0, extra_config=extra_config or {}) + ) + + controller.attach_storage_backend( + "tensorcast", + model_name="model", + storage_backend_extra_config={"tensorcast": {}}, + ) + + assert backend.registered == [anchor] + assert controller.page_get_func == controller._page_get_zero_copy + assert controller.page_set_func == controller._page_set_zero_copy + + +def test_registration_allocator_requires_thin_allocator_and_same_session() -> None: + session = _FakeSession() + other_session = _FakeSession() + config = normalize_tensorcast_config(_wrapped_config(transfer_mode="allocator")) + pool = _make_pool( + "mha", + "page_first", + allocator=TensorcastHostTensorAllocator(cast(object, session)), + ) + registered = _candidate("mha", "page_first", pool=pool, config=config) + + _validate_transfer_mode_registration( + registered, + config=config, + session=cast(object, session), + ) + with pytest.raises(ValueError, match="same Session"): + _validate_transfer_mode_registration( + registered, + config=config, + session=cast(object, other_session), + ) + + pool.allocator = HostTensorAllocator() + with pytest.raises(ValueError, match="TensorcastHostTensorAllocator"): + _validate_transfer_mode_registration( + registered, + config=config, + session=cast(object, session), + ) + + +def test_failure_containment_keeps_unsupported_value_v1_v2_and_clear_apis_explicit() -> ( + None +): + store = _registered_store(session=_FakeSession()) + calls = ( + lambda: store.get("key"), + lambda: store.batch_get(["key"]), + lambda: store.set("key"), + lambda: store.batch_set(["key"]), + lambda: store.batch_exists_v2(["key"]), + lambda: store.batch_get_v2([]), + lambda: store.batch_set_v2([]), + store.clear, + ) + for call in calls: + with pytest.raises(NotImplementedError): + call() + + assert ( + TensorcastStore.register_mem_host_pool_v2 + is HiCacheStorage.register_mem_host_pool_v2 + ) + sidecar = object() + store.register_mem_host_pool_v2(cast(HostKVCache, sidecar), PoolName.MAMBA) + assert store.registered_pools == {PoolName.MAMBA: sidecar} + + +def _store_without_registry( + source: dict[str, object], + *, + session: _FakeSession, + family: str = "mha", +) -> TensorcastStore: + store = object.__new__(TensorcastStore) + store._storage_config = _storage_config(family=family, extra_config=source) + store._tensorcast_config = normalize_tensorcast_config(source) + store._registered = None + store._session = cast(object, session) + store._availability_lock = threading.Lock() + store._disabled = False + store._failure_logged = False + return store + + +def _registered_store(*, session: _FakeSession) -> TensorcastStore: + source = _wrapped_config() + store = _store_without_registry(source, session=session) + store.register_mem_pool_host(_make_pool("mha", "page_first")) + return store