Fix PD decode radix cache double-counting cached_tokens (#25973)
This commit is contained in:
@@ -1436,6 +1436,13 @@ class DecodeTransferQueue:
|
||||
# Success - commit the transfer
|
||||
decode_req.req.output_ids.append(output_id[0].item())
|
||||
decode_req.req.cached_tokens = cached_tokens[0].item()
|
||||
# The prefill node already reported its prefix-cache hit in
|
||||
# cached_tokens[0]. Seed already_computed with it so that
|
||||
# prepare_for_prebuilt's `cached_tokens += pre_len - already_computed`
|
||||
# only adds decode-side reuse *beyond* what prefill counted, instead of
|
||||
# double-counting the shared prompt prefix (which would make
|
||||
# cached_tokens exceed prompt_tokens when decode radix cache is on).
|
||||
decode_req.req.already_computed = decode_req.req.cached_tokens
|
||||
decode_req.req.cached_tokens_device = cached_tokens[1].item()
|
||||
decode_req.req.cached_tokens_host = cached_tokens[2].item()
|
||||
decode_req.req.cached_tokens_storage = cached_tokens[3].item()
|
||||
|
||||
@@ -62,7 +62,13 @@ class ScheduleBatchDisaggregationDecodeMixin:
|
||||
), f"seq_len={seq_len}, pre_len={pre_len}, req.extend_input_len={req.extend_input_len}"
|
||||
|
||||
if not req.retracted_stain:
|
||||
req.cached_tokens += pre_len - req.already_computed
|
||||
# Clamp to avoid double-counting: already_computed is seeded from
|
||||
# the prefill-reported cached_tokens in _commit_transfer_to_req, so
|
||||
# a decode-side prefix shorter than the prefill report must not
|
||||
# subtract from cached_tokens.
|
||||
delta = max(0, pre_len - req.already_computed)
|
||||
req.cached_tokens += delta
|
||||
req.cached_tokens_device += delta
|
||||
req.already_computed = seq_len
|
||||
req.is_retracted = False
|
||||
pre_lens.append(pre_len)
|
||||
|
||||
@@ -344,6 +344,15 @@ def run_multiturn_cache_hit_test(
|
||||
print(msg)
|
||||
|
||||
assert resp.cached_tokens >= expected_cached
|
||||
# Upper bound: cached tokens are a subset of the prompt, so they can
|
||||
# never exceed prompt_len. In PD disaggregation with decode radix
|
||||
# cache, the shared prefix was previously counted on both the prefill
|
||||
# and the decode node, making cached_tokens exceed prompt_len.
|
||||
assert resp.cached_tokens <= resp.prompt_len, (
|
||||
f"Round {round_num}, client {i}: cached_tokens="
|
||||
f"{resp.cached_tokens} exceeds prompt_len={resp.prompt_len} "
|
||||
f"(double-counted prefix across prefill/decode)"
|
||||
)
|
||||
|
||||
# Record this round's prompt_len for next round's expected calc
|
||||
prev_prompt_lens[i] = resp.prompt_len
|
||||
|
||||
Reference in New Issue
Block a user