From 643e1cc7790bd4cdb0906cb891e19358e7e6c493 Mon Sep 17 00:00:00 2001 From: Mick Date: Sun, 28 Jun 2026 14:18:28 +0800 Subject: [PATCH] fix: fix prefill-aware SWA floor tracking (#29520) --- .../snippets/configs/baidu/unlimited-ocr.jsx | 108 +----------------- python/sglang/srt/managers/schedule_batch.py | 4 +- python/sglang/srt/managers/scheduler.py | 7 +- .../vlm/test_unlimited_ocr_server.py | 88 ++++++++++++++ 4 files changed, 96 insertions(+), 111 deletions(-) create mode 100644 test/registered/vlm/test_unlimited_ocr_server.py diff --git a/docs_new/src/snippets/configs/baidu/unlimited-ocr.jsx b/docs_new/src/snippets/configs/baidu/unlimited-ocr.jsx index ea29431ef..0d5406ca0 100644 --- a/docs_new/src/snippets/configs/baidu/unlimited-ocr.jsx +++ b/docs_new/src/snippets/configs/baidu/unlimited-ocr.jsx @@ -3,7 +3,7 @@ export const config = { modelName: "Unlimited-OCR", - supportedHardware: ["h100", "h200", "b200", "b300", "gb200", "gb300"], + supportedHardware: ["h100"], variants: [{ id: "default", label: "Default" }], quantizations: [{ id: "default", label: "Default" }], @@ -44,11 +44,6 @@ export const config = { dockerImages: { h100: "lmsysorg/sglang:dev", - h200: "lmsysorg/sglang:dev", - b200: "lmsysorg/sglang:dev", - b300: "lmsysorg/sglang:dev", - gb200: "lmsysorg/sglang:dev", - gb300: "lmsysorg/sglang:dev", }, github: { @@ -72,106 +67,7 @@ export const config = { strategy: "balanced", nodes: "single", }, - env: [], - flags: [ - "--model-path {{MODEL_NAME}}", - "--attention-backend fa3", - "--page-size 1", - "--context-length 32768", - "--enable-custom-logit-processor", - "--disable-radix-cache", - "--host {{HOST_IP}}", - "--port {{PORT}}", - ], - }, - { - match: { - hw: "h200", - variant: "default", - quant: "default", - strategy: "balanced", - nodes: "single", - }, - env: [], - flags: [ - "--model-path {{MODEL_NAME}}", - "--attention-backend fa3", - "--page-size 1", - "--context-length 32768", - "--enable-custom-logit-processor", - "--disable-radix-cache", - "--host {{HOST_IP}}", - "--port {{PORT}}", - ], - }, - { - match: { - hw: "b200", - variant: "default", - quant: "default", - strategy: "balanced", - nodes: "single", - }, - env: [], - flags: [ - "--model-path {{MODEL_NAME}}", - "--attention-backend fa3", - "--page-size 1", - "--context-length 32768", - "--enable-custom-logit-processor", - "--disable-radix-cache", - "--host {{HOST_IP}}", - "--port {{PORT}}", - ], - }, - { - match: { - hw: "b300", - variant: "default", - quant: "default", - strategy: "balanced", - nodes: "single", - }, - env: [], - flags: [ - "--model-path {{MODEL_NAME}}", - "--attention-backend fa3", - "--page-size 1", - "--context-length 32768", - "--enable-custom-logit-processor", - "--disable-radix-cache", - "--host {{HOST_IP}}", - "--port {{PORT}}", - ], - }, - { - match: { - hw: "gb200", - variant: "default", - quant: "default", - strategy: "balanced", - nodes: "single", - }, - env: [], - flags: [ - "--model-path {{MODEL_NAME}}", - "--attention-backend fa3", - "--page-size 1", - "--context-length 32768", - "--enable-custom-logit-processor", - "--disable-radix-cache", - "--host {{HOST_IP}}", - "--port {{PORT}}", - ], - }, - { - match: { - hw: "gb300", - variant: "default", - quant: "default", - strategy: "balanced", - nodes: "single", - }, + verified: true, env: [], flags: [ "--model-path {{MODEL_NAME}}", diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 021490c35..7efd623b9 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -723,8 +723,8 @@ class Req(ReqDllmMixin): # that preserve length would silently corrupt fill_ids. self.output_ids = array("q") # Full untruncated sequence: origin + output (+ DLLM mask block). - # Kept in sync by _refresh_fill_ids; admission only updates fill_len, - # never mutates this array's length. + # Kept in sync by _refresh_fill_ids; admission only updates + # extend_range, never mutates this array's length. self.full_untruncated_fill_ids = array("q") self.extend_range: Optional[Range] = None self.dllm_initialized: bool = False diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f12843b74..11b16d8d3 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2622,8 +2622,9 @@ class Scheduler( # Stash (cache) the previous chunk only when it produced new KV # beyond what is already cached. A parked chunk (add_chunked_req - # hybrid-SWA early-return) leaves fill_len == len(prefix_indices), - # so there is nothing new to cache and stashing would be a no-op. + # hybrid-SWA early-return) leaves extend_range.end == + # len(prefix_indices), so there is nothing new to cache and + # stashing would be a no-op. if self.chunked_req.extend_range.end > len(self.chunked_req.prefix_indices): self.stash_chunked_request(self.chunked_req) @@ -2970,7 +2971,7 @@ class Scheduler( if self.tp_worker.model_runner.prefill_aware_swa: for req in can_run_list: - req.swa_evict_floor = req.fill_len + req.swa_evict_floor = req.extend_range.end # Record prefill stats for logging after forward. new_batch.prefill_stats = PrefillStats.from_adder( diff --git a/test/registered/vlm/test_unlimited_ocr_server.py b/test/registered/vlm/test_unlimited_ocr_server.py new file mode 100644 index 000000000..af20c0742 --- /dev/null +++ b/test/registered/vlm/test_unlimited_ocr_server.py @@ -0,0 +1,88 @@ +import base64 +import io +import unittest + +import openai +from PIL import Image, ImageDraw, ImageFont + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.vlm_utils import TestOpenAIMLLMServerBase + +register_cuda_ci(est_time=240, stage="base-b", runner_config="1-gpu-large") + + +class TestUnlimitedOCRServer(TestOpenAIMLLMServerBase): + model = "baidu/Unlimited-OCR" + trust_remote_code = False + extra_args = [ + "--attention-backend=fa3", + "--page-size=1", + "--context-length=4096", + "--max-total-tokens=4096", + "--enable-custom-logit-processor", + "--disable-radix-cache", + "--cuda-graph-max-bs-decode=4", + ] + + @staticmethod + def _font(size: int): + for path in ( + "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + "/usr/share/fonts/truetype/liberation2/LiberationSans-Regular.ttf", + ): + try: + return ImageFont.truetype(path, size=size) + except OSError: + pass + return ImageFont.load_default() + + @classmethod + def _make_ocr_image_url(cls) -> str: + img = Image.new("RGB", (640, 360), "white") + draw = ImageDraw.Draw(img) + draw.rectangle((24, 24, 616, 336), outline="black", width=4) + draw.text((72, 92), "SGLang OCR", fill="black", font=cls._font(56)) + draw.text((72, 180), "12345", fill="black", font=cls._font(72)) + + buffer = io.BytesIO() + img.save(buffer, format="PNG") + encoded = base64.b64encode(buffer.getvalue()).decode("ascii") + return f"data:image/png;base64,{encoded}" + + def test_single_image_chat_completion(self): + client = openai.Client(api_key=self.api_key, base_url=self.base_url) + + response = client.chat.completions.create( + model="default", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "document parsing."}, + { + "type": "image_url", + "image_url": {"url": self._make_ocr_image_url()}, + }, + ], + }, + ], + temperature=0, + max_tokens=64, + extra_body={"images_config": {"image_mode": "gundam"}}, + ) + + self.assertEqual(response.choices[0].message.role, "assistant") + text = response.choices[0].message.content + self.assertIsInstance(text, str) + self.assertIn("12345", text) + self.assertIn("sglang", text.lower()) + self.assertGreater(response.usage.prompt_tokens, 0) + self.assertGreater(response.usage.completion_tokens, 0) + self.assertGreater(response.usage.total_tokens, 0) + + +del TestOpenAIMLLMServerBase + + +if __name__ == "__main__": + unittest.main()