[Performance] Update xgrammar-related constrained decoding (#2056)

2024-11-18 09:58:49 +09:00
parent ebaa2f3199
commit 9c745d078e
4 changed files with 47 additions and 23 deletions
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -645,7 +645,7 @@ class ModelRunner:

        # Apply regex vocab_mask
        if sampling_info.vocab_mask is not None:
-            logits = logits.masked_fill(sampling_info.vocab_mask, float("-inf"))
+            sampling_info.apply_mask(logits=logits, vocab_mask=sampling_info.vocab_mask)

        return logits