From 50aa97da45967bb94d58c7259a180f55b392c0be Mon Sep 17 00:00:00 2001 From: linhu-nv <141609318+linhu-nv@users.noreply.github.com> Date: Tue, 7 Jul 2026 15:35:52 +0800 Subject: [PATCH] Feat/flexkv main connector (#29701) --- python/sglang/srt/managers/scheduler.py | 2 +- python/sglang/srt/mem_cache/registry.py | 14 + .../srt/mem_cache/storage/flexkv/README.md | 427 ++++++++ .../srt/mem_cache/storage/flexkv/__init__.py | 87 ++ .../storage/flexkv/example_config_mp.yaml | 38 + .../mem_cache/storage/flexkv/flexkv_comm.py | 662 +++++++++++++ .../storage/flexkv/flexkv_connector.py | 925 ++++++++++++++++++ .../storage/flexkv/flexkv_radix_cache.py | 511 ++++++++++ .../storage/flexkv/verify_outputs.py | 180 ++++ python/sglang/srt/server_args.py | 26 + 10 files changed, 2871 insertions(+), 1 deletion(-) create mode 100644 python/sglang/srt/mem_cache/storage/flexkv/README.md create mode 100644 python/sglang/srt/mem_cache/storage/flexkv/__init__.py create mode 100644 python/sglang/srt/mem_cache/storage/flexkv/example_config_mp.yaml create mode 100644 python/sglang/srt/mem_cache/storage/flexkv/flexkv_comm.py create mode 100644 python/sglang/srt/mem_cache/storage/flexkv/flexkv_connector.py create mode 100644 python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py create mode 100644 python/sglang/srt/mem_cache/storage/flexkv/verify_outputs.py diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index bfdd72bd5..4c7149735 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2759,7 +2759,7 @@ class Scheduler( for req in ready_grammar_requests: self._add_request_to_queue(req) - if self.enable_hierarchical_cache: + if self.enable_hierarchical_cache or self.server_args.enable_flexkv: self.tree_cache.check_hicache_events() if self.enable_priority_preemption or self.is_hybrid_swa: diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index 5b8b9e9de..e4fc6a8c4 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/python/sglang/srt/mem_cache/registry.py @@ -143,6 +143,20 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache: tp_group=ctx.tp_group, ) + if server_args.enable_flexkv: + # Importing the package side-effect registers the explicit + # ``--radix-cache-backend=flexkv`` factory; we then call the + # factory directly so --enable-flexkv stands on its own. + import os + + from sglang.srt.mem_cache.storage.flexkv import _flexkv_factory + + # Honor a CLI --flexkv-config-file by forwarding it via the env + # var that FlexKV's config loader actually reads. + if server_args.flexkv_config_file and not os.environ.get("FLEXKV_CONFIG_PATH"): + os.environ["FLEXKV_CONFIG_PATH"] = server_args.flexkv_config_file + return _flexkv_factory(ctx) + from sglang.srt.mem_cache.radix_cache import RadixCache return RadixCache(params) diff --git a/python/sglang/srt/mem_cache/storage/flexkv/README.md b/python/sglang/srt/mem_cache/storage/flexkv/README.md new file mode 100644 index 000000000..92e55a3d5 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/flexkv/README.md @@ -0,0 +1,427 @@ +# FlexKV ↔ sglang integration + +A `RadixCache` subclass that routes sglang's host-tier KV cache through a +FlexKV [`KVManager`](https://github.com/taco-project/FlexKV) (CPU / SSD / +Remote offload). Same integration pattern as +[`LMCRadixCache`](../lmcache/README.md): `FlexKVRadixCache` overrides +`match_prefix` / `init_load_back` / `cache_finished_req` / `evict`; a +`FlexKVConnector` façade talks to `KVManager`, `KVTPClient`, and a +3-axis (PP × CP × TP) sync context. + +--- + +## Quick start (single H20, single GPU, Qwen3-8B) + +This walks through everything the verification on H20-GPU-11 actually +exercised. Adjust paths / model / GPU as needed. + +### 1. Prereqs + +* `lmsysorg/sglang:dev` (or any sglang container with CUDA 12.x + torch 2.10+). +* This sglang fork (branch `feat/flexkv-main-connector`) and FlexKV + (branch `main`) checked out somewhere reachable from the container + — e.g. `/raid/fly/sglang-connector-dir/{sglang,FlexKV}`. Verified + against FlexKV main at `aa74e39` (PR #184); older commits down to + the layerwise integration also work. + +### 2. Start a container with both repos mounted + +```bash +docker run -d --name flexkv-sglang \ + --gpus all --ipc=host --network host \ + --shm-size=32g --cap-add SYS_NICE --cap-add IPC_LOCK \ + -v /raid/fly:/raid/fly \ + --workdir /raid/fly/sglang-connector-dir \ + --entrypoint "" \ + lmsysorg/sglang:dev sleep infinity + +docker exec flexkv-sglang bash -c " + apt-get update -qq && + apt-get install -y numactl libnuma-dev libxxhash-dev liburing-dev cmake ninja-build +" +``` + +### 3. Install sglang fork (editable) + FlexKV + +```bash +docker exec flexkv-sglang bash -c ' + set -e + git config --global --add safe.directory "*" + + # sglang fork: install in editable mode, replacing the prebuilt sglang + cd /raid/fly/sglang-connector-dir/sglang + pip install --no-deps -e python + + # FlexKV: pin to main, init the xxHash submodule, debug C++ build. + cd /raid/fly/sglang-connector-dir/FlexKV + git checkout main && git pull --ff-only + git submodule update --init third_party/xxHash + pip install -q cython ninja pybind11 + FLEXKV_ENABLE_METRICS=0 bash build.sh --debug + + # Smoke check + python3 -c " +import sglang, flexkv +from flexkv.kvmanager import KVManager +from sglang.srt.mem_cache.storage.flexkv import flexkv_comm +from sglang.srt.mem_cache.registry import registered_radix_cache_backends +import sglang.srt.mem_cache.storage.flexkv # registers +print(\"flexkv ok\", flexkv.__file__) +print(\"sglang ok\", sglang.__file__) +print(\"registered backends:\", registered_radix_cache_backends()) +" +' +``` + +If the build hangs on `pip install sglang-kernel`, see +[Troubleshooting](#troubleshooting). + +### 4. Minimal FlexKV YAML + +```yaml +# /raid/fly/sglang-connector-dir/flexkv_min.yaml +cpu_cache_gb: 16 +``` + +That's enough to enable a 16 GiB CPU offload pool. See +[`example_config_mp.yaml`](example_config_mp.yaml) for SSD / remote / +distributed knobs. + +### 5. Launch the server (MP / synchronous mode) + +```bash +docker exec -d flexkv-sglang bash -c ' + cd /raid/fly/sglang-connector-dir + CUDA_VISIBLE_DEVICES=0 \ + SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK=1 \ + python3 -m sglang.launch_server \ + --model-path /raid/fly/model/Qwen3-8B \ + --port 30000 --tp-size 1 \ + --enable-flexkv \ + --flexkv-config-file /raid/fly/sglang-connector-dir/flexkv_min.yaml \ + --mem-fraction-static 0.45 --max-running-requests 8 \ + > /tmp/sglang.log 2>&1 +' +``` + +`SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK=1` bypasses the prebuilt +`sglang-kernel` version assertion (the `lmsysorg/sglang:dev` image ships +0.4.2.post2; main expects ≥ 0.4.3). Not a FlexKV-specific issue; +remove when the container image is refreshed. + +Wait ~2 min for the model load + CUDA graph capture. Confirm with: + +```bash +docker exec flexkv-sglang bash -c ' + grep -E "fired up|Connector ready" /tmp/sglang.log | tail -2 +' +``` + +Expected (key lines): + +``` +[FlexKV] Connector ready ...: layerwise=False, prefetch=False +The server is fired up and ready to roll! +``` + +### 6. Send a request and observe a cache hit + +```bash +# First call: priming — fresh prefill, FlexKV stores the prefix. +docker exec flexkv-sglang bash -c ' + curl -s http://127.0.0.1:30000/generate -X POST \ + -H "Content-Type: application/json" \ + -d "{\"text\": \"The capital of France is\", + \"sampling_params\": {\"max_new_tokens\": 5, \"temperature\": 0}}" +' + +# Flush the GPU radix (FlexKV CPU pool keeps the data) and re-send. +docker exec flexkv-sglang bash -c ' + curl -s http://127.0.0.1:30000/flush_cache -X POST + curl -s http://127.0.0.1:30000/generate -X POST \ + -H "Content-Type: application/json" \ + -d "{\"text\": \"The capital of France is\", + \"sampling_params\": {\"max_new_tokens\": 5, \"temperature\": 0}}" +' +``` + +Look at the second response's `meta_info`: + +```json +"cached_tokens": 4, +"cached_tokens_details": { "device": 0, "host": 4 }, +``` + +`host: 4` confirms the bytes came back from FlexKV's CPU pool. The +server log should also show a matching D2H/H2D bandwidth line: + +``` +[FLEXKV] ... H2D transfer request: N finished transfer data size: 0.0xx GB ... 30+ GB/s +``` + +### 7. Layerwise mode + +Add `FLEXKV_ENABLE_LAYERWISE_TRANSFER=1` before `python3 -m +sglang.launch_server`. Everything else is identical. On the second +request you'll see `cached_tokens_details: {"device": N, "host": 0}` +(in IP mode the load happens inside `match_prefix` so sglang accounts +for it as device-side) and a log line `LAYERWISE transfer request: N +finished ...`. The startup log will also include +`[FlexKV] Eventfd handshake complete ... counters=3 layers=`. + +--- + +## Correctness verification + +Numerical match against a no-FlexKV baseline (greedy decoding, +deterministic). Scripts are in this repo's testing notes; the canonical +two are reproduced below. + +```bash +# Phase 1: capture the no-FlexKV baseline. +docker exec -d flexkv-sglang bash -c ' + CUDA_VISIBLE_DEVICES=0 SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK=1 \ + python3 -m sglang.launch_server \ + --model-path /raid/fly/model/Qwen3-8B --port 30000 --tp-size 1 \ + --mem-fraction-static 0.45 > /tmp/sglang.log 2>&1 +' +# ... wait until ready ... +docker exec flexkv-sglang python3 /raid/fly/sglang-connector-dir/sglang/python/sglang/srt/mem_cache/storage/flexkv/verify_outputs.py --phase baseline +docker exec flexkv-sglang bash -c "pkill -9 -f launch_server; sleep 3" + +# Phase 2: relaunch with --enable-flexkv and compare. +docker exec -d flexkv-sglang bash -c ' + CUDA_VISIBLE_DEVICES=0 SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK=1 \ + python3 -m sglang.launch_server \ + --model-path /raid/fly/model/Qwen3-8B --port 30000 --tp-size 1 \ + --enable-flexkv --flexkv-config-file /raid/fly/sglang-connector-dir/flexkv_min.yaml \ + --mem-fraction-static 0.45 > /tmp/sglang.log 2>&1 +' +# ... wait until ready ... +docker exec flexkv-sglang python3 /raid/fly/sglang-connector-dir/sglang/python/sglang/srt/mem_cache/storage/flexkv/verify_outputs.py --phase test +``` + +Expected last line: `Total mismatches: 0`. Each prompt is run twice +(R1 fresh / R2 after `flush_cache`); both R1 and R2 outputs must +byte-equal the baseline. + +Repeat the Phase-2 launch with `FLEXKV_ENABLE_LAYERWISE_TRANSFER=1` +to validate the layerwise path. + +--- + +## Selecting the backend + +Two equivalent CLI flags: + +```bash +# Auto-selection chain (matches --enable-lmcache style) +python3 -m sglang.launch_server --enable-flexkv \ + --flexkv-config-file /path/to/flexkv_config.yaml ... + +# Explicit registry path +python3 -m sglang.launch_server --radix-cache-backend flexkv \ + --flexkv-config-file /path/to/flexkv_config.yaml ... +``` + +Either flag also sets `FLEXKV_CONFIG_PATH` so you can omit +`--flexkv-config-file` and configure FlexKV purely through env vars. + +--- + +## Modes + +### MP (synchronous, default) + +* `match_prefix` calls `FlexKVConnector.lookup_kv` only. +* When `host_hit_length > 0`, the scheduler later calls + `init_load_back`, which allocates the uncached slots and fires + `retrieve_kv` (FlexKV `launch` + `wait`). +* `cache_finished_req` runs `put_match` + `launch` and stashes the + in-flight FlexKV task id. Source-node lock is held until + `check_completed_stores` (called from `check_hicache_events` / + `evict`) signals completion. + +This is the path you'll use under any non-trivial deployment topology +(DP > 1, multi-instance, multi-node, ...). + +### IP / layerwise (`FLEXKV_ENABLE_LAYERWISE_TRANSFER=1`) + +* `match_prefix` allocates the uncached slots and fires + `start_load_kv_layerwise` immediately. +* A `FlexKVLayerDoneCounter` is registered onto sglang's KV pool via + `register_layer_transfer_counter`; the per-layer hook blocks each + forward layer on its own eventfd until the FlexKV transfer worker + signals the layer is staged. +* Layerwise mode requires the FlexKV transfer worker's UDS socket + (`/tmp/flexkv_layerwise_eventfd.sock` by default) to be reachable — + the connector handshakes with it at startup. The socket path is + computed by FlexKV's `build_layerwise_eventfd_socket_path` from the + same dp/pp/instance settings, so configuration is taken care of as + long as you launch FlexKV consistently. + +--- + +## Files + +* `flexkv_radix_cache.py` — `FlexKVRadixCache(RadixCache)`. Overrides + `match_prefix`, `init_load_back`, `cache_finished_req`, `evict`, + `check_hicache_events`, `reset`. +* `flexkv_connector.py` — `FlexKVConnector`. Owns the `KVManager`, + `KVTPClient`, and the cross-rank sync context. Public methods: + `lookup_kv`, `retrieve_kv`, `start_load_kv_layerwise`, `store_kv`, + `check_completed_stores`, `prefetch_async`, … +* `flexkv_comm.py` — `FlexKVComm` (3-axis PP × CP × TP sync built on + torch.distributed) + the eventfd / `SCM_RIGHTS` shims used by the + layerwise transfer UDS handshake. **`FlexKVLayerLoadingEvent` here + carries the layerwise correctness fix** (drain stale eventfd + signals on reset, switch `wait` to `select.select` to keep blocking + semantics on a NONBLOCK fd). +* `__init__.py` — registers the `"flexkv"` factory with + `sglang.srt.mem_cache.registry`. + +--- + +## TP / PP / CP / DP + +FlexKV runs one `KVManager` per DP route (= +`instance_id * dp_size + dp_rank`). Every other rank in the same +fan-out is the "sync follower" — `FlexKVComm` broadcasts the +leader's lookup / store decisions via gloo CPU groups so non-leader +ranks know which task ids and slot mappings to use. + +Supported: + +* **TP** (any size) — typical sglang topology. +* **DP** (`dp_size > 1`) and multi-instance — FlexKV automatically + switches its `KVManager` to server-client mode. +* **PP** (`pp_size > 1`) — including cross-node PP. The PP receiver + forwards its slot mappings back to FlexKV's + `TransferManagerOnRemote` via the same ZMQ channel used for GPU + registration. +* **CP** (`attn_cp_size > 1`) — sync handled symmetrically with TP. +* **DP attention** (`enable_dp_attention=True`) — the inner + `attn_tp_size` is what FlexKV uses for register-side routing. + +--- + +## Environment variables + +* `FLEXKV_CONFIG_PATH` — full FlexKV YAML / JSON config (also set + automatically by `--flexkv-config-file`). +* `FLEXKV_ENABLE_LAYERWISE_TRANSFER` — `1` to enable layerwise mode. +* `FLEXKV_LAYERWISE_EVENTFD_SOCKET` — UDS socket path (default + `/tmp/flexkv_layerwise_eventfd.sock`); auto-suffixed per + `(pp_rank, dp_client_id)` when those dims are > 1. +* `FLEXKV_MASTER_HOST` / `FLEXKV_MASTER_PORTS` — multi-node master + endpoint for `TransferManagerOnRemote`. Default + `localhost:5556,5557,5558`. With `nnodes > 1` we also fall back to + `server_args.dist_init_addr`'s host. +* `FLEXKV_KV_CACHE_DTYPE` — override KV dtype when sglang uses + `--kv-cache-dtype auto`. +* `SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK` — bypass the prebuilt + `sglang-kernel` version assertion (not FlexKV-specific). + +--- + +## Troubleshooting + +* **`fatal: not a git repository ... third_party/xxHash`** — FlexKV's + build.sh needs an actual git checkout for the submodule. If you + rsync'd FlexKV without `.git/`, sync it: `rsync -az + /path/to/FlexKV/.git/ :/FlexKV/.git/` then + `git config --global --add safe.directory "*"`. +* **`fatal: detected dubious ownership`** — same fix: + `git config --global --add safe.directory "*"`. +* **`xxhash.h: No such file or directory`** — submodule not init'd. + `cd FlexKV && git submodule update --init third_party/xxHash`. +* **`dist/lease_meta_mempool.h: No such file or directory`** — your + rsync excluded `csrc/dist/`. The directory `FlexKV/csrc/dist/` is + source, not a build artifact; re-sync without `--exclude='dist'`. +* **`No module named 'Cython'`** — install: `pip install cython ninja pybind11`. +* **`sglang-kernel is installed with version 0.4.2.post2, which is + less than the minimum required version 0.4.3`** — either run with + `SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK=1` or refresh + `pip install -U sglang-kernel`. The download is ~600 MB and can + take a long time on slow links. +* **`cudaHostRegister failed with error code 100` (cudaErrorNoDevice)** + — happens when the FlexKV transfer subprocess can't init CUDA on + the assigned device. Usually a stuck previous session; restart + the container. +* **`[FlexKV] Waiting for FlexKV ready` loops > 60 s** — the + KVManager subprocess crashed at boot. Check `/tmp/sglang.log` for + the actual stack (usually a CUDA-init or torch-mp issue). +* **Layerwise mode: server hangs at "Eventfd connected attempts=..."** + — the `LayerwiseTransferWorker` hasn't started yet. Wait — it can + take 20-30 s after `Eventfd server created`. If it never advances, + check the FlexKV-side log lines beginning with `[LayerwiseWorker]`. + +--- + +## Status + +* MP (synchronous) path — verified end-to-end on Qwen3-8B (H20-3e): + output byte-equal to no-FlexKV baseline across short / medium / long + prompts. ~30–46 GB/s observed for D2H stores and ~37 GB/s for H2D + loads. +* IP (layerwise) path — verified end-to-end with the fix in + `flexkv_comm.py`. ~7–12 GB/s per-layer (smaller per-call payload). +* PP / CP / DP / multi-node — code paths driven by `FlexKVComm`, + carried over from the production-validated `BaseKVConnector` + integration. Not exercised in single-GPU smoke tests; needs a + multi-node run before shipping. + +### Known limitations + +* Hybrid models (Mamba / SWA / DSV4 indexer auxiliary pools) are not + supported through this connector — only the primary KV pool is + hooked up. HiCache's multi-pool `batch_*_v2` interface would map + here but requires `PoolTransfer` + `PoolHitPolicy` plumbing in + `FlexKVConnector`. +* Write-back acks are per-request (one `dec_lock_ref` per + `cache_finished_req`), not per-page like HiCache's + `flush_write_through_acks`. +* `--radix-cache-backend=flexkv` and `--enable-flexkv` are + mutually equivalent today; we don't yet emit a deprecation + warning if both are set. + +## Benchmarks + +Setup: Qwen3-8B on 1× H20. Server flags: + + --attention-backend triton --mem-fraction-static 0.32 + --max-running-requests 32 --chunked-prefill-size 16384 + --context-length 32000 + +Workload: 120 prompts sampled from +[`princeton-nlp/SWE-bench_Lite_oracle`](https://huggingface.co/datasets/princeton-nlp/SWE-bench_Lite_oracle) +with input length ≤ 28k tokens (p50 = 7088, max = 27961). Two passes — +pass 1 populates the host cache, pass 2 is the measured run. `qps=2.0`, +`concurrency=24`, `max_new_tokens=32`, `temperature=0`. + +### Warm-pass results + +| Config | TTFT avg / p50 / p90 / p99 | E2E p50 | Throughput | Output tok/s | H2D / D2H | +| --- | --- | --- | --- | --- | --- | +| baseline | 6.86 / 8.04 / 9.88 / 10.89 s | 8.15 s | 1.86 req/s | 37.7 | — | +| `--enable-hierarchical-cache` | **0.04 / 0.04 / 0.06 / 0.06 s** | 0.23 s | 2.02 req/s | 40.8 | — | +| `--enable-flexkv` | **0.05 / 0.05 / 0.07 / 0.08 s** | 0.24 s | 2.02 req/s | 40.8 | 86 / 155 | + +Server-side (via `ReqTimeStats` in the sglang log): 76 / 76 non-EOS-immediate +warm-pass requests have `cached_input_len == input_len` for both `hicache` +and `flexkv` (100 % prefix recovery); baseline stays at ~59 tokens +(system-prompt header only). The 86 `H2D transfer` log lines under `flexkv` +confirm the CPU-tier loadbacks actually fired. + +### Output correctness + +Byte-level diff of generated text across 32 prompts, `temperature=0`: + +* baseline: cold pass == warm pass (32 / 32; fully deterministic without cache) +* `hicache`: warm vs baseline warm — 29 / 32 identical, 3 diverge +* `flexkv`: warm vs baseline warm — 29 / 32 identical, 3 diverge (mostly the same 3 as `hicache`) + +The ~10 % divergence at `temperature=0` is the well-known KV-cache-reuse +artifact caused by floating-point non-associativity between "prefill in +place" and "load pre-computed KV" paths; it affects the mainline +`--enable-hierarchical-cache` at the same rate and is not FlexKV-specific. diff --git a/python/sglang/srt/mem_cache/storage/flexkv/__init__.py b/python/sglang/srt/mem_cache/storage/flexkv/__init__.py new file mode 100644 index 000000000..a2bd129b6 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/flexkv/__init__.py @@ -0,0 +1,87 @@ +"""FlexKV-backed RadixCache integration for sglang. + +Two ways to select this backend at server launch: + +1. ``--enable-flexkv`` (default chain in ``default_radix_cache_factory``) +2. ``--radix-cache-backend=flexkv`` (explicit registry path) + +Importing this package registers the explicit name with the registry, +so the second form is available without further wiring. +""" + +from __future__ import annotations + +import logging + +from sglang.srt.mem_cache.registry import register_radix_cache_backend + +logger = logging.getLogger(__name__) + + +def _flexkv_factory(ctx): + """Build a :class:`FlexKVRadixCache` from a ``TreeCacheBuildContext``. + + ``TreeCacheBuildContext`` carries TP rank/size and the TP group + coordinator, but not PP/CP. We pick those up from the global + accessors in :mod:`sglang.srt.distributed.parallel_state`; FlexKV + needs them to fan out lookup/store decisions across the full TP × CP + × PP topology. + """ + from sglang.srt.distributed.parallel_state import ( + get_attn_cp_group, + get_attn_tp_group, + get_pp_group, + ) + from sglang.srt.mem_cache.storage.flexkv.flexkv_radix_cache import ( + FlexKVRadixCache, + ) + + server_args = ctx.server_args + + # PP group is always available; attn TP / attn CP groups may share + # the regular TP group when attn DP is off — that's fine, the + # connector treats size-1 groups as no-ops. + try: + pp_group = get_pp_group() + except (RuntimeError, AssertionError): + pp_group = None + try: + attn_tp_group = get_attn_tp_group() + except (RuntimeError, AssertionError): + attn_tp_group = ctx.tp_group + try: + attn_cp_group = get_attn_cp_group() + except (RuntimeError, AssertionError): + attn_cp_group = None + + # PP / CP ranks: use the group's own rank_in_group view if available; + # fall back to 0 for single-rank dims. + pp_rank = pp_group.rank_in_group if pp_group is not None else 0 + attn_cp_rank = attn_cp_group.rank_in_group if attn_cp_group is not None else 0 + + return FlexKVRadixCache( + params=ctx.params, + model_config=ctx.model_config, + server_args=server_args, + tp_rank=ctx.tp_rank, + tp_size=ctx.tp_size, + # ``dp_rank`` isn't carried on TreeCacheBuildContext or ServerArgs + # at construction time; the connector normalizes ``None`` to 0 + # for the single-DP-rank case that this factory targets. + dp_rank=None, + pp_rank=pp_rank, + attn_cp_rank=attn_cp_rank, + tp_group=ctx.tp_group, + pp_group=pp_group, + attn_tp_group=attn_tp_group, + attn_cp_group=attn_cp_group, + ) + + +try: + register_radix_cache_backend("flexkv", _flexkv_factory) +except ValueError as exc: + # The registry refuses duplicates. Importing this package twice + # (e.g. via both --enable-flexkv and --radix-cache-backend=flexkv) + # is fine — log and move on. + logger.debug("flexkv backend already registered: %s", exc) diff --git a/python/sglang/srt/mem_cache/storage/flexkv/example_config_mp.yaml b/python/sglang/srt/mem_cache/storage/flexkv/example_config_mp.yaml new file mode 100644 index 000000000..25bca1982 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/flexkv/example_config_mp.yaml @@ -0,0 +1,38 @@ +# Example FlexKV YAML config (passed to sglang via --flexkv-config-file). +# +# Equivalent env vars exist for every field — see flexkv/common/config.py +# (UserConfig.from_env). This file is a minimal CPU-only setup; uncomment +# the SSD / Remote / Redis sections to enable those tiers. + +# ---- CPU host-side cache ---------------------------------------------- +# Size of the FlexKV CPU pool. Used to derive `num_cpu_blocks` together +# with the model dtype, head dim, num kv heads, and page size. +cpu_cache_gb: 64 + +# Optional: pin the CPU pool using transparent huge pages. +# use_hugepage_cpu_buffer: false +# use_hugepage_tmp_buffer: false +# hugepage_size_bytes: 2097152 + +# ---- SSD tier --------------------------------------------------------- +# Set ssd_cache_gb > cpu_cache_gb to enable the SSD spill tier. +# ssd_cache_gb: 256 +# ssd_cache_dir: "/mnt/nvme0/flexkv;/mnt/nvme1/flexkv" # ';'-separated for striping +# enable_gds: false # cuFile / GDS path + +# ---- KV cache dtype override ----------------------------------------- +# When sglang is launched with --kv-cache-dtype auto, FlexKV can't tell +# which dtype the actual KV tensors use. Set explicitly here. +# kv_cache_dtype: bfloat16 + +# ---- Peer / distributed sharing -------------------------------------- +# enable_p2p_cpu: false +# enable_p2p_ssd: false +# enable_3rd_remote: false + +# ---- Redis (for distributed metadata / KV sharing) ------------------- +# redis_host: 127.0.0.1 +# redis_port: 6379 +# redis_password: null +# node_ttl_seconds: 60 +# local_ip: 10.0.0.1 diff --git a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_comm.py b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_comm.py new file mode 100644 index 000000000..00cabf6ae --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_comm.py @@ -0,0 +1,662 @@ +"""Communication helpers for the FlexKV connector. + +FlexKV runs a single KVManager per DP group (typically the TP/CP/PP +sync leader's process). Every other rank in the same KV-cache-sharing +fan-out must be told the leader's decisions: which prefix matched in +FlexKV, which task id the leader allocated, which slot mappings to +send, etc. + +This file provides: + +* ``FlexKVComm`` — a 3-axis (PP × CP × TP) hierarchical sync context + built on torch.distributed (gloo CPU groups). Exposes ``scatter``, + ``scatter_pp``, ``barrier`` and ``all_reduce_min`` plus role flags + (``is_sync_leader`` etc.) that the connector branches on. +* libc / ``eventfd`` shims used by the layerwise transfer worker + socket handshake. +* ``FlexKVLayerLoadingEvent`` and ``FlexKVLayerDoneCounter`` — the + eventfd-backed per-layer completion structures that the FlexKV + layerwise transfer worker signals into. Hooked into sglang's + ``register_layer_transfer_counter`` so each layer's forward waits + for its own host→device copy. +""" + +from __future__ import annotations + +import ctypes +import errno +import logging +import os +import pickle +import socket +import struct +from datetime import timedelta +from typing import Any, Dict, List + +import torch +import torch.distributed as dist + +from sglang.srt.distributed.parallel_state import get_world_group + +logger = logging.getLogger(__name__) + +# PP-channel command tags (used by ``scatter_pp`` payloads). Sender and +# receiver assert on these to catch protocol drift early. +CMD_PUT_META = 2 +CMD_LAYERWISE = 3 +CMD_STORE_COMPLETE = 5 + + +class FlexKVComm: + """3-axis (PP × CP × TP) hierarchical sync for the FlexKV connector. + + Notation: + * "sync leader" is the unique rank that talks to the FlexKV + KVManager: pp_rank=0, attn_cp_rank=0, attn_tp_rank=0. + * "PP stage leader" is the (cp=0, tp=0) rank within a PP stage — + it does cross-PP P2P (``scatter_pp``). + * Every rank participates in collective layers it belongs to. + + Communication strategy: + * P2P (send/recv/isend/irecv) on CPU tensors → ``world_cpu_group`` + (the global gloo group). Sub-group cpu_groups have unreliable + TCP pairs for direct P2P. + * Collectives (all_reduce / barrier) → sglang's sub-group + cpu_groups (fine for collectives). + """ + + # P2P tags. World group is shared with sglang's own P2P, so we pick + # 4-byte tags that won't collide. + _TAG_SCATTER = int.from_bytes(b"FxSc", byteorder="big") + _TAG_PP = int.from_bytes(b"FxPP", byteorder="big") + _TAG_CP = int.from_bytes(b"FxCP", byteorder="big") + _TAG_TP = int.from_bytes(b"FxTP", byteorder="big") + _TAG_PP_AR_MIN = int.from_bytes(b"FxA2", byteorder="big") + _TAG_PP_BARRIER = int.from_bytes(b"FxB2", byteorder="big") + _TAG_PP_BARRIER_BCAST = int.from_bytes(b"FxB3", byteorder="big") + _TAG_AR_BCAST = int.from_bytes(b"FxAR", byteorder="big") + + # Adaptive async-work reaper. gloo's isend Work objects do not auto- + # advance their "completed" state on poll, so a pure poll-based reaper + # leaks. We actively wait() the oldest works with a tiny timeout; + # the watermark grows on stuck reaps (slow / asymmetric peer) and + # shrinks back on clean reaps. + _REAP_HIGH_BASE = 1024 + _REAP_HIGH_MAX = 32768 + _REAP_MAX_DRAIN = 512 + _REAP_PROBE = timedelta(milliseconds=1) + _REAP_LOG_EVERY = 64 + + def __init__( + self, + rank_info, + world_rank: int, + pp_group=None, + attn_tp_group=None, + attn_cp_group=None, + ): + model_config = rank_info.model_config + self.world_rank = world_rank + self._async_works: List = [] + self._reap_high: int = self._REAP_HIGH_BASE + self._reap_calls: int = 0 + self._reap_stuck_total: int = 0 + self._reap_drained_total: int = 0 + + # Accept either GroupCoordinator wrappers (has ``.cpu_group``) or + # raw ProcessGroups. + self.pp_cpu_group = ( + getattr(pp_group, "cpu_group", pp_group) if pp_group is not None else None + ) + self.attn_tp_cpu_group = ( + getattr(attn_tp_group, "cpu_group", attn_tp_group) + if attn_tp_group is not None + else None + ) + self.attn_cp_cpu_group = ( + getattr(attn_cp_group, "cpu_group", attn_cp_group) + if attn_cp_group is not None + else None + ) + + self.pp_size = model_config.pp_size + self.attn_tp_size = model_config.attn_tp_size + self.attn_cp_size = model_config.attn_cp_size + + self.pp_rank = rank_info.pp_rank + self.attn_tp_rank = rank_info.attn_tp_rank + self.attn_cp_rank = rank_info.attn_cp_rank + + self.is_pp_stage_leader = self.attn_tp_rank == 0 and self.attn_cp_rank == 0 + self.is_sync_leader = self.pp_rank == 0 and self.is_pp_stage_leader + self.is_pp_leader = self.pp_rank == 0 and self.is_pp_stage_leader + self.is_cp_leader = self.attn_cp_rank == 0 + self.is_tp_leader = self.attn_tp_rank == 0 + + # P2P routing tables (computed once). + stride = self.attn_tp_size * self.attn_cp_size + self._pp_stage_leader_ranks = [s * stride for s in range(self.pp_size)] + pp_stage_offset = self.pp_rank * stride + self._cp_leader_ranks = ( + [ + pp_stage_offset + cp * self.attn_tp_size + for cp in range(self.attn_cp_size) + ] + if self.attn_cp_size > 1 + else [] + ) + if self.attn_tp_size > 1: + if self.attn_tp_cpu_group is None: + raise RuntimeError( + f"[FlexKV] attn_tp_size={self.attn_tp_size} > 1 but " + f"attn_tp_cpu_group is None — TP CPU group is required " + f"for scatter/collectives." + ) + self._tp_group_ranks = [ + dist.get_global_rank(self.attn_tp_cpu_group, i) + for i in range(self.attn_tp_cpu_group.size()) + ] + else: + self._tp_group_ranks = [] + self._pp_group_global_ranks = ( + [ + dist.get_global_rank(self.pp_cpu_group, i) + for i in range(self.pp_cpu_group.size()) + ] + if self.pp_size > 1 and self.pp_cpu_group is not None + else [] + ) + self._pp_stage_member_ranks = list( + range(pp_stage_offset, pp_stage_offset + stride) + ) + + self.needs_sync = ( + self.pp_size > 1 or self.attn_tp_size > 1 or self.attn_cp_size > 1 + ) + + self._world_cpu_group = get_world_group().cpu_group + + self.pp_group = ( + self.pp_cpu_group + if (self.pp_size > 1 and self.is_pp_stage_leader) + else None + ) + self.is_pp_active = self.pp_size > 1 + self.is_pp_sender = self.is_pp_leader + self.is_pp_receiver = self.is_pp_stage_leader and not self.is_pp_leader + + self.is_cross_node_pp = self.pp_size > rank_info.pp_size_per_node + self.should_send_slot_mapping_to_remote = ( + self.is_pp_receiver and self.is_cross_node_pp + ) + + logger.info( + "[FlexKV] Comm init: rank=%d, pp=%d/%d, tp=%d/%d, cp=%d/%d, " + "sync_leader=%s, stage_leader=%s, cross_node_pp=%s", + world_rank, + self.pp_rank, + self.pp_size, + self.attn_tp_rank, + self.attn_tp_size, + self.attn_cp_rank, + self.attn_cp_size, + self.is_sync_leader, + self.is_pp_stage_leader, + self.is_cross_node_pp, + ) + + # ------------------------------------------------------------------ + # Public collectives + # ------------------------------------------------------------------ + + def scatter(self, data: Any, blocking: bool = False) -> Any: + """Hierarchical fan-out: sync_leader → PP stage leaders → + CP leaders → TP ranks. Returns the leader's payload on every rank. + + ``blocking=False`` queues isends and reaps later — fine for the + hot path; ``True`` blocks until the leader's sends drain (used + on shutdown / barriers). + """ + if self.pp_size > 1 and self.is_pp_stage_leader: + data = self._scatter_group( + data, + self._pp_stage_leader_ranks, + self.is_pp_leader, + self._TAG_PP, + blocking, + ) + if self._cp_leader_ranks: + data = self._scatter_group( + data, + self._cp_leader_ranks, + self.is_cp_leader, + self._TAG_CP, + blocking, + ) + if self._tp_group_ranks: + data = self._scatter_group( + data, + self._tp_group_ranks, + self.is_tp_leader, + self._TAG_TP, + blocking, + ) + return data + + def scatter_pp(self, data: Any) -> Any: + """PP-only fan-out across PP stages (only stage leaders participate).""" + if not self._pp_group_global_ranks: + return data + is_leader = self._pp_group_global_ranks[0] == self.world_rank + return self._scatter_group( + data, + self._pp_group_global_ranks, + is_leader, + self._TAG_SCATTER, + blocking=False, + ) + + def all_reduce_min(self, value: int) -> int: + """Hierarchical all_reduce(MIN) across TP, CP, PP. + + Used to align FlexKV block-count limits across all ranks that + will register GPU buffers (each rank computes the maximum it can + support, and we take the MIN to land on a value everyone can + honor). + """ + tensor = torch.tensor(value, dtype=torch.int64) + if self.attn_tp_size > 1 and self.attn_tp_cpu_group is not None: + dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=self.attn_tp_cpu_group) + if self.attn_cp_size > 1 and self.attn_cp_cpu_group is not None: + dist.all_reduce(tensor, op=dist.ReduceOp.MIN, group=self.attn_cp_cpu_group) + if self.pp_size > 1 and self.is_pp_stage_leader: + self._pp_all_reduce_min_p2p(tensor) + if self.pp_size > 1: + self._bcast_to_stage_members(tensor, self._TAG_AR_BCAST) + return int(tensor.item()) + + def barrier(self) -> None: + if self.attn_tp_size > 1 and self.attn_tp_cpu_group is not None: + dist.barrier(group=self.attn_tp_cpu_group) + if self.attn_cp_size > 1 and self.attn_cp_cpu_group is not None: + dist.barrier(group=self.attn_cp_cpu_group) + if self.pp_size > 1 and self.is_pp_stage_leader: + self._pp_barrier_p2p() + if self.pp_size > 1: + dummy = torch.tensor([0], dtype=torch.int64) + self._bcast_to_stage_members(dummy, self._TAG_PP_BARRIER_BCAST) + + # ------------------------------------------------------------------ + # Internal scatter helper + # ------------------------------------------------------------------ + + def _scatter_group( + self, + data: Any, + group_ranks: List[int], + is_leader: bool, + tag: int, + blocking: bool = False, + ) -> Any: + if not group_ranks or self.world_rank not in group_ranks: + return data + if is_leader: + dsts = [r for r in group_ranks if r != self.world_rank] + works = [] + for dst in dsts: + works.extend(self._isend(dst, data, tag, self._world_cpu_group)) + if blocking: + for w in works: + w.wait() + else: + self._reap_completed_async_works() + self._async_works.extend(works) + return data + return self._recv(group_ranks[0], tag, self._world_cpu_group) + + def _reap_completed_async_works(self) -> None: + n = len(self._async_works) + if n <= self._reap_high: + return + + drained = 0 + stuck = False + for _ in range(self._REAP_MAX_DRAIN): + if not self._async_works: + break + w = self._async_works[0] + try: + w.wait(self._REAP_PROBE) + except RuntimeError: + stuck = True + break + self._async_works.pop(0) + drained += 1 + + self._reap_calls += 1 + self._reap_drained_total += drained + if stuck: + self._reap_stuck_total += 1 + + prev_high = self._reap_high + if stuck: + self._reap_high = min(self._REAP_HIGH_MAX, self._reap_high * 2) + else: + self._reap_high = max(self._REAP_HIGH_BASE, self._reap_high // 2) + if self._reap_high != prev_high: + logger.debug( + "[FlexKV] reap watermark rank=%d %d->%d " + "(stuck=%s drained=%d backlog=%d)", + self.world_rank, + prev_high, + self._reap_high, + stuck, + drained, + n, + ) + if self._reap_calls % self._REAP_LOG_EVERY == 0: + logger.debug( + "[FlexKV] reap stats rank=%d calls=%d drained=%d stuck=%d " + "backlog=%d high=%d", + self.world_rank, + self._reap_calls, + self._reap_drained_total, + self._reap_stuck_total, + len(self._async_works), + self._reap_high, + ) + + # ------------------------------------------------------------------ + # Low-level send / recv on the world cpu group + # ------------------------------------------------------------------ + + def _isend(self, dst: int, data: Any, tag: int = 0, group=None) -> list: + serialized = bytearray(pickle.dumps(data)) + t_size = torch.tensor([len(serialized)], dtype=torch.long) + t_data = torch.frombuffer(serialized, dtype=torch.uint8) + return [ + dist.isend(t_size, dst=dst, tag=tag, group=group), + dist.isend(t_data, dst=dst, tag=tag, group=group), + ] + + def _recv(self, src: int, tag: int = 0, group=None) -> Any: + t_size = torch.tensor([0], dtype=torch.long) + dist.irecv(t_size, src=src, tag=tag, group=group).wait() + size = int(t_size.item()) + if size == 0: + return [] + t_data = torch.empty(size, dtype=torch.uint8) + dist.irecv(t_data, src=src, tag=tag, group=group).wait() + return pickle.loads(t_data.numpy().tobytes()) + + def _send_tensor( + self, tensor: torch.Tensor, dst: int, tag: int = 0, group=None + ) -> None: + dist.send(tensor, dst=dst, tag=tag, group=group) + + def _recv_tensor( + self, tensor: torch.Tensor, src: int, tag: int = 0, group=None + ) -> None: + dist.recv(tensor, src=src, tag=tag, group=group) + + def _bcast_to_stage_members(self, tensor: torch.Tensor, tag: int) -> None: + if not self.is_pp_stage_leader: + self._recv_tensor( + tensor, + src=self._pp_stage_leader_ranks[self.pp_rank], + tag=tag, + group=self._world_cpu_group, + ) + return + for rank in self._pp_stage_member_ranks: + if rank != self.world_rank: + self._send_tensor( + tensor, dst=rank, tag=tag, group=self._world_cpu_group + ) + + def _pp_all_reduce_min_p2p(self, tensor: torch.Tensor) -> None: + leader_rank = self._pp_stage_leader_ranks[0] + other_leaders = self._pp_stage_leader_ranks[1:] + tag = self._TAG_PP_AR_MIN + if self.world_rank == leader_rank: + result = int(tensor.item()) + for src in other_leaders: + other = torch.tensor(0, dtype=torch.int64) + self._recv_tensor(other, src=src, tag=tag, group=self._world_cpu_group) + result = min(result, int(other.item())) + tensor.fill_(result) + for dst in other_leaders: + self._send_tensor(tensor, dst=dst, tag=tag, group=self._world_cpu_group) + else: + self._send_tensor( + tensor, dst=leader_rank, tag=tag, group=self._world_cpu_group + ) + self._recv_tensor( + tensor, src=leader_rank, tag=tag, group=self._world_cpu_group + ) + + def _pp_barrier_p2p(self) -> None: + leader_rank = self._pp_stage_leader_ranks[0] + other_leaders = self._pp_stage_leader_ranks[1:] + tag = self._TAG_PP_BARRIER + dummy = torch.tensor([1], dtype=torch.int64) + if self.world_rank == leader_rank: + for src in other_leaders: + self._recv_tensor(dummy, src=src, tag=tag, group=self._world_cpu_group) + for dst in other_leaders: + self._send_tensor(dummy, dst=dst, tag=tag, group=self._world_cpu_group) + else: + self._send_tensor( + dummy, dst=leader_rank, tag=tag, group=self._world_cpu_group + ) + self._recv_tensor( + dummy, src=leader_rank, tag=tag, group=self._world_cpu_group + ) + + +# ---------------------------------------------------------------------- +# libc / eventfd / SCM_RIGHTS shims for the layerwise UDS handshake +# ---------------------------------------------------------------------- + +_libc = ctypes.CDLL("libc.so.6", use_errno=True) +_libc.eventfd.argtypes = [ctypes.c_uint, ctypes.c_int] +_libc.eventfd.restype = ctypes.c_int +_libc.read.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_size_t] +_libc.read.restype = ctypes.c_ssize_t +_libc.write.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_size_t] +_libc.write.restype = ctypes.c_ssize_t + +EFD_SEMAPHORE = 0x1 +EFD_NONBLOCK = 0x800 + + +def eventfd(initval: int = 0, flags: int = 0) -> int: + fd = _libc.eventfd(ctypes.c_uint(initval), ctypes.c_int(flags)) + if fd == -1: + err = ctypes.get_errno() + raise OSError(err, os.strerror(err)) + return fd + + +def eventfd_write(fd: int, val: int) -> None: + v = ctypes.c_uint64(val) + n = _libc.write(fd, ctypes.byref(v), ctypes.sizeof(v)) + if n != ctypes.sizeof(v): + err = ctypes.get_errno() + raise OSError(err, f"eventfd write failed: {os.strerror(err)}") + + +def eventfd_read(fd: int) -> int: + v = ctypes.c_uint64() + n = _libc.read(fd, ctypes.byref(v), ctypes.sizeof(v)) + if n != ctypes.sizeof(v): + err = ctypes.get_errno() + if err == errno.EAGAIN: + return 0 + raise OSError(err, f"eventfd read failed: {os.strerror(err)}") + return v.value + + +def send_fds(sock: socket.socket, fds: list, extra_data: bytes = b"x") -> None: + """SCM_RIGHTS-send a list of file descriptors over a UDS socket.""" + fds_packed = struct.pack(f"{len(fds)}i", *fds) + ancdata = [(socket.SOL_SOCKET, socket.SCM_RIGHTS, fds_packed)] + sock.sendmsg([extra_data], ancdata) + + +# ---------------------------------------------------------------------- +# Layerwise transfer signaling (eventfd-backed) +# ---------------------------------------------------------------------- + + +class FlexKVLayerLoadingEvent: + """One per producer slot. Holds ``num_layers`` semaphore eventfds — + the FlexKV layerwise worker writes 1 to each as the corresponding + layer's H2D copy completes; the consumer (sglang's + ``register_layer_transfer_counter`` hook) reads to wait for them.""" + + def __init__(self, num_layers: int): + self._num_layers = num_layers + # Semaphore mode so each read consumes exactly one signal. NONBLOCK + # lets ``reset_for_new_transfer`` drain leftover counter values + # without blocking; ``wait`` re-arms the fd to blocking before + # reading so consumers still get the desired blocking semantics. + self.load_event_fds: List[int] = [ + eventfd(0, EFD_SEMAPHORE | EFD_NONBLOCK) for _ in range(num_layers) + ] + self._finished = True + self.wait_remaining: List[int] = [1] * num_layers + + def reset_for_new_transfer(self) -> None: + """Drain any leftover signals from prior transfers, then arm. + + Without this drain, a previous transfer that wrote N eventfd + signals but only had N-K reads (e.g. because the attention + backend skipped a layer's ``get_key_buffer`` call) leaves K + pending. The next transfer's first ``wait(layer)`` returns + immediately reading one of those stale signals, even though + the FlexKV worker hasn't actually finished that layer's H2D + yet — and forward proceeds with wrong KV data. + """ + import os + + for fd in self.load_event_fds: + # The fd is NONBLOCK: read until EAGAIN. Each read is 8 bytes. + while True: + try: + if not os.read(fd, 8): + break + except BlockingIOError: + break + except OSError: + break + self._finished = False + self.wait_remaining = [1] * self._num_layers + + def wait(self, layer_index: int) -> None: + """Block until the FlexKV worker signals layer ``layer_index``. + + The fd was created with EFD_NONBLOCK so reset can drain it. We + re-introduce the blocking semantics with ``select.select`` on a + NONBLOCK fd: the read after select is guaranteed to consume one + signal. + """ + import os + import select + + assert 0 <= layer_index < self._num_layers + fd = self.load_event_fds[layer_index] + while True: + select.select([fd], [], []) + try: + buf = os.read(fd, 8) + if buf: + break + except BlockingIOError: + # Spurious wakeup; loop and re-select. + continue + if layer_index == self._num_layers - 1: + self._finished = True + + def close(self) -> None: + for fd in self.load_event_fds: + try: + os.close(fd) + except Exception: + pass + self.load_event_fds.clear() + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass + + +class FlexKVLayerDoneCounter: + """Triple-buffered slot-based layerwise counter. + + The KV pool calls ``wait_until(layer_id)`` once per layer during + forward. We track which producer slot the current task is using and + block on that slot's ``layer_id``-th eventfd. Producer rotation lets + the next prefetch start before the current one finishes consuming. + """ + + def __init__(self, num_layers: int, num_counters: int = 3): + self.num_layers = num_layers + self.num_counters = num_counters + self.events: List[FlexKVLayerLoadingEvent] = [ + FlexKVLayerLoadingEvent(num_layers) for _ in range(num_counters) + ] + self.producer_index = -1 + self.consumer_index = -1 + self._task_to_producer: Dict[int, int] = {} + + def register_task(self, task_id: int, producer_id: int) -> None: + self._task_to_producer[task_id] = producer_id + + def register_task_with_explicit_counter_id( + self, task_id: int, counter_id: int + ) -> None: + if not 0 <= counter_id < self.num_counters: + raise ValueError( + f"Invalid counter_id={counter_id}, must be in [0, {self.num_counters})" + ) + self._task_to_producer[task_id] = counter_id + self.events[counter_id].reset_for_new_transfer() + + def update_producer(self) -> int: + self.producer_index = (self.producer_index + 1) % self.num_counters + assert self.events[ + self.producer_index + ]._finished, "Producer event should be finished before reuse" + return self.producer_index + + def set_consumer(self, task_id: int) -> None: + if task_id < 0: + self.consumer_index = -1 + return + producer_id = self._task_to_producer.pop(task_id, None) + self.consumer_index = producer_id if producer_id is not None else -1 + + def wait_until(self, threshold: int) -> None: + if self.consumer_index < 0: + return + event = self.events[self.consumer_index] + if event.wait_remaining[threshold] <= 0: + return + event.wait_remaining[threshold] -= 1 + event.wait(threshold) + + def reset(self) -> None: + self.producer_index = -1 + self.consumer_index = -1 + self._task_to_producer.clear() + + def __del__(self) -> None: + try: + for event in self.events: + event.close() + self.events.clear() + except Exception: + pass diff --git a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_connector.py b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_connector.py new file mode 100644 index 000000000..a1fb14462 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_connector.py @@ -0,0 +1,925 @@ +"""Wrapper around FlexKV ``KVManager`` for sglang. + +The public surface is small (see "Public API" below). The class owns: + +* the FlexKV ``KVManager`` (server-client mode when ``dp_size > 1`` or + multi-instance; in-process otherwise — handled by FlexKV itself); +* the per-rank ``KVTPClient`` that registers this rank's GPU KV cache + with the FlexKV TransferManager; +* an optional ``FlexKVLayerDoneCounter`` plus the UDS-side handshake + that wires its eventfds into the FlexKV layerwise transfer worker. + +Cross-rank sync uses :class:`FlexKVComm`. Only the **sync leader** +(rank 0 of every PP × CP × TP axis) talks to ``KVManager``; other +ranks block on broadcast / barrier. + +Modes: + * **MP / synchronous** (default): ``retrieve_kv`` fires ``launch`` + and blocks on ``wait`` so the device slots are ready by the time + sglang's prefill runs. + * **Layerwise** (``FLEXKV_ENABLE_LAYERWISE_TRANSFER=1``): ``launch`` + is fired with ``layerwise_transfer=True`` and the per-layer hook + registered via ``register_layer_transfer_counter`` blocks each + forward layer on its own eventfd. +""" + +from __future__ import annotations + +import logging +import os +import socket +import struct +import time +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +import torch + +from sglang.srt.mem_cache.storage.flexkv.flexkv_comm import ( + CMD_LAYERWISE, + CMD_PUT_META, + CMD_STORE_COMPLETE, + FlexKVComm, + FlexKVLayerDoneCounter, + send_fds, +) + +try: + from flexkv.common.request import KVResponseStatus + from flexkv.common.storage import KVCacheLayout, KVCacheLayoutType + from flexkv.integration.config import FlexKVConfig + from flexkv.kvmanager import KVManager + from flexkv.server.client import KVTPClient + from flexkv.transfer.layerwise import build_layerwise_eventfd_socket_path + from flexkv.transfer_manager import TransferManagerOnRemote +except ImportError as exc: # pragma: no cover - runtime check + raise RuntimeError( + "FlexKV is not installed. Please install the FlexKV package to use " + "--enable-flexkv." + ) from exc + +logger = logging.getLogger(__name__) + + +class FlexKVConnector: + """A FlexKV-side façade used by :class:`FlexKVRadixCache`. + + This class manages connection lifecycle and provides a small, + sgl-friendly contract over FlexKV's task-based API: + + * ``lookup_kv`` — page-aligned hit count + a held task id. + * ``retrieve_kv`` — synchronous load (launch + wait). + * ``start_load_kv_layerwise`` — layerwise async load. + * ``store_kv`` — page-aligned write back. + * ``check_completed_stores`` — drain async store completions. + * ``prefetch_async`` / ``check_prefetch_progress`` / + ``cancel_prefetch`` — opportunistic CPU↔SSD/Remote staging. + * ``release_pending`` — cancel a held task whose load won't run. + * ``reset`` / ``shutdown``. + """ + + def __init__( + self, + *, + sgl_model_config: Any, + server_args: Any, + page_size: int, + kvcache: Any, + tp_rank: int, + dp_rank: Optional[int], + pp_rank: int, + attn_cp_rank: int, + pp_group: Any = None, + attn_tp_group: Any = None, + attn_cp_group: Any = None, + ) -> None: + self.page_size = int(page_size) + + # 1. Resolve FlexKV config from env + sglang server args. + self.flexkv_config = FlexKVConfig.from_env() + self.rank_info = self.flexkv_config.post_init_from_sglang_config( + sglang_config=sgl_model_config, + server_args=server_args, + page_size=self.page_size, + tp_rank=tp_rank, + pp_rank=pp_rank, + dp_rank=dp_rank if dp_rank is not None else 0, + attn_cp_rank=attn_cp_rank, + ) + self.model_config = self.flexkv_config.model_config + self.cache_config = self.flexkv_config.cache_config + self._label = f"[model_config={self.model_config}, rank_info={self.rank_info}]" + + # 2. Cross-rank sync context. + world_rank = ( + torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + ) + self._sync_ctx = FlexKVComm( + rank_info=self.rank_info, + world_rank=world_rank, + pp_group=pp_group, + attn_tp_group=attn_tp_group, + attn_cp_group=attn_cp_group, + ) + + # 3. Align block counts across all ranks (MIN reduce) so each + # rank's KVManager registers compatible sizes. + for attr in ("num_cpu_blocks", "num_ssd_blocks", "num_remote_blocks"): + orig = getattr(self.cache_config, attr, None) + if orig is None or orig <= 0: + continue + aligned = self._sync_ctx.all_reduce_min(int(orig)) + if aligned != orig: + logger.info( + "[FlexKV] Block count MIN alignment '%s': %d -> %d", + attr, + orig, + aligned, + ) + setattr(self.cache_config, attr, aligned) + + # 4. Extract MLA/MHA KV buffers + optional indexer buffers. + indexer_buffers = getattr(kvcache, "index_k_with_scale_buffer", None) + if hasattr(kvcache, "kv_buffer"): + # MLA: K and V share the same buffer (per-layer tensor). + kv_caches = list(kvcache.kv_buffer) + elif hasattr(kvcache, "k_buffer"): + # MHA: K buffers concatenated with V buffers, layer-first. + kv_caches = list(kvcache.k_buffer) + list(kvcache.v_buffer) + else: + raise AttributeError( + f"Unsupported KV cache type {type(kvcache).__name__}: " + f"expected kv_buffer (MLA/NSA) or k_buffer/v_buffer (MHA)." + ) + self._kvcache = kvcache + + # 5. On multi-node setups, every node beyond node 0 needs a + # TransferManagerOnRemote process (FlexKV side) before any rank + # on that node can register GPU buffers. + self._remote_process = None + if ( + self.model_config.nnodes > 1 + and self.rank_info.node_rank > 0 + and self.rank_info.local_rank == 0 + ): + self._remote_process = TransferManagerOnRemote.create_process( + master_host=self.model_config.master_host, + master_ports=self.model_config.master_ports, + ) + logger.info( + "[FlexKV] Launched TransferManagerOnRemote on node_rank=%d %s", + self.rank_info.node_rank, + self._label, + ) + + # 6. Bring up KVManager on the sync leader only. + self.kv_manager: Optional[KVManager] = None + if self._sync_ctx.is_sync_leader: + self.kv_manager = KVManager( + model_config=self.model_config, + cache_config=self.cache_config, + dp_client_id=self.rank_info.dp_client_id, + server_recv_port=self.flexkv_config.server_recv_port, + gpu_register_port=self.flexkv_config.gpu_register_port, + ) + self.kv_manager.start() + + # 7. Per-rank TP client registers this rank's GPU buffers. + self.tp_client = KVTPClient( + self.flexkv_config.gpu_register_port, + dp_client_id=self.rank_info.dp_client_id, + pp_rank=self.rank_info.pp_rank, + device_id=self.rank_info.local_rank, + ) + self._register_with_retry(kv_caches, indexer_buffers) + + # 8. Layerwise transfer plumbing. + self.enable_layerwise = bool( + int(os.environ.get("FLEXKV_ENABLE_LAYERWISE_TRANSFER", "0")) + ) + self._layerwise_socket = build_layerwise_eventfd_socket_path( + dp_client_id=self.rank_info.dp_client_id, + pp_rank=self.rank_info.pp_rank, + model_config=self.model_config, + ) + self._layerwise_eventfd_connect_max_retries = max( + 360, + int(os.environ.get("FLEXKV_LAYERWISE_EVENTFD_CONNECT_MAX_RETRIES", "0")), + ) + self.layer_done_counter: Optional[FlexKVLayerDoneCounter] = None + if self.enable_layerwise: + self.layer_done_counter = FlexKVLayerDoneCounter( + self.rank_info.num_layers_per_pp_stage + ) + self._send_eventfds_to_worker() + + # 9. Wait for the KVManager (and its remote subprocess) to be ready. + if self._sync_ctx.is_sync_leader: + self._wait_kv_manager_ready() + + # 10. Per-rank in-flight tracking. + # Loads + self._pending_lookups: Dict[str, int] = {} # rid -> fkv_task_id + self._inflight_loads: Dict[int, int] = {} # producer_id -> rid hashlike + self._completed_layerwise: List[int] = [] + self._launched_load_tids: List[int] = [] # leader-only, for periodic drain + # Stores + self._inflight_stores: Dict[str, int] = {} # rid -> fkv_task_id + # Prefetches + self._ongoing_prefetches: Dict[str, int] = {} # rid -> fkv_task_id + self._prefetch_enabled = bool( + self.cache_config.enable_ssd + or self.cache_config.enable_remote + or self.cache_config.enable_kv_sharing + ) + + logger.info( + "[FlexKV] Connector ready %s: layerwise=%s, prefetch=%s", + self._label, + self.enable_layerwise, + self._prefetch_enabled, + ) + + # ------------------------------------------------------------------ + # Public API — lookup / load + # ------------------------------------------------------------------ + + def lookup_kv( + self, + token_ids: List[int], + token_mask: torch.Tensor, + rid: Optional[str] = None, + ) -> Tuple[int, int]: + """Page-aligned prefix lookup against FlexKV. + + Args: + token_ids: full token id sequence we'd like to check. + token_mask: 1-D bool tensor or array, True for "this token is + *not* already on GPU and is a candidate for load-back". + rid: if set and hit > 0, the held FlexKV task id is stashed + under this key so a later ``retrieve_kv(rid, slots)`` call + can resolve it. If not set, the held task is cancelled when + hit > 0 and the caller didn't ask to track it. + + Returns: + ``(fkv_task_id, hit_count)``. ``hit_count`` is page-aligned + and may be smaller than the raw FlexKV match if the page + floor truncated it. + """ + fkv_task_id = -1 + hit_length = 0 + + if self._sync_ctx.is_sync_leader and self.kv_manager is not None: + tids_np = np.asarray(token_ids, dtype=np.int64) + mask_np = self._as_numpy_mask(token_mask) + try: + res = self.kv_manager.get_match(token_ids=tids_np, token_mask=mask_np) + except Exception as exc: # noqa: BLE001 + logger.warning("[FlexKV] get_match raised: %s", exc) + res = None + if res is None: + fkv_task_id = -1 + hit_length = 0 + else: + fkv_task_id, matched_mask = res + hit_length = int(matched_mask.sum()) if matched_mask is not None else 0 + + if self._sync_ctx.needs_sync: + payload = self._sync_ctx.scatter( + {"task_id": fkv_task_id, "hit": hit_length} + ) + fkv_task_id = payload["task_id"] + hit_length = payload["hit"] + + # Page-align: FlexKV transfers whole pages. + if hit_length > 0 and self.page_size > 1: + aligned = (hit_length // self.page_size) * self.page_size + if aligned < hit_length: + logger.debug( + "[FlexKV] lookup_kv: page-aligning hit %d -> %d (page=%d)", + hit_length, + aligned, + self.page_size, + ) + hit_length = aligned + + # Decide what to do with the held task. Three cases: + # 1. hit_length > 0 and rid given → stash for retrieve_kv later. + # 2. hit_length > 0 and rid is None → cancel; caller can't use it. + # 3. hit_length == 0 → no work to do; FlexKV already marked the + # empty graph COMPLETED inside get_match, cancel would warn. + if hit_length > 0 and rid is not None and fkv_task_id >= 0: + self._pending_lookups[rid] = fkv_task_id + elif hit_length > 0 and fkv_task_id >= 0 and self._sync_ctx.is_sync_leader: + assert self.kv_manager is not None + self.kv_manager.cancel([fkv_task_id]) + + return fkv_task_id, hit_length + + def release_pending(self, rid: str) -> None: + """Cancel the task held by an earlier ``lookup_kv(rid=...)`` that + won't be followed by a ``retrieve_kv`` (e.g. allocation failed).""" + fkv_task_id = self._pending_lookups.pop(rid, -1) + if fkv_task_id >= 0 and self._sync_ctx.is_sync_leader: + assert self.kv_manager is not None + self.kv_manager.cancel([fkv_task_id]) + + def retrieve_kv( + self, + rid: str, + slot_mapping: torch.Tensor, + ) -> int: + """Synchronous load: ``launch`` + ``wait``. + + Returns the number of slots actually loaded. The caller is + responsible for having allocated ``slot_mapping`` of length + equal to ``hit_length`` from a prior ``lookup_kv``. + """ + fkv_task_id = self._pending_lookups.pop(rid, -1) + if fkv_task_id < 0: + return 0 + + slot_mapping_cpu = self._to_cpu_int64(slot_mapping) + + # Cross-node PP receivers must send their slot mapping back to + # the TransferManagerOnRemote so the remote side knows where to + # land the H2D copies on its own GPUs. + if self._sync_ctx.should_send_slot_mapping_to_remote: + self._send_slot_mapping_to_remote(fkv_task_id, slot_mapping_cpu) + + n = slot_mapping_cpu.numel() + if self._sync_ctx.is_sync_leader and self.kv_manager is not None: + self.kv_manager.launch( + task_ids=[fkv_task_id], + slot_mappings=[slot_mapping_cpu], + as_batch=True, + layerwise_transfer=False, + ) + resp = self.kv_manager.wait([fkv_task_id], timeout=30.0) + if not ( + fkv_task_id in resp + and resp[fkv_task_id].status == KVResponseStatus.SUCCESS + ): + logger.warning( + "[FlexKV] retrieve_kv: task %d failed/timed out", + fkv_task_id, + ) + n = 0 + if self._sync_ctx.needs_sync: + self._sync_ctx.barrier() + return n + + def start_load_kv_layerwise( + self, + rid: str, + slot_mapping: torch.Tensor, + ) -> Tuple[int, int]: + """Layerwise load. Fires ``launch(layerwise_transfer=True)`` and + returns ``(n_slots, producer_id)``. The caller registers + ``producer_id`` with the layer hook so the KV pool blocks on + the right eventfds during forward.""" + assert self.enable_layerwise and self.layer_done_counter is not None, ( + "start_load_kv_layerwise called but layerwise transfer is " + "disabled. Set FLEXKV_ENABLE_LAYERWISE_TRANSFER=1." + ) + fkv_task_id = self._pending_lookups.pop(rid, -1) + if fkv_task_id < 0: + return 0, -1 + + slot_mapping_cpu = self._to_cpu_int64(slot_mapping) + n = slot_mapping_cpu.numel() + + if self._sync_ctx.should_send_slot_mapping_to_remote: + self._send_slot_mapping_to_remote(fkv_task_id, slot_mapping_cpu) + + # Allocate / receive producer slot. + if self._sync_ctx.is_pp_receiver: + payload = self._sync_ctx.scatter_pp(None) + if payload.get("cmd") != CMD_LAYERWISE: + raise RuntimeError( + f"Tag mismatch: expected CMD_LAYERWISE, got " + f"{payload.get('cmd')}" + ) + producer_id = int(payload["counter_id"]) + self.layer_done_counter.register_task_with_explicit_counter_id( + fkv_task_id, producer_id + ) + else: + producer_id = self.layer_done_counter.update_producer() + self.layer_done_counter.events[producer_id].reset_for_new_transfer() + self.layer_done_counter.register_task(fkv_task_id, producer_id) + + if self._sync_ctx.is_pp_sender: + self._sync_ctx.scatter_pp( + { + "cmd": CMD_LAYERWISE, + "fkv_task_id": fkv_task_id, + "counter_id": producer_id, + } + ) + + if self._sync_ctx.is_sync_leader and self.kv_manager is not None: + self.kv_manager.launch( + task_ids=[fkv_task_id], + slot_mappings=[slot_mapping_cpu], + as_batch=True, + layerwise_transfer=True, + counter_id=producer_id, + ) + self._launched_load_tids.append(fkv_task_id) + + # Tell the layer hook which counter slot to wait on. + self.layer_done_counter.set_consumer(fkv_task_id) + return n, producer_id + + def drain_launched_loads(self, threshold: int = 100) -> None: + """Periodic non-blocking sweep on long-lived launched tasks so the + FlexKV pipe doesn't accumulate. No-op on non-leader ranks.""" + if not self._sync_ctx.is_sync_leader or self.kv_manager is None: + return + if len(self._launched_load_tids) < threshold: + return + try: + self.kv_manager.try_wait(task_ids=list(self._launched_load_tids)) + except Exception as exc: # noqa: BLE001 + logger.debug("[FlexKV] drain_launched_loads try_wait: %s", exc) + self._launched_load_tids.clear() + + # ------------------------------------------------------------------ + # Public API — store + # ------------------------------------------------------------------ + + def store_kv( + self, + rid: str, + token_ids: List[int], + kv_indices: torch.Tensor, + ) -> int: + """Schedule a write back from GPU into FlexKV. + + On the sync leader this runs ``put_match`` to discover which + tokens are NOT yet in FlexKV's CPU cache (= the "unmatched" + slice), then ``launch`` on those. On non-leaders the unmatched + mask is received over the PP fan-out so cross-node PP can + forward its slot mappings. + + Returns the FlexKV task id of the in-flight store, or -1 if + nothing needed to be written. + """ + token_ids_np = np.asarray(token_ids, dtype=np.int64) + n = len(token_ids_np) + if n != len(kv_indices): + raise ValueError( + f"store_kv: token_ids has {n} entries but kv_indices " + f"has {len(kv_indices)} entries" + ) + + # Page-align inputs *before* put_match so the FlexKV allocator + # only reserves slots that line up with the slot_mapping we send. + if self.page_size > 1: + aligned_len = (n // self.page_size) * self.page_size + if aligned_len == 0: + self._send_pp_put_meta(-1, []) + return -1 + if aligned_len < n: + token_ids_np = token_ids_np[:aligned_len] + kv_indices = kv_indices[:aligned_len] + + fkv_task_id = -1 + if self._sync_ctx.is_sync_leader and self.kv_manager is not None: + try: + res = self.kv_manager.put_match(token_ids=token_ids_np, token_mask=None) + except Exception as exc: # noqa: BLE001 + logger.warning("[FlexKV] put_match raised: %s", exc) + res = None + if res is None: + self._send_pp_put_meta(-1, []) + return -1 + fkv_task_id, unmatched_mask = res + + self._send_pp_put_meta(fkv_task_id, unmatched_mask) + + if int(unmatched_mask.sum()) > 0: + filtered = kv_indices[unmatched_mask] + slot_mapping_cpu = self._to_cpu_int64(filtered) + self.kv_manager.launch( + task_ids=[fkv_task_id], + slot_mappings=[slot_mapping_cpu], + as_batch=False, + layerwise_transfer=False, + ) + self._inflight_stores[rid] = fkv_task_id + return fkv_task_id + return -1 + + # Non-leader path: receive the unmatched mask + maybe forward + # slot_mapping to the remote-side TransferManager. + if self._sync_ctx.is_pp_receiver: + payload = self._sync_ctx.scatter_pp(None) + if payload.get("cmd") != CMD_PUT_META: + raise RuntimeError( + f"Tag mismatch: expected CMD_PUT_META, got " f"{payload.get('cmd')}" + ) + fkv_task_id = int(payload["fkv_task_id"]) + mask_list = payload.get("unmatched_mask", []) + unmatched_mask = torch.tensor(mask_list, dtype=torch.bool) + if ( + int(unmatched_mask.sum()) > 0 + and fkv_task_id >= 0 + and self._sync_ctx.should_send_slot_mapping_to_remote + ): + filtered = kv_indices[unmatched_mask] + slot_mapping_cpu = self._to_cpu_int64(filtered) + self._send_slot_mapping_to_remote(fkv_task_id, slot_mapping_cpu) + self._inflight_stores[rid] = fkv_task_id + return fkv_task_id + + def check_completed_stores(self) -> List[str]: + """Return rids whose stores have completed since the last call.""" + completed_rids: List[str] = [] + completed_dict: Dict[int, Any] = {} + + if self._sync_ctx.is_sync_leader and self.kv_manager is not None: + if self._inflight_stores: + fk_to_rid = {v: k for k, v in self._inflight_stores.items()} + try: + completed_dict = self.kv_manager.try_wait( + task_ids=list(fk_to_rid.keys()) + ) + except Exception as exc: # noqa: BLE001 + logger.debug("[FlexKV] check_completed_stores: %s", exc) + completed_dict = {} + for fk_tid in completed_dict: + rid = fk_to_rid[fk_tid] + completed_rids.append(rid) + self._inflight_stores.pop(rid, None) + + if self._sync_ctx.is_pp_sender: + self._sync_ctx.scatter_pp( + { + "cmd": CMD_STORE_COMPLETE, + "completed_fk_ids": list(completed_dict), + } + ) + elif self._sync_ctx.is_pp_receiver: + payload = self._sync_ctx.scatter_pp(None) + if payload.get("cmd") != CMD_STORE_COMPLETE: + raise RuntimeError( + f"Tag mismatch: expected CMD_STORE_COMPLETE, got " + f"{payload.get('cmd')}" + ) + fk_ids = payload.get("completed_fk_ids", []) + if fk_ids and self._inflight_stores: + fk_to_rid = {v: k for k, v in self._inflight_stores.items()} + for fk_tid in fk_ids: + if fk_tid in fk_to_rid: + rid = fk_to_rid[fk_tid] + completed_rids.append(rid) + self._inflight_stores.pop(rid, None) + + if self._sync_ctx.needs_sync: + completed_rids = self._sync_ctx.scatter(completed_rids) + return completed_rids + + def wait_store(self, rid: str, timeout: float = 30.0) -> bool: + """Block until a single store task identified by ``rid`` finishes.""" + fkv_task_id = self._inflight_stores.pop(rid, -1) + if fkv_task_id < 0: + return True + if not self._sync_ctx.is_sync_leader or self.kv_manager is None: + return True + try: + resp = self.kv_manager.wait([fkv_task_id], timeout=timeout) + except Exception as exc: # noqa: BLE001 + logger.warning("[FlexKV] wait_store: %s", exc) + return False + return ( + fkv_task_id in resp and resp[fkv_task_id].status == KVResponseStatus.SUCCESS + ) + + # ------------------------------------------------------------------ + # Public API — prefetch + # ------------------------------------------------------------------ + + def prefetch_async(self, rid: str, token_ids: List[int]) -> int: + if not self._prefetch_enabled or not rid: + return -1 + task_id = -1 + if self._sync_ctx.is_sync_leader and self.kv_manager is not None: + try: + task_id = self.kv_manager.prefetch_async( + token_ids=np.asarray(token_ids, dtype=np.int64) + ) + except Exception as exc: # noqa: BLE001 + logger.debug("[FlexKV] prefetch_async: %s", exc) + task_id = -1 + if self._sync_ctx.needs_sync: + payload = self._sync_ctx.scatter({"task_id": task_id}) + task_id = payload["task_id"] + if task_id >= 0: + self._ongoing_prefetches[rid] = task_id + return task_id + + def check_prefetch_progress(self, rid: str) -> bool: + if not self._prefetch_enabled: + return True + task_id = self._ongoing_prefetches.get(rid, -1) + if task_id < 0: + return True + done = False + if self._sync_ctx.is_sync_leader and self.kv_manager is not None: + try: + completed = self.kv_manager.try_wait(task_ids=[task_id]) + except Exception: # noqa: BLE001 + completed = {} + if task_id in completed: + done = True + if self._sync_ctx.needs_sync: + payload = self._sync_ctx.scatter({"done": done}) + done = payload["done"] + if done: + self._ongoing_prefetches.pop(rid, None) + return done + + def cancel_prefetch(self, rid: str) -> None: + self._pending_lookups.pop(rid, None) + # FlexKV doesn't currently support prefetch cancellation, but + # we still drop our tracking entry. + self._ongoing_prefetches.pop(rid, None) + + # ------------------------------------------------------------------ + # Layerwise transfer hooks + # ------------------------------------------------------------------ + + def register_layer_transfer_counter(self, kvcache: Any) -> None: + """Register the FlexKVLayerDoneCounter onto sglang's KV pool so + each forward layer blocks on its eventfd. No-op when layerwise + is disabled.""" + if ( + self.layer_done_counter is None + or kvcache is None + or not hasattr(kvcache, "register_layer_transfer_counter") + ): + return + kvcache.register_layer_transfer_counter(self.layer_done_counter) + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + def reset(self) -> None: + # Drop pending lookups (cancel their held tasks on the leader). + if self._sync_ctx.is_sync_leader and self.kv_manager is not None: + pending = [tid for tid in self._pending_lookups.values() if tid >= 0] + if pending: + try: + self.kv_manager.cancel(pending) + except Exception as exc: # noqa: BLE001 + logger.debug("[FlexKV] reset cancel: %s", exc) + self._pending_lookups.clear() + self._ongoing_prefetches.clear() + self._inflight_loads.clear() + self._completed_layerwise.clear() + self._launched_load_tids.clear() + if self._sync_ctx.is_sync_leader and self.kv_manager is not None: + for fk_tid in list(self._inflight_stores.values()): + if fk_tid >= 0: + try: + self.kv_manager.wait([fk_tid], timeout=20.0) + except Exception: # noqa: BLE001 + pass + self._inflight_stores.clear() + if self.layer_done_counter is not None: + self.layer_done_counter.reset() + + def shutdown(self) -> None: + if self._sync_ctx.is_sync_leader and self.kv_manager is not None: + try: + self.kv_manager.shutdown() + except Exception as exc: # noqa: BLE001 + logger.warning("[FlexKV] kv_manager.shutdown: %s", exc) + if self._remote_process is not None: + try: + self._remote_process.terminate() + self._remote_process.join(timeout=5.0) + if self._remote_process.is_alive(): + self._remote_process.kill() + self._remote_process.join() + except Exception as exc: # noqa: BLE001 + logger.warning("[FlexKV] remote process shutdown: %s", exc) + self._remote_process = None + + # ------------------------------------------------------------------ + # Private helpers + # ------------------------------------------------------------------ + + @staticmethod + def _as_numpy_mask(mask) -> np.ndarray: + if mask is None: + return None + if isinstance(mask, torch.Tensor): + return mask.detach().cpu().numpy() + return np.asarray(mask) + + @staticmethod + def _to_cpu_int64(tensor: torch.Tensor) -> torch.Tensor: + if tensor.is_cuda: + tensor = tensor.cpu() + return tensor.to(torch.int64) + + def _wait_kv_manager_ready(self, poll_interval: float = 10.0) -> None: + assert self.kv_manager is not None + wait_count = 0 + while not self.kv_manager.is_ready(): + time.sleep(poll_interval) + wait_count += 1 + logger.info( + "[FlexKV] Waiting for FlexKV ready %s (waited %.0fs)", + self._label, + wait_count * poll_interval, + ) + logger.info("[FlexKV] FlexKV is ready %s", self._label) + + def _register_with_retry( + self, + kv_caches: List[torch.Tensor], + indexer_buffers: Optional[List[torch.Tensor]] = None, + max_retries: int = 360, + ) -> None: + """Retry GPU registration. On node_rank>0, the + TransferManagerOnRemote may not be ready immediately; retry up + to ~6 minutes.""" + for attempt in range(max_retries): + try: + self._register_to_server(kv_caches, indexer_buffers) + return + except Exception as exc: # noqa: BLE001 + if attempt == max_retries - 1: + raise + if attempt % 30 == 0: + logger.info( + "[FlexKV] GPU register retry %s attempt=%d/%d " "error=%s", + self._label, + attempt + 1, + max_retries, + exc, + ) + time.sleep(1.0) + + def _register_to_server( + self, + kv_caches: List[torch.Tensor], + indexer_buffers: Optional[List[torch.Tensor]] = None, + ) -> None: + assert len(kv_caches) > 0 + assert ( + kv_caches[0].ndim == 3 + ), f"Expected 3D KV cache tensor, got shape={kv_caches[0].shape}" + + is_mla = self.model_config.use_mla + num_blocks, num_kv_heads, head_size = kv_caches[0].shape + + gpu_layout = KVCacheLayout( + type=KVCacheLayoutType.LAYERFIRST, + num_layer=self.rank_info.num_layers_per_pp_stage, + num_block=num_blocks // self.page_size, + tokens_per_block=self.page_size, + num_head=num_kv_heads, + head_size=head_size, + is_mla=is_mla, + ) + + indexer_layout = None + if indexer_buffers is not None and len(indexer_buffers) > 0: + indexer_tensor = indexer_buffers[0] + assert indexer_tensor.ndim == 2, ( + f"Expected 2D indexer tensor (num_pages, page_stride_size), " + f"got shape={indexer_tensor.shape}" + ) + indexer_layout = KVCacheLayout( + type=KVCacheLayoutType.LAYERFIRST, + num_layer=len(indexer_buffers), + num_block=indexer_tensor.shape[0], + tokens_per_block=1, + num_head=1, + head_size=indexer_tensor.shape[1], + is_mla=True, + ) + + self.tp_client.register_to_server( + kv_caches=kv_caches, + kv_layout=gpu_layout, + indexer_buffers=indexer_buffers, + indexer_layout=indexer_layout, + ) + logger.info("[FlexKV] Registered KV caches to server %s", self._label) + + def _send_pp_put_meta(self, fkv_task_id: int, unmatched_mask) -> None: + if not self._sync_ctx.is_pp_active: + return + if hasattr(unmatched_mask, "tolist"): + mask_list = unmatched_mask.tolist() + else: + mask_list = list(unmatched_mask) + self._sync_ctx.scatter_pp( + { + "cmd": CMD_PUT_META, + "fkv_task_id": fkv_task_id, + "unmatched_mask": mask_list, + } + ) + + def _send_slot_mapping_to_remote( + self, task_id: int, slot_mapping_cpu: torch.Tensor + ) -> None: + np_arr = slot_mapping_cpu.numpy() + self.tp_client.set_slot_mapping(task_id, np_arr) + + def _send_eventfds_to_worker(self, retry_interval: float = 1.0) -> None: + """UDS handshake with the FlexKV layerwise transfer worker. + + Sends per-counter eventfd FDs over a unix domain socket using + ``SCM_RIGHTS``. Retries connect (worker may not yet be up) and + retries the whole connect+send sequence on send error. + """ + max_retries = self._layerwise_eventfd_connect_max_retries + max_send_retries = 3 + last_error: Optional[BaseException] = None + + assert self.layer_done_counter is not None + + for send_attempt in range(max_send_retries): + sock: Optional[socket.socket] = None + try: + # Phase 1: connect (worker may not yet be up). + for attempt in range(max_retries): + sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + sock.connect(self._layerwise_socket) + logger.info( + "[FlexKV] Eventfd connected %s socket=%s attempts=%d", + self._label, + self._layerwise_socket, + attempt + 1, + ) + break + except (FileNotFoundError, ConnectionRefusedError) as exc: + sock.close() + sock = None + if attempt == max_retries - 1: + raise RuntimeError( + f"[FlexKV] Failed to connect to eventfd socket " + f"{self._layerwise_socket} after {max_retries} attempts" + ) from exc + time.sleep(retry_interval) + assert sock is not None + + # Phase 2: send 16-byte metadata + per-counter FDs + read ACK. + num_counters = self.layer_done_counter.num_counters + metadata = struct.pack( + "iiii", + self.rank_info.tp_rank_per_node, + self.model_config.tp_size_per_node, + self.rank_info.num_layers_per_pp_stage, + num_counters, + ) + sock.sendall(metadata) + for counter_id in range(num_counters): + fds = self.layer_done_counter.events[counter_id].load_event_fds + send_fds(sock, fds, struct.pack("i", counter_id)) + sock.settimeout(30.0) + try: + ack = sock.recv(1) + except socket.timeout as exc: + raise RuntimeError( + "Timed out waiting for ACK from FlexKV layerwise worker" + ) from exc + if not ack or ack[0] != 1: + raise RuntimeError( + f"FlexKV layerwise worker NACK'd eventfd transfer " + f"(ack={ack!r})" + ) + logger.info( + "[FlexKV] Eventfd handshake complete %s counters=%d layers=%d", + self._label, + num_counters, + self.rank_info.num_layers_per_pp_stage, + ) + return + except Exception as exc: # noqa: BLE001 + last_error = exc + logger.warning( + "[FlexKV] Eventfd handshake send_attempt=%d/%d failed: %s", + send_attempt + 1, + max_send_retries, + exc, + ) + finally: + if sock is not None: + sock.close() + time.sleep(retry_interval) + + raise RuntimeError( + f"[FlexKV] Failed to send eventfds to {self._layerwise_socket} " + f"after {max_send_retries} attempts: {last_error}" + ) diff --git a/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py new file mode 100644 index 000000000..df41c4c82 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py @@ -0,0 +1,511 @@ +"""FlexKV-backed RadixCache for sglang. + +This module exposes :class:`FlexKVRadixCache`, a subclass of +:class:`sglang.srt.mem_cache.radix_cache.RadixCache` that delegates +host-side prefix storage to a FlexKV ``KVManager``. The design mirrors +``LMCRadixCache`` (the LMCache integration) so the scheduler-side +contract is identical: + +* MP (synchronous) mode — the default. + ``match_prefix`` fires only a FlexKV LOOKUP and returns ``host_hit_length``; + the scheduler then calls :meth:`init_load_back` at dispatch time which + allocates slots and fires the FlexKV RETRIEVE. + +* IP (layerwise) mode — enabled with ``FLEXKV_ENABLE_LAYERWISE_TRANSFER=1``. + ``match_prefix`` allocates uncached slots and kicks off a layerwise + load; the per-layer hook registered via + ``register_layer_transfer_counter`` then waits on each layer's + eventfd inside the model's forward pass. + +Selection: ``--enable-flexkv`` on the sglang CLI routes the default +RadixCache factory here. See ``__init__.py`` in this package for the +``register_radix_cache_backend("flexkv", ...)`` entry-point that backs +the explicit ``--radix-cache-backend=flexkv`` form. +""" + +from __future__ import annotations + +import enum +import logging +import threading +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional, Tuple + +import torch + +from sglang.srt.mem_cache.base_prefix_cache import ( + EvictParams, + EvictResult, + InitLoadBackParams, + MatchPrefixParams, + MatchResult, +) +from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode +from sglang.srt.mem_cache.storage.flexkv.flexkv_connector import FlexKVConnector + +if TYPE_CHECKING: + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.managers.schedule_batch import Req + from sglang.srt.mem_cache.cache_init_params import CacheInitParams + from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) + + +class FlexKVMode(enum.Enum): + MP = enum.auto() # synchronous lookup → retrieve in two phases + IP = enum.auto() # in-process layerwise transfer + + +@dataclass +class _LoadBackMarker: + """State carried from a hit-producing ``match_prefix`` to its + matching ``init_load_back``. The detached ``RadixKey`` is a snapshot + of the matched key at lookup time (the live request key aliases + ``req.fill_ids`` which keeps growing).""" + + key: RadixKey + value_numel: int # device tokens already present at lookup time + + +class FlexKVRadixCache(RadixCache): + """RadixCache extended with FlexKV host-tier IO.""" + + def __init__( + self, + params: CacheInitParams, + model_config: Optional[ModelConfig], + server_args: ServerArgs, + tp_rank: int, + tp_size: int, + dp_rank: Optional[int], + pp_rank: int, + attn_cp_rank: int, + tp_group=None, + pp_group=None, + attn_tp_group=None, + attn_cp_group=None, + ) -> None: + super().__init__(params) + + kvcache = self.token_to_kv_pool_allocator.get_kvcache() + # ``tp_group`` and ``attn_tp_group`` are sometimes passed + # interchangeably by sglang's factory; prefer the explicit + # ``attn_tp_group`` when given. + attn_tp_group_eff = attn_tp_group if attn_tp_group is not None else tp_group + + self.flexkv_connector = FlexKVConnector( + sgl_model_config=model_config, + server_args=server_args, + page_size=params.page_size, + kvcache=kvcache, + tp_rank=tp_rank, + dp_rank=dp_rank, + pp_rank=pp_rank, + attn_cp_rank=attn_cp_rank, + pp_group=pp_group, + attn_tp_group=attn_tp_group_eff, + attn_cp_group=attn_cp_group, + ) + + self._mode = ( + FlexKVMode.IP if self.flexkv_connector.enable_layerwise else FlexKVMode.MP + ) + if self._mode is FlexKVMode.IP: + # Register the eventfd counter onto sglang's KV pool so each + # forward layer blocks on its own eventfd. + self.flexkv_connector.register_layer_transfer_counter(kvcache) + + # CUDA streams (mirroring LMCRadixCache). + self.load_stream = torch.cuda.Stream() + self.store_stream = torch.cuda.Stream() + + # Two-phase MP load: stash marker between ``match_prefix`` and + # ``init_load_back``. + self._load_markers: dict[str, _LoadBackMarker] = {} + # ``store_kv`` is async — we keep a lock on the source node + # until FlexKV signals completion, draining in ``evict`` / + # ``check_hicache_events``. + self._inflight_store_nodes: dict[str, TreeNode] = {} + self._node_lock = threading.Lock() + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + + def reset(self) -> None: # type: ignore[override] + super().reset() + if hasattr(self, "_load_markers"): + self._load_markers.clear() + if hasattr(self, "_inflight_store_nodes"): + with self._node_lock: + self._inflight_store_nodes.clear() + if hasattr(self, "flexkv_connector"): + self.flexkv_connector.reset() + + def shutdown(self) -> None: + if hasattr(self, "flexkv_connector"): + self.flexkv_connector.shutdown() + + # ------------------------------------------------------------------ + # match_prefix + # ------------------------------------------------------------------ + + def match_prefix(self, params: MatchPrefixParams) -> MatchResult: # type: ignore[override] + """Look up the longest cached prefix on host KV (FlexKV). + + Dispatches to :meth:`_mp_match_prefix` or :meth:`_ip_match_prefix` + depending on whether layerwise transfer is enabled. + """ + key = params.key + if self.disable or not key: + return super().match_prefix(params) + + # FlexKV operates at page granularity — round the lookup query + # down to a multiple of ``page_size`` so the hit count we report + # back to sglang matches what FlexKV can actually serve. + if self.page_size != 1: + aligned_len = (len(key) // self.page_size) * self.page_size + key = key[:aligned_len] + + base_res = super().match_prefix(params) + if len(key) == 0: + return base_res + + device_value: torch.Tensor = base_res.device_indices + last_node: TreeNode = base_res.last_device_node + + if self._mode is FlexKVMode.MP: + if params.req is None: + return base_res + return self._mp_match_prefix( + key, base_res, device_value, last_node, params.req + ) + return self._ip_match_prefix(key, base_res, device_value, last_node) + + def _mp_match_prefix( + self, + key: RadixKey, + base_res: MatchResult, + device_value: torch.Tensor, + last_node: TreeNode, + req: Req, + ) -> MatchResult: + """LOOKUP-only path. Sets ``host_hit_length`` on the result so + the scheduler later invokes :meth:`init_load_back`.""" + token_ids = key.raw_token_ids() + device_len = int(device_value.numel()) + if device_len >= len(token_ids): + return base_res + + # token_mask=True for tokens NOT on device — FlexKV decides + # which of those it can serve. + token_mask = torch.zeros(len(token_ids), dtype=torch.bool) + token_mask[device_len:] = True + + fkv_task_id, hit = self.flexkv_connector.lookup_kv( + token_ids=token_ids, token_mask=token_mask, rid=req.rid + ) + if hit <= 0: + return base_res + + # Snapshot the matched key (the live key aliases ``req.fill_ids``). + if token_ids is key.token_ids: + token_ids_snap = token_ids[:] + else: + token_ids_snap = token_ids + self._load_markers[req.rid] = _LoadBackMarker( + key=RadixKey(token_ids_snap, key.extra_key, key.is_bigram), + value_numel=device_len, + ) + return MatchResult( + device_indices=device_value, + last_device_node=last_node, + last_host_node=last_node, + best_match_node=last_node, + host_hit_length=hit, + ) + + def _ip_match_prefix( + self, + key: RadixKey, + base_res: MatchResult, + device_value: torch.Tensor, + last_node: TreeNode, + ) -> MatchResult: + """Layerwise path: allocate slots and fire ``start_load_kv_layerwise`` + immediately. Per-layer hook waits during forward.""" + token_ids = key.raw_token_ids() + device_len = int(device_value.numel()) + if device_len >= len(token_ids): + return base_res + + # Quick LOOKUP first to discover how many slots we'd need. + token_mask = torch.zeros(len(token_ids), dtype=torch.bool) + token_mask[device_len:] = True + # No rid here — IP mode self-pops; pass a synthetic stable key. + synthetic_rid = f"_ip_{id(key)}" + _, hit = self.flexkv_connector.lookup_kv( + token_ids=token_ids, token_mask=token_mask, rid=synthetic_rid + ) + if hit <= 0: + return base_res + + result = self._allocate_and_load( + key=key, + value_numel=device_len, + uncached_len=hit, + last_node=last_node, + load_fn=lambda slot_mapping: self.flexkv_connector.start_load_kv_layerwise( + synthetic_rid, slot_mapping + )[0], + ) + if result is None: + return base_res + new_slots, new_node = result + return MatchResult( + device_indices=torch.cat([device_value, new_slots]), + last_device_node=new_node, + last_host_node=new_node, + best_match_node=new_node, + ) + + # ------------------------------------------------------------------ + # init_load_back (MP RETRIEVE) + # ------------------------------------------------------------------ + + def init_load_back( # type: ignore[override] + self, + params: InitLoadBackParams, + ) -> Tuple[torch.Tensor, Optional[TreeNode]]: + """MP RETRIEVE. Allocates uncached slots and fires the FlexKV + load; inserts the resulting TreeNode.""" + req = params.req + last_node: TreeNode = params.best_match_node + marker = self._load_markers.pop(req.rid, None) + if marker is None: + # ``match_prefix`` decided there was no work to do, but the + # scheduler still called us. Release any held task and + # return an empty load. + self.flexkv_connector.release_pending(req.rid) + return ( + torch.empty((0,), dtype=torch.int64, device=self.device), + last_node, + ) + + result = self._allocate_and_load( + key=marker.key, + value_numel=marker.value_numel, + uncached_len=params.host_hit_length, + last_node=last_node, + load_fn=lambda slot_mapping: self.flexkv_connector.retrieve_kv( + req.rid, slot_mapping + ), + ) + if result is None: + # Allocation failed or load returned zero. ``retrieve_kv`` + # already cancels/cleans up on failure paths; release_pending + # is idempotent for the case where allocation failed before + # we even popped the held task. + self.flexkv_connector.release_pending(req.rid) + return ( + torch.empty((0,), dtype=torch.int64, device=self.device), + last_node, + ) + return result + + def _allocate_and_load( + self, + *, + key: RadixKey, + value_numel: int, + uncached_len: int, + last_node: TreeNode, + load_fn, + ) -> Optional[Tuple[torch.Tensor, TreeNode]]: + """Shared allocator + post-load bookkeeping for MP/IP. + + Returns ``(token_slots[:fetched], new_node)`` on success. + ``None`` on either allocation failure or zero retrieved (in + which case all slots are freed). + """ + if uncached_len <= 0: + return None + + # Evict to make room when needed. + if self.token_to_kv_pool_allocator.available_size() < uncached_len: + self.evict(EvictParams(num_tokens=uncached_len)) + token_slots = self.token_to_kv_pool_allocator.alloc(uncached_len) + if token_slots is None: + return None + + # The FlexKV ``launch`` interface takes the slot indices for the + # tokens it will write — no leading ``-1`` padding (FlexKV has + # no concept of "skip these device slots, they're already + # cached"; we pass it exactly the destinations for the + # uncached tail). + num_retrieved = load_fn(token_slots.to(torch.int64)) + + if num_retrieved <= 0: + self.token_to_kv_pool_allocator.free(token_slots) + return None + + # Free the tail of the over-allocation when FlexKV returned + # fewer than expected. + if num_retrieved < uncached_len: + self.token_to_kv_pool_allocator.free(token_slots[num_retrieved:]) + fetched_slots = token_slots[:num_retrieved] + else: + fetched_slots = token_slots + + new_node = TreeNode(priority=last_node.priority) + start = value_numel + end = start + num_retrieved + new_node.key = key[start:end] + new_node.value = fetched_slots + new_node.parent = last_node + last_node.children[new_node.key.child_key(self.page_size)] = new_node + self.evictable_size_ += num_retrieved + self._update_leaf_status(last_node) + self._update_leaf_status(new_node) + + self._record_store_event(new_node.parent) + self._record_store_event(new_node) + + return fetched_slots, new_node + + # ------------------------------------------------------------------ + # cache_finished_req (STORE) + # ------------------------------------------------------------------ + + def cache_finished_req( # type: ignore[override] + self, req: Req, is_insert: bool = True + ) -> None: + """Base cache_finished_req then fire an async FlexKV store.""" + super().cache_finished_req(req, is_insert=is_insert) + if not is_insert: + self._load_markers.pop(req.rid, None) + return + + # Compute the committed prefix mirroring LMCRadixCache's logic. + from sglang.srt.server_args import get_global_server_args + + global_server_args = get_global_server_args() + topk = global_server_args.speculative_eagle_topk + enable_kv_committed_len = topk is None or topk == 1 + if enable_kv_committed_len: + kv_committed_len = req.kv_committed_len + else: + kv_committed_len = len(req.origin_input_ids) + max( + len(req.output_ids) - 1, 0 + ) + + token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len] + if not token_ids: + return + kv_indices = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, :kv_committed_len + ] + + # Anchor on the new last_device_node so FlexKV's lock matches + # the node we'll later unlock when the store completes. + match_result = super().match_prefix( + MatchPrefixParams(key=RadixKey(token_ids, req.extra_key)) + ) + new_last_node = match_result.last_device_node + if new_last_node is None: + return + + self.inc_lock_ref(new_last_node) + try: + with torch.cuda.stream(self.store_stream): + fkv_task_id = self.flexkv_connector.store_kv( + rid=req.rid, + token_ids=list(token_ids), + kv_indices=kv_indices, + ) + except Exception: # noqa: BLE001 + self.dec_lock_ref(new_last_node) + raise + + if fkv_task_id < 0: + # Nothing to write back (either everything already in + # FlexKV, or put_match failed / returned None). + self.dec_lock_ref(new_last_node) + return + + with self._node_lock: + self._inflight_store_nodes[req.rid] = new_last_node + + # ------------------------------------------------------------------ + # evict + completion draining + # ------------------------------------------------------------------ + + def evict(self, params: EvictParams) -> EvictResult: # type: ignore[override] + """Drain completed stores before letting the base evict touch + the source nodes.""" + if self.disable: + return EvictResult() + self._drain_completed_stores() + # Make sure the store stream's GPU work is observed before any + # eviction frees the source slots. + self.store_stream.synchronize() + return super().evict(params) + + def check_hicache_events(self) -> None: # type: ignore[override] + """Periodic non-blocking sweep called by the scheduler tick. + + Drains both store completions (so source nodes get unlocked + quickly) and the launched-load tail (so the FlexKV pipe + doesn't accumulate).""" + self._drain_completed_stores() + self.flexkv_connector.drain_launched_loads() + + def _drain_completed_stores(self) -> None: + completed_rids = self.flexkv_connector.check_completed_stores() + if not completed_rids: + return + with self._node_lock: + for rid in completed_rids: + node = self._inflight_store_nodes.pop(rid, None) + if node is not None: + self.dec_lock_ref(node) + + # ------------------------------------------------------------------ + # Optional pass-throughs used by the scheduler + # ------------------------------------------------------------------ + + def release_aborted_request(self, rid: str) -> None: + """Clean up tracking for an aborted request without invoking FlexKV.""" + self._load_markers.pop(rid, None) + with self._node_lock: + node = self._inflight_store_nodes.pop(rid, None) + if node is not None: + self.dec_lock_ref(node) + self.flexkv_connector.release_pending(rid) + self.flexkv_connector.cancel_prefetch(rid) + + def prefetch_from_storage( + self, rid: str, last_host_node: TreeNode, token_ids + ) -> None: + """Kick off an opportunistic prefetch (SSD/Remote → CPU).""" + try: + self.flexkv_connector.prefetch_async(rid, list(token_ids)) + except Exception as exc: # noqa: BLE001 + logger.debug("[FlexKV] prefetch_from_storage: %s", exc) + + def check_prefetch_progress(self, rid: str) -> bool: + return self.flexkv_connector.check_prefetch_progress(rid) + + def terminate_prefetch(self, rid: str) -> None: + self.flexkv_connector.cancel_prefetch(rid) + + def pop_prefetch_loaded_tokens(self, rid: str) -> int: + # FlexKV doesn't expose per-rid prefetched token counts yet. + return 0 + + @property + def hicache_storage_pass_prefix_keys(self) -> bool: + # We pass token ids, not opaque key strings, so no prefix-key + # accounting in the scheduler. + return False diff --git a/python/sglang/srt/mem_cache/storage/flexkv/verify_outputs.py b/python/sglang/srt/mem_cache/storage/flexkv/verify_outputs.py new file mode 100644 index 000000000..25b7ed26e --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/flexkv/verify_outputs.py @@ -0,0 +1,180 @@ +"""End-to-end correctness check for the FlexKV sglang connector. + +Run twice with different server configurations: + + # 1. Baseline: launch sglang WITHOUT --enable-flexkv first, then: + python verify_outputs.py --phase baseline + + # 2. Restart sglang WITH --enable-flexkv, then: + python verify_outputs.py --phase test + +Each prompt is requested twice in the test phase: + + * R1 (fresh) — first call after server start; FlexKV may still have + state from a previous test run, but match must equal baseline. + * R2 (cached) — after /flush_cache; the GPU radix is empty but + FlexKV's CPU pool keeps the data, so R2 should be a host hit. + +Both R1 and R2 output_ids must byte-equal the baseline. Any mismatch +is reported and exit code is non-zero. Run again with +``FLEXKV_ENABLE_LAYERWISE_TRANSFER=1`` set on the server to exercise +the layerwise path. +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +import urllib.request + +PROMPTS = [ + ( + "PROMPT_SHORT", + "The capital of France is", + 12, + ), + ( + "PROMPT_MEDIUM", + "List the first ten prime numbers in order: 2, 3, 5, ", + 24, + ), + ( + "PROMPT_LONG", + # Long enough to span many KV pages. + ( + "In the year 2025, a research team at a major AI lab released a " + "report describing the architecture of a new large language " + "model. The report had several sections. Section one introduced " + "the model and its training data. Section two covered the " + "attention mechanism in detail, including how the keys and " + "values were managed. Section three discussed deployment, " + "including KV cache offloading to CPU memory and to disk. " + "Section four reported evaluation results on standard " + "benchmarks. Section five concluded with a discussion of " + "future work, including improvements to the offloading layer " + "and to the radix tree used to index cached prefixes. " + "Now, summarize the report in one sentence: " + ), + 60, + ), +] + + +def _post(host: str, path: str, body=None, timeout=120) -> str: + if body is None: + req = urllib.request.Request(f"http://{host}{path}", method="POST") + else: + req = urllib.request.Request( + f"http://{host}{path}", + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=timeout) as resp: + return resp.read().decode() + + +def gen(host: str, text: str, max_new: int) -> dict: + raw = _post( + host, + "/generate", + { + "text": text, + "sampling_params": { + "max_new_tokens": max_new, + "temperature": 0.0, + }, + }, + ) + return json.loads(raw) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--host", + default="127.0.0.1:30000", + help="sglang server host:port (default 127.0.0.1:30000)", + ) + ap.add_argument( + "--phase", + choices=["baseline", "test"], + required=True, + help="baseline: record golden outputs; test: compare against them", + ) + ap.add_argument( + "--baseline-file", + default="/tmp/flexkv_baseline.json", + help="where to write/read the baseline outputs", + ) + args = ap.parse_args() + + if args.phase == "baseline": + result = {} + for name, text, max_new in PROMPTS: + r = gen(args.host, text, max_new) + meta = r["meta_info"] + print( + f"[baseline] {name}: completion={meta['completion_tokens']}, " + f"cached={meta['cached_tokens']}, text={r['text']!r}" + ) + result[name] = { + "text": r["text"], + "output_ids": r["output_ids"], + "completion_tokens": meta["completion_tokens"], + } + with open(args.baseline_file, "w") as f: + json.dump(result, f, indent=2) + print(f"\nWrote baseline to {args.baseline_file}") + return 0 + + with open(args.baseline_file) as f: + baseline = json.load(f) + + errors = 0 + for name, text, max_new in PROMPTS: + b = baseline[name] + + # R1 (fresh): may or may not hit FlexKV depending on prior state. + r1 = gen(args.host, text, max_new) + m1 = r1["meta_info"] + ok1 = r1["output_ids"] == b["output_ids"] + print( + f"[test/{name}] R1 fresh: cached={m1['cached_tokens']}/" + f"{m1['prompt_tokens']}, details={m1.get('cached_tokens_details')}, " + f"output_match={'OK' if ok1 else 'MISMATCH'}" + ) + if not ok1: + print(f" baseline: {b['text']!r}") + print(f" r1 : {r1['text']!r}") + errors += 1 + + # Give the async D2H store a beat to complete before we flush. + time.sleep(2) + _post(args.host, "/flush_cache") + time.sleep(1) + + # R2 (cached): GPU radix is empty; FlexKV must serve the prefix. + r2 = gen(args.host, text, max_new) + m2 = r2["meta_info"] + ok2 = r2["output_ids"] == b["output_ids"] + ratio = m2["cached_tokens"] / max(1, m2["prompt_tokens"]) + print( + f"[test/{name}] R2 cached: cached={m2['cached_tokens']}/" + f"{m2['prompt_tokens']} ({ratio:.1%}), " + f"details={m2.get('cached_tokens_details')}, " + f"output_match={'OK' if ok2 else 'MISMATCH'}" + ) + if not ok2: + print(f" baseline: {b['text']!r}") + print(f" r2 : {r2['text']!r}") + errors += 1 + + print(f"\nTotal mismatches: {errors}") + return 1 if errors else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index ac774b82a..04d96d34d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2227,6 +2227,27 @@ class ServerArgs: "Path to the LMCache YAML configuration file", ] = None + # ------------------------------------------------------------------------- + # FlexKV + # ------------------------------------------------------------------------- + enable_flexkv: A[ + bool, + ( + "Route the default RadixCache through FlexKV's KVManager for " + "host-tier (CPU / SSD / Remote) KV cache offload. Equivalent " + "to --radix-cache-backend=flexkv but also participates in the " + "auto-selection chain alongside --enable-lmcache." + ), + ] = False + flexkv_config_file: A[ + Optional[str], + ( + "Path to the FlexKV YAML / JSON configuration file. " + "Equivalent to setting the FLEXKV_CONFIG_PATH environment " + "variable." + ), + ] = None + # ------------------------------------------------------------------------- # Ktransformers/AMX expert parallelism # ------------------------------------------------------------------------- @@ -5996,6 +6017,11 @@ class ServerArgs: "LMCache is disabled because of using diffusion LLM inference" ) self.enable_lmcache = False + if self.enable_flexkv: + logger.warning( + "FlexKV is disabled because of using diffusion LLM inference" + ) + self.enable_flexkv = False if self.pp_size > 1: logger.warning(