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

Avoid scattered assignment of extend_input_len and fill_len by merging them into Req.extend_range (#27610)

Browse Source
This commit is contained in:
fzyzcjy
2026-06-25 08:33:28 +08:00
committed by GitHub
parent 6c839368e0
commit d0524d6433
14 changed files with 107 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/test_forward_split_prefill.py
+3 -2
View File
@@ -96,9 +96,10 @@ class TestForwardSplitPrefill(CustomTestCase):
sampling_params=sampling_params,
)
req.full_untruncated_fill_ids = req.origin_input_ids
req.fill_len = len(req.full_untruncated_fill_ids)
req.logprob_start_len = -1
req.set_extend_input_len(req.fill_len - len(req.prefix_indices))
req.set_extend_range(
len(req.prefix_indices), len(req.full_untruncated_fill_ids)
)
reqs.append(req)
# Create dummy tree_cache for tests (no prefix caching, just allocation)
test/registered/unit/managers/test_hisparse_unit.py
+16 -4
View File
@@ -15,8 +15,20 @@ from types import SimpleNamespace
import torch
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
from sglang.srt.utils.common import Range
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
class _FakeReq(SimpleNamespace):
@property
def fill_len(self) -> int:
return self.extend_range.end
@property
def extend_input_len(self) -> int:
return self.extend_range.length
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
@@ -42,7 +54,7 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None):
origin_input_ids = list(range(64))
if output_ids is None:
output_ids = []
req = SimpleNamespace(
req = _FakeReq(
rid=rid,
origin_input_ids=origin_input_ids,
output_ids=output_ids,
@@ -57,8 +69,8 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None):
inflight_middle_chunks=0,
)
req.finished = lambda: req.finished_reason is not None
req.set_extend_input_len = lambda extend_input_len: setattr(
req, "extend_input_len", extend_input_len
req.set_extend_range = lambda start, end: setattr(
req, "extend_range", Range(start, end)
)
return req
@@ -220,7 +232,7 @@ class TestHiSparseUnit(unittest.TestCase):
req.kv_allocated_len = fill_len
req.kv_committed_len = fill_len
req.full_untruncated_fill_ids = array("q", range(fill_len))
req.fill_len = fill_len
req.extend_range = Range(0, fill_len)
return kv_loc
# ==================================================================
test/registered/unit/managers/test_prefill_adder.py
+6 -16
View File
@@ -76,7 +76,6 @@ class TestPrefillAdder(CustomTestCase):
req = MagicMock(spec=Req)
req.rid = str(rid)
req.priority = priority
req.extend_input_len = 0
req.prefix_indices = []
req.full_untruncated_fill_ids = []
req.extend_logprob_start_len = 0
@@ -385,11 +384,9 @@ class TestPrefillAdder(CustomTestCase):
# Add a prefill that exactly consumes the chunk budget
req1 = self.create_mock_req("req1", priority=0, max_new_tokens=64)
req1.extend_input_len = 56
req1.host_hit_length = 0
req1.prefix_indices = []
req1.full_untruncated_fill_ids = list(range(56))
req1.fill_len = 56
req1.last_node = MagicMock()
req1.sampling_params.ignore_eos = False
@@ -420,11 +417,9 @@ class TestPrefillAdder(CustomTestCase):
# Same prefill no longer exhausts the chunk budget
req2 = self.create_mock_req("req2", priority=0, max_new_tokens=64)
req2.extend_input_len = 56
req2.host_hit_length = 0
req2.prefix_indices = []
req2.full_untruncated_fill_ids = list(range(56))
req2.fill_len = 56
req2.last_node = MagicMock()
req2.sampling_params.ignore_eos = False
@@ -438,11 +433,9 @@ class TestPrefillAdder(CustomTestCase):
# Fit last small prefill request
req3 = self.create_mock_req("req3", priority=0, max_new_tokens=16)
req3.extend_input_len = 3
req3.host_hit_length = 0
req3.prefix_indices = []
req3.full_untruncated_fill_ids = list(range(3))
req3.fill_len = 3
req3.last_node = MagicMock()
req3.sampling_params.ignore_eos = False
@@ -476,11 +469,9 @@ class TestPrefillAdder(CustomTestCase):
adder.is_hybrid_swa = is_hybrid_swa
req = self.create_mock_req("chunked", priority=0, max_new_tokens=128)
req.extend_input_len = extend_input_len
req.prefix_indices = []
req.full_untruncated_fill_ids = list(range(extend_input_len))
req.fill_len = extend_input_len
req.set_extend_input_len = MagicMock()
req.set_extend_range = MagicMock()
return adder, req
def test_add_chunked_req_hybrid_swa_reserves_page_for_alloc_extend(self):
@@ -497,8 +488,9 @@ class TestPrefillAdder(CustomTestCase):
result = adder.add_chunked_req(req)
self.assertIs(result, req) # truncated → chunked prefill continues
req.set_extend_input_len.assert_called_once()
new_len = req.set_extend_input_len.call_args.args[0]
req.set_extend_range.assert_called_once()
start, end = req.set_extend_range.call_args.args
new_len = end - start
self.assertLessEqual(new_len + PAGE_SIZE, REM_SWA)
self.assertEqual(new_len, REM_SWA - PAGE_SIZE)
@@ -510,13 +502,11 @@ class TestPrefillAdder(CustomTestCase):
adder, req = self._build_hybrid_swa_chunked_req(
page_size=PAGE_SIZE, rem_swa=PAGE_SIZE
)
original_len = req.extend_input_len
result = adder.add_chunked_req(req)
self.assertIs(result, req)
req.set_extend_input_len.assert_not_called()
self.assertEqual(req.extend_input_len, original_len)
req.set_extend_range.assert_not_called()
self.assertEqual(len(adder.can_run_list), 0)
def test_swa_budget_for_req(self):
@@ -554,7 +544,7 @@ class TestPrefillAdder(CustomTestCase):
result = adder.add_chunked_req(req)
self.assertIsNone(result)
req.set_extend_input_len.assert_called_once_with(200)
req.set_extend_range.assert_called_once_with(0, 200)
self.assertIn(req, adder.can_run_list)
test/registered/unit/managers/test_scheduler_chunked_req_gate.py
+2 -2
View File
@@ -15,6 +15,7 @@ maybe_stub_sgl_kernel()
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.managers.scheduler import Scheduler
from sglang.srt.mem_cache.chunk_cache import ChunkCache
from sglang.srt.utils.common import Range
register_cpu_ci(est_time=6, suite="base-a-test-cpu")
@@ -32,10 +33,9 @@ def _make_req(
req.origin_input_ids = array("q", fill_ids)
req.output_ids = array("q")
req.full_untruncated_fill_ids = array("q", fill_ids)
req.fill_len = fill_len
req.prefix_indices = prefix_indices
req.req_pool_idx = req_pool_idx
req.extend_input_len = extend_input_len
req.extend_range = Range(fill_len - extend_input_len, fill_len)
req.inflight_middle_chunks = 0
req.host_hit_length = 0
req.cache_protected_len = 0
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+6 -1
View File
@@ -38,6 +38,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchPrefixParams,
)
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey
from sglang.srt.utils.common import Range
def _make_cache_with_pools(page_size=1):
@@ -68,7 +69,7 @@ class MockReq:
def __init__(self, fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None):
self.full_untruncated_fill_ids = array("q", fill_ids)
self.fill_len = len(self.full_untruncated_fill_ids)
self.extend_range = Range(0, len(self.full_untruncated_fill_ids))
self.origin_input_ids = array(
"q", fill_ids[:-1] if len(fill_ids) > 1 else fill_ids
)
@@ -83,6 +84,10 @@ class MockReq:
self.kv_allocated_len = len(fill_ids)
self.kv_committed_freed = False
@property
def fill_len(self):
return self.extend_range.end
def get_fill_ids(self):
return self.full_untruncated_fill_ids[: self.fill_len]
test/registered/unit/mem_cache/test_unified_radix_cache_bench.py
+3 -1
View File
@@ -639,7 +639,9 @@ def bench_cache_finished(
req.origin_input_ids = array("q", seq)
req.output_ids = array("q")
req.full_untruncated_fill_ids = array("q", seq)
req.fill_len = len(req.full_untruncated_fill_ids)
req.set_extend_range(
len(req.prefix_indices), len(req.full_untruncated_fill_ids)
)
req.last_node = node
req.cache_protected_len = matched_len
req.kv_committed_len = len(seq)
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+15 -5
View File
@@ -898,7 +898,9 @@ class UnifiedRadixCacheSuite:
req.swa_uuid_for_lock = None
req.extra_key = None
req.full_untruncated_fill_ids = array("q", input_ids + output_ids)
req.fill_len = len(req.full_untruncated_fill_ids)
req.set_extend_range(
len(req.prefix_indices), len(req.full_untruncated_fill_ids)
)
if self.cfg.has_mamba:
req.mamba_last_track_seqlen = kv_len
@@ -922,7 +924,9 @@ class UnifiedRadixCacheSuite:
req.origin_input_ids = array("q", prompt_ids)
req.output_ids = array("q", output_ids)
req.full_untruncated_fill_ids = array("q", prompt_ids + output_ids)
req.fill_len = len(req.full_untruncated_fill_ids)
req.set_extend_range(
len(req.prefix_indices), len(req.full_untruncated_fill_ids)
)
kv_len = req.fill_len
kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
@@ -977,7 +981,9 @@ class UnifiedRadixCacheSuite:
req.swa_uuid_for_lock = None
req.extra_key = None
req.full_untruncated_fill_ids = array("q", tokens)
req.fill_len = len(req.full_untruncated_fill_ids)
req.set_extend_range(
len(req.prefix_indices), len(req.full_untruncated_fill_ids)
)
avail_before = allocator.available_size()
tree.cache_finished_req(req, is_insert=False)
@@ -995,7 +1001,9 @@ class UnifiedRadixCacheSuite:
req.origin_input_ids = array("q", tokens)
req.output_ids = array("q")
req.full_untruncated_fill_ids = array("q", tokens)
req.fill_len = len(req.full_untruncated_fill_ids)
req.set_extend_range(
len(req.prefix_indices), len(req.full_untruncated_fill_ids)
)
kv_len = len(tokens)
kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
@@ -1130,7 +1138,9 @@ class UnifiedRadixCacheSuite:
req.swa_uuid_for_lock = None
req.extra_key = None
req.full_untruncated_fill_ids = array("q", input_ids)
req.fill_len = len(req.full_untruncated_fill_ids)
req.set_extend_range(
len(req.prefix_indices), len(req.full_untruncated_fill_ids)
)
if self.cfg.has_mamba:
req.mamba_last_track_seqlen = kv_len