chore: import upstream snapshot with attribution

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=<N>`.
+
+---
+
+## 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/ <remote>:<dir>/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.

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)

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

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

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}"
+        )

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.runtime_context import get_server_args
+
+        global_server_args = get_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

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())