Fix spec decoding with grammar in disagg (#24082)

Co-authored-by: jimmy.shong <jimmy.shong@radixark.ai>
Co-authored-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: Codex <codex@example.com>

test/registered/disaggregation/test_disaggregation_basic.py

@@ -27,7 +27,7 @@ from sglang.test.test_utils import (
     DEFAULT_TARGET_MODEL_EAGLE3,
 )
 
-register_cuda_ci(est_time=700, stage="base-b", runner_config="2-gpu-large")
+register_cuda_ci(est_time=890, stage="base-b", runner_config="2-gpu-large")
 
 
 class TestDisaggregationAccuracy(PauseResumeInPlaceMixin, PDDisaggregationServerBase):
@@ -260,6 +260,104 @@ class TestDisaggregationMooncakeSpec(JSONConstrainedMixin, PDDisaggregationServe
         self.assertGreater(metrics["score"], 0.74)
 
 
+class TestDisaggregationSpecV2Grammar(PDDisaggregationServerBase):
+    """Regression for PD disagg + EAGLE Spec V2 + grammar structured output."""
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.model = DEFAULT_TARGET_MODEL_EAGLE3
+        spec_args = [
+            "--speculative-algorithm",
+            "EAGLE",
+            "--speculative-draft-model-path",
+            DEFAULT_DRAFT_MODEL_EAGLE3,
+            "--speculative-num-steps",
+            "3",
+            "--speculative-eagle-topk",
+            "1",  # Spec V2 only supports topk=1
+            "--speculative-num-draft-tokens",
+            "4",
+            "--grammar-backend",
+            "xgrammar",
+            "--cuda-graph-max-bs",
+            "8",
+            "--dtype=float16",
+            # Cap context to the EAGLE3 draft's native length so the target and
+            # draft ModelConfigs agree (the draft derives 2048 vs the Llama-3.1
+            # target's 131072, which the Spec V2 draft worker otherwise rejects).
+            # 2048 is far above this test's output length.
+            "--context-length",
+            "2048",
+        ]
+        cls.extra_prefill_args = spec_args
+        cls.extra_decode_args = spec_args
+        cls.launch_all()
+
+    @staticmethod
+    def _json_schema() -> str:
+        return json.dumps(
+            {
+                "type": "object",
+                "properties": {
+                    "name": {"type": "string", "pattern": "^[\\w]+$"},
+                    "population": {"type": "integer"},
+                    "country": {"type": "string", "pattern": "^[\\w ]+$"},
+                    "capital": {"type": "string", "pattern": "^[\\w ]+$"},
+                },
+                "required": ["name", "population", "country", "capital"],
+            }
+        )
+
+    def _generate(self, return_logprob: bool):
+        response = requests.post(
+            f"{self.lb_url}/generate",
+            json={
+                "text": "Here is the information of the capital of France in the JSON format.\n",
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": 256,
+                    "json_schema": self._json_schema(),
+                },
+                "return_logprob": return_logprob,
+                "logprob_start_len": 0,
+            },
+        )
+        self.assertEqual(response.status_code, 200, response.text)
+        out = response.json()
+        self.assertGreater(
+            out["meta_info"]["spec_verify_ct"],
+            0,
+            "expected Spec V2 to run (spec_verify_ct > 0)",
+        )
+        return out
+
+    def test_structured_output_no_trailing_tokens(self):
+        """Output is valid JSON with nothing emitted past grammar completion."""
+        out = self._generate(return_logprob=False)
+        text = out["text"]
+        parsed = json.loads(text)
+        for key in ("name", "population", "country", "capital"):
+            self.assertIn(key, parsed)
+        self.assertTrue(
+            text.strip().endswith("}"), f"unexpected trailing tokens: {text!r}"
+        )
+
+    def test_logprob_count_matches_completion_tokens(self):
+        """Trimmed Spec V2 tokens keep logprob count == completion token count."""
+        out = self._generate(return_logprob=True)
+        meta = out["meta_info"]
+        completion_tokens = meta["completion_tokens"]
+        output_logprobs = meta["output_token_logprobs"]
+        self.assertEqual(
+            len(output_logprobs),
+            completion_tokens,
+            "output logprobs must align with retained (trimmed) tokens: "
+            f"got {len(output_logprobs)} logprobs vs {completion_tokens} completion tokens",
+        )
+        json.loads(out["text"])
+
+
 class TestDisaggregationSimulatedRetract(PDDisaggregationServerBase):
     @classmethod
     def setUpClass(cls):