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, 08:30 UTC
[V0][Bugfix] Fix parallel sampling performance regression when guided decoding is enabled (#17731)
Signed-off-by: Madeesh Kannan <shadeMe@users.noreply.github.com> Co-authored-by: Russell Bryant <rbryant@redhat.com>
Madeesh Kannan committed
May 23, 2025 at 12:38 UTC
e493e48524e9e78ab33eafec6461b3940e361189
4 files changed
+40
-8
vllm/model_executor/guided_decoding/guidance_logits_processors.py
+22
-4
@@ -1,4 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+import copy
3
import os
4
from typing import Any
5
@@ -34,9 +35,24 @@ class GuidanceLogitsProcessor:
35
self.grammar = grammar
36
self.tokenizer = tokenizer
37
self.tokenizer_name = tokenizer.name_or_path
38
+ self.ll_tokenizer = None
39
+ self.ll_matcher = None
40
+ self.bitmask = None
41
self.new_sampling = False
42
self.initialized = False
43
44
+ def clone(self) -> "GuidanceLogitsProcessor":
45
+ cloned = copy.copy(self)
46
+ if self.initialized:
47
+ cloned.ll_matcher = llguidance.LLMatcher(
48
+ self.ll_tokenizer, # type: ignore[assignment]
49
+ self.grammar,
50
+ log_level=int(os.environ.get("LLGUIDANCE_LOG_LEVEL", "1")),
51
+ )
52
+ self.bitmask = llguidance.torch.allocate_token_bitmask(
53
+ 1, self.ll_tokenizer.vocab_size) # type: ignore[attr-defined]
54
+ return cloned
55
+
56
def _initialize(self):
57
if self.initialized:
58
return
@@ -56,7 +72,7 @@ class GuidanceLogitsProcessor:
72
73
# create reusable bitmask
74
self.bitmask = llguidance.torch.allocate_token_bitmask(
59
- 1, self.ll_tokenizer.vocab_size)
75
+ 1, self.ll_tokenizer.vocab_size) # type: ignore[attr-defined]
76
77
self.initialized = True
78
@@ -70,15 +86,17 @@ class GuidanceLogitsProcessor:
86
self._initialize()
87
88
if self.new_sampling and len(input_ids) > 0:
73
- self.ll_matcher.consume_token(input_ids[-1])
74
- err = self.ll_matcher.get_error()
89
+ self.ll_matcher.consume_token( # type: ignore[attr-defined]
90
+ input_ids[-1])
91
+ err = self.ll_matcher.get_error() # type: ignore[attr-defined]
92
if err:
93
logger.warning("Error in LLMatcher: %s", err)
94
95
llguidance.torch.fill_next_token_bitmask(self.ll_matcher, self.bitmask,
96
0)
97
llguidance.torch.apply_token_bitmask_inplace(
81
- scores, self.bitmask.to(scores.device))
98
+ scores,
99
+ self.bitmask.to(scores.device)) # type: ignore[attr-defined]
100
101
self.new_sampling = True
102
vllm/model_executor/guided_decoding/outlines_logits_processors.py
+12
@@ -56,6 +56,12 @@ class BaseLogitsProcessor:
56
self._fsm_state: defaultdict[int, Union[int,
57
CFGState]] = defaultdict(int)
58
59
+ def clone(self) -> "BaseLogitsProcessor":
60
+ cloned = copy.copy(self)
61
+ cloned._guide = self._guide.copy()
62
+ cloned._fsm_state = copy.deepcopy(self._fsm_state)
63
+ return cloned
64
+
65
def __call__(self, input_ids: list[int],
66
scores: torch.Tensor) -> torch.Tensor:
67
"""Use the FSM to bias the logits before sampling the next token."""
@@ -218,6 +224,12 @@ class CFGLogitsProcessor(BaseLogitsProcessor):
224
reasoner)
225
self._guide = self._guide.copy()
226
227
+ def clone(self) -> "CFGLogitsProcessor":
228
+ cloned = copy.copy(self)
229
+ cloned._fsm_state = copy.deepcopy(self._fsm_state)
230
+ cloned._guide = self._guide.copy()
231
+ return cloned
232
+
233
234
@lru_cache(maxsize=32)
235
def _adapt_tokenizer(tokenizer: PreTrainedTokenizerBase):
vllm/model_executor/guided_decoding/xgrammar_decoding.py
+5
-3
@@ -302,8 +302,9 @@ class XGrammarLogitsProcessor:
302
prefilled: bool = field(default=False)
303
304
def __post_init__(self):
305
- self.tokenizer_info = self.config.tokenizer_info(
306
- self.config.tokenizer_data)
305
+ if self.tokenizer_info is None:
306
+ self.tokenizer_info = self.config.tokenizer_info(
307
+ self.config.tokenizer_data)
308
309
def __getstate__(self) -> dict[str, Any]:
310
return {'config': self.config, 'reasoner': self.reasoner}
@@ -400,7 +401,8 @@ class XGrammarLogitsProcessor:
401
def clone(self) -> XGrammarLogitsProcessor:
402
"""Create a new instance with shared compiled grammar
403
but separate state"""
403
- new_processor = XGrammarLogitsProcessor(self.config, self.reasoner)
404
+ new_processor = XGrammarLogitsProcessor(self.config, self.reasoner,
405
+ None, self.tokenizer_info)
406
407
# Share the compiled grammar context (immutable after compilation)
408
new_processor.ctx = self.ctx
vllm/sequence.py
+1
-1
@@ -1494,7 +1494,7 @@ class ParallelSampleSequenceGroup(SequenceGroupBase):
1494
for i in range(original_params.n):
1495
request_id_i = f"{request_id}_parallel_sample_{i}"
1496
group.seq_id_to_index[request_id_i] = i
1497
- params = copy.deepcopy(original_params)
1497
+ params = params.clone()
1498
params.n = 1
1499
if params.seed is not None:
1500
params.seed += i