Fix Req array token-id concatenation (#26182)

Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>

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

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

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 []

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

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()),
+            _make_req(1, "a b", [1, 2]),
-            Req(3, "a b c", [1, 2, 3], SamplingParams()),
+            _make_req(3, "a b c", [1, 2, 3]),
-            Req(2, "a", [1], SamplingParams()),
+            _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())
+        r1 = _make_req(1, "a b", [1, 2])
-        r2 = Req(3, "a b c", [1, 2, 3], SamplingParams())
+        r2 = _make_req(3, "a b c", [1, 2, 3])
-        r3 = Req(2, "a", [1], SamplingParams())
+        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())
+        r1 = _make_req(1, "a b", [1, 2])
-        r2 = Req(3, "a b c", [1, 2, 3], SamplingParams())
+        r2 = _make_req(3, "a b c", [1, 2, 3])
-        r3 = Req(2, "a", [1], SamplingParams())
+        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)),
+            _make_req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1000)),
-            Req(3, "a b c", [1, 2, 3], SamplingParams(max_new_tokens=10)),
+            _make_req(3, "a b c", [1, 2, 3], SamplingParams(max_new_tokens=10)),
-            Req(2, "a", [1], SamplingParams(max_new_tokens=100)),
+            _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),
+            _make_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),
+            _make_req(
-            Req(2, "a", [1], SamplingParams(max_new_tokens=100), priority=0),
+                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),
+            _make_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),
+            _make_req(
-            Req(2, "a", [1], SamplingParams(max_new_tokens=100), priority=1),
+                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"),
+            _make_req("r1", "a", [1], routing_key="key_a"),
-            Req("r2", "b", [2], SamplingParams(), routing_key="key_a"),
+            _make_req("r2", "b", [2], routing_key="key_a"),
-            Req("r3", "c", [3], SamplingParams(), routing_key="key_b"),
+            _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"),
+            _make_req("w1", "d", [4], routing_key="key_b"),
-            Req("w2", "e", [5], SamplingParams(), routing_key="key_a"),
+            _make_req("w2", "e", [5], routing_key="key_a"),
-            Req("w3", "f", [6], SamplingParams(), routing_key="key_c"),
+            _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"),
+            _make_req("r1", "a", [1], routing_key="key_b"),
-            Req("r2", "b", [2], SamplingParams(), routing_key="key_a"),
+            _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"),
+            _make_req("w1", "d", [4], routing_key="key_b"),
-            Req("w2", "e", [5], SamplingParams(), routing_key="key_a"),
+            _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"),
+            _make_req("r1", "a", [1], routing_key="key_a"),
-            Req("r2", "b", [2], SamplingParams(), routing_key="key_b"),
+            _make_req("r2", "b", [2], routing_key="key_b"),
-            Req("r3", "c", [3], SamplingParams(), routing_key="key_c"),
+            _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"),
+            _make_req("w1", "d", [4], routing_key="key_d"),
-            Req("w2", "e", [5], SamplingParams(), routing_key="key_e"),
+            _make_req("w2", "e", [5], routing_key="key_e"),
-            Req("w3", "f", [6], SamplingParams(), routing_key="key_c"),
+            _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"),
+            _make_req("w1", "d", [4], routing_key="key_a"),
-            Req("w2", "e", [5], SamplingParams(), routing_key="key_b"),
+            _make_req("w2", "e", [5], routing_key="key_b"),
-            Req("w3", "f", [6], SamplingParams(), routing_key="key_c"),
+            _make_req("w3", "f", [6], routing_key="key_c"),
         ]
 
         policy = SchedulePolicy(

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

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

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.origin_input_ids = array("q", origin_input_ids or [])
-    req.output_ids = output_ids or []
+    req.output_ids = array("q", output_ids or [])
     return req