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

[perf] reduce overhead of fill_ids list reconstruction and decref (#27965)

Browse Source
This commit is contained in:
Qiaolin Yu
2026-06-14 00:41:11 -07:00
committed by GitHub
parent f2d7d67603
commit f293ddf3ce
8 changed files with 363 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+37 -9
View File
@@ -695,11 +695,14 @@ class Req(ReqDllmMixin):
if origin_input_ids_unpadded
else self.origin_input_ids
) # Before image padding
# Each decode stage's output ids
# Each decode stage's output ids. Append-only by contract:
# _refresh_fill_ids infers how many output tokens are already in
# full_untruncated_fill_ids from lengths alone, so in-place rewrites
# that preserve length would silently corrupt fill_ids.
self.output_ids = array("q")
# Full untruncated sequence: origin + output (+ DLLM mask block).
# Rebuilt at the top of each init_next_round_input; admission only
# updates fill_len, never mutates this array's length.
# Kept in sync by _refresh_fill_ids; admission only updates fill_len,
# never mutates this array's length.
self.full_untruncated_fill_ids = array("q")
self.fill_len: int = 0
@@ -1069,6 +1072,27 @@ class Req(ReqDllmMixin):
def get_fill_ids(self) -> array:
return self.full_untruncated_fill_ids[: self.fill_len]
def _refresh_fill_ids(self) -> None:
"""Keep full_untruncated_fill_ids == origin_input_ids + output_ids by
appending only the new output tokens.
Falls back to a full rebuild when the in-place append is invalid:
- aliasing: scheduler_pp_mixin assigns full_untruncated_fill_ids =
origin_input_ids directly, so extending in place would write output
tokens into the origin;
- lengths disagree: fresh req (array still empty), retraction
(output_ids reset to empty), or set_finish_with_abort (origin
replaced by a 1-token stub).
"""
n_have_output = len(self.full_untruncated_fill_ids) - len(self.origin_input_ids)
if (
self.full_untruncated_fill_ids is not self.origin_input_ids
and 0 <= n_have_output <= len(self.output_ids)
):
self.full_untruncated_fill_ids.extend(self.output_ids[n_have_output:])
else:
self.full_untruncated_fill_ids = self.origin_input_ids + self.output_ids
def init_next_round_input(
self,
tree_cache: Optional[BasePrefixCache] = None,
@@ -1078,7 +1102,7 @@ class Req(ReqDllmMixin):
self._init_fill_ids_for_dllm()
self.determine_dllm_phase()
else:
self.full_untruncated_fill_ids = self.origin_input_ids + self.output_ids
self._refresh_fill_ids()
input_len = len(self.full_untruncated_fill_ids)
@@ -1098,14 +1122,16 @@ class Req(ReqDllmMixin):
)
self.logprob_start_len = -1
token_ids_to_match = self.full_untruncated_fill_ids[
: self._compute_max_prefix_len(input_len)
]
# Pass the full array with a raw-token cap (limit) instead of slicing,
# avoiding an O(context) copy per prefill-batch build.
token_ids_to_match = self.full_untruncated_fill_ids
key_limit: Optional[int] = self._compute_max_prefix_len(input_len)
# Disable prefix caching when embed overrides are present: same token IDs
# with different override vectors must not share cached KV values.
if self.positional_embed_overrides is not None:
token_ids_to_match = array("q")
key_limit = None
if tree_cache is not None:
if cow_mamba is None:
@@ -1113,7 +1139,9 @@ class Req(ReqDllmMixin):
match_result = tree_cache.match_prefix(
MatchPrefixParams(
key=RadixKey(
token_ids=token_ids_to_match, extra_key=self.extra_key
token_ids=token_ids_to_match,
extra_key=self.extra_key,
limit=key_limit,
),
req=self,
cow_mamba=cow_mamba,
@@ -2314,7 +2342,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
running_bs = running_batch.batch_size()
for req in running_batch.reqs:
req.full_untruncated_fill_ids = req.origin_input_ids + req.output_ids
req._refresh_fill_ids()
req.fill_len = len(req.full_untruncated_fill_ids)
req.set_extend_input_len(1)
python/sglang/srt/mem_cache/mamba_radix_cache.py
+1 -1
View File
@@ -507,7 +507,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
prev_prefix_len = params.prev_prefix_len
if value is None:
value = torch.tensor([x for x in key.token_ids], dtype=torch.int64)
value = torch.tensor([x for x in key.raw_token_ids()], dtype=torch.int64)
prefix_len, mamba_exist = self._insert_helper(
self.root_node, key, value, mamba_value, params.chunked, prev_prefix_len
)
python/sglang/srt/mem_cache/radix_cache.py
+27 -6
View File
@@ -56,13 +56,14 @@ if TYPE_CHECKING:
class RadixKey:
"""is_bigram=True: token_ids holds raw tokens (N+1 for N bigrams); slices share one boundary token."""
__slots__ = ("token_ids", "extra_key", "is_bigram")
__slots__ = ("token_ids", "extra_key", "is_bigram", "limit")
def __init__(
self,
token_ids: array[int],
extra_key: Optional[str] = None,
is_bigram: bool = False,
limit: Optional[int] = None,
):
# token ids sequence (raw ints in both modes)
self.token_ids = token_ids
@@ -70,21 +71,40 @@ class RadixKey:
self.extra_key = extra_key
# bigram view over token_ids: length = max(0, len(token_ids) - 1)
self.is_bigram = is_bigram
# Optional cap on raw tokens: behave as if token_ids were sliced to
# token_ids[:limit], without the O(n) copy. None = use all tokens.
self.limit = limit
def _raw_len(self) -> int:
n = len(self.token_ids)
if self.limit is not None and self.limit < n:
return self.limit
return n
def raw_token_ids(self) -> array:
"""token_ids honoring `limit` (copies only when capped)."""
n = self._raw_len()
t = self.token_ids
return t if n == len(t) else t[:n]
def __len__(self) -> int:
n = self._raw_len()
if self.is_bigram:
n = len(self.token_ids)
return n - 1 if n > 0 else 0
return len(self.token_ids)
return n
# TODO(Jialin): vectorize with numpy without PyLong boxing
def __iter__(self) -> Iterator:
if self.is_bigram:
t = self.token_ids
for i in range(len(t) - 1):
n = self._raw_len()
if self.is_bigram:
for i in range(n - 1 if n > 0 else 0):
yield (t[i], t[i + 1])
elif n == len(t):
yield from t
else:
yield from self.token_ids
for i in range(n):
yield t[i]
def __getitem__(self, idx: Union[int, slice]) -> RadixKey:
# Normalize int -> 1-element slice so the rest handles one shape.
@@ -166,6 +186,7 @@ class RadixKey:
matched = max(0, min(matched_tokens - 1, len(self), len(other)))
return (matched // page_size) * page_size if page_size > 1 else matched
matched_tokens = min(matched_tokens, len(self), len(other))
if page_size == 1:
return matched_tokens
return (matched_tokens // page_size) * page_size
python/sglang/srt/mem_cache/radix_cache_cpp.py
+1 -1
View File
@@ -101,7 +101,7 @@ class RadixCacheCpp(BasePrefixCache):
def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
key = params.key
device_indices_vec, host_indices_length, node_gpu, node_cpu = (
self.tree.match_prefix(key.token_ids)
self.tree.match_prefix(key.raw_token_ids())
)
return MatchResult(
device_indices=self._merge_tensor(device_indices_vec),
python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py
+9 -4
View File
@@ -46,7 +46,8 @@ class _LMCacheLoadBackMarker:
``match_prefix`` call in MP mode.
"""
key: RadixKey # page-aligned key the scheduler matched on
key: RadixKey # detached snapshot of the matched key (the live query key
# aliases the req's growing fill_ids and must not be retained)
value_numel: int # number of tokens already in radix at match time
@@ -216,14 +217,17 @@ class LMCRadixCache(RadixCache):
LMCache has tokens beyond radix. Otherwise releases
the held read locks and returns the radix-only result.
"""
matched = self.lmcache_connector.lookup_kv(key.token_ids, req.rid)
token_ids = key.raw_token_ids()
matched = self.lmcache_connector.lookup_kv(token_ids, req.rid)
if matched <= value.numel():
# Release the read locks; keep the pending session for end_session.
self.lmcache_connector.release_pending(req.rid)
return base_res
if token_ids is key.token_ids:
token_ids = token_ids[:]
self._mp_load_back_markers[req.rid] = _LMCacheLoadBackMarker(
key=key,
key=RadixKey(token_ids, key.extra_key, key.is_bigram),
value_numel=int(value.numel()),
)
return MatchResult(
@@ -254,13 +258,14 @@ class LMCRadixCache(RadixCache):
if uncached_len == 0:
return base_res
token_ids = key.raw_token_ids()
result = self._load_back(
key=key,
value_numel=int(value.numel()),
uncached_len=uncached_len,
last_node=last_node,
load_fn=lambda sm, pp: self._ip_load_back(
token_ids=key.token_ids,
token_ids=token_ids,
value_numel=int(value.numel()),
slot_mapping=sm,
prefix_pad=pp,
python/sglang/srt/session/session_controller.py
+124 -44
View File
@@ -15,6 +15,7 @@ from __future__ import annotations
import logging
import time
import uuid
from array import array
from typing import TYPE_CHECKING, Dict, Optional
from sglang.srt.managers.io_struct import (
@@ -94,12 +95,111 @@ class Session:
self.req_nodes: Dict[str, SessionReqNode] = {}
self.close_on_finish: bool = False
self._inflight: bool = False
# Token-array lengths of last_req as of its finish_req. The share path
# appends speculatively beyond these; only finish_req confirms them, so
# _share_token_arrays trims back first (heals aborted turns).
self.committed_origin_len: Optional[int] = None
self.committed_unpadded_len: Optional[int] = None
self.committed_fill_len: Optional[int] = None
def is_timed_out(self) -> bool:
if self.timeout is None:
return False
return time.monotonic() - self.last_active_time > self.timeout
@staticmethod
def _strip_bos_token(req: TokenizedGenerateReqInput, tokenizer) -> None:
"""Trim a leading BOS on an appended turn; shift mm offsets to match."""
if not (
tokenizer is not None
and req.input_ids
and req.input_ids[0] == tokenizer.bos_token_id
):
return
req.input_ids = req.input_ids[1:]
if req.mm_inputs:
for item in req.mm_inputs.mm_items:
if item.offsets:
if any(s == 0 for s, _ in item.offsets):
logging.warning(
"mm_item offset starts at 0 (BOS position), "
"clamping to 0 after BOS strip"
)
item.offsets = [
(max(0, s - 1), max(0, e - 1)) for s, e in item.offsets
]
def _share_token_arrays(self, last_req: Req, new_input_ids):
"""Plain streaming append: reuse last_req's token arrays in place.
Trims each array back to its committed length first — an earlier turn
may have appended its tokens and then aborted before finish_req, and
req_nodes still points at last_req, so anything beyond the committed
lengths is unconfirmed. Then extends with last turn's output and the
new input. Returns (input_ids, input_ids_unpadded, carry_fill);
carry_fill (== the new origin) spares the first fill_ids rebuild.
"""
out_tail = last_req.output_ids[: last_req.sampling_params.max_new_tokens]
input_ids = last_req.origin_input_ids
del input_ids[self.committed_origin_len :]
if last_req.origin_input_ids_unpadded is input_ids:
input_ids_unpadded = input_ids
else:
input_ids_unpadded = last_req.origin_input_ids_unpadded
del input_ids_unpadded[self.committed_unpadded_len :]
carry_fill = last_req.full_untruncated_fill_ids
if (
not isinstance(carry_fill, array)
or carry_fill is input_ids
or carry_fill is input_ids_unpadded
):
# Unexpected type or aliased with an origin array (extending it
# below would double-append): let _refresh_fill_ids rebuild.
carry_fill = None
else:
del carry_fill[self.committed_fill_len :]
baked = len(carry_fill) - len(input_ids)
if 0 <= baked <= len(out_tail):
carry_fill.extend(out_tail[baked:])
carry_fill.extend(new_input_ids)
else:
carry_fill = None
input_ids.extend(out_tail)
input_ids.extend(new_input_ids)
if input_ids_unpadded is not input_ids:
input_ids_unpadded.extend(out_tail)
input_ids_unpadded.extend(new_input_ids)
return input_ids, input_ids_unpadded, carry_fill
@staticmethod
def _concat_token_arrays(
last_req: Req, req: TokenizedGenerateReqInput, session_params
):
"""Copy-based assembly for replace/offset/drop_previous_output turns."""
out_tail = last_req.output_ids[: last_req.sampling_params.max_new_tokens]
input_ids = last_req.origin_input_ids + out_tail
if session_params.drop_previous_output:
input_ids = last_req.origin_input_ids[:]
if session_params.offset and session_params.offset != 0:
input_ids = input_ids[: session_params.offset] + req.input_ids
else:
input_ids += req.input_ids
input_ids_unpadded = last_req.origin_input_ids_unpadded + out_tail
if session_params.drop_previous_output:
input_ids_unpadded = last_req.origin_input_ids_unpadded[:]
if session_params.offset and session_params.offset != 0:
input_ids_unpadded = (
input_ids_unpadded[: session_params.offset] + req.input_ids
)
else:
input_ids_unpadded += req.input_ids
return input_ids, input_ids_unpadded
def create_req(
self,
req: TokenizedGenerateReqInput,
@@ -163,54 +263,28 @@ class Session:
abort_message = "Session request is appending to a request that hasn't finished."
logging.warning(abort_message)
carry_fill = None
if last_req is not None:
# trim bos token if it is an append
if (
tokenizer is not None
and req.input_ids
and req.input_ids[0] == tokenizer.bos_token_id
):
req.input_ids = req.input_ids[1:]
# Adjust mm_item offsets since they were computed on
# the pre-strip sequence (with BOS at position 0)
if req.mm_inputs:
for item in req.mm_inputs.mm_items:
if item.offsets:
if any(s == 0 for s, _ in item.offsets):
logging.warning(
"mm_item offset starts at 0 (BOS position), "
"clamping to 0 after BOS strip"
self._strip_bos_token(req, tokenizer)
# In-place sharing is only safe for the plain streaming append:
# streaming sessions allow a single inflight request, last_req has
# finished, and the committed_* lengths recorded by finish_req let
# _share_token_arrays trim away tokens appended by an aborted turn.
# offset / drop_previous_output rewrite history and must copy.
can_share_token_arrays = (
self.streaming
and self.committed_origin_len is not None
and not session_params.drop_previous_output
and not (session_params.offset and session_params.offset != 0)
)
item.offsets = [
(max(0, s - 1), max(0, e - 1)) for s, e in item.offsets
]
input_ids = (
last_req.origin_input_ids
+ last_req.output_ids[: last_req.sampling_params.max_new_tokens]
)
if session_params.drop_previous_output:
input_ids = last_req.origin_input_ids[:]
if session_params.offset and session_params.offset != 0:
input_ids = input_ids[: session_params.offset] + req.input_ids
else:
input_ids += req.input_ids
input_ids_unpadded = (
last_req.origin_input_ids_unpadded
+ last_req.output_ids[: last_req.sampling_params.max_new_tokens]
)
if session_params.drop_previous_output:
input_ids_unpadded = last_req.origin_input_ids_unpadded[:]
if session_params.offset and session_params.offset != 0:
input_ids_unpadded = (
input_ids_unpadded[: session_params.offset] + req.input_ids
if can_share_token_arrays:
input_ids, input_ids_unpadded, carry_fill = self._share_token_arrays(
last_req, req.input_ids
)
else:
input_ids_unpadded += req.input_ids
input_ids, input_ids_unpadded = self._concat_token_arrays(
last_req, req, session_params
)
else:
input_ids = req.input_ids
input_ids_unpadded = req.input_ids
@@ -243,6 +317,8 @@ class Session:
if last_req is not None:
new_req.multimodal_inputs = last_req.multimodal_inputs
new_req.tokenizer = tokenizer
if carry_fill is not None:
new_req.full_untruncated_fill_ids = carry_fill
if abort:
new_req.set_finish_with_abort(abort_message)
@@ -263,6 +339,10 @@ class Session:
prev_node.req.session = None
self.req_nodes.clear()
self.req_nodes[req.rid] = SessionReqNode(req)
# Confirm this req's token arrays as the session's rollback point.
self.committed_origin_len = len(req.origin_input_ids)
self.committed_unpadded_len = len(req.origin_input_ids_unpadded)
self.committed_fill_len = len(req.full_untruncated_fill_ids)
def abort_req(self):
"""Clear inflight flag on abort (req_nodes stays unchanged)."""
python/sglang/srt/session/streaming_session.py
+1 -1
View File
@@ -247,7 +247,7 @@ class StreamingSession(BasePrefixCache):
# token_ids = get_fill_ids()[:input_len-1] (1-token logit reserve
# already applied). min handles retract retry where committed_len
# can exceed len(token_ids) by 1.
prefix_len = min(req.kv_committed_len, len(params.key.token_ids))
prefix_len = min(req.kv_committed_len, len(params.key))
# Streaming sessions are append-only (session_controller rollback
# ensures req_nodes always points to the last successful req).
test/registered/unit/mem_cache/test_session_token_share_unit.py
+163
View File
@@ -0,0 +1,163 @@
"""Unit tests for the streaming-session in-place token-array share protocol
(`Session.create_req` / `finish_req` / `abort_req`):
- token arrays are extended in place and shared across turns (no per-turn copy);
- committed_* lengths recorded at finish_req trim away tokens appended by a
turn that aborted before finishing (mid-turn and first-turn aborts);
- max_new_tokens overshoot falls back to a fill_ids rebuild instead of
carrying an inconsistent array.
"""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=15, suite="base-a-test-cpu")
import unittest
from array import array
from types import SimpleNamespace
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.session.session_controller import Session
from sglang.test.test_utils import CustomTestCase
VOCAB = 1 << 20
def _recv(rid, input_ids, max_new_tokens=8):
return SimpleNamespace(
rid=rid,
input_ids=array("q", input_ids),
mm_inputs=None,
session_params=SimpleNamespace(
id="s", rid=None, offset=None, replace=False, drop_previous_output=False
),
sampling_params=SamplingParams(max_new_tokens=max_new_tokens),
lora_id=None,
custom_logit_processor=None,
stream=False,
return_logprob=False,
top_logprobs_num=0,
token_ids_logprob=None,
require_reasoning=False,
return_hidden_states=False,
return_routed_experts=False,
routed_experts_start_len=0,
priority=None,
routing_key=None,
extra_key=None,
http_worker_ipc=None,
time_stats=None,
)
class TestSessionTokenShare(CustomTestCase):
def setUp(self):
self.session = Session(capacity_of_str_len=0, session_id="s", streaming=True)
def _create(self, rid, input_ids, max_new_tokens=8):
return self.session.create_req(
_recv(rid, input_ids, max_new_tokens=max_new_tokens),
tokenizer=None,
vocab_size=VOCAB,
)
def _decode_and_finish(self, req, output, baked=None):
"""Simulate decode then a successful finish.
`baked` output tokens are folded into the fill array before the rest
arrive (mix_with_running refreshes mid-decode, so the bake is often
partial).
"""
if baked is None:
baked = len(output)
req.output_ids.extend(output[:baked])
req._refresh_fill_ids()
req.output_ids.extend(output[baked:])
self.session.finish_req(req)
def test_normal_multi_turn_share_and_carry(self):
in1, out1 = list(range(100, 110)), [1, 2, 3]
r1 = self._create("r1", in1)
self.assertEqual(list(r1.origin_input_ids), in1)
self._decode_and_finish(r1, out1, baked=2) # partial bake
self.assertEqual(self.session.committed_origin_len, len(in1))
self.assertEqual(self.session.committed_fill_len, len(in1) + 2)
in2, out2 = [7, 8], [4, 5]
r2 = self._create("r2", in2)
# In-place share: same objects, extended to the new prompt.
self.assertIs(r2.origin_input_ids, r1.origin_input_ids)
self.assertEqual(list(r2.origin_input_ids), in1 + out1 + in2)
# Carry: the fill array handed over and equal to the new origin.
self.assertIs(r2.full_untruncated_fill_ids, r1.full_untruncated_fill_ids)
self.assertEqual(list(r2.full_untruncated_fill_ids), list(r2.origin_input_ids))
self._decode_and_finish(r2, out2)
r3 = self._create("r3", [9])
self.assertEqual(list(r3.origin_input_ids), in1 + out1 + in2 + out2 + [9])
self.assertEqual(list(r3.full_untruncated_fill_ids), list(r3.origin_input_ids))
def test_mid_turn_abort_then_continue(self):
in1, out1 = list(range(200, 210)), [1, 2, 3]
r1 = self._create("r1", in1)
self._decode_and_finish(r1, out1)
# Turn 2 extends the shared arrays, decodes a bit, then aborts:
# finish_req never runs, req_nodes still points at r1.
r2 = self._create("r2", [50, 51])
self.assertEqual(list(r2.origin_input_ids), in1 + out1 + [50, 51])
r2.output_ids.extend([6, 7])
r2._refresh_fill_ids()
self.session.abort_req()
self.assertEqual(self.session.committed_origin_len, len(in1))
# Turn 3 must see exactly r1's history — no [50, 51], no doubled out1.
r3 = self._create("r3", [60])
self.assertEqual(list(r3.origin_input_ids), in1 + out1 + [60])
self.assertEqual(list(r3.full_untruncated_fill_ids), list(r3.origin_input_ids))
# Two aborted attempts in a row heal idempotently.
self.session.abort_req()
r4 = self._create("r4", [70])
self.assertEqual(list(r4.origin_input_ids), in1 + out1 + [70])
self.assertEqual(list(r4.full_untruncated_fill_ids), list(r4.origin_input_ids))
def test_first_turn_abort(self):
self._create("r1", [1, 2, 3])
self.assertTrue(self.session._inflight)
self.session.abort_req()
self.assertFalse(self.session._inflight)
# No finish_req ran: nothing committed, next turn starts from scratch.
self.assertIsNone(self.session.committed_origin_len)
r2 = self._create("r2", [4, 5])
self.assertEqual(list(r2.origin_input_ids), [4, 5])
self._decode_and_finish(r2, [9])
r3 = self._create("r3", [6])
self.assertEqual(list(r3.origin_input_ids), [4, 5, 9, 6])
def test_max_new_tokens_overshoot_falls_back(self):
in1 = list(range(300, 310))
r1 = self._create("r1", in1, max_new_tokens=4)
# Spec-decode overshoot: 6 tokens decoded and baked into the fill
# array, then output trimmed to finished_len (like _trim_overshoot)
# before finish.
r1.output_ids.extend([1, 2, 3, 4, 5, 6])
r1._refresh_fill_ids()
del r1.output_ids[4:]
self.session.finish_req(r1)
self.assertEqual(
self.session.committed_fill_len, len(in1) + 6
) # fill kept the overshoot
# Next turn: out_tail is output[:max_new]; the carried fill has more
# baked than out_tail, so the carry is dropped and the fill rebuilds.
r2 = self._create("r2", [50])
self.assertEqual(list(r2.origin_input_ids), in1 + [1, 2, 3, 4] + [50])
self.assertEqual(len(r2.full_untruncated_fill_ids), 0) # carry skipped
r2._refresh_fill_ids()
self.assertEqual(list(r2.full_untruncated_fill_ids), list(r2.origin_input_ids))
if __name__ == "__main__":
unittest.main()