Fix per-tier cached-token re-count for chunked prefill#1
Open
ishandhanani wants to merge 14 commits into
Open
Conversation
Continuation chunks of a chunked prefill re-enter can_run_list via add_chunked_req with prefix_len=0, but from_adder re-read the full req.cached_tokens_* split on every chunk. A 4-chunk request with a 1344-token cached prefix logged: chunk 1: #cached-token: 1344, #cached-token-device: 1344 chunk 2: #cached-token: 0, #cached-token-device: 1344 chunk 3: #cached-token: 0, #cached-token-device: 1344 chunk 4: #cached-token: 0, #cached-token-device: 1344 breaking log_hit_tokens == device + host + storage and re-counting the tier split once per chunk. Track continuation chunks on the adder and credit the split only on the chunk where the prefix is credited. Also drop the formatting-only hunks in schedule_policy.py to keep the diff scoped.
5 tasks
hxieustc
force-pushed
the
add-fine-grain-prefill-cached-token-info
branch
from
June 4, 2026 21:35
77394e6 to
a4be49d
Compare
hxieustc
force-pushed
the
add-fine-grain-prefill-cached-token-info
branch
11 times, most recently
from
June 17, 2026 01:20
f520b04 to
0e6b21c
Compare
hxieustc
force-pushed
the
add-fine-grain-prefill-cached-token-info
branch
2 times, most recently
from
June 23, 2026 21:03
516d152 to
d45be44
Compare
hxieustc
force-pushed
the
add-fine-grain-prefill-cached-token-info
branch
2 times, most recently
from
June 26, 2026 21:09
6ae79bb to
68c830b
Compare
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.
Follow-up to sgl-project#26976 — re-tested the updated branch on 1x L40S (Qwen3-0.6B, HiCache, write_through, page_size 64). All previously-flagged issues are fixed and the tier split is correct for host/device hits. One remaining bug found, fixed here.
Chunked prefill re-counts the tier split
Continuation chunks re-enter
can_run_listviaadd_chunked_reqwithprefix_len=0, butPrefillStats.from_adderre-reads the fullreq.cached_tokens_*split on every chunk. With--chunked-prefill-size 512and a 1344-token cached prefix (4-chunk request):This breaks the invariant documented on
PrefillStats(log_hit_tokens == device + host + storage) and re-counts the split once per chunk for anyone aggregating the log lines. Defaultchunked_prefill_sizeis 4096, so any prompt longer than that hits this in production.Fix
Track continuation chunks on the adder (
continuation_chunk_reqs) and credit the split only on the chunk where the prefix is credited — mirroringlog_hit_tokenssemantics. Also reverts the formatting-only hunks inschedule_policy.pyso the PR diff stays scoped.Verification (same box/config)
#cached-token: 0,#cached-token-device: 13441344 / 1344; chunks 2-4:0 / 0sum #cached-token = 133952 == sum device 121536 + sum host 12416#cached-token: 1536, device: 0, host: 1536on load-back;device: 1536, host: 0after promotionCI States
Latest PR Test (Base): ❌ Run #26969001224
Latest PR Test (Extra): ❌ Run #26969001003