Feat: Add TensorCast storage as a new HiCache backend (#27265)
This commit is contained in:
Yuhan Zhou
@@ -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
|
||||
|
||||
@@ -601,6 +601,7 @@ class HiCacheController:
|
||||
"nixl",
|
||||
"simm",
|
||||
"mori",
|
||||
"tensorcast",
|
||||
]
|
||||
) or (
|
||||
self.storage_backend_type == "dynamic"
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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",
|
||||
)
|
||||
|
||||
@@ -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 <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 <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.
|
||||
@@ -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: ""
|
||||
+389
@@ -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
|
||||
# <daemon_state_dir>/local_handle.sock for same-pod/local SDKs.
|
||||
# daemon_state_dir defaults to:
|
||||
# $TENSORCAST_HOME/hosts/<host_id>/sessions/<session_id>/session
|
||||
# (or ~/.tensorcast/hosts/<host_id>/sessions/<session_id>/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/<host_id>/runtime/daemons/<daemon_id>/local_handle.sock
|
||||
# If the selected socket path exceeds AF_UNIX limits, the daemon falls back to:
|
||||
# $TENSORCAST_HOME/uds/lh-<hash>.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
|
||||
@@ -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)
|
||||
@@ -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
|
||||
+301
@@ -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
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user