support hybrid attention for vllm connector && add integration tests

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",
+    ],
+)

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`:
+
+```
+<TEST_TMPDIR>/kvcm_vllm_e2e/<scenario>/
+  manager/manager.stdout|stderr     # KVCM manager
+  vllm/vllm*.stdout|stderr          # vLLM (EngineCore tracebacks are here)
+  captures/{ref|loaded}_tp{rank}_{token_hash}.pt
+```

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,
+)