Fix(disagg): correct prefill chunk progress and cache owner IDs

[RunningLeon]

Motivation

Fix two correctness issues in PyTorch Prefill/Decode disaggregation (DistServe):

• Chat fails with HTTP 500 because the proxy parses the public chatcmpl-... ID as an integer cache-owner ID. Completions can also use the wrong owner for cache release.
• Long prompts stall because the Prefill-role early return skips chunk-history updates.

Modification

• Advance chunk history before the Prefill-role early return.
• Return an explicit internal cache_session_id from chat/completions and use it consistently for migration and shelf tracking. Reject missing or invalid owners instead of falling back to the public ID.
• Keep the existing forwarding and cleanup structure; lifecycle changes and shared forwarding are deferred to a separate PR.

Validation: 91 CPU tests passed, plus Ruff and commit hooks. On full Qwen3-30B-A3B-Instruct-2507 with two H200 GPUs and real synchronous Mooncake RDMA, 30 PD requests and 4 direct Decode controls passed across both endpoints, streaming modes, chunk boundaries, and concurrency. All 30 migration ACKs matched Prefill session release; P/D scheduler sessions drained to zero and services shut down cleanly.

This verifies functional completion, not numerical equivalence. Cancellation/error-path lifecycle fixes and linear-attention state migration are outside this PR.

BC-breaking (Optional)

Upgrade the proxy and Prefill servers together: the proxy requires cache_session_id. Public response IDs remain unchanged.

Checklist

1. Lint and commit hooks passed.
2. Focused regression tests and real-GPU functional checks passed; validation limits are noted above.
3. No new runtime dependency or downstream version requirement.
4. Code comments clarify cache ownership and chunk progression; compatibility is documented above.

[Copilot]

🟡 Changes recommended

Critical cache-owner correctness and retained-session cleanup issues remain unresolved.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Fixes DistServe cache-owner propagation and long-prefill chunk progress.

Changes:

• Adds internal cache_session_id propagation and validation.
• Advances chunk history before Prefill early returns.
• Adds regression coverage for ownership and chunk progression.

Critical findings:

• completions.py:338 (3 votes): Batched prompts may use the wrong cache owner.
• proxy.py:715 (2 votes): Invalid or missing owners can leak retained Prefill sessions.

File summaries

File	Description
tests/test_lmdeploy/serve/openai/test_distserve.py	Tests cache-owner propagation and validation.
tests/pytorch/engine/test_inputs_maker.py	Tests chunk-history progression.
lmdeploy/serve/proxy/proxy.py	Validates and tracks internal cache owners.
lmdeploy/serve/openai/endpoints/completions.py	Emits cache-owner metadata.
lmdeploy/serve/openai/chat_completions/serving.py	Emits cache-owner metadata.
lmdeploy/pytorch/engine/inputs_maker.py	Updates chunk history for Prefill workers.

Review details

• Files reviewed: 6/6 changed files
• Comments generated: 2
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.