[RFC]: Enable Prefix Caching with Hidden-State I/O (Multi-round / Service Scenarios)

RFC: Enable Prefix Caching with Hidden-State I/O (Multi-round / Service Scenarios)

Context / Motivation

This RFC requests feedback and alignment on extending vLLM prefix caching to support workflows that require hidden state inputs/outputs (e.g., multi-round interactions or production services that reuse intermediate representations). Today, prefix caching is effectively unavailable in these cases, leaving significant reuse/performance potential untapped.

Current Limitations

Prefix caching in vLLM is token-based and assumes the standard autoregressive flow. When hidden states are involved, we hit several blockers:

1. Token-based cache reuse does not cover hidden-state workflows
   Our pipeline may need to consume/produce hidden states across stages. The current prefix caching design does not model these dependencies.

2. Attention backend only returns uncached hidden states
   With prefix caching enabled, the vLLM attention backend returns hidden states only for the uncached portion of the sequence. Our use case requires full hidden states (cached + uncached) for downstream cross-stage logic.

3. KV-cache management ownership mismatch (runner vs scheduler)
   Our cross-stage processing is implemented primarily in the model runner. In contrast, vLLM prefix caching requires KV-cache management and prefix matching to be coordinated in the scheduler. This split makes integration difficult without refactoring dataflow and ownership boundaries.

Observed Performance Costs (Prefix Matching)

Measured on Qwen3 Omni with default mixed-modality inputs:

• Tensor match time: ~0.0096 s
• Hash time: ~0.0508 s
• Model forward time: ~0.0447 s

This suggests the current hashing-based prefix match can be comparable to (or exceed) a single model forward in some cases, so improving the matching strategy is important.

Example: Hidden-State Return Shape With/Without Prefix Caching

Two text requests on Qwen3 Omni:

• Prompt 1: Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words.
• Prompt 2: Explain the workflow of the audio generation pipeline

Hidden-state return shape for prompt 2:

• With prefix caching: [25, 2048]
• Without prefix caching: [57, 2048]

This illustrates that cached-prefix hidden states are not returned, which breaks downstream logic that requires the full sequence hidden states.

Proposed Change.

Proposal: Two-Phase Extension Plan

Phase 1 — Full Hidden States from Attention Backend

Goal: Make prefix caching compatible with hidden-state consumers by enabling retrieval of full hidden states (including cached segments).

Plan A:

• Add support (or an option) for the attention backend to return full hidden states.
• This likely requires collaboration/upstream changes in the vLLM repo.
• Once available, Thinker can enable prefix caching without losing required hidden-state outputs.

Plan B: Extend the vllm prefix caching method to also cache hidden states.

Plan C: Develop a new cache management mechanism similar to the current prefix caching to cache hidden states.

Expected outcome: Prefix caching becomes usable for hidden-state I/O workflows with minimal downstream changes.

Phase 2 — Tensor/Index-Based Prefix Match + Preprocess Dataflow Changes

Goal: Reduce prefix match overhead and align cache management with vLLM’s scheduler-driven design.

• Support prefix matching based on tensor content and index mapping, not only token hashing.
• Modify preprocess/dataflow:
  • Move preprocess responsibilities closer to the input processing path (earlier stage),
  • So prefix identification and cache bookkeeping happen before scheduling/execution,
  • And KV-cache management remains correctly owned by the scheduler.

Expected outcome: Better end-to-end performance and a clean integration path with vLLM’s scheduling and caching model.

Open Questions / Feedback Requested

• API expectations: What should “full hidden states” mean operationally (shape, memory overhead, optional flag, streaming vs materialized)?
• Ownership boundaries: How much cross-stage logic should move out of the model runner to align with scheduler-managed KV-cache?
• Matching strategy: Is tensor/index-based matching acceptable for correctness/performance, and how do we handle collisions/validation?

Success Criteria

• Prefix caching works when hidden-state I/O is required (no functional regressions).
• Prefix match overhead is significantly reduced vs current hashing time (~0.05 s baseline).
• Clear separation of concerns: scheduler manages KV-cache; runner focuses on execution.

Feedback Period.

No response

CC List.

@hsliuustc0106 @fake0fan @DarkLight1337 @Isotr0py @ywang96 @R2-Y @Gaohan123

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.