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>

python/sglang/srt/model_executor/model_runner.py

@@ -2300,10 +2300,19 @@ class ModelRunner(ModelRunnerKVCacheMixin):
     ):
         # NOTE: In overlap mode, the function update_regex_vocab_mask (in sample)
         #       was executed after we processed last batch's results.
-
+        has_grammar = sampling_info.grammars is not None
+        # Avoid compiling grammar masks on every TP rank; only the first rank applies them.
+        apply_vocab_mask = has_grammar and self.sampler.tp_rank == 0
         # Calculate logits bias and apply it to next_token_logits.
-        sampling_info.update_regex_vocab_mask()
-        sampling_info.apply_logits_bias(logits_output.next_token_logits)
+        if apply_vocab_mask:
+            sampling_info.update_regex_vocab_mask()
+        else:
+            sampling_info.vocab_mask = None
+            sampling_info.apply_mask_func = None
+        sampling_info.apply_logits_bias(
+            logits_output.next_token_logits,
+            apply_vocab_mask=apply_vocab_mask,
+        )
 
     def sample(
         self,