Skip to content

fix(distserve): free Prefill cache with scheduler session id - #4989

Closed
gokay-ai wants to merge 1 commit into
InternLM:mainfrom
gokay-ai:cursor/fix-distserve-prefill-free-2445
Closed

gokay-ai wants to merge 1 commit into
InternLM:mainfrom
gokay-ai:cursor/fix-distserve-prefill-free-2445

Conversation

@gokay-ai

Copy link
Copy Markdown

Motivation

DistServe Proxy /v1/completions can leak Prefill scheduler sessions until the Prefill ZMQ engine dies.

After Prefill, the proxy set migration_request.remote_session_id from the OpenAI response id. That field is the client session_id (default -1), not the internal Prefill scheduler session. Decode then cache-frees -1, Prefill logs invalid free, and the preserve_cache sequence stays in scheduler.sessions / seq_manager. Chat has the same mismatch (chatcmpl-* is not an int).

Fixes #4967

Modification

  • Prefill /v1/completions and /v1/chat/completions return cache_session_id (engine/scheduler session id) with the DistServe cache fields.
  • Proxy uses that id for remote_session_id and PD shelf/unshelf. Missing or invalid owners are rejected instead of falling back to the public id.
  • Unit tests cover the public-id miss, response plumbing, proxy migration/free, and no-fallback reject path. No DistServe cluster required.

Related: #4977 also plumbs a cache owner id as part of a broader DistServe change. This PR is scoped to the #4967 free-path mismatch only.

BC-breaking (Optional)

Proxy and Prefill servers should be upgraded together: DistServe proxy now requires cache_session_id from Prefill. Public OpenAI response id is unchanged.

Checklist

  1. Ruff check on the touched files passed; focused pytest: 8 passed.
  2. Unit tests added for the session-id mismatch and proxy plumbing.
  3. N/A
  4. N/A

The DistServe proxy used the OpenAI response id as remote_session_id.
That id is the client session_id (default -1) or a chatcmpl-* string,
not the Prefill scheduler key, so cache-free logged invalid free and
leaked preserved Prefill sessions.

Return cache_session_id from Prefill and use it for migration/free.

Fixes InternLM#4967

Signed-off-by: GokayAI <60583610+gokay-ai@users.noreply.github.com>
@gokay-ai

Copy link
Copy Markdown
Author

Standing down in favor of #4977 (RunningLeon), which already returns cache_session_id from Prefill and uses it for proxy migration/free (same root cause as #4967), plus related DistServe fixes. No need for a second overlapping PR.

@gokay-ai gokay-ai closed this Sep 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] DistServe Proxy requests can leak Prefill scheduler metadata and OOM the Prefill engine

1 participant