diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index 79db0da49..1e7b90609 100644 --- a/python/sglang/bench_one_batch.py +++ b/python/sglang/bench_one_batch.py @@ -56,6 +56,7 @@ import logging import multiprocessing import os import time +from array import array from types import SimpleNamespace from typing import Optional, Tuple @@ -367,7 +368,7 @@ def prepare_inputs_for_correctness_test(bench_args, tokenizer, custom_prompts): req = Req( rid=i, origin_input_text=prompts[i], - origin_input_ids=tmp_input_ids, + origin_input_ids=array("q", tmp_input_ids), sampling_params=sampling_params, ) req.fill_ids = req.origin_input_ids @@ -412,7 +413,7 @@ def prepare_synthetic_inputs_for_latency_test( req = Req( rid=i, origin_input_text="", - origin_input_ids=list(input_ids[i]), + origin_input_ids=array("q", input_ids[i]), sampling_params=sampling_params, ) req.fill_ids = req.origin_input_ids diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index ac55e445c..5078b2697 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -688,9 +688,9 @@ class Req(ReqDllmMixin): ): # Input and output info self.rid = rid - self.origin_input_ids = array("q", origin_input_ids) + self.origin_input_ids = origin_input_ids self.origin_input_ids_unpadded = ( - array("q", origin_input_ids_unpadded) + origin_input_ids_unpadded if origin_input_ids_unpadded else self.origin_input_ids ) # Before image padding diff --git a/python/sglang/srt/sampling/custom_logit_processor.py b/python/sglang/srt/sampling/custom_logit_processor.py index d58a5f6cf..5bd128135 100644 --- a/python/sglang/srt/sampling/custom_logit_processor.py +++ b/python/sglang/srt/sampling/custom_logit_processor.py @@ -1,5 +1,6 @@ import json from abc import ABC, abstractmethod +from array import array from functools import lru_cache from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set @@ -165,7 +166,7 @@ class DeepseekOCRNoRepeatNGramLogitProcessor(CustomLogitProcessor): if ngram_size <= 0 or window_size <= 0: continue - sequence: List[int] = req.origin_input_ids + req.output_ids + sequence = req.origin_input_ids + req.output_ids if len(sequence) < ngram_size: continue @@ -175,14 +176,14 @@ class DeepseekOCRNoRepeatNGramLogitProcessor(CustomLogitProcessor): continue if ngram_size > 1: - current_prefix = tuple(sequence[-(ngram_size - 1) :]) + current_prefix = sequence[-(ngram_size - 1) :] else: - current_prefix = tuple() + current_prefix = array("q") banned_tokens: Set[int] = set() for idx in range(search_start, search_end): ngram = sequence[idx : idx + ngram_size] - if ngram_size == 1 or tuple(ngram[:-1]) == current_prefix: + if ngram_size == 1 or ngram[:-1] == current_prefix: banned_tokens.add(ngram[-1]) whitelist_ids = params.get("whitelist_token_ids") or [] diff --git a/test/manual/test_forward_split_prefill.py b/test/manual/test_forward_split_prefill.py index 2712bfaaa..80d6cc661 100644 --- a/test/manual/test_forward_split_prefill.py +++ b/test/manual/test_forward_split_prefill.py @@ -8,6 +8,7 @@ python3 test_forward_split_prefill.py """ import unittest +from array import array import numpy as np import torch @@ -91,7 +92,7 @@ class TestForwardSplitPrefill(CustomTestCase): req = Req( rid=i, origin_input_text="", - origin_input_ids=list(input_ids[i]), + origin_input_ids=array("q", input_ids[i]), sampling_params=sampling_params, ) req.fill_ids = req.origin_input_ids diff --git a/test/manual/test_schedule_policy.py b/test/manual/test_schedule_policy.py index aa9ecc86c..c04481a46 100644 --- a/test/manual/test_schedule_policy.py +++ b/test/manual/test_schedule_policy.py @@ -1,4 +1,5 @@ import unittest +from array import array from sglang.srt.managers.schedule_batch import Req, ScheduleBatch from sglang.srt.managers.schedule_policy import ( @@ -11,6 +12,18 @@ from sglang.srt.sampling.sampling_params import SamplingParams from sglang.test.test_utils import CustomTestCase +def _make_req(rid, origin_input_text, origin_input_ids, sampling_params=None, **kwargs): + if sampling_params is None: + sampling_params = SamplingParams() + return Req( + rid, + origin_input_text, + array("q", origin_input_ids), + sampling_params, + **kwargs, + ) + + class TestSchedulePolicy(CustomTestCase): def setUp(self): @@ -60,9 +73,9 @@ class TestSchedulePolicy(CustomTestCase): def test_calc_priority_fcfs(self): tree_cache = RadixCache.create_simulated() waiting_queue = [ - Req(1, "a b", [1, 2], SamplingParams()), - Req(3, "a b c", [1, 2, 3], SamplingParams()), - Req(2, "a", [1], SamplingParams()), + _make_req(1, "a b", [1, 2]), + _make_req(3, "a b c", [1, 2, 3]), + _make_req(2, "a", [1]), ] policy = SchedulePolicy( @@ -80,9 +93,9 @@ class TestSchedulePolicy(CustomTestCase): def test_calc_priority_priority_enabled_fcfs_scheduling(self): tree_cache = RadixCache.create_simulated() - r1 = Req(1, "a b", [1, 2], SamplingParams()) - r2 = Req(3, "a b c", [1, 2, 3], SamplingParams()) - r3 = Req(2, "a", [1], SamplingParams()) + r1 = _make_req(1, "a b", [1, 2]) + r2 = _make_req(3, "a b c", [1, 2, 3]) + r3 = _make_req(2, "a", [1]) r1.priority, r1.time_stats.wait_queue_entry_time = 1, 1 r2.priority, r2.time_stats.wait_queue_entry_time = 0, 1 r3.priority, r3.time_stats.wait_queue_entry_time = 0, 0 @@ -107,9 +120,9 @@ class TestSchedulePolicy(CustomTestCase): self, ): tree_cache = RadixCache.create_simulated() - r1 = Req(1, "a b", [1, 2], SamplingParams()) - r2 = Req(3, "a b c", [1, 2, 3], SamplingParams()) - r3 = Req(2, "a", [1], SamplingParams()) + r1 = _make_req(1, "a b", [1, 2]) + r2 = _make_req(3, "a b c", [1, 2, 3]) + r3 = _make_req(2, "a", [1]) r1.priority, r1.time_stats.wait_queue_entry_time = -1, 1 r2.priority, r2.time_stats.wait_queue_entry_time = 0, 1 r3.priority, r3.time_stats.wait_queue_entry_time = 0, 0 @@ -133,9 +146,9 @@ class TestSchedulePolicy(CustomTestCase): tree_cache = RadixCache.create_simulated() waiting_queue = [ - Req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1000)), - Req(3, "a b c", [1, 2, 3], SamplingParams(max_new_tokens=10)), - Req(2, "a", [1], SamplingParams(max_new_tokens=100)), + _make_req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1000)), + _make_req(3, "a b c", [1, 2, 3], SamplingParams(max_new_tokens=10)), + _make_req(2, "a", [1], SamplingParams(max_new_tokens=100)), ] policy = SchedulePolicy( @@ -155,9 +168,11 @@ class TestSchedulePolicy(CustomTestCase): tree_cache = RadixCache.create_simulated() waiting_queue = [ - Req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1), priority=1), - Req(3, "a b c", [1, 2, 3], SamplingParams(max_new_tokens=10), priority=0), - Req(2, "a", [1], SamplingParams(max_new_tokens=100), priority=0), + _make_req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1), priority=1), + _make_req( + 3, "a b c", [1, 2, 3], SamplingParams(max_new_tokens=10), priority=0 + ), + _make_req(2, "a", [1], SamplingParams(max_new_tokens=100), priority=0), ] policy = SchedulePolicy( @@ -179,9 +194,11 @@ class TestSchedulePolicy(CustomTestCase): tree_cache = RadixCache.create_simulated() waiting_queue = [ - Req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1), priority=0), - Req(3, "a b c", [1, 2, 3], SamplingParams(max_new_tokens=10), priority=1), - Req(2, "a", [1], SamplingParams(max_new_tokens=100), priority=1), + _make_req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1), priority=0), + _make_req( + 3, "a b c", [1, 2, 3], SamplingParams(max_new_tokens=10), priority=1 + ), + _make_req(2, "a", [1], SamplingParams(max_new_tokens=100), priority=1), ] policy = SchedulePolicy( @@ -202,16 +219,16 @@ class TestSchedulePolicy(CustomTestCase): tree_cache = RadixCache.create_simulated() running_reqs = [ - Req("r1", "a", [1], SamplingParams(), routing_key="key_a"), - Req("r2", "b", [2], SamplingParams(), routing_key="key_a"), - Req("r3", "c", [3], SamplingParams(), routing_key="key_b"), + _make_req("r1", "a", [1], routing_key="key_a"), + _make_req("r2", "b", [2], routing_key="key_a"), + _make_req("r3", "c", [3], routing_key="key_b"), ] running_batch = ScheduleBatch(reqs=running_reqs) waiting_queue = [ - Req("w1", "d", [4], SamplingParams(), routing_key="key_b"), - Req("w2", "e", [5], SamplingParams(), routing_key="key_a"), - Req("w3", "f", [6], SamplingParams(), routing_key="key_c"), + _make_req("w1", "d", [4], routing_key="key_b"), + _make_req("w2", "e", [5], routing_key="key_a"), + _make_req("w3", "f", [6], routing_key="key_c"), ] policy = SchedulePolicy( @@ -232,14 +249,14 @@ class TestSchedulePolicy(CustomTestCase): tree_cache = RadixCache.create_simulated() running_reqs = [ - Req("r1", "a", [1], SamplingParams(), routing_key="key_b"), - Req("r2", "b", [2], SamplingParams(), routing_key="key_a"), + _make_req("r1", "a", [1], routing_key="key_b"), + _make_req("r2", "b", [2], routing_key="key_a"), ] running_batch = ScheduleBatch(reqs=running_reqs) waiting_queue = [ - Req("w1", "d", [4], SamplingParams(), routing_key="key_b"), - Req("w2", "e", [5], SamplingParams(), routing_key="key_a"), + _make_req("w1", "d", [4], routing_key="key_b"), + _make_req("w2", "e", [5], routing_key="key_a"), ] policy = SchedulePolicy( @@ -259,16 +276,16 @@ class TestSchedulePolicy(CustomTestCase): tree_cache = RadixCache.create_simulated() running_reqs = [ - Req("r1", "a", [1], SamplingParams(), routing_key="key_a"), - Req("r2", "b", [2], SamplingParams(), routing_key="key_b"), - Req("r3", "c", [3], SamplingParams(), routing_key="key_c"), + _make_req("r1", "a", [1], routing_key="key_a"), + _make_req("r2", "b", [2], routing_key="key_b"), + _make_req("r3", "c", [3], routing_key="key_c"), ] running_batch = ScheduleBatch(reqs=running_reqs) waiting_queue = [ - Req("w1", "d", [4], SamplingParams(), routing_key="key_d"), - Req("w2", "e", [5], SamplingParams(), routing_key="key_e"), - Req("w3", "f", [6], SamplingParams(), routing_key="key_c"), + _make_req("w1", "d", [4], routing_key="key_d"), + _make_req("w2", "e", [5], routing_key="key_e"), + _make_req("w3", "f", [6], routing_key="key_c"), ] policy = SchedulePolicy( @@ -291,9 +308,9 @@ class TestSchedulePolicy(CustomTestCase): running_batch = ScheduleBatch(reqs=[]) waiting_queue = [ - Req("w1", "d", [4], SamplingParams(), routing_key="key_a"), - Req("w2", "e", [5], SamplingParams(), routing_key="key_b"), - Req("w3", "f", [6], SamplingParams(), routing_key="key_c"), + _make_req("w1", "d", [4], routing_key="key_a"), + _make_req("w2", "e", [5], routing_key="key_b"), + _make_req("w3", "f", [6], routing_key="key_c"), ] policy = SchedulePolicy( diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index c7a478eb0..04f6d28ec 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -246,7 +246,7 @@ def create_bench_cache( req = Req( rid=_rid[0], origin_input_text="", - origin_input_ids=[], + origin_input_ids=array("q"), sampling_params=SamplingParams(temperature=0, max_new_tokens=1), ) _rid[0] += 1 diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 232b11f46..f9847b873 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -539,7 +539,7 @@ class UnifiedRadixCacheSuite: req = Req( rid=self._rid, origin_input_text="", - origin_input_ids=[], + origin_input_ids=array("q"), sampling_params=sp, ) self._rid += 1 diff --git a/test/registered/unit/sampling/test_custom_logit_processor.py b/test/registered/unit/sampling/test_custom_logit_processor.py index b6b44a825..802b150a6 100644 --- a/test/registered/unit/sampling/test_custom_logit_processor.py +++ b/test/registered/unit/sampling/test_custom_logit_processor.py @@ -7,6 +7,7 @@ register_cpu_ci(est_time=7, suite="base-b-test-cpu") import json import unittest +from array import array from unittest.mock import MagicMock import torch @@ -25,8 +26,8 @@ from sglang.test.test_utils import CustomTestCase # Helper: mock a Req object (used by ThinkingBudget and NGram processors) def _make_req(origin_input_ids=None, output_ids=None): req = MagicMock() - req.origin_input_ids = origin_input_ids or [] - req.output_ids = output_ids or [] + req.origin_input_ids = array("q", origin_input_ids or []) + req.output_ids = array("q", output_ids or []) return req