[Hybrid KV Cache] Retain finalized Mamba decode checkpoints for prefix caching

[qianlihuang]

Purpose

With:

VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0

vLLM uses a sparse retention policy that aims to maximize cache reuse per retained SSM state, providing high prefix-cache hit potential with minimal storage overhead.

Currently, it retains the prompt replay boundary and reactively discovered Marconi-style shared-prefix junctions. Marconi's admission policy also retains the last decoded state for conversation-history reuse.

This PR completes that behavior for sparse retention by privately pinning the latest materialized, scheduler-aligned decode checkpoint and publishing it to the local prefix cache after FINISHED_STOPPED.

Enable with:

VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS=1 vllm serve <model> \
  --enable-prefix-caching \
  --mamba-cache-mode align \
  --prefix-cache-retention-interval 0

Examples

Kimi K3 requires the complete assistant message returned by the API, including reasoning_content and tool_calls, to be passed back in the next request.

History prefix hit examples

History followed by a tool result

... history ...

<|open|>message role="assistant"<|sep|>
<|open|>think<|sep|>...<|close|>think<|sep|>
<|open|>response<|sep|>...<|close|>response<|sep|>
<|open|>tools<|sep|>...<|close|>tools<|sep|>
<|close|>message<|sep|><|end_of_msg|>

<|open|>message role="tool" tool="get_weather" index="1"<|sep|>
Sunny, 31°C
<|close|>message<|sep|><|end_of_msg|>

The hit lands at the deepest scheduler-aligned boundary within the replayed assistant prefix.

History followed by a user message

... history ...

<|open|>message role="assistant"<|sep|>
<|open|>think<|sep|>...<|close|>think<|sep|>
<|open|>response<|sep|>previous answer<|close|>response<|sep|>
<|close|>message<|sep|><|end_of_msg|>

<|open|>message role="user"<|sep|>
Can you explain that further?
<|close|>message<|sep|><|end_of_msg|>

Again, the hit is the deepest aligned checkpoint within the replayed assistant prefix.

Because the completed assistant message is rendered again before the new tool or user message, the latest decode checkpoint provides the deepest reusable prefix.

Test Plan

Server

export VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0
export VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS=<0-or-1>

vllm serve /qwen3p6-35b-a3b-nvfp4 \
  --trust-remote-code \
  -tp 2 \
  --language-model-only \
  --enable-prefix-caching \
  --mamba-cache-mode align \
  --prefix-match-unit 16 \
  --enable-auto-tool-choice \
  --tool-call-parser qwen3_xml \
  --reasoning-parser qwen3 \
  --default-chat-template-kwargs \
    '{"enable_thinking":true,"preserve_thinking":true}' \
  --enable-prompt-tokens-details

Client

#!/usr/bin/env python3
"""Decode-checkpoint retention A/B — one-shot client.

Runs a fixed multi-turn dialog (full-history replay, bs=1, temp=0, seed=42)
against a running vLLM server, prints per-request cached_tokens + output hash.
The FULL assistant message (reasoning + content + tool_calls) is replayed
verbatim each turn, so the replay token sequence matches what was generated —
required for PR #50551 decode checkpoints to hit.

Run once per arm (retain=0, then retain=1) and diff the two prints.
Usage:
  python3 ab_client.py
"""
import hashlib
import json

from openai import OpenAI

MODEL = "/qwen3p6-35b-a3b-nvfp4"
BASE = "http://127.0.0.1:8000/v1"
SEED, TEMP = 42, 0.0

USER_MSGS = [
    "Analyze the trade-offs between cloud-based and on-premises deployment "
    "for a financial services company that must meet strict data residency "
    "requirements. Cover cost, latency, security, compliance, and operational "
    "complexity, then recommend a hybrid approach with a phased migration plan.",
    "Based on the previous analysis, propose a detailed incident-response runbook "
    "for a hybrid cloud outage. Include team roles, escalation path, communication "
    "channels, and step-by-step recovery for both the on-prem and cloud tiers.",
    "Now write a risk register for that runbook: for each of the top 6 risks, "
    "give likelihood, impact, mitigations, and a monitoring metric that would "
    "have caught it early.",
    "Draft the executive summary of this whole plan for the board — under 300 "
    "words, non-technical, focused on business continuity and the expected "
    "downtime budget.",
    "Finally, produce a comparison table as plain text with columns for cloud, "
    "on-prem, and hybrid across 8 dimensions, and a one-line recommendation "
    "justification for each dimension.",
]


def main():
    client = OpenAI(base_url=BASE, api_key="dummy")
    messages = []
    rows = []
    for turn, user_msg in enumerate(USER_MSGS):
        messages.append({"role": "user", "content": user_msg})
        resp = client.chat.completions.create(
            model=MODEL, messages=messages, temperature=TEMP, seed=SEED)
        choice = resp.choices[0]
        if choice.finish_reason != "stop":
            raise RuntimeError(f"turn={turn}: finish_reason={choice.finish_reason!r}")
        usage = resp.usage
        details = getattr(usage, "prompt_tokens_details", None)
        cached = getattr(details, "cached_tokens", 0) or 0 if details else 0

        msg_data = choice.message.model_dump(exclude_none=True)
        content = msg_data.get("content") or ""
        reasoning = (msg_data.get("reasoning")
                     or msg_data.get("reasoning_content") or "")
        tool_calls = msg_data.get("tool_calls") or []

        rows.append({
            "turn": turn,
            "prompt_tokens": usage.prompt_tokens,
            "cached_tokens": cached,
            "completion_tokens": usage.completion_tokens,
            "reasoning_chars": len(reasoning),
            "content_chars": len(content),
            "output_sha": hashlib.sha256(json.dumps(
                {"reasoning": reasoning, "content": content,
                 "tool_calls": tool_calls},
                sort_keys=True, ensure_ascii=False).encode()).hexdigest()[:12],
        })

        # Replay the FULL assistant message, not just content.
        assistant_message = {"role": "assistant", "content": content}
        if reasoning:
            assistant_message["reasoning"] = reasoning
        if tool_calls:
            assistant_message["tool_calls"] = tool_calls
        messages.append(assistant_message)

        print(f"turn={turn} pt={usage.prompt_tokens} ct={usage.completion_tokens} "
              f"cached={cached} reasoning_chars={len(reasoning)}", flush=True)

    print("\nturn | cached/pt/ct | reasoning_chars/content_chars | output_sha")
    for i, r in enumerate(rows):
        print(f"{i} | {r['cached_tokens']}/{r['prompt_tokens']}/{r['completion_tokens']} "
              f"| {r['reasoning_chars']}/{r['content_chars']} | {r['output_sha']}")

    pt = sum(r["prompt_tokens"] for r in rows)
    ca = sum(r["cached_tokens"] for r in rows)
    ct = sum(r["completion_tokens"] for r in rows)
    print(f"TOTAL | cached={ca}/prompt={pt}/completion={ct} | hit_rate={ca/pt:.1%}")


if __name__ == "__main__":
    main()

Test Result

Turn	Retain=0 (cached / prompt / completion)	Retain=1 (cached / prompt / completion)
0	0 / 61 / 3613	0 / 61 / 3518
1	48 / 3729 / 4374	3168 / 3634 / 4080
2	3728 / 8151 / 2950	7392 / 7763 / 3071
3	8144 / 11145 / 1977	10560 / 10878 / 1290
4	11136 / 13168 / 8126	11616 / 12214 / 2935
Total	23056 / 36254 / 21040	32736 / 34550 / 14894
Hit rate	63.6%	94.7%

Relationship to prior work

#37898 added Marconi-style shared-prefix admission while preserving the existing last-state behavior.

#45845 introduced sparse Mamba retention. With VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0, retention became based on replay boundaries and no longer included the finalized decode endpoint.

#47782 preserved shared-prefix junctions under sparse retention, but did not restore the finalized decode endpoint.

This PR restores that remaining last-state behavior.

Future work

A separate follow-up will support hidden-state drafting.

A follow-up can offload finalized Decode-side checkpoints through KV connectors for reuse by Prefill workers in P/D deployments.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

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

[Copilot]

Pull request overview

This PR adds an opt-in mechanism for sparse prefix-cache retention (VLLM_PREFIX_CACHE_RETENTION_INTERVAL=0) to privately pin the latest scheduler-aligned Mamba align-mode decode checkpoint during generation and publish it to the local prefix cache only when the request finishes with FINISHED_STOPPED. This targets multi-turn “full-history replay” workflows (e.g., Kimi K3) by enabling deeper reusable replay boundaries with minimal additional storage.

Changes:

• Introduces private per-request decode-checkpoint candidates in the Mamba align KV cache manager and publishes/discards them at request finalization.
• Adds BlockPool support for caching an exact decode checkpoint hash alias without disturbing other existing hash aliases.
• Wires scheduler and KV cache manager/coordinator logic to update candidates on non-stale outputs and finalize them on request completion; includes validation + env var + comprehensive tests.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
vllm/v1/core/single_type_kv_cache_manager.py	Tracks/pins latest aligned decode checkpoint candidate per request in Mamba align mode; publishes on finalization.
vllm/v1/core/sched/scheduler.py	Updates decode-checkpoint candidates on fresh outputs and finalizes them during request teardown.
vllm/v1/core/kv_cache_manager.py	Computes “materialized” token boundary and forwards candidate updates/finalization to the coordinator.
vllm/v1/core/kv_cache_coordinator.py	Adds config validation and coordinator fan-out for candidate update/finalization across cache groups.
vllm/v1/core/block_pool.py	Adds cache_decode_checkpoint to publish a specific checkpoint hash alias and optionally emit events.
vllm/envs.py	Adds VLLM_PREFIX_CACHE_RETAIN_DECODE_CHECKPOINTS env var and updates retention-interval description.
tests/v1/core/test_prefix_caching.py	Adds unit tests covering candidate selection, pin semantics, promotion behavior, event emission, and config validation.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[zzw09773]

Data point supporting this PR: on 8×MI325X (ROCm, nightly g3ee2df303 + #50597/#50817 for the SITU MoE path), enabling --enable-prefix-caching on Kimi-K3 without these SSM checkpoints looks fine on trivial tests (exact-prefix re-send: 21× latency win, 27k+ cache hits) but dies under real traffic within ~40 minutes: a worker is killed by a signal (exit code None, no Python traceback) while resuming a ~48k-token request from partially-hit cache blocks with 3 concurrent requests. KV usage was only 7% — not memory pressure. Length-independent; partial-hit + concurrency is the trigger. We've reverted to no-prefix-caching until this lands. Happy to test the PR on gfx942.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @qianlihuang.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @qianlihuang.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork