support hybrid attention for vllm connector && add integration test by lpdink · Pull Request #1 · lpdink/tair-kvcache

lpdink · 2026-07-28T11:36:52Z

support hybrid attention for vllm connector
add integration test

…ybrid attention Every kv_cache_group (FullAttentionSpec / MambaSpec) becomes an independent transfer unit with its own location spec (tp{rank}_g{group}), block table and data access strategy: token-granular gather/scatter for attention groups, per-block opaque byte copy for mamba/linear state groups. A full-attention model is simply the one-group case. - GroupMeta/TransferGroup abstractions; ReqState tracks per-group block tables - SupportsHMA (request_finished_all_groups) for hybrid memory allocator models - Adapt to vLLM 0.26.0 packed KV layout (num_blocks, heads, block, 2*head_size) - Strided gather/scatter kernel path for padded/strided page layouts - Null-block detection for unmaterialized mamba boundary state

…ention Two-phase save/load verification driven through a real KVCM manager and a real vLLM OpenAI server: phase 1 saves KV to KVCM and captures references straight from vLLM's paged cache; phase 2 loads via the connector and captures again. Captures use only vLLM's own block-table mapping, breaking the save/load symmetry so per-group translation bugs cannot cancel out. Bit-exact compare with cosine > 99.99% fallback. Scenarios: test_basic (TP=1), test_concurrent (4 reqs), test_tp (TP=2 with cross-block manager/vllm block size mapping for full-attention models). The same targets run against full-attention (Qwen2.5) and hybrid (Qwen3.5) models via KVCM_E2E_MODEL; hybrid runs enable prefix caching (mamba align mode) and restart vLLM between phases so loads come from KVCM, not the local cache. Adds a matrixed GitHub workflow (full-attention + hybrid) for self-hosted GPU runners.

…Error Two fixes in the vLLM connector scheduler path: 1. get_num_new_matched_tokens returned the raw external match count without capping it below the prompt length. This connector loads synchronously (load_kv_async=False), so vLLM schedules num_tokens - num_computed_tokens new tokens and asserts that count is > 0 (vllm 0.26.0 v1/core/sched/scheduler.py waiting-queue loop). A prompt whose token count is an exact multiple of the manager block size with all blocks externally cached made the count 0 and crashed the engine. Drop trailing matched blocks until at least one token remains to recompute, mirroring the fallback in vLLM's own SharedStorage/NIXL connectors. 2. handle_canceled_save_req indexed _alive_requests[req_id] directly, but cancellations arrive from http_executor threads and can race request teardown; use .get() with a warning and skip. Covered by kv_cache_manager/py_connector/test/test_scheduler_state.py and the integration_test/vllm_e2e test_full_hit scenario.

…cheduler state Previously zero unit coverage on the connector's pure logic. New Bazel py_tests under py_connector/test (vllm_stubs.py registers lightweight vLLM / pybind stand-ins in sys.modules so v1_connector imports without a GPU or the compiled client): * test_block_translation: _attn_token_indices / _state_block_ids against an independent brute-force reference, parameterized over ratio=1, ratio>1 and manager_bs != group_bs, plus hand-computed examples. * test_data_transfer_results: MultiResult ordering (in-order, out-of-order, concurrent) and the save/load done callbacks' stride-AND merge, which pins the implicit group-major submission-order contract of _submit_group_tasks; includes the hybrid report_failures=False branch. * test_scheduler_state: get_num_new_matched_tokens (incl. the full-hit cap), parse_block_mask_to_save_indices (offset and bool_masks), _parse_groups (full-attn, hybrid, eagle skip, unsupported spec), and the build_connector_meta state machine (new request, cached deltas with new_block_ids None/non-None, preemption resume via both resumed_req_ids and legacy resumed_from_preemption, save-threshold growth, both request_finished paths, canceled-save races). * test/kernel/test_strided_gather_scatter (GPU): the strided kernel path (block_stride/local_block_size incl. padded pages) added in fc89691 had no coverage; checked element-wise against naive torch indexing, plus a roundtrip and a padding-untouched sentinel check. vllm/BUILD: expose vllm_connector to py_connector subpackages for the tests.

Two production bugs surfaced by the new e2e scenarios: 1. Fail-reschedule loop after a KV load failure. With kv_load_failure_policy=recompute, vLLM reschedules the request and calls get_num_new_matched_tokens again; the manager still advertises the blocks whose storage is gone, so the connector re-matched them and the engine looped load-fail-reschedule forever (request hung). A request that already went through an external load attempt (blocks were allocated) now skips external matching on re-query and recomputes locally. 2. Multi-block hybrid saves always failed. vLLM's mamba 'align' mode only materializes the state block ending the matched region (single_type_kv_cache_manager.MambaManager assigns the null block to interior positions), so every interior manager block has a null (id 0) state target by design. save_task treated that as a failure, the stride-AND merge then dropped those manager blocks from the manager's prefix chain, and hybrid caching silently degraded to the final block only. Null state targets are now transferred vacuously (reported success, nothing copied) on both save and load; load_task previously failed the whole task on any null target for the same reason. Covered by test_scheduler_state (retry guard), test_data_transfer_results (vacuous null-state save/load), and e2e: hybrid basic now verifies 4/4 manager blocks bit-exact; test_load_failure exercises the recompute path end to end.

…rage Harness hardening (task B) plus the P2 coverage fix: * Hybrid prompts were ~578 tokens = 1 x 528 manager block, so multi-block state mapping and incremental save never ran under hybrid. make_base_prompts now emits 140 sentences (~2100 tokens > 3 x 528) for hybrid models. * wait_for_captures timeout raises AssertionError instead of warning. * Expected ref/loaded capture counts are computed per prompt from the actual tokenization (len // mbs, shared-prefix for loads) and asserted as exact lower bounds via assert_report_ok(min_matched=...). * bit-exact comparison is the default; cosine fallback only with KVCM_E2E_ALLOW_COSINE=1. * compare_captures iterates the loaded capture's layers (a loaded layer without a reference is a hard error); mamba align-mode null-state layers may legitimately be absent on either side and test_connector skips capturing null (id 0) state blocks, mirroring the connector's vacuous transfer. * VllmServer/ScenarioEnv: connector_name selection (mutation meta-test), log_level, kv_transfer extra-config overrides, kv_load_failure_policy, key_count_per_file, and a ScenarioEnv helper owning manager/vLLM lifecycle for the new custom scenarios; send_completions accepts token-id prompts and extra payload fields; file-backend root_path gets a trailing slash so block files land inside the dir. * test_connector: drop a duplicated capture line; add MutatedConnector (slot -1 off-by-one, test-side injection only) for the mutation meta-test.

* test_mutation (B4): runs the basic scenario with MutatedConnector (slot -1 in _attn_token_indices) and asserts KV verification FAILS -- proof the capture-based harness catches symmetric translation bugs and is not vacuous. * test_full_hit (C2): prompt trimmed to an exact multiple of the manager block size, resent after being fully saved. Regression for the synchronous full-hit crash (vllm 0.26.0 scheduler.py 'assert num_new_tokens > 0'); asserts the engine survives and 0 < matched < prompt tokens. * test_partial_hit (C1): staged A / A+B / A+B+C prompts with prefix caching on; exercises the non-zero-offset incremental manager query and the incremental save extension, asserts a logged query with offset > 0 and verifies the blocks saved through the incremental path. * test_load_failure (C3): deletes the tail half of the per-block storage files between save and load (key_count_per_file=1, block_per_load_task=1, kv_load_failure_policy=recompute). Full-attn: failures reported to vLLM, surviving head blocks verify bit-exact, mismatches confined to deleted blocks. Hybrid: failure swallowed by design (vLLM invalid-block recovery is single-group only), asserts no hang/crash and the failure log. * test_multi_turn (C4): turn 1 decodes past a manager block boundary (ignore_eos + return_token_ids), asserts the manager committed more blocks than the prompt covers; turn 2 embeds turn 1 prompt+output as token ids and must externally match beyond prompt-only coverage with verified KV. All scenarios pass for both Qwen2.5-7B-Instruct (full-attn) and Qwen3.5-4B (hybrid) alongside the original basic/concurrent/tp regressions.

lpdink · 2026-07-28T11:38:33Z

+  contents: read
+on:
+  pull_request:
+    branches: ["main"]


mmm，这个CI没跑过，我也不了解内部CI机器的setting，要看一下。说不定要适配，我们内部CI机器真的有两张A10么？

lpdink · 2026-07-28T11:40:33Z

+          # Hybrid model: MambaSpec groups + FullAttentionSpec group
+          # (mamba_cache_mode="align").
+          - kind: hybrid-attention
+            model_var: VLLM_E2E_MODEL_HYBRID


这个应该要相应的模型在测试机器上存在，应该在向主仓库的PR中说明要怎么配置仓库环境变量。
或者暂时删除本CI，在下一个里面处理。我总觉得这里不太好搞。要看一下主仓库已有的用内部机器的setting。然后要给对应的机器把模型下好，感觉很麻烦啊。

lpdink · 2026-07-28T11:41:19Z

+          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'


这里抑制了warning，但是我们的CI机器的gcc版本不一定很高。也许可以过，不过fine，也可以接受吧。

lpdink · 2026-07-28T11:42:11Z

+            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


机器不一定有uv。总的来说，整个workflow，都是以我的开发机为环境编写的，我在考虑删除这个CI文件。还不成熟；相关CI建设要的外部资源有点多。

lpdink · 2026-07-28T11:45:15Z

+which is what makes the test able to detect symmetric save/load translation
+bugs. See ``test_connector.py`` for the capture-side details.
+
+Full-attention vs hybrid models


总的来说，在我的本地测试中，我们使用了qwen3.5和qwen2.5，这分别测试了hybrid attention和full attention；现在的实现也可以支持MLA，例如glm-4.7-flash，但是没有做这个集成测试。

lpdink · 2026-07-28T11:46:06Z

+    level=logging.INFO,
+    format="%(asctime)s %(name)s %(levelname)s %(message)s")
+
+MODEL_PATH = os.environ.get("KVCM_E2E_MODEL", "/root/ws/resources/models/Qwen2.5-7B-Instruct")


默认的本地路径，这个要避免的。
这个环境变量应该从CI仓库中来，在这里添加注释说明一般需要一个什么模型；
也要在PR的描述里说清楚。

lpdink · 2026-07-28T11:46:29Z

+    format="%(asctime)s %(name)s %(levelname)s %(message)s")
+
+MODEL_PATH = os.environ.get("KVCM_E2E_MODEL", "/root/ws/resources/models/Qwen2.5-7B-Instruct")
+COSINE_THRESHOLD = 0.9999


我多少有点怀疑这个变量真的用到了吗？我们实际上已经实现了bit级一致的。

lpdink · 2026-07-28T11:46:56Z

+COSINE_THRESHOLD = 0.9999
+# Bit-exact comparison is the default (empirically all scenarios achieve it).
+# Cosine fallback must be explicitly requested.
+ALLOW_COSINE = os.environ.get("KVCM_E2E_ALLOW_COSINE", "0") == "1"


如果没用到，他们都可以被删除。cossim的逻辑也可以删除。

lpdink · 2026-07-28T11:47:50Z

+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:


这在读取对应的模型的config.json，fine，我不确定vllm的模型启动是否也需要，还是有聪明的detect机制；不过对我们测试程序来说，这是可接受的。

lpdink · 2026-07-28T11:48:23Z

+    return (
+        "full_attention_interval" in text_cfg
+        or "linear_conv_kernel_dim" in text_cfg
+        or cfg.get("model_type", "").startswith("qwen3_5")


emmm，感觉这种写法不怎么鲁棒；还不如我们直接写死我们用的测试模型呢？

lpdink · 2026-07-28T11:48:53Z

+# Paths / binaries
+# --------------------------------------------------------------------------- #
+def _runfiles_root() -> Optional[str]:
+    return os.environ.get("RUNFILES_DIR") or os.environ.get("TEST_SRCDIR")


我们用了太多环境变量，但是缺少一个地方说明，应该添加到PR的描述里面。

集中起来，统一起来。

lpdink · 2026-07-28T11:49:26Z

+
+
+def find_python() -> str:
+    return os.environ.get("KVCM_E2E_PYTHON", "/root/ws/env/global_vllm/.venv/bin/python")


本地路径。

lpdink · 2026-07-28T11:51:08Z

+        cfg = {
+            "storage_config": {
+                "type": "file",
+                "global_unique_name": "nfs_01",


对于connector的测试来说，its fine，不过，实际上我们这个测试是可以测试到更广泛的情况的，而不仅仅停留在只是一个connector的测试，也就是可以考虑配置其他存储集群的endpoint来做。
不过可以放到后续。

lpdink · 2026-07-28T11:54:19Z

+        cmd = [
+            find_python(), "-m", "vllm.entrypoints.openai.api_server",
+            "--model", MODEL_PATH,
+            "--served-model-name", "qwen",


也许可以换个通用的名字，回头我们要支持SWA的时候怎么说？

lpdink · 2026-07-28T11:57:04Z

+        env["KVCM_E2E_CAPTURE_DIR"] = self.capture_dir
+        # Keep the connector's KV cache layout matching its expected
+        # [2, num_blocks, block_size, num_kv_heads, head_size] shape.
+        env.setdefault("VLLM_KV_CACHE_LAYOUT", "NHD")


vllm 0.23.0的layout发生了变化，也解放了对hybrid attention的支持...
而此前，这也是为什么这个工作非常难做，耗费了大量的时间，最后发现还是重写是最好的。因为此前的实现面向v0.22.1和full attention做了假设，其数据模型对hybrid attention的实现很不友好。
本次支持是基于刚刚发布的v0.26的，测试也是在v0.26上做的，但是考虑到0.23+的协议稳定，预期目前的实现可以支持v0.23+。

这个行为要验证一下，降级到0.23.0，断言我们添加的集成测试仍然通过。

lpdink · 2026-07-28T11:58:24Z

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


上次支持我们安装了ninja...不过也行吧。我估计都能过，还是有过不了的风险？

lpdink · 2026-07-28T12:09:19Z

+                assert ref_t.shape == got_t.shape, (
+                    f"shape mismatch {layer_name}[{si}]: {ref_t.shape} vs {got_t.shape}"
+                )
+                if not torch.equal(ref_t, got_t):


这里的思路就是收集不同tp, token_count的tensor，然后断言完全一致...额，为什么没有把不同tp的保存到同一个tensor文件，奇怪。难道有多个不同的保存时间点，哦，有道理，还真是。因为这里是异步到达的么？

lpdink · 2026-07-28T12:12:39Z

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


woc，原来是这么解决的吗？amazing。。。

lpdink · 2026-07-28T12:13:26Z

+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.


要验证0.23.0的兼容性。

lpdink · 2026-07-28T12:14:23Z

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


下一步就是支持SWA(gemma-3-4b-it)

lpdink · 2026-07-28T12:20:13Z

额，竟然不能提前导入torch，过不了CI，要解一下。
测试是真的很全面，不过收尾工作也不少。
但是这是一个良好的开始，很快我们就能支持SWA，补全connector在这里的所有能力；

lpdink · 2026-07-29T05:10:06Z

Superseded by alibaba#257, opened against the alibaba/tair-kvcache main repository as requested (so CI runs there). All blocking review comments from this PR are addressed in the new one: torch is no longer required to run the new unit tests, the immature e2e workflow is deleted (e2e CI will be a separate PR), the dev-machine local-path defaults are now required environment variables, the unused cosine fallback is removed, and all e2e environment variables are documented in integration_test/vllm_e2e/README.md plus the PR description. Closing this one.

lpdink added 7 commits July 27, 2026 21:36

lpdink commented Jul 28, 2026

View reviewed changes

lpdink closed this Jul 29, 2026



		def find_python() -> str:
		return os.environ.get("KVCM_E2E_PYTHON", "/root/ws/env/global_vllm/.venv/bin/python")

Conversation

lpdink commented Jul 28, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lpdink commented Jul 28, 2026

Uh oh!

lpdink commented Jul 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant