From b437a0d066a1be5356ec514ed7483a06593b4474 Mon Sep 17 00:00:00 2001 From: Jay Chun Date: Thu, 28 May 2026 12:50:03 +0900 Subject: [PATCH] Fix PD decode radix cache double-counting cached_tokens (#25973) --- python/sglang/srt/disaggregation/decode.py | 7 +++++++ .../srt/disaggregation/decode_schedule_batch_mixin.py | 8 +++++++- python/sglang/test/kits/cache_hit_kit.py | 9 +++++++++ 3 files changed, 23 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 27eb699b2..5138f571a 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -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() diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 86b9aeb3b..5ab9edcb6 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -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) diff --git a/python/sglang/test/kits/cache_hit_kit.py b/python/sglang/test/kits/cache_hit_kit.py index 5e1c9172c..d5117989e 100644 --- a/python/sglang/test/kits/cache_hit_kit.py +++ b/python/sglang/test/kits/cache_hit_kit.py @@ -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