vllm - 💡(How to fix) Fix [Bug][PD] Bidirectional KV transfer produces incorrect results when reasoning traces are stripped between turns [5 comments, 4 participants]

This is an issue to make sure this behavior is tracked and consistently addressed by upper-level routers.

Summary

Bidirectional KV transfer (PR #32553, RFC #32733) can produce incorrect inference results when used with reasoning models (e.g. DeepSeek-R1) whose thinking traces are stripped from the conversation history between turns.

Problem

When D generates a response with thinking traces, its kv_transfer_params records:

• remote_num_tokens = request.num_computed_tokens — covering [prompt | thinking_tokens | response_tokens]
• remote_block_ids — physical blocks for the entire sequence

On the next turn, if the client strips thinking traces before sending, P receives a prompt like [prompt | response_tokens | new_user_msg] — the thinking tokens are missing from the middle.

The block-alignment logic in _apply_prefix_caching (worker.py:2310-2316) does suffix trimming:

remote_block_ids[i] = remote_group[-num_local_blocks:]

This assumes P's prompt is a strict prefix of D's sequence — true in the normal case, but broken when tokens are removed from the middle (for whatever reason, like compacting history or dropping thinking traces).

Result: P loads KV cache computed for different tokens than its actual input, leading to silently incorrect inference.

Suggested Fixes

Router-level detection: (currently suggested approach) A production router can keep the previous turn tokens, indexing by conversation_id. On turn 2, the router can compare prompt tokens. If the prompt is different than expected (traces were stripped), clear the cached kv params and let P fall back to full recomputation.

vllm-level fix: We would need to design a thinking tokens-aware solution, possibly block-aligned so that only non-thinking tokens are pulled. This would require significantly more effort and an RFC first.