minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Spec] Unify decode KV-commit bookkeeping across spec-v2 workers (#28754)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-23 00:53:05 -07:00
committed by GitHub
parent 743ce88bc5
commit 854c688121
5 changed files with 21 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/managers/test_batch_result_processor_spec_grammar.py
+3 -4
View File
@@ -106,9 +106,8 @@ class TestSpecV2GrammarTruncation(CustomTestCase):
predict_tokens = proc._resolve_spec_v2_tokens(result, _FakeBatch([req]))
self.assertEqual(predict_tokens, [[101, 102]])
# EAGLE commits (retained - 1): prepare_for_decode pre-claimed the bonus
# slot, and the dropped suffix is never committed.
self.assertEqual(req.kv_committed_len, 2 - 1)
# No pre-claim: commit the full retained run (no -1 refund).
self.assertEqual(req.kv_committed_len, 2)
def test_resolve_keeps_all_when_grammar_not_terminated(self):
req = _make_req(terminate_after=99)
@@ -118,7 +117,7 @@ class TestSpecV2GrammarTruncation(CustomTestCase):
predict_tokens = proc._resolve_spec_v2_tokens(result, _FakeBatch([req]))
self.assertEqual(predict_tokens, [[201, 202, 203]])
self.assertEqual(req.kv_committed_len, 3 - 1)
self.assertEqual(req.kv_committed_len, 3)
if __name__ == "__main__":
test/registered/unit/spec/test_decode_bookkeeping_ownership.py
+4 -7
View File
@@ -55,16 +55,11 @@ _OWNER_SITES = {
(_SB, "ScheduleBatch.prepare_for_extend", "kv_allocated_len"): 1,
("mem_cache/common.py", "alloc_for_extend", "evict"): 1,
("mem_cache/common.py", "alloc_for_decode", "evict"): 1,
# spec v2: pre-claim in the scheduler-driven mixin, settle in resolve
# spec v2: no pre-claim; resolve commits the full accepted run uniformly.
(*_MIXIN, "decode_batch_idx"): 1,
(*_MIXIN, "evict"): 1,
(*_MIXIN, "kv_committed_len"): 1,
(*_MIXIN, "kv_allocated_len"): 1,
# 3rd resolve mutation: DFLASH settles its full commit_lens here (no
# pre-claim in prepare_for_decode, unlike the EAGLE mixin).
# Spec grammar truncation commits only the retained (pre-termination) length
# here, so the dropped suffix is never over-committed (no later rollback).
(*_RESOLVE, "kv_committed_len"): 3,
(*_RESOLVE, "kv_committed_len"): 1,
(*_RESOLVE, "spec_verify_ct"): 1,
(
"speculative/dflash_info_v2.py",
@@ -92,6 +87,8 @@ _OWNER_SITES = {
(_SS, "StreamingSession._trim_overshoot", "kv_committed_len"): 1,
(_SS, "StreamingSession._trim_overshoot", "kv_allocated_len"): 1,
(_SS, "StreamingSession.try_cache_finished_req", "kv_allocated_len"): 1,
# Inherit the authoritative finished length (not the lagging req clock).
(_SS, "StreamingSession.try_cache_finished_req", "kv_committed_len"): 1,
}