[mem_cache] Settle extend kv_committed_len inside alloc_for_extend (#37085)

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.

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
 

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).