fix(pd): report real cached_tokens instead of always 0 - #1423
Merged
Conversation
In PD mode cached_tokens was always 0 in the response usage and access log. prompt_cache_len is only written into shm_req on the prefill node (in _match_radix_cache); the decode node never sets it, so every decode frame carries 0. api_openai overwrites cached_tokens on every stream chunk, so the final usage takes the last decode frame's 0 and drops the real hit reported by the prefill first frame. Fix: in generate(), capture the first block's prefill-reported prompt_cache_len and stamp it onto every yielded frame (single and multi block). In _wait_to_token_package(), track the max prompt_cache_len across frames for the access log instead of popping the last frame's 0.
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
2 participants
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 模式下响应 usage 和访问日志里的
cached_tokens恒为 0。根因：
prompt_cache_len只在 prefill 节点的_match_radix_cache里写入shm_req，decode 节点从不设置（恒为 0）。而api_openai每个 stream chunk 都会用当前帧覆盖cached_tokens，最终 usage 取的是最后一帧（decode 节点）的 0，prefill 首帧上报的真实命中被冲掉；pd_master 访问日志同理（取最后一帧metadata.pop("prompt_cache_len", 0)）。修复
generate()：捕获首块 prefill 上报的prompt_cache_len，回填到每一帧（单块 / 多块分段均正确，多块只保留首块命中值，符合「首块才是原始 prompt 的真实命中」）。_wait_to_token_package()：访问日志改为取本请求所有帧里prompt_cache_len的最大值，而不是弹出最后一帧的 0。测试
新增
unit_tests/server/httpserver/test_pd_master_cached_tokens.py：test_single_block_prefill_hit_persists_past_decode_zeros：单块场景，prefill 上报命中 30、后续 decode 帧上报 0，断言所有帧都是 30。test_multi_block_keeps_first_block_hit：两段（block0 命中 30、block1 命中 50），断言最终 usage 保留首块 30，不被 50 覆盖。