[charxwu]

Overview

Fix incorrect output from SimpleCPUOffloadConnector when a GPU prefix ends inside an attention chunk. Keep the GPU prefix and recompute the suffix when CPU lookup cannot safely preserve absolute chunk boundaries.

Claims

• Reject CPU suffix lookup when the local prefix is not aligned to every participating chunked-local attention group.
• Preserve CPU reuse for requests without a GPU prefix and for chunk-aligned GPU prefixes.
• Add scheduler and real-LLM regressions and document the recomputation fallback.

Validation

Python source baseline: upstream main at 284c0317c486fc605b915a39a540053ec3a1b59f.

uv run --no-sync .venv/bin/python -m pytest \
  tests/v1/simple_kv_offload/test_scheduler.py \
  -k chunked_local_cpu_lookup_requires_chunk_aligned_prefix -q

Without the fix: 2 failed, 4 passed. Both 16-token local-prefix cases return (64, True) instead of (0, False), whether the required CPU tail block is present or absent.

uv run --no-sync .venv/bin/python -m pytest \
  tests/v1/simple_kv_offload/test_scheduler.py \
  tests/v1/simple_kv_offload/test_kv_events.py -q

Without the fix: 2 failed, 110 passed. With the fix: 112 passed. The six new cases cover local prefixes of 0, 16 and 64 tokens and CPU-tail presence/absence.

uv run --no-sync .venv/bin/python -m pytest \
  tests/v1/simple_kv_offload/test_integration.py \
  -k chunked_local_partial_gpu_hit -q

Without the fix: 1 failed in check_logprobs_close (cold token 487 versus mixed-hit token 48). With the fix: 1 passed. The test saves deterministic tiny Llama4 weights, loads them through vllm_runner, and exercises LLM.generate with one chunked-local and one full-attention layer. A pure CPU hit still reports 80 cached tokens; the misaligned mixed hit retains 16 GPU-cached tokens and recomputes its suffix. Cache state is established through generation and cache-reset APIs, without patching target methods or populating cache internals. Asynchronous stores are drained by condition polling and engine steps, with a timeout and no fixed sleeps.

GPU configuration: one AMD Instinct MI350X, TP1, FP16, eager execution, Triton attention, block size 16, attention chunk size 64 and an 81-token prompt. Both versions use the same prebuilt native extensions: vLLM 0.32.1rc1.dev276+gdf6603f05, Torch 2.13.0+rocm10.0.0, HIP 7.15.26333, Transformers 5.20.2. Complete Python source comes from the baseline above; a fresh native build is not claimed.

Supplementary validation with an independent deterministic saved checkpoint compares all 512 next-token logprobs against cold inference (absolute tolerance 0.005):

Source / GPU prefix	External hit / loaded blocks	Cold → mixed token	Max absolute logprob error
Main / 16 tokens	64 tokens / 4 blocks	60 → 385	4.781295
Fixed / 16 tokens	0 tokens / 0 blocks	60 → 60	0

Pure CPU hits remain correct, with 80 cached tokens, token 60 and maximum logprob error 0.001997. The supplementary checkpoint differs from the Transformers-generated pytest checkpoint, hence the different token IDs. Saved random weights validate cache execution correctness; trained-model quality, throughput, CUDA Graphs and multi-GPU execution were not evaluated.

The existing v1-kv-offload Buildkite step selects both regression suites, with NVIDIA and AMD coverage; the test-tethering check passes. Applicable pre-commit hooks, including Ruff, mypy and Markdown lint, pass. Unmatched actionlint was excluded from local hook initialization because its Go download fails TLS certificate verification. The usage page was rendered with the repository's MkDocs theme and Markdown extensions; the full-site build encountered a PydanticMagicMock type error in the CLI generator's ROCm capability detection.

Details

CPU lookup slices off the GPU-local prefix and passes the remaining hashes and token budget to the coordinator. Chunked-local lookup interprets that budget using absolute chunk positions, so a prefix ending inside a chunk changes which blocks are treated as null padding.

With a 16-token GPU prefix and an 81-token prompt, CPU lookup reports a 64-token suffix hit with four null chunk entries. The claimed absolute hit is 80 tokens, so tokens 64–79 still require a real chunk KV block. The load planner skips that block, producing incorrect inference. The new alignment guard declines this lookup. Its tradeoff is additional prefill for misaligned mixed hits.

Duplicate-work checks were refreshed on 2026-10-23. #57325 introduces broader joint GPU/CPU lookup while retaining the legacy suffix API for fallback paths, including composed connectors. This patch protects that retained API and does not implement joint lookup. #61375 changes OffloadingConnector partial-tail stores; #39702 concerns CPU-hit pinning across lookup/allocation. Neither addresses this chunk-position error.

AI assistance (Codex) was used for investigation, implementation, validation and PR preparation. This PR is ready for maintainer review. The validation results reported above were obtained by Codex.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[mergify]

Documentation preview: https://vllm--63266.org.readthedocs.build/en/63266/

[charxwu]

Hi @shen-shanshan, could you please review this fix for incorrect chunked-local CPU suffix hits in SimpleCPUOffloadConnector?

The issue reproduces on main at 284c0317c4. The fix retains the GPU prefix and recomputes the suffix when the local prefix is not attention-chunk aligned. The real-LLM regression fails before the fix and passes after it on MI350X, and all 112 scheduler/events tests pass with the fix. Details and reproduction commands are in the PR description. Thanks!

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[shen-shanshan]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97054 for commit bbb138c26fe2.

[shen-shanshan]

/ci run

[github-actions]

✅ Triggered Buildkite CI #97083 for commit 6e33e4c98776.