support hybrid attention for vllm connector && add integration test

.github/workflows/test-vllm-e2e.yml

@@ -0,0 +1,111 @@
+name: test-vllm-e2e
+permissions:
+  contents: read
+on:
+  pull_request:
+    branches: ["main"]
+  workflow_dispatch:
+    inputs:
+      runs-on:
+        description: "GPU runner label (needs >= 2 GPUs, e.g. A10)"
+        type: string
+        default: "gpu-a10-x2"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
+jobs:
+  vllm-e2e:
+    name: vllm-e2e (${{ matrix.kind }})
+    # GPU runners are self-hosted; the default label can be overridden per-repo
+    # via the VLLM_E2E_RUNS_ON variable or the workflow_dispatch input.
+    runs-on: ${{ inputs.runs-on || vars.VLLM_E2E_RUNS_ON || 'gpu-a10-x2' }}
+    timeout-minutes: 180
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          # Full-attention model: single FullAttentionSpec kv cache group.
+          - kind: full-attention
+            model_var: VLLM_E2E_MODEL_FULL_ATTN
+          # Hybrid model: MambaSpec groups + FullAttentionSpec group
+          # (mamba_cache_mode="align").
+          - kind: hybrid-attention
+            model_var: VLLM_E2E_MODEL_HYBRID
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: check_gpus
+        run: |
+          nvidia-smi --query-gpu=index,name,memory.total --format=csv,noheader
+          GPU_COUNT=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
+          if [ "$GPU_COUNT" -lt 2 ]; then
+            echo "::error::This test needs at least 2 GPUs, found $GPU_COUNT"
+            exit 1
+          fi
+
+      - name: resolve_runner_config
+        # Self-hosted runner provisioning: a vLLM (>= 0.26.0) venv and the model
+        # checkouts are pre-cached on the runner and exposed via repo variables:
+        #   VLLM_E2E_PYTHON            venv python with vllm installed
+        #   VLLM_E2E_MODEL_FULL_ATTN   e.g. a local Qwen2.5-7B-Instruct checkout
+        #   VLLM_E2E_MODEL_HYBRID      e.g. a local Qwen3.5-4B checkout
+        env:
+          VENV_PY: ${{ vars.VLLM_E2E_PYTHON }}
+          MODEL: ${{ vars[matrix.model_var] }}
+        run: |
+          if [ ! -x "$VENV_PY" ]; then
+            echo "::error::VLLM_E2E_PYTHON ($VENV_PY) not found; provision the runner"
+            exit 1
+          fi
+          "$VENV_PY" -c 'import vllm; v = vllm.__version__; print("vllm", v)'
+          if [ ! -f "$MODEL/config.json" ]; then
+            echo "::error::model not found at $MODEL; provision the runner"
+            exit 1
+          fi
+          echo "KVCM_E2E_PYTHON=$VENV_PY" >> "$GITHUB_ENV"
+          echo "KVCM_E2E_MODEL=$MODEL" >> "$GITHUB_ENV"
+
+      - name: build_binaries
+        run: |
+          set -x
+          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'
+
+      - name: install_wheels
+        # Bazel wheel filenames contain unstamped {STABLE_*} template variables;
+        # read the real version from the wheel METADATA and rename before install.
+        run: |
+          set -x
+          mkdir -p /tmp/kvcm_whl && rm -f /tmp/kvcm_whl/*.whl
+          for whl in bazel-bin/kv_cache_manager/client/pybind/kvcm_py_client-*.whl \
+                     bazel-bin/kv_cache_manager/py_connector/vllm/kvcm_vllm_connector-*.whl; do
+            pkg=$(basename "$whl" | sed 's/-{STABLE.*//')
+            ver=$(unzip -p "$whl" "*.dist-info/METADATA" | awk '/^Version:/{print $2; exit}')
+            cp "$whl" "/tmp/kvcm_whl/${pkg}-${ver}-cp312-cp312-manylinux_2_32_x86_64.whl"
+          done
+          "$KVCM_E2E_PYTHON" -m pip install --no-deps --force-reinstall /tmp/kvcm_whl/*.whl || \
+            uv pip install --python "$KVCM_E2E_PYTHON" --no-deps --force-reinstall /tmp/kvcm_whl/*.whl
+
+      - name: run_e2e_tests
+        run: |
+          set -x
+          bazelisk test //integration_test/vllm_e2e/... \
+            --cache_test_results=no --test_output=errors \
+            --test_env=KVCM_E2E_PYTHON="$KVCM_E2E_PYTHON" \
+            --test_env=KVCM_E2E_MODEL="$KVCM_E2E_MODEL" \
+            --per_file_copt='external/jsoncpp_git/.*@-Wno-error'
+
+      - name: upload_logs
+        if: failure()
+        uses: actions/upload-artifact@v6
+        with:
+          name: vllm-e2e-logs-${{ matrix.kind }}
+          path: |
+            /tmp/kvcm_vllm_e2e/**/*.stdout
+            /tmp/kvcm_vllm_e2e/**/*.stderr
+            bazel-out/*-opt/testlogs/integration_test/vllm_e2e/**
+          if-no-files-found: ignore

integration_test/vllm_e2e/BUILD

@@ -0,0 +1,143 @@
+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.
+py_library(
+    name = "e2e_lib",
+    srcs = [
+        "e2e_lib.py",
+        "test_connector.py",
+    ],
+    imports = ["."],
+    tags = ["no-remote-exec"],
+)
+
+py_test(
+    name = "test_basic",
+    srcs = ["test_basic.py"],
+    data = [
+        "//kv_cache_manager:kv_cache_manager_bin",
+    ],
+    imports = ["."],
+    tags = [
+        "no-remote-exec",
+        "gpu",  # requires 1+ GPU
+        "exclusive",  # GPU tests must run serially to avoid CUDA OOM contention
+    ],
+    timeout = "eternal",
+    deps = [":e2e_lib"],
+)
+
+py_test(
+    name = "test_concurrent",
+    srcs = ["test_concurrent.py"],
+    data = [
+        "//kv_cache_manager:kv_cache_manager_bin",
+    ],
+    imports = ["."],
+    tags = [
+        "no-remote-exec",
+        "gpu",  # requires 1+ GPU
+        "exclusive",  # GPU tests must run serially to avoid CUDA OOM contention
+    ],
+    timeout = "eternal",
+    deps = [":e2e_lib"],
+)
+
+py_test(
+    name = "test_tp",
+    srcs = ["test_tp.py"],
+    data = [
+        "//kv_cache_manager:kv_cache_manager_bin",
+    ],
+    imports = ["."],
+    tags = [
+        "no-remote-exec",
+        "gpu",  # requires 1+ GPU
+        "exclusive",  # GPU tests must run serially to avoid CUDA OOM contention
+    ],
+    timeout = "eternal",
+    deps = [":e2e_lib"],
+)
+
+py_test(
+    name = "test_full_hit",
+    srcs = ["test_full_hit.py"],
+    data = [
+        "//kv_cache_manager:kv_cache_manager_bin",
+    ],
+    imports = ["."],
+    tags = [
+        "no-remote-exec",
+        "gpu",  # requires 1+ GPU
+        "exclusive",  # GPU tests must run serially to avoid CUDA OOM contention
+    ],
+    timeout = "eternal",
+    deps = [":e2e_lib"],
+)
+
+py_test(
+    name = "test_partial_hit",
+    srcs = ["test_partial_hit.py"],
+    data = [
+        "//kv_cache_manager:kv_cache_manager_bin",
+    ],
+    imports = ["."],
+    tags = [
+        "no-remote-exec",
+        "gpu",  # requires 1+ GPU
+        "exclusive",  # GPU tests must run serially to avoid CUDA OOM contention
+    ],
+    timeout = "eternal",
+    deps = [":e2e_lib"],
+)
+
+py_test(
+    name = "test_load_failure",
+    srcs = ["test_load_failure.py"],
+    data = [
+        "//kv_cache_manager:kv_cache_manager_bin",
+    ],
+    imports = ["."],
+    tags = [
+        "no-remote-exec",
+        "gpu",  # requires 1+ GPU
+        "exclusive",  # GPU tests must run serially to avoid CUDA OOM contention
+    ],
+    timeout = "eternal",
+    deps = [":e2e_lib"],
+)
+
+py_test(
+    name = "test_multi_turn",
+    srcs = ["test_multi_turn.py"],
+    data = [
+        "//kv_cache_manager:kv_cache_manager_bin",
+    ],
+    imports = ["."],
+    tags = [
+        "no-remote-exec",
+        "gpu",  # requires 1+ GPU
+        "exclusive",  # GPU tests must run serially to avoid CUDA OOM contention
+    ],
+    timeout = "eternal",
+    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",
+    srcs = ["test_mutation.py"],
+    data = [
+        "//kv_cache_manager:kv_cache_manager_bin",
+    ],
+    imports = ["."],
+    tags = [
+        "no-remote-exec",
+        "gpu",  # requires 1+ GPU
+        "exclusive",  # GPU tests must run serially to avoid CUDA OOM contention
+    ],
+    timeout = "eternal",
+    deps = [":e2e_lib"],
+)

integration_test/vllm_e2e/README.md

@@ -0,0 +1,102 @@
+# 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.26.0.
+
+## 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: bit-exact preferred, cosine
+   similarity > 99.99% as fallback.
+
+## 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 |
+
+## Running
+
+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/... \
+  --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:
+
+| Variable | Meaning |
+|---|---|
+| `KVCM_E2E_PYTHON` | Python interpreter with vLLM + both KVCM wheels installed |
+| `KVCM_E2E_MODEL` | Model path; hybrid models are auto-detected from `config.json` |
+
+## 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
+```