A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 22:15 UTC
[v0][Core] Use xgrammar shared context to avoid copy overhead for offline engine (#13837)
Signed-off-by: Seth Kimmel <seth.kimmel3@gmail.com>
Seth Kimmel committed
Feb 25, 2025 at 22:58 UTC
e206b5433109d298e53451015465b2bf8f03ef0a
1 file changed
+23
-3
vllm/model_executor/guided_decoding/xgrammar_decoding.py
+23
-3
@@ -3,7 +3,6 @@
3
# noqa: UP007
4
from __future__ import annotations
5
6
-import copy
6
import json
7
import re
8
from dataclasses import dataclass, field
@@ -348,5 +347,26 @@ class XGrammarLogitsProcessor:
347
return scores
348
349
def clone(self) -> XGrammarLogitsProcessor:
351
- """Deepcopy due to per-sequence state in the matchers"""
352
- return copy.deepcopy(self)
350
+ """Create a new instance with shared compiled grammar
351
+ but separate state"""
352
+ new_processor = XGrammarLogitsProcessor(self.config)
353
+
354
+ # Share the compiled grammar context (immutable after compilation)
355
+ new_processor.ctx = self.ctx
356
+
357
+ # Create fresh matchers for the new sequence
358
+ if self.ctx is not None:
359
+ new_processor.matchers = [
360
+ xgr.GrammarMatcher(self.ctx) for _ in range(self.batch_size)
361
+ ]
362
+
363
+ # Create a new token bitmask with the same size
364
+ if hasattr(self, 'token_bitmask') and self.token_bitmask is not None:
365
+ new_processor.token_bitmask = self.token_bitmask
366
+
367
+ # Copy simple attributes
368
+ new_processor.batch_size = self.batch_size
369
+ # Reset prefilled state for new sequence
370
+ new_processor.prefilled = False
371
+
372
+ return new_processor