[PD]: Bidirectional KV transfer

[S1ro1]

Sibling to PrimeIntellect-ai/prime-rl#2522

Summary

• Adds opt-in Decode -> Prefill KV metadata reuse for vLLM P/D routing behind pd_kv_cache_ttl_secs; default is 0, so the existing path does no extra conversation parsing or decode KV work.
• Keys reuse only by x-conversation-id, caches only object-valued decode kv_transfer_params, and prunes expired entries before reads and writes.
• Wires the TTL through Python router args, Rust config, CLI, and the PyO3 binding.

Validation

• cargo build --all-targets
• cargo test
• source .env && uv sync --all-extras from the sibling prime-rl branch after pinning this router commit

Note

Medium Risk
Adds a new in-memory, TTL-governed cache keyed by x-conversation-id that influences request payloads in the vLLM PD router; incorrect TTL sizing or high conversation churn could increase memory use or cause subtle routing/behavior changes in PD flows.

Overview
Enables optional bidirectional KV transfer in vLLM Prefill/Decode mode by caching Decode-side kv_transfer_params per conversation and reusing them on subsequent Prefill requests.

Introduces pd_kv_cache_ttl_secs (default 0 = disabled) wired through Python args/docs, Rust RouterConfig, CLI (--pd-kv-cache-ttl-secs), and the PyO3 binding. When enabled, VllmPDRouter stores kv_transfer_params extracted from Decode responses (including SSE streaming bodies) in a TTL cache keyed by x-conversation-id, evicts expired entries on access, and injects adjusted params into Prefill requests to point back at the Decode instance.

^{Reviewed by Cursor Bugbot for commit 1a441d6. Bugbot is set up for automated code reviews on this repo. Configure here.}

Note

Add bidirectional KV transfer support to VllmPDRouter with conversation-scoped caching

• Adds a decode_kv_cache to VllmPDRouter that stores kv_transfer_params from decode responses, keyed by x-conversation-id header, with a configurable TTL (pd_kv_cache_ttl_secs, default 0/disabled).
• On subsequent prefill requests within the same conversation, cached decode-side KV params are injected into the prefill request, enabling bidirectional P/D KV transfer.
• Exposes pd_kv_cache_ttl_secs through the Rust RouterConfig, Python RouterArgs, CLI (--pd-kv-cache-ttl-secs), and the PyO3 Router constructor.
• Cache eviction runs before each insert, removing entries whose TTL has elapsed.
• Behavioral Change: when pd_kv_cache_ttl_secs > 0 and requests include x-conversation-id, prefill kv_transfer_params will differ from default based on prior decode responses.

^{Macroscope summarized 1a441d6.}

[macroscopeapp]

Approvability

Verdict: Needs human review

1 blocking correctness issue found. This PR introduces substantial new functionality for bidirectional KV cache transfer, adding new caching logic and modifying request processing behavior. Additionally, there is an unresolved review comment identifying a potential bug where non-object cached values could break prefill-decode coordination.

^{You can customize Macroscope's approvability policy. Learn more.}

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 645c4c0. Configure here.}