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
|
# Success - commit the transfer
|
||||||
decode_req.req.output_ids.append(output_id[0].item())
|
decode_req.req.output_ids.append(output_id[0].item())
|
||||||
decode_req.req.cached_tokens = cached_tokens[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_device = cached_tokens[1].item()
|
||||||
decode_req.req.cached_tokens_host = cached_tokens[2].item()
|
decode_req.req.cached_tokens_host = cached_tokens[2].item()
|
||||||
decode_req.req.cached_tokens_storage = cached_tokens[3].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}"
|
), f"seq_len={seq_len}, pre_len={pre_len}, req.extend_input_len={req.extend_input_len}"
|
||||||
|
|
||||||
if not req.retracted_stain:
|
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.already_computed = seq_len
|
||||||
req.is_retracted = False
|
req.is_retracted = False
|
||||||
pre_lens.append(pre_len)
|
pre_lens.append(pre_len)
|
||||||
|
|||||||
@@ -344,6 +344,15 @@ def run_multiturn_cache_hit_test(
|
|||||||
print(msg)
|
print(msg)
|
||||||
|
|
||||||
assert resp.cached_tokens >= expected_cached
|
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
|
# Record this round's prompt_len for next round's expected calc
|
||||||
prev_prompt_lens[i] = resp.prompt_len
|
prev_prompt_lens[i] = resp.prompt_len
|
||||||
|
|||||||
Reference in New Issue
Block a user