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

Refactor Req.fill_ids into full_untruncated_fill_ids + fill_len with equivalence (#26637)

Browse Source
This commit is contained in:
fzyzcjy
2026-06-08 14:52:18 +08:00
committed by GitHub
parent 4201de11de
commit 259a2da3e0
29 changed files with 136 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/test_forward_split_prefill.py
+3 -2
View File
@@ -95,9 +95,10 @@ class TestForwardSplitPrefill(CustomTestCase):
origin_input_ids=array("q", input_ids[i]),
sampling_params=sampling_params,
)
req.fill_ids = req.origin_input_ids
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(len(req.fill_ids) - len(req.prefix_indices))
req.set_extend_input_len(req.fill_len - len(req.prefix_indices))
reqs.append(req)
# Create dummy tree_cache for tests (no prefix caching, just allocation)
test/registered/input_embedding/test_input_embeds_chunked.py
+6 -6
View File
@@ -3,14 +3,14 @@
Covers two bugs with the same crash signature
(RuntimeError: shape mismatch in set_kv_buffer) but opposite polarity:
- Chunked prefill truncation (#20376): PrefillAdder truncates fill_ids and
- Chunked prefill truncation (#20376): PrefillAdder shrinks fill_len and
extend_input_len on chunk overflow but not input_embeds, so the full array
flows through while out_cache_loc is sized for the truncated length.
Polarity: cache_k > loc.
- Retraction with output_ids (#14110): after retraction, fill_ids includes
accumulated output_ids but input_embeds only covers origin_input_ids.
Polarity: cache_k < loc.
- Retraction with output_ids (#14110): after retraction, get_fill_ids()
includes accumulated output_ids but input_embeds only covers
origin_input_ids. Polarity: cache_k < loc.
"""
import unittest
@@ -162,8 +162,8 @@ class TestInputEmbedsChunkedAndRetract(CustomTestCase):
SGLANG_TEST_RETRACT forces retraction every few scheduler iterations.
Combined with ignore_eos and a reasonable max_new_tokens, at least one
request is retracted mid-decode with non-empty output_ids, then
re-prefilled. Pre-#14110 this crashes (cache_k < loc) because fill_ids
includes output_ids but input_embeds does not.
re-prefilled. Pre-#14110 this crashes (cache_k < loc) because the
filled token sequence includes output_ids but input_embeds does not.
"""
text = "The quick brown fox jumps over the lazy dog. " * 4
embeds = _embeds_for(text)
test/registered/unit/managers/test_hisparse_unit.py
+3 -1
View File
@@ -9,6 +9,7 @@ Tests cover:
import os
import unittest
from array import array
from types import SimpleNamespace
import torch
@@ -210,7 +211,8 @@ class TestHiSparseUnit(unittest.TestCase):
self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc)
req.kv_allocated_len = fill_len
req.kv_committed_len = fill_len
req.fill_ids = list(range(fill_len))
req.full_untruncated_fill_ids = array("q", range(fill_len))
req.fill_len = fill_len
return kv_loc
# ==================================================================
test/registered/unit/managers/test_prefill_adder.py
+8 -4
View File
@@ -386,7 +386,8 @@ class TestPrefillAdder(CustomTestCase):
req1.extend_input_len = 56
req1.host_hit_length = 0
req1.prefix_indices = []
req1.fill_ids = list(range(56))
req1.full_untruncated_fill_ids = list(range(56))
req1.fill_len = 56
req1.last_node = MagicMock()
req1.sampling_params.ignore_eos = False
@@ -420,7 +421,8 @@ class TestPrefillAdder(CustomTestCase):
req2.extend_input_len = 56
req2.host_hit_length = 0
req2.prefix_indices = []
req2.fill_ids = list(range(56))
req2.full_untruncated_fill_ids = list(range(56))
req2.fill_len = 56
req2.last_node = MagicMock()
req2.sampling_params.ignore_eos = False
@@ -437,7 +439,8 @@ class TestPrefillAdder(CustomTestCase):
req3.extend_input_len = 3
req3.host_hit_length = 0
req3.prefix_indices = []
req3.fill_ids = list(range(3))
req3.full_untruncated_fill_ids = list(range(3))
req3.fill_len = 3
req3.last_node = MagicMock()
req3.sampling_params.ignore_eos = False
@@ -473,7 +476,8 @@ class TestPrefillAdder(CustomTestCase):
req = self.create_mock_req("chunked", priority=0, max_new_tokens=128)
req.extend_input_len = extend_input_len
req.prefix_indices = []
req.fill_ids = list(range(extend_input_len))
req.full_untruncated_fill_ids = list(range(extend_input_len))
req.fill_len = extend_input_len
req.set_extend_input_len = MagicMock()
return adder, req
test/registered/unit/managers/test_scheduler_chunked_req_gate.py
+2 -1
View File
@@ -30,7 +30,8 @@ def _make_req(
req.rid = "test-req"
req.origin_input_ids = array("q", fill_ids)
req.output_ids = array("q")
req.fill_ids = array("q", fill_ids)
req.full_untruncated_fill_ids = array("q", fill_ids)
req.fill_len = len(req.full_untruncated_fill_ids)
req.prefix_indices = prefix_indices
req.req_pool_idx = req_pool_idx
req.extend_input_len = extend_input_len
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+5 -1
View File
@@ -67,7 +67,8 @@ class MockReq:
"""Minimal mock Req with fields needed by cache_unfinished/finished_req."""
def __init__(self, fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None):
self.fill_ids = array("q", fill_ids)
self.full_untruncated_fill_ids = array("q", fill_ids)
self.fill_len = len(self.full_untruncated_fill_ids)
self.origin_input_ids = array(
"q", fill_ids[:-1] if len(fill_ids) > 1 else fill_ids
)
@@ -82,6 +83,9 @@ class MockReq:
self.kv_allocated_len = len(fill_ids)
self.kv_committed_freed = False
def get_fill_ids(self):
return self.full_untruncated_fill_ids[: self.fill_len]
def pop_committed_kv_cache(self):
self.kv_committed_freed = True
return self.kv_committed_len
test/registered/unit/mem_cache/test_unified_radix_cache_bench.py
+2 -1
View File
@@ -637,7 +637,8 @@ def bench_cache_finished(
req = env.make_req()
req.origin_input_ids = array("q", seq)
req.output_ids = array("q")
req.fill_ids = array("q", seq)
req.full_untruncated_fill_ids = array("q", seq)
req.fill_len = 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
+11 -6
View File
@@ -798,7 +798,8 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
req.extra_key = None
req.fill_ids = array("q", input_ids + output_ids)
req.full_untruncated_fill_ids = array("q", input_ids + output_ids)
req.fill_len = len(req.full_untruncated_fill_ids)
if self.cfg.has_mamba:
req.mamba_last_track_seqlen = kv_len
@@ -821,8 +822,9 @@ class UnifiedRadixCacheSuite:
output_ids = self._make_seq(2000, 7)
req.origin_input_ids = array("q", prompt_ids)
req.output_ids = array("q", output_ids)
req.fill_ids = array("q", prompt_ids + output_ids)
kv_len = len(req.fill_ids)
req.full_untruncated_fill_ids = array("q", prompt_ids + output_ids)
req.fill_len = 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)
req.kv_committed_len = kv_len
@@ -875,7 +877,8 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
req.extra_key = None
req.fill_ids = array("q", tokens)
req.full_untruncated_fill_ids = array("q", tokens)
req.fill_len = len(req.full_untruncated_fill_ids)
avail_before = allocator.available_size()
tree.cache_finished_req(req, is_insert=False)
@@ -892,7 +895,8 @@ class UnifiedRadixCacheSuite:
tokens = self._make_seq(1, 3)
req.origin_input_ids = array("q", tokens)
req.output_ids = array("q")
req.fill_ids = array("q", tokens)
req.full_untruncated_fill_ids = array("q", tokens)
req.fill_len = 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)
@@ -1026,7 +1030,8 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
req.extra_key = None
req.fill_ids = array("q", input_ids)
req.full_untruncated_fill_ids = array("q", input_ids)
req.fill_len = len(req.full_untruncated_fill_ids)
if self.cfg.has_mamba:
req.mamba_last_track_seqlen = kv_len
test/registered/unit/utils/test_common.py
+2 -2
View File
@@ -14,8 +14,8 @@ register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-small")
class TestFlattenArraysToInt64Tensor(CustomTestCase):
"""`flatten_arrays_to_int64_tensor` is invoked by `prepare_for_extend`
to build the per-batch input_ids tensor (pinned, async H2D) from a
list of array.array('q') per-req fill_ids slices. Tests the full
matrix of (device, pin) the production code paths through.
list of array.array('q') per-req get_fill_ids() slices. Tests the
full matrix of (device, pin) the production code paths through.
"""
DEVICES = ("cpu", "cuda")