From 07f326c184257e903d1d6446417a62ea5e9d3d86 Mon Sep 17 00:00:00 2001
From: Bowen Wang <89072276+bowenwan6@users.noreply.github.com>
Date: Thu, 4 Jun 2026 23:27:43 +0100
Subject: [PATCH] Fix multimodal synthetic benchmark prompt generation to
 exclude special tokens (#26864)

Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
---
 python/sglang/benchmark/datasets/common.py    | 17 ++++++++--
 .../bench_fn/test_benchmark_datasets_api.py   | 34 ++++++++++++++++++-
 2 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/python/sglang/benchmark/datasets/common.py b/python/sglang/benchmark/datasets/common.py
index d6eaad4ef..29bafff6b 100644
--- a/python/sglang/benchmark/datasets/common.py
+++ b/python/sglang/benchmark/datasets/common.py
@@ -81,10 +81,21 @@ def gen_prompt(tokenizer, token_num):
     return tokenizer.decode(selected_tokens)
 
 
+@lru_cache(maxsize=1)
+def get_available_multimodal_text_tokens(tokenizer, image_pad_id):
+    """Get valid token ids for synthetic multimodal text prompts."""
+    excluded_token_ids = set(getattr(tokenizer, "all_special_ids", []) or [])
+    if image_pad_id is not None:
+        excluded_token_ids.add(image_pad_id)
+    return [
+        token_id
+        for token_id in get_available_tokens(tokenizer)
+        if token_id not in excluded_token_ids
+    ]
+
+
 def gen_mm_prompt(tokenizer, image_pad_id, token_num):
     """Generate a random prompt of specified token length using tokenizer vocabulary."""
-    all_available_tokens = list(tokenizer.get_vocab().values())
-    if image_pad_id:
-        all_available_tokens.remove(image_pad_id)
+    all_available_tokens = get_available_multimodal_text_tokens(tokenizer, image_pad_id)
     selected_tokens = random.choices(all_available_tokens, k=token_num)
     return tokenizer.decode(selected_tokens)
diff --git a/test/registered/bench_fn/test_benchmark_datasets_api.py b/test/registered/bench_fn/test_benchmark_datasets_api.py
index d56babfee..0e1a0f816 100644
--- a/test/registered/bench_fn/test_benchmark_datasets_api.py
+++ b/test/registered/bench_fn/test_benchmark_datasets_api.py
@@ -19,7 +19,7 @@ from tokenizers.pre_tokenizers import Whitespace
 from transformers import PreTrainedTokenizerFast
 
 from sglang.benchmark.datasets import DATASET_MAPPING, get_dataset
-from sglang.benchmark.datasets.common import DatasetRow
+from sglang.benchmark.datasets.common import DatasetRow, gen_mm_prompt
 from sglang.benchmark.datasets.custom import sample_custom_requests
 from sglang.benchmark.datasets.generated_shared_prefix import (
     GeneratedSharedPrefixDataset,
@@ -384,6 +384,38 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase):
         self.assertTrue(all(isinstance(row, DatasetRow) for row in rows))
         self.assertTrue(all(row.image_data for row in rows))
 
+    def test_gen_mm_prompt_excludes_special_tokens(self):
+        tokenizer = create_lightweight_tokenizer()
+        multimodal_special_tokens = [
+            "<|image_pad|>",
+            "<|video_pad|>",
+            "<|vision_start|>",
+            "<|vision_end|>",
+            "<|vision_pad|>",
+        ]
+        tokenizer.add_special_tokens(
+            {"additional_special_tokens": multimodal_special_tokens}
+        )
+        special_token_ids = set(
+            tokenizer.convert_tokens_to_ids(multimodal_special_tokens)
+        )
+        image_pad_id = tokenizer.convert_tokens_to_ids("<|image_pad|>")
+        captured_population = {}
+
+        def fake_choices(population, k):
+            captured_population["tokens"] = population
+            return population[:k]
+
+        with patch(
+            "sglang.benchmark.datasets.common.random.choices",
+            side_effect=fake_choices,
+        ):
+            gen_mm_prompt(tokenizer, image_pad_id, token_num=8)
+
+        sampled_pool = set(captured_population["tokens"])
+        self.assertFalse(special_token_ids & sampled_pool)
+        self.assertTrue(sampled_pool)
+
     def test_mmmu_sampler(self):
         fake_records = [
             {"image_1": Image.new("RGB", (4, 4), color="white"), "question": "q1"},