support hybrid attention for vllm connector && add integration tests - #257
support hybrid attention for vllm connector && add integration tests#257lpdink wants to merge 56 commits into
Conversation
…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.
The open-source CI image has no torch/triton/orjson/zmq/requests, so the three pure-logic py_tests failed at import time. vllm_stubs now registers stand-ins for missing third-party deps (real modules always win), the connector uses typing.TYPE_CHECKING instead of typing_extensions, and the GPU-only Triton kernel test is tagged manual like the other GPU suites.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
…ness - delete the immature test-vllm-e2e workflow (runner/model/venv/uv are all dev-machine specific); e2e CI will land as a separate PR - KVCM_E2E_MODEL / KVCM_E2E_PYTHON become required env vars with clear errors instead of local-path defaults - remove the unused cosine-similarity fallback: every scenario is verified bit-exact - tag the e2e targets manual (GPU + vLLM venv + model required) and add an explicit :e2e_tests suite - document all environment variables in the README
f673b14 to
729e95f
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f673b149a7
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # not a failure: report it saved vacuously. Failing it would | ||
| # stride-AND the whole manager block out of the manager's prefix | ||
| # chain and kill multi-block hybrid caching entirely. | ||
| valid = [i for i in range(n) if block_ids[i] != 0] |
There was a problem hiding this comment.
For hybrid/Mamba none or align runs, a manager block can have block_id == 0 here when it was an interior/skipped state block, so this path reports the block saved without writing any bytes for that state group. If a later request only matches through that same manager block, vLLM may allocate a real nonzero final state target for it; load_task will then try to read the URI that was never written, and the hybrid load path reports no invalid block ids, so the request can continue with missing recurrent state instead of recomputing. These null-source blocks should not be advertised as cacheable unless the required state was actually materialized.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 3979907 (connector), c0a1a48 (manager-side validation fix this uncovered), with the reproducing e2e scenario in 0d70939.
You were right, and the failure mode is exactly the one you described. The fix expresses the absence to the manager instead of encoding it as a success, using the location_spec_group_names mechanism the server already implements (FilterWriteCache even documents the intent: "allowing complementary locations (e.g. KV-only and Mamba-only) to coexist in the same block").
Save side. Registration now declares two LocationSpecGroups — full (every group's spec) and attn (attention specs only). build_connector_meta derives each manager block's state completeness from vLLM's block table (in the scheduler loop, since later steps mutate that table) and start_write_cache announces it per key, so the manager allocates and advertises exactly the specs that will hold data. Transfers no longer conflate "nothing to do" with "done": a group abstains for blocks it provably has no data for, and a block's verdict is the AND over the groups that did carry data — a block no group wrote is not published at all. Full-attention models declare no groups and their requests are byte-identical to before.
Load side. getCacheLocation reports the real per-block spec coverage, so get_num_new_matched_tokens truncates an external match to the last state-complete block. A hybrid request can only resume where the recurrent state ends; the attention KV of a longer prefix is unusable without it. This also means the case you found is no longer merely detected — it is never started, which narrows the reachable surface of the "hybrid load failures cannot be reported" limitation to genuine storage failures.
Manager fix. StartWriteCache validated location_spec_group_names.size() against keys.size(), but a token-only request (empty block_keys) derives its keys further down in the same function, so keys.size() was 0 and every such request carrying group names was rejected with size not match, expect[0], real[N]. The proto contract says the list is per block, so the block count is now derived from the tokens when block_keys is empty. sglang never hit this because it always sends explicit block_keys.
Evidence — the new test_cross_request_prefix scenario. It saves request A, measures which of A's blocks really hold a state independently (from the harness's own reference captures, not from the code under test), then runs a strict token prefix B whose match would end on a state-less block: once on a fresh instance id (ground truth) and once against A's cache.
Before the fix, with A's states materialized only at blocks [2,5,6]:
A state materialized per block (from captures): [False, False, True, False, False, True, True]
A per-block coverage: ['full','full','full','full','full','full','full']
B reference output: [7189, 279, 271, 248068, 198, 90700, ...]
B cached output: [7189, 279, 869, 220, 17, 17, ...] <- silently different
After:
A per-block coverage: [missing, missing, full, missing, missing, full, full] <- matches reality
req:... truncated external match from 5 to 3 blocks: later blocks carry no recurrent state
B cached output == B reference output (32 tokens, token for token)
Regression: 9/9 e2e scenarios pass on vLLM 0.26.0 against both Qwen3.5-4B (hybrid) and Qwen2.5-7B-Instruct (full attention); test_basic + test_cross_request_prefix also pass on 0.23.0. Unit tests grew 76 -> 95 cases (the two test_*_vacuously_succeed cases that pinned the old behaviour are replaced by tests for the new dispositions, truncation and per-key group names).
Known cost, disclosed in the PR body: cross-request hybrid prefix reuse is now truncated to the last state-complete block, which costs hits (observed 3 of 7 blocks carry state). Correctness is prioritized over hit rate; back-filling the missing states via complementary writes — which the manager already supports — is a possible follow-up.
| import uuid | ||
| from typing import Optional | ||
|
|
||
| import requests |
There was a problem hiding this comment.
As added, every e2e test imports e2e_lib, and this top-level import runs before any KVCM_E2E_* validation. In the repo test Python, and in Bazel sandboxes unless @pip_cpu//requests is included in the target deps, requests is not present, so the manual e2e targets fail at collection with ModuleNotFoundError before they can launch the vLLM venv specified by KVCM_E2E_PYTHON. Add the pip dependency to the BUILD target or avoid importing it in the Bazel runner.
Useful? React with 👍 / 👎.
| table = load_req.all_block_ids[0] | ||
| gbs = self._group_metas[0].block_size | ||
| report_ids = [table[(mb * self._manager_block_size) // gbs] | ||
| for mb in load_req.manager_block_idxes] |
There was a problem hiding this comment.
When an EAGLE/drafter group is present and skipped by _parse_groups, _num_groups can still be 1 while the only transferred group has group_idx == 1. In that case a single-group load failure is reported from all_block_ids[0], which is the skipped EAGLE table rather than the table used by _plan_group_transfers; vLLM then invalidates the wrong blocks and the failed external load can be retried or treated inconsistently. Use the surviving group's group_idx when choosing the table for report_ids.
Useful? React with 👍 / 👎.
| elif isinstance(spec, FullAttentionSpec): | ||
| # Attention KV is token-granular; scale from the spec's page size | ||
| # to the manager block size. | ||
| per_token_bytes = spec.page_size_bytes // spec.block_size | ||
| metas.append(GroupMeta( | ||
| group_idx=idx, | ||
| is_attention=True, | ||
| layer_names=list(group.layer_names), | ||
| block_size=spec.block_size, | ||
| per_block_bytes=per_token_bytes * self._manager_block_size * len(group.layer_names), | ||
| )) |
There was a problem hiding this comment.
When vLLM supplies a FullAttentionSpec that carries sliding_window or attention_chunk_size metadata, this branch accepts it as ordinary full attention and publishes its blocks to KVCM. For SWA/chunked-local models (called out as unsupported by this change), that lets later requests reuse external blocks under full-prefix semantics even though the model runner only treats a moving window/chunk as valid context, which can produce incorrect KV reuse instead of failing fast. Check those spec fields here and raise until the connector implements the matching/clipping rules.
Useful? React with 👍 / 👎.
…tion PR #257's rewrite only accepted vLLM >= 0.26.0's packed 4-D KV layout (num_blocks, H, block, 2D), regressing full-attention support on older vLLM. Detect all three flash_attn layouts from the tensor shape itself (never from version strings) and normalize each into token-major per-pointer views: * 4-D packed (vllm >= 0.26.0): one pointer per layer, as before. * 5-D N-first (num_blocks, 2, block, H, D) (vllm 0.23.0 - 0.25.x): two pointers per layer, K/V interleaved per block -> kernel strided path. * 5-D KV-first (2, num_blocks, block, H, D) (vllm <= 0.22.1): two flat pointers per layer. The Triton gather/scatter kernel already addresses through a flat pointer array ([K0, V0, K1, V1, ...] for non-MLA), so it needs zero changes; TransferGroup grows num_kv_ptrs (pointer count, = layer_num for packed, 2x for split layouts) and the staging buffer views in data_transfer use it. per_block_bytes is identical across layouts, so the manager, storage and transfer protocol are unaffected. Unrecognized layouts still fail fast at startup. Hybrid (mamba) models on vllm <= 0.22.x are rejected with a clear NotImplementedError: those schedulers assert num_external_computed_tokens == 0 in _mamba_block_aligned_split, so the first external hit would crash mid-flight. The gate probes the installed scheduler for that blocking assert (capability check, not a version comparison). Note the saved byte layout differs between the packed and split-K/V eras, so KV cache is not portable across vLLM upgrades; instance_id isolation already prevents such mixing in practice. New unit tests cover view construction for all three layouts (shape / stride / pointer math on stub tensors, no torch required), fail-fast on unrecognized and ambiguous shapes, and the hybrid gate on old/new/ unprobeable schedulers. The e2e VerifyingConnector reuses attn_kv_views so its captures follow the same normalization.
Multi-version validation report (vLLM 0.22.1 / 0.23.0 / 0.26.0)As promised, here are the cross-version e2e results after commit 9dc5961 (shape-based KV layout detection). Environment: 2x NVIDIA A10, one dedicated venv per vLLM version ( vLLM 0.22.1 — split K/V layout
|
| Scenario | Qwen2.5-7B (full attn) | Qwen3.5-4B (hybrid) |
|---|---|---|
| test_basic | ✅ PASSED, 36/36 | ❌ by design (see below) |
| test_concurrent | ✅ PASSED, 147/147 | ❌ by design |
| test_tp (TP=2) | ✅ PASSED, 72/72 | ❌ by design |
Hybrid on 0.22.1 fails fast, not mid-flight: vLLM <= 0.22.x asserts num_external_computed_tokens == 0 in Scheduler._mamba_block_aligned_split ("External KV connector is not verified yet"), so the first external hit would crash the engine. The connector now probes the installed scheduler for that blocking assert (a capability check — no version-string comparison anywhere in the detection logic) and raises at startup:
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.
The gate behavior is verified against all three real venvs (raises on 0.22.1, passes on 0.23.0/0.26.0) and unit-tested against stubbed old/new/unprobeable schedulers.
vLLM 0.23.0 — split K/V layout (num_blocks, 2, block, H, D)
| Scenario | Qwen2.5-7B (full attn) | Qwen3.5-4B (hybrid) |
|---|---|---|
| test_basic | ✅ PASSED, 36/36 | ✅ PASSED, 4/4 |
| test_concurrent | ✅ PASSED, 147/147 | ✅ PASSED, 16/16 |
| test_tp (TP=2) | ✅ PASSED, 72/72 | ✅ PASSED, 16/16 |
The K/V halves of one kernel page are interleaved in this layout, exercised through the Triton kernel's existing strided path (block_stride spans both halves). 0.24.x/0.25.x share this layout and API surface (checked per source tag), so they are expected-compatible although not e2e-run.
vLLM 0.26.0 — packed layout (num_blocks, H, block, 2D) — full regression, no test cache
All 8 scenarios x both models, --cache_test_results=no:
| Scenario | Qwen2.5-7B (full attn) | Qwen3.5-4B (hybrid) |
|---|---|---|
| test_basic | ✅ PASSED, 36/36 | ✅ PASSED, 4/4 |
| test_concurrent | ✅ PASSED, 147/147 | ✅ PASSED, 16/16 |
| test_tp (TP=2) | ✅ PASSED, 72/72 | ✅ PASSED, 16/16 |
| test_partial_hit | ✅ PASSED | ✅ PASSED |
| test_full_hit | ✅ PASSED | ✅ PASSED |
| test_multi_turn | ✅ PASSED, 57/57 | ✅ PASSED, 6/6 |
| test_load_failure | ✅ PASSED | ✅ PASSED |
| test_mutation (meta) | ✅ PASSED | ✅ PASSED |
Implementation notes
- Layout detection is purely shape-based (
attn_kv_viewsinv1_connector.py): 4-D packed -> one pointer per layer; 5-D KV-first (t[0]/t[1]) and 5-D N-first (t[:,0]/t[:,1]) -> two pointers per layer,[K0, V0, K1, V1, ...]. Ambiguous or unrecognized shapes fail fast at startup. - Zero changes to the Triton kernel — its pointer-array addressing already supports split K/V;
git diffonbatch_gather_scatter_helper.pyis empty for this commit. per_block_bytesis identical across layouts, so the manager, storage backends and transfer protocol are untouched.- KV saved under different layout eras is byte-incompatible;
instance_idisolation prevents cross-era mixing in practice (noted in the PR description). - Unit suite grows to 53 cases (new
test_kv_layouts.py: view construction for the three layouts, fail-fast cases, hybrid gate), all runnable without torch as before.
…tion PR #257's rewrite only accepted vLLM >= 0.26.0's packed 4-D KV layout (num_blocks, H, block, 2D), regressing full-attention support on older vLLM. Detect all three flash_attn layouts from the tensor shape itself (never from version strings) and normalize each into token-major per-pointer views: * 4-D packed (vllm >= 0.26.0): one pointer per layer, as before. * 5-D N-first (num_blocks, 2, block, H, D) (vllm 0.23.0 - 0.25.x): two pointers per layer, K/V interleaved per block -> kernel strided path. * 5-D KV-first (2, num_blocks, block, H, D) (vllm <= 0.22.1): two flat pointers per layer. The Triton gather/scatter kernel already addresses through a flat pointer array ([K0, V0, K1, V1, ...] for non-MLA), so it needs zero changes; TransferGroup grows num_kv_ptrs (pointer count, = layer_num for packed, 2x for split layouts) and the staging buffer views in data_transfer use it. per_block_bytes is identical across layouts, so the manager, storage and transfer protocol are unaffected. Unrecognized layouts still fail fast at startup. Hybrid (mamba) models on vllm <= 0.22.x are rejected with a clear NotImplementedError: those schedulers assert num_external_computed_tokens == 0 in _mamba_block_aligned_split, so the first external hit would crash mid-flight. The gate probes the installed scheduler for that blocking assert (capability check, not a version comparison). Note the saved byte layout differs between the packed and split-K/V eras, so KV cache is not portable across vLLM upgrades; instance_id isolation already prevents such mixing in practice. New unit tests cover view construction for all three layouts (shape / stride / pointer math on stub tensors, no torch required), fail-fast on unrecognized and ambiguous shapes, and the hybrid gate on old/new/ unprobeable schedulers. The e2e VerifyingConnector reuses attn_kv_views so its captures follow the same normalization.
1f92800 to
9dc5961
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9dc5961e63
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| elif isinstance(spec, FullAttentionSpec): | ||
| # Attention KV is token-granular; scale from the spec's page size | ||
| # to the manager block size. | ||
| per_token_bytes = spec.page_size_bytes // spec.block_size |
There was a problem hiding this comment.
When a FullAttentionSpec has page_size_padded/block-strided pages, this uses spec.page_size_bytes, which includes the per-page padding gap, to size the registered location and staging buffer. The transfer group later gathers only compact token data with shape num_kv_ptrs × manager_block_size × per_token_dim, so padded pages make gpu_buffer.view(...).view(...) in save/load require fewer elements than group.per_block_bytes and the transfer fails before any block is written; size this from the unpadded/real page bytes instead of the padded page size.
Useful? React with 👍 / 👎.
| def num_allocated_blocks(self) -> int: | ||
| if not self.block_ids_per_group: | ||
| return 0 | ||
| return min(len(b) for b in self.block_ids_per_group) |
There was a problem hiding this comment.
When vLLM includes a skipped EAGLE/MTP group, that group's block table can intentionally lag or drop the last block while _parse_groups excludes it from transfer. Taking the minimum across every table in block_ids_per_group caps target_save_num by a group KVCM never saves, so fully computed main-model blocks beyond the drafter table are never submitted to start_write_cache; base this count on the transferred group indices instead.
Useful? React with 👍 / 👎.
| @@ -0,0 +1,837 @@ | |||
| """Orchestration for the KVCM <-> vLLM end-to-end KV cache verification test. | |||
There was a problem hiding this comment.
这个文件做的事情太多了，要拆一下。不要都放到一个文件里。
There was a problem hiding this comment.
Done in 9418a10: e2e_lib.py is now a re-exporting facade over two modules with one concern each -- lib_utils.py (pure helpers: detection, paths, tokenization, queries, traffic, capture comparison; no process state) and servers.py (process lifecycle: manager binary, vLLM server, ScenarioEnv). Scenario files keep their imports.
| "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", |
There was a problem hiding this comment.
Kept for now: the manual tags are what keeps these targets out of OSS CI (no GPU runner there). They come off together with the CI-wiring PR that actually provides the runner.
| return ( | ||
| "full_attention_interval" in text_cfg | ||
| or "linear_conv_kernel_dim" in text_cfg | ||
| or cfg.get("model_type", "").startswith("qwen3_5") |
There was a problem hiding this comment.
Agreed, it was a shape-sniffing pile. Rewritten to prefer what the model actually declares: architectures first (Qwen3Next/Zamba/FalconH1/Samba/Jamba markers), falling back to the layer-knob heuristics only when the config declares no architectures (9418a10).
| # --------------------------------------------------------------------------- # | ||
| # vLLM server | ||
| # --------------------------------------------------------------------------- # | ||
| class VllmServer: |
There was a problem hiding this comment.
拆一下，结构清晰一点，后续对vllmServer/kvcmServer均可能有其他setting。utils放到utils下面。
There was a problem hiding this comment.
Covered by the same split (9418a10): utils live in lib_utils.py, the server classes in servers.py with room for further per-server settings.
| "--served-model-name", SERVED_MODEL_NAME, | ||
| "--port", str(self.port), | ||
| "--tensor-parallel-size", str(self.tp_size), | ||
| "--max-model-len", "4096", |
There was a problem hiding this comment.
比如这里就hardcode了，考虑更好的方法来传入vllm庞大数量的参数。
There was a problem hiding this comment.
The serving knobs are no longer hardcoded CLI fragments: they move into an override-able vllm_args dict on VllmServer (max-model-len, gpu-memory-utilization, ...), so a scenario can widen them without editing the harness (9418a10).
| binary = find_manager_binary(repo_root) | ||
| cmd = [ | ||
| binary, | ||
| "--env", f"kvcm.service.rpc_port={self.rpc_port}", |
There was a problem hiding this comment.
The manager side already writes its config to a JSON file (_write_config); the vLLM side is now centralized in the vllm_args dict mentioned above. A scenario-file format (YAML/TOML) is the next step if the scenarios keep growing -- noted as follow-up.
_alive_requests was an unconditional per-request grab bag -- statistics, block-table mirrors, burn flags and the save ledger all lived on one ReqState created for every request the match hook ever saw. It is now three structures with disjoint lifetimes: - _tracked: RequestLedger per request, created at the FIRST allocation (update_state_after_alloc) and dropped at retirement. It holds only what hooks do not re-provide: the accumulated per-group block tables, the save water-mark, and the save-session ledger. The token stream is read from the live vLLM Request (all_token_ids); only its length is tracked. Requests never allocated are never tracked. - _load_failed / _load_attempted: the external-match discipline, as two small sets cleaned at retirement instead of booleans on every request. - The local/remote match statistics are gone (vLLM already knows its computed tokens; nobody consumed them). get_num_new_matched_tokens becomes a pure query per its contract: it never registers state, and the LoadRequest moves to update_state_after_alloc -- the moment both halves of its address (the clamped match from the query cache + the freshly allocated block tables) exist. The query cache drops its 1s TTL for a request-scoped produce/consume lifecycle (produced by the match hook, consumed at allocation, invalidated at retirement), so a request that waits many steps between query and allocation no longer loses its answer. Unit tests rewritten around the ledger and the FakeLocationQueries produce/consume double; 97 cases pass.
…h asserted role slots The role split files are connector_scheduler.py / connector_worker.py (ConnectorScheduler / ConnectorWorker), and the shell owns one slot per role -- connector_scheduler on scheduler-role instances, connector_worker on worker-role instances, the other None -- with every hook asserting the slot it needs (the mooncake pattern), so calling a scheduler hook on a worker instance fails loudly instead of surprising. TairKvCacheConnectorExtraConfig becomes a Pydantic model: unknown extra_config keys are rejected at startup (typos fail fast), defaults document every knob in one place, and model_dump() feeds the manager client factory. request_finished stays as a shim for the pre-SupportsHMA dispatch (upstream plans to deprecate that path -- link in the comment) and will be removed once the minimum supported vLLM dispatches via SupportsHMA only. 97 unit tests pass.
…e split e2e_lib.py (839 lines) carried everything; it is now three modules with one concern each, behind a re-exporting facade so scenario files keep their imports: * lib_utils.py -- pure helpers (model detection, paths, tokenization, manager queries, traffic, capture comparison); no process state; * servers.py -- process lifecycle: the real manager binary, the real vLLM OpenAI server and ScenarioEnv; * e2e_lib.py -- compatibility facade only. Review follow-ups baked in: * the hybrid model detection prefers declared architectures (Qwen3Next/Zamba/FalconH1/Samba/Jamba markers) and falls back to the layer-knob heuristics only when the config declares none; * the vLLM serving knobs (max-model-len, gpu-memory-utilization, ...) move into an override-able vllm_args dict instead of hardcoded CLI fragments -- a scenario can widen them without editing the harness; * the VLLM_KV_CACHE_LAYOUT=NHD comment now answers the question asked in review: it is required, not cosmetic -- the connector asserts token-major pages at register_kv_caches and refuses to start on an HND-default vLLM without it. test_connector.py adapts to the connector's new shape: the worker keeps no request mirror any more, so the VerifyingConnector's scheduler side snapshots the ledger's live token streams onto the envelope (token_snapshots) for the worker-side captures; the mutated meta-test swaps connector_worker instead of _core; captures dispatch on AttentionGroupMeta and the (views, layout) attn_kv_views signature.
The strided path existed only as inline comments inside the kernel body; the module docstring now states why it exists (the N-first split layout interleaves K/V per block and padded allocations leave gaps), gives the offset decomposition explicitly, and clarifies that every pointer is its own base. Also records that this kernel serves the vLLM connector only: the sglang connector moves data through sglang's own hicache path and never touches this module, so extending it for vLLM's layouts burdens no other connector.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 93fcb63414
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
| 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}" |
There was a problem hiding this comment.
During a rolling upgrade where vLLM 0.25 and 0.26 workers share an instance_id, this name remains identical even though attn_kv_views serializes split-K/V and packed pages in different byte orders. The registration metadata contains model/dtype/size but no KVLayout, so the manager accepts both workers and a new worker can load an old worker's bytes into the packed layout, silently corrupting inference. Include the layout/format version in the registered namespace or reject mixed layouts for one instance.
AGENTS.md reference: AGENTS.md:L40-L42
Useful? React with 👍 / 👎.
| if self._async_get_cache_location: | ||
| self._query_async(request, computed_blocks) | ||
| return None |
There was a problem hiding this comment.
With the default asynchronous query mode, vLLM calls the match hook again every scheduler step while the manager request is pending. Because this branch invokes _query_async even when entry already exists with in_flight=True, every retry submits another identical HTTP request; the four-worker executor accumulates stale queries, amplifies manager load, and delays save-location work sharing the same executor. Only submit when this call creates or replaces the entry, and otherwise return None for the existing request.
Useful? React with 👍 / 👎.
| first_attn = next(kv_caches[name] | ||
| for meta in self._group_metas | ||
| if isinstance(meta, AttentionGroupMeta) | ||
| for name in meta.layer_names) |
There was a problem hiding this comment.
When vLLM supplies a pure Mamba/state-space model with one or more MambaSpec groups but no FullAttentionSpec, parse_groups accepts the configuration and the connector advertises HMA support, but this next(...) has no attention tensor and raises StopIteration during worker initialization. Derive the device from a state cache as well, or explicitly reject state-only configurations during parsing with an actionable error.
Useful? React with 👍 / 👎.
…ferences Four fixes surfaced by running the e2e harness against the local tree: - _version_info is stamped into the wheel at build time; a source checkout has none, so the connector falls back to dev markers instead of failing to import (running from a checkout is a supported mode -- it is how the e2e harness exercises the code under review). - data_transfer still passed the deleted kv_stride field to the gather/ scatter kernels (AttributeError at save time); both call sites drop it. - the e2e harness resolved worker-role attributes (_data_transfer, _tp_rank, ...) on the shell, where the role split moved them to ConnectorWorker; properties forward them now. _capture_pending_loaded also referenced its caller's meta variable instead of its parameter. - the harness passes the repository root ahead of any installed wheel in the vLLM process's PYTHONPATH, so connector modules resolve to the checkout under test; run_e2e moves to servers.py (it orchestrates ScenarioEnv) and servers.py imports the utils it consumes. test_basic passes bit-exact (36/36 blocks) against Qwen2.5-7B-Instruct on vLLM 0.26.0 with the real manager binary.
During decode vLLM's all_token_ids lags the scheduled-token count by one: the token scheduled in a step is appended to the request only once sampled. The ledger kept token_len (the scheduled count) and the save water-mark derived its block count from it while slicing the hash material from all_token_ids -- so every decode-time save announced one block more than the token list could key, the manager returned one location fewer than announced, and the worker's strict alignment check dropped the whole session. Decode-produced blocks never registered; prefill has no lag, which is why only the first block survived. The pre-mirror code was safe by accident: its request mirror appended the same lagging slice it counted from. The fix restores the invariant explicitly: the block count derives from len(all_token_ids) -- the key material itself -- keeping the allocated-KV cap for the chunked-prefill case where all_token_ids runs ahead of the computed KV. A one-block shortfall self-heals on the next step. Regression test pins the lag scenario (token unknown -> no announcement; token lands -> 48 tokens announced). Found by the multi_turn e2e (bit-exact 4/4 after the fix); the unit fakes kept the two counters in sync, which is exactly why they missed it.
…tion swap Two fixes from the hybrid e2e run: - is_hybrid_model short-circuited on a non-empty architectures list: Qwen3.5-4B declares its own architecture (Qwen3_5ForConditionalGeneration) which matches no known marker, so the model was classified full-attention. With that, prefix caching stayed off and the connector ran in the explicitly unsupported mamba_cache_mode=none regime -- every loaded state diverged bit-exactly while attention KV verified. The signals are ORed now; a non-matching architecture string never shadows the layer knobs. - MutatedConnector swapped connector_worker on scheduler-role instances too, where the slot is None (AttributeError at startup). The swap is guarded to worker-role instances only. All nine scenarios pass bit-exact against Qwen3.5-4B (hybrid) on vLLM 0.26.0 with the real manager binary, including the mutation meta-test (the injected off-by-one is caught as expected).
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b8ddf632a0
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
| from typing import Any, Dict | ||
|
|
||
| from pydantic import BaseModel, Field |
There was a problem hiding this comment.
All four new connector unit tests import v1_connector, which imports this module during collection, but kv_cache_manager/py_connector/test/BUILD depends only on :vllm_stubs and the @pip_cpu requirements contain no Pydantic package. Consequently, the advertised Bazel tests fail at collection with ModuleNotFoundError: No module named 'pydantic' in the repository's dependency-only test environment. Add Pydantic to the open-source lockfile and Bazel deps, or extend the dependency-free stubs to cover it.
Useful? React with 👍 / 👎.
| if getattr(model_config, "use_mla", False): | ||
| raise NotImplementedError("MLA models are not supported by TairKvCacheConnector") |
There was a problem hiding this comment.
When model_config.use_mla is true, this unconditional rejection prevents the connector from starting for every MLA deployment. The previous implementation explicitly supported that case by setting _use_mla, registering the one-component per-rank layout, and propagating use_mla in the deployment metadata, while this change's stated supported scope still says MLA follows the attention inheritance path. Preserve the existing MLA path (including the correct registration metadata), or stop advertising it as supported.
Useful? React with 👍 / 👎.
The scheduler re-asks get_num_new_matched_tokens every step while an async location query is on the wire, and get_locations_for_query only short-circuited *completed* cache entries: an in-flight entry fell through to _query_async unconditionally, so every re-ask submitted another full getCacheLocation. With ~12ms query latency that is one duplicate HTTP query (and one duplicate manager-side per-block lookup of the whole request prefix) per engine step until the answer lands -- measured on the vLLM 0.22.1 e2e perf harness: 24 requests produced 163 getCacheLocation calls and 23,472 manager block queries, vs 24 calls / 3,456 blocks for the pre-rewrite connector at the same commit base (the old manager-side query cache's RUNNING entry absorbed re-asks). Client-side effect: +7.7% TTFT p50 from http-pool contention. Fix: while an entry is in flight, return None (wait) regardless of the asked offset; a completed entry with a stale offset still re-issues exactly once. Re-measured: 24 calls / 3,456 blocks -- exact parity with the pre-rewrite connector -- and TTFT p50 back to baseline. The full e2e suite stays bit-exact (9/9 scenarios, both model kinds). New test_location_query.py pins the fan-out contract with a gated fake manager client: in-flight re-asks issue exactly one fetch, a stale offset re-issues once after the previous answer lands, sync mode answers inline, and a failed query drops the entry so the next ask retries instead of waiting forever.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 05714af000
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| 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) |
There was a problem hiding this comment.
When the documented Bazel :e2e_tests command runs, this import executes in the Bazel test interpreter before the separately configured KVCM_E2E_PYTHON child process is launched. However, integration_test/vllm_e2e/BUILD declares only @pip_cpu//requests, and the pip_cpu lockfile contains no Torch package, so the manual suite fails immediately with ModuleNotFoundError: No module named 'torch' in a dependency-only environment. Add Torch to the test runtime or perform hashing/comparison without importing it in the driver.
Useful? React with 👍 / 👎.
| # Re-issued with a different offset meanwhile. | ||
| 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) | ||
| entry.locations = locations | ||
| entry.in_flight = False |
There was a problem hiding this comment.
When a canceled request's request_id is reused while its old manager query is still running, invalidate() removes the old entry and the new request installs another in-flight entry under the same ID; this callback then finds that new entry and stores the old prompt's locations into it because it verifies neither entry identity nor the queried offset. The new request can consequently load KV belonging to the canceled prompt, while its own later callback is discarded because the entry has already been marked complete. Capture the entry or a generation token when submitting the query and only update that exact entry.
Useful? React with 👍 / 👎.
The lifecycle rewrite collapsed the query cache key to req_id, which changed reuse semantics vs origin/main in a non-obvious way: a re-ask at a grown offset (chunked prefill, re-ask after partial compute) is a *different* query, but the req_id key either served it the older offset's answer or made it wait behind the older offset's in-flight query. Restore origin/main's key identity (req_id, query_type, token_length, computed_blocks) under the lifecycle cache: entries for several offsets of one request coexist, an in-flight query never blocks a different offset, and consume still pops exactly the answer the match hook last returned (tracked in _last_answered).
Per-task torch.empty staging made the connector's VRAM footprint unbounded and racing the engine for its own allocation: at high load on low-headroom GPUs (doc085 @ gpu-mem-util 0.92) the 98 MiB per load task allocations OOMed the engine where origin/main -- whose gather kernel writes a pre-allocated pinned host pool directly -- ran fine. Replace the dynamic allocs with one _StagingPool per transfer group (save and load share it): a fixed, configurable HBM + pinned reservation (staging_pool_blocks, default 128, validated >= the largest task batch) handed out as contiguous runs, with blocking acquire as backpressure when exhausted. Bulk D2H/H2D and kernel views are unchanged -- only the buffer source differs. Exception paths drain the stream before the slots go back so a failed task cannot leave enqueued work against a reused view. This also restores origin/main's no-dynamic-allocation invariant the revert of _PinnedBudget (7bd0413) dropped, without reintroducing the connector-level lifetime governance that #280's deadline chain owns.
…path The pool fix (a0d6289) bounded the connector's staging but kept a device-side mirror per transfer group: staging_pool_blocks x per_block_bytes of permanently reserved HBM (112 MiB/group at the default on Qwen2.5-7B TP1, once per hybrid group) that competes with the engine's KV cache and batch headroom -- exactly the reviewer concern on low-end cards. The mirror only existed to route transfers through a bulk D2H/H2D copy, and that copy is not needed: the strided gather/scatter kernel addresses host pinned memory directly over PCIe (UVA zero-copy, as in origin/main), and state-group copies are plain copy_ between the GPU tensors and the pinned slices. Make _StagingPool pinned-host-only and delete the GPU buffer, the gpu_view, and both bulk copies. Save gathers HBM -> pinned slot and hands the slot to SaveKvCaches; load hands the slot to LoadKvCaches and scatters pinned -> HBM. Backpressure, contiguous-run allocation, exception-path stream drains and the capacity guard are unchanged; the kernel and the scheduler/worker instruction flow are untouched. staging_pool_blocks now sizes host RAM only. Trade-off: a slot is reused once its task reports, so after an SDK timeout/error a background DMA may still scribble on a reused slot -- the same exposure origin/main's CopyBufferAllocator always had. The proper fix is #280's deadline contract; taking the exposure now is what buys zero VRAM without delaying this PR.
wangxiyu191
left a comment
wangxiyu191
left a comment
There was a problem hiding this comment.
我基于当前 head 05714af0 完成了这一轮 review。整体上，这次重构把 Scheduler / Worker / Common 的职责拆开后结构清晰了很多；per-block spec coverage 正确修复了空 Mamba state 被分配、被当作写成功并发布的问题，默认异步 query 对同一 in-flight 请求的重复提交在当前 head 也已修复。
建议合入前重点处理下面的 inline comments，尤其是：load failure recovery 的能力判断应基于 block-table 形状，而不是 is_hybrid；以及恢复以 bytes 为维度的 staging capacity/backpressure，并评估额外 GPU bounce buffer 的必要性。
现有 thread 中仍需跟进但这里不重复评论的有：异步 query callback 需要绑定 entry/generation；纯 Mamba 配置应在初始化前给出明确拒绝；MLA 是恢复支持还是更新支持矩阵需要形成显式决策；Pydantic 与 e2e Torch 的 Bazel 依赖闭包需要补齐。
几个跨文件结论：
- Hybrid 模型真实 I/O load failure 的安全 recovery 受限于上游 vLLM 的 multi-group 能力；本 PR 暂不实现可以接受，但必须作为 Known Limitation 明确提示，并保留足够醒目的 error/metric，因为当前后果可能是基于未正确加载的 state 继续推理。
- 删除 worker-side request mirror 的性能验证目标是不回退：在显存足以让 L1 全命中、外部加载为 0 的条件下，对比 Connector off 与 KVCM + TairMempool，关注吞吐、尾延迟和 GPU bubble；不要求证明删除 mirror 一定带来收益。
- KV layout / serialization version 暂可依赖发布流程避免混用，不作为 blocker；仍建议把稳定的 format generation（例如
kv_serialization=vllm-hma-v1）写入model_deployment.extra，只在不兼容格式变化时递增，使 RegisterInstance 能直接拒绝新旧 instance 混用。 - 当前 PR 与最新
main冲突，请在合入前 rebase；同时更新 PR 的支持矩阵和实际测试状态，避免实现与描述不一致。
| # output. Remove report_failures gating once upstream supports | ||
| # multi-group invalid-block recovery. | ||
| report_ids = [] | ||
| if not self.is_hybrid: |
There was a problem hiding this comment.
这里判断的应该是上游 recovery 能处理的 block-table 形状，而不是模型是否 hybrid。vLLM 0.26 的 _update_requests_with_invalid_blocks 会把 get_block_ids(req_id) 解包成恰好一个 table；纯 Attention 也可能有多个 table（多个 Attention group，或仍保留被 connector 跳过的 EAGLE/MTP drafter group）。这时 not self.is_hybrid 仍为 true，失败 block 被上报后，上游可能在 tuple unpack 处报错，而不是进入 recompute。
建议用 len(load_req.all_block_ids) == 1 决定 report_failures；进入该分支后再断言唯一传输组为 Attention，并按它的 group_idx 取 report_ids。请覆盖单 Attention、多 Attention group、skipped drafter、Attention + Mamba 四种 table 形状。
| 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" | ||
| cpu_buffer = torch.empty(len(valid) * group.per_block_bytes, dtype=torch.uint8, |
There was a problem hiding this comment.
旧实现的 CopyBufferAllocator 同时承担了容量 backpressure；重构后每个并发 task 都会动态申请 len(valid) * per_block_bytes 的 pinned CPU buffer，save 还会申请同尺寸 GPU scratch，load 的 .to(device) 也会产生 GPU buffer。结合 32 个 I/O worker 和默认 128 blocks/task，staging 峰值可能达到 GiB 级。SDK deadline 只约束 buffer 生命周期，不能限制同时 in-flight 的总字节数。
建议恢复以 bytes 为维度的有界、预分配且可复用的 pinned CPU pool，并对超大 task 做切分；同时评估完整 GPU bounce buffer 是否必要，若不能消除也应做成有界、可复用的 staging。
|
|
||
| if num_external_tokens <= 0: | ||
| return | ||
| locations, computed_blocks = self._location_query_manager.consume_locations(req_id) |
There was a problem hiding this comment.
consume_locations() 的签名是 Optional[Tuple[...]]，返回 None 时这里会先在 tuple unpack 处抛 TypeError，所以下面的 locations is None ... 检查实际上到不了。请先保存并检查 result，再解包。
另外，此时 vLLM 已经按正的 external hit 分配了 blocks；如果缓存结果真的缺失，静默 return 也不是安全降级。建议将它作为 invariant violation fail closed（带 request id 和匹配/分配上下文），或者走明确受支持的 invalid-block recovery。
| for vllm_req in scheduler_output.scheduled_new_reqs: | ||
| ledger = self._tracked.get(vllm_req.req_id) | ||
| if ledger is None: | ||
| continue # never allocated (defensive): nothing to record |
There was a problem hiding this comment.
请确认“scheduled request 找不到 ledger”是否存在正常生命周期路径。按 vLLM 0.26 的常规 hook 顺序，update_state_after_alloc() 会在 scheduled_new_reqs 产生前创建 ledger，我暂时没有找到合法路径；下面 cached request 的同类 continue 也一样。
如果 cancellation / preemption / resume 或某个受支持版本确实可能触发，请把原因写清楚；否则可以保留防御性 continue，但至少记录 warning/error，并带上 request id、new/cached/resumed 状态和 scheduled token 数，避免生命周期契约被破坏后静默丢状态。
|
|
||
| if ledger.scheduled_saving_count == ledger.sent_saving_count: | ||
| self._retire_request(req_id) | ||
| return True, None |
There was a problem hiding this comment.
旧实现会在这里返回 local_matched_token_num / remote_matched_token_num，vLLM 会把第二个返回值透传到 EngineCoreOutput.kv_transfer_params。统一返回 None 后，按请求关联 L1/local 与 KVCM/remote 命中、TTFT 和请求属性的能力消失了，aggregate counter 无法还原这层信息。
建议把两个计数保留在 RequestLedger 并继续返回原有结构；如果确实要移除，请把它作为显式兼容性决策，并提供等价的 per-request 替代观测。
There was a problem hiding this comment.
已修复：5f033efe。
调查结论：删除发生在 ledger 重写（cbe1bc4d）时——两个计数和 request mirror 一起被清，属于把被消费的观测契约当成死代码顺手删了，并非 hybrid/group 模型的技术障碍（两个计数都是 group 无关的 token 粒度总量，与 per-block spec coverage 正交）。
已确认 vLLM 侧透传链三个版本（0.22.1 / 0.23.0 / 0.26.0）一致：request_finished 的 extra_info → EngineCoreOutput.kv_transfer_params → RequestOutput.kv_transfer_params → OpenAI CompletionResponse.kv_transfer_params（协议字段，还会经 extra_args 回传用于 PD 预热），确实是被消费的用户可见行为。
实现照抄旧语义：
- match hook 每条应答路径记录（local = 询问时的
num_computed_tokens，remote = clamp 后的应答）到RequestLedger； - burned-match 重查报告 local 照旧、remote = 0（match 已花掉）；
_finish_request两个分支（立即 retire / 延迟 finish）都返回该 dict。
关于您提到的断言缺失：完全同意，此前删了 4 个 commit 没有任何测试报警。新增 4 个契约测试 pin 住它：两个 finish 分支必须返回 dict（含类型断言）、re-ask 以最后一次应答覆盖、burned 路径 remote 归零而 local 保留。
| # Per kv_cache_group block table, in each group's own block_size units. | ||
| block_ids_per_group: List[List[int]] | ||
| # Tokens accounted for by the ledger (prompt + scheduled decode steps). | ||
| token_len: int |
There was a problem hiding this comment.
小建议：token_len 目前看起来已经成为 dead state——生产代码会初始化/递增它，但 save watermark 已改为读取 len(vllm_request.all_token_ids)，没有实际消费者。若没有计划用它做 telemetry/invariant，建议删除字段、更新逻辑和只验证该字段的测试；若要保留，请改成能表达真实语义的名字（例如 num_scheduled_tokens），并说明它相对 all_token_ids 的一 token lag。
|
|
||
| locations = response["locations"] | ||
| write_session_id = response["write_session_id"] | ||
| logger.info("req:%s save session %s: block_mask=%s locations=%d", |
There was a problem hiding this comment.
小问题：block_mask 可能包含覆盖整个目标 prefix 的 bool_masks.values。长请求或非连续写入时，每次 incremental save 都在 INFO 打完整向量，日志量会很大，累计开销也可能随 prefix 增长。建议降到 DEBUG，INFO 只记录 mask 类型、offset、total/existing/missing counts。
|
|
||
| # Spec group names advertised at registration and used per key in | ||
| # start_write_cache. See build_spec_groups for the semantics. | ||
| ATTN_SPEC_GROUP = "attn" |
There was a problem hiding this comment.
这里的命名容易反转读者的直觉：attn 实际是基础 Attention KV，而 full 是 Attention KV + 所有 recurrent state 的并集；KVCM 现有术语更接近原子组 Full（普通 Attention KV）/ Linear（线性 Attention / recurrent state）。建议统一为内部 FullAttentionGroupMeta / MambaStateGroupMeta，KVCM 原子 spec group 沿用 Full / Linear，并集用 All（比 Complete 更直接）。至少应避免再形成一套 attn/full 含义相反的协议词汇。
…hybridness
vLLM's invalid-block recovery unpacks a single-group block table
('(req_block_ids,) = get_block_ids(req_id)' in
_update_requests_with_invalid_blocks), so report_failures must key on
the table count, not on whether a state group exists: a multi-group
attention-only model (unmerged sw+full, or a skipped EAGLE/MTP drafter
group) is not hybrid yet still breaks the unpack and crashes the
scheduler.
Worker: report_failures = len(load_req.all_block_ids) == 1, with the
single-table branch asserting the transferred group is attention group 0
(a lone state group would feed state block ids into upstream's
token-granular recovery math).
Scheduler: _external_match_burned applies the one-shot burn by the same
shape, read from the block-table snapshot recorded at allocation time
(mirrors kv_cache_manager.get_block_ids; the group count is static).
Tests: worker-side table shapes (single attention / multi attention /
skipped drafter / hybrid / lone state group) and scheduler-side burn
shapes (skipped drafter and multi-attention now burn like hybrid).
…lots earlier The pool fix bounded staging at 128 blocks (one full task): with the synchronous SDK transfer holding slots, that serializes the SDK feed to a single task and queues burst loads behind saves -- measured -3.5% ab throughput and +54% vs +65% TP2 hit throughput against 512 blocks. The GPU mirror is gone (zero-VRAM data path), so the pool costs pinned host RAM only: ~896 MiB per attention group at 1024, which is origin/main's validated concurrency (8 full tasks in flight). Host RAM is cheap; restore 1024 as the default and document the sizing rationale in the config comment. Also wait for the engine's forward event *before* acquiring a slot: the slot is only needed for gather + SDK transfer, so holding it during the forward pass only extends pool occupancy.
…ne new_block_ids Two hot-path fixes from the static-check pass (ty/ruff): 1. update_state_after_alloc unpacked consume_locations() into a tuple whose Optional return was never checked -- a None would raise TypeError from the unpack itself, and the dead 'is None' branch below would silently let vLLM run the request on KV that was never loaded. vLLM only allocates for an external hit after the match hook answered positively, so None here is a contract violation: fail closed with request context instead of corrupting the request's output without a traceback. 2. _ingest_scheduled_reqs replaced a resumed request's block table unconditionally ([list(b) for b in new_block_ids]), but upstream CachedRequestData.new_block_ids may be None (get_block_ids with allow_none=True when no group got fresh blocks, PR #23262): the replace would crash on None. Skip None entries -- nothing was allocated to record -- and let the incremental branch narrow the type for the zip. Also document the MultiResult flatten invariant with an explicit assert (ty cannot see the count check guarantees every slot is filled).
…k, identity-guarded writes The four-tuple cache key (67e1a98) restored origin/main's key identity, but origin/main's key serves a TTL cache that is *read* by its single monolithic consumer; the split connector *pops* answers across two hooks (match has the offset, alloc does not), and multi-entry coexistence then needs a routing patch (_last_answered) whose correctness silently depends on match/alloc adjacency. vLLM's contract makes coexistence pointless: the alloc hook always consumes the answer of the *last* match hook -- it passes the matched count back verbatim -- so only the newest ask's answer can ever be consumed. Replace the per-key entry dict and the _last_answered router with 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 and the old ask's answer -- arrived or in flight -- can no longer be consumed. This also removes the 05714af-era serialization where a grown offset waited behind the older offset's in-flight query; * late answers never beat newer asks: each ask captures the slot object it created, and the async callback writes only if the slot is still that object (identity is the version; a monotonic counter is kept for logs only). Superseded/consumed/invalidated slots discard the answer. store_result gets simpler too: the slot is necessarily the query the hook just answered (the scheduler thread is the only superseder), so the clamp writes the slot directly -- the mis-addressed-write failure mode disappears with the router. Tests: rewrite test_location_query.py around the supersession contract (same-offset dedupe, immediate RPC for a grown offset, superseded slots not servable, late answers dropped in both orders, consumed/invalidated slots discarding in-flight answers).
…n_len, tame mask logs Three review follow-ups in the scheduler-role connector: * scheduled_new_reqs / cached_reqs with no ledger were silently skipped. update_state_after_alloc creates the ledger before vLLM can schedule anything, so this only fires on a broken hook-order contract: log it loudly (with resumed state and scheduled tokens) instead of dropping the block table invisibly. * RequestLedger.token_len was dead state: initialized and incremented but never read (save sizing counts blocks from all_token_ids, whose one-token decode lag was the reason token_len-derived counts were dropped). Remove the field, its maintenance, and its tests. * start_write_cache logged the full block_mask vector at INFO on every incremental save. Summarize (mask kind, offset, total/existing) at INFO and keep the full vector at DEBUG.
…p constants Two review follow-ups in group parsing: * a mamba-only model passed parse_groups (its MambaSpec groups are legal) and then died in register_kv_caches with an obscure StopIteration on the 'first attention tensor' lookup. Refuse it explicitly before init: pure-mamba / attention-free models are not supported, full-attention and hybrid (attention + mamba) only. The empty-group case (all groups skipped) also raises instead of asserting. * ATTN_SPEC_GROUP / FULL_SPEC_GROUP read backwards: 'attn' is attention-only coverage and 'full' is the union of attention KV plus every recurrent state group. Rename the Python constants to ATTN_ONLY_SPEC_GROUP / ALL_SPEC_GROUP. 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 -- so only the code-side names move.
…2e venv model The connector imports pydantic (config) and orjson (tp_coordinator) at runtime, but the bazel dependency closure never declared them: the vllm_connector py_library listed only in-repo deps, and the py312 unit tests only passed because their runtime python (a venv) supplied the imports. Close the gap: * add pydantic==2.13.4 and orjson==3.12.0 to the pip_cpu requirements base and regenerate the lock (cp312 hashes); * declare @pip_cpu//pydantic and @pip_cpu//orjson in vllm_connector, and the wheel's PyPI requires so pip resolves them on install; * document in integration_test/vllm_e2e/BUILD that vLLM/torch/triton are venv-provided by design (manual + GPU-tagged targets, connector under test runs from source via PYTHONPATH).
Review 意见处理汇报（2026-08-25）谢谢 @wangxiyu191 的详细 review。基于 已修复（commit 对应）
顺带发现并修复
未决事项
单测：py_connector 全套 5/5 绿（test_scheduler_state / test_location_query / test_data_transfer_results / test_kv_layouts / test_block_translation）。 |
…params The ledger rewrite (cbe1bc4) dropped ReqState's local/remote matched token counts together with the mirror, and _finish_request started returning None instead of the accounting dict -- silently removing an observable contract: vLLM threads request_finished's extra_info into EngineCoreOutput.kv_transfer_params -> RequestOutput -> the OpenAI response, where consumers attribute per-request TTFT and hits to the local (vLLM prefix cache) vs remote (KVCM) sources. Found only by human review four commits later; nothing in the unit suite pinned it. Restore it on the ledger, group-agnostic by design (both counters are token-granular totals, orthogonal to per-block spec coverage): * the match hook records (local = the num_computed_tokens it was asked at, remote = the clamped answer) on every answer path; * the burned-match re-ask reports local as asked and remote = 0 (the match is spent), matching the pre-rewrite semantics; * _finish_request returns the dict from both branches (immediate retire and delayed-finish). New tests pin the contract: both finish branches must return the dict, re-asks overwrite with the last answer, and the burned path zeroes the remote half while keeping the local hit.
Resolve the one conflict in cache_manager_test.cc by keeping both sides' additions: our TestStartWriteCacheSpecGroupNamesWithTokenIdsOnly (token-only writes with per-block spec-group names, hybrid regression) and main's TestStartWriteCacheRecordWriteBytes (write-bytes metric accounting) are independent tests registered against different instances.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d265c8cc62
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| for name in ("torch", "triton", "triton.language", "zmq", "requests"): | ||
| if name not in sys.modules and not _importable(name): | ||
| sys.modules[name] = MagicMock(__name__=name) |
There was a problem hiding this comment.
When requests is absent—as intended for the dependency-light unit-test environment—this installs a MagicMock as the module, so requests.HTTPError and requests.RequestException are mock objects rather than exception classes. Importing manager_client.py then fails while defining KvCacheManagerHTTPError(requests.HTTPError, AssertionError) with a metaclass conflict, and every new connector test aborts during collection. Provide a minimal requests stub with real Exception subclasses or declare the real requests dependency for these tests.
Useful? React with 👍 / 👎.
| self._cpu = torch.empty(total, dtype=torch.uint8, device="cpu", | ||
| pin_memory=(device.type == "cuda")) |
There was a problem hiding this comment.
When the connector runs on the explicitly supported MUSA path selected by _get_device_module, this condition allocates ordinary pageable CPU memory even though the new attention transfer kernels directly dereference the staging pointer from the accelerator via UVA. Unlike CUDA-pinned—or MUSA-pinned—host storage, pageable memory is not device-addressable, so attention save/load tasks fail or fault as soon as the gather/scatter kernel uses the pool. Allocate accelerator-pinned host memory for MUSA as well, while retaining pageable memory only for CPU/test devices.
Useful? React with 👍 / 👎.
| # 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): |
There was a problem hiding this comment.
During normal asynchronous accelerator execution, ready_event.wait() without an explicit stream inserts the dependency into the IO executor thread's current stream, but the gather/state copies are subsequently enqueued on the separate _save_stream. The save stream can therefore read paged KV/state before the model's forward writes represented by ready_event have completed, publishing stale or partially written cache bytes. Insert the wait while _save_stream is current, or host-synchronize the event before gathering.
Useful? React with 👍 / 👎.
…ock count The 1024-block default (a973a01) was derived for full-attention blocks (~0.875 MiB each: 1024 blocks ~= 896 MiB of pinned host RAM per group). Hybrid blocks are ~17.3 MiB, so the same count pins ~17.3 GiB per group; four groups then die in the pinned allocator (cudaHostAlloc OOM) at engine start -- found by the final-validation e2e smoke on Qwen3.5-4B. Derive the per-group block count from a byte ceiling (staging_pool_max_bytes_per_group, default 1 GiB): effective = clamp(staging_pool_blocks, byte_cap // block_bytes, floor = one full task batch) Full-attention sizing is unchanged (1024 blocks, the validated 8-task concurrency); hybrid falls back to the batch size (128 blocks), the configuration the staging-removal campaign validated. One task batch always fits contiguously; the cap can only shrink, never grow.
池容量修复与 merge main（最终更新）
当前 PR 状态：wangxiyu191 的 8 条 inline 意见全部落地（#1 block-table 形状 |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0db89c9a33
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| env["PYTHONPATH"] = os.pathsep.join([ | ||
| os.path.dirname(os.path.abspath(__file__)), | ||
| find_repo_root(), | ||
| env.get("PYTHONPATH", ""), |
There was a problem hiding this comment.
In the documented e2e setup, adding the checkout root ahead of site-packages causes Python to load the source tree's regular kv_cache_manager package, whose client/pybind directory contains no kvcm_py_client extension; the compiled module exists only in the installed wheel/Bazel output. Consequently, importing connector_worker.py fails before vLLM starts, so the e2e targets cannot exercise the source connector as intended. Stage the extension alongside the source package or extend the package path instead of shadowing the installed package wholesale.
Useful? React with 👍 / 👎.
…ed fix) The merge brought in main's manager_client exceptions, which subclass requests' (KvCacheManagerHTTPError(requests.HTTPError, AssertionError)). In the open-source CI -- no requests installed -- vllm_stubs replaced requests with a MagicMock module, and a MagicMock attribute cannot be a base class: class creation died with 'metaclass conflict' before a single test ran. Give the requests stand-in real exception classes mirroring requests' own hierarchy (RequestException(IOError) -> HTTPError/ConnectionError/ Timeout); Session/post stay MagicMocks because the tests patch them. Record stand-ins in vllm_stubs.STUBBED so tests can skip on real- behaviour requirements: test_views_slice_the_same_run needs an actual torch tensor and now skips when torch is stubbed. Also import vllm_stubs before the bare 'import torch' in test_data_transfer_results (the stub must be in sys.modules first in the CI environment), and decouple the zero-VRAM allocation guard from torch.device -- a SimpleNamespace(type='cuda') drives the same contract check wherever torch itself is a stand-in. Verified in both environments: the no-deps simulation (meta-path blocker over the five CI targets: 132 tests, 0 failures) and the full bazel suite on the dev venv (5/5).
3 participants
Summary
Rewrite the vLLM connector around per-group KV cache transfer, adding support for hybrid attention models (mamba/linear state layers + full attention, e.g. Qwen3.5) on vLLM 0.23.0+, together with a unit-test suite and a two-phase bit-exact end-to-end verification harness.
Supported scope:
FullAttentionSpecgroup — the one-group case of the same code path.MambaSpecgroups +FullAttentionSpecgroup, covering themamba_cache_modenone/align/allbehaviors.vLLM version support matrix (e2e-verified)
vLLM's flash_attn backend shipped three KV cache layouts over time; the connector detects them from the tensor shape itself (never from version strings) and normalizes each into token-major transfer pointers, so all three eras work through the same Triton kernel with zero kernel changes:
(2, num_blocks, block, H, D)split K/Vnum_external_computed_tokens == 0for mamba align mode; the connector fails fast at startup with a clearNotImplementedErrorpointing to vLLM >= 0.23.0(num_blocks, 2, block, H, D)split K/V(num_blocks, H, block, 2D)packed0.24.x / 0.25.x share the 0.23.0 layout and connector API surface (verified against the vLLM source per tag), so they are expected-compatible although not e2e-run.
Motivation
The previous connector was written against vLLM <= 0.22.1 and assumed a single full-attention KV cache group. That data model is fundamentally unfriendly to hybrid models:
kv_cache_groups, each with its own block table, block size and storage layout (token-granular KV for attention vs per-block opaque state for mamba/linear).Adapting the old implementation would have meant grafting a multi-group model onto single-group assumptions everywhere; a rewrite around "one group = one independent transfer unit" is simpler and strictly more general (a full-attention model is just the one-group case, there is no separate hybrid path).
Details
tp{rank}_g{group}), its own block-table translation and its own transfer strategy:alignmode materializes a recurrent state only at segment boundaries, so the interior manager blocks of a request have attention KV but no state — that sparsity is published to the manager through per-blocklocation_spec_group_names(groupsfull/attn), never encoded as a successful write, and the load side truncates an external match to the last state-complete block.attn_kv_viewsnormalizes each attention layer's paged cache (packed 4-D, or split K/V 5-D in either KV-first or N-first order) into token-major per-pointer views; split layouts contribute two transfer pointers per layer ([K0, V0, K1, V1, ...]), the interleaved N-first layout rides the kernel's existing strided path, and unrecognized layouts fail fast at startup.per_block_bytesis identical across layouts, so the manager, storage backends and transfer protocol are layout-agnostic.num_new_tokens > 0), canceled saves no longer raiseKeyError, and failed loads no longer enter a retry loop.Changes
kv_cache_manager/py_connector/vllm/v1_connector.pykv_cache_manager/py_connector/vllm/data_transfer.pykv_cache_manager/py_connector/vllm/metadata.pykv_cache_manager/py_connector/common/types.pyTransferGroup/KVCacheInfoshared data modelkv_cache_manager/py_connector/kernel/batch_gather_scatter_helper.pykv_cache_manager/py_connector/test/{test_block_translation,test_data_transfer_results,test_scheduler_state,test_kv_layouts}.pykv_cache_manager/py_connector/test/vllm_stubs.pykv_cache_manager/py_connector/test/kernel/test_strided_gather_scatter.pymanual)integration_test/vllm_e2e/*kv_cache_manager/manager/cache_manager.ccStartWriteCache: validatelocation_spec_group_namesagainst the block count, so token-only writes (no explicitblock_keys) can carry per-block group namesKnown Limitations / Risks
manual(they need a GPU runner, a prepared vLLM venv and local model checkouts); wiring them into CI needs runner/repo-variable setup and will land as a separate PR. The pure-logic unit tests do run in the existing open-source CI without torch.FullAttentionSpecgroups carryingsliding_window/attention_chunk_sizeat startup instead of silently publishing windowed KV as full-prefix caches.Scheduler._update_requests_with_invalid_blocks,vllm/v1/core/sched/scheduler.py) only understands single-group block tables (upstreamTODO (davidb): add support for hybrid memory allocator), so the connector passesreport_failures=Falsefor multi-group (hybrid) models: a failed hybrid load is logged but not reported, and vLLM may decode from partially loaded KV, potentially producing corrupt output until upstream adds multi-group invalid-block recovery. Single-group (pure attention) models report failures normally and recover via recompute.The reachable surface of this limitation is now limited to genuine storage failures: the case that used to walk into it deterministically — an external match ending on a block whose recurrent state was never written — is prevented at scheduling time by the per-block spec coverage and match truncation described above (
test_cross_request_prefix), so such a load is never started.mamba_cache_mode="align"). With prefix caching disabled (mamba_cache_mode="none") vLLM keeps a single resident mamba state block per request, so every manager block maps to that same block: saves would publish the current state under earlier prefix keys and loads would overwrite one block repeatedly (last write wins). Do not serve hybrid models through this connector with prefix caching off.alignmode materializes states sparsely (observed: 3 of 7 manager blocks). Correctness is prioritized over hit rate here — the attention KV of a longer prefix is unusable without its boundary state. The manager already supports complementary writes (a block's missing state specs can be filled in by a later write), so back-filling the missing states is a possible follow-up optimization.Testing
bazelisk test //kv_cache_manager/py_connector/...tair-kvcache-kvcm-dev, no torch installed).manual)://kv_cache_manager/py_connector/test/kernel:test_strided_gather_scatter— element-wise comparison of the strided gather/scatter against a naive torch reference.manual; needs 1-2 GPUs, a vLLM >= 0.26 venv with both KVCM wheels, and a local model — seeintegration_test/vllm_e2e/README.mdfor setup and the full environment-variable table,KVCM_E2E_MODEL/KVCM_E2E_PYTHONare required):bazelisk test //integration_test/vllm_e2e:e2e_tests \ --test_env=KVCM_E2E_PYTHON=/path/to/vllm-venv/bin/python \ --test_env=KVCM_E2E_MODEL=/path/to/modeltest_basic/test_concurrent/test_tppass bit-exact on vLLM 0.22.1 (full attention) and 0.23.0 (full attention + hybrid) — full per-version results in the validation report comment below.