fix: fix prefill-aware SWA floor tracking (#29520)

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: [],
+      verified: true,
-      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",
-      },
       env: [],
       flags: [
         "--model-path {{MODEL_NAME}}",

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,
+        # Kept in sync by _refresh_fill_ids; admission only updates
-        # never mutates this array's length.
+        # 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

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),
+            # hybrid-SWA early-return) leaves extend_range.end ==
-            # so there is nothing new to cache and stashing would be a no-op.
+            # 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(

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()