minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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/disaggregation/decode_schedule_batch_mixin.py
+2 -1
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import logging
from array import array
from http import HTTPStatus
from typing import TYPE_CHECKING, List
@@ -71,7 +72,7 @@ class ScheduleBatchDisaggregationDecodeMixin:
# Set fields
self.input_ids = torch.tensor(
sum(input_ids, []), dtype=torch.int32, device=self.device
sum(input_ids, array("q")), dtype=torch.int32, device=self.device
)
self.req_pool_indices = torch.tensor(
req_pool_indices, dtype=torch.int64, device=self.device
python/sglang/srt/disaggregation/encode_receiver.py
+2 -1
View File
@@ -7,6 +7,7 @@ import threading
import time
import uuid
from abc import ABC, abstractmethod
from array import array
from collections import OrderedDict, defaultdict
from enum import IntEnum
from http import HTTPStatus
@@ -588,7 +589,7 @@ class WaitingImageRequest:
**self.recv_embedding_data.get_mm_extra_meta(),
)
self.recv_req.mm_inputs = mm_inputs
self.recv_req.input_ids = mm_inputs.input_ids
self.recv_req.input_ids = array("q", mm_inputs.input_ids)
self.status = WaitingImageRequestStatus.SUCCESS
self.recv_socket.close()
python/sglang/srt/dllm/mixin/req.py
+2 -1
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import enum
from array import array
from typing import TYPE_CHECKING, Optional
from sglang.srt.dllm.config import DllmConfig
@@ -62,7 +63,7 @@ class ReqDllmMixin:
self.fill_ids = (
self.origin_input_ids
+ self.output_ids
+ [self.dllm_config.mask_id] * self.dllm_config.block_size
+ array("q", [self.dllm_config.mask_id] * self.dllm_config.block_size)
)
def _update_block_offset_for_dllm(self):
python/sglang/srt/dllm/mixin/scheduler.py
+2 -1
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import logging
from array import array
from typing import TYPE_CHECKING, List, Optional, Set, Union
from sglang.srt.dllm.config import DllmConfig
@@ -79,7 +80,7 @@ class SchedulerDllmMixin:
if new_tokens == 0:
continue
req.fill_ids[-new_tokens:] = next_token_ids[:]
req.fill_ids[-new_tokens:] = array("q", next_token_ids)
self.metrics_reporter.num_generated_tokens += new_tokens
req.output_ids.extend(next_token_ids)
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)
python/sglang/srt/mem_cache/events.py
+1 -1
View File
@@ -65,7 +65,7 @@ class KVCacheEventMixin:
if is_bigram:
page_tokens = [(raw[j], raw[j + 1]) for j in range(start, end)]
else:
page_tokens = raw[start:end]
page_tokens = list(raw[start:end])
block_hash = hash_str_to_int64(node.hash_value[page_index])
python/sglang/srt/mem_cache/hiradix_cache.py
-4
View File
@@ -804,10 +804,6 @@ class HiRadixCache(RadixCache):
def evictable_size(self):
return self.evictable_size_
def _to_radix_key(self, token_ids: List[int]) -> RadixKey:
"""Convert raw token_ids to a RadixKey; must be list (not tuple) for paged match."""
return RadixKey(token_ids=list(token_ids))
def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult:
if self.disable:
return IncLockRefResult(delta=0)
python/sglang/srt/mem_cache/mamba_radix_cache.py
+2 -1
View File
@@ -20,6 +20,7 @@ The radix tree data structure for managing the hybrid (full and Mamba) KV cache.
"""
import heapq
from array import array
from collections import defaultdict
from functools import lru_cache
from typing import TYPE_CHECKING, List, Optional, Tuple
@@ -456,7 +457,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
def reset(self) -> None:
self.root_node = TreeNode()
self.root_node.key = RadixKey([], None)
self.root_node.key = RadixKey(array("q"), None)
self.root_node.value = []
self.root_node.hash_value = []
self.root_node.full_lock_ref = 1
python/sglang/srt/mem_cache/radix_cache.py
+12 -14
View File
@@ -26,6 +26,7 @@ import heapq
import logging
import sys
import time
from array import array
from collections import defaultdict
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union
@@ -70,7 +71,7 @@ class RadixKey:
def __init__(
self,
token_ids: List[int],
token_ids: array[int],
extra_key: Optional[str] = None,
is_bigram: bool = False,
):
@@ -87,6 +88,7 @@ class RadixKey:
return n - 1 if n > 0 else 0
return len(self.token_ids)
# TODO(Jialin): vectorize with numpy without PyLong boxing
def __iter__(self) -> Iterator:
if self.is_bigram:
t = self.token_ids
@@ -110,7 +112,7 @@ class RadixKey:
if self.is_bigram:
# bigrams [start, stop) span raw tokens [start, stop + 1);
# empty slice -> empty raw tokens (not a dangling boundary token).
raw = self.token_ids[start : stop + 1] if stop > start else []
raw = self.token_ids[start : stop + 1] if stop > start else array("q")
return RadixKey(raw, self.extra_key, is_bigram=True)
return RadixKey(self.token_ids[start:stop], self.extra_key)
@@ -144,6 +146,7 @@ class RadixKey:
f"{self.extra_key=} != {other.extra_key=}"
)
# TODO(Jialin): replace zip with numpy to skip per-element PyLong boxing
def match(self, other: "RadixKey", page_size: int = 1) -> int:
"""Logical-unit prefix length shared with ``other``. Result is rounded down to ``page_size``."""
self._check_compatible(other)
@@ -337,7 +340,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
def reset(self):
# Initialize root with minimum priority so any real priority overrides it
self.root_node = TreeNode(priority=-sys.maxsize)
self.root_node.key = RadixKey(token_ids=[], extra_key=None)
self.root_node.key = RadixKey(token_ids=array("q"), extra_key=None)
self.root_node.value = []
self.root_node.host_value = []
self.root_node.lock_ref = 1
@@ -811,20 +814,15 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
if __name__ == "__main__":
tree = RadixCache.create_simulated()
# Example token id sequences (as lists of ints)
tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 3], extra_key=None)))
tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 3], extra_key=None)))
tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 4, 5], extra_key=None)))
tree.insert(
InsertParams(key=RadixKey(token_ids=[1, 2, 4, 5, 6, 7], extra_key=None))
)
tree.insert(
InsertParams(key=RadixKey(token_ids=[8, 9, 10, 11, 12], extra_key=None))
)
tree.insert(InsertParams(key=RadixKey(token_ids=array("q", [1, 2, 3]))))
tree.insert(InsertParams(key=RadixKey(token_ids=array("q", [1, 2, 3]))))
tree.insert(InsertParams(key=RadixKey(token_ids=array("q", [1, 2, 4, 5]))))
tree.insert(InsertParams(key=RadixKey(token_ids=array("q", [1, 2, 4, 5, 6, 7]))))
tree.insert(InsertParams(key=RadixKey(token_ids=array("q", [8, 9, 10, 11, 12]))))
tree.pretty_print()
print(
tree.match_prefix(
MatchPrefixParams(key=RadixKey(token_ids=[1, 2, 3, 13, 14], extra_key=None))
MatchPrefixParams(key=RadixKey(token_ids=array("q", [1, 2, 3, 13, 14])))
)
)
python/sglang/srt/mem_cache/unified_radix_cache.py
+2 -1
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import logging
import threading
import time
from array import array
from collections import defaultdict
from functools import partial
from typing import TYPE_CHECKING, Any, Optional
@@ -257,7 +258,7 @@ class UnifiedRadixCache(BasePrefixCache):
def _reset_full(self) -> None:
"""Full reset: destroy entire tree and all state."""
self.root_node = UnifiedTreeNode(self.tree_components)
self.root_node.key = RadixKey([], None)
self.root_node.key = RadixKey(array("q"), None)
self.root_node.component_data[BASE_COMPONENT_TYPE].value = []
for ct in self.tree_components:
self.root_node.component_data[ct].lock_ref = 1
python/sglang/srt/models/llava.py
+8 -2
View File
@@ -13,8 +13,11 @@
# ==============================================================================
"""Inference-only LLaVa model compatible with HuggingFace weights."""
from __future__ import annotations
import math
import re
from array import array
from functools import lru_cache
from typing import Dict, Iterable, List, Optional, Tuple, Type, Union
@@ -73,7 +76,9 @@ class LlavaBaseForCausalLM(nn.Module):
return "pad"
return "anyres"
def pad_input_ids(self, input_ids: List[int], image_inputs: MultimodalInputs):
def pad_input_ids(
self, input_ids: array[int], image_inputs: MultimodalInputs
) -> array[int]:
image_sizes = flatten_nested_list(
[item.image_sizes for item in image_inputs.mm_items]
)
@@ -125,9 +130,10 @@ class LlavaBaseForCausalLM(nn.Module):
except ValueError:
offset = 0
# old_len + pad_len - 1, because we need to remove image_token_id
pad_token = pad_values[image_idx % len(pad_values)]
input_ids = (
input_ids[:offset]
+ [pad_values[image_idx % len(pad_values)]] * new_image_feature_len
+ array("q", [pad_token]) * new_image_feature_len
+ input_ids[offset + 1 :]
)
offset_list.append(offset)
python/sglang/srt/models/llavavid.py
+8 -3
View File
@@ -13,7 +13,10 @@
# ==============================================================================
"""Inference-only LLaVa video model compatible with HuggingFace weights."""
from typing import Iterable, List, Optional, Tuple
from __future__ import annotations
from array import array
from typing import Iterable, Optional, Tuple
import numpy as np
import torch
@@ -57,8 +60,10 @@ class LlavaVidForCausalLM(nn.Module):
torch.empty(config.text_config.hidden_size, dtype=torch.float16)
)
def pad_input_ids(self, input_ids: List[int], image_inputs: MultimodalInputs):
pad_values = [item.pad_value for item in image_inputs.mm_items]
def pad_input_ids(
self, input_ids: array[int], image_inputs: MultimodalInputs
) -> array[int]:
pad_values = array("q", (item.pad_value for item in image_inputs.mm_items))
new_image_feature_len = self.image_feature_len
pad_ids = pad_values * (
python/sglang/srt/models/mllama.py
+7 -2
View File
@@ -4,7 +4,10 @@
# https://github.com/vllm-project/vllm/blob/7193774b1ff8603ad5bf4598e5efba0d9a39b436/vllm/model_executor/models/mllama.py
"""PyTorch Mllama model."""
from __future__ import annotations
import math
from array import array
from typing import Iterable, List, Optional, Tuple, Union
import torch
@@ -823,9 +826,11 @@ class MllamaForConditionalGeneration(nn.Module):
)
self.logits_processor = LogitsProcessor(config.text_config)
def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
def pad_input_ids(
self, input_ids: array[int], mm_inputs: MultimodalInputs
) -> array[int]:
pixel_values = torch.cat([item.feature for item in mm_inputs.mm_items], dim=0)
pad_values = [item.pad_value for item in mm_inputs.mm_items]
pad_values = array("q", (item.pad_value for item in mm_inputs.mm_items))
num_concurrent_media, num_tiles = pixel_values.shape[1:3]
num_patches = self.vision_model.num_patches
python/sglang/srt/models/moss_vl.py
+7 -4
View File
@@ -3,6 +3,7 @@
from __future__ import annotations
import logging
from array import array
from functools import partial
from typing import Iterable, List, Optional, Tuple
@@ -1122,15 +1123,17 @@ class MossVLForConditionalGeneration(nn.Module):
return total_len
def _build_encoder_prefix_pad_ids(self, mm_inputs: MultimodalInputs) -> List[int]:
def _build_encoder_prefix_pad_ids(self, mm_inputs: MultimodalInputs) -> array[int]:
encoder_len = self._get_encoder_len(mm_inputs)
if encoder_len == 0 or not mm_inputs.mm_items:
return []
return array("q")
pad_value = mm_inputs.mm_items[0].pad_value
return [pad_value] * encoder_len
return array("q", [pad_value]) * encoder_len
def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
def pad_input_ids(
self, input_ids: array[int], mm_inputs: MultimodalInputs
) -> array[int]:
encoder_len = self._get_encoder_len(mm_inputs)
mm_inputs.num_image_tokens = encoder_len
if encoder_len == 0:
python/sglang/srt/models/whisper.py
+8 -4
View File
@@ -1,4 +1,7 @@
from typing import Any, Iterable, List, Optional, Tuple
from __future__ import annotations
from array import array
from typing import Any, Iterable, Optional, Tuple
import torch
from transformers import WhisperConfig
@@ -418,14 +421,15 @@ class WhisperForConditionalGeneration(torch.nn.Module):
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
def pad_input_ids(
self, input_ids: array[int], mm_inputs: MultimodalInputs
) -> array[int]:
# Prepend dummy encoder tokens so that prepare_encoder_info_extend
# correctly allocates encoder KV cache locations in the KV pool.
# These dummy tokens are stripped before the model forward receives input_ids.
encoder_len = self.config.max_source_positions
mm_inputs.num_image_tokens = encoder_len
pad_ids = [0] * encoder_len
return pad_ids + input_ids
return array("q", [0]) * encoder_len + input_ids
def forward(
self,
python/sglang/srt/utils/common.py
+16
View File
@@ -45,6 +45,7 @@ import traceback
import types
import uuid
import warnings
from array import array
from collections import OrderedDict, defaultdict
from contextlib import contextmanager
from dataclasses import dataclass
@@ -101,6 +102,21 @@ logger = logging.getLogger(__name__)
torch_release = pkg_version.parse(torch.__version__).release
def flatten_arrays_to_int64_tensor(
parts: List[array[int]], device, pin: bool
) -> torch.Tensor:
"""Flatten a list of array.array('q') buffers into one int64 tensor.
Uses NumPy here to speed up the conversion by using memcpy
instead of a per-element PyLong-to-int64 walk.
"""
combined = np.concatenate([np.frombuffer(p, dtype=np.int64) for p in parts])
cpu_t = torch.from_numpy(combined)
if pin:
cpu_t = cpu_t.pin_memory()
return cpu_t.to(device, non_blocking=True)
# https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip
@lru_cache(maxsize=1)
def is_hip() -> bool: