diff --git a/integration_test/vllm_e2e/BUILD b/integration_test/vllm_e2e/BUILD new file mode 100644 index 000000000..7c3fa9134 --- /dev/null +++ b/integration_test/vllm_e2e/BUILD @@ -0,0 +1,196 @@ +package(default_visibility = ["//integration_test:__subpackages__"]) + +# Shared library: orchestration (manager + vLLM + driver + comparison) and the +# verifying connector injected into vLLM via kv_connector_module_path. +# +# NOTE: vLLM / torch / triton are NOT bazel deps: the e2e harness runs the +# tests inside a dedicated venv (one per supported vLLM era) and puts the +# worktree root first on PYTHONPATH -- the connector under test runs from +# source, and torch/vLLM resolve from the venv. These targets are manual + +# GPU-tagged for the same reason; bazel only stages the driver and the +# manager binary. +py_library( + name = "e2e_lib", + srcs = [ + "e2e_lib.py", + "lib_utils.py", + "servers.py", + "test_connector.py", + ], + imports = ["."], + tags = ["no-remote-exec"], + deps = ["@pip_cpu//requests"], +) + +py_test( + name = "test_basic", + timeout = "eternal", + srcs = ["test_basic.py"], + data = [ + "//kv_cache_manager:kv_cache_manager_bin", + ], + imports = ["."], + tags = [ + "exclusive", # GPU tests must run serially to avoid CUDA OOM contention + "gpu", # requires 1+ GPU + "manual", # needs a GPU machine + vLLM venv + model; see README.md + "no-remote-exec", + ], + deps = [":e2e_lib"], +) + +py_test( + name = "test_concurrent", + timeout = "eternal", + srcs = ["test_concurrent.py"], + data = [ + "//kv_cache_manager:kv_cache_manager_bin", + ], + imports = ["."], + tags = [ + "exclusive", # GPU tests must run serially to avoid CUDA OOM contention + "gpu", # requires 1+ GPU + "manual", # needs a GPU machine + vLLM venv + model; see README.md + "no-remote-exec", + ], + deps = [":e2e_lib"], +) + +py_test( + name = "test_tp", + timeout = "eternal", + srcs = ["test_tp.py"], + data = [ + "//kv_cache_manager:kv_cache_manager_bin", + ], + imports = ["."], + tags = [ + "exclusive", # GPU tests must run serially to avoid CUDA OOM contention + "gpu", # requires 1+ GPU + "manual", # needs a GPU machine + vLLM venv + model; see README.md + "no-remote-exec", + ], + deps = [":e2e_lib"], +) + +py_test( + name = "test_full_hit", + timeout = "eternal", + srcs = ["test_full_hit.py"], + data = [ + "//kv_cache_manager:kv_cache_manager_bin", + ], + imports = ["."], + tags = [ + "exclusive", # GPU tests must run serially to avoid CUDA OOM contention + "gpu", # requires 1+ GPU + "manual", # needs a GPU machine + vLLM venv + model; see README.md + "no-remote-exec", + ], + deps = [":e2e_lib"], +) + +py_test( + name = "test_partial_hit", + timeout = "eternal", + srcs = ["test_partial_hit.py"], + data = [ + "//kv_cache_manager:kv_cache_manager_bin", + ], + imports = ["."], + tags = [ + "exclusive", # GPU tests must run serially to avoid CUDA OOM contention + "gpu", # requires 1+ GPU + "manual", # needs a GPU machine + vLLM venv + model; see README.md + "no-remote-exec", + ], + deps = [":e2e_lib"], +) + +py_test( + name = "test_load_failure", + timeout = "eternal", + srcs = ["test_load_failure.py"], + data = [ + "//kv_cache_manager:kv_cache_manager_bin", + ], + imports = ["."], + tags = [ + "exclusive", # GPU tests must run serially to avoid CUDA OOM contention + "gpu", # requires 1+ GPU + "manual", # needs a GPU machine + vLLM venv + model; see README.md + "no-remote-exec", + ], + deps = [":e2e_lib"], +) + +py_test( + name = "test_multi_turn", + timeout = "eternal", + srcs = ["test_multi_turn.py"], + data = [ + "//kv_cache_manager:kv_cache_manager_bin", + ], + imports = ["."], + tags = [ + "exclusive", # GPU tests must run serially to avoid CUDA OOM contention + "gpu", # requires 1+ GPU + "manual", # needs a GPU machine + vLLM venv + model; see README.md + "no-remote-exec", + ], + deps = [":e2e_lib"], +) + +py_test( + name = "test_cross_request_prefix", + timeout = "eternal", + srcs = ["test_cross_request_prefix.py"], + data = [ + "//kv_cache_manager:kv_cache_manager_bin", + ], + imports = ["."], + tags = [ + "exclusive", # GPU tests must run serially to avoid CUDA OOM contention + "gpu", # requires 1+ GPU + "manual", # needs a GPU machine + vLLM venv + model; see README.md + "no-remote-exec", + ], + deps = [":e2e_lib"], +) + +# Meta-test: injects an off-by-one into the connector's token translation and +# asserts the KV verification FAILS -- proof the harness is not vacuous. +py_test( + name = "test_mutation", + timeout = "eternal", + srcs = ["test_mutation.py"], + data = [ + "//kv_cache_manager:kv_cache_manager_bin", + ], + imports = ["."], + tags = [ + "exclusive", # GPU tests must run serially to avoid CUDA OOM contention + "gpu", # requires 1+ GPU + "manual", # needs a GPU machine + vLLM venv + model; see README.md + "no-remote-exec", + ], + deps = [":e2e_lib"], +) + +# Wildcards skip "manual" targets; run the whole suite explicitly with +# bazelisk test //integration_test/vllm_e2e:e2e_tests +test_suite( + name = "e2e_tests", + tags = ["manual"], # keep wildcard builds from expanding into the GPU tests + tests = [ + ":test_basic", + ":test_concurrent", + ":test_cross_request_prefix", + ":test_full_hit", + ":test_load_failure", + ":test_multi_turn", + ":test_mutation", + ":test_partial_hit", + ":test_tp", + ], +) diff --git a/integration_test/vllm_e2e/README.md b/integration_test/vllm_e2e/README.md new file mode 100644 index 000000000..484984f0c --- /dev/null +++ b/integration_test/vllm_e2e/README.md @@ -0,0 +1,121 @@ +# vLLM <-> KVCM End-to-End KV Cache Verification + +End-to-end integration tests for the KVCM vLLM connector +(`kv_cache_manager/py_connector/vllm`). Each test starts a real KVCM manager +(local-file storage backend) and a real vLLM OpenAI server, drives prompts +through the OpenAI API and verifies that the KV cache data saved to / loaded +from KVCM is correct. + +Requires 1-2 GPUs and vLLM (0.22.1, 0.23.0 and 0.26.0 are e2e-verified; the +connector detects the KV cache layout of each era from the tensor shape). + +## What is verified + +The connector translates between three block spaces per `kv_cache_group`: + +``` +KVCM manager block idx -> global token idx -> group logical block + (step 1, connector-only) (step 2/3, shared with vLLM) +``` + +A bug in step 1 is *symmetric*: save gathers from the wrong slots and load +scatters back to the same wrong slots, so a transport round trip alone cannot +detect it. The test breaks the symmetry with `VerifyingConnector` +(`test_connector.py`), a subclass of the production connector that +independently captures KV data from vLLM's paged cache using only vLLM's own +block-table mapping: + +1. **Phase 1** — fresh prompts: prefill -> connector saves to KVCM. The saved + token ranges are captured from the paged cache (**reference** captures). +2. **Phase 2** — same prompts + suffix: connector reports an external match and + loads from KVCM. The loaded blocks are captured (**loaded** captures). +3. The driver (`e2e_lib.py`) matches loaded captures against references by + token content and compares per layer, requiring bit-exact equality (the + transfer is a verbatim byte round trip; all scenarios achieve it). + +## Model coverage + +The same test targets run against either model kind, selected by +`KVCM_E2E_MODEL`: + +| Kind | Example | Groups | Orchestration | +|---|---|---|---| +| Full attention | Qwen2.5-7B-Instruct | 1 `FullAttentionSpec` | prefix caching off, one server for both phases | +| Hybrid | Qwen3.5-4B | 3 `MambaSpec` + 1 `FullAttentionSpec` | prefix caching on (`mamba_cache_mode="align"`), server restarted between phases so phase 2 loads from KVCM instead of the local prefix cache | + +Hybrid specifics verified: + +* Per-group location specs (`tp{rank}_g{group}`) and per-group block tables. +* Attention groups: token-granular gather/scatter through the Triton kernel. +* Mamba/linear groups: per-block opaque state copy, where a manager block's + *last* token selects the state block (`_state_block_ids`). + +## Scenarios + +| Test | TP | Prompts | Notes | +|---|---|---|---| +| `test_basic` | 1 | 1 | Minimal save -> load round trip | +| `test_concurrent` | 1 | 4 | Concurrent requests: ReqState tracking, per-request block attribution | +| `test_tp` | 2 | 2 | TP coordination; for full-attention models also `preferred_block_size=32` != vLLM block size (16), forcing real cross-block translation | +| `test_partial_hit` | 1 | 1 | Phase 2 extends the prompt mid-block: partial external hit | +| `test_full_hit` | 1 | 1 | Phase 2 resends the identical prompt: full-prompt hit is capped so >= 1 token is recomputed | +| `test_multi_turn` | 1 | 1 | Growing conversation: each turn loads the previous turns' blocks and saves new ones | +| `test_cross_request_prefix` | 1 | 2 | Request B is a strict token prefix of saved request A, ending inside one of A's blocks. Hybrid: B's match must be truncated to the last block whose recurrent state was really materialized, and B's output must be token-identical to a no-cache reference | +| `test_load_failure` | 1 | 1 | Storage files deleted between phases: load fails, retry loop must not spin, request still completes | +| `test_mutation` | 1 | 1 | Meta-test: injected off-by-one in the slot translation must make verification FAIL (proves the harness is not vacuous) | + +## Running + +These targets are tagged `manual`: they need a GPU machine with a prepared +vLLM venv and a local model, so `bazelisk test //integration_test/...` skips +them and they must be requested explicitly (see below). + +Build prerequisites (from the repo root): + +```bash +bazelisk build //kv_cache_manager:kv_cache_manager_bin \ + //kv_cache_manager/client/pybind:kvcm_py_client_lib_wheel \ + //kv_cache_manager/py_connector/vllm:kvcm_vllm_connector_wheel \ + --per_file_copt='external/jsoncpp_git/.*@-Wno-error' +``` + +Install both wheels into the vLLM venv (rename them first: the Bazel output +name contains unstamped `{STABLE_*}` template variables; read the real version +from the wheel's `METADATA`). + +Run (tagged `exclusive`, so they execute serially): + +```bash +bazelisk test //integration_test/vllm_e2e:e2e_tests \ + --cache_test_results=no --test_output=errors \ + --test_env=KVCM_E2E_PYTHON=/path/to/vllm-venv/bin/python \ + --test_env=KVCM_E2E_MODEL=/path/to/model \ + --per_file_copt='external/jsoncpp_git/.*@-Wno-error' +``` + +## Environment variables + +All environment variables used by the e2e harness: + +| Variable | Required | Meaning | +|---|---|---| +| `KVCM_E2E_MODEL` | yes | Path to a local HF model directory (`config.json` + weights). Full-attention coverage needs a plain attention model (e.g. Qwen2.5-7B-Instruct); hybrid coverage needs a mamba/linear + attention model (e.g. Qwen3.5-4B). Hybrid models are auto-detected from `config.json`. | +| `KVCM_E2E_PYTHON` | yes | Python interpreter of a venv with vLLM (any supported version, see above) and both KVCM wheels (`kvcm_py_client`, `kvcm_vllm_connector`) installed. | +| `KVCM_E2E_CAPTURE_DIR` | internal | Set by the driver for the vLLM subprocess; tells `VerifyingConnector` where to write `.pt` captures. Do not set manually. | + +The driver also sets vLLM knobs for the spawned server (`VLLM_KV_CACHE_LAYOUT=NHD`, +`VLLM_ATTENTION_BACKEND=FLASH_ATTN`, `VLLM_USE_FLASHINFER_SAMPLER=0`, +`FLASHINFER_DISABLE_VERSION_CHECK=1`) via `env.setdefault`, so a value you +export yourself wins. + +## Debugging + +Bazel's `test.log` only shows the driver's view (e.g. HTTP 500). The real +tracebacks live in the scenario workdir under `$TEST_TMPDIR`: + +``` +/kvcm_vllm_e2e// + manager/manager.stdout|stderr # KVCM manager + vllm/vllm*.stdout|stderr # vLLM (EngineCore tracebacks are here) + captures/{ref|loaded}_tp{rank}_{token_hash}.pt +``` diff --git a/integration_test/vllm_e2e/e2e_lib.py b/integration_test/vllm_e2e/e2e_lib.py new file mode 100644 index 000000000..458cd1e99 --- /dev/null +++ b/integration_test/vllm_e2e/e2e_lib.py @@ -0,0 +1,17 @@ +"""Compatibility facade over the split harness modules. + +Pure helpers live in lib_utils.py, process lifecycle (manager binary, +vLLM server, ScenarioEnv) in servers.py; this module re-exports the +original surface so existing scenario files keep their imports. +""" + +from lib_utils import ( # noqa: F401 + is_hybrid_model, _runfiles_root, find_repo_root, find_manager_binary, + find_python, free_port, wait_http, tokenize, get_manager_block_size, + block_token_hash, full_block_hashes, wait_for_prefix_cached, + send_completions, count_captures, wait_for_captures, compare_captures, + assert_report_ok, make_base_prompts, shared_token_prefix_len, +) +from servers import ( # noqa: F401 + ManagerProcess, VllmServer, ScenarioEnv, run_e2e, +) diff --git a/integration_test/vllm_e2e/lib_utils.py b/integration_test/vllm_e2e/lib_utils.py new file mode 100644 index 000000000..c3ee5c4f6 --- /dev/null +++ b/integration_test/vllm_e2e/lib_utils.py @@ -0,0 +1,423 @@ +"""Pure helpers for the vLLM e2e harness: model detection, paths, +tokenization, manager queries, prompt/answer traffic and capture +comparison. No process state lives here.""" + + +import glob +import json +import logging +import os +import shutil +import socket +import subprocess +import time +import uuid +from typing import Optional + +import requests + +logger = logging.getLogger("vllm_e2e") +# This module only runs inside test drivers; make the orchestration evidence +# (block counts, verification report, log-scan results) visible in test.log. +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(name)s %(levelname)s %(message)s") + +# Required: path to a local HF model directory (config.json + weights). +# Full-attention coverage needs a plain attention model (e.g. +# Qwen2.5-7B-Instruct); hybrid coverage needs a mamba/linear + attention model +# (e.g. Qwen3.5-4B). Hybrid models are auto-detected from config.json. +MODEL_PATH = os.environ.get("KVCM_E2E_MODEL") +if not MODEL_PATH: + raise RuntimeError( + "KVCM_E2E_MODEL is not set. Point it at a local model directory, e.g. " + "--test_env=KVCM_E2E_MODEL=/path/to/Qwen2.5-7B-Instruct (full " + "attention) or /path/to/Qwen3.5-4B (hybrid). See " + "integration_test/vllm_e2e/README.md.") + +# Model-agnostic alias the driver uses in OpenAI API requests, so the tests do +# not depend on the model directory name. +SERVED_MODEL_NAME = "e2e-model" + + + +def is_hybrid_model(model_path: str) -> bool: + """Detect a hybrid (mamba/linear + full attention) model from its config.""" + try: + with open(os.path.join(model_path, "config.json")) as f: + cfg = json.load(f) + except Exception: + return False + text_cfg = cfg.get("text_config", cfg) + # Any known hybrid marker wins: architecture names cover the families we + # know (e.g. Qwen3NextForCausalLM), the layer knobs cover configs that + # interleave linear/mamba layers with full attention without naming a + # known family (e.g. Qwen3_5ForConditionalGeneration declares its own + # architecture). The signals are ORed -- a non-matching architecture + # string must never shadow the knobs. + archs = ", ".join(cfg.get("architectures", []) or []) + arch_markers = ("Qwen3Next", "Zamba", "FalconH1", "Samba", "Jamba") + return ( + any(m in archs for m in arch_markers) + or "full_attention_interval" in text_cfg + or "linear_conv_kernel_dim" in text_cfg + or str(cfg.get("model_type", "")).startswith("qwen3_") + ) + + +# --------------------------------------------------------------------------- # +# Paths / binaries +# --------------------------------------------------------------------------- # + +def _runfiles_root() -> Optional[str]: + return os.environ.get("RUNFILES_DIR") or os.environ.get("TEST_SRCDIR") + + +def find_repo_root() -> str: + """Locate the KVCM repository root (works under Bazel runfiles and plain).""" + here = os.path.dirname(os.path.abspath(__file__)) + # integration_test/vllm_e2e/e2e_lib.py -> repo root is two levels up. + candidate = os.path.abspath(os.path.join(here, "..", "..")) + if os.path.exists(os.path.join(candidate, "WORKSPACE")): + return candidate + runfiles = _runfiles_root() + if runfiles: + cand = os.path.join(runfiles, "kv_cache_manager") + if os.path.exists(os.path.join(cand, "WORKSPACE")): + return cand + return candidate + + +def find_manager_binary(repo_root: str) -> str: + candidates = [ + os.path.join(repo_root, "bazel-bin/kv_cache_manager/kv_cache_manager_bin"), + os.path.join(repo_root, "bazel-out/k8-opt/bin/kv_cache_manager/kv_cache_manager_bin"), + ] + runfiles = _runfiles_root() + if runfiles: + candidates.append( + os.path.join(runfiles, "kv_cache_manager", "kv_cache_manager", + "kv_cache_manager_bin") + ) + for c in candidates: + if os.path.exists(c): + return c + raise RuntimeError( + "kv_cache_manager_bin not found; build it with: " + "bazelisk build //kv_cache_manager:kv_cache_manager_bin" + ) + + +def find_python() -> str: + # Required: python interpreter of a venv with vLLM >= 0.26.0 and both KVCM + # wheels (kvcm_py_client, kvcm_vllm_connector) installed; see README.md. + python = os.environ.get("KVCM_E2E_PYTHON") + if not python: + raise RuntimeError( + "KVCM_E2E_PYTHON is not set. Point it at the python of a vLLM " + "venv with the KVCM wheels installed, e.g. " + "--test_env=KVCM_E2E_PYTHON=/path/to/venv/bin/python. See " + "integration_test/vllm_e2e/README.md.") + return python + + +def free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + # Loopback only: everything in this harness is single-machine. + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def wait_http(url: str, timeout: float, post_body: Optional[dict] = None) -> bool: + deadline = time.time() + timeout + while time.time() < deadline: + try: + if post_body is not None: + r = requests.post(url, json=post_body, timeout=3) + else: + r = requests.get(url, timeout=3) + if r.status_code < 500: + return True + except Exception: + pass + time.sleep(1.0) + return False + + +# --------------------------------------------------------------------------- # +# KVCM manager + +# --------------------------------------------------------------------------- # +def tokenize(base_url: str, prompt: str) -> list[int]: + r = requests.post(f"{base_url}/tokenize", + json={"model": SERVED_MODEL_NAME, "prompt": prompt}, timeout=60) + r.raise_for_status() + return r.json()["tokens"] + + +def get_manager_block_size(manager_uri: str, instance_id: str) -> int: + """Ask the manager for the registered instance's manager block size.""" + r = requests.post(f"{manager_uri}/api/getInstanceInfo", + json={"trace_id": "e2e_bs", "instance_id": instance_id}, + timeout=10) + r.raise_for_status() + block_size = r.json()["instance_info"]["block_size"] + assert block_size > 0, f"bad manager block size: {block_size}" + return block_size + + +def block_token_hash(token_ids: list[int]) -> str: + """Token-content hash used in capture file names (mirrors test_connector).""" + import hashlib + import torch + return hashlib.sha256( + torch.tensor(token_ids, dtype=torch.int64).numpy().tobytes() + ).hexdigest()[:16] + + +def full_block_hashes(token_ids: list[int], manager_block_size: int) -> list[str]: + """Per-manager-block capture hashes for the full blocks of a token stream.""" + n = len(token_ids) // manager_block_size + return [ + block_token_hash(token_ids[i * manager_block_size:(i + 1) * manager_block_size]) + for i in range(n) + ] + + +def wait_for_prefix_cached(manager_uri: str, instance_id: str, + token_ids: list[int], min_blocks: int, + timeout: float = 120.0) -> bool: + """Poll the manager until at least min_blocks of the prefix are committed. + + Mirrors what the connector's get_num_new_matched_tokens queries + (query_type=QT_PREFIX_MATCH, block_mask offset=0 for a fresh request), so it + guarantees phase 2 will actually hit the external cache for all min_blocks. + """ + deadline = time.time() + timeout + payload = { + "trace_id": "e2e_probe", + "token_ids": token_ids, + "instance_id": instance_id, + "query_type": "QT_PREFIX_MATCH", + "block_mask": {"offset": 0}, + } + while time.time() < deadline: + try: + r = requests.post(f"{manager_uri}/api/getCacheLocation", + json=payload, timeout=10) + if r.status_code == 200: + data = r.json() + if data.get("header", {}).get("status", {}).get("code") == "OK": + locs = data.get("locations", []) + if len(locs) >= min_blocks: + logger.info("prefix cached: %d location(s)", len(locs)) + return True + except Exception: + pass + time.sleep(1.0) + logger.warning("timed out waiting for prefix to be cached") + return False + + +def send_completions(base_url: str, prompts: list, max_tokens: int = 4, + temperature: float = 0.0, **extra_payload) -> list[dict]: + """Send prompts concurrently and return the OpenAI responses. + + Each prompt may be a string or a list of token ids (the completions API + accepts both). extra_payload is merged into the request body (e.g. + return_token_ids=True).""" + from concurrent.futures import ThreadPoolExecutor + + client_url = f"{base_url}/v1/completions" + + def _one(prompt) -> dict: + payload = { + "model": SERVED_MODEL_NAME, + "prompt": prompt, + "max_tokens": max_tokens, + "temperature": temperature, + **extra_payload, + } + r = requests.post(client_url, json=payload, timeout=300) + r.raise_for_status() + return r.json() + + with ThreadPoolExecutor(max_workers=max(1, len(prompts))) as ex: + return list(ex.map(_one, prompts)) + + +# --------------------------------------------------------------------------- # +# Capture comparison + +# --------------------------------------------------------------------------- # +def count_captures(capture_dir: str, kind: str) -> int: + return len(glob.glob(os.path.join(capture_dir, f"{kind}_*.pt"))) + + +def wait_for_captures(capture_dir: str, kind: str, expected: int, + timeout: float = 120.0) -> int: + """Wait until at least ``expected`` captures of ``kind`` exist. + + Raises AssertionError on timeout: a missing capture means the connector + never exercised the code path under test, so the scenario must fail rather + than silently verify fewer blocks. + """ + deadline = time.time() + timeout + while time.time() < deadline: + n = count_captures(capture_dir, kind) + if n >= expected: + logger.info("saw %d/%d %s captures", n, expected, kind) + return n + time.sleep(1.0) + n = count_captures(capture_dir, kind) + raise AssertionError( + f"timed out waiting for {kind} captures: got {n}, want {expected}") + + +def compare_captures(capture_dir: str, tp_size: int) -> dict: + """Compare loaded captures against reference captures. + + Every block that was *loaded* from KVCM must correspond to a *reference* + capture (same tp rank + token content) with matching KV data. The direction + matters: saves are incremental, so some saved blocks may legitimately not be + reloaded (e.g. the tokenization boundary block) -- but every loaded block + must match something that was saved. + + Returns a report dict; the caller asserts on it. + """ + import torch + + refs = {} + loaded = {} + for path in glob.glob(os.path.join(capture_dir, "*.pt")): + name = os.path.basename(path)[:-3] # strip .pt + parts = name.split("_") + kind, tp, token_hash = parts[0], parts[1], "_".join(parts[2:]) + key = (tp, token_hash) + (refs if kind == "ref" else loaded)[key] = path + + report = { + "num_refs": len(refs), + "num_loaded": len(loaded), + "matched": 0, + "bit_exact": 0, + "failures": [], + "loaded_without_ref": [], + "matched_keys": [], + } + + for key, loaded_path in sorted(loaded.items()): + if key not in refs: + report["loaded_without_ref"].append(key) + continue + ref = torch.load(refs[key], map_location="cpu", weights_only=True) + got = torch.load(loaded_path, map_location="cpu", weights_only=True) + + assert ref["token_ids"] == got["token_ids"], f"token id mismatch for {key}" + + all_bit_exact = True + # Compare every layer present in the *loaded* capture: each one was + # actually written by the connector and must match its reference. + # Mamba "align" state layers can legitimately be absent on either side + # (vLLM materializes states only at segment boundaries; interior blocks + # get the null block, which is neither saved nor loaded) -- but a + # loaded layer without a reference is a hard error. + for layer_name, got_kv in got["kv"].items(): + assert layer_name in ref["kv"], ( + f"loaded layer {layer_name} of {key} has no reference capture") + ref_kv = ref["kv"][layer_name] + # Attention groups are a single Tensor; mamba/linear/gdn groups are a + # list[Tensor] (e.g. [conv_state, ssm_state]). Compare uniformly. + if isinstance(ref_kv, (list, tuple)): + ref_parts = list(ref_kv) + got_parts = list(got_kv) + assert len(ref_parts) == len(got_parts), ( + f"state count mismatch {layer_name}: " + f"{len(ref_parts)} vs {len(got_parts)}" + ) + else: + ref_parts = [ref_kv] + got_parts = [got_kv] + + for si, (ref_t, got_t) in enumerate(zip(ref_parts, got_parts)): + assert ref_t.shape == got_t.shape, ( + f"shape mismatch {layer_name}[{si}]: {ref_t.shape} vs {got_t.shape}" + ) + # The transfer is a verbatim byte round trip, so the loaded data + # must be bit-identical to what was saved -- no tolerance. + if not torch.equal(ref_t, got_t): + all_bit_exact = False + report["failures"].append({ + "key": key, + "layer": f"{layer_name}[{si}]", + "mismatched_elems": int((ref_t != got_t).sum()), + "num_elems": ref_t.numel(), + }) + + report["matched"] += 1 + report["matched_keys"].append(key) + if all_bit_exact: + report["bit_exact"] += 1 + + return report + + +def assert_report_ok(report: dict, min_matched: int = 1): + """Assert the comparison succeeded. + + min_matched is the exact lower bound of (ref, loaded) capture pairs computed + from the prompts' tokenization (num full manager blocks x tp ranks); a lower + count means some blocks were silently never saved or never loaded. + """ + problems = [] + if report["loaded_without_ref"]: + problems.append( + f"loaded captures with no matching reference: {report['loaded_without_ref']}" + ) + if report["failures"]: + problems.append(f"bit-exact failures: {report['failures']}") + if report["matched"] < min_matched: + problems.append( + f"matched {report['matched']} loaded captures, expected >= {min_matched}" + ) + if problems: + raise AssertionError("KV verification failed: " + "; ".join(problems)) + logger.info( + "KV verification OK: matched=%d bit_exact=%d (refs=%d loaded=%d)", + report["matched"], report["bit_exact"], + report["num_refs"], report["num_loaded"], + ) + + +# --------------------------------------------------------------------------- # +# Scenario runner + +def make_base_prompts(num_prompts: int, hybrid: bool) -> list[str]: + """Distinct, deterministic prompts. Each sentence carries a unique counter + so every manager block has unique token content -- this avoids hash + collisions between blocks with identical text but different KV (RoPE is + position-dependent). + + Hybrid models pin the manager block size to the scheduler block size (528), + so their prompts must be much longer to span multiple manager blocks + (empirically 140 sentences ~ 2100 tokens > 3 x 528). Full-attention models + use manager blocks of 16/32 tokens, where 40 sentences (~580 tokens) already + span dozens of blocks. + """ + num_sentences = 140 if hybrid else 40 + return [ + f"Prompt number {i}. " + " ".join( + f"Sentence {j} of prompt {i} has value {j * 7 + i * 131}." + for j in range(num_sentences) + ) + for i in range(num_prompts) + ] + + +def shared_token_prefix_len(a: list[int], b: list[int]) -> int: + n = 0 + for x, y in zip(a, b): + if x != y: + break + n += 1 + return n diff --git a/integration_test/vllm_e2e/servers.py b/integration_test/vllm_e2e/servers.py new file mode 100644 index 000000000..dc484f4d8 --- /dev/null +++ b/integration_test/vllm_e2e/servers.py @@ -0,0 +1,467 @@ +"""Process lifecycle for the vLLM e2e harness: the real KVCM manager +binary and the real vLLM OpenAI server, plus ScenarioEnv which owns +their lifetimes for one scenario.""" + +from lib_utils import ( + MODEL_PATH, assert_report_ok, block_token_hash, compare_captures, + count_captures, find_manager_binary, find_python, find_repo_root, + free_port, full_block_hashes, get_manager_block_size, is_hybrid_model, + make_base_prompts, send_completions, shared_token_prefix_len, + tokenize, wait_for_captures, wait_for_prefix_cached, wait_http, +) + + +import glob +import json +import logging +import os +import shutil +import socket +import subprocess +import time +import uuid +from typing import Optional + +import requests + +logger = logging.getLogger("vllm_e2e") +# This module only runs inside test drivers; make the orchestration evidence +# (block counts, verification report, log-scan results) visible in test.log. +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(name)s %(levelname)s %(message)s") + +# Required: path to a local HF model directory (config.json + weights). +# Full-attention coverage needs a plain attention model (e.g. +# Qwen2.5-7B-Instruct); hybrid coverage needs a mamba/linear + attention model +# (e.g. Qwen3.5-4B). Hybrid models are auto-detected from config.json. +MODEL_PATH = os.environ.get("KVCM_E2E_MODEL") +if not MODEL_PATH: + raise RuntimeError( + "KVCM_E2E_MODEL is not set. Point it at a local model directory, e.g. " + "--test_env=KVCM_E2E_MODEL=/path/to/Qwen2.5-7B-Instruct (full " + "attention) or /path/to/Qwen3.5-4B (hybrid). See " + "integration_test/vllm_e2e/README.md.") + +# Model-agnostic alias the driver uses in OpenAI API requests, so the tests do +# not depend on the model directory name. +SERVED_MODEL_NAME = "e2e-model" + +# --------------------------------------------------------------------------- # +class ManagerProcess: + def __init__(self, workdir: str, storage_root: str, key_count_per_file: int = 8): + self.workdir = workdir + os.makedirs(workdir, exist_ok=True) + self.rpc_port = free_port() + self.http_port = free_port() + self.admin_rpc_port = free_port() + self.admin_http_port = free_port() + self.storage_root = storage_root + self.key_count_per_file = key_count_per_file + self.proc: Optional[subprocess.Popen] = None + self.config_path = os.path.join(workdir, "startup_config.json") + + def manager_uri(self) -> str: + return f"http://127.0.0.1:{self.http_port}" + + def _write_config(self): + cfg = { + "storage_config": { + "type": "file", + "global_unique_name": "nfs_01", + "storage_spec": { + # The backend concatenates root_path + key with no + # separator; the trailing slash keeps files inside the dir. + "root_path": self.storage_root.rstrip("/") + "/", + "key_count_per_file": self.key_count_per_file, + }, + }, + "instance_group": { + "name": "default", + "storage_candidates": ["nfs_01"], + "global_quota_group_name": "default_quota_group", + "max_instance_count": 100, + "quota": { + "capacity": 30000000000, + "quota_config": [ + {"storage_type": "file", "capacity": 10000000000}, + {"storage_type": "hf3fs", "capacity": 10000000000}, + {"storage_type": "pace", "capacity": 10000000000}, + ], + }, + "cache_config": { + "reclaim_strategy": { + "reclaim_policy": 1, + "trigger_strategy": {"used_percentage": 0.8}, + "delay_before_delete_ms": 1000, + }, + "cache_prefer_strategy": 2, + "meta_indexer_config": { + "max_key_count": 1000000, + "mutex_shard_num": 16, + "batch_key_size": 16, + "meta_storage_backend_config": { + "storage_type": "local", + "storage_uri": "", + }, + "meta_cache_policy_config": { + "type": "LRU", + "capacity": 10000, + "cache_shard_bits": 0, + "high_pri_pool_ratio": 0.0, + }, + }, + }, + "user_data": '{"description": "vllm e2e test instance group"}', + "version": 1, + }, + } + with open(self.config_path, "w") as f: + json.dump(cfg, f, indent=2) + + def start(self, repo_root: str): + self._write_config() + binary = find_manager_binary(repo_root) + cmd = [ + binary, + "--env", f"kvcm.service.rpc_port={self.rpc_port}", + "--env", f"kvcm.service.http_port={self.http_port}", + "--env", f"kvcm.service.admin_rpc_port={self.admin_rpc_port}", + "--env", f"kvcm.service.admin_http_port={self.admin_http_port}", + "--env", f"kvcm.startup_config={self.config_path}", + "--env", "kvcm.logger.log_level=5", + ] + logger.info("starting manager: %s (cwd=%s)", " ".join(cmd), self.workdir) + self.proc = subprocess.Popen( + cmd, + cwd=self.workdir, + stdout=open(os.path.join(self.workdir, "manager.stdout"), "w"), + stderr=open(os.path.join(self.workdir, "manager.stderr"), "w"), + ) + if not wait_http( + f"{self.manager_uri()}/api/getClusterInfo", + timeout=60, + post_body={"trace_id": "probe", "instance_id": "probe"}, + ): + raise RuntimeError("manager did not become ready; see manager.stderr") + logger.info("manager ready at %s", self.manager_uri()) + + def stop(self): + if self.proc and self.proc.poll() is None: + self.proc.terminate() + try: + self.proc.wait(timeout=10) + except subprocess.TimeoutExpired: + self.proc.kill() + + +# --------------------------------------------------------------------------- # +# vLLM server + +# --------------------------------------------------------------------------- # +class VllmServer: + def __init__(self, workdir: str, capture_dir: str, manager_uri: str, + tp_size: int, coordinator_base_port: int, + instance_id: str, preferred_block_size: int, + enable_prefix_caching: bool, + connector_name: str = "VerifyingConnector", + log_level: str = "INFO", + extra_config_overrides: Optional[dict] = None, + kv_load_failure_policy: Optional[str] = None): + self.workdir = workdir + os.makedirs(workdir, exist_ok=True) + self.capture_dir = capture_dir + os.makedirs(capture_dir, exist_ok=True) + self.port = free_port() + self.manager_uri = manager_uri + self.tp_size = tp_size + self.coordinator_base_port = coordinator_base_port + self.instance_id = instance_id + self.preferred_block_size = preferred_block_size + self.enable_prefix_caching = enable_prefix_caching + self.connector_name = connector_name + self.log_level = log_level + self.extra_config_overrides = extra_config_overrides or {} + self.kv_load_failure_policy = kv_load_failure_policy + self.proc: Optional[subprocess.Popen] = None + + def base_url(self) -> str: + return f"http://127.0.0.1:{self.port}" + + def start(self, repo_root: str, log_suffix: str = ""): + extra_config = { + "manager_uri": self.manager_uri, + "coordinator_base_port": self.coordinator_base_port, + "instance_group": "default", + "instance_id": self.instance_id, + "preferred_block_size": self.preferred_block_size, + "log_level": self.log_level, + } + extra_config.update(self.extra_config_overrides) + kv_transfer_config = { + "kv_connector": self.connector_name, + "kv_role": "kv_both", + "kv_connector_module_path": "test_connector", + "kv_connector_extra_config": extra_config, + } + if self.kv_load_failure_policy: + kv_transfer_config["kv_load_failure_policy"] = self.kv_load_failure_policy + # Serving knobs, override-able via vllm_args (a scenario can widen + # max-model-len or raise gpu-memory-utilization without editing the + # harness); keys are vLLM CLI flags without the leading "--". + vllm_args = { + "max-model-len": "4096", + "gpu-memory-utilization": "0.85", + "enforce-eager": None, + "max-num-seqs": "16", + **getattr(self, "vllm_args", {}), + } + cmd = [ + find_python(), "-m", "vllm.entrypoints.openai.api_server", + "--model", MODEL_PATH, + "--served-model-name", SERVED_MODEL_NAME, + "--host", "127.0.0.1", # loopback only: single-machine harness + "--port", str(self.port), + "--tensor-parallel-size", str(self.tp_size), + ] + for flag, value in vllm_args.items(): + cmd += [f"--{flag}"] if value is None else [f"--{flag}", str(value)] + cmd += ["--kv-transfer-config", json.dumps(kv_transfer_config)] + if self.enable_prefix_caching: + # Hybrid models need prefix caching to expose per-group block tables + # (mamba_cache_mode="align"); align mode requires chunked prefill. + cmd += ["--enable-prefix-caching", "--enable-chunked-prefill"] + else: + cmd += ["--no-enable-prefix-caching"] + env = os.environ.copy() + # The connector modules must resolve to THIS checkout, ahead of any + # kv_cache_manager wheel installed in the vLLM venv. + env["PYTHONPATH"] = os.pathsep.join([ + os.path.dirname(os.path.abspath(__file__)), + find_repo_root(), + env.get("PYTHONPATH", ""), + ]) + env["KVCM_E2E_CAPTURE_DIR"] = self.capture_dir + # Required, not cosmetic: the connector asserts token-major pages + # (NHD memory order) at register_kv_caches; on a vLLM whose default + # layout is HND the connector refuses to start without this, and a + # capture comparison across layouts would be meaningless anyway. + env.setdefault("VLLM_KV_CACHE_LAYOUT", "NHD") + # Force FlashAttention for the full-attention layers: it produces the + # [2, num_blocks, block_size, num_kv_heads, head_size] layout the + # connector expects, and avoids the flashinfer backend entirely. + env.setdefault("VLLM_ATTENTION_BACKEND", "FLASH_ATTN") + # Use the PyTorch-native sampler; the flashinfer sampler JIT-compiles + # with ninja, which is not available in the test environment. + env.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0") + # The test venv ships mismatched flashinfer / flashinfer-cubin wheels; + # skip the version check so importing vLLM's attention registry does not + # crash before the FlashAttention backend is selected. + env.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1") + logger.info("starting vllm: %s", " ".join(cmd)) + self.proc = subprocess.Popen( + cmd, + cwd=self.workdir, + env=env, + stdout=open(os.path.join(self.workdir, f"vllm{log_suffix}.stdout"), "w"), + stderr=open(os.path.join(self.workdir, f"vllm{log_suffix}.stderr"), "w"), + ) + if not wait_http(f"{self.base_url()}/health", timeout=600): + raise RuntimeError("vllm did not become ready; see vllm.stderr") + logger.info("vllm ready at %s", self.base_url()) + + def stop(self): + if self.proc and self.proc.poll() is None: + self.proc.terminate() + try: + self.proc.wait(timeout=15) + except subprocess.TimeoutExpired: + self.proc.kill() + + +# --------------------------------------------------------------------------- # +# Request driver + +# --------------------------------------------------------------------------- # +class ScenarioEnv: + """Owns one scenario's manager + vLLM server lifecycle and scratch dirs. + + Custom scenarios (partial-hit, full-hit, load-failure, multi-turn) share + this; run_e2e keeps its own two-phase flow on top of the same pieces. + """ + + def __init__(self, scenario: str, tp_size: int = 1, + preferred_block_size: int = 0, + enable_prefix_caching: Optional[bool] = None, + connector_name: str = "VerifyingConnector", + log_level: str = "INFO", + extra_config_overrides: Optional[dict] = None, + key_count_per_file: int = 8, + kv_load_failure_policy: Optional[str] = None): + self.scenario = scenario + self.tp_size = tp_size + self.hybrid = is_hybrid_model(MODEL_PATH) + self.preferred_block_size = 0 if self.hybrid else preferred_block_size + self.enable_prefix_caching = (self.hybrid if enable_prefix_caching is None + else enable_prefix_caching) + self.connector_name = connector_name + self.log_level = log_level + self.extra_config_overrides = extra_config_overrides + self.kv_load_failure_policy = kv_load_failure_policy + + self.repo_root = find_repo_root() + scratch_root = (os.environ.get("TEST_TMPDIR") + or os.environ.get("TMPDIR") or "/tmp") + self.base_workdir = os.path.join(scratch_root, "kvcm_vllm_e2e", scenario) + if os.path.exists(self.base_workdir): + shutil.rmtree(self.base_workdir) + self.storage_root = os.path.join(self.base_workdir, "nfs") + self.capture_dir = os.path.join(self.base_workdir, "captures") + self.vllm_dir = os.path.join(self.base_workdir, "vllm") + os.makedirs(self.storage_root, exist_ok=True) + + self.instance_id = f"e2e-{scenario}-{uuid.uuid4().hex[:8]}" + self.manager = ManagerProcess( + os.path.join(self.base_workdir, "manager"), self.storage_root, + key_count_per_file=key_count_per_file) + self.vllm: Optional[VllmServer] = None + + def start_manager(self): + self.manager.start(self.repo_root) + + def start_vllm(self, log_suffix: str = "") -> VllmServer: + self.vllm = VllmServer( + self.vllm_dir, self.capture_dir, self.manager.manager_uri(), + self.tp_size, coordinator_base_port=free_port(), + instance_id=self.instance_id, + preferred_block_size=self.preferred_block_size, + enable_prefix_caching=self.enable_prefix_caching, + connector_name=self.connector_name, + log_level=self.log_level, + extra_config_overrides=self.extra_config_overrides, + kv_load_failure_policy=self.kv_load_failure_policy, + ) + self.vllm.start(self.repo_root, log_suffix=log_suffix) + return self.vllm + + def restart_vllm(self, log_suffix: str = "") -> VllmServer: + """Restart vLLM to drop its local prefix cache (KVCM state persists).""" + if self.vllm: + self.vllm.stop() + return self.start_vllm(log_suffix) + + def manager_block_size(self) -> int: + return get_manager_block_size(self.manager.manager_uri(), self.instance_id) + + def scan_connector_logs(self, pattern: str) -> list: + """Regex-scan all vLLM std streams; returns list of match groups.""" + import re + out = [] + for path in glob.glob(os.path.join(self.vllm_dir, "vllm*.std*")): + with open(path, errors="replace") as f: + for line in f: + m = re.search(pattern, line) + if m: + out.append(m.groups() if m.groups() else m.group(0)) + return out + + def stop(self): + if self.vllm: + self.vllm.stop() + self.manager.stop() + + +def run_e2e(scenario: str, tp_size: int, num_prompts: int, + preferred_block_size: int, connector_name: str = "VerifyingConnector", + expect_verification_failure: bool = False): + """Run one full save-then-load verification scenario. + + Full-attention models: prefix caching off, one server across both phases. + Hybrid models: prefix caching on (align mode -> per-group block tables), the + vLLM server is restarted between phases so phase 2 loads from KVCM instead of + hitting the local prefix cache. + + connector_name selects the connector class inside test_connector.py; the + mutation meta-test passes "MutatedConnector" and sets + expect_verification_failure=True to prove the harness detects an injected + off-by-one in the token translation. + """ + import torch # noqa: F401 (ensure torch importable early for clear errors) + + env = ScenarioEnv(scenario, tp_size=tp_size, + preferred_block_size=preferred_block_size, + connector_name=connector_name) + hybrid = env.hybrid + capture_dir = env.capture_dir + logger.info("scenario=%s model=%s hybrid=%s tp=%d prompts=%d preferred_bs=%d", + scenario, MODEL_PATH, hybrid, tp_size, num_prompts, + env.preferred_block_size) + + try: + env.start_manager() + vllm = env.start_vllm(log_suffix="" if not hybrid else "_p1") + + base_prompts = make_base_prompts(num_prompts, hybrid) + suffixes = [f" Now answer question {i}: what is 2+2?" for i in range(num_prompts)] + phase2_prompts = [p + s for p, s in zip(base_prompts, suffixes)] + + # Compute per-prompt expected block counts from the actual tokenization + # so a silently dropped prompt (or block) fails the run. + mbs = env.manager_block_size() + base_tokens = [tokenize(vllm.base_url(), p) for p in base_prompts] + phase2_tokens = [tokenize(vllm.base_url(), p) for p in phase2_prompts] + expected_save_blocks = [len(t) // mbs for t in base_tokens] + # Loads cover the shared token prefix (the base/suffix boundary token may + # re-merge under tokenization, shortening the shared prefix by one). + expected_load_blocks = [ + min(shared_token_prefix_len(b, p2) // mbs, s) + for b, p2, s in zip(base_tokens, phase2_tokens, expected_save_blocks) + ] + assert all(n >= 1 for n in expected_load_blocks), ( + f"prompts too short to span a manager block (mbs={mbs}): " + f"{expected_load_blocks}") + logger.info("mbs=%d expected save blocks=%s load blocks=%s", + mbs, expected_save_blocks, expected_load_blocks) + + # ---- Phase 1: fresh prefill -> connector saves -> reference capture. + logger.info("phase 1: sending %d fresh prompts", num_prompts) + send_completions(vllm.base_url(), base_prompts) + wait_for_captures(capture_dir, "ref", + expected=tp_size * sum(expected_save_blocks), timeout=180) + + # The save is committed to the manager asynchronously after the ref + # capture (which fires when the save is submitted). Wait until the manager + # actually has the prefix, otherwise phase 2 would find no match. + for toks, blocks in zip(base_tokens, expected_save_blocks): + if not wait_for_prefix_cached(env.manager.manager_uri(), env.instance_id, + toks, min_blocks=blocks): + raise AssertionError("save was not committed to the manager in time") + + # Hybrid models keep prefix caching on, which also populates the local + # prefix cache; restart vLLM so phase 2 loads from KVCM, not locally. + if hybrid: + logger.info("restarting vLLM before phase 2 (clear local prefix cache)") + vllm = env.restart_vllm(log_suffix="_p2") + + # ---- Phase 2: same prefix + suffix -> connector loads -> loaded capture. + logger.info("phase 2: sending %d prefix+suffix prompts", num_prompts) + send_completions(vllm.base_url(), phase2_prompts) + wait_for_captures(capture_dir, "loaded", + expected=tp_size * sum(expected_load_blocks), timeout=180) + + report = compare_captures(capture_dir, tp_size) + min_matched = tp_size * sum(expected_load_blocks) + if expect_verification_failure: + try: + assert_report_ok(report, min_matched=min_matched) + except AssertionError as e: + logger.info("verification failed as expected: %s", e) + return + raise AssertionError( + "mutated connector passed KV verification; the harness is blind") + assert_report_ok(report, min_matched=min_matched) + logger.info("scenario %s PASSED: %s", scenario, json.dumps( + {k: v for k, v in report.items() if k not in ("failures", "matched_keys")}, + default=str)) + finally: + env.stop() diff --git a/integration_test/vllm_e2e/test_basic.py b/integration_test/vllm_e2e/test_basic.py new file mode 100644 index 000000000..2c4d19e12 --- /dev/null +++ b/integration_test/vllm_e2e/test_basic.py @@ -0,0 +1,27 @@ +"""test_basic: single-request save/load KV verification (TP=1). + +Sends one prompt (prefill + save -> reference capture), then the same prompt +with a suffix (load + prefill -> loaded capture), and verifies the prefix KV +data matches bit-exactly. + +Works for both full-attention and hybrid models (selected via KVCM_E2E_MODEL); +see e2e_lib.run_e2e for the per-model orchestration differences. +""" + +import unittest + +from e2e_lib import run_e2e + + +class TestBasic(unittest.TestCase): + def test_basic(self): + run_e2e( + scenario="basic", + tp_size=1, + num_prompts=1, + preferred_block_size=0, # manager block size == vllm block size + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/integration_test/vllm_e2e/test_concurrent.py b/integration_test/vllm_e2e/test_concurrent.py new file mode 100644 index 000000000..4ea7501de --- /dev/null +++ b/integration_test/vllm_e2e/test_concurrent.py @@ -0,0 +1,28 @@ +"""test_concurrent: multiple concurrent requests save/load KV verification. + +Sends several distinct prompts concurrently (all prefill + save -> reference +captures), then the same prompts each with their own suffix concurrently (all +load + prefill -> loaded captures), and verifies each request's prefix KV data +matches. This exercises ReqState tracking, per-request block attribution and +async task races. + +Works for both full-attention and hybrid models (selected via KVCM_E2E_MODEL). +""" + +import unittest + +from e2e_lib import run_e2e + + +class TestConcurrent(unittest.TestCase): + def test_concurrent(self): + run_e2e( + scenario="concurrent", + tp_size=1, + num_prompts=4, + preferred_block_size=0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/integration_test/vllm_e2e/test_connector.py b/integration_test/vllm_e2e/test_connector.py new file mode 100644 index 000000000..7fbf1e327 --- /dev/null +++ b/integration_test/vllm_e2e/test_connector.py @@ -0,0 +1,370 @@ +"""A verification wrapper around the production KVCM vLLM connector. + +This connector is injected via vLLM's ``kv_connector_module_path`` and subclasses +the production ``TairKvCacheConnector`` without modifying it. Its purpose is to +independently capture the KV data that lives in vLLM's paged KV cache so that the +test driver can verify the connector's save/load translation layer. + +Why this catches translation bugs +--------------------------------- +The production connector is built around *per-group* transfer. Every +``kv_cache_group`` (a ``FullAttentionSpec`` group for pure-attention models, or +several ``MambaSpec`` groups plus one ``FullAttentionSpec`` group for hybrid +models) is a self-contained transfer unit with its own block table and its own +translation: + + KVCM manager block idx -> global token idx -> group logical block + (step 1, connector-only) (step 2/3, shared with vLLM) + +Step 1 is connector-only logic. A bug there makes *save* gather from the wrong +physical slots and *load* scatter to the wrong physical slots. Because save and +load share the same translation, a transport-level round trip still "matches" +(the bug is symmetric). + +To break the symmetry we capture KV data using ONLY the position -> physical +slot mapping that vLLM itself uses (its slot_mapping kernel), a pure function of +the group's block table and the token position. This reference is independent of +the connector's step-1 logic, so a step-1 bug makes the captured data diverge +from what the connector saved/loaded. + +Group kinds +----------- +* Attention groups -> ``torch.Tensor`` of shape + ``[2, num_blocks, kernel_block_size, num_kv_heads, head_size]``. Captured + per-token using the three-tier mapping (group logical block -> kernel physical + block) because the scheduler's group block size may exceed the kernel block + size. +* Mamba/linear/gdn groups -> ``list[Tensor]`` (e.g. ``[conv_state, ssm_state]``). + The state is stored **per group block**; we capture the whole state slice for + the group block that each manager block maps to (mirroring the connector's + ``_state_block_ids``: a manager block's *last* token selects the block). + +Capture points +-------------- +* Reference (save path): in ``wait_for_save`` we read the KV of the saved token + range straight out of the paged cache (the forward pass has completed and the + slots are not modified by the parent's async gather). +* Loaded (load path): loads are async, so the load step has no forward pass and + the worker does not yet know the request's token ids. We record the load's + per-group block tables in ``start_load_kv`` and emit the capture in a later + ``wait_for_save`` once the token ids have arrived. The loaded KV persists in + the paged cache (its blocks are allocated to the request). + +Captures are written to ``$KVCM_E2E_CAPTURE_DIR`` as ``.pt`` files named +``{ref|loaded}_tp{rank}_{token_hash}.pt`` so the out-of-process driver can match +reference vs loaded by content (the captured token ids). +""" + +import hashlib +import os +import threading +import typing + +import torch + +from kv_cache_manager.py_connector.common.logger import logger +from kv_cache_manager.py_connector.vllm.metadata import TairKvCacheConnectorMetadata +from kv_cache_manager.py_connector.vllm.v1_connector import ( + TairKvCacheConnector, attn_kv_views) +from kv_cache_manager.py_connector.vllm.connector_worker import ConnectorWorker +from kv_cache_manager.py_connector.vllm.vllm_common import AttentionGroupMeta + +CAPTURE_DIR_ENV = "KVCM_E2E_CAPTURE_DIR" + + +def _worker_attr(name): + """Property forwarding to the ConnectorWorker's attributes: the role + split moved them off the shell, and the capture hooks run on the + worker-role instance.""" + return property(lambda self: getattr(self.connector_worker, name)) + + +class VerifyingConnector(TairKvCacheConnector): + """Production connector + independent per-group KV capture for e2e.""" + + _tp_rank = _worker_attr("_tp_rank") + _device_mod = _worker_attr("_device_mod") + _device = _worker_attr("_device") + _kv_caches = _worker_attr("_kv_caches") + _data_transfer = _worker_attr("_data_transfer") + + # ------------------------------------------------------------------ # + # Setup + # ------------------------------------------------------------------ # + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): + super().register_kv_caches(kv_caches) + + self._capture_dir = os.environ.get(CAPTURE_DIR_ENV, "") + if self._capture_dir: + os.makedirs(self._capture_dir, exist_ok=True) + + # Snapshot the static per-group description into a capture-friendly form. + # Each entry: (group_idx, is_attention, layer_names, group_block_size, + # kernel_block_size). kernel_block_size is read straight off the tensor. + self._cap_groups = [] + for meta in self._group_metas: + if isinstance(meta, AttentionGroupMeta): + ref = kv_caches[meta.layer_names[0]] + kernel_bs = attn_kv_views(ref)[0][0].shape[1] + else: + kernel_bs = 0 + self._cap_groups.append( + (meta.group_idx, isinstance(meta, AttentionGroupMeta), + list(meta.layer_names), meta.block_size, kernel_bs)) + + # Track completion of async load scatters. The parent's load task already + # CPU-synchronizes its own scatter before reporting the task result, so a + # threading.Event set from the done callback is sufficient to know the + # scatter is globally visible. + self._load_done_events: dict[str, list[threading.Event]] = {} + self._load_events_lock = threading.Lock() + + # Loads are async and their step has no forward pass, so the worker does + # not yet have the request's token ids. Record the load's per-group block + # tables here and emit the capture once the token ids arrive. + # req_id -> (manager_block_idxes, block_ids_per_group) + self._pending_loaded: dict[str, tuple[list, list]] = {} + + orig_factory = self._data_transfer.create_load_done_callback + + def tracking_factory(req_id, *args, **kwargs): + orig_cb = orig_factory(req_id, *args, **kwargs) + evt = threading.Event() + with self._load_events_lock: + self._load_done_events.setdefault(req_id, []).append(evt) + + def cb(task_results): + try: + orig_cb(task_results) + finally: + evt.set() + + return cb + + self._data_transfer.create_load_done_callback = tracking_factory + logger.warning( + "VerifyingConnector enabled, capture_dir=%s tp_rank=%s groups=%s " + "vllm_bs=%s manager_bs=%s", + self._capture_dir, self._tp_rank, + [(g[0], "attn" if g[1] else "state", g[3], g[4]) for g in self._cap_groups], + self._vllm_block_size, self._manager_block_size, + ) + + # ------------------------------------------------------------------ # + # Load hook: record pending loaded captures + # ------------------------------------------------------------------ # + def start_load_kv(self, forward_context, **kwargs) -> None: + meta = typing.cast(TairKvCacheConnectorMetadata, self._get_connector_metadata()) + load_reqs = [ + (lr.req_id, list(lr.manager_block_idxes), + [list(b) for b in lr.all_block_ids]) + for lr in meta.to_load_requests + if lr.all_block_ids and lr.need_load_locations + ] + + super().start_load_kv(forward_context, **kwargs) + + if getattr(self, "_capture_dir", "") and load_reqs: + for req_id, mbis, bpg in load_reqs: + self._pending_loaded[req_id] = (mbis, bpg) + logger.warning( + "VerifyingConnector recorded %d pending loaded capture(s)", + len(load_reqs)) + + # ------------------------------------------------------------------ # + # Scheduler-side: ship token snapshots for the worker's captures + # ------------------------------------------------------------------ # + def build_connector_meta(self, scheduler_output): + meta = super().build_connector_meta(scheduler_output) + # The captures below identify a block by its token content; the + # worker no longer mirrors token streams (self-contained + # instructions), so hand it the live streams from the ledger. + scheduler = self.connector_scheduler + meta.token_snapshots = { + req_id: ledger.vllm_request.all_token_ids + for req_id, ledger in scheduler._tracked.items() + } + return meta + + # ------------------------------------------------------------------ # + # Save hook: reference captures + emit pending loaded captures + # ------------------------------------------------------------------ # + def wait_for_save(self): + meta = typing.cast(TairKvCacheConnectorMetadata, self._get_connector_metadata()) + + if getattr(self, "_capture_dir", "") and getattr(self, "_kv_caches", None): + try: + self._capture_refs(meta) + self._capture_pending_loaded(meta) + except Exception as e: # never break inference for a capture error + logger.warning("VerifyingConnector capture failed: %s", e, exc_info=True) + + super().wait_for_save() + + def _capture_refs(self, meta: TairKvCacheConnectorMetadata): + if not meta.to_save_requests: + return + # Make all forward-pass KV writes visible before reading the paged cache. + self._device_mod.synchronize() + tokens = getattr(meta, "token_snapshots", {}) + for save_req in meta.to_save_requests: + token_ids = tokens.get(save_req.req_id) + if token_ids is None or not save_req.all_block_ids: + continue + self._capture_range( + kind="ref", + token_ids=token_ids, + block_ids_per_group=save_req.all_block_ids, + manager_block_idxes=save_req.manager_block_idxes, + ) + + def _capture_pending_loaded(self, meta): + if not self._pending_loaded: + return + done = [] + tokens = getattr(meta, "token_snapshots", {}) + for req_id, (mbis, bpg) in self._pending_loaded.items(): + token_ids = tokens.get(req_id) + if token_ids is None: + # token ids have not arrived on this worker yet; wait for a + # later step in which the request is scheduled. + continue + with self._load_events_lock: + evts = list(self._load_done_events.get(req_id, [])) + for evt in evts: + evt.wait(timeout=120) + self._device_mod.synchronize() + self._capture_range( + kind="loaded", + token_ids=token_ids, + block_ids_per_group=bpg, + manager_block_idxes=mbis, + ) + done.append(req_id) + for req_id in done: + del self._pending_loaded[req_id] + + # ------------------------------------------------------------------ # + # Capture helpers + # ------------------------------------------------------------------ # + def _capture_range(self, kind, token_ids, block_ids_per_group, manager_block_idxes): + if not manager_block_idxes or not block_ids_per_group: + return + # One record per manager block. Saves are batched incrementally while + # loads arrive all-at-once, so per-block records let the driver match + # reference vs loaded captures by each block's token content. + for b in manager_block_idxes: + self._capture_block(kind, token_ids, block_ids_per_group, b) + + def _attn_token_slot(self, pos, block_table, group_bs, kernel_bs): + """Map a global token position to its flat slot in an attention group. + + Mirrors vLLM's own slot_mapping kernel expressed with the three-tier + block hierarchy (group logical block -> kernel physical block). Works for + pure-attention groups (group_bs == kernel_bs, ratio 1) and hybrid + attention groups (group block larger than kernel block). Independent of + the connector's step-1 (manager-block) logic, which is what we verify. + """ + ratio = group_bs // kernel_bs + logical = pos // group_bs + off = pos % group_bs + physical = block_table[logical] * ratio + off // kernel_bs + return physical * kernel_bs + off % kernel_bs + + def _capture_block(self, kind, token_ids, block_ids_per_group, manager_block_idx): + mbs = self._manager_block_size + + # Global token positions covered by this manager block. + positions = list(range(manager_block_idx * mbs, (manager_block_idx + 1) * mbs)) + if positions[-1] >= len(token_ids): + positions = [p for p in positions if p < len(token_ids)] + if not positions: + return + + captured_token_ids = [token_ids[p] for p in positions] + kv_by_layer = {} + + for group_idx, is_attention, layer_names, group_bs, kernel_bs in self._cap_groups: + block_table = block_ids_per_group[group_idx] + if is_attention: + slots = [self._attn_token_slot(p, block_table, group_bs, kernel_bs) + for p in positions] + slot_tensor = torch.tensor(slots, dtype=torch.long, device=self._device) + for layer_name in layer_names: + kv_cache = self._kv_caches[layer_name] + # Normalize the layout (packed 4-D or split K/V 5-D) into + # token-major views via the production helper and gather the + # whole per-token vector by (block, token) advanced indexing + # -- split K/V views are non-contiguous, so flattening them + # first would copy the entire cache tensor. Split views are + # concatenated on the content dim, so a capture is + # comparable across save/load within one run. + parts = [] + for v in attn_kv_views(kv_cache)[0]: + blk = slot_tensor // kernel_bs + tok = slot_tensor % kernel_bs + parts.append(v[blk, tok].reshape(len(slots), -1)) + gathered = (parts[0] if len(parts) == 1 + else torch.cat(parts, dim=-1)).contiguous() + kv_by_layer[layer_name] = gathered.cpu() + else: + # State stored once per group block; the manager block's last + # token selects the block (mirrors _state_block_ids). vLLM's + # mamba "align" mode materializes states only at segment + # boundaries -- interior blocks hold the null block (id 0) and + # carry no state to capture (the connector skips them too). + logical = ((manager_block_idx + 1) * mbs - 1) // group_bs + block_id = block_table[logical] + if block_id == 0: + continue + for layer_name in layer_names: + states = self._kv_caches[layer_name] # list[Tensor] + kv_by_layer[layer_name] = [s[block_id].detach().cpu() for s in states] + + token_hash = hashlib.sha256( + torch.tensor(captured_token_ids, dtype=torch.int64).numpy().tobytes() + ).hexdigest()[:16] + path = os.path.join(self._capture_dir, f"{kind}_tp{self._tp_rank}_{token_hash}.pt") + torch.save({"token_ids": captured_token_ids, "kv": kv_by_layer}, path) + logger.warning( + "VerifyingConnector captured %s block=%d tokens=%d..%d tp=%s -> %s", + kind, manager_block_idx, positions[0], positions[-1], self._tp_rank, path) + + +class MutatedWorkerCore(ConnectorWorker): + """Off-by-one in the attention token translation (see MutatedConnector).""" + + def _attn_token_indices(self, group, manager_block_idxes, block_table): + out = super()._attn_token_indices(group, manager_block_idxes, block_table) + return [[slot - 1 for slot in block] for block in out] + + +class MutatedConnector(VerifyingConnector): + """Meta-test connector: injects an off-by-one into the attention token + translation (every gathered/scattered slot shifted by -1). + + The shift is symmetric between save and load, so with contiguous block + tables a transport round trip cancels it in the interior of the loaded + range (slot(t)-1 == slot(t-1)); the leak is at the boundary: the last + loaded token's true slot is never written and keeps stale (uninitialized) + data. The capture-based verification reads the cache through vLLM's own + slot mapping and must observe that divergence -- the mutation e2e test + asserts that verification FAILS with this connector. + + -1 (not +1) keeps every shifted slot in bounds: vLLM reserves physical + block 0 as the null block, so real slots are >= kernel_block_size and + slot-1 >= 0, while slot+1 of the cache's last block would read/write out + of bounds. Only reachable through the test-side ``kv_connector_module_path`` + injection; never part of the production wheel. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # The translation lives on the ConnectorWorker since the role split; + # swap the worker-role instance (state included) for the mutated + # subclass. Scheduler-role instances have no worker to mutate. + if self.connector_worker is None: + return + mutated = MutatedWorkerCore.__new__(MutatedWorkerCore) + mutated.__dict__.update(self.connector_worker.__dict__) + self.connector_worker = mutated diff --git a/integration_test/vllm_e2e/test_cross_request_prefix.py b/integration_test/vllm_e2e/test_cross_request_prefix.py new file mode 100644 index 000000000..bf6b1d3d7 --- /dev/null +++ b/integration_test/vllm_e2e/test_cross_request_prefix.py @@ -0,0 +1,254 @@ +"""test_cross_request_prefix: a shorter request must not resume from a prefix +whose recurrent (mamba) boundary state was never written. + +Why this scenario exists +----------------------- +Every other e2e scenario re-sends the *same or a longer* prompt, so the block +that ends an external match is always a block whose mamba state vLLM actually +materialized. Hybrid ``mamba_cache_mode="align"`` materializes a state only at +segment boundaries: the interior manager blocks of a saved request map to +vLLM's null block and hold no state at all. + +This scenario is the missing cross-request case: + +* **Phase 1** -- request A (long, several manager blocks) is saved. Which of A's + blocks actually carry state is measured *independently* of the connector's + publication logic, from the harness's own reference captures (the verifying + connector captures state layers only for non-null state blocks). +* **Phase 2** -- request B (a token-level *prefix* of A ending inside one of A's + blocks) runs against a *fresh instance id* on a restarted server, so nothing + is cached anywhere: its greedy output is the ground truth. +* **Phase 3** -- a restarted server (empty local prefix cache) runs B against + A's cached blocks, so B's external match ends on an *interior* block of A. + +If interior blocks are published as fully cached even though their state was +never written, phase 3 loads a never-written state URI. The hybrid load failure +cannot be reported to vLLM (single-group invalid-block recovery upstream), so B +silently decodes from a garbage recurrent state and its output diverges from +phase 2's ground truth. That divergence is the primary assertion. + +Correct behaviour: the manager knows -- through per-key +``location_spec_group_names`` -- exactly which blocks carry state, and the +connector truncates the external match to the last state-complete block. B then +resumes from a real boundary state (or recomputes), producing the ground-truth +output. + +Full-attention models have no state groups: they must show *uniform* spec +coverage and *no* truncation, which makes this scenario a regression guard for +them too. +""" + +import glob +import logging +import os +import unittest + +import requests + +from e2e_lib import ( + ScenarioEnv, full_block_hashes, send_completions, tokenize, + wait_for_captures, wait_for_prefix_cached, +) + +logger = logging.getLogger("vllm_e2e") + + +def per_block_spec_names(manager_uri: str, instance_id: str, + token_ids: list[int]) -> list[list[str]]: + """Per-manager-block location spec names, in block order, as the manager + reports them to a prefix-match query -- exactly what the connector's load + path sees in ``get_num_new_matched_tokens``.""" + r = requests.post(f"{manager_uri}/api/getCacheLocation", json={ + "trace_id": "e2e_specs", + "token_ids": token_ids, + "instance_id": instance_id, + "query_type": "QT_PREFIX_MATCH", + "block_mask": {"offset": 0}, + }, timeout=10) + r.raise_for_status() + return [sorted(s["name"] for s in loc.get("location_specs", [])) + for loc in r.json().get("locations", [])] + + +def captured_state_presence(capture_dir: str, token_ids: list[int], + manager_block_size: int) -> list[bool]: + """Per-manager-block "vLLM materialized a recurrent state here", measured + from the harness's reference captures. + + ``test_connector.VerifyingConnector`` captures a state group's layers only + when the manager block maps to a real (non-null) state block, and stores + them as a ``list[Tensor]`` (attention layers are a single Tensor). So a + capture containing a list-valued layer is a block whose state exists. + + This is independent of the connector's manager-facing publication logic, + which is what the scenario verifies. + """ + import torch + + presence = [] + for idx, token_hash in enumerate(full_block_hashes(token_ids, manager_block_size)): + matches = glob.glob(os.path.join(capture_dir, f"ref_tp0_{token_hash}.pt")) + assert matches, f"no reference capture for manager block {idx}" + rec = torch.load(matches[0], map_location="cpu", weights_only=True) + presence.append(any(isinstance(v, (list, tuple)) + for v in rec["kv"].values())) + return presence + + +def long_prompt(num_sentences: int) -> str: + """Deterministic prompt whose every sentence is unique, so no two manager + blocks share token content (and therefore no two share a cache key).""" + return "Report A. " + " ".join( + f"Item {j} of the ledger records the value {j * 17 + 3} and the tag " + f"{(j * 31) % 97}." + for j in range(num_sentences)) + + +def completion_token_ids(base_url: str, token_ids: list[int], + max_tokens: int) -> list[int]: + """Greedy continuation of an explicit token-id prompt, as token ids.""" + resp = send_completions(base_url, [token_ids], max_tokens=max_tokens, + ignore_eos=True, return_token_ids=True)[0] + out = list(resp["choices"][0]["token_ids"]) + assert len(out) == max_tokens, f"short completion: {len(out)}" + return out + + +class TestCrossRequestPrefix(unittest.TestCase): + # Long enough that a corrupted recurrent state cannot plausibly reproduce + # the reference continuation token for token. + MAX_TOKENS = 32 + # Mirrors the harness's --max-model-len; prompt + output must fit. + MAX_MODEL_LEN = 4096 + + def test_cross_request_prefix(self): + env = ScenarioEnv("cross_request_prefix", log_level="DEBUG") + real_instance_id = env.instance_id + try: + env.start_manager() + + # ---- Phase 1: save request A, long enough to span several manager + # blocks so interior (state-less) blocks exist. + vllm = env.start_vllm(log_suffix="_p1") + mbs = env.manager_block_size() + # Hybrid models pin the manager block to the scheduler block (528), + # so they need a far longer prompt to span several blocks. + toks_a = tokenize(vllm.base_url(), long_prompt(260 if env.hybrid else 60)) + # Cap the length so prompt + MAX_TOKENS stays inside the model len. + blocks_a = min(len(toks_a) // mbs, + (self.MAX_MODEL_LEN - self.MAX_TOKENS) // mbs) + self.assertGreaterEqual( + blocks_a, 4, + f"only {blocks_a} manager blocks available (mbs={mbs}); this " + f"scenario needs interior blocks") + # Send A as explicit token ids so its cache keys are exactly toks_a. + toks_a = toks_a[:blocks_a * mbs] + logger.info("A: %d tokens = %d x %d", len(toks_a), blocks_a, mbs) + completion_token_ids(vllm.base_url(), toks_a, max_tokens=4) + wait_for_captures(env.capture_dir, "ref", expected=blocks_a, timeout=180) + self.assertTrue(wait_for_prefix_cached( + env.manager.manager_uri(), env.instance_id, toks_a, + min_blocks=blocks_a), "A was not committed to the manager") + + # ---- Ground truth: which of A's blocks really have a state? + has_state = captured_state_presence(env.capture_dir, toks_a, mbs) + logger.info("A state materialized per block (from captures): %s", + has_state) + # ---- And what does the manager advertise per block? + specs = per_block_spec_names(env.manager.manager_uri(), + env.instance_id, toks_a) + self.assertEqual(len(specs), blocks_a) + union = sorted(set().union(*[set(s) for s in specs])) + complete = [s == union for s in specs] + logger.info("A per-block coverage (union=%s): %s", union, + [("full" if c else "missing:%s" % sorted(set(union) - set(s))) + for c, s in zip(complete, specs)]) + + # ---- Choose B: a strict prefix of A ending mid-block whose last + # *candidate* block (where the external match would end) carries no + # state. That is the case the bug corrupts. The null-state layout + # depends on how vLLM's save batches lined up with segment + # boundaries, so derive it from the measured ground truth. + if env.hybrid: + choices = [n for n in range(2, blocks_a) if not has_state[n - 1]] + self.assertTrue( + choices, + f"no state-less interior block in A (state: {has_state}); " + f"the scenario cannot exercise truncation") + candidates = max(choices) + else: + candidates = blocks_a - 1 + toks_b = toks_a[:candidates * mbs + mbs // 2] + logger.info("B: %d tokens, %d candidate blocks (last has_state=%s)", + len(toks_b), candidates, has_state[candidates - 1]) + + # ---- Phase 2: ground truth for B -- fresh instance id (no KVCM + # entries) on a restarted server (no local prefix cache). + env.instance_id = real_instance_id + "-ref" + vllm = env.restart_vllm(log_suffix="_ref") + ref_out = completion_token_ids(vllm.base_url(), toks_b, self.MAX_TOKENS) + logger.info("B reference output: %s", ref_out) + + # ---- Phase 3: B against A's cached blocks. + env.instance_id = real_instance_id + vllm = env.restart_vllm(log_suffix="_p3") + got_out = completion_token_ids(vllm.base_url(), toks_b, self.MAX_TOKENS) + logger.info("B cached output: %s", got_out) + + matched = [int(g[0]) for g in + env.scan_connector_logs(r"matched (\d+) external tokens")] + logger.info("connector external match(es): %s", matched) + + # ---- Primary assertion: reusing A's prefix must not change B's + # output. A mismatch means B resumed from state it never had. + self.assertEqual( + got_out, ref_out, + "B's output changed when it reused A's cached prefix: the " + "external match ended on a block whose recurrent state was " + f"never written (external matches: {matched})") + + # ---- A load failure must never be silently swallowed. + self.assertFalse( + env.scan_connector_logs(r"load task failed"), + "a load failed in phase 3: a published block was not readable") + self.assertFalse( + env.scan_connector_logs(r"load failed for \d+/\d+ blocks"), + "a hybrid load failure was swallowed in phase 3") + + # ---- Structural evidence. + matched_blocks = max(matched) // mbs if matched else 0 + if not env.hybrid: + # No state groups: uniform coverage, nothing to truncate. + self.assertTrue(all(complete), + f"full-attention coverage is not uniform: {specs}") + self.assertEqual( + matched_blocks, candidates, + f"full-attention match was truncated: {matched_blocks} of " + f"{candidates} blocks") + return + + # Hybrid: what the manager advertises must match reality... + self.assertEqual( + complete, has_state, + f"advertised spec coverage {complete} does not match the " + f"blocks that really have a state {has_state}: state-less " + f"blocks are published as fully cached") + # ...and the match must stop at the last state-complete block. + expected = 0 + for i in range(candidates): + if has_state[i]: + expected = i + 1 + self.assertEqual( + matched_blocks, expected, + f"match not truncated to the last state-complete block: got " + f"{matched_blocks} blocks, expected {expected} (state of " + f"candidates: {has_state[:candidates]})") + self.assertTrue( + env.scan_connector_logs(r"truncated external match"), + "the truncation was not logged") + finally: + env.stop() + + +if __name__ == "__main__": + unittest.main() diff --git a/integration_test/vllm_e2e/test_full_hit.py b/integration_test/vllm_e2e/test_full_hit.py new file mode 100644 index 000000000..39c214ecb --- /dev/null +++ b/integration_test/vllm_e2e/test_full_hit.py @@ -0,0 +1,82 @@ +"""test_full_hit: full-prompt external hit must not crash the engine. + +Regression test for the synchronous-load full-hit bug: this connector reports +external matches with load_kv_async=False, so vLLM schedules +``num_tokens - num_computed_tokens`` new tokens and asserts that count is > 0 +(vllm/v1/core/sched/scheduler.py, waiting-queue loop: ``assert num_new_tokens +> 0``). Without capping, a prompt whose token count is an exact multiple of the +manager block size and whose blocks are all externally cached would make the +count 0 and kill the engine. + +Phase 1 saves a prompt of exactly N manager blocks; phase 2 resends the very +same prompt (as explicit token ids, so tokenization cannot shift the length). +Asserts: the engine survives, the completion is well-formed, and the connector +reports 0 < matched < prompt tokens (the cap dropped at least the last block). + +Runs against both full-attention and hybrid models via $KVCM_E2E_MODEL. +""" + +import logging +import unittest + +from e2e_lib import ( + ScenarioEnv, make_base_prompts, send_completions, tokenize, + wait_for_prefix_cached, +) + +logger = logging.getLogger("vllm_e2e") + + +class TestFullHit(unittest.TestCase): + def test_full_hit(self): + env = ScenarioEnv("full_hit") + try: + env.start_manager() + vllm = env.start_vllm(log_suffix="_p1" if env.hybrid else "") + + mbs = env.manager_block_size() + # Trim a long-enough prompt's token ids to an exact multiple of the + # manager block size (>= 2 blocks so the cap has room to drop one). + toks = tokenize(vllm.base_url(), make_base_prompts(1, env.hybrid)[0]) + num_blocks = len(toks) // mbs + self.assertGreaterEqual( + num_blocks, 2, f"prompt too short: {len(toks)} tokens, mbs={mbs}") + prompt_ids = toks[:num_blocks * mbs] + logger.info("full-hit prompt: %d tokens = %d x %d", + len(prompt_ids), num_blocks, mbs) + + # Phase 1: fresh prefill -> all blocks saved. + resp1 = send_completions(vllm.base_url(), [prompt_ids])[0] + self.assertTrue(resp1["choices"][0]["text"]) + self.assertTrue(wait_for_prefix_cached( + env.manager.manager_uri(), env.instance_id, prompt_ids, + min_blocks=num_blocks)) + + if env.hybrid: + # Clear the local prefix cache so phase 2 goes external. + vllm = env.restart_vllm(log_suffix="_p2") + + # Phase 2: the exact same prompt -> full external hit. Without the + # cap this crashes the engine (assert num_new_tokens > 0). + resp2 = send_completions(vllm.base_url(), [prompt_ids])[0] + self.assertTrue(resp2["choices"][0]["text"]) + + # The engine must still be alive and serving. + resp3 = send_completions(vllm.base_url(), ["sanity check prompt"])[0] + self.assertTrue(resp3["choices"][0]["text"]) + + # Connector-side evidence: matched > 0 (external hit happened) and + # matched < prompt tokens (the cap left tokens to recompute). + matched = [int(g[0]) for g in + env.scan_connector_logs(r"matched (\d+) external tokens")] + self.assertTrue(matched, "no 'matched N external tokens' log found") + hit = [m for m in matched if m > 0] + self.assertTrue(hit, f"no positive external match in {matched}") + self.assertTrue(all(m < len(prompt_ids) for m in hit), + f"match not capped below prompt len: {matched}") + finally: + env.stop() + + +if __name__ == "__main__": + unittest.main() diff --git a/integration_test/vllm_e2e/test_load_failure.py b/integration_test/vllm_e2e/test_load_failure.py new file mode 100644 index 000000000..694483bc3 --- /dev/null +++ b/integration_test/vllm_e2e/test_load_failure.py @@ -0,0 +1,180 @@ +"""test_load_failure: storage loss between save and load must degrade, not kill. + +Phase 1 saves a long prompt; the test then deletes the storage files of the +*tail half* of the manager blocks (resolved through the manager's ordered +getCacheLocation response, one file per block via key_count_per_file=1). +Phase 2 reloads the same prefix with ``block_per_load_task=1`` so every block +fails or succeeds independently. + +Full-attention models (single group, report_failures=True): the connector +reports the failed blocks' vLLM block ids; with +``kv_load_failure_policy="recompute"`` (vLLM 0.26.0 defaults to "fail", which +turns any load failure into a 500) vLLM truncates the computed-token count at +the first invalid block and recomputes from there +(vllm/v1/core/sched/scheduler.py::_handle_invalid_blocks / +_update_requests_with_invalid_blocks). Asserts: the request returns a normal +completion, the failures were logged and reported, every *surviving* head +block's loaded KV is bit-exact, and any mismatching capture belongs to a +deleted (recomputed) block. Recomputed blocks are not held to bit-exactness: +they contain freshly recomputed KV whose numerics depend on prefill chunking, +which is vLLM's business, not the connector's. + +Hybrid models (multiple groups, report_failures=False): vLLM's invalid-block +recovery only supports single-group block tables, so the connector only logs +the failure. Asserts: the request still returns (no hang, no crash) and the +failure was logged. KV content is NOT verified: with the failure swallowed +the affected blocks keep garbage by design. + +This scenario also regression-tests the fail-reschedule loop fix: a request +whose external load failed must not re-match external blocks on requeue +(v1_connector.get_num_new_matched_tokens retry guard), otherwise the engine +loops load-fail-reschedule forever and the request hangs. +""" + +import logging +import os +import unittest +from urllib.parse import urlparse + +import requests + +from e2e_lib import ( + ScenarioEnv, compare_captures, full_block_hashes, make_base_prompts, + send_completions, tokenize, wait_for_captures, wait_for_prefix_cached, +) + +logger = logging.getLogger("vllm_e2e") + + +def get_block_files(manager_uri: str, instance_id: str, token_ids: list[int], + spec_name: str | None = None) -> list[str]: + """Per-manager-block storage file paths, in block order, from the manager's + getCacheLocation response. + + Defaults to the spec every block is guaranteed to have: hybrid models + publish per-block spec coverage, so a *state* group's spec is absent from + the blocks where vLLM materialized no recurrent state. The attention spec + is the one present on every published block -- and it is also the one whose + loss makes a load fail, which is what this scenario needs. + """ + r = requests.post(f"{manager_uri}/api/getCacheLocation", json={ + "trace_id": "e2e_block_files", + "token_ids": token_ids, + "instance_id": instance_id, + "query_type": "QT_PREFIX_MATCH", + "block_mask": {"offset": 0}, + }, timeout=10) + r.raise_for_status() + locations = r.json().get("locations", []) + if spec_name is None: + # The spec present in *every* location is the attention one; state specs + # are sparse. Intersect to find it without knowing the group layout. + common = None + for location in locations: + names = {s["name"] for s in location.get("location_specs", [])} + common = names if common is None else (common & names) + assert common, f"no spec is common to all {len(locations)} locations" + spec_name = sorted(common)[0] + logger.info("using spec %s (present on all %d blocks)", spec_name, + len(locations)) + files = [] + for location in locations: + for spec in location.get("location_specs", []): + if spec["name"] == spec_name: + # uri: file://?size=... + files.append(urlparse(spec["uri"]).path) + return files + + +class TestLoadFailure(unittest.TestCase): + def test_load_failure(self): + env = ScenarioEnv( + "load_failure", + extra_config_overrides={"block_per_load_task": 1}, + key_count_per_file=1, # one file per block -> per-block failures + kv_load_failure_policy="recompute", + ) + try: + env.start_manager() + vllm = env.start_vllm(log_suffix="_p1" if env.hybrid else "") + mbs = env.manager_block_size() + + prompt = make_base_prompts(1, env.hybrid)[0] + suffix = " Now answer: what is 2+2?" + toks = tokenize(vllm.base_url(), prompt) + save_blocks = len(toks) // mbs + self.assertGreaterEqual(save_blocks, 2) + + # ---- Phase 1: save everything. + send_completions(vllm.base_url(), [prompt]) + wait_for_captures(env.capture_dir, "ref", expected=save_blocks, + timeout=180) + self.assertTrue(wait_for_prefix_cached( + env.manager.manager_uri(), env.instance_id, toks, + min_blocks=save_blocks)) + + # ---- Sabotage: delete the tail half of the blocks' files. The + # head blocks stay loadable, so vLLM truncates at the first deleted + # block and the surviving loads remain verifiable. + files = get_block_files(env.manager.manager_uri(), env.instance_id, + toks) + self.assertEqual(len(files), save_blocks) + keep = save_blocks // 2 + for path in files[keep:]: + os.remove(path) + logger.info("deleted %d/%d block files (kept blocks 0..%d)", + save_blocks - keep, save_blocks, keep - 1) + + if env.hybrid: + vllm = env.restart_vllm(log_suffix="_p2") + + # ---- Phase 2: load with holes. The request must return normally. + resp = send_completions(vllm.base_url(), [prompt + suffix])[0] + self.assertTrue(resp["choices"][0]["text"]) + + # The engine must survive and keep serving. + resp2 = send_completions(vllm.base_url(), ["engine alive?"])[0] + self.assertTrue(resp2["choices"][0]["text"]) + + failed_tasks = env.scan_connector_logs(r"load task failed") + self.assertTrue(failed_tasks, "no load failure was logged; the " + "sabotage did not break any loaded block") + + if env.hybrid: + # report_failures=False path: swallowed but logged. + swallowed = env.scan_connector_logs( + r"load failed for \d+/\d+ blocks .*hybrid") + self.assertTrue(swallowed, + "hybrid load failure was not logged") + return + + # Full-attention: vLLM was told about the invalid blocks... + reported = env.scan_connector_logs(r"block_ids_with_load_errors") + self.assertTrue(reported, "failed loads were not reported to vLLM") + + # ...and every surviving head block's loaded KV is bit-exact, + # while any mismatch belongs to a deleted (recomputed) block. + wait_for_captures(env.capture_dir, "loaded", expected=keep, + timeout=180) + report = compare_captures(env.capture_dir, tp_size=1) + hashes = full_block_hashes(toks, mbs) + kept_keys = {("tp0", h) for h in hashes[:keep]} + deleted_keys = {("tp0", h) for h in hashes[keep:]} + failed_keys = {f["key"] for f in report["failures"]} + self.assertFalse( + failed_keys & kept_keys, + f"surviving loaded blocks mismatched: {failed_keys & kept_keys}") + self.assertTrue( + failed_keys <= deleted_keys, + f"mismatches outside the deleted blocks: " + f"{failed_keys - deleted_keys}") + matched_kept = kept_keys & set(report["matched_keys"]) + self.assertEqual( + len(matched_kept), keep, + f"only {len(matched_kept)}/{keep} surviving blocks verified") + finally: + env.stop() + + +if __name__ == "__main__": + unittest.main() diff --git a/integration_test/vllm_e2e/test_multi_turn.py b/integration_test/vllm_e2e/test_multi_turn.py new file mode 100644 index 000000000..5912b0895 --- /dev/null +++ b/integration_test/vllm_e2e/test_multi_turn.py @@ -0,0 +1,112 @@ +"""test_multi_turn: decode-time incremental save feeds the next turn's hit. + +Turn 1 sends a prompt and generates enough output tokens (max_tokens crossing +at least one manager block boundary) that the save threshold in +``build_connector_meta`` fires again during decode: blocks composed of +generated tokens are saved incrementally. Turn 2 sends prompt + turn-1 output +as its prompt (a real multi-turn conversation) and must externally match +*more* blocks than the turn-1 prompt alone covers -- proving decode-produced +blocks were saved -- and their loaded KV must verify against the references +captured during decode. + +Full-attention (mbs=16): three decode blocks, same server both turns (prefix +caching off, the external hit is directly observable). +Hybrid (mbs=528): one decode block (528+ generated tokens); the server is +restarted before turn 2 because prefix caching must stay on for hybrid models +and would otherwise mask the external hit with a local one. +""" + +import logging +import unittest + +from e2e_lib import ( + ScenarioEnv, assert_report_ok, compare_captures, send_completions, + shared_token_prefix_len, tokenize, wait_for_captures, + wait_for_prefix_cached, +) + +logger = logging.getLogger("vllm_e2e") + + +class TestMultiTurn(unittest.TestCase): + def test_multi_turn(self): + env = ScenarioEnv("multi_turn") + try: + env.start_manager() + vllm = env.start_vllm(log_suffix="_t1" if env.hybrid else "") + mbs = env.manager_block_size() + + turn1_prompt = ("A short story request. Please write a long, " + "detailed story about a robot that learns to paint.") + prompt_tokens = tokenize(vllm.base_url(), turn1_prompt) + prompt_blocks = len(prompt_tokens) // mbs + + # ---- Turn 1: generate output crossing >= 1 manager block + # boundary (3 blocks for full-attn's mbs=16; 1 block for hybrid's + # mbs=528 to keep decode time bounded). + max_tokens = (mbs + 32) if env.hybrid else (mbs * 3 + 5) + resp = send_completions( + vllm.base_url(), [turn1_prompt], max_tokens=max_tokens, + ignore_eos=True, return_token_ids=True)[0] + choice = resp["choices"][0] + output_ids = choice["token_ids"] + self.assertEqual(len(output_ids), max_tokens) + turn1_ids = choice["prompt_token_ids"] + output_ids + # The connector tracks tokens when they are *scheduled as input*; + # the very last sampled token never re-enters a step, so at most + # (len - 1) // mbs blocks can have been committed. + turn1_blocks = (len(turn1_ids) - 1) // mbs + self.assertGreater( + turn1_blocks, prompt_blocks, + "turn 1 output did not cross a manager block boundary") + logger.info("turn1: %d prompt + %d output tokens = %d blocks " + "(prompt alone: %d)", len(choice["prompt_token_ids"]), + len(output_ids), turn1_blocks, prompt_blocks) + + # Decode-produced blocks must be committed: the manager holds the + # full prompt+output prefix, more blocks than the prompt covers. + self.assertTrue(wait_for_prefix_cached( + env.manager.manager_uri(), env.instance_id, turn1_ids, + min_blocks=turn1_blocks)) + wait_for_captures(env.capture_dir, "ref", expected=turn1_blocks, + timeout=180) + + if env.hybrid: + # Prefix caching is on for hybrid; restart so turn 2's hit + # comes from KVCM, not the local prefix cache. + vllm = env.restart_vllm(log_suffix="_t2") + + # ---- Turn 2: conversation continues; prompt embeds turn 1's + # prompt + output as token ids (immune to detokenization drift). + turn2_suffix = tokenize(vllm.base_url(), + " Now summarize the story in one word.") + turn2_ids = turn1_ids + turn2_suffix + shared_blocks = min( + shared_token_prefix_len(turn1_ids, turn2_ids) // mbs, + turn1_blocks) + self.assertGreater(shared_blocks, prompt_blocks, + "turn 2 shares no decode-produced block") + resp2 = send_completions(vllm.base_url(), [turn2_ids])[0] + self.assertTrue(resp2["choices"][0]["text"]) + + # The external hit must cover decode-produced blocks. + matched = [int(g[0]) for g in + env.scan_connector_logs(r"matched (\d+) external tokens")] + best = max(matched, default=0) + self.assertGreater( + best, prompt_blocks * mbs, + f"external hit ({best} tokens) does not exceed the prompt-only " + f"coverage ({prompt_blocks * mbs} tokens): decode-time saves " + f"were not used") + + # And the loaded decode-block KV must verify. + wait_for_captures(env.capture_dir, "loaded", + expected=shared_blocks, timeout=180) + report = compare_captures(env.capture_dir, tp_size=1) + assert_report_ok(report, min_matched=shared_blocks) + finally: + env.stop() + + +if __name__ == "__main__": + unittest.main() diff --git a/integration_test/vllm_e2e/test_mutation.py b/integration_test/vllm_e2e/test_mutation.py new file mode 100644 index 000000000..12c4e1e18 --- /dev/null +++ b/integration_test/vllm_e2e/test_mutation.py @@ -0,0 +1,35 @@ +"""test_mutation: meta-test proving the e2e KV verification is not vacuous. + +Runs the standard basic scenario with ``MutatedConnector`` (defined in +test_connector.py, injected only through the test-side +``kv_connector_module_path``), which shifts every attention slot produced by +``_attn_token_indices`` by one -- a symmetric off-by-one: save gathers token +t's KV from the shifted slot and load scatters it back there, so with +contiguous block tables a transport round trip cancels the bug in the interior +of the loaded range. It cannot cancel at the range boundary: one loaded +token's true slot is never written and keeps stale uninitialized data. The +capture comparison reads the cache through vLLM's own slot mapping and must +observe the divergence; run_e2e(expect_verification_failure=True) asserts the +verification FAILS. If the mutated run verifies clean, the harness is blind +and this test fails. +""" + +import unittest + +from e2e_lib import run_e2e + + +class TestMutation(unittest.TestCase): + def test_mutated_connector_is_caught(self): + run_e2e( + scenario="mutation", + tp_size=1, + num_prompts=1, + preferred_block_size=0, + connector_name="MutatedConnector", + expect_verification_failure=True, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/integration_test/vllm_e2e/test_partial_hit.py b/integration_test/vllm_e2e/test_partial_hit.py new file mode 100644 index 000000000..399bd5eb9 --- /dev/null +++ b/integration_test/vllm_e2e/test_partial_hit.py @@ -0,0 +1,96 @@ +"""test_partial_hit: incremental query + incremental save on a non-zero prefix. + +The standard scenarios always start requests from zero computed blocks, so the +``computed_blocks > 0`` branch of ``get_num_new_matched_tokens`` (manager query +with a non-zero ``block_mask.offset``) and the ``block_mask.offset`` branch of +the save path (manager skipping already-stored blocks) are never exercised. +This scenario forces both, with prefix caching enabled for all model types: + +* Stage 1: save prefix A (fresh prefill). +* Stage 2 (same server): send A+B. vLLM locally hits A, so the connector + queries the manager with offset = locally computed blocks (> 0), then + extends the save; the manager's start_write_cache response skips the + already-stored A blocks via a non-zero block_mask offset. +* Stage 3 (restarted server, local cache cleared): send A+B+C. The connector + externally matches A+B -- proving stage 2's incremental save committed -- + loads it, and the loaded KV is verified against the reference captures. + +Log evidence asserted: a getCacheLocation request carrying a non-zero offset +(requires connector DEBUG logging, enabled via log_level). +""" + +import logging +import unittest + +from e2e_lib import ( + ScenarioEnv, assert_report_ok, compare_captures, make_base_prompts, + send_completions, shared_token_prefix_len, tokenize, wait_for_captures, + wait_for_prefix_cached, +) + +logger = logging.getLogger("vllm_e2e") + + +class TestPartialHit(unittest.TestCase): + def test_partial_hit(self): + env = ScenarioEnv("partial_hit", enable_prefix_caching=True, + log_level="DEBUG") + try: + env.start_manager() + vllm = env.start_vllm(log_suffix="_s12") + mbs = env.manager_block_size() + + base = make_base_prompts(1, env.hybrid)[0] + # Three nested prompts: A < A+B < A+B+C. + prompt_a = base + prompt_ab = base + " Continuation section B. " + " ".join( + f"Extra sentence {j} carries value {j * 13 + 7}." + for j in range(90 if env.hybrid else 30)) + prompt_abc = prompt_ab + " Final question: what is 2+2?" + + toks_a = tokenize(vllm.base_url(), prompt_a) + toks_ab = tokenize(vllm.base_url(), prompt_ab) + toks_abc = tokenize(vllm.base_url(), prompt_abc) + blocks_a = len(toks_a) // mbs + blocks_ab = len(toks_ab) // mbs + shared_abc = shared_token_prefix_len(toks_ab, toks_abc) // mbs + self.assertGreaterEqual(blocks_a, 1) + self.assertGreater(blocks_ab, blocks_a, + "B must add at least one manager block") + logger.info("blocks: A=%d AB=%d shared(AB,ABC)=%d", + blocks_a, blocks_ab, shared_abc) + + # ---- Stage 1: save A. + send_completions(vllm.base_url(), [prompt_a]) + wait_for_captures(env.capture_dir, "ref", expected=blocks_a, timeout=180) + self.assertTrue(wait_for_prefix_cached( + env.manager.manager_uri(), env.instance_id, toks_a, + min_blocks=blocks_a)) + + # ---- Stage 2: A hits the local prefix cache -> incremental + # external query (non-zero offset) + incremental save of B. + send_completions(vllm.base_url(), [prompt_ab]) + self.assertTrue(wait_for_prefix_cached( + env.manager.manager_uri(), env.instance_id, toks_ab, + min_blocks=blocks_ab)) + wait_for_captures(env.capture_dir, "ref", expected=blocks_ab, timeout=180) + + offsets = [int(g[0]) for g in env.scan_connector_logs( + r"get_kvcache_location request:.*'offset': (\d+)")] + self.assertTrue(any(o > 0 for o in offsets), + f"no incremental query with non-zero offset: {offsets}") + + # ---- Stage 3: restart (clear local cache) and load A+B. + vllm = env.restart_vllm(log_suffix="_s3") + send_completions(vllm.base_url(), [prompt_abc]) + wait_for_captures(env.capture_dir, "loaded", + expected=shared_abc, timeout=180) + + report = compare_captures(env.capture_dir, tp_size=1) + assert_report_ok(report, min_matched=shared_abc) + finally: + env.stop() + + +if __name__ == "__main__": + unittest.main() diff --git a/integration_test/vllm_e2e/test_tp.py b/integration_test/vllm_e2e/test_tp.py new file mode 100644 index 000000000..a44c572f4 --- /dev/null +++ b/integration_test/vllm_e2e/test_tp.py @@ -0,0 +1,33 @@ +"""test_tp: TP=2 save/load KV verification with a non-trivial block translation. + +Runs the save/load verification under tensor parallelism (TP=2), where each rank +has an independent forward context, slot mapping and capture, and the connector's +ZMQ-based TP coordination is fully exercised. + +For full-attention models it also sets ``preferred_block_size=32`` while vLLM +uses its default block size (16), forcing the connector's manager-block <-> +group-block translation (``_attn_token_indices``) to do real cross-block +mapping -- the code path most prone to symmetric save/load bugs. For hybrid +models the manager block size is pinned to the scheduler block size (mamba state +is per scheduler block), so run_e2e ignores preferred_block_size there. + +Works for both full-attention and hybrid models (selected via KVCM_E2E_MODEL). +""" + +import unittest + +from e2e_lib import run_e2e + + +class TestTp(unittest.TestCase): + def test_tp(self): + run_e2e( + scenario="tp", + tp_size=2, + num_prompts=2, + preferred_block_size=32, # != vllm block size (16) for full-attn models + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/kv_cache_manager/manager/cache_manager.cc b/kv_cache_manager/manager/cache_manager.cc index 166487f34..cf9a7fa9d 100644 --- a/kv_cache_manager/manager/cache_manager.cc +++ b/kv_cache_manager/manager/cache_manager.cc @@ -1097,8 +1097,19 @@ CacheManager::StartWriteCache(RequestContext *request_context, const std::string &trace_id = request_context->trace_id(); auto *service_metrics_collector = dynamic_cast(request_context->metrics_collector()); if (!location_spec_group_names.empty()) { + // The group names are per *block*. A token-only request carries no + // block keys (they are generated from the tokens below), so the block + // count must come from the tokens in that case -- otherwise a valid + // token-only request with group names is always rejected against a + // key count of 0. + size_t block_count = keys.size(); + if (block_count == 0 && !tokens.empty()) { + auto [ec_bs, block_size] = GetBlockSize(request_context, instance_id); + RETURN_IF_EC_NOT_OK_WITH_TYPE_LOG(WARN, ec_bs, StartWriteCacheInfo, "start write cache failed"); + block_count = tokens.size() / block_size; + } auto check_ec = - CheckLocationSpecGroupNames(request_context, instance_id, keys.size(), location_spec_group_names); + CheckLocationSpecGroupNames(request_context, instance_id, block_count, location_spec_group_names); RETURN_IF_EC_NOT_OK_WITH_TYPE(check_ec, StartWriteCacheInfo); } auto [ec, meta_searcher] = CheckInputAndGetMetaSearcher(request_context, instance_id, keys, tokens); diff --git a/kv_cache_manager/manager/test/cache_manager_test.cc b/kv_cache_manager/manager/test/cache_manager_test.cc index 7f6e67501..88cf42e12 100644 --- a/kv_cache_manager/manager/test/cache_manager_test.cc +++ b/kv_cache_manager/manager/test/cache_manager_test.cc @@ -1079,6 +1079,54 @@ TEST_F(CacheManagerTest, TestStartWriteDuplicateCache) { } } +// A token-only StartWriteCache (block_keys empty, keys derived from token_ids) +// must accept per-block location_spec_group_names: the size check has to count +// blocks, not the empty block_keys vector. Regression for hybrid connectors, +// which announce per-block spec coverage on token-only writes. +TEST_F(CacheManagerTest, TestStartWriteCacheSpecGroupNamesWithTokenIdsOnly) { + constexpr int64_t kBlockSize = 4; + std::vector location_spec_infos = { + LocationSpecInfo("tp0_attn", 512), + LocationSpecInfo("tp0_state", 512), + }; + std::vector location_spec_groups = { + LocationSpecGroup("attn", {"tp0_attn"}), + LocationSpecGroup("full", {"tp0_attn", "tp0_state"}), + }; + auto expected = std::pair(EC_OK, default_storage_configs); + ASSERT_EQ(expected, + cache_manager_->RegisterInstance(request_context_.get(), + "default", + "token_only_sg", + kBlockSize, + location_spec_infos, + createModelDeployment(), + location_spec_groups)); + + // 3 blocks of tokens, no block_keys: the middle block carries no state. + CacheManager::TokenIdsVector tokens; + for (int64_t i = 0; i < 3 * kBlockSize; ++i) { + tokens.push_back(i); + } + const std::vector group_names{"full", "attn", "full"}; + auto [ec, info] = cache_manager_->StartWriteCache( + request_context_.get(), "token_only_sg", {}, tokens, group_names, 100000000); + ASSERT_EQ(EC_OK, ec); + const auto &locs = info.locations().cache_locations_view(); + ASSERT_EQ(3u, locs.size()); + // Each block is allocated exactly the specs of its announced group, so the + // state-less block is never published as holding a state. + ASSERT_EQ(2, locs[0].spec_size()); + ASSERT_EQ(1, locs[1].spec_size()); + ASSERT_EQ(2, locs[2].spec_size()); + ASSERT_EQ("tp0_attn", locs[1].location_specs()[0].name()); + + // A mismatched length must still be rejected. + auto [ec_bad, info_bad] = cache_manager_->StartWriteCache( + request_context_.get(), "token_only_sg", {}, tokens, {"full"}, 100000000); + ASSERT_NE(EC_OK, ec_bad); +} + TEST_F(CacheManagerTest, TestStartWriteCacheRecordWriteBytes) { auto expected = std::pair(EC_OK, default_storage_configs); ASSERT_EQ(expected, diff --git a/kv_cache_manager/py_connector/common/types.py b/kv_cache_manager/py_connector/common/types.py deleted file mode 100644 index bc341e7e7..000000000 --- a/kv_cache_manager/py_connector/common/types.py +++ /dev/null @@ -1,22 +0,0 @@ -from enum import Enum -from typing import Tuple, Dict, Optional, Any - -import attrs -import torch - - -@attrs.define(frozen=True) -class KVCacheInfo: - tp_rank: int - world_size: int - kvcaches: Dict[str, torch.Tensor] - kvcache_ptr_tensor_cpu: torch.Tensor - kvcache_ptr_tensor_gpu: torch.Tensor - all_kvcache_ptr_tensor_gpu: torch.Tensor - layer_num: int - local_token_num: int - per_manager_block_shape: Tuple[int, ...] - per_manager_block_byte_size: int - per_token_per_layer_dim_size: int - device: torch.device - dtype: torch.dtype diff --git a/kv_cache_manager/py_connector/kernel/batch_gather_scatter_helper.py b/kv_cache_manager/py_connector/kernel/batch_gather_scatter_helper.py index ebb297505..1b83a83bf 100644 --- a/kv_cache_manager/py_connector/kernel/batch_gather_scatter_helper.py +++ b/kv_cache_manager/py_connector/kernel/batch_gather_scatter_helper.py @@ -1,3 +1,35 @@ +"""Batch gather/scatter between paged KV caches and flat host buffers. + +Consumed by the vLLM connector only (the sglang connector moves its data +through sglang's own hicache path and never touches this module). + +Two addressing modes per pointer: + +* flat (``block_stride == 0``): kernel pages are contiguous, so a token's + element offset is simply ``flat_token_idx * NUM_DIMS_PER_TOKEN``. This is + the whole story for the packed 4-D layout (vLLM >= 0.26.0) and the + KV-first split layout (vLLM <= 0.22.1), whose halves are flat. +* strided (``block_stride != 0``): consecutive kernel pages of one pointer + are ``block_stride`` elements apart. Two layouts need it: the N-first + split layout (vLLM 0.23.0-0.25.x) interleaves K and V per block + (``t[:, 0]`` / ``t[:, 1]`` halves), and page_size_padded allocations + leave a gap the kernel must skip. The offset decomposes the flat token + index into (kernel page, position in page): + + kv_block_idx = flat_token_idx // local_block_size + token_in_kv_block = flat_token_idx % local_block_size + offset = kv_block_idx * block_stride + + token_in_kv_block * NUM_DIMS_PER_TOKEN + + (local_block_size is the tensor's page size, which may differ from the + manager block size on padded allocations.) + +Every pointer in the array is its own base (K and V halves included), so +there is no K->V stride to add. The layout each pointer came from travels +on AttentionTransferGroup.kv_layout; the GPU test in test/kernel checks +both modes element-wise against naive torch indexing. +""" + from typing import List, Optional, Union import torch @@ -42,8 +74,15 @@ def kv_cache_batch_gather_kernel( NUM_KVCACHE_PTRS: tl.constexpr, # num_layers * kv_count BLOCK_SIZE: tl.constexpr, # 隐藏维度分块大小 DTYPE: tl.constexpr = tl.float16, + kv_stride: tl.constexpr = 0, # stride between K and V (for V pointers) + block_stride: tl.constexpr = 0, # stride between blocks (0 = use flat indexing) + local_block_size: tl.constexpr = 0, # actual block size in tensor (0 = use NUM_TOKENS_PER_BLOCK) ): NUM_DIMS_PER_BLOCK = NUM_TOKENS_PER_BLOCK * NUM_DIMS_PER_TOKEN + + # Determine if using strided layout + USE_STRIDED: tl.constexpr = (block_stride != 0) + EFFECTIVE_LOCAL_BLOCK_SIZE: tl.constexpr = local_block_size if local_block_size > 0 else NUM_TOKENS_PER_BLOCK pid = tl.program_id(0) grid_size = tl.num_programs(0) # 实际grid大小 (如3) @@ -64,6 +103,8 @@ def kv_cache_batch_gather_kernel( # 3. 遍历所有KV缓存指针 (k/v for each layer) for ptr_idx in tl.range(NUM_KVCACHE_PTRS): # 3.1 加载当前层的KV缓存基地址 + # Note: For non-MLA, pointer array is [K0, V0, K1, V1, ...] + # V pointer is already V's base (tensor[1].data_ptr()), no need to add kv_stride kvcache_ptr = tl.load(kv_cache_ptrs_ptr + ptr_idx).to(tl.pointer_type(DTYPE)) # 3.2 计算当前层在dst中的基础偏移 @@ -90,7 +131,16 @@ def kv_cache_batch_gather_kernel( # 从HBM的KV缓存加载数据 # 计算源指针: [BLOCK_SIZE] - src_ptrs = kvcache_ptr + global_token_idx * NUM_DIMS_PER_TOKEN + dim_idx_in_token + if USE_STRIDED: + # Strided layout: convert flat token index to strided offset + # V pointer already includes kv_stride offset, so no need to add it again + kv_block_idx = global_token_idx // EFFECTIVE_LOCAL_BLOCK_SIZE + token_in_kv_block = global_token_idx % EFFECTIVE_LOCAL_BLOCK_SIZE + strided_offset = kv_block_idx * block_stride + token_in_kv_block * NUM_DIMS_PER_TOKEN + src_ptrs = kvcache_ptr + strided_offset + dim_idx_in_token + else: + # Contiguous layout: flat indexing + src_ptrs = kvcache_ptr + global_token_idx * NUM_DIMS_PER_TOKEN + dim_idx_in_token load_mask = mask & token_gather_mask data = tl.load(src_ptrs, mask=load_mask, other=0.0) # 大块连续写入 host memory (PCIe优化) @@ -107,7 +157,10 @@ def batch_gather_kv_caches( dst_block_indices: List[int], # List of dst block indices num_tokens_per_block: int, dim_size_per_token_per_layer: int, - sm_count: int = 3 + sm_count: int = 3, + kv_stride: int = 0, # stride between K and V (for V pointers) + block_stride: int = 0, # stride between blocks (0 = use flat indexing) + local_block_size: int = 0, # actual block size in tensor (0 = use num_tokens_per_block) ): # 配置参数 total_blocks = len(dst_block_indices) @@ -132,6 +185,9 @@ def batch_gather_kv_caches( BLOCK_SIZE=2048, DTYPE=pytorch_dtype_to_triton_dtype(dst_tensor.dtype), num_warps=32, + kv_stride=kv_stride, + block_stride=block_stride, + local_block_size=local_block_size if local_block_size > 0 else num_tokens_per_block, ) # TODO autotune num_warps and BLOCK_SIZE @@ -148,8 +204,15 @@ def kv_cache_batch_scatter_kernel( NUM_KVCACHE_PTRS: tl.constexpr, # num_layers * kv_count BLOCK_SIZE: tl.constexpr, # 隐藏维度分块大小 DTYPE: tl.constexpr = tl.float16, + kv_stride: tl.constexpr = 0, # stride between K and V (for V pointers) + block_stride: tl.constexpr = 0, # stride between blocks (0 = use flat indexing) + local_block_size: tl.constexpr = 0, # actual block size in tensor (0 = use NUM_TOKENS_PER_BLOCK) ): NUM_DIMS_PER_BLOCK = NUM_TOKENS_PER_BLOCK * NUM_DIMS_PER_TOKEN + + # Determine if using strided layout + USE_STRIDED: tl.constexpr = (block_stride != 0) + EFFECTIVE_LOCAL_BLOCK_SIZE: tl.constexpr = local_block_size if local_block_size > 0 else NUM_TOKENS_PER_BLOCK pid = tl.program_id(0) grid_size = tl.num_programs(0) # 实际grid大小 (如3) @@ -170,6 +233,8 @@ def kv_cache_batch_scatter_kernel( # 3. 遍历所有KV缓存指针 (k/v for each layer) for ptr_idx in range(NUM_KVCACHE_PTRS): # 3.1 加载当前层的KV缓存基地址 + # Note: For non-MLA, pointer array is [K0, V0, K1, V1, ...] + # V pointer is already V's base (tensor[1].data_ptr()), no need to add kv_stride kvcache_ptr = tl.load(kv_cache_ptrs_ptr + ptr_idx).to(tl.pointer_type(DTYPE)) # 3.2 计算当前层在src中的基础偏移 @@ -201,7 +266,16 @@ def kv_cache_batch_scatter_kernel( # 向HBM的KV缓存写入数据 # 计算目的指针: [BLOCK_SIZE] - dst_ptrs = kvcache_ptr + global_token_idx * NUM_DIMS_PER_TOKEN + dim_idx_in_token + if USE_STRIDED: + # Strided layout: convert flat token index to strided offset + # V pointer already includes kv_stride offset, so no need to add it again + kv_block_idx = global_token_idx // EFFECTIVE_LOCAL_BLOCK_SIZE + token_in_kv_block = global_token_idx % EFFECTIVE_LOCAL_BLOCK_SIZE + strided_offset = kv_block_idx * block_stride + token_in_kv_block * NUM_DIMS_PER_TOKEN + dst_ptrs = kvcache_ptr + strided_offset + dim_idx_in_token + else: + # Contiguous layout: flat indexing + dst_ptrs = kvcache_ptr + global_token_idx * NUM_DIMS_PER_TOKEN + dim_idx_in_token tl.store(dst_ptrs, data, mask=load_mask) @@ -214,7 +288,10 @@ def batch_scatter_kv_caches( src_block_indices: List[int], # List of src block indices num_tokens_per_block: int, dim_size_per_token_per_layer: int, - sm_count: int = 3 + sm_count: int = 3, + kv_stride: int = 0, # stride between K and V (for V pointers) + block_stride: int = 0, # stride between blocks (0 = use flat indexing) + local_block_size: int = 0, # actual block size in tensor (0 = use num_tokens_per_block) ): # 配置参数 total_blocks = len(src_block_indices) @@ -245,4 +322,7 @@ def batch_scatter_kv_caches( BLOCK_SIZE=2048, DTYPE=pytorch_dtype_to_triton_dtype(src_tensor.dtype), num_warps=32, + kv_stride=kv_stride, + block_stride=block_stride, + local_block_size=local_block_size if local_block_size > 0 else num_tokens_per_block, ) diff --git a/kv_cache_manager/py_connector/test/BUILD b/kv_cache_manager/py_connector/test/BUILD new file mode 100644 index 000000000..2ce352abb --- /dev/null +++ b/kv_cache_manager/py_connector/test/BUILD @@ -0,0 +1,50 @@ +load("@rules_python//python:py_library.bzl", "py_library") +load("@rules_python//python:py_test.bzl", "py_test") + +# Stubs that make v1_connector importable without vLLM / CUDA / the compiled +# kvcm_py_client. Tests import this module before anything under vllm/. +py_library( + name = "vllm_stubs", + srcs = [ + "__init__.py", + "vllm_stubs.py", + ], + deps = [ + "//kv_cache_manager/py_connector/vllm:vllm_connector", + ], +) + +py_test( + name = "test_block_translation", + srcs = ["test_block_translation.py"], + tags = ["no-remote-exec"], + deps = [":vllm_stubs"], +) + +py_test( + name = "test_data_transfer_results", + srcs = ["test_data_transfer_results.py"], + tags = ["no-remote-exec"], + deps = [":vllm_stubs"], +) + +py_test( + name = "test_kv_layouts", + srcs = ["test_kv_layouts.py"], + tags = ["no-remote-exec"], + deps = [":vllm_stubs"], +) + +py_test( + name = "test_scheduler_state", + srcs = ["test_scheduler_state.py"], + tags = ["no-remote-exec"], + deps = [":vllm_stubs"], +) + +py_test( + name = "test_location_query", + srcs = ["test_location_query.py"], + tags = ["no-remote-exec"], + deps = [":vllm_stubs"], +) diff --git a/kv_cache_manager/py_connector/test/kernel/BUILD b/kv_cache_manager/py_connector/test/kernel/BUILD new file mode 100644 index 000000000..bfb46942b --- /dev/null +++ b/kv_cache_manager/py_connector/test/kernel/BUILD @@ -0,0 +1,19 @@ +load("@rules_python//python:py_test.bzl", "py_test") + +py_test( + name = "test_strided_gather_scatter", + srcs = [ + "__init__.py", + "test_strided_gather_scatter.py", + ], + main = "test_strided_gather_scatter.py", + tags = [ + "exclusive", # GPU tests run serially to avoid CUDA contention + "gpu", # requires 1 GPU + "manual", # needs torch/triton + GPU; not runnable in open-source CI + "no-remote-exec", + ], + deps = [ + "//kv_cache_manager/py_connector/kernel", + ], +) diff --git a/kv_cache_manager/py_connector/test/kernel/test_strided_gather_scatter.py b/kv_cache_manager/py_connector/test/kernel/test_strided_gather_scatter.py new file mode 100644 index 000000000..b1ca1caa2 --- /dev/null +++ b/kv_cache_manager/py_connector/test/kernel/test_strided_gather_scatter.py @@ -0,0 +1,186 @@ +"""GPU tests for the strided path of the batch gather/scatter Triton kernel. + +The flat path (block_stride=0) is covered by test_batch_gather_scatter.py. +Here we cover the strided path added for vLLM's paged layout, where the flat +token index is decomposed as (kv_block, token_in_block) and the block starts +``block_stride`` elements apart -- including padded pages where +``block_stride > local_block_size * dims_per_token`` leaves a gap between +blocks that must be skipped, not walked. + +Every case is checked element-wise against a naive torch reference that +performs the same (kv_block, token) decomposition with plain indexing. +""" + +import unittest + +import torch + +from kv_cache_manager.py_connector.kernel.batch_gather_scatter_helper import ( + batch_gather_kv_caches, + batch_scatter_kv_caches, +) + + +def _make_paged_caches(num_layers, num_blocks, local_block_size, dims_per_token, + pad_tokens, device, dtype, fill_random=True): + """Per-layer paged caches shaped (num_blocks, padded_tokens, dims) where + padded_tokens = local_block_size + pad_tokens. block_stride (in elements) + is padded_tokens * dims_per_token.""" + caches = [] + for _ in range(num_layers): + t = torch.randn(num_blocks, local_block_size + pad_tokens, dims_per_token, + device=device, dtype=dtype) if fill_random else \ + torch.zeros(num_blocks, local_block_size + pad_tokens, dims_per_token, + device=device, dtype=dtype) + caches.append(t) + return caches + + +def _ref_slot(cache, flat_token_idx, local_block_size): + blk = flat_token_idx // local_block_size + tok = flat_token_idx % local_block_size + return cache[blk, tok, :] + + +class TestStridedGatherScatter(unittest.TestCase): + # (local_block_size, pad_tokens, tokens_per_manager_block) + CASES = [ + (16, 0, 16), # strided == flat geometry (stride still exercised) + (16, 4, 16), # padded pages: gap between blocks + (64, 0, 528), # hybrid attention: manager block spans many kv blocks + (64, 8, 48), # padded + manager block not aligned to kv block + ] + + def setUp(self): + if not torch.cuda.is_available(): + self.skipTest("requires a GPU") + torch.manual_seed(7) + self.device = "cuda" + self.dtype = torch.bfloat16 + self.num_layers = 3 + self.dims = 128 + self.num_kv_blocks = 64 + + def _indices(self, num_manager_blocks, tokens_per_block, local_block_size): + total_tokens = self.num_kv_blocks * local_block_size + need = num_manager_blocks * tokens_per_block + assert need <= total_tokens, "test setup: not enough kv slots" + perm = torch.randperm(total_tokens)[:need] + return perm.tolist() + + def test_gather_strided_matches_reference(self): + for local_bs, pad, tokens_per_block in self.CASES: + with self.subTest(local_bs=local_bs, pad=pad, tpb=tokens_per_block): + caches = _make_paged_caches( + self.num_layers, self.num_kv_blocks, local_bs, self.dims, + pad, self.device, self.dtype) + block_stride = caches[0].stride(0) + self.assertEqual(block_stride, (local_bs + pad) * self.dims) + ptrs = torch.tensor([c.data_ptr() for c in caches], + device=self.device, dtype=torch.int64) + num_mb = 4 + token_indices = self._indices(num_mb, tokens_per_block, local_bs) + dst_block_indices = [2, 0, 3, 1] + dst = torch.zeros(num_mb, self.num_layers, tokens_per_block, + self.dims, device="cpu", dtype=self.dtype, + pin_memory=True) + batch_gather_kv_caches( + ptrs, dst, token_indices, dst_block_indices, + tokens_per_block, self.dims, + block_stride=block_stride, local_block_size=local_bs) + torch.cuda.synchronize() + + caches_cpu = [c.cpu() for c in caches] + for mb in range(num_mb): + for pos in range(tokens_per_block): + flat_idx = token_indices[mb * tokens_per_block + pos] + for layer in range(self.num_layers): + want = _ref_slot(caches_cpu[layer], flat_idx, local_bs) + got = dst[dst_block_indices[mb], layer, pos, :] + torch.testing.assert_close( + got, want, + msg=f"gather mismatch mb={mb} pos={pos} " + f"layer={layer} flat={flat_idx}") + + def test_scatter_strided_matches_reference(self): + for local_bs, pad, tokens_per_block in self.CASES: + with self.subTest(local_bs=local_bs, pad=pad, tpb=tokens_per_block): + caches = _make_paged_caches( + self.num_layers, self.num_kv_blocks, local_bs, self.dims, + pad, self.device, self.dtype, fill_random=False) + # Sentinel in the padding region: scatter must never touch it. + sentinel = 123.0 + if pad: + for c in caches: + c[:, local_bs:, :] = sentinel + block_stride = caches[0].stride(0) + ptrs = torch.tensor([c.data_ptr() for c in caches], + device=self.device, dtype=torch.int64) + num_mb = 4 + token_indices = self._indices(num_mb, tokens_per_block, local_bs) + src_block_indices = [1, 3, 0, 2] + src = torch.randn(num_mb, self.num_layers, tokens_per_block, + self.dims, dtype=self.dtype).pin_memory() + batch_scatter_kv_caches( + ptrs, src, token_indices, src_block_indices, + tokens_per_block, self.dims, + block_stride=block_stride, local_block_size=local_bs) + torch.cuda.synchronize() + + caches_cpu = [c.cpu() for c in caches] + for mb in range(num_mb): + for pos in range(tokens_per_block): + flat_idx = token_indices[mb * tokens_per_block + pos] + for layer in range(self.num_layers): + got = _ref_slot(caches_cpu[layer], flat_idx, local_bs) + want = src[src_block_indices[mb], layer, pos, :] + torch.testing.assert_close( + got, want, + msg=f"scatter mismatch mb={mb} pos={pos} " + f"layer={layer} flat={flat_idx}") + if pad: + for layer, c in enumerate(caches_cpu): + self.assertTrue( + bool((c[:, local_bs:, :] == sentinel).all()), + f"scatter wrote into the padding of layer {layer}") + + def test_gather_scatter_roundtrip_strided(self): + """Scattering gathered data into zeroed caches must reproduce exactly + the gathered slots (and only them).""" + local_bs, pad, tokens_per_block = 64, 8, 48 + src_caches = _make_paged_caches( + self.num_layers, self.num_kv_blocks, local_bs, self.dims, + pad, self.device, self.dtype) + dst_caches = _make_paged_caches( + self.num_layers, self.num_kv_blocks, local_bs, self.dims, + pad, self.device, self.dtype, fill_random=False) + block_stride = src_caches[0].stride(0) + src_ptrs = torch.tensor([c.data_ptr() for c in src_caches], + device=self.device, dtype=torch.int64) + dst_ptrs = torch.tensor([c.data_ptr() for c in dst_caches], + device=self.device, dtype=torch.int64) + num_mb = 3 + token_indices = self._indices(num_mb, tokens_per_block, local_bs) + buf = torch.zeros(num_mb, self.num_layers, tokens_per_block, self.dims, + device="cpu", dtype=self.dtype, pin_memory=True) + batch_gather_kv_caches( + src_ptrs, buf, token_indices, list(range(num_mb)), + tokens_per_block, self.dims, + block_stride=block_stride, local_block_size=local_bs) + torch.cuda.synchronize() + batch_scatter_kv_caches( + dst_ptrs, buf, token_indices, list(range(num_mb)), + tokens_per_block, self.dims, + block_stride=block_stride, local_block_size=local_bs) + torch.cuda.synchronize() + src_cpu = [c.cpu() for c in src_caches] + dst_cpu = [c.cpu() for c in dst_caches] + for flat_idx in token_indices: + for layer in range(self.num_layers): + torch.testing.assert_close( + _ref_slot(dst_cpu[layer], flat_idx, local_bs), + _ref_slot(src_cpu[layer], flat_idx, local_bs)) + + +if __name__ == "__main__": + unittest.main() diff --git a/kv_cache_manager/py_connector/test/test_block_translation.py b/kv_cache_manager/py_connector/test/test_block_translation.py new file mode 100644 index 000000000..e1ee4da6f --- /dev/null +++ b/kv_cache_manager/py_connector/test/test_block_translation.py @@ -0,0 +1,126 @@ +"""Unit tests for the connector's manager-block -> physical-slot translation. + +Covers ``_attn_token_indices`` (attention groups: token-granular three-tier +mapping) and ``_state_block_ids`` (mamba/state groups: manager block's last +token selects the group block), verifying against an independent brute-force +reference implementation, token by token. +""" + +import unittest + +from kv_cache_manager.py_connector.test.vllm_stubs import make_connector +from kv_cache_manager.py_connector.vllm.transfer_types import ( + AttentionTransferGroup, KVLayout, StateTransferGroup) + + +def _make_group(group_bs, kernel_bs=0, is_attention=True): + common = dict(group_idx=0, spec_name="tp0_g0", layer_names=["layer0"], + block_size=group_bs, per_block_bytes=0, layer_num=1) + if is_attention: + return AttentionTransferGroup( + kv_layout=KVLayout.PACKED_4D, kvcache_ptr_tensor_gpu=None, + num_kv_ptrs=1, per_token_dim=8, kernel_block_size=kernel_bs, + block_stride=0, **common) + return StateTransferGroup( + block_view_tensors=[], page_size_bytes=0, **common) + + +def _ref_attn_token_indices(manager_bs, group_bs, kernel_bs, manager_block_idxes, + block_table): + """Brute-force reference: walk every token of every manager block and map it + through the block hierarchy step by step.""" + out = [] + for mb in manager_block_idxes: + slots = [] + for tok in range(mb * manager_bs, (mb + 1) * manager_bs): + group_block = tok // group_bs # logical block in group table + tok_in_group = tok - group_block * group_bs + kernel_in_group = tok_in_group // kernel_bs + tok_in_kernel = tok_in_group - kernel_in_group * kernel_bs + physical = block_table[group_block] * (group_bs // kernel_bs) + kernel_in_group + slots.append(physical * kernel_bs + tok_in_kernel) + out.append(slots) + return out + + +def _ref_state_block_ids(manager_bs, group_bs, manager_block_idxes, block_table): + """Brute-force reference: the state covering a manager block is the state of + the group block containing the manager block's last token.""" + out = [] + for mb in manager_block_idxes: + last_token = (mb + 1) * manager_bs - 1 + out.append(block_table[last_token // group_bs]) + return out + + +class TestAttnTokenIndices(unittest.TestCase): + # (manager_bs, group_bs, kernel_bs): ratio=1, ratio>1, manager != group. + CASES = [ + (16, 16, 16), # full attention default: all equal + (32, 16, 16), # preferred_block_size > vllm block size + (528, 528, 64), # hybrid: group block spans several kernel blocks + (528, 528, 528), # hybrid with kernel == group + (48, 16, 8), # manager > group > kernel + ] + + def test_against_reference(self): + for manager_bs, group_bs, kernel_bs in self.CASES: + with self.subTest(manager_bs=manager_bs, group_bs=group_bs, + kernel_bs=kernel_bs): + conn = make_connector(manager_block_size=manager_bs) + group = _make_group(group_bs, kernel_bs) + # Enough non-trivially permuted blocks for 4 manager blocks. + needed = 4 * manager_bs // group_bs + 1 + block_table = [(i * 7 + 3) % 97 for i in range(needed)] + mbis = [0, 1, 3] + got = conn._attn_token_indices(group, mbis, block_table) + want = _ref_attn_token_indices( + manager_bs, group_bs, kernel_bs, mbis, block_table) + self.assertEqual(got, want) + + def test_manual_example(self): + # manager_bs=4, group_bs=2, kernel_bs=2; block_table maps logical + # blocks 0..3 -> physical 5,2,9,0. Manager block 1 covers tokens 4..7 -> + # logical blocks 2,3 -> physical 9,0 -> slots 18,19,0,1. + conn = make_connector(manager_block_size=4) + group = _make_group(group_bs=2, kernel_bs=2) + got = conn._attn_token_indices(group, [1], [5, 2, 9, 0]) + self.assertEqual(got, [[18, 19, 0, 1]]) + + def test_out_of_range_asserts(self): + conn = make_connector(manager_block_size=16) + group = _make_group(group_bs=16, kernel_bs=16) + with self.assertRaises(AssertionError): + conn._attn_token_indices(group, [1], [0]) # table too short + + +class TestStateBlockIds(unittest.TestCase): + def test_against_reference(self): + for manager_bs, group_bs in [(528, 528), (16, 16), (16, 32), (48, 16)]: + with self.subTest(manager_bs=manager_bs, group_bs=group_bs): + conn = make_connector(manager_block_size=manager_bs) + group = _make_group(group_bs, is_attention=False) + needed = 4 * manager_bs // group_bs + 1 + block_table = [(i * 11 + 5) % 89 for i in range(needed)] + mbis = [0, 1, 3] + got = conn._state_block_ids(group, mbis, block_table) + want = _ref_state_block_ids(manager_bs, group_bs, mbis, block_table) + self.assertEqual(got, want) + + def test_manual_example(self): + # manager_bs=4, group_bs=8: manager blocks 0 and 1 both end inside group + # block 0; manager block 2 ends in group block 1. + conn = make_connector(manager_block_size=4) + group = _make_group(group_bs=8, is_attention=False) + got = conn._state_block_ids(group, [0, 1, 2], [7, 3]) + self.assertEqual(got, [7, 7, 3]) + + def test_out_of_range_asserts(self): + conn = make_connector(manager_block_size=16) + group = _make_group(group_bs=16, is_attention=False) + with self.assertRaises(AssertionError): + conn._state_block_ids(group, [2], [0, 1]) + + +if __name__ == "__main__": + unittest.main() diff --git a/kv_cache_manager/py_connector/test/test_data_transfer_results.py b/kv_cache_manager/py_connector/test/test_data_transfer_results.py new file mode 100644 index 000000000..4029ab741 --- /dev/null +++ b/kv_cache_manager/py_connector/test/test_data_transfer_results.py @@ -0,0 +1,452 @@ +"""Unit tests for MultiResult flattening and the save/load done callbacks. + +The done callbacks decode a flat result list whose layout is an implicit +contract with ``_submit_group_tasks``: tasks are submitted group-major +(group0's blocks, then group1's blocks, ...), so a manager block's success is +the stride-AND ``flat[i % num_blocks]``. These tests pin that contract with +hand-computed expectations. +""" + +import threading +import time +import unittest +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import MagicMock + +# vllm_stubs must be imported before torch: in the open-source CI (no torch +# installed) it registers the MagicMock stand-in that the bare `import torch` +# below then resolves to. +from kv_cache_manager.py_connector.test import vllm_stubs # noqa: F401 (stubs) + +import torch +from kv_cache_manager.py_connector.vllm.data_transfer import ( + DataTransferManager, MultiResult) +from kv_cache_manager.py_connector.vllm.transfer_types import KVLayout +from kv_cache_manager.py_connector.common.tp_coordinator import ( + CoordinateMsgSerializer) + + +class TestMultiResult(unittest.TestCase): + def test_flatten_in_submission_order(self): + got = [] + mr = MultiResult(3, got.extend) + mr.submit_result(0, [True, False]) + mr.submit_result(1, [False]) + mr.submit_result(2, [True, True, True]) + self.assertEqual(got, [True, False, False, True, True, True]) + + def test_out_of_order_submit(self): + got = [] + mr = MultiResult(3, got.extend) + mr.submit_result(2, ["c"]) + mr.submit_result(0, ["a"]) + self.assertEqual(got, []) # callback must not fire early + mr.submit_result(1, ["b"]) + self.assertEqual(got, ["a", "b", "c"]) + + def test_duplicate_submit_asserts(self): + mr = MultiResult(2, lambda flat: None) + mr.submit_result(0, [True]) + with self.assertRaises(AssertionError): + mr.submit_result(0, [True]) + + def test_concurrent_submit(self): + n = 64 + results = [] + done = threading.Event() + + def cb(flat): + results.append(flat) + done.set() + + mr = MultiResult(n, cb) + barrier = threading.Barrier(n) + + def worker(i): + barrier.wait() + mr.submit_result(i, [i]) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(n)] + for t in threads: + t.start() + for t in threads: + t.join() + self.assertTrue(done.wait(timeout=5)) + self.assertEqual(len(results), 1) # callback fires exactly once + self.assertEqual(results[0], list(range(n))) + + +def _make_dtm(): + """DataTransferManager with only the state the callbacks touch.""" + dtm = DataTransferManager.__new__(DataTransferManager) + dtm._coordinator_client = MagicMock() + return dtm + + +def _sent_event(dtm): + (payload,), _ = dtm._coordinator_client.send.call_args + return CoordinateMsgSerializer.loads(payload).content + + +class TestSaveDoneCallback(unittest.TestCase): + def test_multi_group_stride_and(self): + # 3 blocks x 2 groups, flat = group0[b0,b1,b2] + group1[b0,b1,b2]. + # Block b is saved only if both groups succeeded for b. + dtm = _make_dtm() + cb = dtm.create_save_done_callback("req", 0, "sess", num_blocks=3) + cb([True, True, False, # group 0 + True, False, True]) # group 1 + evt = _sent_event(dtm) + self.assertEqual(evt.type, "SendBlockFinishedEvent") + self.assertEqual(evt.write_session_id, "sess") + self.assertEqual(evt.is_success_list, [True, False, False]) + + def test_single_group_passthrough(self): + dtm = _make_dtm() + cb = dtm.create_save_done_callback("req", 1, "sess", num_blocks=2) + cb([False, True]) + self.assertEqual(_sent_event(dtm).is_success_list, [False, True]) + + +class TestLoadDoneCallback(unittest.TestCase): + def test_multi_group_failure_merge(self): + dtm = _make_dtm() + cb = dtm.create_load_done_callback( + "req", 0, epoch=7, block_ids=[10, 20, 30], num_blocks=3) + cb([True, False, True, # group 0 + True, True, False]) # group 1 + evt = _sent_event(dtm) + self.assertEqual(evt.type, "LoadBlockFinishedEvent") + self.assertEqual(evt.epoch, 7) + # blocks 1 and 2 each failed in one group -> report their table ids. + self.assertEqual(evt.failed_block_idxs, [20, 30]) + + def test_all_success_reports_empty(self): + dtm = _make_dtm() + cb = dtm.create_load_done_callback( + "req", 0, epoch=0, block_ids=[10, 20], num_blocks=2) + cb([True, True, True, True]) + self.assertEqual(_sent_event(dtm).failed_block_idxs, []) + + def test_report_failures_false_hybrid(self): + # Hybrid models cannot report invalid block ids to vLLM: the failure + # must be swallowed (empty failed list) but the finished event still sent. + dtm = _make_dtm() + cb = dtm.create_load_done_callback( + "req", 0, epoch=1, block_ids=[], num_blocks=2, report_failures=False) + cb([False, True]) + evt = _sent_event(dtm) + self.assertEqual(evt.type, "LoadBlockFinishedEvent") + self.assertEqual(evt.failed_block_idxs, []) + + +class TestNullStateBlocks(unittest.TestCase): + """Mamba 'align' mode: vLLM materializes a recurrent state only at segment + boundaries, so a state group's null (id 0) target carries no state. + + The connector must never turn that absence into a *success*: the block would + be published as fully cached and a later, shorter request would resume from + a state URI nobody ever wrote. Instead the state group abstains for such a + block (None = "no data of mine here"), and the block's verdict is decided by + the groups that did carry data -- the missing state is expressed to the + manager through the block's spec coverage (see + v1_connector._spec_groups). These paths do no GPU work, so they run on CPU. + """ + + @staticmethod + def _state_group(): + from kv_cache_manager.py_connector.vllm.transfer_types import StateTransferGroup + return StateTransferGroup( + group_idx=0, spec_name="tp0_g0", + layer_names=["m0"], block_size=528, per_block_bytes=1024, + layer_num=1, block_view_tensors=[], page_size_bytes=1024) + + @staticmethod + def _attn_group(): + from kv_cache_manager.py_connector.vllm.transfer_types import AttentionTransferGroup + return AttentionTransferGroup( + group_idx=1, spec_name="tp0_g1", + layer_names=["a0"], block_size=528, per_block_bytes=1024, + layer_num=1, kv_layout=KVLayout.PACKED_4D, + kvcache_ptr_tensor_gpu=None, num_kv_ptrs=1, per_token_dim=8, + kernel_block_size=528, block_stride=0) + + def _run(self, method, **kwargs): + dtm = _make_dtm() + results = {} + mr = MultiResult(1, lambda flat: results.setdefault("flat", flat)) + getattr(dtm, method)(mr, 0, self._state_group(), **kwargs) + return results["flat"] + + def test_save_null_state_blocks_abstain_not_succeed(self): + # No state and (consistently) no location for it: the group abstains. + flat = self._run("save_task", + remote_uris=[None, None], + block_token_indices=None, + block_ids=[0, 0], + ready_event=None) + self.assertEqual(flat, [None, None]) + + def test_save_null_state_with_location_fails(self): + # The manager allocated a state location but vLLM has no state to put + # there: publishing it would advertise unwritten bytes. + flat = self._run("save_task", + remote_uris=["u0", None], + block_token_indices=None, + block_ids=[0, 0], + ready_event=None) + self.assertEqual(flat, [False, None]) + + def test_save_real_state_without_location_fails(self): + # A state exists but was not announced: it cannot be published. + flat = self._run("save_task", + remote_uris=[None], + block_token_indices=None, + block_ids=[7], + ready_event=None) + self.assertEqual(flat, [False]) + + def test_load_null_state_targets_abstain(self): + # vLLM does not need a state for these blocks (only the block ending + # the reused prefix does), whatever the manager published. + flat = self._run("load_task", + remote_uris=["u0", "u1", "u2"], + block_token_indices=None, + block_ids=[0, 0, 0]) + self.assertEqual(flat, [None, None, None]) + + def test_load_real_target_without_location_fails(self): + # vLLM needs this state but nothing was published for it: the request + # must not run on an unwritten state. + flat = self._run("load_task", + remote_uris=[None], + block_token_indices=None, + block_ids=[9]) + self.assertEqual(flat, [False]) + + def test_attention_block_without_location_fails(self): + # Attention KV is never sparse: a missing location is a failure, and + # the block must be kept out of the staging batch so the remaining + # buffers stay aligned with the URI list. + dtm = _make_dtm() + skipped, failed = dtm._save_dispositions( + self._attn_group(), remote_uris=["u0", None, "u2"], + block_ids=None, n=3) + self.assertEqual((skipped, failed), (set(), {1})) + + def test_save_dispositions_state_group(self): + dtm = _make_dtm() + # block0: no state, nothing published -> abstain + # block1: state + location -> transfer + # block2: no state but published -> fail + # block3: state but unpublished -> fail + skipped, failed = dtm._save_dispositions( + self._state_group(), remote_uris=[None, "u1", "u2", None], + block_ids=[0, 5, 0, 6], n=4) + self.assertEqual(skipped, {0}) + self.assertEqual(failed, {2, 3}) + + +class TestTaskCrashReporting(unittest.TestCase): + """A task that dies mid-transfer must still report, all-failed. + + submit_task drops the future, so an escaping exception is silently + swallowed: the MultiResult callback never fires, the save session hangs + (SendBlockFinishedEvent never sent) and -- worse -- a load leaves vLLM + believing KV it never received under the connector's synchronous-load + contract. Both tasks wrap their body and report every block as failed.""" + + _state_group = staticmethod(TestNullStateBlocks._state_group) + + def _run_crashing(self, method): + from unittest.mock import patch + dtm = _make_dtm() + results = {} + mr = MultiResult(1, lambda flat: results.setdefault("flat", flat)) + group = self._state_group() + kwargs = dict(remote_uris=["u0", "u1"], + block_token_indices=None, + block_ids=[5, 6]) + crash = "_%s_valid_blocks" % method.split("_")[0] + with patch.object(dtm, crash, side_effect=RuntimeError("boom")): + if method == "save_task": + kwargs["ready_event"] = None + getattr(dtm, method)(mr, 0, group, **kwargs) + return dtm, results["flat"] + + def test_save_task_crash_reports_all_failed(self): + dtm, flat = self._run_crashing("save_task") + self.assertEqual(flat, [False, False]) + + def test_load_task_crash_reports_all_failed(self): + dtm, flat = self._run_crashing("load_task") + self.assertEqual(flat, [False, False]) + + +class TestAbstainedVerdicts(unittest.TestCase): + """A block's verdict is the AND over the groups that carried data for it. + A group that abstained (None) must neither pass nor fail the block, and a + block no group wrote at all must not be published.""" + + def test_save_abstain_does_not_mask_other_group(self): + dtm = _make_dtm() + cb = dtm.create_save_done_callback("req", 0, "sess", num_blocks=3) + cb([None, None, True, # state group: only block 2 had a state + True, False, True]) # attention group + # Block 0 rides on attention alone, block 1 fails there, block 2 both. + self.assertEqual(_sent_event(dtm).is_success_list, [True, False, True]) + + def test_save_all_groups_abstain_is_not_published(self): + dtm = _make_dtm() + cb = dtm.create_save_done_callback("req", 0, "sess", num_blocks=2) + cb([None, None]) + self.assertEqual(_sent_event(dtm).is_success_list, [False, False]) + + def test_load_abstain_does_not_mask_other_group(self): + dtm = _make_dtm() + cb = dtm.create_load_done_callback( + "req", 0, epoch=3, block_ids=[10, 20, 30], num_blocks=3) + cb([None, None, False, # state group: block 2 needed a state, failed + True, True, True]) # attention group + self.assertEqual(_sent_event(dtm).failed_block_idxs, [30]) + + def test_load_all_groups_abstain_counts_as_failure(self): + dtm = _make_dtm() + cb = dtm.create_load_done_callback( + "req", 0, epoch=0, block_ids=[11], num_blocks=1) + cb([None]) + self.assertEqual(_sent_event(dtm).failed_block_idxs, [11]) + + +if __name__ == "__main__": + unittest.main() + + +class TestStagingPool(unittest.TestCase): + """_StagingPool: contiguous-run slot management with backpressure. + + The pool stages transfers in pinned host memory only (the kernel reaches + it directly over PCIe); these tests pin the run bookkeeping: exact fit, + fragmentation and re-merge on release, blocking acquire, and the capacity + guard. + """ + + def _pool(self, max_blocks=8, block_bytes=16): + from kv_cache_manager.py_connector.vllm.data_transfer import _StagingPool + return _StagingPool(torch.device("cpu"), block_bytes, max_blocks) + + def test_roundtrip_and_merge(self): + pool = self._pool() + a = pool.acquire(3) + b = pool.acquire(5) # exact fit of the remainder + self.assertEqual((a, b), (0, 3)) + pool.release(a, 3) + pool.release(b, 5) # neighbours must merge back to one run + self.assertEqual(pool._runs, [[0, 8]]) + self.assertEqual(pool.acquire(8), 0) # full capacity usable again + + def test_fragmentation_blocks_then_merge_wakes(self): + pool = self._pool() + a, b, c = pool.acquire(2), pool.acquire(2), pool.acquire(2) + self.assertEqual((a, b, c), (0, 2, 4)) + pool.release(a, 2) # free: [0,2) and [6,8) + pool.release(c, 2) + # 5 free blocks in total but no contiguous run of 5: blocks. + with ThreadPoolExecutor(max_workers=1) as ex: + fut = ex.submit(pool.acquire, 5) + time.sleep(0.2) + self.assertFalse(fut.done(), "fragmented pool must block") + pool.release(b, 2) # glues [0,8) back together + self.assertEqual(fut.result(timeout=5), 0) + + def test_blocking_acquire_wakes_on_release(self): + pool = self._pool(max_blocks=4) + held = pool.acquire(3) + with ThreadPoolExecutor(max_workers=1) as ex: + fut = ex.submit(pool.acquire, 3) + time.sleep(0.2) + self.assertFalse(fut.done(), "acquire must block while exhausted") + pool.release(held, 3) + self.assertEqual(fut.result(timeout=5), 0) + + def test_oversized_acquire_raises(self): + pool = self._pool(max_blocks=4) + with self.assertRaises(ValueError): + pool.acquire(5) + + @unittest.skipIf("torch" in vllm_stubs.STUBBED, + "needs a real torch tensor (stubbed in the open-source CI)") + def test_views_slice_the_same_run(self): + pool = self._pool(max_blocks=8, block_bytes=16) + start = pool.acquire(3) + cpu = pool.cpu_view(start, 3) + self.assertEqual(cpu.numel(), 48) + cpu.zero_() + self.assertTrue(bool((pool._cpu[0:48] == 0).all())) + + def test_pool_allocates_host_memory_only(self): + """Zero-VRAM regression guard: even for a CUDA device the pool makes + exactly one allocation, and it is pinned host memory -- no device-side + mirror may come back. The fake device keeps the contract testable + wherever torch itself is a stand-in.""" + from types import SimpleNamespace + from unittest.mock import patch + from kv_cache_manager.py_connector.vllm.data_transfer import _StagingPool + with patch("torch.empty") as empty: + _StagingPool(SimpleNamespace(type="cuda"), + per_block_bytes=16, max_blocks=8) + self.assertEqual(empty.call_count, 1, + "pool must own exactly one backing allocation") + kwargs = empty.call_args.kwargs + self.assertEqual(kwargs.get("device"), "cpu") + self.assertTrue(kwargs.get("pin_memory"), + "pool backing memory must be pinned for zero-copy") + + +class TestPoolByteCap(unittest.TestCase): + """_effective_pool_blocks: per-group sizing under the pinned-RAM ceiling. + + Found by the final-validation e2e smoke: a fixed 1024-block count was + tuned for full-attention blocks (~0.875 MiB each) but hybrid blocks are + ~17.3 MiB, so four groups pinned ~68 GiB of host RAM and engine start + died in the pinned allocator. Pinned here: the count is derived per + group from the byte cap, and one full task batch always fits. + """ + + def _f(self, configured, need, block_bytes, max_bytes): + from kv_cache_manager.py_connector.vllm.data_transfer import ( + _effective_pool_blocks) + return _effective_pool_blocks(configured, need, block_bytes, max_bytes) + + def test_full_attention_blocks_keep_the_configured_count(self): + # 1024 x 0.875 MiB ~= 896 MiB <= 1 GiB cap: unchanged sizing (the + # validated origin/main concurrency of 8 full tasks in flight). + self.assertEqual( + self._f(1024, 128, 917_504, 2**30), 1024) + + def test_hybrid_blocks_cap_by_bytes(self): + # 17.3 MiB blocks: 1024 would pin ~17.3 GiB per group; a cap that + # allows more than one task batch caps to cap // block_bytes. + block_bytes = 17_301_504 + self.assertEqual( + self._f(1024, 128, block_bytes, 8 * 2**30), + 8 * 2**30 // block_bytes) + + def test_hybrid_default_cap_floors_at_one_task_batch(self): + # With the 1 GiB default the byte cap alone would allow only 62 + # blocks (< one 128-block task batch); the batch must still fit + # contiguously, so the effective size is the batch -- the same + # configuration the staging-removal campaign validated for hybrid. + self.assertEqual( + self._f(1024, 128, 17_301_504, 2**30), 128) + + def test_one_task_batch_always_fits(self): + # Even when the byte cap is smaller than one task batch, the batch + # must still fit contiguously (contiguity invariant); the caller + # warns that the cap was exceeded. + self.assertEqual( + self._f(1024, 128, 17_301_504, 128 * 1024), 128) + + def test_byte_cap_can_only_shrink_not_grow(self): + # A configured count below the byte cap is authoritative. + self.assertEqual(self._f(256, 128, 917_504, 2**30), 256) diff --git a/kv_cache_manager/py_connector/test/test_kv_layouts.py b/kv_cache_manager/py_connector/test/test_kv_layouts.py new file mode 100644 index 000000000..88852fec6 --- /dev/null +++ b/kv_cache_manager/py_connector/test/test_kv_layouts.py @@ -0,0 +1,300 @@ +"""Unit tests for the multi-version KV cache layout detection. + +``attn_kv_views`` must recognize the three flash_attn layouts vLLM has shipped +(detected from the tensor shape, never from version strings) and reject +anything else: + +* 4-D packed ``(num_blocks, H, block, 2D)`` -- vLLM >= 0.26.0 +* 5-D N-first ``(num_blocks, 2, block, H, D)`` -- vLLM 0.23.0 - 0.25.x +* 5-D KV-first ``(2, num_blocks, block, H, D)`` -- vLLM <= 0.22.1 + +``_build_transfer_group`` must derive the transfer pointers / strides from the +normalized views, and ``ensure_hybrid_supported`` must fail fast when the +installed vLLM's scheduler rejects external KV loads for hybrid models +(vLLM <= 0.22.x). + +Runs without torch: a minimal FakeTensor models the strided-view semantics +(shape / stride / offset / data_ptr) that the code under test reads. +""" + +import sys +import types +import unittest + +from kv_cache_manager.py_connector.test.vllm_stubs import make_connector +from kv_cache_manager.py_connector.vllm.vllm_common import AttentionGroupMeta +from kv_cache_manager.py_connector.vllm.v1_connector import ( + attn_kv_views, ensure_hybrid_supported, GroupMeta) +from kv_cache_manager.py_connector.vllm.transfer_types import KVLayout + +ITEMSIZE = 2 # bf16/fp16 +BASE_PTR = 1 << 20 + + +class FakeTensor: + """Minimal strided tensor: only what attn_kv_views / _build_transfer_group + read (dim/shape/stride/permute/indexing/data_ptr).""" + + def __init__(self, shape, strides, offset=0, base=BASE_PTR): + self.shape = tuple(shape) + self._strides = tuple(strides) + self._offset = offset + self._base = base + + @classmethod + def contiguous(cls, shape, base=BASE_PTR): + strides, acc = [], 1 + for s in reversed(shape): + strides.append(acc) + acc *= s + return cls(shape, tuple(reversed(strides)), base=base) + + def dim(self): + return len(self.shape) + + def stride(self, i=None): + return self._strides if i is None else self._strides[i] + + def data_ptr(self): + return self._base + self._offset * ITEMSIZE + + def permute(self, *dims): + return FakeTensor([self.shape[d] for d in dims], + [self._strides[d] for d in dims], + self._offset, self._base) + + def __getitem__(self, idx): + if isinstance(idx, int): # t[i]: drop dim 0 + return FakeTensor(self.shape[1:], self._strides[1:], + self._offset + idx * self._strides[0], self._base) + if isinstance(idx, tuple) and idx[0] == slice(None) and isinstance(idx[1], int): + # t[:, i]: drop dim 1 + return FakeTensor(self.shape[:1] + self.shape[2:], + self._strides[:1] + self._strides[2:], + self._offset + idx[1] * self._strides[1], self._base) + raise TypeError(f"unsupported index {idx!r}") + + +def packed_4d(n=10, h=4, b=16, d2=256, base=BASE_PTR): + """vLLM >= 0.26.0: NHD memory is (n, b, h, d2) contiguous; the registered + tensor is its (n, h, b, d2) permuted view.""" + return FakeTensor.contiguous([n, b, h, d2], base=base).permute(0, 2, 1, 3) + + +def kv_first_5d(n=10, b=16, h=4, d=128, base=BASE_PTR): + """vLLM <= 0.22.1: (2, n, b, h, d) contiguous.""" + return FakeTensor.contiguous([2, n, b, h, d], base=base) + + +def n_first_5d(n=10, b=16, h=4, d=128, base=BASE_PTR): + """vLLM 0.23.0 - 0.25.x: (n, 2, b, h, d) contiguous.""" + return FakeTensor.contiguous([n, 2, b, h, d], base=base) + + +class TestAttnKvViews(unittest.TestCase): + def test_packed_4d(self): + views, layout = attn_kv_views(packed_4d()) + self.assertIs(layout, KVLayout.PACKED_4D) + self.assertEqual(len(views), 1) + v = views[0] + self.assertEqual(v.shape, (10, 16, 4, 256)) # (n, b, h, 2d) + self.assertEqual(v.stride(), (16 * 4 * 256, 4 * 256, 256, 1)) + self.assertEqual(v.data_ptr(), BASE_PTR) # storage base + + def test_kv_first_5d(self): + views, layout = attn_kv_views(kv_first_5d()) + self.assertIs(layout, KVLayout.SPLIT_KV_5D_KV_FIRST) + self.assertEqual(len(views), 2) + k, v = views + for view in (k, v): + self.assertEqual(view.shape, (10, 16, 4, 128)) + self.assertEqual(view.stride(), (16 * 4 * 128, 4 * 128, 128, 1)) + self.assertEqual(k.data_ptr(), BASE_PTR) + # V base = K base + num_blocks * block * h * d elements. + self.assertEqual(v.data_ptr() - k.data_ptr(), + 10 * 16 * 4 * 128 * ITEMSIZE) + + def test_n_first_5d(self): + views, layout = attn_kv_views(n_first_5d()) + self.assertIs(layout, KVLayout.SPLIT_KV_5D_N_FIRST) + self.assertEqual(len(views), 2) + k, v = views + for view in (k, v): + self.assertEqual(view.shape, (10, 16, 4, 128)) + # K and V of one block are interleaved: the block stride covers + # both halves while the inner page stays token-major. + self.assertEqual(view.stride(), (2 * 16 * 4 * 128, 4 * 128, 128, 1)) + self.assertEqual(v.data_ptr() - k.data_ptr(), + 16 * 4 * 128 * ITEMSIZE) + + def test_unrecognized_layouts_fail_fast(self): + bad = [ + FakeTensor.contiguous([10, 16, 4]), # 3-D + FakeTensor.contiguous([10, 2, 16, 4, 128, 2]), # 6-D + FakeTensor.contiguous([10, 16, 2, 4, 128]), # 5-D, K/V dim misplaced + ] + for t in bad: + with self.subTest(shape=t.shape): + with self.assertRaises(NotImplementedError): + attn_kv_views(t) + + def test_ambiguous_layout_fails_fast(self): + # num_blocks == 2 in a KV-first shape is indistinguishable from a + # two-block N-first shape; refusing beats guessing. + with self.assertRaises(NotImplementedError): + attn_kv_views(FakeTensor.contiguous([2, 2, 16, 4, 128])) + + +def _make_group_conn(): + conn = make_connector(manager_block_size=16) + conn._self_spec_names = ["tp0_g0"] + conn._device = "cpu" + return conn + + +def _attn_meta(layer_names, block_size=16): + return AttentionGroupMeta(group_idx=0, layer_names=layer_names, + block_size=block_size, per_block_bytes=0) + + +class TestBuildTransferGroup(unittest.TestCase): + """Pointer construction per layout. Layer tensors get distinct bases so the + interleaving [K0, V0, K1, V1, ...] is observable. The pointer list is + captured by patching ``torch.tensor`` (works with both the stubbed and a + real torch: no tensor math happens on the captured value).""" + + def _build(self, kv_caches): + import unittest.mock as mock + import kv_cache_manager.py_connector.vllm.connector_worker as wc + conn = _make_group_conn() + captured = [] + + def fake_tensor(data, **kw): + captured[:] = list(data) + t = mock.MagicMock() + t.to.return_value = t + return t + + with mock.patch.object(wc.torch, "tensor", side_effect=fake_tensor): + g = conn._build_attention_group( + _attn_meta(list(kv_caches.keys())), kv_caches) + return g, captured + + def test_packed_one_ptr_per_layer(self): + kv = {"l0": packed_4d(base=BASE_PTR), "l1": packed_4d(base=2 * BASE_PTR)} + g, ptrs = self._build(kv) + self.assertEqual(g.num_kv_ptrs, 2) + self.assertEqual(g.layer_num, 2) + self.assertEqual(g.per_token_dim, 4 * 256) + self.assertEqual(g.kernel_block_size, 16) + self.assertEqual(g.block_stride, 0) # flat + self.assertEqual(ptrs, [BASE_PTR, 2 * BASE_PTR]) + + def test_kv_first_two_ptrs_per_layer(self): + kv = {"l0": kv_first_5d(base=BASE_PTR), "l1": kv_first_5d(base=2 * BASE_PTR)} + g, ptrs = self._build(kv) + self.assertEqual(g.num_kv_ptrs, 4) + self.assertEqual(g.layer_num, 2) + self.assertEqual(g.per_token_dim, 4 * 128) + self.assertEqual(g.block_stride, 0) # each half is flat token-major + v_off = 10 * 16 * 4 * 128 * ITEMSIZE + self.assertEqual(ptrs, [BASE_PTR, BASE_PTR + v_off, + 2 * BASE_PTR, 2 * BASE_PTR + v_off]) + + def test_n_first_strided_blocks(self): + kv = {"l0": n_first_5d(base=BASE_PTR)} + g, ptrs = self._build(kv) + self.assertEqual(g.num_kv_ptrs, 2) + self.assertEqual(g.per_token_dim, 4 * 128) + # K/V interleaved per block -> kernel must walk the strided path. + self.assertEqual(g.block_stride, 2 * 16 * 4 * 128) + v_off = 16 * 4 * 128 * ITEMSIZE + self.assertEqual(ptrs, [BASE_PTR, BASE_PTR + v_off]) + + def test_unrecognized_layout_fails_fast(self): + with self.assertRaises(NotImplementedError): + self._build({"l0": FakeTensor.contiguous([10, 16, 4])}) + + +class _BlockedScheduler: + """Mimics vLLM <= 0.22.x: external loads are asserted away.""" + + def _mamba_block_aligned_split(self, request, num_new_tokens, + num_new_local_computed_tokens=0, + num_external_computed_tokens=0): + assert num_external_computed_tokens == 0, ( + "External KV connector is not verified yet" + ) + + +class _OpenScheduler: + """Mimics vLLM >= 0.23.0: the split handles external tokens.""" + + def _mamba_block_aligned_split(self, request, num_new_tokens, + num_new_local_computed_tokens=0, + num_external_computed_tokens=0): + return num_new_tokens + + +# Method exists but inspect.getsource fails (frozen / bytecode-only vLLM): +# compiled from a string, so there is no source file to read. +_exec_ns = {} +exec(compile("def _mamba_block_aligned_split(self, *a, **kw):\n pass\n", + "", "exec"), _exec_ns) +_SourcelessScheduler = type( + "_SourcelessScheduler", (), + {"_mamba_block_aligned_split": _exec_ns["_mamba_block_aligned_split"]}) + + +class TestHybridGate(unittest.TestCase): + MOD = "vllm.v1.core.sched.scheduler" + + def _with_scheduler(self, cls): + mod = types.ModuleType(self.MOD) + mod.Scheduler = cls + old = sys.modules.get(self.MOD) + sys.modules[self.MOD] = mod + self.addCleanup(lambda: (sys.modules.pop(self.MOD, None), + old and sys.modules.__setitem__(self.MOD, old))) + + def test_old_vllm_hybrid_raises_gracefully(self): + self._with_scheduler(_BlockedScheduler) + with self.assertRaises(NotImplementedError) as ctx: + ensure_hybrid_supported() + # The message must tell the operator what to do. + self.assertIn("vLLM >= 0.23.0", str(ctx.exception)) + self.assertIn("hybrid", str(ctx.exception)) + + def test_new_vllm_hybrid_passes(self): + self._with_scheduler(_OpenScheduler) + ensure_hybrid_supported() # must not raise + + def test_method_removed_does_not_block(self): + # Future vLLM refactors _mamba_block_aligned_split away: the blocking + # assert went with it, so hybrid must not be blocked. + self._with_scheduler(object) # no _mamba_block_aligned_split at all + ensure_hybrid_supported() # must not raise + + def test_sourceless_method_fails_closed(self): + # Method exists but its source is unavailable: the vllm <= 0.22.x + # blocking assert cannot be ruled out, so the gate must fail closed + # with an actionable override hint. + self._with_scheduler(_SourcelessScheduler) + with self.assertRaises(NotImplementedError) as ctx: + ensure_hybrid_supported() + self.assertIn("force_hybrid_support", str(ctx.exception)) + + def test_sourceless_method_force_override(self): + self._with_scheduler(_SourcelessScheduler) + ensure_hybrid_supported(force=True) # must not raise + + def test_force_does_not_unblock_known_bad_vllm(self): + # force only bypasses the *inconclusive* probe; a positively detected + # blocking assert still raises. + self._with_scheduler(_BlockedScheduler) + with self.assertRaises(NotImplementedError): + ensure_hybrid_supported(force=True) + + +if __name__ == "__main__": + unittest.main() diff --git a/kv_cache_manager/py_connector/test/test_location_query.py b/kv_cache_manager/py_connector/test/test_location_query.py new file mode 100644 index 000000000..e64c3c2ee --- /dev/null +++ b/kv_cache_manager/py_connector/test/test_location_query.py @@ -0,0 +1,258 @@ +"""Unit tests for LocationQueryManager's query fan-out and supersession. + +The scheduler re-asks ``get_num_new_matched_tokens`` every step while an +async location query is in flight. ``get_locations_for_query`` must absorb +those re-asks: issuing one full ``getCacheLocation`` per engine step while +the answer is on the wire multiplies manager load ~7x per request in the +vLLM 0.22.1 e2e perf harness (24 requests -> 163 HTTP queries). + +The cache holds one slot per request. A re-ask at the same offset waits +for (or serves) the slot; a re-ask at a different offset *supersedes* it +-- the newest ask wins, and the old ask's answer may arrive late but must +never write into the new slot (late answers only beat older asks, never +newer ones). +""" + +import threading +import time +import unittest +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace + +from kv_cache_manager.py_connector.test.vllm_stubs import ( # noqa: F401 installs stubs + _install_stubs, +) +from kv_cache_manager.py_connector.vllm.location_query_manager import ( + LocationQueryManager, +) + + +class GatedClient: + """Fake manager client: each call blocks on its own gate and answers + from a per-call answer list (the last answer repeats for extra calls). + ``open=True`` pre-sets every gate (sync-mode tests, where the fetch + runs inline and nobody else can release it).""" + + def __init__(self, answers=None, open=False): + self.calls = 0 + self.lock = threading.Lock() + self.gates = [] + self.answers = list(answers or [["loc0", "loc1"]]) + self.open = open + + def get_cache_location(self, request): + with self.lock: + self.calls += 1 + idx = self.calls - 1 + gate = threading.Event() + if self.open: + gate.set() + self.gates.append(gate) + gate.wait(timeout=10) + return {"locations": list(self.answers[min(idx, len(self.answers) - 1)])} + + +class LocationQueryFanoutTest(unittest.TestCase): + def setUp(self): + self.client = GatedClient() + self.executor = ThreadPoolExecutor(max_workers=4) + self.req = SimpleNamespace(request_id="r1", prompt_token_ids=[1, 2, 3]) + + def tearDown(self): + self.executor.shutdown(wait=True) + + def _manager(self, async_mode=True): + return LocationQueryManager( + self.client, self.executor, "inst", + async_get_cache_location=async_mode) + + def _await_answer(self, lqm, offset): + """Pump the hook until the slot answers; assert the client was hit.""" + deadline = time.time() + 10 + while time.time() < deadline: + result = lqm.get_locations_for_query(self.req, offset) + if result is not None: + return result + time.sleep(0.01) + self.fail(f"query at offset {offset} never answered") + + def _gate(self, idx): + """Wait for the idx-th call's gate to exist, then return it.""" + deadline = time.time() + 5 + while time.time() < deadline: + with self.client.lock: + if len(self.client.gates) > idx: + return self.client.gates[idx] + time.sleep(0.005) + self.fail(f"gate {idx} was never created (calls={self.client.calls})") + + # ------------------------------------------------------------------ # + # Dedupe: same offset + # ------------------------------------------------------------------ # + def test_in_flight_reasks_do_not_fan_out(self): + lqm = self._manager() + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) # submit + for _ in range(5): # scheduler re-asks while the query is on the wire + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) + time.sleep(0.2) + self.assertEqual(self.client.calls, 1, + "re-asks while in flight must not issue new queries") + + self._gate(0).set() + self.assertEqual(self._await_answer(lqm, 0), ["loc0", "loc1"]) + self.assertEqual(self.client.calls, 1, + "cached answer must be served without a new query") + + # ------------------------------------------------------------------ # + # Supersession: a different offset is a different query + # ------------------------------------------------------------------ # + def test_superseding_offset_reissues_once_answered(self): + lqm = self._manager() + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) # submit + self._gate(0).set() + self.assertEqual(self._await_answer(lqm, 0), ["loc0", "loc1"]) + self.assertEqual(self.client.calls, 1) + # A new offset invalidates the cached answer and re-issues exactly one + # query (the documented re-ask path for grown prefixes). + self.assertIsNone(lqm.get_locations_for_query(self.req, 5)) + self.assertIsNone(lqm.get_locations_for_query(self.req, 5)) # in flight + time.sleep(0.2) + self.assertEqual(self.client.calls, 2) + + def test_superseding_offset_starts_immediately(self): + # A grown offset must issue its own RPC at once, not wait behind the + # older offset's in-flight query -- and once superseded, the older + # offset is no longer addressable: asking it again supersedes back. + lqm = self._manager() + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) # RPC #1 + self.assertIsNone(lqm.get_locations_for_query(self.req, 8)) # RPC #2 + self.assertEqual(self.client.calls, 2) + self.assertIsNone(lqm.get_locations_for_query(self.req, 8)) # same key: dedup + self.assertEqual(self.client.calls, 2) + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) # older offset + self.assertEqual(self.client.calls, 3, # supersedes back + "asking the superseded offset must re-issue, not serve") + + def test_superseded_answer_is_not_served_at_its_offset(self): + # offset 0 answered, then superseded by offset 8: even when offset 8 + # answered too, asking offset 0 again starts a fresh query instead of + # serving the dead slot's answer. + lqm = self._manager() + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) # submit + self._gate(0).set() + self.assertEqual(self._await_answer(lqm, 0), ["loc0", "loc1"]) + self.assertIsNone(lqm.get_locations_for_query(self.req, 8)) # supersede + self._gate(1).set() + self.assertEqual(self._await_answer(lqm, 8), ["loc0", "loc1"]) + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) # fresh RPC + self._gate(2) # wait until the fresh query reaches the client + self.assertEqual(self.client.calls, 3) + + # ------------------------------------------------------------------ # + # Late answers never beat newer asks + # ------------------------------------------------------------------ # + def test_late_answer_of_old_ask_does_not_write_new_slot(self): + # Old ask answered first: its answer must be dropped, the new slot + # stays in flight until its own answer lands. + lqm = self._manager() + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) # RPC #1 + self.assertIsNone(lqm.get_locations_for_query(self.req, 8)) # RPC #2 + self._gate(0).set() # old answers + deadline = time.time() + 5 + while time.time() < deadline and self.client.calls < 2: + time.sleep(0.01) + time.sleep(0.1) + # The new slot must still be in flight (not answered by the old ask). + self.assertIsNone(lqm.get_locations_for_query(self.req, 8)) + self._gate(1).set() + self.assertEqual(self._await_answer(lqm, 8), ["loc0", "loc1"]) + + def test_late_answer_of_old_ask_does_not_overwrite_new_answer(self): + # New ask answered first, then the old answer lands: it must not + # overwrite the new slot's answer. + lqm = self._manager() + self.client.answers = [["old_answers"], ["new_answers"]] + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) # RPC #1 + self.assertIsNone(lqm.get_locations_for_query(self.req, 8)) # RPC #2 + self._gate(1).set() # new answers + self.assertEqual(self._await_answer(lqm, 8), ["new_answers"]) + self._gate(0).set() # old answers late + time.sleep(0.2) + self.assertEqual(lqm.get_locations_for_query(self.req, 8), ["new_answers"], + "the superseded ask must not overwrite the new slot") + + # ------------------------------------------------------------------ # + # Failure and sync mode + # ------------------------------------------------------------------ # + def test_sync_mode_answers_inline_and_caches(self): + self.client = GatedClient(open=True) + lqm = self._manager(async_mode=False) + self.assertEqual(lqm.get_locations_for_query(self.req, 0), ["loc0", "loc1"]) + self.assertEqual(lqm.get_locations_for_query(self.req, 0), ["loc0", "loc1"]) + self.assertEqual(self.client.calls, 1) + + def test_failed_query_drops_slot_for_retry(self): + lqm = self._manager() + + def boom(request): + raise RuntimeError("manager down") + + self.client.get_cache_location = boom + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) + deadline = time.time() + 10 # failure pops the slot asynchronously + while time.time() < deadline: + with lqm._lock: + if not lqm._queries.get("r1"): + break + time.sleep(0.01) + # The next ask re-issues instead of waiting forever on a dead slot. + restore = GatedClient.get_cache_location + self.client.get_cache_location = \ + lambda request: restore(self.client, request) + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) # re-issue + self._gate(0).set() + self.assertEqual(self._await_answer(lqm, 0), ["loc0", "loc1"]) + + # ------------------------------------------------------------------ # + # Consume / store / invalidate + # ------------------------------------------------------------------ # + def test_consume_returns_answers_and_empties(self): + lqm = self._manager() + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) # submit + self._gate(0).set() + self.assertEqual(self._await_answer(lqm, 0), ["loc0", "loc1"]) + # The hook clamps the answer before the allocation consumes it. + lqm.store_result("r1", ["loc0"]) + self.assertEqual(lqm.consume_locations("r1"), (["loc0"], 0)) + self.assertIsNone(lqm.consume_locations("r1")) + + def test_consumed_slot_discards_the_inflight_answer(self): + # consume pops even an in-flight slot; the late answer must not + # resurrect a slot for a request vLLM already moved past. + lqm = self._manager() + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) + self.assertIsNone(lqm.consume_locations("r1")) # pops the in-flight slot + self._gate(0).set() + time.sleep(0.2) + with lqm._lock: + self.assertNotIn("r1", lqm._queries) + # The next ask is a fresh query. + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) + self._gate(1) # wait until the fresh query reaches the client + self.assertEqual(self.client.calls, 2) + + def test_invalidate_drops_slot_and_discards_late_answer(self): + lqm = self._manager() + self.assertIsNone(lqm.get_locations_for_query(self.req, 0)) + lqm.invalidate("r1") + with lqm._lock: + self.assertNotIn("r1", lqm._queries) + self._gate(0).set() + time.sleep(0.2) + with lqm._lock: + self.assertNotIn("r1", lqm._queries, + "a late answer must not resurrect the slot") + + +if __name__ == "__main__": + unittest.main() diff --git a/kv_cache_manager/py_connector/test/test_scheduler_state.py b/kv_cache_manager/py_connector/test/test_scheduler_state.py new file mode 100644 index 000000000..f3b99730e --- /dev/null +++ b/kv_cache_manager/py_connector/test/test_scheduler_state.py @@ -0,0 +1,1055 @@ +"""Unit tests for the connector's scheduler-side logic. + +Covers, against fake vLLM SchedulerOutput / Request objects: + +* ``get_num_new_matched_tokens`` -- including the full-prompt external hit cap + (a fully cached prompt must leave >= 1 token to recompute, otherwise vLLM's + synchronous-load scheduling path asserts ``num_new_tokens > 0``); +* ``parse_block_mask_to_save_indices`` -- ``offset`` and ``bool_masks`` forms; +* ``_parse_groups`` -- full-attention single group, hybrid multi group, eagle + group skip, unsupported spec error; +* ``build_connector_meta`` -- new request, cached deltas (``new_block_ids`` + None / non-None), preemption resume via both the 0.26 ``resumed_req_ids`` and + the legacy ``resumed_from_preemption`` interfaces, save-threshold trigger, + and the two ``request_finished`` paths (saves landed / in-flight). +""" + +import unittest +from dataclasses import dataclass, field +from types import SimpleNamespace +from unittest.mock import MagicMock + +from kv_cache_manager.py_connector.test.vllm_stubs import ( + make_connector, make_connector_scheduler, GroupMeta) +from kv_cache_manager.py_connector.vllm.connector_scheduler import RequestLedger +from kv_cache_manager.py_connector.vllm.vllm_common import ( + AttentionGroupMeta, StateGroupMeta, parse_groups) +from kv_cache_manager.py_connector.vllm.v1_connector import TairKvCacheConnector +from kv_cache_manager.py_connector.vllm.metadata import ( + SaveRequest, LoadRequest, TairKvCacheConnectorMetadata) + + +# --------------------------------------------------------------------------- # +# Fakes +# --------------------------------------------------------------------------- # +@dataclass +class FakeRequest: + request_id: str + prompt_token_ids: list + output_token_ids: list = field(default_factory=list) + + @property + def num_tokens(self): + return len(self.prompt_token_ids) + len(self.output_token_ids) + + @property + def all_token_ids(self): + return self.prompt_token_ids + self.output_token_ids + + +def make_scheduler_connector(mbs=16, vllm_bs=None, locations=None, + num_groups=1, num_state_groups=0, tp_size=1): + """ConnectorScheduler with the scheduler-loop state and a mocked query manager.""" + return make_connector_scheduler( + manager_block_size=mbs, vllm_block_size=vllm_bs, + num_groups=num_groups, num_state_groups=num_state_groups, + tp_size=tp_size, locations=locations) + + +def fake_scheduler_output(new_reqs=(), cached_req_ids=(), num_scheduled=None, + new_block_ids=(), resumed_req_ids=frozenset(), + legacy_resumed=None): + """Build a fake SchedulerOutput. legacy_resumed switches the cached-reqs + container to the pre-0.26 interface (resumed_from_preemption list, no + resumed_req_ids attribute).""" + if legacy_resumed is not None: + cached = SimpleNamespace( + req_ids=list(cached_req_ids), + resumed_from_preemption=list(legacy_resumed), + new_block_ids=list(new_block_ids), + ) + else: + cached = SimpleNamespace( + req_ids=list(cached_req_ids), + resumed_req_ids=set(resumed_req_ids), + new_block_ids=list(new_block_ids), + ) + return SimpleNamespace( + scheduled_new_reqs=list(new_reqs), + scheduled_cached_reqs=cached, + num_scheduled_tokens=dict(num_scheduled or {}), + ) + + +def make_locations(n): + return [{"location_specs": [{"name": "tp0_g0", "uri": f"file://blk{i}"}]} + for i in range(n)] + + +def alloc(conn, req, ext_tokens, blocks): + """Simulate update_state_after_alloc: vLLM allocated blocks (per group) + for a match of ext_tokens.""" + conn.update_state_after_alloc( + req, SimpleNamespace(get_block_ids=lambda: [list(b) for b in blocks]), + ext_tokens) + + +# --------------------------------------------------------------------------- # +# get_num_new_matched_tokens +# --------------------------------------------------------------------------- # +class TestGetNumNewMatchedTokens(unittest.TestCase): + MBS = 16 + + def _run(self, prompt_len, num_computed, num_locations): + conn = make_scheduler_connector( + mbs=self.MBS, locations=make_locations(num_locations)) + conn._location_query_manager.last_computed_blocks = num_computed // self.MBS + req = FakeRequest("r0", list(range(prompt_len))) + matched, async_load = conn.get_num_new_matched_tokens(req, num_computed) + return conn, matched, async_load + + def test_partial_hit_no_cap(self): + conn, matched, async_load = self._run(4 * self.MBS + 5, 0, 4) + self.assertEqual(matched, 4 * self.MBS) + self.assertTrue(async_load) # a pending load is reported as async + self.assertEqual(conn._waiting_to_load_requests, []) + req = FakeRequest("r0", list(range(4 * self.MBS + 5))) + alloc(conn, req, matched, [[100, 101, 102, 103]]) + self.assertEqual(conn._waiting_to_load_requests[0].manager_block_idxes, + [0, 1, 2, 3]) + self.assertEqual(conn._waiting_to_load_requests[0].all_block_ids, + [[100, 101, 102, 103]]) + + def test_full_hit_capped_to_leave_one_token(self): + # Prompt is exactly 4 manager blocks, all externally cached: the last + # block must be dropped so vLLM still schedules >= 1 new token. + conn, matched, _ = self._run(4 * self.MBS, 0, 4) + self.assertEqual(matched, 3 * self.MBS) + req = FakeRequest("r0", list(range(4 * self.MBS))) + alloc(conn, req, matched, [[100, 101, 102]]) + self.assertEqual(conn._waiting_to_load_requests[0].manager_block_idxes, + [0, 1, 2]) + # The ledger counts only the blocks actually treated as hit. + self.assertEqual(conn._tracked["r0"].has_saved_block_num, 3) + + def test_full_hit_with_local_prefix(self): + # 2 blocks locally computed + 2 remote = whole prompt -> drop one remote. + conn, matched, _ = self._run(4 * self.MBS, 2 * self.MBS, 2) + self.assertEqual(matched, self.MBS) + req = FakeRequest("r0", list(range(4 * self.MBS))) + alloc(conn, req, matched, [[100, 101, 102]]) + self.assertEqual(conn._waiting_to_load_requests[0].manager_block_idxes, [2]) + + def test_single_block_full_hit_degrades_to_zero(self): + conn, matched, async_load = self._run(self.MBS, 0, 1) + self.assertEqual(matched, 0) + self.assertFalse(async_load) + self.assertEqual(conn._waiting_to_load_requests, []) + + def test_no_locations(self): + conn, matched, async_load = self._run(100, 0, 0) + self.assertEqual(matched, 0) + self.assertFalse(async_load) + + def test_query_in_flight_answers_none(self): + # The manager query runs async; until it lands the connector answers + # None, vLLM re-asks next step instead of blocking the scheduler loop. + conn = make_scheduler_connector(mbs=self.MBS, locations=None) # in flight + req = FakeRequest("r0", list(range(4 * self.MBS + 5))) + matched, async_load = conn.get_num_new_matched_tokens(req, 0) + self.assertIsNone(matched) + self.assertFalse(async_load) + self.assertEqual(conn._waiting_to_load_requests, []) + # No request state exists while the query is pending: the match hook + # is a pure query now; the ledger starts at the first allocation. + self.assertNotIn("r0", conn._tracked) + + def test_requery_after_load_failure_skips_external(self): + # A failed load comes back to the scheduler as invalid block ids + # (update_connector_output); the re-query under kv_load_failure_policy + # =recompute must not re-match: the manager still advertises the + # blocks whose bytes are gone, and re-matching loops + # fail -> reschedule forever. + conn = make_scheduler_connector(mbs=self.MBS, locations=make_locations(2)) + req = FakeRequest("r0", list(range(4 * self.MBS + 5))) + matched, _ = conn.get_num_new_matched_tokens(req, 0) + self.assertEqual(matched, 2 * self.MBS) + # vLLM allocates blocks for the load attempt; the load then fails. + alloc(conn, req, matched, [[100, 101, 102]]) + conn.update_connector_output(SimpleNamespace(invalid_block_ids={101})) + self.assertIn("r0", conn._load_failed) + # Retry: same request re-enters the waiting queue with 0 computed. + matched2, async2 = conn.get_num_new_matched_tokens(req, 0) + self.assertEqual(matched2, 0) + self.assertFalse(async2) + self.assertEqual(len(conn._waiting_to_load_requests), 1) # no new load + + def test_preempted_requery_keeps_external_match(self): + # Preemption alone is not a failure: the loaded KV is healthy, only + # the scheduling position was lost, so the re-query matches again + # (single-table shapes only -- multi-table cannot tell failure from + # preemption, see test_hybrid_load_attempt_burns_the_match). + conn = make_scheduler_connector(mbs=self.MBS, locations=make_locations(2)) + req = FakeRequest("r0", list(range(4 * self.MBS + 5))) + matched, _ = conn.get_num_new_matched_tokens(req, 0) + alloc(conn, req, matched, [[100, 101, 102]]) + conn.update_connector_output(SimpleNamespace(invalid_block_ids=set())) + matched2, _ = conn.get_num_new_matched_tokens(req, 0) + self.assertEqual(matched2, matched) # re-matched, not fast-failed + alloc(conn, req, matched2, [[100, 101, 102]]) + self.assertEqual(len(conn._waiting_to_load_requests), 2) + + def test_hybrid_load_attempt_burns_the_match(self): + # Hybrid load failures cannot be reported to vLLM (single-group + # invalid-block recovery only), so no explicit signal exists: any + # allocation for an external hint burns the match conservatively. + conn = make_scheduler_connector( + mbs=self.MBS, num_groups=1, num_state_groups=1, + locations=hybrid_locations([True, True])) + req = FakeRequest("r0", list(range(4 * self.MBS + 5))) + matched, _ = conn.get_num_new_matched_tokens(req, 0) + self.assertEqual(matched, 2 * self.MBS) + alloc(conn, req, matched, [[100, 101], [50, 51]]) + # Even with no failure signal at all, the re-query fast-fails. + matched2, async2 = conn.get_num_new_matched_tokens(req, 0) + self.assertEqual(matched2, 0) + self.assertFalse(async2) + + def test_skipped_drafter_shape_burns_the_match(self): + # A skipped drafter group leaves two vLLM block tables even for an + # attention-only model: failures cannot be reported (single-group + # recovery unpack), so the conservative one-shot burn applies by + # block-table shape, not by hybridness. + conn = make_scheduler_connector(mbs=self.MBS, locations=make_locations(2)) + req = FakeRequest("r0", list(range(4 * self.MBS + 5))) + matched, _ = conn.get_num_new_matched_tokens(req, 0) + self.assertEqual(matched, 2 * self.MBS) + # Two tables: group 0 (skipped drafter) + group 1 (transferred). + alloc(conn, req, matched, [[900], [100, 101, 102]]) + matched2, async2 = conn.get_num_new_matched_tokens(req, 0) + self.assertEqual(matched2, 0) + self.assertFalse(async2) + + def test_multi_attention_group_shape_burns_the_match(self): + # Two transferred attention groups: the same multi-table shape, the + # same one-shot burn, despite the model being attention-only. + conn = make_scheduler_connector( + mbs=self.MBS, num_groups=2, locations=make_locations(2)) + req = FakeRequest("r0", list(range(4 * self.MBS + 5))) + matched, _ = conn.get_num_new_matched_tokens(req, 0) + self.assertEqual(matched, 2 * self.MBS) + alloc(conn, req, matched, [[100, 101], [200, 201]]) + matched2, async2 = conn.get_num_new_matched_tokens(req, 0) + self.assertEqual(matched2, 0) + self.assertFalse(async2) + + def test_load_attempted_flag_tracks_the_one_external_load(self): + # load_attempted is explicit state, not inference from other fields: + # set exactly when vLLM allocates blocks for an external hit. It only + # burns the re-query for multi-table shapes (which get no failure + # signal); single-table requests burn theirs through load_failed. + conn = make_scheduler_connector(mbs=self.MBS, locations=make_locations(2)) + req = FakeRequest("r0", list(range(4 * self.MBS + 5))) + conn.get_num_new_matched_tokens(req, 0) + # Query done, but no allocation yet (vLLM may still re-ask us). + self.assertNotIn("r0", conn._load_attempted) + # An allocation with no external hit is not a load attempt. + alloc(conn, req, 0, [[100]]) + self.assertNotIn("r0", conn._load_attempted) + # Blocks allocated for an external hit: the load attempt begins. + alloc(conn, req, 2 * self.MBS, [[100, 101, 102]]) + self.assertIn("r0", conn._load_attempted) + + +# --------------------------------------------------------------------------- # +# Per-block spec coverage (hybrid sparse recurrent state) +# --------------------------------------------------------------------------- # +def hybrid_locations(coverage, tp_size=1, num_attn=1, num_state=1): + """Manager locations whose per-block spec set encodes ``coverage``: + True -> every group's spec (state included), False -> attention specs only. + """ + locs = [] + for complete in coverage: + names = [f"tp{r}_g{g}" for r in range(tp_size) for g in range(num_attn)] + if complete: + names += [f"tp{r}_g{num_attn + g}" + for r in range(tp_size) for g in range(num_state)] + locs.append({"location_specs": [{"name": n, "uri": f"u_{n}"} + for n in names]}) + return locs + + +class TestSpecGroups(unittest.TestCase): + """Registration must advertise the two spec groups a hybrid model needs to + express per-block state sparsity -- and must stay silent for models that + have no sparsity (byte-identical requests, old-manager compatible).""" + + def test_full_attention_declares_no_groups(self): + conn = make_connector_scheduler(num_groups=1, tp_size=2) + self.assertEqual(conn._spec_groups(), []) + + def test_hybrid_declares_attn_and_full(self): + conn = make_connector_scheduler(num_groups=1, num_state_groups=2, tp_size=2) + groups = {g["name"]: g["spec_names"] for g in conn._spec_groups()} + self.assertEqual(sorted(groups), ["attn", "full"]) + # attn: the attention spec of every rank; full: every group of every rank. + self.assertEqual(groups["attn"], ["tp0_g0", "tp1_g0"]) + self.assertEqual(groups["full"], + ["tp0_g0", "tp0_g1", "tp0_g2", + "tp1_g0", "tp1_g1", "tp1_g2"]) + + +class TestStateCompleteMask(unittest.TestCase): + """Which manager blocks have a materialized recurrent state is read from + vLLM's block table: a state group pointing at the null block (id 0) has + none. This mask is what start_write_cache announces per key.""" + + def _req(self, tables): + return RequestLedger(vllm_request=FakeRequest("r0", []), + block_ids_per_group=tables, + has_saved_block_num=0) + + def test_full_attention_is_always_complete(self): + conn = make_connector_scheduler(manager_block_size=16, num_groups=1) + req = self._req([[7, 0, 9]]) + self.assertEqual(conn._state_complete_mask(req, range(3)), + [True, True, True]) + + def test_null_state_blocks_are_incomplete(self): + conn = make_connector_scheduler(manager_block_size=16, num_groups=1, + num_state_groups=1) + # State table: blocks 0 and 2 are null (no state), block 1 is real. + req = self._req([[100, 101, 102], [0, 55, 0]]) + self.assertEqual(conn._state_complete_mask(req, range(3)), + [False, True, False]) + + def test_all_state_groups_must_have_state(self): + conn = make_connector_scheduler(manager_block_size=16, num_groups=1, + num_state_groups=2) + # Block 1 has a state in group 1 but not in group 2 -> incomplete. + req = self._req([[100, 101], [7, 8], [7, 0]]) + self.assertEqual(conn._state_complete_mask(req, range(2)), + [True, False]) + + def test_short_state_table_is_incomplete(self): + # A state table that does not reach the block cannot prove a state. + conn = make_connector_scheduler(manager_block_size=16, num_groups=1, + num_state_groups=1) + req = self._req([[100, 101], [55]]) + self.assertEqual(conn._state_complete_mask(req, range(2)), + [True, False]) + + +class TestExternalHitTruncation(unittest.TestCase): + """A hybrid request can only resume where the recurrent state ends, so an + external match must be cut back to the last state-complete block. The + attention KV of a longer prefix is worthless without that state.""" + + MBS = 16 + + def _matched(self, coverage, prompt_len=None, num_state_groups=1, + tp_size=1): + conn = make_scheduler_connector( + mbs=self.MBS, num_state_groups=num_state_groups, tp_size=tp_size, + locations=hybrid_locations(coverage, tp_size=tp_size, + num_state=num_state_groups)) + req = FakeRequest("r0", list(range(prompt_len or + (len(coverage) + 2) * self.MBS))) + matched, _ = conn.get_num_new_matched_tokens(req, 0) + if matched: + alloc(conn, req, matched, [[100 + i for i in range( + matched // self.MBS + 1)]]) + return conn, matched + + def test_truncates_to_last_state_complete_block(self): + conn, matched = self._matched([True, True, False, False]) + self.assertEqual(matched, 2 * self.MBS) + self.assertEqual(conn._waiting_to_load_requests[0].manager_block_idxes, + [0, 1]) + + def test_interior_gap_is_kept(self): + # Only the *end* of the match must carry state; earlier state-less + # blocks are fine (their state is never read). + conn, matched = self._matched([True, False, True, False]) + self.assertEqual(matched, 3 * self.MBS) + self.assertEqual(conn._waiting_to_load_requests[0].manager_block_idxes, + [0, 1, 2]) + + def test_no_state_anywhere_drops_the_match(self): + conn, matched = self._matched([False, False, False]) + self.assertEqual(matched, 0) + alloc(conn, FakeRequest("r0", list(range(5 * 16))), 0, [[100]]) + self.assertEqual(conn._waiting_to_load_requests, []) + + def test_all_complete_is_untouched(self): + conn, matched = self._matched([True, True, True]) + self.assertEqual(matched, 3 * self.MBS) + + def test_every_rank_must_have_the_state(self): + # tp2: block 1 has the state spec of rank 0 only -- rank 1 would read + # nothing, so the block cannot end the match. + conn = make_scheduler_connector(mbs=self.MBS, num_state_groups=1, + tp_size=2) + locs = hybrid_locations([True, True], tp_size=2, num_state=1) + locs[1]["location_specs"] = [ + s for s in locs[1]["location_specs"] if s["name"] != "tp1_g1"] + conn._location_query_manager.locations = locs + conn._location_query_manager.in_flight = False + req = FakeRequest("r0", list(range(6 * self.MBS))) + matched, _ = conn.get_num_new_matched_tokens(req, 0) + self.assertEqual(matched, self.MBS) + + def test_full_attention_never_truncates(self): + # Single group: the state specs do not exist, so coverage is uniform + # and the match must be untouched (no hit-rate regression). + conn = make_scheduler_connector(mbs=self.MBS, locations=make_locations(3)) + req = FakeRequest("r0", list(range(6 * self.MBS))) + matched, _ = conn.get_num_new_matched_tokens(req, 0) + self.assertEqual(matched, 3 * self.MBS) + + def test_truncation_runs_before_the_full_hit_cap(self): + # Prompt is exactly 3 blocks; coverage allows 3 but the last carries no + # state -> truncate to 2, and the full-hit cap then has nothing to drop. + conn, matched = self._matched([True, True, False], + prompt_len=3 * self.MBS) + self.assertEqual(matched, 2 * self.MBS) + + def test_full_hit_cap_retruncates_to_state_complete(self): + # The cap drops trailing blocks without looking at their coverage, so + # it can move the match end onto a state-less block: coverage + # [True, False, True] truncates to all 3, the cap (prompt == 3 blocks) + # drops block 2, and the new match end (block 1) has no state. The + # match must be re-truncated -- loading it would end a hybrid request + # on a state nobody wrote, unreportably (report_failures=False). + conn, matched = self._matched([True, False, True], + prompt_len=3 * self.MBS) + self.assertEqual(matched, self.MBS) + self.assertEqual(conn._waiting_to_load_requests[0].manager_block_idxes, + [0]) + end = max(conn._waiting_to_load_requests[0].manager_block_idxes) + self.assertTrue(conn._location_covers_states( + conn._waiting_to_load_requests[0].need_load_locations[end])) + + +class TestStartWriteCacheSpecGroups(unittest.TestCase): + """start_write_cache must tell the manager, per key, which specs the block + will really hold -- that is how "no state here" becomes visible instead of + being encoded as a successful write.""" + + def _conn(self, num_state_groups): + conn = make_scheduler_connector(mbs=16, num_state_groups=num_state_groups) + conn._extra_config = SimpleNamespace( + instance_id="inst", write_timeout_seconds=30) + conn._manager_client = MagicMock() + conn._manager_client.start_write_cache.return_value = { + "locations": [], "write_session_id": "sess"} + return conn + + def _request_sent(self, conn): + (req,), _ = conn._manager_client.start_write_cache.call_args + return req + + def test_hybrid_sends_per_key_group_names(self): + conn = self._conn(num_state_groups=1) + conn.start_save_kvcache_async("r0", list(range(48)), 3, + [True, False, True]) + self.assertEqual(self._request_sent(conn)["location_spec_group_names"], + ["full", "attn", "full"]) + + def test_full_attention_omits_group_names(self): + conn = self._conn(num_state_groups=0) + conn.start_save_kvcache_async("r0", list(range(32)), 2, [True, True]) + self.assertNotIn("location_spec_group_names", self._request_sent(conn)) + + def test_mask_length_is_checked(self): + conn = self._conn(num_state_groups=1) + with self.assertRaises(AssertionError): + conn.start_save_kvcache_async("r0", list(range(48)), 3, [True]) + + +# --------------------------------------------------------------------------- # +# parse_block_mask_to_save_indices +# --------------------------------------------------------------------------- # +class TestParseBlockMask(unittest.TestCase): + def setUp(self): + self.conn = make_connector_scheduler(manager_block_size=16) + + def test_offset_branch(self): + resp = {"block_mask": {"offset": 2}} + self.assertEqual( + self.conn.parse_block_mask_to_save_indices(resp, 5), [2, 3, 4]) + + def test_offset_zero(self): + resp = {"block_mask": {"offset": 0}} + self.assertEqual( + self.conn.parse_block_mask_to_save_indices(resp, 3), [0, 1, 2]) + + def test_bool_masks_branch(self): + resp = {"block_mask": {"bool_masks": {"values": [True, False, True, False]}}} + self.assertEqual( + self.conn.parse_block_mask_to_save_indices(resp, 4), [1, 3]) + + def test_missing_mask(self): + self.assertEqual(self.conn.parse_block_mask_to_save_indices({}, 3), []) + + +# --------------------------------------------------------------------------- # +# _parse_groups +# --------------------------------------------------------------------------- # +class TestParseGroups(unittest.TestCase): + def _kv_cache_config(self, groups): + return SimpleNamespace(kv_cache_groups=groups) + + def _parse(self, groups, mbs): + return parse_groups(self._kv_cache_config(groups), mbs) + + def _attn_group(self, layers, block_size=16, page_size_bytes=32768, + page_size_padded=None): + from vllm.v1.kv_cache_interface import FullAttentionSpec + return SimpleNamespace( + layer_names=layers, + kv_cache_spec=FullAttentionSpec(block_size, page_size_bytes, + page_size_padded=page_size_padded)) + + def _mamba_group(self, layers, block_size=528, page_size_bytes=1024): + from vllm.v1.kv_cache_interface import MambaSpec + return SimpleNamespace( + layer_names=layers, + kv_cache_spec=MambaSpec(block_size, page_size_bytes)) + + def test_full_attention_single_group(self): + mbs = 32 + metas = self._parse( + [self._attn_group(["l0", "l1"], block_size=16, page_size_bytes=32768)], mbs) + self.assertEqual(len(metas), 1) + m = metas[0] + self.assertIsInstance(m, AttentionGroupMeta) + self.assertEqual(m.group_idx, 0) + self.assertEqual(m.block_size, 16) + # per_token = 32768 // 16 = 2048; per_block = 2048 * 32 (manager) * 2 layers + self.assertEqual(m.per_block_bytes, 2048 * 32 * 2) + + def test_pure_mamba_is_refused_before_init(self): + # A mamba-only model would otherwise reach register_kv_caches and + # die on the 'first attention tensor' lookup with an obscure + # StopIteration; refuse it explicitly in parse_groups instead. + with self.assertRaisesRegex( + NotImplementedError, "pure-mamba / attention-free models"): + self._parse([self._mamba_group(["m0"])], mbs=528) + + def test_hybrid_multi_group(self): + mbs = 528 + metas = self._parse([ + self._mamba_group(["m0", "m1"], page_size_bytes=1000), + self._mamba_group(["m2"], page_size_bytes=2000), + self._attn_group(["a0"], block_size=528, page_size_bytes=528 * 64), + ], mbs) + self.assertEqual([m.group_idx for m in metas], [0, 1, 2]) + self.assertEqual([type(m).__name__ for m in metas], + ['StateGroupMeta', 'StateGroupMeta', 'AttentionGroupMeta']) + self.assertEqual(metas[0].per_block_bytes, 1000 * 2) # page * layers + self.assertEqual(metas[1].per_block_bytes, 2000) + self.assertEqual(metas[2].per_block_bytes, 64 * 528) # per_token * mbs + + def test_eagle_group_skipped(self): + mbs = 16 + eagle = self._attn_group(["drafter"]) + eagle.is_eagle_group = True + metas = self._parse( + [eagle, self._attn_group(["a0"])], mbs) + self.assertEqual(len(metas), 1) + self.assertEqual(metas[0].layer_names, ["a0"]) + self.assertEqual(metas[0].group_idx, 1) # group_idx keeps vLLM numbering + + def test_padded_attention_page_uses_compact_size(self): + # page_size_padded inflates spec.page_size_bytes with an allocation + # gap the gather kernel never copies; per_block_bytes must come from + # the compact real_page_size_bytes. + mbs = 16 + metas = self._parse( + [self._attn_group(["l0", "l1"], block_size=16, + page_size_bytes=32768, page_size_padded=40960)], mbs) + # per_token = 32768 // 16 = 2048 (not 40960 // 16 = 2560). + self.assertEqual(metas[0].per_block_bytes, 2048 * 16 * 2) + + def test_padded_attention_without_compact_size_raises(self): + # A padded spec that exposes no real_page_size_bytes cannot be sized + # correctly -- must refuse, not silently over-allocate. + mbs = 16 + group = self._attn_group(["l0"], block_size=16, + page_size_bytes=32768, page_size_padded=40960) + del group.kv_cache_spec.real_page_size_bytes + with self.assertRaises(NotImplementedError): + self._parse([group], mbs) + + def test_unpadded_attention_without_compact_size_falls_back(self): + # No padding + no real_page_size_bytes: page_size_bytes is already + # compact, use it. + mbs = 16 + group = self._attn_group(["l0"], block_size=16, page_size_bytes=32768) + del group.kv_cache_spec.real_page_size_bytes + metas = self._parse([group], mbs) + self.assertEqual(metas[0].per_block_bytes, 2048 * 16) + + def test_windowed_attention_spec_rejected(self): + # vLLM can merge SWA / chunked-attention layers into a + # FullAttentionSpec that keeps sliding_window / attention_chunk_size + # set; such blocks are not full-prefix KV and must be refused. + for window_field in ("sliding_window", "attention_chunk_size"): + with self.subTest(field=window_field): + mbs = 16 + group = self._attn_group(["l0"]) + setattr(group.kv_cache_spec, window_field, 1024) + with self.assertRaises(NotImplementedError) as cm: + self._parse([group], mbs) + self.assertIn(window_field, str(cm.exception)) + + def test_windowed_fields_none_accepted(self): + # Real FullAttentionSpec objects carry the fields as None; that is the + # ordinary full-attention case and must still parse. + mbs = 16 + group = self._attn_group(["l0"]) + group.kv_cache_spec.sliding_window = None + group.kv_cache_spec.attention_chunk_size = None + metas = self._parse([group], mbs) + self.assertEqual(len(metas), 1) + + def test_unsupported_spec_raises(self): + mbs = 16 + bad = SimpleNamespace(layer_names=["x"], kv_cache_spec=object()) + with self.assertRaises(NotImplementedError): + self._parse([bad], mbs) + + def test_no_usable_groups_is_refused(self): + mbs = 16 + with self.assertRaisesRegex(NotImplementedError, + "no usable kv cache groups"): + self._parse([], mbs) + + +# --------------------------------------------------------------------------- # +# Skipped (EAGLE/MTP drafter) groups: block tables indexed by vLLM group idx +# --------------------------------------------------------------------------- # +class TestSkippedGroupIndexing(unittest.TestCase): + """When parse_groups skips a group (EAGLE/MTP drafter), its block table is + still present in block_ids_per_group / all_block_ids at its vLLM group + index. Consumers must index by GroupMeta.group_idx, never assume the + transferred groups start at 0 or include every table.""" + + MBS = 16 + + def _skipped_group0_connector(self): + """ConnectorScheduler where vLLM group 0 is a skipped drafter and group 1 + is the transferred attention group.""" + conn = make_connector_scheduler(manager_block_size=self.MBS) + conn._group_metas = [AttentionGroupMeta( + group_idx=1, layer_names=["a0"], + block_size=self.MBS, per_block_bytes=0)] + conn._num_groups = 1 + return conn + + def test_num_allocated_blocks_ignores_skipped_group(self): + conn = self._skipped_group0_connector() + ledger = RequestLedger( + vllm_request=FakeRequest("r0", list(range(64))), + # Drafter table (group 0) lags with 1 block; attention has 4. + block_ids_per_group=[[100], [200, 201, 202, 203]], + has_saved_block_num=0) + self.assertEqual(conn._num_allocated_blocks(ledger), 4) + + def test_num_allocated_blocks_still_mins_transferred_groups(self): + # Two transferred groups (1 and 2), one skipped drafter (0): min is + # taken over the transferred ones only. + conn = make_connector_scheduler(manager_block_size=self.MBS) + conn._group_metas = [ + AttentionGroupMeta(group_idx=1, layer_names=["a0"], + block_size=self.MBS, per_block_bytes=0), + StateGroupMeta(group_idx=2, layer_names=["m0"], + block_size=self.MBS, per_block_bytes=0, + page_size_bytes=0), + ] + conn._num_groups = 2 + ledger = RequestLedger( + vllm_request=FakeRequest("r0", []), + block_ids_per_group=[[9], [1, 2, 3], [4, 5]], + has_saved_block_num=0) + self.assertEqual(conn._num_allocated_blocks(ledger), 2) + + def test_num_allocated_blocks_empty(self): + conn = self._skipped_group0_connector() + ledger = RequestLedger(vllm_request=FakeRequest("r0", []), + block_ids_per_group=[], + has_saved_block_num=0) + self.assertEqual(conn._num_allocated_blocks(ledger), 0) + + def test_single_group_reports_failures_against_group0_table(self): + # Exactly one vLLM block table: the transferred group is 0 and the + # failure report maps manager blocks into its table. + conn = make_connector(manager_block_size=self.MBS) + conn._extra_config = SimpleNamespace(block_per_load_task=8) + conn._data_transfer = MagicMock() + conn._plan_group_transfers = MagicMock(return_value=None) + meta = TairKvCacheConnectorMetadata(epoch=0) + meta.add_load_request(LoadRequest( + req_id="r0", manager_block_idxes=[0, 1], + need_load_locations=[{"location_specs": []}] * 2, + all_block_ids=[[10, 11, 12]])) + conn.start_load_kv(MagicMock(), meta) + args, kwargs = conn._data_transfer.create_load_done_callback.call_args + self.assertEqual(args[3], [10, 11]) # report_ids from group 0's table + self.assertTrue(kwargs["report_failures"]) + + def test_skipped_drafter_shape_disables_failure_reporting(self): + # Two vLLM block tables (skipped drafter group 0 + transferred + # attention group 1): upstream recovery unpacks exactly one table, + # so reporting would crash the scheduler. Failures must NOT be + # reported for this shape even though the model is attention-only. + conn = make_connector(manager_block_size=self.MBS) + conn._group_metas = [AttentionGroupMeta( + group_idx=1, layer_names=["a0"], + block_size=self.MBS, per_block_bytes=0)] + conn._num_groups = 1 + conn._extra_config = SimpleNamespace(block_per_load_task=8) + conn._data_transfer = MagicMock() + conn._plan_group_transfers = MagicMock(return_value=None) + meta = TairKvCacheConnectorMetadata(epoch=0) + meta.add_load_request(LoadRequest( + req_id="r0", manager_block_idxes=[0, 1], + need_load_locations=[{"location_specs": []}] * 2, + # Group 0 (drafter) has a lagging 1-entry table. + all_block_ids=[[999], [10, 11]])) + conn.start_load_kv(MagicMock(), meta) + args, kwargs = conn._data_transfer.create_load_done_callback.call_args + self.assertEqual(args[3], []) # no report ids when not reporting + self.assertFalse(kwargs["report_failures"]) + + def test_multi_attention_group_shape_disables_failure_reporting(self): + # Two transferred attention groups (e.g. unmerged sw + full): two + # vLLM block tables, attention-only, still breaks the single-table + # recovery unpack -- must not be reported. + conn = make_connector(manager_block_size=self.MBS, num_groups=2) + conn._extra_config = SimpleNamespace(block_per_load_task=8) + conn._data_transfer = MagicMock() + conn._plan_group_transfers = MagicMock(return_value=None) + meta = TairKvCacheConnectorMetadata(epoch=0) + meta.add_load_request(LoadRequest( + req_id="r0", manager_block_idxes=[0], + need_load_locations=[{"location_specs": []}], + all_block_ids=[[10], [20]])) + conn.start_load_kv(MagicMock(), meta) + args, kwargs = conn._data_transfer.create_load_done_callback.call_args + self.assertEqual(args[3], []) + self.assertFalse(kwargs["report_failures"]) + + def test_hybrid_shape_disables_failure_reporting(self): + # Attention + mamba: two vLLM block tables, the original hybrid case. + conn = make_connector(manager_block_size=self.MBS, + num_groups=1, num_state_groups=1) + conn._extra_config = SimpleNamespace(block_per_load_task=8) + conn._data_transfer = MagicMock() + conn._plan_group_transfers = MagicMock(return_value=None) + meta = TairKvCacheConnectorMetadata(epoch=0) + meta.add_load_request(LoadRequest( + req_id="r0", manager_block_idxes=[0], + need_load_locations=[{"location_specs": []}], + all_block_ids=[[10], [20]])) + conn.start_load_kv(MagicMock(), meta) + args, kwargs = conn._data_transfer.create_load_done_callback.call_args + self.assertEqual(args[3], []) + self.assertFalse(kwargs["report_failures"]) + + def test_lone_state_group_refuses_to_report(self): + # A single vLLM block table that is not attention (pure mamba): the + # token-granular recovery math upstream cannot consume state block + # ids; the worker must fail loudly instead of reporting nonsense. + conn = make_connector(manager_block_size=self.MBS) + conn._group_metas = [StateGroupMeta( + group_idx=0, layer_names=["m0"], + block_size=self.MBS, per_block_bytes=0, page_size_bytes=0)] + conn._num_groups = 1 + conn._extra_config = SimpleNamespace(block_per_load_task=8) + conn._data_transfer = MagicMock() + conn._plan_group_transfers = MagicMock(return_value=None) + meta = TairKvCacheConnectorMetadata(epoch=0) + meta.add_load_request(LoadRequest( + req_id="r0", manager_block_idxes=[0], + need_load_locations=[{"location_specs": []}], + all_block_ids=[[10]])) + with self.assertRaises(AssertionError): + conn.start_load_kv(MagicMock(), meta) + + +# --------------------------------------------------------------------------- # +# build_connector_meta +# --------------------------------------------------------------------------- # +class TestBuildConnectorMeta(unittest.TestCase): + MBS = 16 + + def _new_request(self, conn, req_id, num_tokens, num_blocks, + num_locations=0): + """Simulate the scheduler flow for a fresh request: query, alloc, then + one build_connector_meta step.""" + conn._location_query_manager.locations = make_locations(num_locations) + conn._location_query_manager.in_flight = False + req = FakeRequest(req_id, list(range(num_tokens))) + matched, _ = conn.get_num_new_matched_tokens(req, 0) + block_ids = [list(range(100, 100 + num_blocks))] + conn.update_state_after_alloc( + req, SimpleNamespace(get_block_ids=lambda: block_ids), matched) + out = fake_scheduler_output( + new_reqs=[SimpleNamespace(req_id=req_id, block_ids=block_ids)]) + return req, conn.build_connector_meta(out) + + def test_new_request_full_state(self): + conn = make_scheduler_connector(mbs=self.MBS) + req, meta = self._new_request(conn, "r0", 40, 3) + # The ledger absorbed the first allocation's block table. + self.assertEqual(conn._tracked["r0"].block_ids_per_group, + [[100, 101, 102]]) + # 40 tokens / 3 blocks -> min(40, 48)//16 = 2 blocks to save. + conn._http_executor.submit.assert_called_once() + args = conn._http_executor.submit.call_args[0] + # The per-block state-completeness mask is computed here, in the + # scheduler loop, and handed to the http thread: it is read off vLLM's + # block table, which later steps mutate. + self.assertEqual(args[1:], ("r0", list(range(32)), 2, [True, True])) + self.assertEqual(conn._tracked["r0"].has_saved_block_num, 2) + + def test_load_request_emitted_after_alloc(self): + conn = make_scheduler_connector(mbs=self.MBS) + req, meta = self._new_request(conn, "r0", 40, 3, num_locations=2) + self.assertEqual(len(meta.to_load_requests), 1) + lr = meta.to_load_requests[0] + self.assertEqual(lr.manager_block_idxes, [0, 1]) + self.assertEqual(lr.all_block_ids, [[100, 101, 102]]) + # Externally hit blocks are not re-saved. + conn._http_executor.submit.assert_not_called() + + def test_cached_delta_with_and_without_new_blocks(self): + conn = make_scheduler_connector(mbs=self.MBS) + req, _ = self._new_request(conn, "r0", 40, 3) + # Step 2: 8 decode tokens, no new blocks (PR #23262: may be None). + req.output_token_ids = list(range(1000, 1008)) + out = fake_scheduler_output( + cached_req_ids=["r0"], num_scheduled={"r0": 8}, new_block_ids=[None]) + conn.build_connector_meta(out) + # Step 3: 2 more tokens with a new block -> table grows. + req.output_token_ids = list(range(1000, 1010)) + out = fake_scheduler_output( + cached_req_ids=["r0"], num_scheduled={"r0": 2}, + new_block_ids=[[[103]]]) + conn.build_connector_meta(out) + self.assertEqual(conn._tracked["r0"].block_ids_per_group, + [[100, 101, 102, 103]]) + + def _preempted_step(self, conn, req, use_legacy): + kwargs = dict(cached_req_ids=["r0"], num_scheduled={"r0": 0}, + new_block_ids=[[[200, 201]]]) + if use_legacy: + kwargs["legacy_resumed"] = [True] + else: + kwargs["resumed_req_ids"] = {"r0"} + return conn.build_connector_meta(fake_scheduler_output(**kwargs)) + + def test_resumed_from_preemption_both_interfaces(self): + for use_legacy in (False, True): + with self.subTest(legacy=use_legacy): + conn = make_scheduler_connector(mbs=self.MBS) + req, _ = self._new_request(conn, "r0", 40, 3) + self._preempted_step(conn, req, use_legacy) + # Resume replaces (not extends) the block table. + self.assertEqual(conn._tracked["r0"].block_ids_per_group, + [[200, 201]]) + + def test_resumed_with_none_new_blocks_keeps_the_recorded_table(self): + # A resumed request whose new_block_ids is None (upstream + # get_block_ids with allow_none=True: no group got fresh blocks) + # must not crash the table replace -- and must not touch the + # recorded table either. + for use_legacy in (False, True): + with self.subTest(legacy=use_legacy): + conn = make_scheduler_connector(mbs=self.MBS) + req, _ = self._new_request(conn, "r0", 40, 3) + before = [list(t) for t in + conn._tracked["r0"].block_ids_per_group] + kwargs = dict(cached_req_ids=["r0"], num_scheduled={"r0": 0}, + new_block_ids=[None]) + if use_legacy: + kwargs["legacy_resumed"] = [True] + else: + kwargs["resumed_req_ids"] = {"r0"} + conn.build_connector_meta(fake_scheduler_output(**kwargs)) + self.assertEqual(conn._tracked["r0"].block_ids_per_group, + before) + + def test_save_threshold_grows_incrementally(self): + conn = make_scheduler_connector(mbs=self.MBS) + req, _ = self._new_request(conn, "r0", 40, 3) # saved 2 blocks + conn._http_executor.submit.reset_mock() + # 8 more tokens -> 48 total, table full at 3 blocks -> third block saves. + req.output_token_ids = list(range(1000, 1008)) + out = fake_scheduler_output( + cached_req_ids=["r0"], num_scheduled={"r0": 8}, new_block_ids=[[[103]]]) + conn.build_connector_meta(out) + args = conn._http_executor.submit.call_args[0] + self.assertEqual(args[3], 3) # target_save_num + self.assertEqual(conn._tracked["r0"].has_saved_block_num, 3) + + def test_save_threshold_counts_key_material_not_scheduled(self): + # During real decode all_token_ids lags the scheduled count by one: + # the token scheduled in this step is appended to it only once + # sampled. A block whose last token -- and thus cache key -- is not + # known yet must not be announced (the manager would return one + # location fewer than announced and the worker would drop the whole + # session). Regression test: the scheduled-count-derived block + # count (a since-removed ledger field) did exactly that. + conn = make_scheduler_connector(mbs=self.MBS) + req, _ = self._new_request(conn, "r0", 40, 3) # saved 2 blocks + conn._http_executor.submit.reset_mock() + # 8 more tokens scheduled (48 total) but only 7 appended to the + # request: block 2's last token id is still unknown. + req.output_token_ids = list(range(1000, 1007)) + out = fake_scheduler_output( + cached_req_ids=["r0"], num_scheduled={"r0": 8}, new_block_ids=[[[103]]]) + conn.build_connector_meta(out) + conn._http_executor.submit.assert_not_called() + self.assertEqual(conn._tracked["r0"].has_saved_block_num, 2) + + # The lagging token lands one step later: now the third block saves, + # with a token list that really holds 48 ids. + req.output_token_ids = list(range(1000, 1008)) + out = fake_scheduler_output( + cached_req_ids=["r0"], num_scheduled={"r0": 1}, new_block_ids=[[[]]]) + conn.build_connector_meta(out) + args = conn._http_executor.submit.call_args[0] + self.assertEqual(args[3], 3) # target_save_num + self.assertEqual(len(args[2]), 3 * self.MBS) # token_ids sent + self.assertEqual(conn._tracked["r0"].has_saved_block_num, 3) + + def test_save_request_drain_and_finish_paths(self): + conn = make_scheduler_connector(mbs=self.MBS) + req, _ = self._new_request(conn, "r0", 40, 3) + state = conn._tracked["r0"] + self.assertEqual(state.scheduled_saving_count, 1) + + # Finish while the save is still in flight: request must stay alive. + keep, extra = conn.request_finished(req, []) + self.assertTrue(keep) + self.assertTrue(state.need_report_after_saving_finished) + self.assertIn("r0", conn._tracked) + + # The async save lands: drained into to_save_requests and, because the + # request already finished, a FinishRequest is emitted and state dropped. + with conn._waiting_to_save_requests_lock: + conn._waiting_to_save_requests.append( + SaveRequest("r0", make_locations(2), [0, 1], "sess")) + meta = conn.build_connector_meta(fake_scheduler_output()) + self.assertEqual(len(meta.to_save_requests), 1) + self.assertEqual([f.req_id for f in meta.to_finish_requests], ["r0"]) + self.assertNotIn("r0", conn._tracked) + + def test_request_finished_when_saves_landed(self): + conn = make_scheduler_connector(mbs=self.MBS) + req, _ = self._new_request(conn, "r0", 40, 3) + with conn._waiting_to_save_requests_lock: + conn._waiting_to_save_requests.append( + SaveRequest("r0", make_locations(2), [0, 1], "sess")) + conn.build_connector_meta(fake_scheduler_output()) + keep, extra = conn.request_finished(req, []) + self.assertTrue(keep) + self.assertNotIn("r0", conn._tracked) + meta = conn.build_connector_meta(fake_scheduler_output()) + self.assertEqual([f.req_id for f in meta.to_finish_requests], ["r0"]) + + # ------------------------------------------------------------------ # + # Hit accounting contract (kv_transfer_params) + # ------------------------------------------------------------------ # + # These counts travel to clients through vLLM's kv_transfer_params + # (EngineCoreOutput -> RequestOutput -> the OpenAI response) and were + # silently dropped in an earlier refactor (found only by human review); + # the contract is pinned here so a future refactor cannot remove it + # without a test failing. + def test_finish_reports_hit_accounting(self): + conn = make_scheduler_connector(mbs=self.MBS) + req, _ = self._new_request(conn, "r0", 4 * self.MBS + 5, 4, + num_locations=3) + # No saves in flight: finish reports the accounting immediately. + conn._tracked["r0"].scheduled_saving_count = 0 + conn._tracked["r0"].sent_saving_count = 0 + keep, extra = conn.request_finished(req, []) + self.assertTrue(keep) + self.assertEqual(extra, { + "local_matched_token_num": 0, + "remote_matched_token_num": 3 * self.MBS, + }) + self.assertIsInstance(extra["local_matched_token_num"], int) + self.assertIsInstance(extra["remote_matched_token_num"], int) + + def test_finish_reports_hit_accounting_with_saves_in_flight(self): + # The delayed-finish branch must report the same accounting: vLLM + # consumes kv_transfer_params from whichever output finishes the + # request. + conn = make_scheduler_connector(mbs=self.MBS) + req, _ = self._new_request(conn, "r0", 4 * self.MBS + 5, 4, + num_locations=3) + keep, extra = conn.request_finished(req, []) + self.assertTrue(keep) + self.assertEqual(extra, { + "local_matched_token_num": 0, + "remote_matched_token_num": 3 * self.MBS, + }) + + def test_hit_accounting_follows_the_last_match_answer(self): + # A re-ask with a grown local hit overwrites the accounting (the + # last answer is what the request runs with), and the burned path + # zeroes the remote half while keeping the local hit. + conn = make_scheduler_connector(mbs=self.MBS) + req = FakeRequest("r0", list(range(4 * self.MBS + 5))) + conn._location_query_manager.locations = make_locations(3) + conn._location_query_manager.in_flight = False + conn.get_num_new_matched_tokens(req, 0) # answer: 3 blocks + self.assertEqual(conn._tracked["r0"].remote_matched_token_num, + 3 * self.MBS) + # Re-ask at a grown offset: local hit counted, remote re-answered. + conn._location_query_manager.locations = make_locations(2) + conn._location_query_manager.in_flight = False + conn.get_num_new_matched_tokens(req, self.MBS) + self.assertEqual(conn._tracked["r0"].local_matched_token_num, self.MBS) + self.assertEqual(conn._tracked["r0"].remote_matched_token_num, + 2 * self.MBS) + + def test_hit_accounting_burned_match_zeroes_remote(self): + conn = make_scheduler_connector( + mbs=self.MBS, num_groups=1, num_state_groups=1, + locations=hybrid_locations([True, True])) + req = FakeRequest("r0", list(range(4 * self.MBS + 5))) + conn.get_num_new_matched_tokens(req, 0) + alloc(conn, req, 2 * self.MBS, [[100, 101], [50, 51]]) + # The re-ask after the burned match reports zero remote, and the + # local hit (the offset vLLM re-asked at) stands. + matched, _ = conn.get_num_new_matched_tokens(req, self.MBS) + self.assertEqual(matched, 0) + self.assertEqual(conn._tracked["r0"].local_matched_token_num, self.MBS) + self.assertEqual(conn._tracked["r0"].remote_matched_token_num, 0) + + def test_canceled_save_unknown_request_no_crash(self): + # Cancellations arrive from http_executor threads and may race request + # teardown; an unknown req_id must be skipped, not KeyError. + conn = make_scheduler_connector(mbs=self.MBS) + with conn._canceled_save_request_ids_lock: + conn._canceled_save_request_ids.append("ghost") + conn.build_connector_meta(fake_scheduler_output()) # must not raise + + def test_canceled_save_finishes_request(self): + conn = make_scheduler_connector(mbs=self.MBS) + req, _ = self._new_request(conn, "r0", 40, 3) + conn.request_finished(req, []) # save in flight -> delayed finish + with conn._canceled_save_request_ids_lock: + conn._canceled_save_request_ids.append("r0") + meta = conn.build_connector_meta(fake_scheduler_output()) + self.assertEqual([f.req_id for f in meta.to_finish_requests], ["r0"]) + self.assertNotIn("r0", conn._tracked) + + +if __name__ == "__main__": + unittest.main() diff --git a/kv_cache_manager/py_connector/test/vllm_stubs.py b/kv_cache_manager/py_connector/test/vllm_stubs.py new file mode 100644 index 000000000..b587ad531 --- /dev/null +++ b/kv_cache_manager/py_connector/test/vllm_stubs.py @@ -0,0 +1,308 @@ +"""Shared test stubs: make ``v1_connector`` importable without vLLM/CUDA/pybind. + +``v1_connector`` imports vLLM, the compiled ``kvcm_py_client`` and several +third-party runtime deps (torch, triton, orjson, zmq, requests) at module +level. For pure-logic unit tests we register lightweight stand-ins in +``sys.modules`` *before* the first import, then build connector instances via +``__new__`` with only the attributes the code under test reads. No production +module is modified; real modules are preferred whenever they are importable +(e.g. on a dev machine with a full vLLM venv). +""" + +import importlib.util +import json +import sys +import types +from typing import Optional +from unittest.mock import MagicMock + +#: Modules this run replaced with stand-ins (empty when the real deps are +#: importable). Tests that need real behaviour (e.g. an actual torch tensor) +#: skip on membership instead of failing against MagicMocks. +STUBBED: set = set() + + +def _module(name: str) -> types.ModuleType: + mod = sys.modules.get(name) + if mod is None: + mod = types.ModuleType(name) + sys.modules[name] = mod + return mod + + +def _importable(name: str) -> bool: + try: + return importlib.util.find_spec(name) is not None + except (ImportError, ValueError): + return False + + +def _stub_third_party(): + """Register stand-ins for third-party deps missing from the environment + (the open-source CI runs these tests without torch/triton/orjson/zmq/ + requests installed). Real modules always win.""" + # Pure-attribute deps: a MagicMock module is enough because the pure-logic + # tests never execute tensor/socket/http work at module import time. + for name in ("torch", "triton", "triton.language", "zmq"): + if name not in sys.modules and not _importable(name): + sys.modules[name] = MagicMock(__name__=name) + STUBBED.add(name.split(".")[0]) + + # requests needs real exception classes, not a MagicMock: production code + # subclasses them (manager_client defines KvCacheManagerHTTPError( + # requests.HTTPError, AssertionError)), and a MagicMock attribute cannot + # serve as a base class (metaclass conflict at class-creation time). The + # stand-in mirrors requests' own hierarchy (RequestException(IOError)); + # the functional entry points stay MagicMocks because the tests patch + # them (requests.post / requests.Session) before any call. + if "requests" not in sys.modules and not _importable("requests"): + req = _module("requests") + STUBBED.add("requests") + + class _RequestException(IOError): + """Stand-in for requests.exceptions.RequestException.""" + + class _HTTPError(_RequestException): + pass + + class _ConnectionError(_RequestException): + pass + + class _Timeout(_RequestException): + pass + + req.RequestException = _RequestException + req.HTTPError = _HTTPError + req.ConnectionError = _ConnectionError + req.Timeout = _Timeout + req.Session = MagicMock(name="requests.Session") + req.post = MagicMock(name="requests.post") + req.get = MagicMock(name="requests.get") + + # orjson is used functionally (CoordinateMsgSerializer round trips), so + # the stand-in must actually (de)serialize; stdlib json handles the + # dataclass payloads via __dict__. + if "orjson" not in sys.modules and not _importable("orjson"): + orjson = _module("orjson") + orjson.dumps = lambda obj: json.dumps( + obj, default=lambda o: o.__dict__).encode() + orjson.loads = json.loads + + +def _install_stubs(): + _stub_third_party() + existing = sys.modules.get("vllm") + if existing is not None: + # Either our stub is already in place or the real vLLM is importable; + # in both cases the connector import will succeed as-is. + return + + # ---- kv_cache_manager.client.pybind (compiled extension) ---- + pybind = _module("kv_cache_manager.client.pybind") + kvcm_py_client = MagicMock() + kvcm_py_client.ClientErrorCode.ER_OK = 0 + pybind.kvcm_py_client = kvcm_py_client + + # ---- kv_cache_manager.py_connector.common._version_info (generated) ---- + version = _module("kv_cache_manager.py_connector.common._version_info") + version.FULL_VERSION = "0.0.0-test" + version.GIT_COMMIT = "test" + version.BUILD_TIME = "test" + + # ---- vllm ---- + vllm = _module("vllm") + vllm._kvcm_test_stub = True + + config = _module("vllm.config") + config.VllmConfig = MagicMock + vllm.config = config + + distributed = _module("vllm.distributed") + distributed.get_tensor_model_parallel_rank = lambda: 0 + vllm.distributed = distributed + _module("vllm.distributed.kv_transfer") + _module("vllm.distributed.kv_transfer.kv_connector") + _module("vllm.distributed.kv_transfer.kv_connector.v1") + base = _module("vllm.distributed.kv_transfer.kv_connector.v1.base") + + class KVConnectorRole: + SCHEDULER = 0 + WORKER = 1 + + class KVConnectorMetadata: + pass + + class KVConnectorBase_V1: + def __init__(self, vllm_config, role, kv_cache_config=None): + self._connector_metadata = None + + def _get_connector_metadata(self): + return self._connector_metadata + + class SupportsHMA: + pass + + base.KVConnectorBase_V1 = KVConnectorBase_V1 + base.KVConnectorMetadata = KVConnectorMetadata + base.KVConnectorRole = KVConnectorRole + base.SupportsHMA = SupportsHMA + + utils = _module("vllm.utils") + torch_utils = _module("vllm.utils.torch_utils") + torch_utils.get_kv_cache_torch_dtype = MagicMock() + network_utils = _module("vllm.utils.network_utils") + network_utils.get_ip = lambda: "127.0.0.1" + utils.torch_utils = torch_utils + utils.network_utils = network_utils + + v1 = _module("vllm.v1") + kv_cache_interface = _module("vllm.v1.kv_cache_interface") + + class FullAttentionSpec: + def __init__(self, block_size, page_size_bytes, page_size_padded=None): + self.block_size = block_size + self.page_size_padded = page_size_padded + self.real_page_size_bytes = page_size_bytes + # Mirror vLLM's AttentionSpec: page_size_bytes returns the padded + # size when padding is set. + self.page_size_bytes = (page_size_padded if page_size_padded + is not None else page_size_bytes) + + class MambaSpec: + def __init__(self, block_size, page_size_bytes): + self.block_size = block_size + self.page_size_bytes = page_size_bytes + + kv_cache_interface.FullAttentionSpec = FullAttentionSpec + kv_cache_interface.MambaSpec = MambaSpec + + _module("vllm.v1.core") + sched = _module("vllm.v1.core.sched") + output = _module("vllm.v1.core.sched.output") + output.SchedulerOutput = MagicMock + sched.output = output + + outputs = _module("vllm.v1.outputs") + outputs.KVConnectorOutput = MagicMock + v1.kv_cache_interface = kv_cache_interface + v1.outputs = outputs + + +_install_stubs() + +# Import after stubs are in place. +from kv_cache_manager.py_connector.vllm.vllm_common import ( # noqa: E402 + AttentionGroupMeta, GroupMeta, StateGroupMeta) +from kv_cache_manager.py_connector.vllm.connector_scheduler import ConnectorScheduler # noqa: E402 +from kv_cache_manager.py_connector.vllm.connector_worker import ConnectorWorker # noqa: E402 + + +def _make_group_metas(num_groups: int, num_state_groups: int, + block_size: int) -> list: + """Attention groups first, then mamba-style state groups; group_idx is the + vLLM group index (what block tables are indexed by).""" + return [ + AttentionGroupMeta(group_idx=i, layer_names=[f"l{i}"], + block_size=block_size, per_block_bytes=0) + for i in range(num_groups) + ] + [ + StateGroupMeta(group_idx=num_groups + i, layer_names=[f"m{i}"], + block_size=block_size, per_block_bytes=0, + page_size_bytes=0) + for i in range(num_state_groups) + ] + + +def make_connector(manager_block_size: int = 16, + vllm_block_size: Optional[int] = None, + num_groups: int = 1, + num_state_groups: int = 0, + tp_size: int = 1) -> ConnectorWorker: + """Build a bare ConnectorWorker (no __init__) with the minimal state used by + the pure translation logic under test (block index translation, transfer + group building). + + ``num_state_groups`` appends that many mamba-style (non-attention) groups + after the ``num_groups`` attention groups, which is what turns on the + hybrid-only logic (spec groups, per-block state completeness, hit + truncation).""" + conn = ConnectorWorker.__new__(ConnectorWorker) + conn._manager_block_size = manager_block_size + conn._vllm_block_size = vllm_block_size or manager_block_size + conn._tp_size = tp_size + conn._tp_rank = 0 + conn._self_spec_names = {} + conn._device = "cpu" + conn._group_metas = _make_group_metas( + num_groups, num_state_groups, conn._vllm_block_size) + conn._num_groups = len(conn._group_metas) + conn._state_group_idxs = [m.group_idx for m in conn._group_metas + if isinstance(m, StateGroupMeta)] + return conn + + +class FakeLocationQueries: + """Test double for LocationQueryManager: same produce/consume contract. + + ``locations`` is the manager's answer (None = in flight). The match + hook's clamped result (store_result) is what the allocation consumes; + the offset recorded at get time travels with it.""" + + def __init__(self, locations=None): + self.locations = locations + self.in_flight = locations is None + self._stored = None + self.last_computed_blocks = 0 + + def get_locations_for_query(self, request, computed_blocks): + self.last_computed_blocks = computed_blocks + return None if self.in_flight else list(self.locations) + + def store_result(self, req_id, locations): + self._stored = list(locations) + + def consume_locations(self, req_id): + if self.in_flight or self._stored is None: + return None + result, self._stored = self._stored, None + return result, self.last_computed_blocks + + def invalidate(self, req_id): + self._stored = None + + +def make_connector_scheduler(manager_block_size: int = 16, + vllm_block_size: Optional[int] = None, + num_groups: int = 1, + num_state_groups: int = 0, + tp_size: int = 1, + locations=None) -> ConnectorScheduler: + """Build a bare ConnectorScheduler (no __init__) with the scheduler-loop state + build_connector_meta and friends need, plus a FakeLocationQueries + answering ``locations`` (None means "still in flight").""" + from unittest.mock import MagicMock + core = ConnectorScheduler.__new__(ConnectorScheduler) + core._manager_block_size = manager_block_size + core._vllm_block_size = vllm_block_size or manager_block_size + core._tp_size = tp_size + core._group_metas = _make_group_metas( + num_groups, num_state_groups, core._vllm_block_size) + core._num_groups = len(core._group_metas) + core._state_group_idxs = [m.group_idx for m in core._group_metas + if isinstance(m, StateGroupMeta)] + core._epoch = 0 + core._tracked = {} + core._load_failed = set() + core._load_attempted = set() + core._waiting_to_load_requests = [] + import threading + core._waiting_to_save_requests_lock = threading.Lock() + core._waiting_to_save_requests = [] + core._waiting_to_finish_requests = [] + core._canceled_save_request_ids_lock = threading.Lock() + core._canceled_save_request_ids = [] + core._http_executor = MagicMock() + core._manager_client = MagicMock() + core._coordinator_client = MagicMock() + core._location_query_manager = FakeLocationQueries(locations) + return core diff --git a/kv_cache_manager/py_connector/vllm/BUILD b/kv_cache_manager/py_connector/vllm/BUILD index 2fc64b26c..23d957714 100644 --- a/kv_cache_manager/py_connector/vllm/BUILD +++ b/kv_cache_manager/py_connector/vllm/BUILD @@ -7,9 +7,12 @@ load("@python_platform//:platform.bzl", "python_platform") py_library( name = "vllm_connector", srcs = glob(["*.py"]), + visibility = ["//kv_cache_manager/py_connector:__subpackages__"], deps = ["//kv_cache_manager/py_connector/common:common", "//kv_cache_manager/py_connector/kernel:kernel", - "//kv_cache_manager/client/pybind:kvcm_py_client_lib"] + "//kv_cache_manager/client/pybind:kvcm_py_client_lib", + "@pip_cpu//pydantic", + "@pip_cpu//orjson"] ) @@ -39,6 +42,10 @@ py_wheel( # }, python_requires = ">=3.9", python_tag = python_abi(), + requires = [ + "pydantic>=2.11.4", + "orjson>=3", + ], stamp = 1, version = "{STABLE_KVCM_VERSION}+{STABLE_BUILD_TIMESTAMP}.{STABLE_GIT_COMMIT}", deps = [ diff --git a/kv_cache_manager/py_connector/vllm/config.py b/kv_cache_manager/py_connector/vllm/config.py index 1481e40a2..e4a0deb6c 100644 --- a/kv_cache_manager/py_connector/vllm/config.py +++ b/kv_cache_manager/py_connector/vllm/config.py @@ -1,38 +1,89 @@ -from typing import Any - - -class TairKvCacheConnectorExtraConfig: - def __init__(self, extra_config: dict[str, Any]): - self.manager_uri: str = extra_config["manager_uri"] - self.coordinator_base_port: int = extra_config["coordinator_base_port"] - self.instance_group: str = extra_config["instance_group"] - self.instance_id: str = extra_config["instance_id"] - self.preferred_block_size: int = extra_config.get("preferred_block_size", 0) - self.storage_configs: dict[str, dict] = extra_config.get("storage_configs", {}) - - self.write_timeout_seconds: int = extra_config.get("write_timeout_seconds", 30) - self.sdk_thread_num = extra_config.get("sdk_thread_num", 32) - self.sdk_queue_size = extra_config.get("sdk_queue_size", 1000) - self.sdk_get_timeout_ms = extra_config.get("sdk_get_timeout_ms", 15000) - self.sdk_put_timeout_ms = extra_config.get("sdk_put_timeout_ms", 15000) - - self.read_iov_block_size = extra_config.get("read_iov_block_size", 0) - self.write_iov_block_size = extra_config.get("write_iov_block_size", 0) - self.hf3fs_concurrent_io_block_count = extra_config.get("hf3fs_concurrent_io_block_count", 32) - - self.block_per_save_task = extra_config.get("block_per_save_task", 128) - self.block_per_load_task = extra_config.get("block_per_load_task", 128) - - self.async_get_cache_location = extra_config.get("async_get_cache_location", True) - # TODO: add async and try wait - # self.async_get_cache_location_wait_time = extra_config.get("async_get_cache_location_wait_time", 0) - - # Leader discovery - self.auto_discover_leader: bool = extra_config.get("auto_discover_leader", False) - self.leader_retry_count: int = extra_config.get("leader_retry_count", 1) - self.leader_retry_base_interval_seconds: float = extra_config.get("leader_retry_base_interval_seconds", 0.005) - self.discovery_refresh_interval_seconds: int = extra_config.get("discovery_refresh_interval_seconds", 30) - self.min_discover_interval_seconds: int = extra_config.get("min_discover_interval_seconds", 1) - self.request_timeout_seconds: float = extra_config.get("request_timeout_seconds", 1.0) - - self.log_level: str = extra_config.get("log_level", "") +"""kv_connector_extra_config, as a validated Pydantic model. + +Field defaults double as the documentation of every knob; unknown keys are +rejected (typos in extra_config fail at startup instead of silently +defaulting).""" + +from typing import Any, Dict + +from pydantic import BaseModel, Field + + +class TairKvCacheConnectorExtraConfig(BaseModel): + # --- Identity / registration --- + manager_uri: str + coordinator_base_port: int + instance_group: str + instance_id: str + # Manager block size override; 0 keeps the vLLM scheduler block size. + # Ignored for hybrid models (mamba state is per scheduler block). + preferred_block_size: int = 0 + storage_configs: Dict[str, Any] = Field(default_factory=dict) + + # --- Write sessions --- + write_timeout_seconds: int = 30 + + # --- Transfer SDK --- + sdk_thread_num: int = 32 + sdk_queue_size: int = 1000 + sdk_get_timeout_ms: int = 15000 + sdk_put_timeout_ms: int = 15000 + read_iov_block_size: int = 0 + write_iov_block_size: int = 0 + hf3fs_concurrent_io_block_count: int = 32 + + # --- Transfer task granularity --- + block_per_save_task: int = 128 + block_per_load_task: int = 128 + + # --- Staging buffers --- + # Pre-allocated contiguous staging slots per transfer group (shared by + # save and load). The pool is *pinned host memory only* -- the gather/ + # scatter kernel and the state copies reach it directly over PCIe, so + # the connector's device-memory footprint is zero and this knob sizes + # host RAM, not VRAM. An exhausted pool blocks the task (backpressure). + # + # The pool is the *concurrency* of staging: tasks hold their slots for + # gather + the synchronous SDK transfer, so the pool size caps how many + # tasks feed the SDK at once. At 128 (one full task) the SDK transfer + # serializes and burst loads queue behind saves: measured -3.5% ab + # throughput and +54% vs +65% TP2 hit throughput against a 512-block + # pool (perf 2026-08-20, doc protocol). 1024 restores origin/main's + # concurrency (8 full tasks) for ~896 MiB pinned host RAM per attention + # group -- host RAM is cheap, and the old 1024-slot behaviour is what + # the merge-base AB baseline ran. Shrink only together with + # block_per_save_task/block_per_load_task (must be >= the larger one, + # one task stages its whole batch contiguously). + staging_pool_blocks: int = 1024 + + # Per-group ceiling on the staging pool's pinned host RAM. The block + # count above was derived for full-attention blocks (~0.875 MiB each: + # 1024 blocks ~= 896 MiB); hybrid blocks are ~17.3 MiB, and the same + # count would pin ~17 GiB per group -- four groups then die in the + # pinned allocator at engine start on an ordinary host. Effective blocks + # per group = min(staging_pool_blocks, this_cap // block_bytes), never + # below one full task batch (the contiguity invariant). Raise this cap + # (not just the block count) to give hybrid deployments more in-flight + # transfer concurrency. + staging_pool_max_bytes_per_group: int = 2**30 + + # --- Manager queries --- + async_get_cache_location: bool = True + + # --- Leader discovery / HTTP --- + auto_discover_leader: bool = False + leader_retry_count: int = 1 + leader_retry_base_interval_seconds: float = 0.005 + discovery_refresh_interval_seconds: int = 30 + min_discover_interval_seconds: int = 1 + request_timeout_seconds: float = 1.0 + + log_level: str = "" + + # Escape hatch for the hybrid capability probe: when + # Scheduler._mamba_block_aligned_split exists but its source cannot be + # inspected (frozen/bytecode-only vLLM) the connector fails closed; + # set this to true to force-enable hybrid models there. + force_hybrid_support: bool = False + + model_config = {"extra": "forbid"} diff --git a/kv_cache_manager/py_connector/vllm/connector_scheduler.py b/kv_cache_manager/py_connector/vllm/connector_scheduler.py new file mode 100644 index 000000000..05552ad5f --- /dev/null +++ b/kv_cache_manager/py_connector/vllm/connector_scheduler.py @@ -0,0 +1,685 @@ +"""Scheduler side of the connector: matching, saving orchestration, finishing. + +The scheduler role owns the request ledger (``_tracked``): per-request state +that compensates for information vLLM only provides inside hooks -- the +accumulated block tables (only increments arrive after the first +allocation), the save water-mark and the save-session ledger that decides +when a finished request's blocks may be freed. Entries are created on the +request's first allocation and dropped at retirement; requests that were +never allocated are never tracked. + +Two small side tables carry the external-match discipline: ``_load_failed`` +(requests whose load came back invalid) and ``_load_attempted`` (requests +that spent an external allocation -- relevant for hybrid models whose load +failures cannot be reported). The worker side lives in connector_worker; both +speak the vllm_common vocabulary. +""" + +import threading +import time +from dataclasses import dataclass, field +from concurrent.futures import ThreadPoolExecutor +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple + +from kv_cache_manager.py_connector.common.logger import logger +from kv_cache_manager.py_connector.common.tp_coordinator import ( + CoordinateMsgSerializer, CoordinateMessage, SendBlockStartEvent, TpCoordinatorClient) +from kv_cache_manager.py_connector.vllm.location_query_manager import LocationQueryManager +from kv_cache_manager.py_connector.vllm.metadata import ( + FinishRequest, LoadRequest, SaveRequest, TairKvCacheConnectorMetadata) +from kv_cache_manager.py_connector.vllm.vllm_common import ( + ATTN_ONLY_SPEC_GROUP, ALL_SPEC_GROUP, GroupMeta, StateGroupMeta, + build_spec_groups, spec_name) + +if TYPE_CHECKING: + from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.outputs import KVConnectorOutput + from vllm.v1.request import Request + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + + +@dataclass +class RequestLedger: + """Per-request scheduler state, created on the request's first allocation. + + ``vllm_request`` is the live vLLM Request object; the token stream is + read from it on demand (all_token_ids) and only its length is tracked + here. ``has_saved_block_num`` anchors the incremental saves: blocks are + saved once as the computed prefix crosses manager-block boundaries.""" + vllm_request: "Request" + # Per kv_cache_group block table, in each group's own block_size units. + block_ids_per_group: List[List[int]] + # Manager blocks already saved (or covered by an external hit). + has_saved_block_num: int + # Hit accounting reported to vLLM at finish (kv_transfer_params -> the + # OpenAI response): local = vLLM's own prefix-cache hit when the match + # hook last answered, remote = the external hit it returned (clamped). + # Written on the match hook's answer paths (initial / re-ask / burned), + # read exactly once by _finish_request. Group-agnostic by design: both + # are token-granular totals. + local_matched_token_num: int = 0 + remote_matched_token_num: int = 0 + # Save-session ledger: sessions started vs sessions handed to the worker. + scheduled_saving_count: int = 0 + sent_saving_count: int = 0 + # Set when the request finished while sessions were still in flight. + need_report_after_saving_finished: bool = False + + +class ConnectorScheduler: + """State and hooks for the scheduler-role connector instance.""" + + def __init__(self, extra_config, group_metas: List[GroupMeta], + manager_block_size: int, vllm_block_size: int, tp_size: int, + manager_client, coordinator_client: TpCoordinatorClient): + self._extra_config = extra_config + self._group_metas = group_metas + self._num_groups = len(group_metas) + self._state_group_idxs = [m.group_idx for m in group_metas + if isinstance(m, StateGroupMeta)] + self._manager_block_size = manager_block_size + self._vllm_block_size = vllm_block_size + self._tp_size = tp_size + self._manager_client = manager_client + self._coordinator_client = coordinator_client + + self._epoch = 0 + self._http_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="kvcm_http_") + self._location_query_manager = LocationQueryManager( + manager_client, self._http_executor, extra_config.instance_id, + extra_config.async_get_cache_location) + + self._tracked: Dict[str, RequestLedger] = {} + self._waiting_to_load_requests: List[LoadRequest] = [] + self._waiting_to_save_requests_lock = threading.Lock() + self._waiting_to_save_requests: List[SaveRequest] = [] + self._waiting_to_finish_requests: List[FinishRequest] = [] + self._canceled_save_request_ids_lock = threading.Lock() + self._canceled_save_request_ids: List[str] = [] + # External-match discipline; cleaned at retirement. + self._load_failed: set = set() + self._load_attempted: set = set() + + def shutdown(self): + self._location_query_manager.shutdown() + self._http_executor.shutdown(wait=False) + + # ------------------------------------------------------------------ # + # Hybrid state coverage + # ------------------------------------------------------------------ # + def _spec_groups(self) -> List[dict]: + """LocationSpecGroups for registration; see vllm_common.build_spec_groups.""" + return build_spec_groups(self._group_metas, self._tp_size) + + def _group_block_size(self, group_idx: int) -> int: + for meta in self._group_metas: + if meta.group_idx == group_idx: + return meta.block_size + raise KeyError(f"no such transferred group: {group_idx}") + + def _state_complete_mask(self, ledger: RequestLedger, manager_block_idxes) -> List[bool]: + """Per manager block: does *every* state group hold a real state? + + vLLM's block table is the ground truth: in "align" mode a manager block + whose state block is the null block (id 0) has no materialized state. + Attention-only models return all True (nothing can be missing). + """ + if not self._state_group_idxs: + return [True] * len(manager_block_idxes) + mbs = self._manager_block_size + mask = [] + for mb in manager_block_idxes: + complete = True + for group_idx in self._state_group_idxs: + table = ledger.block_ids_per_group[group_idx] + # State covers the prefix ending at the block's last token. + logical = ((mb + 1) * mbs - 1) // self._group_block_size(group_idx) + if logical >= len(table) or table[logical] == 0: + complete = False + break + mask.append(complete) + return mask + + def _num_allocated_blocks(self, ledger: RequestLedger) -> int: + """Min allocated block-table length across the *transferred* groups. + + ``block_ids_per_group`` is indexed by the vLLM group index and includes + groups skipped by ``parse_groups`` (EAGLE/MTP drafters). A drafter's + block table can lag behind the target model's, so including it in the + min would permanently understate how many blocks are saveable.""" + if not ledger.block_ids_per_group: + return 0 + return min(len(ledger.block_ids_per_group[meta.group_idx]) + for meta in self._group_metas) + + # ------------------------------------------------------------------ # + # External matching + # ------------------------------------------------------------------ # + def _location_covers_states(self, location: dict) -> bool: + """Does this published block carry every state group's spec for a rank? + + A hybrid block saved without a materialized recurrent state is published + under the attention-only spec group, so its location simply has no spec + for the state groups. ``getCacheLocation`` reports the real coverage, + which is what makes the sparsity visible here. + """ + names = {spec.get("name") for spec in location.get("location_specs", [])} + return all(spec_name(rank, group_idx) in names + for rank in range(self._tp_size) + for group_idx in self._state_group_idxs) + + def _external_match_burned(self, req_id: str) -> bool: + """Has this request lost its option of an external match? + + * single-group block table: load failures are reported to vLLM + (report_failures=True in the worker's start_load_kv) and come back + as invalid block ids in update_connector_output -- the explicit + signal. A mere preemption re-query keeps its match: the loaded KV + is healthy, only the scheduling position was lost. + * multi-group block table (hybrid, or a skipped EAGLE/MTP drafter + group alongside the transferred attention group): vLLM's + invalid-block recovery unpacks a single group (upstream TODO), so + failures cannot be reported and no signal ever comes back. The + request instead gets one conservative shot: any allocation for an + external hit burns the match, because a failed block cannot be + told apart from a healthy one afterwards. + + The shape is read from the block-table snapshot recorded at + allocation time (``ledger.block_ids_per_group``): it mirrors + kv_cache_manager.get_block_ids, the tuple the recovery path unpacks, + and its length -- the group count -- is model-static. + """ + if req_id in self._load_failed: + return True + if req_id not in self._load_attempted: + return False + ledger = self._tracked.get(req_id) + if ledger is None or not ledger.block_ids_per_group: + # Defensive: every load attempt is recorded through + # update_state_after_alloc, which snapshots the block table. + # Without a snapshot the recovery shape is unknown; keep the + # match rather than burn it blindly. + return False + return len(ledger.block_ids_per_group) > 1 + + def get_num_new_matched_tokens(self, request: "Request", + num_computed_tokens: int) -> Tuple[Optional[int], bool]: + """Answer vLLM's per-request question: beyond the ``num_computed_tokens`` + it already has, how many more tokens can the external KV supply? + + A pure query: it never blocks (an in-flight manager query returns + None and vLLM re-asks next step) and never mutates request state. + The answer is cached per request; ``update_state_after_alloc`` + consumes it and turns it into the LoadRequest once vLLM allocates + blocks for the hit. + + Returns ``(external_tokens, load_kv_async)``: ``None`` while the + query is in flight; ``0`` when there is no hit, or the request's + match is burned (see ``_external_match_burned``); ``>0`` for a + block-aligned count clamped to a prefix vLLM can safely resume from + (``_safe_external_prefix``). + """ + req_id = request.request_id + if self._external_match_burned(req_id): + # vLLM re-asks whenever a request falls back to the waiting queue: + # preemption, or a failed load under kv_load_failure_policy= + # recompute. For this request the external match is burned -- + # see _external_match_burned for which signal burned it. Whatever + # is left, vLLM recomputes locally. + logger.warning("req:%s re-queried after an external load attempt, " + "skip external match", req_id) + # Hit accounting: this re-ask answers "no external match"; + # the local hit stands, the remote hit is spent. + ledger = self._tracked.get(req_id) + if ledger is not None: + ledger.local_matched_token_num = num_computed_tokens + ledger.remote_matched_token_num = 0 + return 0, False + + computed_blocks = num_computed_tokens // self._manager_block_size + need_load_locations = self._location_query_manager.get_locations_for_query( + request, computed_blocks) + if need_load_locations is None: + return None, False + + need_load_locations = self._safe_external_prefix( + req_id, need_load_locations, + num_computed_tokens, request.num_tokens) + # Cache the clamped answer: the allocation consumes exactly this. + self._location_query_manager.store_result(req_id, need_load_locations) + new_matched_count = len(need_load_locations) * self._manager_block_size + # Hit accounting: record the answer the request will finish with + # (local = vLLM's hit when asked, remote = this clamped answer). + ledger = self._tracked.get(req_id) + if ledger is None: + ledger = RequestLedger( + vllm_request=request, + block_ids_per_group=[], + has_saved_block_num=0, + ) + self._tracked[req_id] = ledger + ledger.local_matched_token_num = num_computed_tokens + ledger.remote_matched_token_num = new_matched_count + logger.info("req:%s matched %d external tokens", req_id, new_matched_count) + return new_matched_count, new_matched_count > 0 + + def _safe_external_prefix(self, req_id: str, locations: List[dict], + num_computed_tokens: int, num_tokens: int) -> List[dict]: + """Clamp an external match to the longest prefix vLLM can resume from. + + Two constraints, both only ever trimming the tail, so the answer is + the longest prefix satisfying both at once: + + * the match must end on a state-complete block -- a hybrid request + resumes from the recurrent state ending the reused prefix, so an + ending block without one is unloadable however much attention KV + precedes it; + * at least one token must stay uncomputed -- logits are not part of + the KV cache, so the model still needs one token to sample from, + and vLLM's synchronous-load path asserts num_new_tokens > 0 + (vLLM's own connectors apply the same cap). + + Full-attention models have no state groups and only feel the cap. + """ + # Cap: how many leading blocks may be matched at all. + limit = len(locations) + while limit and num_computed_tokens + limit * self._manager_block_size >= num_tokens: + limit -= 1 + # Within that allowance the match may only end on a block carrying + # the recurrent state: scan for the last one. + keep = 0 + for i, location in enumerate(locations[:limit]): + if self._location_covers_states(location): + keep = i + 1 + if keep < len(locations): + logger.info("req:%s truncated external match from %d to %d blocks " + "(full-hit cap / no recurrent state at the end)", + req_id, len(locations), keep) + return locations[:keep] + + def update_state_after_alloc(self, request: "Request", blocks: "KVCacheBlocks", + num_external_tokens: int): + """First (or re-) allocation: record the ledger and ship the load. + + The block tables arrive here whole, once per allocation; increments + come later through the scheduling output. When vLLM allocated for an + external hit, the query the match hook cached is consumed and turned + into the LoadRequest -- at this point, and only here, both halves of + its address are known (manager locations + physical slots).""" + req_id = request.request_id + ledger = self._tracked.get(req_id) + if ledger is None: + ledger = RequestLedger( + vllm_request=request, + block_ids_per_group=[], + has_saved_block_num=0, + ) + self._tracked[req_id] = ledger + ledger.block_ids_per_group = [list(b) for b in blocks.get_block_ids()] + + if num_external_tokens <= 0: + return + consumed = self._location_query_manager.consume_locations(req_id) + if consumed is None: + # Invariant violation: vLLM allocated blocks for an external hit + # (num_external_tokens > 0) only after the match hook answered + # positively, and that answer was stored for exactly this + # consumption. No cached answer here means the match/alloc + # contract is broken -- vLLM would run the request on KV that + # was never loaded. Fail closed: a traceback beats silently + # corrupting the request's output. + raise RuntimeError( + f"req {req_id}: allocated for {num_external_tokens} external " + f"tokens but no cached location query exists; the match hook " + f"did not answer, or its answer was consumed/invalidated " + f"before this allocation") + locations, computed_blocks = consumed + # Blocks were allocated for an external hit: the request is now + # spending its external load (burns hybrid re-queries). + self._load_attempted.add(req_id) + if not locations: + return + total_remote_blocks = computed_blocks + len(locations) + ledger.has_saved_block_num = total_remote_blocks + self._waiting_to_load_requests.append(LoadRequest( + req_id=req_id, + manager_block_idxes=list(range(computed_blocks, total_remote_blocks)), + need_load_locations=locations, + all_block_ids=[list(b) for b in ledger.block_ids_per_group], + )) + + def update_connector_output(self, connector_output: "KVConnectorOutput"): + """Consume the worker's step output: mark requests whose external + load failed, at request granularity. + + vLLM reports invalid blocks, not requests; a block id is matched + against the block tables recorded in update_state_after_alloc. One + failed block is enough -- the request recomputes as a whole, which + blocks to recompute exactly is vLLM's decision. + """ + invalid = getattr(connector_output, "invalid_block_ids", None) + if not invalid: + return + for req_id, ledger in self._tracked.items(): + if any(b in invalid + for group_ids in ledger.block_ids_per_group for b in group_ids): + self._load_failed.add(req_id) + logger.warning("req:%s external load failed (invalid blocks " + "reached its block table)", req_id) + + # ------------------------------------------------------------------ # + # Per-step metadata assembly and saving orchestration + # ------------------------------------------------------------------ # + def build_connector_meta(self, scheduler_output: "SchedulerOutput") -> TairKvCacheConnectorMetadata: + """Assemble one engine step's envelope (see TairKvCacheConnectorMetadata). + + Two sources feed the envelope: vLLM's scheduling output for this step + (block-table increments for continuing requests) and the residue of + earlier steps' async work (admitted loads, resolved save sessions, + retirements). Only the last stage depends on the others: save + settlement may retire a request within this step, so finishes flush + last. + """ + meta = TairKvCacheConnectorMetadata(self._epoch) + self._epoch += 1 + + self._ingest_scheduled_reqs(scheduler_output) + self._dispatch_incremental_saves() + self._collect_load_instructions(meta) + self._collect_save_instructions(meta) + self._collect_finish_instructions(meta) + return meta + + def _ingest_scheduled_reqs(self, scheduler_output: "SchedulerOutput") -> None: + """Absorb vLLM's scheduling list into the request ledger. + + The block tables are the scheduler's ledger: vLLM hands them over + once per first allocation (update_state_after_alloc / + scheduled_new_reqs), then only as increments (cached_reqs. + new_block_ids). Accumulating them here is what lets later stages + translate manager blocks into physical slots for the worker's + instructions. A preemption resume replaces the whole table + (re-allocation maps the same logical blocks to fresh slots).""" + for vllm_req in scheduler_output.scheduled_new_reqs: + ledger = self._tracked.get(vllm_req.req_id) + if ledger is None: + # update_state_after_alloc creates the ledger before vLLM + # can schedule the request, so this only fires on a broken + # hook-order contract: log it loudly instead of silently + # dropping the block table. + logger.warning( + "scheduled new req %s has no ledger (hook-order " + "contract broken?); its block table is not recorded", + vllm_req.req_id) + continue + ledger.block_ids_per_group = [list(b) for b in vllm_req.block_ids] + + cached_reqs = scheduler_output.scheduled_cached_reqs + for idx, req_id in enumerate(cached_reqs.req_ids): + ledger = self._tracked.get(req_id) + if ledger is None: + if hasattr(cached_reqs, "resumed_req_ids"): + resumed = req_id in cached_reqs.resumed_req_ids + else: + resumed = cached_reqs.resumed_from_preemption[idx] + logger.warning( + "scheduled cached req %s has no ledger (resumed=%s, " + "scheduled_tokens=%d); its block increments are not " + "recorded", req_id, resumed, + scheduler_output.num_scheduled_tokens[req_id]) + continue + + if hasattr(cached_reqs, "resumed_req_ids"): + resumed = req_id in cached_reqs.resumed_req_ids + else: + resumed = cached_reqs.resumed_from_preemption[idx] + + new_block_ids = cached_reqs.new_block_ids[idx] + if new_block_ids is None: + # https://github.com/vllm-project/vllm/pull/23262: None when + # no group got new blocks this step (get_block_ids with + # allow_none=True). Keep the recorded table as-is: a resumed + # request with nothing allocated has no fresh slots to record. + continue + if resumed: + ledger.block_ids_per_group = [list(b) for b in new_block_ids] + else: + for group_ids, new_ids in zip(ledger.block_ids_per_group, + new_block_ids): + group_ids.extend(new_ids) + + def _dispatch_incremental_saves(self) -> None: + """Start async StartWriteCache calls for requests whose computed + prefix crossed a manager-block boundary since the last step. + + This stage only acquires write locations (the HTTP call resolves tens + of ms later, off the scheduler thread); the data itself moves later, + gathered by the worker out of HBM. The session the http thread + produces surfaces in _collect_save_instructions on a later step. + """ + for ledger in self._tracked.values(): + # Count blocks by the key material (all_token_ids), never by the + # scheduled token count: during decode all_token_ids lags one + # token behind (the token scheduled in this step is appended only + # once sampled), so a scheduled-count-derived block count + # announces blocks whose last token -- and thus cache key -- is + # not known yet. The manager then returns one location fewer + # than announced and the worker's strict alignment check drops + # the whole session. The allocated cap still bounds the other + # way: under chunked prefill all_token_ids runs ahead of the + # computed KV. + target_save_num = min( + len(ledger.vllm_request.all_token_ids), + self._num_allocated_blocks(ledger) * self._vllm_block_size) \ + // self._manager_block_size + if target_save_num > ledger.has_saved_block_num: + logger.info("req:%s incremental save: %d -> %d blocks " + "(tokens=%d, allocated=%d)", + ledger.vllm_request.request_id, + ledger.has_saved_block_num, target_save_num, + len(ledger.vllm_request.all_token_ids), + self._num_allocated_blocks(ledger)) + ledger.scheduled_saving_count += 1 + # Per-block state completeness must be read here, in the + # scheduler loop: it comes from vLLM's block table, which the + # http_executor thread would race against later steps. + self._http_executor.submit( + self.start_save_kvcache_async, ledger.vllm_request.request_id, + ledger.vllm_request.all_token_ids[:target_save_num * self._manager_block_size], + target_save_num, + self._state_complete_mask(ledger, range(target_save_num))) + ledger.has_saved_block_num = target_save_num + + def _collect_load_instructions(self, meta: TairKvCacheConnectorMetadata) -> None: + """Hand the worker the loads whose blocks were allocated. + + LoadRequests are built by update_state_after_alloc, at the moment + both halves of their address (manager locations + physical slots) + became known; this stage only moves them into the envelope.""" + for load_req in self._waiting_to_load_requests: + meta.add_load_request(load_req) + self._waiting_to_load_requests = [] + + def _collect_save_instructions(self, meta: TairKvCacheConnectorMetadata) -> None: + """Hand the worker the save sessions whose write locations have + arrived, and settle finished requests whose saving is now complete. + + Canceled sessions (start_write_cache failed on the http thread) + settle here too: they count as sent, so a request waiting only on + them can be retired. + """ + with self._waiting_to_save_requests_lock: + new_save_reqs = self._waiting_to_save_requests + self._waiting_to_save_requests = [] + for save_req in new_save_reqs: + ledger = self._tracked.get(save_req.req_id) + if ledger is None: + logger.warning("request %s is not tracked, skip saving", save_req.req_id) + continue + # Snapshot the ledger for the worker: the gather translates + # manager blocks into physical slots through these tables, and + # the worker keeps no mirror of its own. + save_req.all_block_ids = [list(b) for b in ledger.block_ids_per_group] + meta.add_save_request(save_req) + ledger.sent_saving_count += 1 + if (ledger.need_report_after_saving_finished and + ledger.scheduled_saving_count == ledger.sent_saving_count): + self._retire_request(save_req.req_id) + + self.handle_canceled_save_req() + + def _collect_finish_instructions(self, meta: TairKvCacheConnectorMetadata) -> None: + """Flush retirements queued since the last step. Runs last on purpose: + save settlement above may retire a request within this very step.""" + for finish_req in self._waiting_to_finish_requests: + meta.add_finish_request(finish_req) + self._waiting_to_finish_requests = [] + + def _retire_request(self, req_id: str) -> None: + """The request's save obligations are settled: tell the worker to drop + its bookkeeping and stop tracking the request ourselves.""" + self._waiting_to_finish_requests.append(FinishRequest(req_id)) + self._tracked.pop(req_id, None) + self._load_failed.discard(req_id) + self._load_attempted.discard(req_id) + self._location_query_manager.invalidate(req_id) + + def start_save_kvcache_async(self, req_id, token_ids, target_save_num, + state_complete_mask): + """Ask the manager for write locations for a request's first + ``target_save_num`` manager blocks. + + ``state_complete_mask[i]`` says whether manager block i has a + materialized recurrent state. Blocks without one are announced under + the attention-only spec group, so the manager allocates (and later + advertises) only the specs that will really hold data -- absence is + never encoded as a successful write. + """ + request = { + "trace_id": "%s_%d" % (req_id, self._epoch), + "instance_id": self._extra_config.instance_id, + "block_keys": [], + "token_ids": token_ids, + "write_timeout_seconds": self._extra_config.write_timeout_seconds, + } + if self._state_group_idxs: + assert len(state_complete_mask) == target_save_num, ( + f"state mask {len(state_complete_mask)} != {target_save_num} blocks") + request["location_spec_group_names"] = [ + ALL_SPEC_GROUP if complete else ATTN_ONLY_SPEC_GROUP + for complete in state_complete_mask] + if not all(state_complete_mask): + logger.info("req:%s saving %d/%d blocks without a recurrent " + "state (attention specs only)", req_id, + state_complete_mask.count(False), target_save_num) + try: + response = self._manager_client.start_write_cache(request) + except Exception as e: + logger.warning("start_write_cache error, skip saving: %s", e) + with self._canceled_save_request_ids_lock: + self._canceled_save_request_ids.append(req_id) + return + + locations = response["locations"] + write_session_id = response["write_session_id"] + mask = response.get("block_mask") or {} + if "bool_masks" in mask: + values = mask["bool_masks"].get("values", []) + mask_summary = (f"bool_masks offset={mask['bool_masks'].get('offset')} " + f"total={len(values)} " + f"existing={sum(1 for v in values if v)}") + else: + mask_summary = f"offset={mask.get('offset')}" + logger.info("req:%s save session %s: block_mask %s locations=%d", + req_id, write_session_id[:8], mask_summary, len(locations)) + logger.debug("req:%s save session %s: block_mask=%s locations=%d", + req_id, write_session_id[:8], + response.get("block_mask"), len(locations)) + + if not locations: + try: + self._manager_client.finish_write_cache({ + "trace_id": "finish_%s" % write_session_id[:8], + "instance_id": self._extra_config.instance_id, + "write_session_id": write_session_id, + "success_blocks": {"bool_masks": {"offset": 0}}, + }) + except Exception as e: + logger.warning("finish_write_cache failed, session: %s, error: %s", + write_session_id, e) + with self._canceled_save_request_ids_lock: + self._canceled_save_request_ids.append(req_id) + return + + need_block_idx = self.parse_block_mask_to_save_indices(response, target_save_num) + message = CoordinateMessage(time.time(), SendBlockStartEvent( + request_id=req_id, write_session_id=write_session_id, locations=locations)) + self._coordinator_client.send(CoordinateMsgSerializer.dumps(message)) + + with self._waiting_to_save_requests_lock: + self._waiting_to_save_requests.append(SaveRequest( + req_id, locations, need_block_idx, write_session_id)) + + def handle_canceled_save_req(self): + with self._canceled_save_request_ids_lock: + canceled = self._canceled_save_request_ids + self._canceled_save_request_ids = [] + for req_id in canceled: + # Cancellations come from http_executor threads; the request may + # already have been finished and removed by the scheduler loop. + ledger = self._tracked.get(req_id) + if ledger is None: + logger.warning("canceled save for unknown request %s, skip", req_id) + continue + ledger.sent_saving_count += 1 + if (ledger.need_report_after_saving_finished and + ledger.scheduled_saving_count == ledger.sent_saving_count): + self._retire_request(req_id) + + def get_finished_count(self): + # Only rank0 reports finished requests. + return 1 + + def parse_block_mask_to_save_indices(self, response: dict, target_save_num: int) -> List[int]: + block_mask = response.get("block_mask", {}) + if "offset" in block_mask: + return list(range(block_mask["offset"], target_save_num)) + values = block_mask.get("bool_masks", {}).get("values", []) + return [idx for idx, saved in enumerate(values) if not saved] + + # ------------------------------------------------------------------ # + # Request finish + # ------------------------------------------------------------------ # + def request_finished_all_groups( + self, request: "Request", + block_ids: Tuple[List[int], ...]) -> Tuple[bool, Optional[dict]]: + return self._finish_request(request) + + def request_finished(self, request: "Request", + block_ids: List[int]) -> Tuple[bool, Optional[dict]]: + return self._finish_request(request) + + def _finish_request(self, request: "Request") -> Tuple[bool, Optional[dict]]: + req_id = request.request_id + ledger = self._tracked.get(req_id) + if ledger is None: + logger.info("request_finished for untracked request: %s", req_id) + self._location_query_manager.invalidate(req_id) + return False, None + + # Hit accounting travels to the client through vLLM's + # kv_transfer_params (EngineCoreOutput -> RequestOutput -> the + # OpenAI response); consumers attribute TTFT and per-request hits + # to local (vLLM prefix cache) vs remote (KVCM) sources with it. + extra_info = { + "local_matched_token_num": ledger.local_matched_token_num, + "remote_matched_token_num": ledger.remote_matched_token_num, + } + + if ledger.scheduled_saving_count == ledger.sent_saving_count: + self._retire_request(req_id) + return True, extra_info + + # Saves still in flight; delay freeing the blocks until they land. + ledger.need_report_after_saving_finished = True + return True, extra_info diff --git a/kv_cache_manager/py_connector/vllm/connector_worker.py b/kv_cache_manager/py_connector/vllm/connector_worker.py new file mode 100644 index 000000000..5c8654c63 --- /dev/null +++ b/kv_cache_manager/py_connector/vllm/connector_worker.py @@ -0,0 +1,544 @@ +"""Worker side of the connector: translation and data-plane transfer. + +Owns the transfer client and the paged-cache views (TransferGroups). Every +instruction arriving through the metadata is self-contained (LoadRequest / +SaveRequest carry their own block tables), so the worker keeps no mirrored +request state -- tp0's per-request save-session ledger is the only +request-level bookkeeping. Per-step metadata is passed in explicitly by the +shell. +""" + +import copy +import json +import typing +from typing import List, Optional, Tuple + +from kv_cache_manager.client.pybind import kvcm_py_client + +import torch +from vllm.distributed import get_tensor_model_parallel_rank + +from kv_cache_manager.py_connector.common.logger import logger +from kv_cache_manager.py_connector.common.tp_coordinator import ( + SaveContext, TpCoordinatorClient, TpCoordinatorServer) +from kv_cache_manager.py_connector.vllm.data_transfer import ( + MultiResult, DataTransferManager, _get_device_module) +from kv_cache_manager.py_connector.vllm.metadata import ( + FinishRequest, TairKvCacheConnectorMetadata) +from kv_cache_manager.py_connector.vllm.transfer_types import ( + AttentionTransferGroup, KVCacheInfo, StateTransferGroup, TransferGroup, + TransferPlan) +from kv_cache_manager.py_connector.vllm.vllm_common import ( + AttentionGroupMeta, GroupMeta, StateGroupMeta, attn_kv_views, spec_name) + +if typing.TYPE_CHECKING: + from vllm.forward_context import ForwardContext + from vllm.attention import AttentionMetadata + + +class ConnectorWorker: + """State and hooks for the worker-role connector instance (one per TP rank).""" + + def __init__(self, extra_config, group_metas: List[GroupMeta], + manager_block_size: int, tp_size: int, host_ip: str, + manager_client, coordinator_client: TpCoordinatorClient, + register_response: dict): + self._extra_config = extra_config + self._group_metas = group_metas + self._num_groups = len(group_metas) + self._manager_block_size = manager_block_size + self._tp_size = tp_size + self._manager_client = manager_client + self._coordinator_client = coordinator_client + + # Per-request in-flight save sessions (write_session granularity). + # tp0 only; this is the worker's whole request-level state -- there is + # no mirrored request table (every instruction is self-contained). + self._pending_saves: dict = {} + self._finish_pending: set = set() + + self._tp_rank = get_tensor_model_parallel_rank() + self._device_mod = None + port = extra_config.coordinator_base_port + if self._tp_rank == 0: + self._coordinator_server = TpCoordinatorServer( + host_ip, port, tp_size, self.on_save_finished) + + self._self_spec_names = { + meta.group_idx: spec_name(self._tp_rank, meta.group_idx) + for meta in self._group_metas + } + max_group_bytes = max(m.per_block_bytes for m in self._group_metas) + self._iov_size = max_group_bytes * extra_config.hf3fs_concurrent_io_block_count + + self._storage_configs = register_response["storage_configs"] + sdk_backend_configs = self.parse_hf3fs_configs(self._storage_configs) + transfer_client_json = { + "instance_group": extra_config.instance_group, + "instance_id": extra_config.instance_id, + "block_size": self._manager_block_size, + "sdk_config": { + "thread_num": extra_config.sdk_thread_num, + "queue_size": extra_config.sdk_queue_size, + "sdk_backend_configs": sdk_backend_configs, + "timeout_config": { + "get_timeout_ms": extra_config.sdk_get_timeout_ms, + "put_timeout_ms": extra_config.sdk_put_timeout_ms, + }, + }, + "location_spec_infos": { + self._self_spec_names[meta.group_idx]: meta.per_block_bytes + for meta in self._group_metas + }, + } + init_params = kvcm_py_client.InitParams() + init_params.role_type = kvcm_py_client.RoleType.WORKER + init_params.self_location_spec_name = self._self_spec_names[self._group_metas[0].group_idx] + init_params.storage_configs = f"{self._storage_configs}" + transfer_client_config = json.dumps(transfer_client_json) + logger.info("transfer_client_config: %s", transfer_client_config) + self._transfer_client = kvcm_py_client.TransferClient.Create( + transfer_client_config, init_params) + assert self._transfer_client is not None, "kvcm_py_client.TransferClient.Create failed" + logger.warning( + "TairKvCacheConnector worker inited, tp rank: %d/%d, host: %s:%d, groups: %d", + self._tp_rank, self._tp_size, host_ip, port, self._num_groups) + + # ------------------------------------------------------------------ # + # Storage config plumbing + # ------------------------------------------------------------------ # + def parse_hf3fs_configs(self, storage_configs): + hf3fs_configs = [] + storage_configs_json = json.loads(storage_configs) + for storage_config in storage_configs_json: + if storage_config["type"] == "vcns_hf3fs": + storage_config["type"] = "hf3fs" + if storage_config["type"] == "hf3fs" and storage_config["is_available"]: + hf3fs_configs.append({ + "type": storage_config["type"], + "mountpoint": storage_config["storage_spec"]["mountpoint"], + "root_dir": storage_config["storage_spec"]["root_dir"], + "read_iov_block_size": self._extra_config.read_iov_block_size, + "read_iov_size": self._iov_size, + "write_iov_block_size": self._extra_config.write_iov_block_size, + "write_iov_size": self._iov_size, + }) + self._storage_configs = json.dumps(storage_configs_json) + return hf3fs_configs + + # ------------------------------------------------------------------ # + # KV cache registration + # ------------------------------------------------------------------ # + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): + self._kv_caches = kv_caches + first_attn = next(kv_caches[name] + for meta in self._group_metas + if isinstance(meta, AttentionGroupMeta) + for name in meta.layer_names) + self._dtype = first_attn.dtype + self._device = first_attn.device + self._device_mod = _get_device_module(self._device) + + groups: List[TransferGroup] = [] + for meta in self._group_metas: + if isinstance(meta, AttentionGroupMeta): + groups.append(self._build_attention_group(meta, kv_caches)) + else: + groups.append(self._build_state_group(meta, kv_caches)) + + self._kvcache_info = KVCacheInfo( + tp_rank=self._tp_rank, + world_size=self._tp_size, + groups=groups, + device=self._device, + dtype=self._dtype, + ) + self._data_transfer = DataTransferManager( + self._kvcache_info, self._manager_block_size, + self._transfer_client, self._coordinator_client, self._extra_config) + + logger.warning("register_kv_caches done: %s", [ + (g.spec_name, type(g).__name__, g.layer_num, g.per_block_bytes) + for g in groups]) + + def _build_attention_group(self, meta: AttentionGroupMeta, + kv_caches) -> AttentionTransferGroup: + spec = self._self_spec_names[meta.group_idx] + tensors = [kv_caches[name] for name in meta.layer_names] + ref = tensors[0] + for t in tensors: + assert t.shape == ref.shape and t.stride() == ref.stride(), \ + "attention layers in one group must share shape/stride" + # Normalize the layout into per-pointer token-major views of shape + # (num_blocks, kernel_block_size, heads, content_dim); everything + # below is layout-independent. Split-K/V layouts (vllm <= 0.25.x) + # yield two views (= two transfer pointers) per layer, the packed + # layout (vllm >= 0.26.0) yields one. + ref_views, kv_layout = attn_kv_views(ref) + view = ref_views[0] + kernel_block_size = view.shape[1] + assert meta.block_size % kernel_block_size == 0, \ + f"group block size {meta.block_size} not a multiple of kernel " \ + f"block size {kernel_block_size}" + per_token_dim = view.shape[2] * view.shape[3] # heads * content_dim + # The gather/scatter kernel needs token-major memory inside a page: + # logical dims (blk, tok, head, dim) contiguous within a block. This + # is vLLM's NHD order; HND would interleave heads across tokens. + for v in ref_views: + assert v.stride()[1:] == (per_token_dim, v.shape[3], 1), \ + f"kv cache page not token-major: shape={tuple(v.shape)} " \ + f"stride={v.stride()}; set VLLM_KV_CACHE_LAYOUT=NHD" + # Non-flat block layouts (page_size_padded gaps, or split K/V + # interleaved per block as in the 5-D N-first layout) go through the + # kernel's strided path. Stride 0 = fast flat indexing. + flat = view.stride(0) == kernel_block_size * per_token_dim + block_stride = 0 if flat else view.stride(0) + # Pointer array ordered [K0, V0, K1, V1, ...] for split layouts and + # [L0, L1, ...] for the packed layout; each view's data_ptr() is its + # own storage base, so the kernel never adds a K->V offset. + ptrs = [v.data_ptr() for t in tensors for v in attn_kv_views(t)[0]] + ptr_tensor = torch.tensor(ptrs, dtype=torch.int64, device="cpu").to(self._device) + return AttentionTransferGroup( + group_idx=meta.group_idx, + spec_name=spec, + layer_names=meta.layer_names, + block_size=meta.block_size, + per_block_bytes=meta.per_block_bytes, + layer_num=len(meta.layer_names), + kv_layout=kv_layout, + kvcache_ptr_tensor_gpu=ptr_tensor, + num_kv_ptrs=len(ptrs), + per_token_dim=per_token_dim, + kernel_block_size=kernel_block_size, + block_stride=block_stride, + ) + + def _build_state_group(self, meta: StateGroupMeta, + kv_caches) -> StateTransferGroup: + # Mamba/state group: each layer is a list[Tensor] sharing one storage; + # rebuild a (num_blocks, page_size_bytes) byte view for opaque copy. + spec = self._self_spec_names[meta.group_idx] + block_views = [] + for name in meta.layer_names: + states = kv_caches[name] + assert isinstance(states, (list, tuple)) and len(states) > 0, \ + f"state layer {name} should be a list of tensors" + storage = states[0].untyped_storage() + for st in states[1:]: + assert st.untyped_storage().data_ptr() == storage.data_ptr(), \ + f"state layer {name}: tensors do not share storage" + num_blocks = states[0].shape[0] + need = num_blocks * meta.page_size_bytes + assert storage.nbytes() >= need, \ + f"state layer {name}: storage {storage.nbytes()} < {need}" + byte_view = torch.tensor([], dtype=torch.uint8, device=self._device).set_(storage) + block_views.append(byte_view[:need].view(num_blocks, meta.page_size_bytes)) + return StateTransferGroup( + group_idx=meta.group_idx, + spec_name=spec, + layer_names=meta.layer_names, + block_size=meta.block_size, + per_block_bytes=meta.per_block_bytes, + layer_num=len(meta.layer_names), + block_view_tensors=block_views, + page_size_bytes=meta.page_size_bytes, + ) + + # ------------------------------------------------------------------ # + # Block index translation + # ------------------------------------------------------------------ # + def _attn_token_indices(self, group: AttentionTransferGroup, manager_block_idxes, + block_table) -> List[List[int]]: + """Map manager blocks to flat token slots of one attention group. + + Three-tier hierarchy: + manager block (KVCM unit) -> global token idx + -> group block (block_table unit, group.block_size tokens) + -> kernel physical block (tensor unit; ratio physical per group block). + """ + mbs = self._manager_block_size + gbs = group.block_size + kbs = group.kernel_block_size + ratio = gbs // kbs + out = [] + for mb in manager_block_idxes: + idxs = [] + base = mb * mbs + for i in range(mbs): + tok = base + i + logical = tok // gbs + assert logical < len(block_table), ( + f"group block {logical} out of range (len={len(block_table)})") + off = tok % gbs + phys = block_table[logical] * ratio + off // kbs + idxs.append(phys * kbs + off % kbs) + out.append(idxs) + return out + + def _state_block_ids(self, group: StateTransferGroup, manager_block_idxes, + block_table) -> List[int]: + """Map manager blocks to block ids of a state (mamba) group. + + State is stored once per group block and covers the whole prefix up to + that block, so the manager block's last token selects the block. + + KNOWN LIMITATION (mamba_cache_mode="none"): without prefix caching + vLLM keeps only one resident state block per request, so every manager + block resolves to that same block id here. Saves then publish the + *current* state bytes under every manager block's key (mislabeled as + earlier prefixes) and loads overwrite one block repeatedly (last write + wins). Hybrid serving is only supported with prefix caching enabled + (mamba_cache_mode="align"), where the table is position-indexed.""" + mbs = self._manager_block_size + gbs = group.block_size + out = [] + for mb in manager_block_idxes: + logical = ((mb + 1) * mbs - 1) // gbs + assert logical < len(block_table), ( + f"group block {logical} out of range (len={len(block_table)})") + out.append(block_table[logical]) + return out + + # ------------------------------------------------------------------ # + # Load / save + # ------------------------------------------------------------------ # + def _iter_task_chunks(self, plans: List[TransferPlan], per_task: int): + """Slice each plan's blocks into per-task chunks. + + Pure generator: the task index is the consumer's concern (enumerate + there); this only reads the plans.""" + for plan in plans: + n = len(plan.uris) + for i in range(0, n, per_task): + end = min(n, i + per_task) + yield (plan.group, + plan.uris[i:end], + plan.token_indices[i:end] if plan.token_indices is not None else None, + plan.block_ids[i:end] if plan.block_ids is not None else None) + + def _plan_group_transfers(self, locations, manager_block_idxes, + block_ids_per_group) -> Optional[List[TransferPlan]]: + """Build the per-group TransferPlans for a set of manager blocks. + + ``uris`` is positionally aligned with ``manager_block_idxes`` and may + contain ``None`` where a block carries no data for that group's spec + (hybrid per-block spec coverage, see ``build_spec_groups``). Deciding + what a hole means is the transfer task's job: for attention groups a + hole is a failure, for state groups it must agree with the block + table (a null state block). Returns None only if the manager returned + the wrong number of locations altogether. + """ + num_blocks = len(manager_block_idxes) + if len(locations) != num_blocks: + logger.warning("%d locations for %d blocks, skip transfer", + len(locations), num_blocks) + return None + # One pass over the locations: per block, map spec name -> uri. The + # per-group extraction below then reads by name instead of rescanning + # every location's spec list for every group. + per_block_uri_maps = [ + {s["name"]: s["uri"] for s in location.get("location_specs", [])} + for location in locations] + self._check_block_table_covers(block_ids_per_group) + plans = [] + for group in self._kvcache_info.groups: + uris = [m.get(group.spec_name) for m in per_block_uri_maps] + # block_ids_per_group is indexed by the vLLM group index. + block_table = block_ids_per_group[group.group_idx] + if isinstance(group, AttentionTransferGroup): + plans.append(TransferPlan( + group=group, uris=uris, + token_indices=self._attn_token_indices( + group, manager_block_idxes, block_table))) + else: + plans.append(TransferPlan( + group=group, uris=uris, + block_ids=self._state_block_ids( + group, manager_block_idxes, block_table))) + return plans + + def _check_block_table_covers(self, block_ids_per_group) -> None: + """``block_ids_per_group`` is indexed by the raw vLLM group index -- + a guarantee vLLM provides today + (https://github.com/vllm-project/vllm/blob/v0.26.0/vllm/v1/core/sched/output.py, + CachedRequestData.new_block_ids: one entry per kv cache group) but + that we must not silently rely on. Fail loudly when the table cannot + address every transferred group.""" + if not block_ids_per_group: + return + need = max(m.group_idx for m in self._group_metas) + assert len(block_ids_per_group) > need, ( + f"block table has {len(block_ids_per_group)} groups but the " + f"connector transfers vLLM group {need}; the group-index " + f"alignment between vLLM and this connector is broken") + + def start_load_kv(self, forward_context: "ForwardContext", + meta: TairKvCacheConnectorMetadata, **kwargs) -> None: + for load_req in meta.to_load_requests: + if not load_req.need_load_locations: + continue + num_blocks = len(load_req.manager_block_idxes) + plans = self._plan_group_transfers( + load_req.need_load_locations, load_req.manager_block_idxes, + load_req.all_block_ids) + + # Report failures against the block table vLLM can act on: map each + # manager block to the logical block holding its first token. vLLM + # truncates computed tokens at the first invalid block, so this is + # sufficient for recovery. + # + # UPSTREAM LIMITATION: vLLM's invalid-block recovery + # (Scheduler._update_requests_with_invalid_blocks, vllm/v1/core/ + # sched/scheduler.py) unpacks a single-group block table -- + # "(req_block_ids,) = ...get_block_ids(req_id)" under + # "TODO (davidb): add support for hybrid memory allocator" -- so + # for multi-group models (hybrid attn+mamba, or any model whose + # vLLM block table has more than one entry -- skipped EAGLE/MTP + # drafter groups still count) a failed load CANNOT be reported + # and is only logged; vLLM then decodes from whatever bytes the + # partial load left in the paged cache, which can produce corrupt + # output. Gate on the block-table shape (all_block_ids is the + # snapshot of kv_cache_manager.get_block_ids -- the very table + # the recovery path unpacks), not on is_hybrid: a multi-group + # attention-only model is not hybrid yet still breaks the unpack. + # Remove the gating once upstream supports multi-group + # invalid-block recovery. + report_ids = [] + report_failures = len(load_req.all_block_ids) == 1 + if report_failures: + # With exactly one vLLM block table the transferred group is + # group 0, and it must be an attention group: a lone state + # group would feed state block ids into upstream's + # token-granular recovery math (it divides by block_size). + # Both hold for every supported single-group model today; + # assert so a future unsupported shape fails loudly instead + # of reporting nonsense. + only = self._group_metas[0] + assert isinstance(only, AttentionGroupMeta), ( + f"single vLLM block table but the transferred group is " + f"{type(only).__name__}, not attention; cannot report " + f"invalid blocks") + assert only.group_idx == 0, ( + f"single vLLM block table but the transferred group is " + f"group {only.group_idx}; group-index alignment broken") + table = load_req.all_block_ids[0] + gbs = only.block_size + report_ids = [table[(mb * self._manager_block_size) // gbs] + for mb in load_req.manager_block_idxes] + done_cb = self._data_transfer.create_load_done_callback( + load_req.req_id, self._tp_rank, meta.epoch, + copy.copy(report_ids), num_blocks, + report_failures=report_failures) + + if plans is None: + # Nothing submitted; report the whole load as failed. + mr = MultiResult(1, done_cb) + mr.submit_result(0, [False] * num_blocks * self._num_groups) + continue + + chunks = list(self._iter_task_chunks( + plans, self._extra_config.block_per_load_task)) + multi_result = MultiResult(len(chunks), done_cb) + for task_idx, chunk in enumerate(chunks): + self._data_transfer.submit_task( + self._data_transfer.load_task, multi_result, task_idx, *chunk) + + def wait_for_layer_load(self, layer_name: str) -> None: + pass + + def save_kv_layer(self, layer_name: str, kv_layer: torch.Tensor, + attn_metadata: "AttentionMetadata", **kwargs) -> None: + pass + + def wait_for_save(self, meta: TairKvCacheConnectorMetadata): + if not meta.to_save_requests: + return + ready_event = self._device_mod.Event() + ready_event.record(self._device_mod.current_stream()) + + for save_req in meta.to_save_requests: + num_blocks = len(save_req.manager_block_idxes) + plans = self._plan_group_transfers( + save_req.target_locations, save_req.manager_block_idxes, + save_req.all_block_ids) + + done_cb = self._data_transfer.create_save_done_callback( + save_req.req_id, self._tp_rank, save_req.write_session_id, num_blocks) + + if plans is None: + mr = MultiResult(1, done_cb) + mr.submit_result(0, [False] * num_blocks * self._num_groups) + continue + + chunks = list(self._iter_task_chunks( + plans, self._extra_config.block_per_save_task)) + multi_result = MultiResult(len(chunks), done_cb) + for task_idx, chunk in enumerate(chunks): + self._data_transfer.submit_task( + self._data_transfer.save_task, multi_result, task_idx, + *chunk, ready_event) + if self._tp_rank == 0: + # One session in flight; the coordinator's completion event + # for this req decrements it (see get_finished). + self._pending_saves[save_req.req_id] = \ + self._pending_saves.get(save_req.req_id, 0) + 1 + + def on_save_finished(self, write_session_id: str, save_context: SaveContext): + for block_idx in range(len(save_context.locations)): + fully_saved = all(save_context.result_per_rank[rank][block_idx] + for rank in range(self._tp_size)) + save_context.success_mask.append(fully_saved) + logger.debug("finish_write_cache mask:%s session:%s", + save_context.success_mask, write_session_id) + try: + self._manager_client.finish_write_cache({ + "trace_id": "finish_%s" % write_session_id[:8], + "instance_id": self._extra_config.instance_id, + "write_session_id": write_session_id, + "success_blocks": {"bool_masks": {"values": save_context.success_mask}}, + }) + except Exception as e: + logger.warning("finish_write_cache failed, session: %s, error: %s", + write_session_id, e) + + def get_finished(self, finished_req_ids: set, + meta: TairKvCacheConnectorMetadata) -> Tuple[Optional[set], Optional[set]]: + """Report request completion to vLLM. + + A request is reported once its last save session has settled on the + coordinator (tp0 aggregates the per-rank verdicts into + finish_write_cache). The scheduler's FinishRequest -- sent when its + own ledger says the request is settled -- either reports immediately + or defers until the last session lands.""" + if self._tp_rank != 0: + for finish_req in meta.to_finish_requests: + self._pending_saves.pop(finish_req.req_id, None) + self._finish_pending.discard(finish_req.req_id) + return None, None + + finished_saving = [] + finished_saving_tasks, finished_loading_tasks = self._coordinator_server.get_finished_tasks() + for req_id in finished_saving_tasks: + remaining = self._pending_saves.get(req_id, 0) - 1 + if remaining > 0: + self._pending_saves[req_id] = remaining + continue + self._pending_saves.pop(req_id, None) + if req_id in self._finish_pending: + self._finish_pending.discard(req_id) + finished_saving.append(req_id) + + for finish_req in meta.to_finish_requests: + if finish_req.req_id not in self._pending_saves: + finished_saving.append(finish_req.req_id) + else: + self._finish_pending.add(finish_req.req_id) + return set(finished_saving), set(finished_loading_tasks) + + def get_block_ids_with_load_errors(self) -> set: + if self._tp_rank != 0: + return set() + failed = self._coordinator_server.get_failed_loading_block_idxs() + if failed: + logger.warning("block_ids_with_load_errors: %s", failed) + return failed diff --git a/kv_cache_manager/py_connector/vllm/data_transfer.py b/kv_cache_manager/py_connector/vllm/data_transfer.py index 36a16323b..dd824a76d 100644 --- a/kv_cache_manager/py_connector/vllm/data_transfer.py +++ b/kv_cache_manager/py_connector/vllm/data_transfer.py @@ -1,15 +1,42 @@ -import time -import threading -from concurrent.futures.thread import ThreadPoolExecutor +"""Per-group KV cache transfer between vLLM's paged cache and KVCM storage. + +Each ``TransferGroup`` is an independent transfer unit: -from typing import Any +* Attention groups store token-granular KV; a manager block is gathered/scattered + through the strided Triton kernel (``batch_gather_scatter_helper``) which handles + both the contiguous full-attention layout and the block-strided hybrid layout. +* Mamba/linear/gdn groups store per-block opaque state; a manager block maps to a + single logical block whose raw bytes are copied verbatim. + +The transport itself is layout-agnostic and **zero-VRAM**: for every manager +block we hand the SDK a ``BlockBuffer`` (a pinned CPU region) and the block's +remote URI. Save gathers HBM -> pinned host directly and ``SaveKvCaches``; +load ``LoadKvCaches`` -> pinned host, then scatters pinned -> HBM directly. +The gather/scatter kernel addresses host pinned memory over PCIe (UVA +zero-copy), and state-group copies use plain ``copy_`` between the GPU tensors +and the pinned slices -- no device-side staging buffer exists anywhere on the +data path, so the connector competes with the engine for exactly zero HBM. +""" + +import threading +import time +from concurrent.futures import ThreadPoolExecutor import torch from kv_cache_manager.client.pybind import kvcm_py_client +from kv_cache_manager.py_connector.common.tp_coordinator import ( + CoordinateMsgSerializer, TpCoordinatorClient, CoordinateMessage, + SendBlockFinishedEvent, LoadBlockFinishedEvent, +) +from kv_cache_manager.py_connector.common.logger import logger +from kv_cache_manager.py_connector.vllm.transfer_types import ( + AttentionTransferGroup, KVCacheInfo, StateTransferGroup, TransferGroup) +from kv_cache_manager.py_connector.kernel import batch_gather_scatter_helper + def _get_device_module(device=None): - """Return the device module matching the runtime device.""" + """Return the torch device module matching the runtime device.""" if device is not None and hasattr(torch, "get_device_module"): return torch.get_device_module(device) try: @@ -20,25 +47,96 @@ def _get_device_module(device=None): pass return torch.cuda -from kv_cache_manager.py_connector.common.tp_coordinator import CoordinateMsgSerializer, TpCoordinatorClient, \ - CoordinateMessage, SendBlockFinishedEvent, LoadBlockFinishedEvent -from kv_cache_manager.py_connector.common.logger import logger -from kv_cache_manager.py_connector.common.types import KVCacheInfo -from kv_cache_manager.py_connector.kernel import batch_gather_scatter_helper -from kv_cache_manager.py_connector.kernel.gather_scatter_helper import CopyBufferAllocator +class _StagingPool: + """Bounded, pre-allocated *pinned host* slots for one TransferGroup. -class MultiResult: - """多任务结果管理类 - - 用于管理多个异步任务的结果, 当所有任务完成时触发回调 + Save and load share one pool. The slots are both the SDK's I/O buffers + and the kernel's gather/scatter target: the strided kernel and plain + ``copy_`` read/write host pinned memory directly over PCIe (UVA + zero-copy), so the pool needs no device-side mirror and the connector + reserves zero HBM. An exhausted pool blocks the acquiring task -- + backpressure -- instead of failing. Slots are handed out as *contiguous* + runs because the kernel view needs one piece of memory. + + Slots are reused once their task reports. This is the historical + origin/main behaviour: after an SDK timeout/error a background DMA may + still touch the buffer for a while (the deadline contract is not + upstream yet), and a slot reused in that window can be scribbled on -- + an accepted trade-off for removing the GPU staging copy. """ + + def __init__(self, device, per_block_bytes: int, max_blocks: int): + if max_blocks <= 0: + raise ValueError("staging pool must have at least one block slot") + self.block_bytes = per_block_bytes + self.max_blocks = max_blocks + self._cond = threading.Condition() + total = max_blocks * per_block_bytes + # Pinned host memory needs a CUDA context; on other devices (tests, + # CPU-only runs) fall back to pageable memory. + self._cpu = torch.empty(total, dtype=torch.uint8, device="cpu", + pin_memory=(device.type == "cuda")) + # Free runs as [start, start+len) block ranges, kept sorted by start. + self._runs = [[0, max_blocks]] + + def acquire(self, n: int) -> int: + """Block until a contiguous run of ``n`` block slots is free; return + its starting block index.""" + if n <= 0: + raise ValueError("must acquire at least one block slot") + if n > self.max_blocks: + raise ValueError( + f"staging task of {n} blocks exceeds the pool capacity " + f"{self.max_blocks}; raise staging_pool_blocks or shrink " + f"block_per_save_task/block_per_load_task") + with self._cond: + while True: + for i, (start, length) in enumerate(self._runs): + if length >= n: + rest = length - n + if rest: + self._runs[i] = [start + n, rest] + else: + self._runs.pop(i) + return start + self._cond.wait() + + def release(self, start: int, n: int) -> None: + with self._cond: + pos, run = 0, [start, n] + for pos, (s, _len) in enumerate(self._runs): + if s > start: + break + else: + pos = len(self._runs) + self._runs.insert(pos, run) + # Merge with the neighbours the release just glued together. + merged = [] + for s, l in self._runs: + if merged and merged[-1][0] + merged[-1][1] == s: + merged[-1][1] += l + else: + merged.append([s, l]) + self._runs = merged + self._cond.notify_all() + + def cpu_view(self, start: int, n: int) -> torch.Tensor: + b = self.block_bytes + return self._cpu[start * b:(start + n) * b] + + +class MultiResult: + """Collect the per-block success flags of several async tasks and fire a + callback once every task has reported. Each result is a list[bool] aligned + with the manager blocks the task handled (in submission order).""" + def __init__(self, size: int, callback): - self._size: int = size + self._size = size self._results = [None] * size self._lock = threading.Lock() - self._finished_num: int = 0 - self._finished_callback = callback + self._finished_num = 0 + self._callback = callback def submit_result(self, idx: int, result): with self._lock: @@ -46,247 +144,429 @@ def submit_result(self, idx: int, result): self._results[idx] = result self._finished_num += 1 if self._finished_num == self._size: - self._finished_callback(self._results) + # Flatten in submission order. Every slot is filled by the + # count check (submission asserts the slot was None), so the + # per-part guard only documents the invariant for readers. + flat = [] + for part in self._results: + assert part is not None + flat.extend(part) + self._callback(flat) -class DataTransferManager: - """KVCache数据传输核心类 - - 负责实际的KV缓存保存和加载操作, 包括: - 1. 保存任务（save_task） - 2. 加载任务（load_task） - 3. 回调创建（_create_save_done_callback, _create_load_done_callback） +def _effective_pool_blocks(configured, need, block_bytes, max_bytes): + """Blocks per staging-pool group under the per-group pinned-RAM ceiling. + + min(configured, max_bytes // block_bytes), floored at one full task + batch: a task stages its whole batch contiguously, so the pool can never + be smaller than the batch -- even when the byte cap alone would say so + (the caller warns in that case). Pure so the sizing contract is + unit-pinnable. """ - - def __init__(self, - kvcache_info: KVCacheInfo, - manager_block_size: int, - copy_buffer_allocator: CopyBufferAllocator, - transfer_client: Any, - coordinator_client: TpCoordinatorClient, - extra_config: Any): - """ - 初始化KV数据传输器 - - Args: - kvcache_info: KV缓存信息 - manager_block_size: instance的block_size - copy_buffer_allocator: 复制缓冲区分配器 - transfer_client: 传输客户端 - coordinator_client: 协调器客户端 - extra_config: 额外配置 - """ - self._kvcache_info = kvcache_info + by_bytes = max_bytes // block_bytes if block_bytes > 0 else configured + return max(need, min(configured, by_bytes)) + + +class DataTransferManager: + def __init__(self, kvcache_info: KVCacheInfo, manager_block_size: int, + transfer_client, coordinator_client: TpCoordinatorClient, extra_config): + self._info = kvcache_info self._manager_block_size = manager_block_size - self._copy_buffer_allocator = copy_buffer_allocator self._transfer_client = transfer_client self._coordinator_client = coordinator_client self._extra_config = extra_config - self._device_mod = _get_device_module(self._kvcache_info.device) - - # 创建内部线程池执行器 - self._io_executor = self._create_io_executor() - - # 保存和加载流 + self._device = kvcache_info.device + self._device_mod = _get_device_module(self._device) self._save_stream = self._device_mod.Stream() self._load_stream = self._device_mod.Stream() - - def _create_io_executor(self) -> ThreadPoolExecutor: - """创建IO线程池执行器""" - from concurrent.futures import ThreadPoolExecutor - - # 初始化线程池，设置线程名和初始化函数 - def init_worker(): - import torch - self._device_mod.set_device(self._kvcache_info.device) - - return ThreadPoolExecutor( - max_workers=32, - thread_name_prefix="kvcm_io_", - initializer=init_worker - ) - - def submit_task(self, func, *args, **kwargs): - """提交任务到内部线程池 - - Args: - func: 要执行的函数 - *args: 函数参数 - **kwargs: 函数关键字参数 - - Returns: - Future对象 - """ - return self._io_executor.submit(func, *args, **kwargs) - def load_task(self, multi_result: MultiResult, task_idx, remote_uris, block_token_indices): - """加载任务 - - Args: - multi_result: 多任务结果管理器 - task_idx: 任务索引 - remote_uris: 远程URI列表 - block_token_indices: 块令牌索引列表 - """ - logger.debug("load remote_uris:%s, block_token_indices:%s", remote_uris, block_token_indices) + pool_blocks = extra_config.staging_pool_blocks + need = max(extra_config.block_per_save_task, + extra_config.block_per_load_task) + if pool_blocks < need: + raise ValueError( + f"staging_pool_blocks={pool_blocks} is smaller than the " + f"largest task batch ({need}); one task stages its whole " + f"batch contiguously, so the pool must cover it -- raise " + f"staging_pool_blocks or shrink block_per_save_task/" + f"block_per_load_task") + # One pool per group: block shapes differ between attention and state + # groups. The pool is pinned host memory only -- the kernel reaches it + # over PCIe -- so the connector's device-memory footprint is zero. + # The byte cap keeps the configured block count honest for groups + # with large blocks (hybrid state/attention blocks are ~17.3 MiB vs + # ~0.875 MiB for full attention): the same count would otherwise pin + # ~17 GiB of host RAM per group and engine start dies in the pinned + # allocator. + self._pools = {} + for g in kvcache_info.groups: + blocks = _effective_pool_blocks( + pool_blocks, need, g.per_block_bytes, + extra_config.staging_pool_max_bytes_per_group) + if blocks < pool_blocks: + logger.warning( + "staging pool %s: %d blocks x %d bytes would pin " + "%.1f GiB; capped to %d blocks (%.1f MiB) by " + "staging_pool_max_bytes_per_group=%d", + g.spec_name, pool_blocks, g.per_block_bytes, + pool_blocks * g.per_block_bytes / 2**30, blocks, + blocks * g.per_block_bytes / 2**20, + extra_config.staging_pool_max_bytes_per_group) + self._pools[g.spec_name] = _StagingPool( + self._device, g.per_block_bytes, blocks) + for name, pool in self._pools.items(): + logger.info("staging pool %s: %d blocks x %d bytes " + "(pinned %.1f MiB, GPU 0)", + name, pool.max_blocks, + pool.block_bytes, + pool.max_blocks * pool.block_bytes / 2**20) + + def _init_worker(): + self._device_mod.set_device(self._device) - copy_buffer_indices = self._copy_buffer_allocator.alloc_buffer_idx_blocking(len(remote_uris)) - copy_buffers = self._copy_buffer_allocator.get_buffer_by_idx(copy_buffer_indices) + self._io_executor = ThreadPoolExecutor( + max_workers=32, thread_name_prefix="kvcm_io_", initializer=_init_worker) + + def submit_task(self, func, *args, **kwargs): + return self._io_executor.submit(func, *args, **kwargs) + # ------------------------------------------------------------------ # + # BlockBuffer helper + # ------------------------------------------------------------------ # + @staticmethod + def _make_block_buffers(base_ptr: int, per_block_bytes: int, count: int): buffers = [] - for copy_buffer in copy_buffers: - buffer = kvcm_py_client.BlockBuffer() - iovs = [] + for i in range(count): + buf = kvcm_py_client.BlockBuffer() iov = kvcm_py_client.Iov() iov.type = kvcm_py_client.MemoryType.CPU - iov.base = copy_buffer.data_ptr() - iov.size = copy_buffer.nbytes + iov.base = base_ptr + i * per_block_bytes + iov.size = per_block_bytes iov.ignore = False - iovs.append(iov) - buffer.iovs = iovs - buffers.append(buffer) - logger.debug("start transfer") - transfer_result = self._transfer_client.LoadKvCaches(remote_uris, buffers) - logger.debug("done transfer,result:%s", transfer_result) - if transfer_result == kvcm_py_client.ClientErrorCode.ER_OK: - with self._device_mod.stream(self._load_stream): - batch_gather_scatter_helper.batch_scatter_kv_caches( - self._kvcache_info.all_kvcache_ptr_tensor_gpu, - self._copy_buffer_allocator._raw_buffer, - block_token_indices, - copy_buffer_indices, - self._manager_block_size, - self._kvcache_info.per_token_per_layer_dim_size, - ) - - copy_done_event = self._device_mod.Event() - copy_done_event.record(self._load_stream) - copy_done_event.synchronize() - - logger.debug("done scatter") - else: - logger.warning("load task failed, remote_uris:%s, block_token_indices:%s, transfer_result:%s", - remote_uris, - block_token_indices, transfer_result) - self._copy_buffer_allocator.free_buffer(copy_buffer_indices) - multi_result.submit_result(task_idx, [transfer_result] * len(remote_uris)) - - def create_load_done_callback(self, req_id, tp_rank, epoch, local_block_ids): - """创建加载完成回调函数 - - Args: - req_id: 请求ID - tp_rank: TP rank - epoch - local_block_ids: 本地块ID列表 - - Returns: - 回调函数 + buf.iovs = [iov] + buffers.append(buf) + return buffers + + # ------------------------------------------------------------------ # + # Save + # ------------------------------------------------------------------ # + def save_task(self, multi_result: MultiResult, task_idx, group: TransferGroup, + remote_uris, block_token_indices, block_ids, ready_event): + """Gather one group's manager blocks from HBM and save them. + + block_token_indices: attention -> list[list[int]] flat token slots per block. + block_ids: state -> list[int] block id per manager block. + remote_uris: positionally aligned with the manager blocks; None + where the manager allocated no location for this + group's spec (see ``_spec_groups``). + + A block is reported successful only when its data was actually written. + Where a state group has no state (vLLM's null block in mamba "align" + mode) the manager was already told so at start_write_cache time -- the + block simply carries no spec for this group -- so nothing is written and + nothing is claimed: the block is *excluded* from this group's verdict + rather than reported as a success. """ - def generate_message(task_results): - failed_block_idxs = [] - idx = 0 - for task_result in task_results: - for block_result in task_result: - if block_result != kvcm_py_client.ClientErrorCode.ER_OK: - failed_block_idxs.append(local_block_ids[idx]) - idx += 1 - - msg = CoordinateMessage( - time.time(), - LoadBlockFinishedEvent(request_id=req_id, tp_rank=tp_rank, - epoch=epoch, failed_block_idxs=failed_block_idxs) - ) - self._coordinator_client.send(CoordinateMsgSerializer.dumps(msg)) + n = len(remote_uris) + # Single exit: whatever happens below, the task reports, otherwise the + # MultiResult callback never fires (submit_task drops the future, so an + # escaping exception is silently swallowed) and the save session hangs. + ok_mask = [False] * n + try: + # Three dispositions per block: abstain (this group holds no data + # for it by design), transfer, or fail outright. + skipped, failed = self._save_dispositions(group, remote_uris, block_ids, n) + valid = [i for i in range(n) if i not in skipped and i not in failed] + ok_mask = [None if i in skipped else False for i in range(n)] + if valid: + self._save_valid_blocks(group, remote_uris, + block_token_indices, block_ids, + ready_event, valid, ok_mask) + except Exception: + # Fail the whole task: partially transferred blocks are unknown, so + # report conservatively -- never publish a block we cannot vouch + # for (abstentions included; that only costs hit rate, not truth). + logger.exception("save task crashed group=%s blocks=%d, failing them all", + group.spec_name, n) + ok_mask = [False] * n + multi_result.submit_result(task_idx, ok_mask) + + def _save_dispositions(self, group: TransferGroup, remote_uris, block_ids, n): + """Split the task's blocks into (abstained, failed); the rest transfer. + + A state group abstains for a manager block exactly when vLLM's block + table points at the null block *and* the manager allocated no URI for + the group's spec. The two must agree, because the scheduler derived the + announced spec coverage from that very block table. Either disagreement + fails the block: + + * location but no state -- writing the null block's bytes would publish + a state the model never produced; + * state but no location -- the state cannot be published at all. - return generate_message - - def save_task(self, multi_result: MultiResult, task_idx, remote_uris, block_token_indices, - kvcache_ready_event): - """保存任务 - - Args: - multi_result: 多任务结果管理器 - task_idx: 任务索引 - remote_uris: 远程URI列表 - block_token_indices: 块令牌索引列表 - kvcache_ready_event: KV缓存就绪事件 + Attention KV is never sparse, so a missing location simply fails. """ - logger.debug("save remote_uris:%s, block_token_indices:%s", remote_uris, block_token_indices) - - with self._device_mod.stream(self._save_stream): - kvcache_ready_event.wait() - copy_buffer_indices = self._copy_buffer_allocator.alloc_buffer_idx_blocking(len(remote_uris)) - batch_gather_scatter_helper.batch_gather_kv_caches( - self._kvcache_info.all_kvcache_ptr_tensor_gpu, - self._copy_buffer_allocator._raw_buffer, - block_token_indices, - copy_buffer_indices, - self._manager_block_size, - self._kvcache_info.per_token_per_layer_dim_size, - ) - copy_done_event = self._device_mod.Event() - copy_done_event.record(self._save_stream) - - copy_done_event.synchronize() - - logger.debug("done gather") - - copy_buffers = self._copy_buffer_allocator.get_buffer_by_idx(copy_buffer_indices) - buffers = [] - for copy_buffer in copy_buffers: - buffer = kvcm_py_client.BlockBuffer() - iovs = [] - iov = kvcm_py_client.Iov() - iov.type = kvcm_py_client.MemoryType.CPU - iov.base = copy_buffer.data_ptr() - iov.size = copy_buffer.nbytes - iov.ignore = False - iovs.append(iov) - buffer.iovs = iovs - buffers.append(buffer) - logger.debug("start transfer") - - transfer_result = self._transfer_client.SaveKvCaches(remote_uris, buffers) - logger.debug("done transfer,result:%s", transfer_result) - if transfer_result[0] != kvcm_py_client.ClientErrorCode.ER_OK: - logger.warning("save task failed, remote_uris:%s, block_token_indices:%s, transfer_result:%s", remote_uris, - block_token_indices, transfer_result) - - self._copy_buffer_allocator.free_buffer(copy_buffer_indices) - # TODO: submit uri when enable local alloc - multi_result.submit_result(task_idx, [transfer_result[0]] * len(remote_uris)) - - def create_save_done_callback(self, req_id, tp_rank, write_session_id): - """创建保存完成回调函数 - - Args: - req_id: 请求ID - tp_rank: TP rank - write_session_id: 写入会话ID - - Returns: - 回调函数 + if isinstance(group, AttentionTransferGroup): + failed = {i for i in range(n) if remote_uris[i] is None} + if failed: + logger.warning("save group %s: %d/%d attention blocks have no " + "location, failing them", group.spec_name, + len(failed), n) + return set(), failed + skipped, failed = set(), set() + for i in range(n): + is_null = block_ids[i] == 0 + has_uri = remote_uris[i] is not None + if is_null and not has_uri: + skipped.add(i) + elif is_null: + failed.add(i) + logger.warning("save group %s: block %d has a location but no " + "materialized state; failing it instead of " + "publishing bytes the model never produced", + group.spec_name, i) + elif not has_uri: + failed.add(i) + logger.warning("save group %s: block %d has a materialized " + "state but no location to write it to; failing it", + group.spec_name, i) + if skipped: + logger.debug("save group %s: %d/%d blocks carry no state " + "(not published)", group.spec_name, len(skipped), n) + return skipped, failed + + @staticmethod + def _load_skipped_blocks(group: TransferGroup, remote_uris, block_ids, n) -> set: + """Blocks this group has nothing to load into. + + Asymmetric with the save side on purpose: on load, a null target means + vLLM does not *need* this group's data for that block (in mamba "align" + mode only the block ending the reused prefix needs its state), so the + block is skipped whatever the manager published. The reverse -- a real + target with nothing published -- is a genuine failure: the request would + run on an unwritten state. """ - def generate_message(task_results): - is_successes = [] - # TODO: report uri when enable local alloc - # remote_uris = [] - for task_result in task_results: - for block_result in task_result: - if block_result != kvcm_py_client.ClientErrorCode.ER_OK: - is_successes.append(False) - # remote_uris.append(None) - else: - is_successes.append(True) - # remote_uris.extend(future_result[1]) - - msg = CoordinateMessage( - time.time(), - SendBlockFinishedEvent(request_id=req_id, tp_rank=tp_rank, - write_session_id=write_session_id, - is_success_list=is_successes) - ) + if isinstance(group, AttentionTransferGroup): + missing = sum(uri is None for uri in remote_uris) + if missing: + logger.warning("load group %s: %d/%d attention blocks have no " + "location, failing them", group.spec_name, missing, n) + return set() + skipped = {i for i in range(n) if block_ids[i] == 0} + if skipped: + logger.debug("load group %s: %d/%d blocks need no state", + group.spec_name, len(skipped), n) + return skipped + + def _save_valid_blocks(self, group, remote_uris, + block_token_indices, block_ids, ready_event, + valid, ok_mask): + uris = [remote_uris[i] for i in valid] + assert all(uri is not None for uri in uris), \ + f"group {group.spec_name}: save batch contains a block without a " \ + f"location; _save_dispositions must have failed it" + # Wait for the engine's forward pass *before* taking a pool slot: + # the slot is needed only for the gather + SDK transfer, so holding + # it while the model still runs only extends the pool occupancy and + # queues concurrent loads behind a save that is not even staging. + ready_event.wait() + pool = self._pools[group.spec_name] + start = pool.acquire(len(valid)) + try: + cpu_buffer = pool.cpu_view(start, len(valid)) + with self._device_mod.stream(self._save_stream): + # Gather straight into the pinned host slot: the kernel + # (attention) and copy_ (state) write host pinned memory + # directly over PCIe; no device-side staging copy exists. + if isinstance(group, AttentionTransferGroup): + view = cpu_buffer.view(self._info.dtype).view( + len(valid), group.num_kv_ptrs, + self._manager_block_size, group.per_token_dim) + batch_gather_scatter_helper.batch_gather_kv_caches( + group.kvcache_ptr_tensor_gpu, view, + [block_token_indices[i] for i in valid], + list(range(len(valid))), self._manager_block_size, + group.per_token_dim, + block_stride=group.block_stride, + local_block_size=group.kernel_block_size) + else: + for out_i, i in enumerate(valid): + for layer_idx in range(group.layer_num): + dst = (out_i * group.layer_num + layer_idx) * group.page_size_bytes + cpu_buffer[dst:dst + group.page_size_bytes].copy_( + group.block_view_tensors[layer_idx][block_ids[i]], + non_blocking=True) + done = self._device_mod.Event() + done.record(self._save_stream) + done.synchronize() + + buffers = self._make_block_buffers( + cpu_buffer.data_ptr(), group.per_block_bytes, len(valid)) + result = self._transfer_client.SaveKvCaches(uris, buffers) + ok = (result[0] == kvcm_py_client.ClientErrorCode.ER_OK) + if not ok: + logger.warning("save task failed group=%s uris=%d result=%s", + group.spec_name, len(uris), result) + for i in valid: + ok_mask[i] = ok + except BaseException: + # Drain the stream before the slots go back: a failed task may + # have left kernel/copy work enqueued against the staging views. + self._save_stream.synchronize() + raise + finally: + pool.release(start, len(valid)) + + def create_save_done_callback(self, req_id, tp_rank, write_session_id, num_blocks): + """block success = AND across all groups that had data for the block. + + Task results are ordered group0[blocks], group1[blocks], ... so a + block's verdict is the stride-AND ``flat[i % num_blocks]``. ``None`` + entries mean "this group holds no data for this block by design" (see + save_task) and are skipped: they neither pass nor fail the block. A + block whose every group is None was never written at all and must not + be published. + """ + def cb(flat): + is_success = [None] * num_blocks + for i, ok in enumerate(flat): + if ok is None: + continue + b = i % num_blocks + is_success[b] = ok if is_success[b] is None else (is_success[b] and ok) + # No group wrote anything for the block -> nothing to publish. + is_success = [bool(ok) for ok in is_success] + msg = CoordinateMessage(time.time(), SendBlockFinishedEvent( + request_id=req_id, tp_rank=tp_rank, + write_session_id=write_session_id, is_success_list=is_success)) self._coordinator_client.send(CoordinateMsgSerializer.dumps(msg)) + return cb - return generate_message + # ------------------------------------------------------------------ # + # Load + # ------------------------------------------------------------------ # + def load_task(self, multi_result: MultiResult, task_idx, group: TransferGroup, + remote_uris, block_token_indices, block_ids): + """Load one group's manager blocks from storage into HBM. + + Mirror of ``save_task``: ``remote_uris`` is positionally aligned with + the manager blocks and holds None where the published block carries no + data for this group's spec. A state group's block is skipped only when + vLLM's target is the null block *and* nothing was published -- i.e. the + state is neither needed nor available. Everything else must transfer. + """ + n = len(remote_uris) + # Single exit, as in save_task: an escaping exception would silently + # swallow the MultiResult callback, which under the connector's + # synchronous-load contract leaves vLLM believing KV it never received. + ok_mask = [False] * n + try: + skipped = self._load_skipped_blocks(group, remote_uris, block_ids, n) + # A block we must restore but nothing was published for cannot be + # loaded; fail it without letting it shift the staging batch. + failed = {i for i in range(n) + if i not in skipped and remote_uris[i] is None} + valid = [i for i in range(n) if i not in skipped and i not in failed] + ok_mask = [None if i in skipped else False for i in range(n)] + if valid: + ok = self._load_valid_blocks(group, remote_uris, + block_token_indices, block_ids, + valid) + for i in valid: + ok_mask[i] = ok + except Exception: + logger.exception("load task crashed group=%s blocks=%d, failing them all", + group.spec_name, n) + ok_mask = [False] * n + multi_result.submit_result(task_idx, ok_mask) + + def _load_valid_blocks(self, group, remote_uris, block_token_indices, + block_ids, valid) -> bool: + pool = self._pools[group.spec_name] + start = pool.acquire(len(valid)) + try: + cpu_buffer = pool.cpu_view(start, len(valid)) + buffers = self._make_block_buffers(cpu_buffer.data_ptr(), + group.per_block_bytes, len(valid)) + uris = [remote_uris[i] for i in valid] + assert all(uri is not None for uri in uris), \ + f"group {group.spec_name}: load batch contains a block without a " \ + f"location; load_task must have failed it" + result = self._transfer_client.LoadKvCaches(uris, buffers) + ok = (result == kvcm_py_client.ClientErrorCode.ER_OK) + if ok: + with self._device_mod.stream(self._load_stream): + # Scatter straight out of the pinned host slot: the + # kernel (attention) and copy_ (state) read host pinned + # memory directly over PCIe; no device-side staging copy. + if isinstance(group, AttentionTransferGroup): + view = cpu_buffer.view(self._info.dtype).view( + len(valid), group.num_kv_ptrs, + self._manager_block_size, group.per_token_dim) + batch_gather_scatter_helper.batch_scatter_kv_caches( + group.kvcache_ptr_tensor_gpu, view, + [block_token_indices[i] for i in valid], + list(range(len(valid))), self._manager_block_size, + group.per_token_dim, + block_stride=group.block_stride, + local_block_size=group.kernel_block_size) + else: + for out_i, i in enumerate(valid): + for layer_idx in range(group.layer_num): + src = (out_i * group.layer_num + layer_idx) * group.page_size_bytes + group.block_view_tensors[layer_idx][block_ids[i]].copy_( + cpu_buffer[src:src + group.page_size_bytes], + non_blocking=True) + done = self._device_mod.Event() + done.record(self._load_stream) + done.synchronize() + else: + logger.warning("load task failed group=%s uris=%d result=%s", + group.spec_name, len(uris), result) + return ok + except BaseException: + # Drain the stream before the slots go back (as in save). + self._load_stream.synchronize() + raise + finally: + pool.release(start, len(valid)) + + def create_load_done_callback(self, req_id, tp_rank, epoch, block_ids, num_blocks, + report_failures=True): + """A manager block is loaded only if every group that had data for it + succeeded. + + ``None`` entries mean "this group has no data for this block by design" + (mamba "align" interior blocks, see load_task) and are skipped, exactly + as in the save callback. + + block_ids is the block table used to report vLLM-visible invalid block + ids. vLLM's invalid-block recovery unpacks a single-group block table + (https://github.com/vllm-project/vllm/blob/releases/v0.26.0/vllm/v1/core/sched/scheduler.py#L2693, + "TODO (davidb): add support for hybrid memory allocator"); for a + multi-group (hybrid) model the unpack raises ValueError and crashes + the scheduler, so hybrid connectors pass report_failures=False: the + failure is logged, vLLM keeps the partially loaded KV, and the request + may produce corrupt output -- the contained alternative to an + engine-wide crash. See start_load_kv for the full trade-off.""" + def cb(flat): + merged = [None] * num_blocks + for i, ok in enumerate(flat): + if ok is None: + continue + b = i % num_blocks + merged[b] = ok if merged[b] is None else (merged[b] and ok) + # A block no group loaded anything for was not restored. + merged = [bool(ok) for ok in merged] + failed = [] + if report_failures: + failed = [block_ids[i] for i in range(min(num_blocks, len(block_ids))) + if not merged[i]] + elif not all(merged): + logger.warning("load failed for %d/%d blocks of req %s (hybrid: " + "not reporting invalid block ids)", + merged.count(False), num_blocks, req_id) + msg = CoordinateMessage(time.time(), LoadBlockFinishedEvent( + request_id=req_id, tp_rank=tp_rank, epoch=epoch, failed_block_idxs=failed)) + self._coordinator_client.send(CoordinateMsgSerializer.dumps(msg)) + return cb diff --git a/kv_cache_manager/py_connector/vllm/location_query_manager.py b/kv_cache_manager/py_connector/vllm/location_query_manager.py index 0d5ce5ba2..0e62ea281 100644 --- a/kv_cache_manager/py_connector/vllm/location_query_manager.py +++ b/kv_cache_manager/py_connector/vllm/location_query_manager.py @@ -1,193 +1,190 @@ -"""缓存管理模块，包含TairKvCacheConnector的本地缓存管理逻辑""" +"""Manager-side location queries for the external-match hook. + +``get_num_new_matched_tokens`` must never block the scheduler loop, so +queries run on the http executor and the answer is cached per request +until consumed. The cache is request-lifecycle scoped -- produced by the +match hook, consumed by ``update_state_after_alloc`` (which turns the +answer into a LoadRequest) and invalidated when the request retires. No +TTL: the producer and the consumer are two hooks of the same request's +life, so nothing can outlive its usefulness for long. + +**One slot per request, superseded by the newest ask.** The scheduler +asks twice per request's life: ``get_num_new_matched_tokens`` (with the +offset) and ``update_state_after_alloc`` (without it). vLLM's contract +fills the gap: the alloc hook always consumes the answer of the *last* +match hook that returned a hit -- it passes the matched token count back +verbatim. So the cache needs no per-offset addressing; a new ask at a +different offset is a *different* query that supersedes the old one, and +only the newest ask's answer may ever be consumed. + +**Object identity is the version.** A superseded (or invalidated / +consumed) query must not write its late answer into the new slot. Each +ask captures the slot object it created; the async callback compares +identity before writing. A monotonic version number is kept for logs +only -- correctness never depends on it. +""" -import enum -import time import threading from dataclasses import dataclass -from typing import Any, Tuple, Dict +from typing import Dict, Optional, Tuple + from kv_cache_manager.py_connector.common.manager_client import KvCacheManagerClient from kv_cache_manager.py_connector.common.logger import logger +QUERY_TYPE = "QT_PREFIX_MATCH" + @dataclass(frozen=True) class QueryCacheKey: + """Identity of one location query; different offsets are different queries.""" + req_id: str query_type: str token_length: int - computed_manager_block_size: int - - -class QueryCacheStatus(enum.Enum): - RUNNING = 0 - FINISHED = 1 - NOT_FOUND = 2 + computed_blocks: int @dataclass -class QueryCacheValue: - locations: list[Any] - query_time: float - is_done: bool = False +class _Query: + """The request's one active query. ``locations is None`` = in flight.""" + + key: QueryCacheKey + version: int = 0 + locations: list = None # None until the manager answered class LocationQueryManager: - """Location请求管理器，负责Location请求和Location信息缓存的管理""" + """Per-request location query cache: produce on match, consume on alloc. - def __init__(self, manager_client: KvCacheManagerClient, http_executor, instance_id: str, - async_get_cache_location: bool): - """ - 初始化缓存管理器 - - Args: - manager_client: Manager客户端 - http_executor: HTTP执行器 - instance_id: 实例ID - async_get_cache_location: 是否启用异步GetCacheLocation - """ + One ``_Query`` slot per request. Re-asking the slot's own offset + deduplicates (in flight -> wait, answered -> serve); asking a different + offset *supersedes* the slot (the newest ask wins, older asks are + dropped wherever they are in their lifetime). + """ + + def __init__(self, manager_client: KvCacheManagerClient, http_executor, + instance_id: str, async_get_cache_location: bool): self._manager_client = manager_client self._http_executor = http_executor self._instance_id = instance_id self._async_get_cache_location = async_get_cache_location - - self._local_query_cache_lock = threading.Lock() - self._local_query_cache: Dict[QueryCacheKey, QueryCacheValue] = {} - - self._cleanup_stop_event = threading.Event() - self._cleanup_thread = threading.Thread(target=self._cleanup_expired_local_cache, daemon=True) - self._cleanup_thread.start() + self._lock = threading.Lock() + # req_id -> the request's one active query slot. + self._queries: Dict[str, _Query] = {} def shutdown(self): - """关闭缓存管理器""" - self._cleanup_stop_event.set() - if self._cleanup_thread.is_alive(): - self._cleanup_thread.join(timeout=2.0) - - def _cleanup_expired_local_cache(self): - """清理过期的本地缓存条目""" - while not self._cleanup_stop_event.wait(1.0): - try: - current_time = time.time() - expired_keys = [] - - with self._local_query_cache_lock: - for key, value in self._local_query_cache.items(): - # TODO: editable ttl - if current_time - value.query_time > 1: - expired_keys.append(key) - for key in expired_keys: - self._local_query_cache.pop(key) - - if expired_keys: - logger.debug("Cleaned up %d expired query cache entries", len(expired_keys)) - except Exception as e: - logger.warning("Error during query cache cleanup: %s", e) + pass - def _get_cache_from_manager(self, request, computed_manager_block_size: int, query_key: QueryCacheKey): - """ - 从管理器获取缓存位置 - - Args: - request: 请求对象 - computed_manager_block_size: 本地已命中的block数量 - query_key: 查询键 - """ - try: - get_request = { - "trace_id": request.request_id, - "token_ids": request.prompt_token_ids, - "instance_id": self._instance_id, - "query_type": "QT_PREFIX_MATCH", - "block_mask": { - "offset": computed_manager_block_size - } - } - logger.debug("get_kvcache_location request: %s", get_request) - result = self._manager_client.get_cache_location(get_request) - logger.debug("get_kvcache_location result: %s", result) - need_load_locations = result["locations"] - with self._local_query_cache_lock: - if query_key not in self._local_query_cache: - logger.warning("_local_query_cache not found %s when request finished", request.request_id) - return - self._local_query_cache[query_key].locations = need_load_locations - self._local_query_cache[query_key].is_done = True - except Exception as e: - logger.warning("get_cache_location error, request_id: %s, error: %s", request.request_id, e) - with self._local_query_cache_lock: - if query_key in self._local_query_cache: - self._local_query_cache.pop(query_key) - - def _try_get_locations_from_local_cache(self, query_key: QueryCacheKey) -> Tuple[QueryCacheStatus, list]: - """ - 尝试从本地缓存获取位置 - - Args: - query_key: 查询键 - - Returns: - (状态, 位置列表) - """ - with self._local_query_cache_lock: - if query_key not in self._local_query_cache: - return QueryCacheStatus.NOT_FOUND, [] - query_value = self._local_query_cache[query_key] - # TODO: editable ttl - if time.time() - query_value.query_time > 1: - # cache timeout - self._local_query_cache.pop(query_key) - return QueryCacheStatus.NOT_FOUND, [] - if query_value.is_done: - return QueryCacheStatus.FINISHED, query_value.locations - else: - return QueryCacheStatus.RUNNING, [] - - def get_locations_for_query(self, request, computed_manager_block_size: int) -> Tuple[bool, list]: - """ - 获取查询的位置信息 - - Args: - request: 请求对象 - computed_manager_block_size: 本地已命中的block数量 - - Returns: - (查询是否完成, 位置列表) - """ - # TODO: async get_kvcache_location - query_key = QueryCacheKey( + @staticmethod + def _key(request, computed_blocks: int) -> QueryCacheKey: + return QueryCacheKey( req_id=request.request_id, - query_type="QT_PREFIX_MATCH", + query_type=QUERY_TYPE, token_length=len(request.prompt_token_ids), - computed_manager_block_size=computed_manager_block_size) - try: - status, locations = self._try_get_locations_from_local_cache(query_key) - if status == QueryCacheStatus.RUNNING: - return False, [] - elif status == QueryCacheStatus.FINISHED: - return True, locations - - # status == QueryCacheStatus.NOT_FOUND - # insert new cache - with self._local_query_cache_lock: - self._local_query_cache[query_key] = QueryCacheValue( - locations=[], query_time=time.time(), is_done=False) - if self._async_get_cache_location: - future = self._http_executor.submit(self._get_cache_from_manager, request, computed_manager_block_size, - query_key) - return False, [] - else: - self._get_cache_from_manager(request, computed_manager_block_size, query_key) - # only sync call need check again - status, locations = self._try_get_locations_from_local_cache(query_key) - if status == QueryCacheStatus.NOT_FOUND: - # do_get_cache_locations error, bypass load - return True, [] - elif status == QueryCacheStatus.RUNNING: - return False, [] - elif status == QueryCacheStatus.FINISHED: - return True, locations + computed_blocks=computed_blocks) + + def _fetch_from_manager(self, request, computed_blocks: int): + """Run the actual GetCacheLocation call (http thread or inline).""" + get_request = { + "trace_id": request.request_id, + "token_ids": request.prompt_token_ids, + "instance_id": self._instance_id, + "query_type": QUERY_TYPE, + "block_mask": {"offset": computed_blocks}, + } + logger.debug("get_kvcache_location request: %s", get_request) + result = self._manager_client.get_cache_location(get_request) + logger.debug("get_kvcache_location result: %s", result) + return result["locations"] + + def _query_async(self, request, key: QueryCacheKey, q: _Query) -> None: + def run(): + try: + locations = self._fetch_from_manager(request, key.computed_blocks) + except Exception as e: + logger.warning("get_cache_location error, request_id: %s, error: %s", + request.request_id, e) + with self._lock: + # Drop the slot: the next match hook re-issues the query. + if self._queries.get(request.request_id) is q: + self._queries.pop(request.request_id, None) + return + with self._lock: + if self._queries.get(request.request_id) is q: + q.locations = locations else: - logger.warning("unknown local cache query status: %s", status) - return True, [] + # Superseded / consumed / invalidated meanwhile: the + # answer of a dead ask must never write into the new + # slot (late answers only win against older asks, never + # against newer ones). + logger.debug("get_cache_location answer dropped: request %s " + "query v%d was superseded", + request.request_id, q.version) + + self._http_executor.submit(run) + + def get_locations_for_query(self, request, computed_blocks: int) -> Optional[list]: + """Ask (or re-ask) for the request's external match at this offset. + + Returns the locations when the answer is already cached for this + exact query key, None while a query is in flight for it (the + scheduler re-asks next step), and [] when the query answered "no + hit". An ask at a different offset supersedes the slot: a new + query starts immediately, and the old one's answer -- arrived or + still in flight -- can no longer be consumed. A failed query drops + its slot so the next hook call re-issues. + """ + req_id = request.request_id + key = self._key(request, computed_blocks) + with self._lock: + q = self._queries.get(req_id) + if q is not None and q.key == key: + # Same query: wait while in flight (dedupe -- the scheduler + # re-asks every step), serve the cached answer otherwise. + return None if q.locations is None else q.locations + # First ask, or a different offset: the newest ask supersedes. + q = _Query(key=key, version=q.version + 1 if q is not None else 0) + self._queries[req_id] = q + + if self._async_get_cache_location: + self._query_async(request, key, q) + return None + try: + locations = self._fetch_from_manager(request, computed_blocks) except Exception as e: - logger.warning("get_locations_for_query error, request_id: %s, error: %s", request.request_id, e) - return True, [] + logger.warning("get_cache_location error, request_id: %s, error: %s", + req_id, e) + with self._lock: + if self._queries.get(req_id) is q: + self._queries.pop(req_id, None) + return [] + with self._lock: + if self._queries.get(req_id) is q: + q.locations = locations + return locations + + def store_result(self, req_id: str, locations: list) -> None: + """Overwrite the cached answer (the match hook clamps it to a + vLLM-safe prefix before the allocation consumes it). The slot is + necessarily the query the hook just got its answer from: the hook + runs on the scheduler thread and only it supersedes slots.""" + with self._lock: + q = self._queries.get(req_id) + if q is not None: + q.locations = locations + + def consume_locations(self, req_id: str) -> Optional[Tuple[list, int]]: + """Pop the request's answered query: (locations, computed_blocks), + or None when nothing is cached (no query was issued / still in + flight / already consumed).""" + with self._lock: + q = self._queries.pop(req_id, None) + if q is None or q.locations is None: + return None + return q.locations, q.key.computed_blocks + + def invalidate(self, req_id: str) -> None: + """Drop the request's query: it is going away, and any answer still + in flight must not write into a future slot.""" + with self._lock: + self._queries.pop(req_id, None) diff --git a/kv_cache_manager/py_connector/vllm/metadata.py b/kv_cache_manager/py_connector/vllm/metadata.py index d89d872b8..c30ea1ddd 100644 --- a/kv_cache_manager/py_connector/vllm/metadata.py +++ b/kv_cache_manager/py_connector/vllm/metadata.py @@ -1,59 +1,55 @@ from dataclasses import dataclass, field +from typing import List + from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata @dataclass class SaveRequest: req_id: str - target_locations: list[dict] - manager_block_idxes: list + # CacheLocation dicts returned by the manager (one per manager block to save), + # each carrying location_specs for every registered spec name. + target_locations: List[dict] + # Manager block indices (into the request's token stream) being saved. + manager_block_idxes: List[int] write_session_id: str + # Per-group block table snapshot taken by the scheduler when the save + # instruction is handed to the worker: the gather needs it to translate + # manager blocks into physical slots, and the scheduler's ledger is the + # only place it exists. The worker keeps no request mirror. + all_block_ids: List[List[int]] = field(default_factory=list) -@dataclass() +@dataclass class LoadRequest: req_id: str - manager_block_idxes: list - need_load_locations: list[dict] - local_block_ids: list = field(default_factory=list) + manager_block_idxes: List[int] + need_load_locations: List[dict] + # Per-group block tables: all_block_ids[group_idx] is the list of local + # block ids for that kv_cache_group (filled by the scheduler once vLLM + # allocated the blocks). Length 1 for pure-attention models. + all_block_ids: List[List[int]] = field(default_factory=list) -@dataclass() +@dataclass class FinishRequest: req_id: str -@dataclass -class ReqStateToWorker: - """发送给工作节点的请求状态数据结构""" - - req_id: str - has_saved_block_num: int - new_tokens_ids: list = field(default_factory=list) - new_local_block_ids: list = field(default_factory=list) - resumed_from_preemption: bool = False - is_delta: bool = True - @dataclass class TairKvCacheConnectorMetadata(KVConnectorMetadata): - """TairKvCacheConnector的元数据类，用于在调度器和工作节点之间传递状态""" - requests: list[ReqStateToWorker] + """Scheduler -> worker metadata for one engine step. - def __init__(self, epoch: int): - """ - 初始化元数据 + Three instruction kinds, each consumed by its own worker hook: + to_load_requests (start_load_kv), to_save_requests (wait_for_save) and + to_finish_requests (get_finished). Every instruction is self-contained + -- the worker keeps no per-request mirror of scheduler state.""" - Args: - epoch: 当前epoch编号 - """ + def __init__(self, epoch: int): self.epoch = epoch - self.requests: list[ReqStateToWorker] = [] - self.to_load_requests: list[LoadRequest] = [] - self.to_save_requests: list[SaveRequest] = [] - self.to_finish_requests: list[FinishRequest] = [] - - def add_req_state_to_worker(self, request: ReqStateToWorker): - self.requests.append(request) + self.to_load_requests: List[LoadRequest] = [] + self.to_save_requests: List[SaveRequest] = [] + self.to_finish_requests: List[FinishRequest] = [] def add_load_request(self, request: LoadRequest): self.to_load_requests.append(request) @@ -65,5 +61,7 @@ def add_finish_request(self, finish_request: FinishRequest): self.to_finish_requests.append(finish_request) def __repr__(self): - return f"TairKvCacheConnectorMetadata(requests={self.requests})" - + return (f"TairKvCacheConnectorMetadata(epoch={self.epoch}, " + f"load={len(self.to_load_requests)}, " + f"save={len(self.to_save_requests)}, " + f"finish={len(self.to_finish_requests)})") diff --git a/kv_cache_manager/py_connector/vllm/transfer_types.py b/kv_cache_manager/py_connector/vllm/transfer_types.py new file mode 100644 index 000000000..f143e2c1c --- /dev/null +++ b/kv_cache_manager/py_connector/vllm/transfer_types.py @@ -0,0 +1,121 @@ +"""Worker-side transfer data model, one type per KV cache group kind. + +The vLLM connector treats every ``kv_cache_group`` as an independent +transfer unit. Attention groups and state (mamba/linear) groups transfer +dissimilar data through dissimilar mechanics -- token-granular gather / +scatter through the Triton kernel vs per-block opaque byte copies -- so +they are modelled as two explicit subclasses of :class:`TransferGroup` +and dispatched with ``isinstance``. The dispatch is the extension point +for future group kinds (e.g. sliding-window attention would add its own +subclass instead of overloading a boolean). + +Nothing here is consumed by the sglang connector; per the placement rule +("common only when shared") these types live under ``vllm/``. +""" + +from dataclasses import dataclass, field +from enum import Enum +from typing import List, Optional + +import torch + + +class KVLayout(Enum): + """The flash_attn paged-KV tensor layouts this connector understands. + + Detected from the tensor *shape* (never from version strings) in + ``vllm_common.attn_kv_views``. One layout per vLLM era: + + * vLLM <= 0.22.1 returned ``(2, num_blocks, block, H, D)``: + https://github.com/vllm-project/vllm/blob/v0.22.1/vllm/v1/attention/backends/flash_attn.py#L149 + * vLLM 0.23.0 - 0.25.x returned ``(num_blocks, 2, block, H, D)``: + https://github.com/vllm-project/vllm/blob/v0.23.0/vllm/v1/attention/backends/flash_attn.py#L149 + * vLLM >= 0.26.0 returns the packed 4-D ``(num_blocks, H, block, 2D)``: + https://github.com/vllm-project/vllm/blob/v0.26.0/vllm/v1/attention/backends/flash_attn.py#L141 + + The saved byte layout differs between the split-K/V and packed eras, so + KV cache is not portable across vLLM major upgrades; instance_id + isolation prevents such mixing in practice. + """ + + SPLIT_KV_5D_KV_FIRST = "split_kv_5d_kv_first" # (2, num_blocks, block, H, D) + SPLIT_KV_5D_N_FIRST = "split_kv_5d_n_first" # (num_blocks, 2, block, H, D) + PACKED_4D = "packed_4d" # (num_blocks, H, block, 2D) + + +@dataclass(frozen=True) +class TransferGroup: + """Base: fields every kv cache group carries regardless of kind.""" + + group_idx: int + # Location spec name registered with the manager, e.g. "tp0_g3". + spec_name: str + layer_names: List[str] + # The group's own block table granularity in tokens (spec.block_size). + block_size: int + # Bytes stored per manager block for this whole group (all its layers). + per_block_bytes: int + layer_num: int + + +@dataclass(frozen=True) +class AttentionTransferGroup(TransferGroup): + """Token-granular KV, moved through the strided gather/scatter kernel.""" + + # Which vLLM-era layout the pointers below were normalized from. + kv_layout: KVLayout + # int64 tensor of transfer-pointer bases on the compute device. For the + # packed layout one pointer per layer [L0, L1, ...]; for split K/V layouts + # two per layer [K0, V0, K1, V1, ...] -- each view's data_ptr() is its own + # base, so the kernel never adds a K->V offset. + kvcache_ptr_tensor_gpu: torch.Tensor + # Number of transfer pointers (staging buffer rows per block): + # layer_num for the packed layout, 2 * layer_num for split K/V. + num_kv_ptrs: int + # heads * content dim per pointer (content dim is 2*D packed, D split). + per_token_dim: int + # Tokens per kernel page (may differ from block_size on padded pages). + kernel_block_size: int + # Element stride between kernel pages of one pointer; 0 => the pages are + # contiguous and the kernel uses flat indexing. + block_stride: int + + +@dataclass(frozen=True) +class StateTransferGroup(TransferGroup): + """Per-block opaque state bytes, copied verbatim (mamba/linear/gdn).""" + + # Per layer, (num_blocks, page_size_bytes) uint8 views into the state + # storage (all of a layer's state tensors share one storage). + block_view_tensors: List[torch.Tensor] = field(default_factory=list) + # Bytes per block per state layer (spec.page_size_bytes). + page_size_bytes: int = 0 + + +@dataclass(frozen=True) +class TransferPlan: + """One group's slice of a save (or load): what to move, from/to where. + + Built by the worker from the manager's locations plus the request's + block tables; consumed by the transfer tasks. ``uris`` is positionally + aligned with the manager blocks and may contain ``None`` where a block + carries no data for this group's spec (hybrid sparse coverage) -- what a + hole means is the task's disposition logic (see data_transfer).""" + + group: TransferGroup + uris: List + # Attention groups: flat token slots per manager block (gather/scatter). + token_indices: Optional[List[List[int]]] = None + # State groups: source/target state block id per manager block. + block_ids: Optional[List[int]] = None + + +@dataclass(frozen=True) +class KVCacheInfo: + """Worker-side registered KV cache description (all groups).""" + + tp_rank: int + world_size: int + groups: List[TransferGroup] + device: torch.device + dtype: torch.dtype diff --git a/kv_cache_manager/py_connector/vllm/v1_connector.py b/kv_cache_manager/py_connector/vllm/v1_connector.py index 3a0de4d98..439f2204c 100644 --- a/kv_cache_manager/py_connector/vllm/v1_connector.py +++ b/kv_cache_manager/py_connector/vllm/v1_connector.py @@ -1,25 +1,32 @@ -import copy -import json -import math -import time -import typing -import inspect -import threading +"""KVCM vLLM connector (v1): a thin shell over the per-role implementations. + +vLLM instantiates one connector per role (a scheduler instance and one +worker instance per TP rank) from a single registered class, so this shell +is the plugin entry point: it performs the role-agnostic setup (config +parsing, manager registration) and delegates every hook to the role object +it owns -- ``connector_scheduler`` on scheduler-role instances, +``connector_worker`` on worker-role instances; the other slot stays None, +and every hook asserts the slot it needs (the mooncake pattern): -from dataclasses import dataclass, field -from typing import Any, Optional, List, Dict, Tuple +* ``ConnectorScheduler`` -- matching, saving orchestration, request + finishing (connector_scheduler.py); +* ``ConnectorWorker`` -- block translation and data-plane transfer + (connector_worker.py). -from concurrent.futures import ThreadPoolExecutor -from kv_cache_manager.client.pybind import kvcm_py_client +Shared vocabulary (GroupMeta / spec naming / KV layout normalization / +hybrid gate) lives in vllm_common.py. +""" + +import typing + +from typing import Optional -import torch -import typing_extensions from vllm.config import VllmConfig -from vllm.distributed import get_tensor_model_parallel_rank from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole, + SupportsHMA, ) try: @@ -31,859 +38,228 @@ from vllm.utils import get_kv_cache_torch_dtype, get_ip from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.kv_cache_interface import MambaSpec from vllm.v1.outputs import KVConnectorOutput from kv_cache_manager.py_connector.common.manager_client import KvCacheManagerClient -from kv_cache_manager.py_connector.common.tp_coordinator import CoordinateMsgSerializer, TpCoordinatorServer, \ - TpCoordinatorClient, SendBlockStartEvent, CoordinateMessage, SaveContext +from kv_cache_manager.py_connector.common.tp_coordinator import TpCoordinatorClient from kv_cache_manager.py_connector.common.logger import logger, configure_log_level -from kv_cache_manager.py_connector.common._version_info import FULL_VERSION, GIT_COMMIT, BUILD_TIME -from kv_cache_manager.py_connector.common.types import KVCacheInfo -from kv_cache_manager.py_connector.kernel.gather_scatter_helper import CopyBufferAllocator -from kv_cache_manager.py_connector.vllm.metadata import SaveRequest, LoadRequest, FinishRequest, ReqStateToWorker, \ - TairKvCacheConnectorMetadata -from kv_cache_manager.py_connector.vllm.config import TairKvCacheConnectorExtraConfig -from kv_cache_manager.py_connector.vllm.location_query_manager import LocationQueryManager -from kv_cache_manager.py_connector.vllm.data_transfer import MultiResult, DataTransferManager, _get_device_module +try: + # Stamped into the wheel at build time; absent in a source checkout. + from kv_cache_manager.py_connector.common._version_info import ( + FULL_VERSION, GIT_COMMIT, BUILD_TIME) +except ImportError: + FULL_VERSION, GIT_COMMIT, BUILD_TIME = "dev", "source", "source" -if typing_extensions.TYPE_CHECKING: +from kv_cache_manager.py_connector.vllm.config import TairKvCacheConnectorExtraConfig +from kv_cache_manager.py_connector.vllm.connector_scheduler import ConnectorScheduler +from kv_cache_manager.py_connector.vllm.connector_worker import ConnectorWorker +from kv_cache_manager.py_connector.vllm.metadata import TairKvCacheConnectorMetadata +from kv_cache_manager.py_connector.vllm.vllm_common import ( + GroupMeta, StateGroupMeta, attn_kv_views, build_spec_groups, + ensure_hybrid_supported, parse_groups, spec_name) + +if typing.TYPE_CHECKING: from vllm.forward_context import ForwardContext from vllm.attention import AttentionMetadata from vllm.v1.request import Request from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.kv_cache_interface import KVCacheConfig - -@dataclass -class ReqState: - """请求状态类，跟踪单个请求的状态信息""" - - # TODO: split this class to ReqStateInScheduler and ReqStateInWorker - req_id: str - token_ids: list[int] - local_block_ids: list[int] - has_saved_block_num: int - local_matched_token_num: int - remote_matched_token_num: int - - # vllm_request only avail in scheduler - vllm_request: Optional["Request"] - - # scheduled_saving_count, sent_saving_count, need_report_after_saving_finished: - # not sync between scheduler and worker and have different meaning - # only available in scheduler and tp0 worker - scheduled_saving_count: int = 0 - sent_saving_count: int = 0 - need_report_after_saving_finished: bool = False - - @staticmethod - def create_from_delta(req_state_delta: 'ReqStateToWorker'): - """从ReqStateToWorker创建ReqState实例""" - return ReqState( - req_id=req_state_delta.req_id, - token_ids=req_state_delta.new_tokens_ids, - local_block_ids=req_state_delta.new_local_block_ids, - has_saved_block_num=req_state_delta.has_saved_block_num, - local_matched_token_num=0, - remote_matched_token_num=0, - vllm_request=None - ) - - def update_from_delta(self, req_state_delta: 'ReqStateToWorker'): - """使用ReqStateToWorker更新当前状态""" - self.token_ids.extend(req_state_delta.new_tokens_ids) - - if req_state_delta.resumed_from_preemption: - self.local_block_ids = req_state_delta.new_local_block_ids - else: - self.local_block_ids.extend(req_state_delta.new_local_block_ids) - - -@dataclass -class TransferTaskArgs: - blocks_idx: List[List[int]] = field(default_factory=list) - remote_uris: List[str] = field(default_factory=list) +# Compatibility re-exports: tests and the e2e harness import these from +# v1_connector (their original home before the role split). +__all__ = [ + "TairKvCacheConnector", "attn_kv_views", "ensure_hybrid_supported", + "GroupMeta", "spec_name", "build_spec_groups", "parse_groups", +] -class TairKvCacheConnector(KVConnectorBase_V1): - def _tp_rank_to_spec_name(self, tp_rank: int) -> str: - """Convert TP rank to location spec name.""" - return f"tp{tp_rank}" +class TairKvCacheConnector(KVConnectorBase_V1, SupportsHMA): - def __init__(self, - vllm_config: "VllmConfig", - role: KVConnectorRole, - kv_cache_config: Optional["KVCacheConfig"] = None, - ): + # ------------------------------------------------------------------ # + # Init / registration (role-agnostic) + # ------------------------------------------------------------------ # + def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole, + kv_cache_config: Optional["KVCacheConfig"] = None): + super().__init__(vllm_config, role, kv_cache_config) + assert kv_cache_config is not None, \ + "TairKvCacheConnector requires vLLM to pass kv_cache_config (vllm >= 0.11.1)" - init_params = inspect.signature(KVConnectorBase_V1.__init__).parameters - if len(init_params) == 3: - # vllm <= 0.11.0 - super().__init__(vllm_config, role) - else: - # vllm >= 0.11.1 - super().__init__(vllm_config, role, kv_cache_config) - - logger.warning("KVCM vllm connector version: %s (commit: %s, build: %s)", FULL_VERSION, GIT_COMMIT, BUILD_TIME) - - connector_extra_config = vllm_config.kv_transfer_config.kv_connector_extra_config - self._extra_config = TairKvCacheConnectorExtraConfig(connector_extra_config) - - # Apply log level with priority: env var > startup param > default - configure_log_level(self._extra_config.log_level) + logger.warning("KVCM vllm connector version: %s (commit: %s, build: %s)", + FULL_VERSION, GIT_COMMIT, BUILD_TIME) - self._kv_caches: Optional[dict[str, torch.Tensor]] = None - self._local_block_size = vllm_config.cache_config.block_size + extra_config = TairKvCacheConnectorExtraConfig( + **vllm_config.kv_transfer_config.kv_connector_extra_config) + configure_log_level(extra_config.log_level) model_config = vllm_config.model_config + assert vllm_config.parallel_config.pipeline_parallel_size == 1 + if getattr(model_config, "use_mla", False): + raise NotImplementedError("MLA models are not supported by TairKvCacheConnector") - self._use_mla = (hasattr(model_config, "use_mla") and - isinstance(model_config.use_mla, bool) and - model_config.use_mla) - - manager_block_size = self._local_block_size - if self._extra_config.preferred_block_size != 0: - manager_block_size = self._extra_config.preferred_block_size - + self._vllm_block_size = vllm_config.cache_config.block_size self._tp_size = vllm_config.parallel_config.tensor_parallel_size - kv_dtype = get_kv_cache_torch_dtype(vllm_config.cache_config.cache_dtype, model_config.dtype) - num_layer = model_config.get_num_layers(vllm_config.parallel_config) - per_tp_rank_kv_head_num = model_config.get_num_kv_heads(vllm_config.parallel_config) - head_size = model_config.get_head_size() - per_manager_location_spec_shape = [num_layer, 1 if self._use_mla else 2, manager_block_size, - per_tp_rank_kv_head_num, - head_size] + self._kv_dtype = get_kv_cache_torch_dtype( + vllm_config.cache_config.cache_dtype, model_config.dtype) + + # Manager block size: attention KV is token-granular and can be re-blocked, + # but mamba state exists once per scheduler block, so hybrid models must + # keep manager block == scheduler block. + manager_block_size = self._vllm_block_size + self._has_state_groups = any( + isinstance(g.kv_cache_spec, MambaSpec) for g in kv_cache_config.kv_cache_groups) + if self._has_state_groups: + ensure_hybrid_supported(force=extra_config.force_hybrid_support) + if extra_config.preferred_block_size != 0: + if self._has_state_groups: + if extra_config.preferred_block_size != self._vllm_block_size: + logger.warning( + "preferred_block_size=%d ignored for hybrid model: mamba state is " + "per scheduler block (%d)", extra_config.preferred_block_size, + self._vllm_block_size) + else: + manager_block_size = extra_config.preferred_block_size + self._manager_block_size = manager_block_size + + self._group_metas = parse_groups(kv_cache_config, manager_block_size) - assert vllm_config.parallel_config.pipeline_parallel_size == 1 deployment = { "model_name": model_config.served_model_name, - "dtype": str(kv_dtype)[6:], # remove "torch." - "use_mla": self._use_mla, - "tp_size": vllm_config.parallel_config.tensor_parallel_size, + "dtype": str(self._kv_dtype)[6:], # strip "torch." + "use_mla": False, + "tp_size": self._tp_size, "dp_size": vllm_config.parallel_config.data_parallel_size, "pp_size": vllm_config.parallel_config.pipeline_parallel_size, } - logger.info(deployment) + logger.info("deployment: %s, groups: %s", deployment, self._group_metas) self._manager_client = KvCacheManagerClient.from_connector_config( - vars(self._extra_config) - ) - self._manager_block_size = manager_block_size - - self._alive_requests: dict[str, ReqState] = {} - self._waiting_to_load_requests: List[LoadRequest] = [] - self._waiting_to_save_requests_lock = threading.Lock() - self._waiting_to_save_requests: List[SaveRequest] = [] - self._waiting_to_finish_requests: List[FinishRequest] = [] + extra_config.model_dump()) + host_ip = get_ip() - self._canceled_save_request_ids_lock = threading.Lock() - self._canceled_save_request_ids: List[str] = [] - - # TODO: add coordinator host auto detection, maybe use data parallel host - # TODO: add DP support - self._host_ip = get_ip() - port = self._extra_config.coordinator_base_port - - register_response = self._manager_client.register_instance({ - "trace_id": "trace_trace", - "instance_group": self._extra_config.instance_group, - "instance_id": self._extra_config.instance_id, + register_request = { + "trace_id": "register_%s" % extra_config.instance_id, + "instance_group": extra_config.instance_group, + "instance_id": extra_config.instance_id, "model_deployment": deployment, "block_size": manager_block_size, - "location_spec_infos": [{ - "name": self._tp_rank_to_spec_name(rank), - "size": math.prod(per_manager_location_spec_shape) * kv_dtype.itemsize - } for rank in range(self._tp_size)], - }) - # TODO: check conflict and update - self._iov_size = math.prod( - per_manager_location_spec_shape) * kv_dtype.itemsize * self._extra_config.hf3fs_concurrent_io_block_count - + "location_spec_infos": [ + {"name": spec_name(rank, meta.group_idx), "size": meta.per_block_bytes} + for rank in range(self._tp_size) for meta in self._group_metas + ], + } + spec_groups = build_spec_groups(self._group_metas, self._tp_size) + if spec_groups: + # Hybrid models publish per-block spec coverage + # (see vllm_common.build_spec_groups). + register_request["location_spec_groups"] = spec_groups + register_response = self._manager_client.register_instance(register_request) + + # One role object per instance; the other slot stays None and every + # hook asserts the slot it needs. + self.connector_scheduler: Optional[ConnectorScheduler] = None + self.connector_worker: Optional[ConnectorWorker] = None if role == KVConnectorRole.SCHEDULER: - self._epoch = 0 - self._coordinator_client = TpCoordinatorClient(self._host_ip, port) - self._http_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="kvcm_http_") - self._location_query_manager = LocationQueryManager(self._manager_client, self._http_executor, - self._extra_config.instance_id, - self._extra_config.async_get_cache_location) - - logger.warning( - "TairKvCacheConnector in scheduler inited, kv_connector_extra_config: %r," - " server block size: %d, vllm block size: %d,", - self._extra_config.__dict__, - self._manager_block_size, - self._local_block_size, - ) - - elif role == KVConnectorRole.WORKER: - self._tp_rank = get_tensor_model_parallel_rank() - self._device_mod = None - self._save_stream = None - self._load_stream = None - + self.connector_scheduler = ConnectorScheduler( + extra_config, self._group_metas, manager_block_size, + self._vllm_block_size, self._tp_size, self._manager_client, + TpCoordinatorClient(host_ip, extra_config.coordinator_base_port)) logger.warning( - "TairKvCacheConnector in worker inited, tp rank: %d, tp size: %d, host_ip: %s, port: %d" % ( - self._tp_rank, self._tp_size, self._host_ip, port) - ) - - if self._tp_rank == 0: - # start coordinator - self._coordinator_server = TpCoordinatorServer(self._host_ip, port, self._tp_size, - self.on_save_finished) - - self._coordinator_client = TpCoordinatorClient(self._host_ip, port) - - self._storage_configs = register_response["storage_configs"] - # data transfer setup - self._location_spec_name = self._tp_rank_to_spec_name(self._tp_rank) - self._write_timeout_seconds = self._extra_config.write_timeout_seconds - - sdk_backend_configs = [] - - hf3fs_configs = self.parse_hf3fs_configs(self._storage_configs) - sdk_backend_configs.extend(hf3fs_configs) - logger.debug(sdk_backend_configs) - transfer_client_json = { - "instance_group": self._extra_config.instance_group, - "instance_id": self._extra_config.instance_id, - "block_size": self._manager_block_size, - "sdk_config": { - "thread_num": self._extra_config.sdk_thread_num, - "queue_size": self._extra_config.sdk_queue_size, - "sdk_backend_configs": sdk_backend_configs, - "timeout_config": { - "get_timeout_ms": self._extra_config.sdk_get_timeout_ms, - "put_timeout_ms": self._extra_config.sdk_put_timeout_ms, - }, - }, - "location_spec_infos": { - self._location_spec_name: math.prod(per_manager_location_spec_shape) * kv_dtype.itemsize, - }, - } - self._transfer_client_config = json.dumps(transfer_client_json) - - self._init_params = kvcm_py_client.InitParams() - self._init_params.role_type = kvcm_py_client.RoleType.WORKER - self._init_params.self_location_spec_name = self._location_spec_name - self._init_params.storage_configs = f"{self._storage_configs}" - - logger.info("_transfer_client_config:%s, _init_params:%s", self._transfer_client_config, self._init_params) - - self._transfer_client = kvcm_py_client.TransferClient.Create( - self._transfer_client_config, self._init_params - ) - assert self._transfer_client is not None, "kvcm_py_client.TransferClient.Create failed" + "TairKvCacheConnector scheduler inited, extra_config: %r, " + "manager block size: %d, vllm block size: %d, groups: %d", + extra_config.model_dump(), manager_block_size, + self._vllm_block_size, len(self._group_metas)) + else: + self.connector_worker = ConnectorWorker( + extra_config, self._group_metas, manager_block_size, + self._tp_size, host_ip, self._manager_client, + TpCoordinatorClient(host_ip, extra_config.coordinator_base_port), + register_response) def shutdown(self): - # TODO: stop background threads and cleanup transfer client + if self.connector_scheduler is not None: + self.connector_scheduler.shutdown() self._manager_client.close() return None - def parse_hf3fs_configs(self, storage_configs): - hf3fs_configs = [] - storage_configs_json = json.loads(storage_configs) - for storage_config in storage_configs_json: - if storage_config["type"] == "vcns_hf3fs": - storage_config["type"] = "hf3fs" - if storage_config["type"] == "hf3fs" and storage_config["is_available"]: - hf3fs_config = { - "type": storage_config["type"], - "mountpoint": storage_config["storage_spec"]["mountpoint"], - "root_dir": storage_config["storage_spec"]["root_dir"], - "read_iov_block_size": self._extra_config.read_iov_block_size, - "read_iov_size": self._iov_size, - "write_iov_block_size": self._extra_config.write_iov_block_size, - "write_iov_size": self._iov_size, - } - hf3fs_configs.append(hf3fs_config) - self._storage_configs = json.dumps(storage_configs_json) - return hf3fs_configs - - def generate_blocks(self, token_ids, block_size, max_token_length) -> list[dict[str, Any]]: - results = [] - token_length = min(len(token_ids), max_token_length) - # token_length = len(token_ids) - for i in range(0, token_length, block_size): - if i + block_size > token_length: - break - results.append({ - "token_ids": token_ids[i:i + block_size], - "unique_id": None, - "location": None - }) - return results - - # ============================== - # Worker-side methods - # ============================== - - def generate_blocks_idx(self, manager_block_idxes, local_block_ids): - blocks_idx = [] - for manager_block_idx in manager_block_idxes: - # get kvcache index list - block_idx = [] - for i in range(self._manager_block_size): - now_token_idx = manager_block_idx * self._manager_block_size + i - assert now_token_idx // self._local_block_size < len(local_block_ids) - local_block_id = local_block_ids[now_token_idx // self._local_block_size] - token_offset = now_token_idx % self._local_block_size - block_idx.append(local_block_id * self._local_block_size + token_offset) - blocks_idx.append(block_idx) - return blocks_idx - - def on_save_finished(self, write_session_id: str, save_context: SaveContext): - logger.debug(save_context.result_per_rank) - for block_idx in range(len(save_context.locations)): - # TODO: report uri when enable local alloc - # location_specs = [] - is_fully_saved = True - for rank in range(self._tp_size): - is_success = save_context.result_per_rank[rank][block_idx] - if not is_success: - # this spec is not fully saved, report failed - is_fully_saved = False - # else: - # # Convert the spec to include name field instead of tp_rank - # location_specs.append({ - # "name": self._tp_rank_to_spec_name(rank), - # "uri": spec - # }) - if is_fully_saved: - # save_context.locations[block_idx]["location_specs"] = location_specs - save_context.success_mask.append(True) - else: - save_context.success_mask.append(False) - logger.debug("finish_write_cache blocks:%s mask:%s write_session_id:%s", save_context.locations, - save_context.success_mask, write_session_id) - try: - self._manager_client.finish_write_cache({ - "trace_id": "test_test", - "instance_id": self._extra_config.instance_id, - "write_session_id": write_session_id, - "success_blocks": { - "bool_masks": { - "values": save_context.success_mask - } - } - }) - except Exception as e: - logger.warning("finish_write_cache failed, write_session_id: %s, error: %s", write_session_id, e) - - def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): - _, first_layer_kvcache = next(iter(kv_caches.items())) - self._kv_caches = kv_caches - # TODO: support MLA - - assert self._local_block_size == first_layer_kvcache.shape[2], "kv cache shape error" - for layer_name, kvcache in kv_caches.items(): - assert kvcache.is_contiguous(), "kv cache must be contiguous" - - # torch.Size([2, block_num, block_size, kv_head_num, kv_dim]) - # 2 -> key, value - self._local_block_num = first_layer_kvcache.shape[1] - self._local_token_num = self._local_block_num * self._local_block_size - - self._dtype = first_layer_kvcache.dtype - self._device = first_layer_kvcache.device - self._device_mod = _get_device_module(self._device) - self._save_stream = self._device_mod.Stream() - self._load_stream = self._device_mod.Stream() - self._per_manager_location_spec_layer_shape = [first_layer_kvcache.shape[0], - self._manager_block_size, - first_layer_kvcache.shape[3] * first_layer_kvcache.shape[4]] - self._per_manager_location_spec_layer_byte_size = math.prod( - self._per_manager_location_spec_layer_shape) * self._dtype.itemsize - self._per_layer_token_key_dim_size = first_layer_kvcache.shape[3] * first_layer_kvcache.shape[4] - self._per_layer_token_key_byte_size = (first_layer_kvcache.shape[3] * - first_layer_kvcache.shape[4] * self._dtype.itemsize) - assert self._per_layer_token_key_byte_size == first_layer_kvcache[0][0][1].data_ptr() - \ - first_layer_kvcache[0][0][0].data_ptr(), "kv cache shape error" - assert self._per_manager_location_spec_layer_byte_size == 2 * self._manager_block_size * self._per_layer_token_key_byte_size - - self._per_manager_location_spec_shape = [len(self._kv_caches)] + self._per_manager_location_spec_layer_shape - self._per_manager_location_spec_byte_size = math.prod( - self._per_manager_location_spec_shape) * self._dtype.itemsize - - self._kvcache_ptr_tensor_cpu = torch.tensor( - [self._kv_caches[name].data_ptr() for name in self._kv_caches], - dtype=torch.int64, - device="cpu" - ) - self._kvcache_ptr_tensor_gpu = self._kvcache_ptr_tensor_cpu.to(self._device) - if self._use_mla: - self._all_kvcache_ptr_tensor_cpu = torch.tensor( - [self._kv_caches[name].data_ptr() for name in self._kv_caches], - dtype=torch.int64, - device="cpu" - ) - else: - kvcache_ptrs = [] - for name in self._kv_caches: - kvcache_ptrs.append(self._kv_caches[name][0].data_ptr()) - kvcache_ptrs.append(self._kv_caches[name][1].data_ptr()) - self._all_kvcache_ptr_tensor_cpu = torch.tensor( - kvcache_ptrs, - dtype=torch.int64, - device="cpu" - ) - self._all_kvcache_ptr_tensor_gpu = self._all_kvcache_ptr_tensor_cpu.to(self._device) - - self._kvcache_info = KVCacheInfo( - self._tp_rank, - self._tp_size, - self._kv_caches, - self._kvcache_ptr_tensor_cpu, - self._kvcache_ptr_tensor_gpu, - self._all_kvcache_ptr_tensor_gpu, - len(self._kv_caches), - self._local_token_num, - tuple(self._per_manager_location_spec_shape), - self._per_manager_location_spec_byte_size, - self._per_layer_token_key_dim_size, - self._device, - self._dtype - ) - self._copy_buffer_allocator = CopyBufferAllocator(torch.device("cpu"), self._dtype, - self._per_manager_location_spec_shape, 1024) - - # 初始化DataTransferManager实例 - self._data_transfer = DataTransferManager( - self._kvcache_info, - self._manager_block_size, - self._copy_buffer_allocator, - self._transfer_client, - self._coordinator_client, - self._extra_config, - ) - - logger.warning("register_kv_caches, _per_manager_location_spec_layer_shape: %s", - self._per_manager_location_spec_layer_shape) + # ------------------------------------------------------------------ # + # Scheduler hooks (called on the scheduler-role instance) + # ------------------------------------------------------------------ # + def get_num_new_matched_tokens(self, request: "Request", + num_computed_tokens: int): + assert self.connector_scheduler is not None + return self.connector_scheduler.get_num_new_matched_tokens( + request, num_computed_tokens) + + def update_state_after_alloc(self, request: "Request", blocks: "KVCacheBlocks", + num_external_tokens: int): + assert self.connector_scheduler is not None + self.connector_scheduler.update_state_after_alloc( + request, blocks, num_external_tokens) + + def update_connector_output(self, connector_output: KVConnectorOutput): + assert self.connector_scheduler is not None + self.connector_scheduler.update_connector_output(connector_output) + + def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata: + assert self.connector_scheduler is not None + return self.connector_scheduler.build_connector_meta(scheduler_output) + + def request_finished_all_groups(self, request: "Request", + block_ids) -> tuple: + assert self.connector_scheduler is not None + return self.connector_scheduler.request_finished_all_groups(request, block_ids) + + def request_finished(self, request: "Request", block_ids) -> tuple: + # Only reached on vLLM versions without SupportsHMA dispatch + # (https://github.com/vllm-project/vllm/blob/v0.26.0/vllm/v1/core/sched/scheduler.py#L2513 + # -- upstream plans to deprecate this path). Kept as a shim for the + # three vLLM eras this connector supports; remove once the minimum + # supported version dispatches via SupportsHMA only. + assert self.connector_scheduler is not None + return self.connector_scheduler.request_finished(request, block_ids) + def get_finished_count(self): + assert self.connector_scheduler is not None + return self.connector_scheduler.get_finished_count() + + # ------------------------------------------------------------------ # + # Worker hooks (called on each worker-role instance) + # ------------------------------------------------------------------ # + def register_kv_caches(self, kv_caches: dict): + assert self.connector_worker is not None + self.connector_worker.register_kv_caches(kv_caches) + + def bind_connector_metadata(self, connector_metadata: KVConnectorMetadata) -> None: + # The worker consumes each instruction directly from the metadata + # (start_load_kv / wait_for_save receive it explicitly); no mirror + # to replay here. + super().bind_connector_metadata(connector_metadata) def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: - meta = typing.cast(TairKvCacheConnectorMetadata, self._get_connector_metadata()) - - for load_req in meta.to_load_requests: - if len(load_req.need_load_locations) == 0: - continue - - block_token_indices = self.generate_blocks_idx(load_req.manager_block_idxes, load_req.local_block_ids) - all_remote_uris = self.get_self_uris(load_req.need_load_locations) - - per_task_size = self._extra_config.block_per_load_task - task_num = math.ceil(len(block_token_indices) / per_task_size) - done_callback = self._data_transfer.create_load_done_callback( - load_req.req_id, - self._kvcache_info.tp_rank, - meta.epoch, - copy.copy(load_req.local_block_ids) - ) - multi_result = MultiResult(task_num, done_callback) - - task_idx = 0 - for i in range(0, len(block_token_indices), per_task_size): - end_idx = min(len(block_token_indices), i + per_task_size) - task_remote_uris = all_remote_uris[i:end_idx] - task_block_token_indices = block_token_indices[i:end_idx] - self._data_transfer.submit_task(self._data_transfer.load_task, multi_result, task_idx, task_remote_uris, - task_block_token_indices) - task_idx += 1 + assert self.connector_worker is not None + self.connector_worker.start_load_kv( + forward_context, self._get_connector_metadata(), **kwargs) def wait_for_layer_load(self, layer_name: str) -> None: - # logger.warning("wait_for_layer_load, layer_name: %s", layer_name) - pass + assert self.connector_worker is not None + self.connector_worker.wait_for_layer_load(layer_name) - def save_kv_layer(self, layer_name: str, kv_layer: torch.Tensor, attn_metadata: "AttentionMetadata", - **kwargs) -> None: - # logger.warning("save_kv_layer, layer_name: %s", layer_name) - pass + def save_kv_layer(self, layer_name: str, kv_layer, + attn_metadata: "AttentionMetadata", **kwargs) -> None: + assert self.connector_worker is not None + self.connector_worker.save_kv_layer(layer_name, kv_layer, attn_metadata, **kwargs) def wait_for_save(self): - meta = typing.cast(TairKvCacheConnectorMetadata, self._get_connector_metadata()) - # logger.warning("wait_for_save, meta: %r", meta) - - kvcache_ready_event = None - if len(meta.to_save_requests) > 0: - kvcache_ready_event = self._device_mod.Event() - kvcache_ready_event.record(self._device_mod.current_stream()) - - for req_save in meta.to_save_requests: - req = self._alive_requests[req_save.req_id] - - # get idx - blocks_idx = self.generate_blocks_idx(req_save.manager_block_idxes, req.local_block_ids) - all_remote_uris = self.get_self_uris(req_save.target_locations) - - per_task_size = self._extra_config.block_per_save_task - task_num = math.ceil(len(blocks_idx) / per_task_size) - done_callback = self._data_transfer.create_save_done_callback( - req.req_id, - self._kvcache_info.tp_rank, - req_save.write_session_id - ) - multi_result = MultiResult(task_num, done_callback) - - task_idx = 0 - for i in range(0, len(blocks_idx), per_task_size): - end_idx = min(len(blocks_idx), i + per_task_size) - task_remote_uris = all_remote_uris[i:end_idx] - task_block_token_indices = blocks_idx[i:end_idx] - self._data_transfer.submit_task(self._data_transfer.save_task, multi_result, task_idx, task_remote_uris, - task_block_token_indices, - kvcache_ready_event) - task_idx += 1 - if self._tp_rank == 0: - req.scheduled_saving_count += 1 - - def get_self_uris(self, locations): - all_remote_uris = [] - for idx, location in enumerate(locations): - for location_spec in location["location_specs"]: - # Match by location spec name instead of tp_rank - if self._tp_rank_to_spec_name(self._kvcache_info.tp_rank) == location_spec["name"]: - all_remote_uris.append(location_spec["uri"]) - return all_remote_uris - - def get_finished( - self, finished_req_ids: set[str] - ) -> Tuple[Optional[set[str]], Optional[set[str]]]: - meta = typing.cast(TairKvCacheConnectorMetadata, self._get_connector_metadata()) - - if self._tp_rank != 0: - for finish_req in meta.to_finish_requests: - req_id = finish_req.req_id - if req_id in self._alive_requests: - self._alive_requests.pop(req_id) - return None, None - - # self._tp_rank == 0 - finished_saving_reqs = [] - # check if any request is saving kvcache - (finished_saving_tasks, finished_loading_tasks) = self._coordinator_server.get_finished_tasks() - for req_id in finished_saving_tasks: - req = self._alive_requests[req_id] - req.sent_saving_count += 1 - - assert req.sent_saving_count <= req.scheduled_saving_count - if (req.need_report_after_saving_finished and - req.sent_saving_count == req.scheduled_saving_count): - finished_saving_reqs.append(req_id) - self._alive_requests.pop(req_id) - - for finish_req in meta.to_finish_requests: - req_id = finish_req.req_id - if req_id not in self._alive_requests: - # called get_num_new_matched_tokens but never scheduled - continue - req = self._alive_requests[req_id] - if req.sent_saving_count == req.scheduled_saving_count: - finished_saving_reqs.append(req_id) - self._alive_requests.pop(req_id) - else: - self._alive_requests[req_id].need_report_after_saving_finished = True - return set(finished_saving_reqs), set(finished_loading_tasks) - - def get_block_ids_with_load_errors(self) -> set[int]: - if self._tp_rank != 0: - return set() - failed_set = self._coordinator_server.get_failed_loading_block_idxs() - if len(failed_set) > 0: - logger.warning("block_ids_with_load_errors: %s", failed_set) - return failed_set - - def bind_connector_metadata( - self, connector_metadata: KVConnectorMetadata) -> None: - self._connector_metadata = connector_metadata - meta = typing.cast(TairKvCacheConnectorMetadata, self._get_connector_metadata()) - - for req_state_delta in meta.requests: - if req_state_delta.req_id not in self._alive_requests: - assert not req_state_delta.is_delta - if not req_state_delta.is_delta: - self._alive_requests[req_state_delta.req_id] = ReqState.create_from_delta(req_state_delta) - else: - self._alive_requests[req_state_delta.req_id].update_from_delta(req_state_delta) - - # ============================== - # Scheduler-side methods - # ============================== - def get_num_new_matched_tokens(self, request: "Request", num_computed_tokens: int) -> Tuple[int, bool]: - # logger.warning("get matched token ids: %s, id: %s", request.prompt_token_ids, request.request_id) - - bypass_match = False - # TODO: add arrival_time to req_id in order to handle same request id - if request.request_id in self._alive_requests: - # bypass remote match for alive requests - # possible cases: - # 1. reschedule when all kvcache loading failed - # 2. TODO: no enough hbm to schedule the request - # bypass_match = True - # logger.warning("bypass match for alive request, req_id: %s", request.request_id) - pass - - computed_manager_block_size = num_computed_tokens // self._manager_block_size - all_calced_remote_block_num = computed_manager_block_size - new_matched_count = 0 - - if not bypass_match: - is_query_done, need_load_locations = ( - self._location_query_manager.get_locations_for_query(request, computed_manager_block_size)) - if not is_query_done: - # async get_cache_location - return None, False - new_matched_count = len(need_load_locations) * self._manager_block_size - logger.info("req:%s, new_matched_count:%d", request.request_id, new_matched_count) - - all_calced_remote_block_num = computed_manager_block_size + len(need_load_locations) - - if new_matched_count != 0: - self._waiting_to_load_requests.append(LoadRequest( - req_id=request.request_id, - manager_block_idxes=[i for i in range(computed_manager_block_size, all_calced_remote_block_num)], - need_load_locations=need_load_locations, - )) - - new_req_meta = ReqState(request.request_id, copy.copy(request.prompt_token_ids), [], - all_calced_remote_block_num, - num_computed_tokens, - new_matched_count, - request) - - self._alive_requests[request.request_id] = new_req_meta - return new_matched_count, new_matched_count > 0 - - def update_state_after_alloc(self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int): - if request.request_id not in self._alive_requests: - return - req_state = self._alive_requests[request.request_id] - # blocks_ids[0]: only one KV cache groups for now - # refer to vllm/v1/core/kv_cache_manager.py:35 - req_state.local_block_ids = copy.copy(blocks.get_block_ids()[0]) + assert self.connector_worker is not None + self.connector_worker.wait_for_save(self._get_connector_metadata()) - def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata: - meta = TairKvCacheConnectorMetadata(self._epoch) - self._epoch += 1 - - for load_req in self._waiting_to_load_requests: - request = self._alive_requests[load_req.req_id] - if len(request.local_block_ids) == 0: - # ignore load_req if vllm has not called update_state_after_alloc, - # vllm will call get_num_new_matched_tokens again - continue - load_req.local_block_ids = request.local_block_ids - meta.add_load_request(load_req) - self._waiting_to_load_requests = [] - - for vllm_req in scheduler_output.scheduled_new_reqs: - request = self._alive_requests[vllm_req.req_id] - request.local_block_ids = copy.copy(vllm_req.block_ids[0]) - - state_to_worker = ReqStateToWorker(req_id=request.req_id, - has_saved_block_num=request.has_saved_block_num, - new_tokens_ids=request.token_ids, - new_local_block_ids=request.local_block_ids, - is_delta=False - ) - meta.add_req_state_to_worker(state_to_worker) - logger.info("new request: %s, block_ids_len: %d", vllm_req.req_id, len(vllm_req.block_ids[0])) - - cached_reqs = scheduler_output.scheduled_cached_reqs - for idx, req_id in enumerate(cached_reqs.req_ids): - request = self._alive_requests[req_id] - vllm_req = request.vllm_request - num_new_tokens = scheduler_output.num_scheduled_tokens[req_id] - num_current_tokens = len(request.token_ids) - - new_token_ids = vllm_req.all_token_ids[ - num_current_tokens: num_current_tokens + num_new_tokens - ] - state_to_worker = ReqStateToWorker(req_id=request.req_id, - has_saved_block_num=request.has_saved_block_num) - - request.token_ids.extend(new_token_ids) - state_to_worker.new_tokens_ids = new_token_ids - - resumed_from_preemption = False - if hasattr(cached_reqs, "resumed_req_ids"): - # vllm >= 0.11.1 - resumed_from_preemption = req_id in cached_reqs.resumed_req_ids - else: - # vllm <= 0.11.0 - resumed_from_preemption = cached_reqs.resumed_from_preemption[idx] + def get_finished(self, finished_req_ids: set): + assert self.connector_worker is not None + return self.connector_worker.get_finished( + finished_req_ids, self._get_connector_metadata()) - if resumed_from_preemption: - request.local_block_ids = copy.copy(cached_reqs.new_block_ids[idx][0]) - state_to_worker.resumed_from_preemption = True - state_to_worker.new_local_block_ids = request.local_block_ids - else: - if cached_reqs.new_block_ids[idx] is None: - # https://github.com/vllm-project/vllm/pull/23262 - continue - new_block_ids = cached_reqs.new_block_ids[idx][0] - request.local_block_ids.extend(new_block_ids) - state_to_worker.new_local_block_ids = new_block_ids - meta.add_req_state_to_worker(state_to_worker) - - for req in self._alive_requests.values(): - target_save_num = min(len(req.token_ids), - len(req.local_block_ids) * self._local_block_size) // self._manager_block_size - if target_save_num > req.has_saved_block_num: - req.scheduled_saving_count += 1 - self._http_executor.submit( - self.start_save_kvcache_async, - req.req_id, - req.token_ids[:target_save_num * self._manager_block_size], - target_save_num - ) - req.has_saved_block_num = target_save_num - - new_save_reqs: List[SaveRequest] = [] - with self._waiting_to_save_requests_lock: - new_save_reqs = self._waiting_to_save_requests - self._waiting_to_save_requests = [] - for save_req in new_save_reqs: - if save_req.req_id not in self._alive_requests: - # TODO: should not happen anymore - logger.warning("request %s is not alive, skip saving", save_req.req_id) - continue - req = self._alive_requests[save_req.req_id] - meta.add_save_request(save_req) - - req.sent_saving_count += 1 - if (req.need_report_after_saving_finished and - req.scheduled_saving_count == req.sent_saving_count): - self._waiting_to_finish_requests.append(FinishRequest(req.req_id)) - self._alive_requests.pop(req.req_id) - - self.handle_canceled_save_req() - - for finish_req in self._waiting_to_finish_requests: - meta.add_finish_request(finish_req) - self._waiting_to_finish_requests = [] - - # logger.warning("build_connector_meta: %r", meta) - return meta - - def start_save_kvcache_async(self, req_id, token_ids, target_save_num): - request = { - "trace_id": "%s_%d" % (req_id, self._epoch), - "instance_id": self._extra_config.instance_id, - "block_keys": [], - "token_ids": token_ids, - "write_timeout_seconds": 30 - } - logger.debug("start_write_cache req: %s", request) - try: - response = self._manager_client.start_write_cache(request) - except Exception as e: - logger.warning("start_write_cache error, skip this saving, exception: %s", e) - with self._canceled_save_request_ids_lock: - self._canceled_save_request_ids.append(req_id) - return - # call manager start write - logger.debug("start_write_cache resp: %s", response) - locations = response["locations"] - write_session_id = response["write_session_id"] - # check if success - - if len(locations) == 0: - try: - self._manager_client.finish_write_cache({ - "trace_id": "test_test", - "instance_id": self._extra_config.instance_id, - "write_session_id": write_session_id, - "success_blocks": { - "bool_masks": { - "offset": 0 - } - } - }) - except Exception as e: - logger.warning("finish_write_cache failed, write_session_id: %s, error: %s", write_session_id, e) - with self._canceled_save_request_ids_lock: - self._canceled_save_request_ids.append(req_id) - return - - need_block_idx = self.parse_block_mask_to_save_indices(response, target_save_num) - logger.debug("target_save_num: %s, need_block_idx: %s", target_save_num, need_block_idx) - message = CoordinateMessage(time.time(), SendBlockStartEvent(request_id=req_id, - write_session_id=write_session_id, - locations=locations)) - self._coordinator_client.send(CoordinateMsgSerializer.dumps(message)) - - with self._waiting_to_save_requests_lock: - self._waiting_to_save_requests.append(SaveRequest( - req_id, - locations, - need_block_idx, - write_session_id - )) - - def handle_canceled_save_req(self): - canceled_save_req_ids = [] - with self._canceled_save_request_ids_lock: - canceled_save_req_ids = self._canceled_save_request_ids - self._canceled_save_request_ids = [] - for canceled_req_id in canceled_save_req_ids: - req = self._alive_requests[canceled_req_id] - req.sent_saving_count += 1 - if (req.need_report_after_saving_finished and - req.scheduled_saving_count == req.sent_saving_count): - self._waiting_to_finish_requests.append(FinishRequest(req.req_id)) - self._alive_requests.pop(req.req_id) - - def get_finished_count(self): - # only rank0 will return finished - return 1 - - def update_connector_output(self, connector_output: KVConnectorOutput): - """ - Update KVConnector state from worker-side connectors output. - - Args: - connector_output (KVConnectorOutput): the worker-side - connectors output. - """ - - return - - def parse_block_mask_to_save_indices(self, response: dict, target_save_num: int) -> list[int]: - # 从response中提取block_mask - block_mask = response.get("block_mask", {}) - save_indices = [] - if "offset" in block_mask: - offset = block_mask["offset"] - for idx in range(offset, target_save_num): - save_indices.append(idx) - else: - bool_masks = block_mask.get("bool_masks", {}).get("values", []) - # 找出所有为False的索引（需要保存的block） - for idx, is_saved in enumerate(bool_masks): - if not is_saved: # False表示需要保存 - save_indices.append(idx) - - return save_indices - - def request_finished( - self, - request: "Request", - block_ids: list[int], - ) -> Tuple[bool, Optional[dict[str, Any]]]: - if request.request_id not in self._alive_requests: - logger.info("request_finished not alive request: %s", request.request_id) - return False, {} - - req = self._alive_requests[request.request_id] - extra_info = {"local_matched_token_num": req.local_matched_token_num, - "remote_matched_token_num": req.remote_matched_token_num} - - if req.scheduled_saving_count == req.sent_saving_count: - self._waiting_to_finish_requests.append(FinishRequest(req.req_id)) - self._alive_requests.pop(req.req_id) - return True, extra_info - - # This request still has some save requests waiting to be issued or canceled, - # delay finishing this request - req.need_report_after_saving_finished = True - - return True, extra_info + def get_block_ids_with_load_errors(self) -> set: + assert self.connector_worker is not None + return self.connector_worker.get_block_ids_with_load_errors() diff --git a/kv_cache_manager/py_connector/vllm/vllm_common.py b/kv_cache_manager/py_connector/vllm/vllm_common.py new file mode 100644 index 000000000..b22222421 --- /dev/null +++ b/kv_cache_manager/py_connector/vllm/vllm_common.py @@ -0,0 +1,281 @@ +"""Shared vocabulary between the scheduler side and the worker side. + +Everything here is role-agnostic: the data model both cores speak +(GroupMeta), the spec naming scheme shared with the manager +registration, the KV layout normalization, the hybrid capability gate and +the kv_cache_config parsing. The two cores (scheduler_core / worker_core) +and the thin connector shell (v1_connector) build on this module; nothing +here may import them. +""" + +from dataclasses import dataclass, field +from typing import List, Optional + +import torch + +from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec + +from kv_cache_manager.py_connector.common.logger import logger +from kv_cache_manager.py_connector.vllm.transfer_types import ( + KVLayout, +) + +# Spec group names advertised at registration and used per key in +# start_write_cache. See build_spec_groups for the semantics. +# +# NOTE: the wire strings are frozen protocol: the manager keys on the +# "full" prefix (meta_searcher.cc) and its tests / the optimizer client +# emit these literals. Only the Python-side names below are free to move; +# "attn" reads as attention-only coverage, "full" as attention + every +# recurrent state group (the union, i.e. *all* specs). +ATTN_ONLY_SPEC_GROUP = "attn" +ALL_SPEC_GROUP = "full" + + +def spec_name(tp_rank: int, group_idx: int) -> str: + """Location spec name for one (tp rank, kv cache group) shard.""" + return f"tp{tp_rank}_g{group_idx}" + + +def build_spec_groups(group_metas: List["GroupMeta"], tp_size: int) -> List[dict]: + """LocationSpecGroups describing which specs a block may carry. + + Hybrid (mamba "align") models write a *sparse* set of recurrent states: + vLLM materializes a state only at segment boundaries, so the interior + manager blocks of a request have attention KV but no state. Declaring + two groups lets ``start_write_cache`` say, per block, which specs that + block will actually hold: + + * ``full`` -- *all* specs: attention KV + every recurrent state; + * ``attn`` -- attention specs only (no state was materialized). + + The manager then stores exactly the advertised specs, reports the real + per-block coverage in ``getCacheLocation``, and later lets a + complementary write fill in a block's missing state specs. + + Full-attention models have nothing to be sparse about: they declare no + groups at all, which keeps their requests byte-identical to before (and + compatible with managers that predate spec groups). + """ + state_groups = [m for m in group_metas if isinstance(m, StateGroupMeta)] + if not state_groups: + return [] + attn_specs = sorted( + spec_name(rank, meta.group_idx) + for rank in range(tp_size) + for meta in group_metas if isinstance(meta, AttentionGroupMeta)) + all_specs = sorted( + spec_name(rank, meta.group_idx) + for rank in range(tp_size) for meta in group_metas) + return [{"name": ATTN_ONLY_SPEC_GROUP, "spec_names": attn_specs}, + {"name": ALL_SPEC_GROUP, "spec_names": all_specs}] + + +@dataclass(frozen=True) +class GroupMeta: + """Static description of one kv_cache_group, derived from KVCacheConfig + (see parse_groups). Available in both scheduler and worker roles (before + tensors exist). Kind-specific subclasses carry the kind-specific sizing.""" + + group_idx: int + layer_names: List[str] + # The group's block table granularity in tokens (spec.block_size). + block_size: int + # Bytes stored per manager block for the whole group. + per_block_bytes: int + + +@dataclass(frozen=True) +class AttentionGroupMeta(GroupMeta): + """FullAttentionSpec group: token-granular KV, re-blockable to the + manager block size. Sizing derives from the *compact* page size (see + parse_groups).""" + + +@dataclass(frozen=True) +class StateGroupMeta(GroupMeta): + """MambaSpec group: one opaque state per block, verbatim byte copies. + page_size_bytes is per layer; per_block_bytes = page_size_bytes * layers.""" + + # Bytes per block per state layer (spec.page_size_bytes). + page_size_bytes: int = 0 + + +def parse_groups(kv_cache_config, manager_block_size: int) -> List[GroupMeta]: + """Derive the transferable GroupMeta list from vLLM's KVCacheConfig + (https://github.com/vllm-project/vllm/blob/v0.26.0/vllm/v1/kv_cache_interface.py#L952: + kv_cache_groups holds one KVCacheGroupSpec per block table, each with its + kv_cache_spec -- FullAttentionSpec at L227, MambaSpec at L690).""" + metas = [] + for idx, group in enumerate(kv_cache_config.kv_cache_groups): + if getattr(group, "is_eagle_group", False): + logger.warning("skip eagle group %d (%d layers)", idx, len(group.layer_names)) + continue + spec = group.kv_cache_spec + if isinstance(spec, MambaSpec): + metas.append(StateGroupMeta( + group_idx=idx, + layer_names=list(group.layer_names), + block_size=spec.block_size, + per_block_bytes=spec.page_size_bytes * len(group.layer_names), + page_size_bytes=spec.page_size_bytes, + )) + elif isinstance(spec, FullAttentionSpec): + # FullAttentionSpec doubles as the merged spec of hybrid + # SWA/chunked-attention models (vLLM merges window layers into + # it, keeping sliding_window/attention_chunk_size set). Those + # blocks hold windowed KV, not the full prefix -- publishing + # them as prefix caches would corrupt reuse. Refuse explicitly. + for window_field in ("sliding_window", "attention_chunk_size"): + if getattr(spec, window_field, None) is not None: + raise NotImplementedError( + f"group {idx}: FullAttentionSpec has {window_field}=" + f"{getattr(spec, window_field)}; sliding-window / " + f"chunked attention KV is not full-prefix and is " + f"not yet supported by TairKvCacheConnector") + # Attention KV is token-granular; scale from the spec's page size + # to the manager block size. Use the *compact* page size: + # spec.page_size_bytes returns page_size_padded when set, which + # includes an allocation-alignment gap the gather kernel never + # copies -- sizing locations/staging buffers with it would break + # the staging view() and waste storage. real_page_size_bytes is + # exactly the raw KV bytes (2 * block * heads * head_dim * dtype). + compact_page_bytes = getattr(spec, "real_page_size_bytes", None) + if compact_page_bytes is None: + if getattr(spec, "page_size_padded", None) is not None: + raise NotImplementedError( + f"group {idx}: page_size_padded=" + f"{spec.page_size_padded} but this vLLM exposes no " + f"real_page_size_bytes to recover the compact page " + f"size; padded attention layouts are unsupported here") + compact_page_bytes = spec.page_size_bytes + per_token_bytes = compact_page_bytes // spec.block_size + metas.append(AttentionGroupMeta( + group_idx=idx, + layer_names=list(group.layer_names), + block_size=spec.block_size, + per_block_bytes=per_token_bytes * manager_block_size * len(group.layer_names), + )) + else: + raise NotImplementedError( + f"Unsupported kv cache spec {type(spec).__name__} in group {idx}") + if not metas: + # Every group was skipped (all-EAGLE config or an empty group list): + # nothing to transfer, refuse explicitly instead of asserting. + raise NotImplementedError( + "no usable kv cache groups (all groups skipped?)") + if not any(isinstance(m, AttentionGroupMeta) for m in metas): + # Pure-mamba / attention-free models have no attention KV to + # transfer; the register_kv_caches path would fail obscurely later + # (it requires an attention layer tensor). Refuse before init. + raise NotImplementedError( + "pure-mamba / attention-free models are not supported: " + "TairKvCacheConnector transfers full-attention or hybrid " + "(attention + mamba) KV caches only") + return metas + + +def attn_kv_views(ref: torch.Tensor) -> tuple: + """Normalize one attention layer's paged KV cache into per-pointer views. + + vLLM's flash_attn backend changed ``get_kv_cache_shape`` twice; the three + layouts (see KVLayout for the per-version source links) are detected from + the tensor shape itself (never from version strings): + + * 4-D ``(num_blocks, H, block, 2*D)`` -- K/V packed into the content dim + (vLLM >= 0.26.0). One transfer pointer per layer. + * 5-D ``(num_blocks, 2, block, H, D)`` -- N-first split K/V + (vLLM 0.23.0 - 0.25.x). Two pointers per layer: ``t[:, 0]`` / ``t[:, 1]``. + * 5-D ``(2, num_blocks, block, H, D)`` -- KV-first split K/V + (vLLM <= 0.22.1). Two pointers per layer: ``t[0]`` / ``t[1]``. + + Returns ``(views, layout)``. Every view has the logical shape + ``(num_blocks, kernel_block_size, heads, content_dim)`` matching the NHD + memory order, so all downstream math (per-token dim, token-major check, + block stride, data_ptr) is layout-independent -- the layout travels along + only for traceability. Unrecognized layouts raise. + """ + if ref.dim() == 4: + # Packed content dim; permute to token-major logical order. The permuted + # view shares storage, data_ptr() is the storage base. + return [ref.permute(0, 2, 1, 3)], KVLayout.PACKED_4D + if ref.dim() == 5: + kv_first = ref.shape[0] == 2 + n_first = ref.shape[1] == 2 + if kv_first and n_first: + raise NotImplementedError( + f"ambiguous kv layout {tuple(ref.shape)}: cannot tell the K/V " + f"dim from a num_blocks dim of size 2") + if kv_first: + return [ref[0], ref[1]], KVLayout.SPLIT_KV_5D_KV_FIRST + if n_first: + return [ref[:, 0], ref[:, 1]], KVLayout.SPLIT_KV_5D_N_FIRST + raise NotImplementedError( + f"unrecognized kv cache layout {tuple(ref.shape)}; expected the packed " + f"4-D (vllm >= 0.26.0) or one of the split K/V 5-D layouts " + f"(vllm <= 0.25.x)") + + +def _hybrid_external_load_supported() -> Optional[bool]: + """vLLM <= 0.22.x cannot combine mamba align mode with a KV connector: + ``Scheduler._mamba_block_aligned_split`` asserts + ``num_external_computed_tokens == 0`` ("External KV connector is not + verified yet"), so the first external match would crash the scheduler. + Probe the installed vLLM for that blocking assert (a capability check, + not a version-string comparison). + + Returns: + True -- supported (assert absent, or the method was removed by a + newer vLLM: the assert went away with it); + False -- unsupported (the blocking assert is present); + None -- the method exists but its source is unavailable (frozen / + bytecode-only install), so the assert cannot be ruled out. + """ + try: + from vllm.v1.core.sched.scheduler import Scheduler + method = Scheduler._mamba_block_aligned_split + except (ImportError, AttributeError): + # No such method: the blocking assert was removed/refactored away. + return True + try: + import inspect + src = inspect.getsource(method) + except Exception: + return None # method exists but cannot be inspected + return "External KV connector is not verified yet" not in src + + +def ensure_hybrid_supported(force: bool = False): + """Fail fast with a clear message when a hybrid (mamba) model is served on + a vLLM whose scheduler rejects external KV loads (vllm <= 0.22.x). + + When the probe is inconclusive (method present but source unavailable) the + gate fails closed: a wrong guess would crash the scheduler on the first + external match. ``force`` (extra_config ``force_hybrid_support``) bypasses + the inconclusive case for source-restricted environments.""" + supported = _hybrid_external_load_supported() + if supported: + return + if supported is None: + if force: + logger.warning( + "force_hybrid_support=true: skipping the hybrid external-load " + "capability probe; if this vLLM's scheduler still asserts " + "'External KV connector is not verified yet' the first " + "external match will crash it") + return + raise NotImplementedError( + "TairKvCacheConnector: cannot verify that this vLLM supports " + "hybrid (mamba) models with an external KV connector -- " + "Scheduler._mamba_block_aligned_split exists but its source is " + "unavailable, so the vllm <= 0.22.x blocking assert cannot be " + "ruled out. If you know this vLLM is >= 0.23.0, set " + "kv_connector_extra_config {\"force_hybrid_support\": true} to " + "bypass this check.") + raise NotImplementedError( + "TairKvCacheConnector: this vLLM version cannot combine hybrid " + "(mamba) models with an external KV connector -- its scheduler " + "asserts num_external_computed_tokens == 0 in " + "_mamba_block_aligned_split ('External KV connector is not " + "verified yet'). Upgrade to vLLM >= 0.23.0 for hybrid model " + "support; full-attention models are unaffected.") diff --git a/open_source/deps/requirements_base.txt b/open_source/deps/requirements_base.txt index d372ae37c..3d4a19893 100644 --- a/open_source/deps/requirements_base.txt +++ b/open_source/deps/requirements_base.txt @@ -3,3 +3,5 @@ grpcio==1.62.0 grpcio-tools==1.62.0 protobuf==4.25 requests==2.32.5 +pydantic==2.13.4 +orjson==3.12.0 diff --git a/open_source/deps/requirements_lock_cpu.txt b/open_source/deps/requirements_lock_cpu.txt index c55d8e9fc..fb4a8bd8d 100644 --- a/open_source/deps/requirements_lock_cpu.txt +++ b/open_source/deps/requirements_lock_cpu.txt @@ -6,6 +6,10 @@ # --index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple +annotated-types==0.8.0 \ + --hash=sha256:13b2beaad985e05e2d6407ee4c4f35590b11f8d693a258a561055cac8f64cab7 \ + --hash=sha256:f072f4d804ea359e4eaf198b1af7a8b0943881a87f31bb764f8bf219bb9419e0 + # via pydantic certifi==2025.11.12 \ --hash=sha256:97de8790030bbd5c2d96b7ec782fc2f7820ef8dba6db909ccf95449f2d062d4b \ --hash=sha256:d8ab5478f2ecd78af242878415affce761ca6bc54a22a27e026d7c25357c3316 @@ -243,6 +247,73 @@ idna==3.11 \ --hash=sha256:771a87f49d9defaf64091e6e6fe9c18d4833f140bd19464795bc32d966ca37ea \ --hash=sha256:795dafcc9c04ed0c1fb032c2aa73654d8e8c5023a7df64a53f39190ada629902 # via requests +orjson==3.12.0 \ + --hash=sha256:010811c1b69773450a01cef97727a67b223242f350b77d4ca000e59a9ef2155a \ + --hash=sha256:01efac2074fffb4cb1ea3fab7861e9d0f2a26913854a972f5ac760525dbdaf6e \ + --hash=sha256:03091c8a64db4be38746597ceea68f33c238e27acd9bfe99fb59420224ae7a55 \ + --hash=sha256:08231552159be266a7269555bd9f7c016aee7d9ad6dab06eb58796c5ccb7101c \ + --hash=sha256:0b1ac5bf6609b2716c7954011c5fef6254922df029f45d032ee4ebf5d363cbed \ + --hash=sha256:103b5db66aa53c1f9e88c2524be4f383e831ba7dfd5f9f5af6336a177c622f11 \ + --hash=sha256:1192a7021b6d071aaf909864f6e924d6a2675ca360485b972b8401749311750b \ + --hash=sha256:11edb4660a6680abee9788a3a9072208a2c96538cc1322bd79542065229d8e54 \ + --hash=sha256:18a87929f31d94a77f7dc93cf527e91f39ce7fe7813d588a4de2507efd32a387 \ + --hash=sha256:1c680706fc8396d95e7c4c1f9482563f552137aef91b57237a3ad5aaf64629df \ + --hash=sha256:2b7bcefb9f40fa242fa6b06377232c048e655747790829609168c01162f60578 \ + --hash=sha256:2bb3ce43203936072dd8b4917b01d3aecfc02329bfb42510cb7cfb24708adc9c \ + --hash=sha256:2d3a9da945a4d96ae758fdaaca56742e6b73b6fd554c5d8876f252a6dad70b83 \ + --hash=sha256:2eb5c56e534127b2b8fa38d2363c8b1b8190367ee0d1d16c041517d880843b94 \ + --hash=sha256:31ed278a36304390adc3eec5d7f6fd593a7c3e99e5a06cd07866396c4b1b4710 \ + --hash=sha256:33efefcf5d88eaf400b47e2eba02f91f319bb9951be61ca500b7d536d3f2079d \ + --hash=sha256:3bb17a06f9bd15237b3216c044209fe92597379124018cfc196fbb846cde64df \ + --hash=sha256:3dbce9b6b3074b31a5d5dd322a9c4e5b16f206091ece4194c2e36952847a105e \ + --hash=sha256:40f92192227505acca4e2533ce565f8e6b9535f7d0d09b0968452f18b7376b38 \ + --hash=sha256:477ecaf6b9f88f873341b91fcc736119ca81b5e002a9f7f308ff5b4f2ce2a70e \ + --hash=sha256:50fae885cb073eac7556353ff3df93312b0d5137b0a5056b2bb63f97ed9a93c7 \ + --hash=sha256:532ff8cd4bd59a327a953a7dcde922c7fc25b85e29721bb8633265430d3a3873 \ + --hash=sha256:53c0c474a9d9aff9aebfc0c88de1f28f843d940e6e3a80729abdf6a20274356f \ + --hash=sha256:58c58e1de0006ffb580368d6793c36c7b0b021db066479cf281bf5061e732328 \ + --hash=sha256:5a0fdbc216388f653d3752ff310e710f59253bd4ed6a2bfb3f4f06b84714bbd8 \ + --hash=sha256:61318b6de893c7a9d9f3e5ecbadccbfc26a7eb417ccc7bbf0771de3b4d72f868 \ + --hash=sha256:644d005bc82f917337a95ce270c9f6f92f9834c2bed7b1477572f8db00784222 \ + --hash=sha256:6a2a79c89984dc719817d388c8709e0efc2a2795a934eaa746b4882eb6045adc \ + --hash=sha256:6a31348d7dfa64cd9c78bd1f510ff44c48fe64d71094e6b90e364dba3b55949e \ + --hash=sha256:747843254519dd43b93eee3153a19e5a509334320c4d2f823ec879232db5c796 \ + --hash=sha256:784106539f4b9d4b930e0b4eb8d45168507dae001945e71b4675a367f1e5e806 \ + --hash=sha256:7c2ad193c8004254f34b499f3bd2c80f043d10754aff2b38f93da574f4883f98 \ + --hash=sha256:83445adc40cba26d6d621185a45128ce455b766af368cad2ab64b970603a7978 \ + --hash=sha256:859fc4196855890150bb08e649b30d2c93b249b3e3edd0d3bb2231abf8aa8adc \ + --hash=sha256:8c3bb86dd10f39b3fbf434b7d5dc7cac77d6fc8ac572ae30a10731ede2c4b647 \ + --hash=sha256:8e29957429c35bbb5a185a119c523aa2428b7bbf1a293724c7b9375ed8f892a3 \ + --hash=sha256:8e386b0bc0ddd7cd2056f884b5a0af33592bd01ac66a7ca4b42a65a7e7774a13 \ + --hash=sha256:92ffc09e07233a6ab6d4e067f7841edcbcc134cb4812155cf171ea5255a421d7 \ + --hash=sha256:9a36ec60f1796f9a3f13e3b98390295e17a1c7c10155b448d264098bf9ee5900 \ + --hash=sha256:9caf3d09f47c3c70c4451ada20ef9bc4a4cdffa26f49862cf0a253b329aae2d5 \ + --hash=sha256:9e6fee342a48760e854d743e7a81534d8e2925a6f46e09f750cf56b50fd1de5d \ + --hash=sha256:a15f9a891bce5f5cc5d210e3ad8614d4d1b489a56448c099d6d2a7168b2d954a \ + --hash=sha256:a696529ec96a90d9a5f9570207efe403c8b08f8e4aa2783ee3403511e2fdfa10 \ + --hash=sha256:a6cf4b18e7de173f209f2084ffbd736dd72389a396326ee80a7022168be232e5 \ + --hash=sha256:a791f793b287bbc135b8e87c34e35c8bfc693e2a8a620fab1ae682b925f9a32e \ + --hash=sha256:a94f0f0c6fcbb2b5bd9734c57a489c7584a732bbdf04a39e8c83b861e9d03e92 \ + --hash=sha256:aa3e43a6846e91d7bde3d5a9c66090fcd8744f569a9b6cffc5e1ca38f6a461c0 \ + --hash=sha256:ad0422b92d5195443a39f80c3bcf731cc2e00f153bd32063a47b73b057bd0f03 \ + --hash=sha256:ad29eece0c601737f2a60edc2752a84e7a0785df3efb62e3012834700a5afe0d \ + --hash=sha256:b85931be5b6763c31283805c9bdaae1ca03ad9f6f12a15f1cbf6745b907932c2 \ + --hash=sha256:b9dca132b1fda5565088e65a6b6e742285e0aeceb6fae549fa8863e16c7d3998 \ + --hash=sha256:bc7a872f03522d90e0429e6c0c5cd23084f767bedcb4c58048eec19294613344 \ + --hash=sha256:bd57d79aefa3f84eec851d6de7a366795b9345cfaf17f82b4820430a7a5fa241 \ + --hash=sha256:bf44e374aadde77b1f6109f1030be51433eb61984379852766b6f4e187db7b1e \ + --hash=sha256:c6b11be792c3d2c6a4be2af4ebf97a68d0bf5f580aca6e86a418a354f6cc846a \ + --hash=sha256:d14203fb1aae2ad9b3d52f8a0e82aeb10197ef1c9bc61da7f358bd70b00123d5 \ + --hash=sha256:d39f3f5c3927e2dc0913fe5bbc1a2f6b1b9d1bba1de6358340d0ad0d0c00ca92 \ + --hash=sha256:d8e78d3d93705e3d27cc17cdb209e44d7a8ea203010cac6ce9c7ffc1ae1996f1 \ + --hash=sha256:dce0166feb0a737ab84f598c9a338cbc0b764a036617aa686194f53c7eba0c3e \ + --hash=sha256:e4ac5059baab4b3acbd99485de019ff8cda0fdf34b61fa74f7197a53db78bfe8 \ + --hash=sha256:e9683ee9ea0659da64f36574ef675b8a86330c34c19ea75db1fb93c3ff99e0ef \ + --hash=sha256:ed4ca42bd55955aa34deedcfdfd0e0c31abf51143aae158ae2bc3520b626e517 \ + --hash=sha256:f06dd838d1e07d9b1de0932ec0485ec92c4d5f5d1ad4817a656268c3e88be1e1 \ + --hash=sha256:f3c0683136acdc29afdf88a5bc2f7d3d0e34087788d1d63c0144b805a87a196f \ + --hash=sha256:fb2539159dfe8d371914f354360fa50e4a577cc89222a3828b9650a5e5040252 + # via -r open_source/deps/./requirements_base.txt protobuf==4.25.0 \ --hash=sha256:1a3ba712877e6d37013cdc3476040ea1e313a6c2e1580836a94f76b3c176d575 \ --hash=sha256:1a53d6f64b00eecf53b65ff4a8c23dc95df1fa1e97bb06b8122e5a64f49fc90a \ @@ -258,10 +329,147 @@ protobuf==4.25.0 \ # via # -r open_source/deps/./requirements_base.txt # grpcio-tools +pydantic==2.13.4 \ + --hash=sha256:45a282cde31d808236fd7ea9d919b128653c8b38b393d1c4ab335c62924d9aba \ + --hash=sha256:c40756b57adaa8b1efeeced5c196f3f3b7c435f90e84ea7f443901bec8099ef6 + # via -r open_source/deps/./requirements_base.txt +pydantic-core==2.46.4 \ + --hash=sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0 \ + --hash=sha256:0186750b482eefa11d7f435892b09c5c606193ef3375bcf94aa00ae6bfb66262 \ + --hash=sha256:041bde0a48fd37cf71cab1c9d56d3e8625a3793fef1f7dd232b3ff37e978ecda \ + --hash=sha256:0c563b08bca408dc7f65f700633d8442fffb2421fc47b8101377e9fd65051ff0 \ + --hash=sha256:0cbe8b01f948de4286c74cdd6c667aceb38f5c1e26f0693b3983d9d74887c65e \ + --hash=sha256:0ce40cd7b21210e99342afafbd4d0f76d784eb5b1d60f3bdc566be4983c6c73b \ + --hash=sha256:0e96592440881c74a213e5ad528e2b24d3d4f940de2766bed9010ab1d9e51594 \ + --hash=sha256:10e17cbb10a330363733efc4d7c4d0dd827ac0909b8f6a6542298fed1ea62f29 \ + --hash=sha256:133878133d271ade3d41d1bfb2a45ec38dbdbda40bc065921c6b04e4630127e2 \ + --hash=sha256:14d4edf427bdcf950a8a02d7cb44a08614388dd6e1bdcbf4f67504fa7887da9c \ + --hash=sha256:14f4c5d6db102bd796a627bbb3a17b4cf4574b9ae861d8b7c9a9661c6dd3362d \ + --hash=sha256:17299feefe090f2caa5b8e37222bb5f663e4935a8bfa6931d4102e5df1a9f398 \ + --hash=sha256:184c081504d17f1c1066e430e117142b2c77d9448a97f7b65c6ac9fd9aee238d \ + --hash=sha256:18e5ceec2ab67e6d5f1a9085e5a24c9c4e2ac4545730bfe668680bca05e555f3 \ + --hash=sha256:19e51f073cd3df251856a8a4189fbdf1de4012c3ebacfb1884f94f1eb406079f \ + --hash=sha256:1a7dd0b3ee80d90150e3495a3a13ac34dbcbfd4f012996a6a1d8900e91b5c0fb \ + --hash=sha256:1d8ba486450b14f3b1d63bc521d410ec7565e52f887b9fb671791886436a42f7 \ + --hash=sha256:2108ba5c1c1eca18030634489dc544844144ee36357f2f9f780b93e7ddbb44b5 \ + --hash=sha256:228ee9bae8bef5b1e97ec58302f80357c37199e0d0a99174e138d28e6957b9d9 \ + --hash=sha256:23ace664830ee0bfe014a0c7bc248b1f7f25ed7ad103852c317624a1083af462 \ + --hash=sha256:2412e734dcb48da14d4e4006b82b46b74f2518b8a26ee7e58c6844a6cd6d03c4 \ + --hash=sha256:29c61fc04a3d840155ff08e475a04809278972fe6aef51e2720554e96367e34b \ + --hash=sha256:2f84c03c8607173d16b5a854ec68a2f9079ae03237a54fb506d13af47e1d018d \ + --hash=sha256:3009f12e4e90b7f88b4f9adb1b0c4a3d58fe7820f3238c190047209d148026df \ + --hash=sha256:3245406455a5d98187ec35530fd772b1d799b26667980872c8d4614991e2c4a2 \ + --hash=sha256:3447661d99f75a3683a4cf5c87da72f2161964611864dbbeac7fbb118bb4bfc0 \ + --hash=sha256:372429a130e469c9cd698925ce5fc50940b7a1336b0d82038e63d5bbc4edc519 \ + --hash=sha256:395aebd9183f9d112f569aeb5b2214d1a10a33bec8456447f7fbdfa51d38d4cd \ + --hash=sha256:3a233125ac121aa3ffba9a2b59edfc4a985a76092dc8279586ab4b71390875e7 \ + --hash=sha256:3be77f45df024d789a672ae34f8b06fb346c4f9f46ea714956660ea4862e89ac \ + --hash=sha256:3bf92c5d0e00fefaab325a4d27828fe6b6e2a21848686b5b60d2d9eeb09d76c6 \ + --hash=sha256:3ecbc122d18468d06ca279dc26a8c2e2d5acb10943bb35e36ae92096dc3b5565 \ + --hash=sha256:3fb702cd90b0446a3a1c5e470bfa0dd23c0233b676a9099ddcc964fa6ca13898 \ + --hash=sha256:428e04521a40150c85216fc8b85e8d39fece235a9cf5e383761238c7fa9b96fb \ + --hash=sha256:432c179df7874eeb73307aad2df0755e1ae0efa61ff0ea89b93e194411ae3928 \ + --hash=sha256:4a05d69cba51d852c5c3e92758653245a50c0b646ced0cf05bd793ed592839d6 \ + --hash=sha256:4c63ebc82684aa89d9a3bcbd13d515b3be44250dc68dd3bd81526c1cb31286c3 \ + --hash=sha256:4fc73cb559bdb54b1134a706a2802a4cddd27a0633f5abb7e53056268751ac6a \ + --hash=sha256:4fcbe087dbc2068af7eda3aa87634eba216dbda64d1ae73c8684b621d33f6596 \ + --hash=sha256:56cb4851bcaf3d117eddcef4fe66afd750a50274b0da8e22be256d10e5611987 \ + --hash=sha256:5855698a4856556d86e8e6cd8434bc3ac0314ee8e12089ae0e143f64c6256e4e \ + --hash=sha256:5a4330cdbc57162e4b3aa303f588ba752257694c9c9be3e7ebb11b4aca659b5d \ + --hash=sha256:5b712b53160b79a5850310b912a5ef8e57e56947c8ad690c227f5c9d7e561712 \ + --hash=sha256:5d5902252db0d3cedf8d4a1bc68f70eeb430f7e4c7104c8c476753519b423008 \ + --hash=sha256:617d7e2ca7dcb8c5cf6bcb8c59b8832c94b36196bbf1cbd1bfb56ed341905edd \ + --hash=sha256:62f875393d7f270851f20523dd2e29f082bcc82292d66db2b64ea71f64b6e1c1 \ + --hash=sha256:633147d34cf4550417f12e2b1a0383973bdf5cdfde212cb09e9a581cf10820be \ + --hash=sha256:66ce7632c22d837c95301830e111ad0128a32b8207533b60896a96c4915192ea \ + --hash=sha256:6b3ace8194b0e5204818c92802dcdca7fc6d88aabbb799d7c795540d9cd6d292 \ + --hash=sha256:6f2eeda33a839975441c86a4119e1383c50b47faf0cbb5176985565c6bb02c33 \ + --hash=sha256:7027560ee92211647d0d34e3f7cd6f50da56399d26a9c8ad0da286d3869a53f3 \ + --hash=sha256:7283d57845ecf5a163403eb0702dfc220cc4fbdd18919cb5ccea4f95ee1cdab4 \ + --hash=sha256:7a5f930472650a82629163023e630d160863fce524c616f4e5186e5de9d9a49b \ + --hash=sha256:7bfb192b3f4b9e8a89b6277b6ce787564f62cfd272055f6e685726b111dc7826 \ + --hash=sha256:811ff8e9c313ab425368bcbb36e5c4ebd7108c2bbf4e4089cfbb0b01eff63fac \ + --hash=sha256:8233f2947cf85404441fd7e0085f53b10c93e0ee78611099b5c7237e36aacbf7 \ + --hash=sha256:82cf5301172168103724d49a1444d3378cb20cdee30b116a1bd6031236298a5d \ + --hash=sha256:8358a950c8909158e3df31538a7e4edc2d7265a7c54b47f0864d9e5bae9dcebf \ + --hash=sha256:85bb3611ff1802f3ee7fdd7dbff26b56f343fb432d57a4728fdd49b6ef35e2f4 \ + --hash=sha256:86e1a4418c6cd97d60c95c71164158eaf7324fae7b0923264016baa993eba6fc \ + --hash=sha256:8b9bab013d1c7a79d3501ff86d0bc9c31bf587db4551677b96bec07df78c6b15 \ + --hash=sha256:8c5dac79fa1614d1e06ca695109c6105923bd9c7d1d6c918d4e637b7e6b32fd3 \ + --hash=sha256:8d0820e8192167f80d88d64038e609c31452eeca865b4e1d9950a27a4609b00b \ + --hash=sha256:8daafc69c93ee8a0204506a3b6b30f586ef54028f52aeeeb5c4cfc5184fd5914 \ + --hash=sha256:9037063db01f09b09e237c282b6792bd4da634b5402c4e7f0c61effed7701a04 \ + --hash=sha256:905a0ed8ea6f2d61c1738835f99b699348d7857379083e5fc497fa0c967a407c \ + --hash=sha256:90884113d8b48f760e9587002789ddd741e76ab9f89518cd1e43b1f1a52ec44b \ + --hash=sha256:91a06d2e259ecfbd8c901d70c3c507900458498142b3026a296b7de4d1322cc9 \ + --hash=sha256:926c9541b14b12b1681dca8a0b75feb510b06c6341b70a8e500c2fdcff837cce \ + --hash=sha256:9401557acd873c3a7f3eb9383edef8ac4968f9510e340f4808d427e75667e7b4 \ + --hash=sha256:9551187363ffc0de2a00b2e47c25aeaeb1020b69b668762966df15fc5659dd5a \ + --hash=sha256:962ccbab7b642487b1d8b7df90ef677e03134cf1fd8880bf698649b22a69371f \ + --hash=sha256:97e7cf2be5c77b7d1a9713a05605d49460d02c6078d38d8bef3cbe323c548424 \ + --hash=sha256:9aa768456404a8bf48a4406685ac2bec8e72b62c69313734fa3b73cf33b3a894 \ + --hash=sha256:9bc519fbf2b7578398853d815009ae5e4d4603d12f4e3f91da8c06852d3da3e9 \ + --hash=sha256:9d56801be94b86a9da183e5f3766e6310752b99ff647e38b09a9500d88e46e76 \ + --hash=sha256:9f444c499b3eefd3a92e348059471ea0c3a6e303d9c1cec09fa748fd9f895201 \ + --hash=sha256:9fa8ae11da9e2b3126c6426f147e0fba88d96d65921799bb30c6abd1cb2c97fb \ + --hash=sha256:a0f62d0a58f4e7da165457e995725421e0064f2255d8eccebc49f41bbc23b109 \ + --hash=sha256:a396dcc17e5a0b164dbe026896245a4fa9ff402edca1dff0be3d53a517f74de4 \ + --hash=sha256:aaa2a54443eff1950ba5ddc6b6ccda0d9c84a364276a62f969bdf2a390650848 \ + --hash=sha256:ad785e92e6dc634c21555edc8bd6b64957ab844541bcb96a1366c202951ae526 \ + --hash=sha256:af8244b2bef6aaad6d92cda81372de7f8c8d36c9f0c3ea36e827c60e7d9467a0 \ + --hash=sha256:b078afbc25f3a1436c7a1d2cd3e322497ee99615ba97c563566fdf46aff1ee01 \ + --hash=sha256:b2f69dec1725e79a012d920df1707de5caf7ed5e08f3be4435e25803efc47458 \ + --hash=sha256:b8458003118a712e66286df6a707db01c52c0f52f7db8e4a38f0da1d3b94fc4e \ + --hash=sha256:bb63e0198ca18aad131c089b9204c23079c3afa95487e561f4c522d519e55aba \ + --hash=sha256:bfec22eab3c8cc2ceec0248aec886624116dc079afa027ecc8ad4a7e62010f8a \ + --hash=sha256:c1747f85cee84c26985853c6f3d9bd3e75da5212912443fa111c113b9c246f39 \ + --hash=sha256:c1b3f518abeca3aa13c712fd202306e145abf59a18b094a6bafb2d2bbf59192c \ + --hash=sha256:c50f2528cf200c5eed56faf3f4e22fcd5f38c157a8b78576e6ba3168ec35f000 \ + --hash=sha256:c68fcd102d71ea85c5b2dfac3f4f8476eff42a9e078fd5faefff6d145063536b \ + --hash=sha256:c7a7bd4e39e8e4c12c39cd480356842b6a8a06e41b23a55a5e3e191718838ddf \ + --hash=sha256:c94f0688e7b8d0a67abf40e57a7eaaecd17cc9586706a31b76c031f63df052b4 \ + --hash=sha256:cbaf13819775b7f769bf4a1f066cb6df7a28d4480081a589828ef190226881cd \ + --hash=sha256:cd2213145bcc2ba85884d0ac63d222fece9209678f77b9b4d76f054c561adb28 \ + --hash=sha256:ce5c1d2a8b27468f433ca974829c44060b8097eedc39933e3c206a90ee49c4a9 \ + --hash=sha256:d396ec2b979760aaf3218e76c24e65bd0aca24983298653b3a9d7a45f9e47b30 \ + --hash=sha256:d51026d73fcfd93610abc7b27789c26b313920fcfb20e27462d74a7f8b06e983 \ + --hash=sha256:d80ee3d731373b24cebbc10d689ca4ee1875caf0d5703a245db18efd4dd37fc1 \ + --hash=sha256:d995260fdf4e1db774581b4900e0f832abe3c7c84996726bbc161b19c8f29e76 \ + --hash=sha256:da4b951fe36dc7c3a1ccb4e3cd1747c3542b8c9ceede8fc86cae054e764485f5 \ + --hash=sha256:daa27d92c36f24388fe3ad306b174781c747627f134452e4f128ea00ce1fe8c4 \ + --hash=sha256:db06ffe51636ffe9ca531fe9023dd64bdd794be8754cb5df57c5498ae5b518a7 \ + --hash=sha256:e0d65b8c354be7fb5f720c3caa8bc940bc2d20ce749c8e06135f07f8ed95dd7c \ + --hash=sha256:e68b7a074f65a2fd746c52a7ce6142ab7006074ac269ace0c25cd8ba171f8066 \ + --hash=sha256:e739fee756ba1010f8bcccb534252e85a35fe45ae92c295a06059ce58b74ccd3 \ + --hash=sha256:e846ae7835bf0703ae43f534ab79a867146dadd59dc9ca5c8b53d5c8f7c9ef02 \ + --hash=sha256:e9c26f834c65f5752f3f06cb08cb86a913ceb7274d0db6e267808a708b46bc89 \ + --hash=sha256:ea793e075b70290d89d8142074262885d3f7da19634845135751bd6344f73b50 \ + --hash=sha256:f027324c56cd5406ca49c124b0db10e56c69064fec039acc571c29020cc87c76 \ + --hash=sha256:f13a646d65d09fbf1bc6b3a9635d30095c8e7e5cc419ff35ecc563c5fd04cd49 \ + --hash=sha256:f47286a97f0bc9b8859519809077b91b2cefe4ae47fcbf5e466a009c1c5d742b \ + --hash=sha256:f747929cf940cddb5b3668a390056ddd5ba2e5010615ea2dcf4f9c4f3ab8791d \ + --hash=sha256:f99626688942fb746e545232e7726926f3be91b5975f8b55327665fafda991c7 \ + --hash=sha256:f9fa868638bf362d3d138ea55829cefb3d5f4b0d7f142234382a15e2485dbec4 \ + --hash=sha256:fbdb89b3e1c94a30cc5edfce477c6e6a5dc4d8f84665b455c27582f211a1c72c \ + --hash=sha256:fc010ab034c8c7452522748bf937df58020d256ccae0874463d1f4d01758af8e \ + --hash=sha256:fc3e9034a63de20e15e8ade85358bc6efc614008cab72898b4b4952bea0509ff \ + --hash=sha256:fd8b3d9fd264be37976686c7f65cd52a83f5e84f4bfd2adf9c1d469676bbb6ae + # via pydantic requests==2.32.5 \ --hash=sha256:2462f94637a34fd532264295e186976db0f5d453d1cdd31473c85a6a161affb6 \ --hash=sha256:dbba0bac56e100853db0ea71b82b4dfd5fe2bf6d3754a8893c3af500cec7d7cf # via -r open_source/deps/./requirements_base.txt +typing-extensions==4.16.0 \ + --hash=sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8 \ + --hash=sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5 + # via + # pydantic + # pydantic-core + # typing-inspection +typing-inspection==0.4.4 \ + --hash=sha256:547274fa6b0a561ccf549cc9524b999a578e737d015d8709d021f9d0d13bea47 \ + --hash=sha256:65b8397ba37ccbce054456aaccddfc91e6e3083c92824df348d96ca832f3f147 + # via pydantic urllib3==2.5.0 \ --hash=sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760 \ --hash=sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc