fix(pd): preserve prefill cache metadata in first-token race - #1424
Closed
sufubao wants to merge 1 commit into
Closed
fix(pd): preserve prefill cache metadata in first-token race#1424sufubao wants to merge 1 commit into
sufubao wants to merge 1 commit into
Conversation
Both the prefill node and the decode node can emit the first token (count_output_tokens == 1) for a PD request, and they race. Whichever arrives first was yielded immediately; the duplicate from the other node was dropped via `continue`. prompt_cache_len is reported only on the prefill node's first token, so when the decode node's first token won the race the prefill first token (and its prompt_cache_len) was dropped, leaving the stream's cache metadata at 0. Fix: hold the winning first token (and buffer subsequent tokens) until the prefill node's first token arrives, then stamp its prompt_cache_len onto the held first token and flush the held tokens in order.
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
问题
PD 模式下 prefill 节点和 decode 节点都会为一个请求上报首 token(
count_output_tokens == 1),两者存在竞态。原来的处理是谁先到谁 yield,另一个当重复 token 用continue丢弃。但
prompt_cache_len(缓存命中长度)只在 prefill 节点的首 token 上携带。当 decode 节点的首 token 抢先到达并被 yield,prefill 节点的首 token(带着真实prompt_cache_len)就被当作重复 token 丢弃,整条流的缓存命中元信息变成 0,导致 usage / 访问日志的cached_tokens不准(见配套 PR #1423 的另一面)。修复
fetch_pd_stream:先到的首 token 不立即 yield,而是暂存在first_token_package,其后的 token(无论是重复首 token 还是output_index > 1)暂存进pending_token_list,直到 prefill 节点的首 token 到达、拿到prefill_prompt_cache_len;然后把它盖到暂存的首 token 上,再按原顺序 flush 出去。顺序保持不变:暂存期间所有后续 token 一并进
pending_token_list,不会插队。测试
新增
unit_tests/server/httpserver/test_pd_master_token_race.py:构造「decode 首 token 先到、decode 第二 token 紧随、prefill 首 token 后到(带prompt_cache_len=7)」的竞态序列,断言:prompt_cache_len被盖成 7;原实现下该用例的首 token
prompt_cache_len会是 0。1 passed。