Skip to content

fix(pd): preserve prefill cache metadata in first-token race - #1424

Closed
sufubao wants to merge 1 commit into
ModelTC:mainfrom
sufubao:pd-token-race-cache-metadata
Closed

fix(pd): preserve prefill cache metadata in first-token race#1424
sufubao wants to merge 1 commit into
ModelTC:mainfrom
sufubao:pd-token-race-cache-metadata

Conversation

@sufubao

@sufubao sufubao commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator

问题

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)」的竞态序列,断言:

  • 第一个 yield 的是 decode 首 token,且 prompt_cache_len 被盖成 7;
  • 第二个 yield 的是 decode 第二 token(顺序未乱)。

原实现下该用例的首 token prompt_cache_len 会是 0。1 passed

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.
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@shihaobai shihaobai closed this Aug 4, 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.

2 participants