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

perf: migrate Req token-id storage to array.array('q') in Scheduler (#25098)

Browse Source 
Co-authored-by: jialino <jialino@fb.com>
This commit is contained in:
Jialin Ouyang
2026-05-22 10:51:07 -07:00
committed by GitHub
co-authored by jialino
parent 5e9bd21979
commit 06c23d55b5
34 changed files with 833 additions and 299 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/detokenizer_manager.py
+1 -1
View File
@@ -238,7 +238,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
if rid not in self.decode_status:
s = DecodeStatus(
decoded_text=recv_obj.decoded_texts[i],
decode_ids=recv_obj.decode_ids[i],
decode_ids=list(recv_obj.decode_ids[i]),
surr_offset=0,
read_offset=recv_obj.read_offsets[i],
)
python/sglang/srt/managers/io_struct.py
+5 -4
View File
@@ -21,6 +21,7 @@ from __future__ import annotations
import copy
import uuid
from abc import ABC
from array import array
from collections import Counter
from dataclasses import dataclass, field
from enum import Enum
@@ -712,7 +713,7 @@ class TokenizedGenerateReqInput(BaseReq):
# The input text
input_text: str
# The input token ids
input_ids: List[int]
input_ids: Optional[array[int]]
# The multimodal inputs
mm_inputs: object
# The sampling parameters
@@ -1027,7 +1028,7 @@ class TokenizedEmbeddingReqInput(BaseReq):
# The input text
input_text: str
# The input token ids
input_ids: List[int]
input_ids: array[int]
# The image inputs
image_inputs: dict
# The token type ids
@@ -1075,10 +1076,10 @@ class BatchTokenIDOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin):
finished_reasons: List[BaseFinishReason]
# For incremental decoding
decoded_texts: List[str]
decode_ids: List[int]
decode_ids: List[array[int]]
read_offsets: List[int]
# Only used when `--skip-tokenizer-init` is on
output_ids: Optional[List[int]]
output_ids: Optional[List[array[int]]]
# Detokenization configs
skip_special_tokens: List[bool]
spaces_between_special_tokens: List[bool]
python/sglang/srt/managers/schedule_batch.py
+22 -17
View File
@@ -2,7 +2,11 @@ from __future__ import annotations
from sglang.srt.dllm.config import DllmConfig
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.utils.common import ceil_align, is_pin_memory_available
from sglang.srt.utils.common import (
ceil_align,
flatten_arrays_to_int64_tensor,
is_pin_memory_available,
)
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -35,11 +39,11 @@ import copy
import dataclasses
import logging
import re
from array import array
from concurrent.futures import Future
from enum import Enum, auto
from functools import lru_cache
from http import HTTPStatus
from itertools import chain
from typing import (
TYPE_CHECKING,
Any,
@@ -611,14 +615,14 @@ class Req(ReqDllmMixin):
self,
rid: str,
origin_input_text: str,
origin_input_ids: List[int],
origin_input_ids: array[int],
sampling_params: SamplingParams,
return_logprob: bool = False,
top_logprobs_num: int = 0,
dllm_config: Optional[DllmConfig] = None,
token_ids_logprob: List[int] = None,
stream: bool = False,
origin_input_ids_unpadded: Optional[Tuple[int]] = None,
origin_input_ids_unpadded: Optional[array[int]] = None,
lora_id: Optional[str] = None,
input_embeds: Optional[List[List[float]]] = None,
positional_embed_overrides: Optional[PositionalEmbeds] = None,
@@ -659,9 +663,10 @@ class Req(ReqDllmMixin):
)
self.origin_input_ids = origin_input_ids
# Each decode stage's output ids
self.output_ids = []
self.output_ids = array("q")
# fill_ids = origin_input_ids + output_ids. Updated if chunked.
self.fill_ids = []
self.fill_ids = array("q")
self.session = session
self.input_embeds = input_embeds
self.positional_embed_overrides = positional_embed_overrides
@@ -948,7 +953,7 @@ class Req(ReqDllmMixin):
return self.sampling_params.max_new_tokens == 0 and spec_alg is None
@property
def output_ids_through_stop(self) -> List[int]:
def output_ids_through_stop(self) -> array[int]:
"""Get the output ids through the stop condition. Stop position is included."""
if self.finished_len is not None:
return self.output_ids[: self.finished_len]
@@ -1037,7 +1042,7 @@ class Req(ReqDllmMixin):
# 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 = []
token_ids_to_match = array("q")
if tree_cache is not None:
if cow_mamba is None:
@@ -1305,7 +1310,7 @@ class Req(ReqDllmMixin):
# Therefore, we discard the generated output_ids and restart prefill and generation
# to ensure shape consistency in KV cache.
if self.input_embeds is not None:
self.output_ids = []
self.output_ids = array("q")
def offload_kv_cache(self, req_to_token_pool, token_to_kv_pool_allocator):
token_indices = req_to_token_pool.req_to_token[
@@ -1370,7 +1375,9 @@ class Req(ReqDllmMixin):
logger.error(f"{error_msg}, {self.rid=}")
self.multimodal_inputs = None
self.grammar = None
self.origin_input_ids = [0] # set it to one token to skip the long prefill
self.origin_input_ids = array(
"q", [0]
) # set it to one token to skip the long prefill
self.return_logprob = False
self.logprob_start_len = -1
self.to_finish = FINISH_ABORT(
@@ -1629,7 +1636,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
def is_dllm(self):
return self.dllm_config is not None
def prepare_encoder_info_extend(self, input_ids: List[int], seq_lens: List[int]):
def prepare_encoder_info_extend(
self, input_ids: List[array[int]], seq_lens: List[int]
):
_pin = is_pin_memory_available(self.device)
self.encoder_lens_cpu = []
self.encoder_cached = []
@@ -1678,9 +1687,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
pt += req.extend_input_len
# Reassign
self.input_ids = torch.tensor(
sum(input_ids, []), dtype=torch.int64, pin_memory=_pin
).to(self.device, non_blocking=True)
self.input_ids = flatten_arrays_to_int64_tensor(input_ids, self.device, _pin)
self.seq_lens = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=_pin).to(
self.device, non_blocking=True
)
@@ -1783,9 +1790,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
]
_pin = is_pin_memory_available(self.device)
input_ids_tensor = torch.tensor(
list(chain.from_iterable(input_ids)), dtype=torch.int64, pin_memory=_pin
).to(self.device, non_blocking=True)
input_ids_tensor = flatten_arrays_to_int64_tensor(input_ids, self.device, _pin)
seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=_pin).to(
self.device, non_blocking=True
)
python/sglang/srt/managers/schedule_policy.py
+2 -1
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import logging
from array import array
from sglang.srt.environ import envs
from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor
@@ -84,7 +85,7 @@ IGNORE_EOS_RESERVE_TOKENS = 1
def match_prefix_for_req(
tree_cache: BasePrefixCache,
req: Req,
token_ids: Optional[List[int]] = None,
token_ids: Optional[array[int]] = None,
*,
cow_mamba: bool = False,
include_req: bool = False,
python/sglang/srt/managers/scheduler.py
+8 -6
View File
@@ -20,6 +20,7 @@ import os
import signal
import sys
import time
from array import array
from collections import deque
from contextlib import contextmanager, nullcontext
from functools import partial
@@ -959,6 +960,7 @@ class Scheduler(
)
self.dp_tp_cpu_group = self.dp_tp_group.cpu_group
# TODO(Jialin): Migrate pad_input_ids implementations to return array.
self.pad_input_ids_func = self.tp_worker.get_pad_input_ids_func()
set_random_seed(self.random_seed)
@@ -1782,8 +1784,7 @@ class Scheduler(
if recv_req.input_embeds is not None:
# Generate fake input_ids based on the length of input_embeds
seq_length = len(recv_req.input_embeds)
fake_input_ids = [1] * seq_length
recv_req.input_ids = fake_input_ids
recv_req.input_ids = array("q", [1]) * seq_length
if recv_req.bootstrap_port is None:
# Use default bootstrap port
@@ -1909,8 +1910,8 @@ class Scheduler(
# Expand a single image token into multiple dummy tokens for receiving image embeddings.
# The pad function is model-specific and can be None for some backends.
if self.pad_input_ids_func:
req.origin_input_ids = self.pad_input_ids_func(
req.origin_input_ids, image_inputs
req.origin_input_ids = array(
"q", self.pad_input_ids_func(req.origin_input_ids, image_inputs)
)
req.extend_image_inputs(image_inputs)
self._maybe_compute_mrope_positions(req)
@@ -2182,8 +2183,9 @@ class Scheduler(
# embedding models or models not requiring special padding.
# If None, `req.origin_input_ids` is expected to be correctly populated already.
if self.pad_input_ids_func:
req.origin_input_ids = self.pad_input_ids_func(
req.origin_input_ids, image_inputs
# See companion call site above for the array.array wrap rationale.
req.origin_input_ids = array(
"q", self.pad_input_ids_func(req.origin_input_ids, image_inputs)
)
req.extend_image_inputs(image_inputs)
python/sglang/srt/managers/scheduler_pp_mixin.py
+8 -4
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import logging
import math
import time
from array import array
from collections import defaultdict, deque
from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
@@ -556,7 +557,7 @@ class SchedulerPPMixin:
if self.pp_group.is_first_rank:
model_runner = self.tp_worker.model_runner
model_config = model_runner.model_config
input_ids_list = []
input_ids_list: List[array[int]] = []
for i in range(128):
chunk_size = int(
self.chunked_prefill_size * 1.25
@@ -564,9 +565,12 @@ class SchedulerPPMixin:
)
if chunk_size <= 0:
break
input_ids = np.random.randint(
0, 10000, size=chunk_size, dtype=np.int64
).tolist()
input_ids = array(
"q",
np.random.randint(
0, 10000, size=chunk_size, dtype=np.int64
).tobytes(),
)
input_ids_list.append(input_ids)
sampling_params = SamplingParams(
python/sglang/srt/managers/tokenizer_manager.py
+13 -7
View File
@@ -13,6 +13,8 @@
# ==============================================================================
"""TokenizerManager is a process that tokenizes the text."""
from __future__ import annotations
import asyncio
import copy
import dataclasses
@@ -24,6 +26,7 @@ import signal
import socket
import sys
import threading
from array import array
from collections import deque
from contextlib import nullcontext
from datetime import datetime
@@ -987,12 +990,15 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
self,
obj: Union[GenerateReqInput, EmbeddingReqInput],
input_text: str,
input_ids: List[int],
input_ids: Optional[List[int]],
input_embeds: Optional[Union[List[float], None]] = None,
mm_inputs=None,
token_type_ids: Optional[List[int]] = None,
) -> Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput]:
"""Create a tokenized request object from common parameters."""
input_ids_arr: Optional[array[int]] = (
array("q", input_ids) if input_ids is not None else None
)
# Parse sampling parameters
# Note: if there are preferred sampling params, we use them if they are not
# explicitly passed in sampling_params
@@ -1020,7 +1026,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
tokenized_obj = TokenizedGenerateReqInput(
input_text,
input_ids,
input_ids_arr,
mm_inputs,
sampling_params,
obj.return_logprob,
@@ -1062,12 +1068,12 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
and obj.embed_override_token_id is not None
):
positional_embed_overrides = self._resolve_embed_overrides(
input_ids, obj.embed_override_token_id, obj.embed_overrides
input_ids_arr, obj.embed_override_token_id, obj.embed_overrides
)
tokenized_obj = TokenizedEmbeddingReqInput(
input_text,
input_ids,
input_ids_arr,
mm_inputs,
token_type_ids,
sampling_params,
@@ -1088,7 +1094,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
@staticmethod
def _resolve_embed_overrides(
input_ids: List[int],
input_ids: array[int],
token_id: int,
embeds: List[torch.Tensor],
) -> PositionalEmbeds:
@@ -1787,7 +1793,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
self.server_args.incremental_streaming_output and is_stream
)
delta_text = recv_obj.output_strs[i]
delta_output_ids = recv_obj.output_ids[i]
delta_output_ids = list(recv_obj.output_ids[i])
output_offset = state.last_output_offset
state.append_text(delta_text)
state.output_ids.extend(delta_output_ids)
@@ -1830,7 +1836,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
incremental = (
self.server_args.incremental_streaming_output and is_stream
)
delta_output_ids = recv_obj.output_ids[i]
delta_output_ids = list(recv_obj.output_ids[i])
output_offset = state.last_output_offset
state.output_ids.extend(delta_output_ids)