feat: reduce constrained-decoding overhead in TP (#13947)

Signed-off-by: Raayan Dhar raayan.dhar@gmail.com <raayan.dhar@gmail.com>
Signed-off-by: raayandhar <raayan.dhar@gmail.com>

test/registered/constrained_decoding/test_constrained_decoding.py

@@ -5,6 +5,8 @@ register_amd_ci(est_time=179, suite="stage-b-test-small-1-gpu-amd")
 
 import unittest
 
+import torch
+
 from sglang.srt.utils import kill_process_tree
 from sglang.test.kits.ebnf_constrained_kit import TestEBNFConstrainedMinxin
 from sglang.test.kits.json_constrained_kit import TestJSONConstrainedMixin
@@ -21,6 +23,7 @@ from sglang.test.test_utils import (
 class ServerWithGrammar(CustomTestCase):
     backend = "xgrammar"
     disable_overlap = False
+    tp_size = 1
 
     @classmethod
     def setUpClass(cls):
@@ -33,6 +36,9 @@ class ServerWithGrammar(CustomTestCase):
             cls.backend,
         ]
 
+        if cls.tp_size > 1:
+            launch_args += ["--tp-size", str(cls.tp_size)]
+
         if cls.disable_overlap:
             launch_args += ["--disable-overlap-schedule"]
 
@@ -70,5 +76,34 @@ class TestLLGuidanceBackend(
     backend = "llguidance"
 
 
+# Test with TP=2
+@unittest.skipIf(torch.cuda.device_count() < 2, "Requires at least 2 GPUs")
+class TestXGrammarBackendTP2(
+    ServerWithGrammar,
+    TestJSONConstrainedMixin,
+    TestEBNFConstrainedMinxin,
+    TestRegexConstrainedMixin,
+):
+    backend = "xgrammar"
+    tp_size = 2
+
+
+@unittest.skipIf(torch.cuda.device_count() < 2, "Requires at least 2 GPUs")
+class TestOutlinesBackendTP2(ServerWithGrammar, TestJSONConstrainedMixin):
+    backend = "outlines"
+    tp_size = 2
+
+
+@unittest.skipIf(torch.cuda.device_count() < 2, "Requires at least 2 GPUs")
+class TestLLGuidanceBackendTP2(
+    ServerWithGrammar,
+    TestJSONConstrainedMixin,
+    TestEBNFConstrainedMinxin,
+    TestRegexConstrainedMixin,
+):
+    backend = "llguidance"
+    tp_size = 2
+
+
 if __name__ == "__main__":
     unittest.main()