@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e206b54331
Fork 0

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