diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 656093296..eaa943855 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2509,12 +2509,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): req.extend_batch_idx += 1 - # update req-level memory management fields - # TODO(th4): co-locate this req.kv bookkeeping with the real KV - # allocation in alloc_for_extend above; they are currently a few - # steps apart and should become one owned-kv allocation step. - req.kv.kv_committed_len = seq_len - # If input_embeds are available, store them if req.input_embeds is not None: # Slice to match extend_input_len — PrefillAdder truncates @@ -3221,10 +3215,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # the allocator, triggered from mem_cache/common.py.) self.out_cache_loc = alloc_for_decode(self, token_per_req=1) - # Update req-level memory management fields for req in self.reqs: req.decode_batch_idx += 1 - req.kv.kv_committed_len += 1 # New-tensor avoids racing model_worker_batch refs queued for # overlap forward. diff --git a/python/sglang/srt/mem_cache/allocation.py b/python/sglang/srt/mem_cache/allocation.py index 211183b05..5e3044757 100644 --- a/python/sglang/srt/mem_cache/allocation.py +++ b/python/sglang/srt/mem_cache/allocation.py @@ -387,6 +387,7 @@ def alloc_for_extend( for req, seq_len in zip(batch.reqs, batch.seq_lens_cpu.tolist()): req.kv.kv_allocated_len = seq_len + req.kv.kv_committed_len = seq_len return out_cache_loc, req_pool_indices_device, req_pool_indices_cpu @@ -581,6 +582,7 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor: for req in batch.reqs: req.kv.kv_allocated_len += token_per_req + req.kv.kv_committed_len += token_per_req return out_cache_loc diff --git a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py index c5cb85193..4c6b7f576 100644 --- a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py +++ b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py @@ -53,14 +53,14 @@ _SS = "session/streaming_session.py" _OWNER_SITES = { # non-spec scheduler (_SB, "ScheduleBatch.prepare_for_decode", "decode_batch_idx"): 1, - (_SB, "ScheduleBatch.prepare_for_decode", "kv_committed_len"): 1, (_SB, "ScheduleBatch.prepare_for_extend", "extend_batch_idx"): 1, - (_SB, "ScheduleBatch.prepare_for_extend", "kv_committed_len"): 1, # kv_allocated_len is settled inside the owned-kv alloc functions (op28). ("mem_cache/allocation.py", "alloc_for_extend", "evict"): 1, ("mem_cache/allocation.py", "alloc_for_extend", "kv_allocated_len"): 1, + ("mem_cache/allocation.py", "alloc_for_extend", "kv_committed_len"): 1, ("mem_cache/allocation.py", "alloc_for_decode", "evict"): 1, ("mem_cache/allocation.py", "alloc_for_decode", "kv_allocated_len"): 1, + ("mem_cache/allocation.py", "alloc_for_decode", "kv_committed_len"): 1, # spec v2: no pre-claim; resolve commits the full accepted run uniformly. # kv_allocated_len for spec v2 draft decode (eagle + dflash) is settled # inside the owned-kv alloc_for_spec_decode function (op42).