@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 34120f5acd
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, 12:15 UTC

[V1][Feature] Enable Speculative Decoding with Structured Outputs (#14702)

Signed-off-by: Benjamin Chislett <benjamin.chislett@centml.ai> Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com>

Benjamin Chislett committed Apr 29, 2025 at 17:02 UTC 34120f5acd693924cc783dfaa33733afbe9ae8b0
9 files changed +209 -59
benchmarks/backend_request_func.py
+1
@@ -260,6 +260,7 @@ async def async_request_openai_completions(
260 if request_func_input.model_name else request_func_input.model,
261 "prompt": request_func_input.prompt,
262 "temperature": 0.0,
263 + "repetition_penalty": 1.0,
264 "max_tokens": request_func_input.output_len,
265 "logprobs": request_func_input.logprobs,
266 "stream": True,
benchmarks/benchmark_serving_structured_output.py
+6 -3
@@ -123,6 +123,8 @@ def sample_requests(tokenizer: PreTrainedTokenizerBase,
123 copy.deepcopy(schema) for _ in range(args.num_prompts)
124 ]
125 for i in range(len(json_schemas)):
126 + if "properties" not in json_schemas[i]:
127 + json_schemas[i]["properties"] = {}
128 json_schemas[i]["properties"][
129 f"__optional_field_{uuid.uuid4()}"] = {
130 "type":
@@ -134,7 +136,7 @@ def sample_requests(tokenizer: PreTrainedTokenizerBase,
136 json_schemas = [schema] * args.num_prompts
137
138 def gen_prompt(index: int):
137 - return f"Generate an example of a user profile given the following schema: {json.dumps(get_schema(index))}" # noqa: E501
139 + return f"Generate an example of a brief user profile given the following schema: {json.dumps(get_schema(index))}" # noqa: E501
140
141 def get_schema(index: int):
142 return json_schemas[index % len(json_schemas)]
@@ -231,7 +233,8 @@ def sample_requests(tokenizer: PreTrainedTokenizerBase,
233 idx -= len_dataset
234 schema = dataset["schema"][idx]
235 prompt = tokenizer.apply_chat_template(dataset["prompt"][idx],
234 - tokenize=False)
236 + tokenize=False,
237 + add_generation_prompt=True)
238 input_len = len(tokenizer(prompt).input_ids)
239 completion = dataset["completion"][idx]
240
@@ -849,7 +852,7 @@ if __name__ == "__main__":
852 'json', 'json-unique', 'grammar', 'regex',
853 'choice', 'xgrammar_bench'
854 ])
852 - parser.add_argument("--json_schema_path",
855 + parser.add_argument("--json-schema-path",
856 type=str,
857 default=None,
858 help="Path to json schema.")
tests/v1/entrypoints/llm/test_struct_output_generate.py
+28 -7
@@ -16,13 +16,31 @@ from vllm.outputs import RequestOutput
16 from vllm.platforms import current_platform
17 from vllm.sampling_params import GuidedDecodingParams, SamplingParams
18
19 +NGRAM_SPEC_CONFIG = {
20 + "model": "[ngram]",
21 + "num_speculative_tokens": 5,
22 + "prompt_lookup_max": 5,
23 + "prompt_lookup_min": 1,
24 +}
25 +
26 +EAGLE_SPEC_CONFIG = {
27 + "method": "eagle",
28 + "model": "yuhuili/EAGLE-LLaMA3.1-Instruct-8B",
29 + "num_speculative_tokens": 5,
30 +}
31 +
32 PARAMS_MODELS_BACKENDS_TOKENIZER_MODE = [
20 - ("mistralai/Ministral-8B-Instruct-2410", "xgrammar", "auto"),
21 - ("mistralai/Ministral-8B-Instruct-2410", "guidance", "auto"),
22 - ("mistralai/Ministral-8B-Instruct-2410", "xgrammar", "mistral"),
23 - ("Qwen/Qwen2.5-1.5B-Instruct", "xgrammar", "auto"),
33 + ("mistralai/Ministral-8B-Instruct-2410", "xgrammar", "auto", None),
34 + ("mistralai/Ministral-8B-Instruct-2410", "guidance", "auto", None),
35 + ("mistralai/Ministral-8B-Instruct-2410", "xgrammar", "mistral", None),
36 + ("Qwen/Qwen2.5-1.5B-Instruct", "xgrammar", "auto", None),
37 #FIXME: This test is flaky on CI thus disabled
38 #("Qwen/Qwen2.5-1.5B-Instruct", "guidance", "auto"),
39 + ("mistralai/Ministral-8B-Instruct-2410", "guidance", "auto",
40 + NGRAM_SPEC_CONFIG),
41 + ("Qwen/Qwen2.5-1.5B-Instruct", "xgrammar", "auto", NGRAM_SPEC_CONFIG),
42 + ("meta-llama/Meta-Llama-3.1-8B-Instruct", "xgrammar", "auto",
43 + EAGLE_SPEC_CONFIG)
44 ]
45
46 PARAMS_MODELS_TOKENIZER_MODE = [
@@ -45,8 +63,9 @@ class CarDescription(BaseModel):
63
64
65 @pytest.mark.skip_global_cleanup
48 -@pytest.mark.parametrize("model_name, guided_decoding_backend, tokenizer_mode",
49 - PARAMS_MODELS_BACKENDS_TOKENIZER_MODE)
66 +@pytest.mark.parametrize(
67 + "model_name, guided_decoding_backend, tokenizer_mode, speculative_config",
68 + PARAMS_MODELS_BACKENDS_TOKENIZER_MODE)
69 def test_structured_output(
70 monkeypatch: pytest.MonkeyPatch,
71 sample_json_schema: dict[str, Any],
@@ -58,6 +77,7 @@ def test_structured_output(
77 guided_decoding_backend: str,
78 tokenizer_mode: str,
79 model_name: str,
80 + speculative_config: dict[str, Any],
81 ):
82 monkeypatch.setenv("VLLM_USE_V1", "1")
83
@@ -71,7 +91,8 @@ def test_structured_output(
91 max_model_len=1024,
92 guided_decoding_backend=guided_decoding_backend,
93 guided_decoding_disable_any_whitespace=True,
74 - tokenizer_mode=tokenizer_mode)
94 + tokenizer_mode=tokenizer_mode,
95 + speculative_config=speculative_config)
96
97 #
98 # Test 1: Generate JSON output based on a provided schema
vllm/v1/core/sched/scheduler.py
+12 -5
@@ -441,7 +441,7 @@ class Scheduler(SchedulerInterface):
441 grammar_bitmask = self.structured_output_manager.grammar_bitmask(
442 self.requests,
443 structured_output_request_ids,
444 - len(self.running),
444 + scheduled_spec_decode_tokens,
445 )
446 # Construct the scheduler output.
447 new_reqs_data = [
@@ -682,10 +682,6 @@ class Scheduler(SchedulerInterface):
682 self.encoder_cache_manager.free_encoder_input(
683 request, input_id)
684
685 - # Add newly generated spec token ids to the request.
686 - if spec_token_ids is not None:
687 - request.spec_token_ids = spec_token_ids[req_index]
688 -
685 stopped = False
686 new_logprobs = None
687 new_token_ids = generated_token_ids
@@ -717,6 +713,17 @@ class Scheduler(SchedulerInterface):
713 request.structured_output_request.grammar.accept_tokens( # type: ignore[union-attr]
714 req_id, new_token_ids)
715
716 + # Add newly generated spec token ids to the request.
717 + if spec_token_ids is not None:
718 + if request.use_structured_output:
719 + metadata = request.structured_output_request
720 + assert metadata is not None and metadata.grammar is not None
721 + # Needs to happen after new_token_ids are accepted.
722 + request.spec_token_ids = metadata.grammar.validate_tokens(
723 + spec_token_ids[req_index])
724 + else:
725 + request.spec_token_ids = spec_token_ids[req_index]
726 +
727 # Get prompt logprobs for this request.
728 prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
729 if new_token_ids:
vllm/v1/structured_output/__init__.py
+46 -13
@@ -27,6 +27,7 @@ class StructuredOutputManager:
27 def __init__(self, vllm_config: VllmConfig):
28 self.backend: Optional[StructuredOutputBackend] = None
29 self.vllm_config = vllm_config
30 +
31 self._grammar_bitmask: Optional[torch.Tensor] = None
32
33 # The default max_workers if not specified is the number of CPUs * 5,
@@ -80,7 +81,7 @@ class StructuredOutputManager:
81 self,
82 requests: dict[str, Request],
83 structured_output_request_ids: dict[str, int],
83 - batch_len: int,
84 + scheduled_spec_decode_tokens: dict[str, list[int]],
85 ) -> Optional[npt.NDArray[np.int32]]:
86 # Prepare the structured output bitmask for this batch.
87 if not structured_output_request_ids:
@@ -88,20 +89,52 @@ class StructuredOutputManager:
89
90 if self._grammar_bitmask is None:
91 assert self.backend is not None
91 - self._grammar_bitmask = self.backend.allocate_token_bitmask(
92 - self.vllm_config.scheduler_config.max_num_seqs)
93 -
94 - # Fill the bitmask using the index of each request equal to its
95 - # position in the batch. Resize the bitmask down to the size of
96 - # the batch.
97 - bitmask_tensor = self._grammar_bitmask
98 - for req_id, batch_index in structured_output_request_ids.items():
92 + max_batch_size = self.vllm_config.scheduler_config.max_num_seqs
93 + if self.vllm_config.speculative_config is not None:
94 + max_num_spec_tokens = self.vllm_config.\
95 + speculative_config.num_speculative_tokens
96 + else:
97 + max_num_spec_tokens = 0
98 +
99 + # Allocate a bitmask for each token needing to be checked:
100 + # one for each speculative position, and one more for the
101 + # bonus token / non-speculative token.
102 + self._grammar_bitmask = \
103 + self.backend.allocate_token_bitmask(
104 + max_batch_size * (1 + max_num_spec_tokens))
105 +
106 + # Generate a batched bitmask for all structured output requests.
107 + # When speculative decoding is enabled, we need to include multiple
108 + # masks for each request, one for each possible bonus token position.
109 + # These are stored inline in the tensor and unpacked by the gpu runner.
110 + cumulative_index = 0
111 + ordered_seq = sorted(structured_output_request_ids.items(),
112 + key=lambda x: x[1])
113 + # NOTE: This outer loop can likely be parallelized to improve
114 + # performance of bitmask generation for large batches.
115 + for req_id, _ in ordered_seq:
116 request = requests[req_id].structured_output_request
117 assert request is not None and request.grammar is not None
101 - if not request.grammar.is_terminated():
102 - request.grammar.fill_bitmask(bitmask_tensor, batch_index)
103 - if batch_len < self._grammar_bitmask.shape[0]:
104 - bitmask_tensor = self._grammar_bitmask[:batch_len]
118 + state_advancements = 0
119 + req_tokens = scheduled_spec_decode_tokens.get(req_id, []) + [None]
120 + for i, token in enumerate(req_tokens):
121 + if not request.grammar.is_terminated():
122 + request.grammar.fill_bitmask(self._grammar_bitmask,
123 + cumulative_index)
124 + if token is not None:
125 + # In order to generate the correct bitmask for each
126 + # position in the speculative sequence, we advance
127 + # the FSM state for each speculative token and rollback
128 + # to restore the previous state when we are finished.
129 + assert request.grammar.accept_tokens(req_id, [token])
130 + state_advancements += 1
131 + cumulative_index += 1
132 + if state_advancements > 0:
133 + request.grammar.rollback(state_advancements)
134 +
135 + bitmask_tensor = self._grammar_bitmask
136 + if cumulative_index < self._grammar_bitmask.shape[0]:
137 + bitmask_tensor = self._grammar_bitmask[:cumulative_index]
138
139 # After finishing with the xgrammar operations, we convert to
140 # np.ndarray, because that is much more efficient for serialization
vllm/v1/structured_output/backend_guidance.py
+21
@@ -144,6 +144,27 @@ class GuidanceGrammar(StructuredOutputGrammar):
144
145 return r
146
147 + def validate_tokens(self, tokens: list[int]) -> list[int]:
148 + """Checks if the list of tokens are accepted by the parser in sequence.
149 + Will not advance the parser.
150 +
151 + Returns the prefix list of tokens that are accepted by the parser.
152 + """
153 + if len(tokens) == 0:
154 + return []
155 + if self.ll_matcher.is_stopped():
156 + return []
157 +
158 + num_tokens = self.ll_matcher.validate_tokens(tokens)
159 +
160 + self.check_error()
161 +
162 + return tokens[:num_tokens]
163 +
164 + def rollback(self, num_tokens: int) -> None:
165 + self.ll_matcher.rollback(num_tokens)
166 + self.check_error()
167 +
168 def fill_bitmask(self, bitmask: torch.Tensor, idx: int) -> None:
169 # this will automatically return [EOS] mask if the matcher is stopped
170 # or otherwise in an error state
vllm/v1/structured_output/backend_types.py
+24
@@ -35,6 +35,30 @@ class StructuredOutputGrammar(ABC):
35 bool: True if the tokens are accepted, False otherwise.
36 """
37
38 + @abstractmethod
39 + def validate_tokens(self, tokens: list[int]) -> list[int]:
40 + """
41 + Validates the provided tokens against the grammar.
42 + Will not advance the FSM.
43 +
44 + Args:
45 + tokens (list[int]): A list of token IDs to validate.
46 +
47 + Returns:
48 + list[int]: A list of accepted token IDs. Will be a prefix
49 + of the input tokens, and empty if none are accepted.
50 + """
51 +
52 + @abstractmethod
53 + def rollback(self, num_tokens: int) -> None:
54 + """
55 + Rolls back the state of the grammar by a specified number of tokens.
56 + Will also revert counters for the number of processed tokens.
57 +
58 + Args:
59 + num_tokens (int): The number of tokens to roll back.
60 + """
61 +
62 @abstractmethod
63 def fill_bitmask(self, bitmask: torch.Tensor, batch_index: int) -> None:
64 """
vllm/v1/structured_output/backend_xgrammar.py
+30 -2
@@ -40,6 +40,11 @@ class XgrammarBackend(StructuredOutputBackend):
40 self.disable_any_whitespace = \
41 vllm_config.decoding_config.disable_any_whitespace
42
43 + self.num_speculative_tokens = 0
44 + if self.vllm_config.speculative_config is not None:
45 + self.num_speculative_tokens = \
46 + self.vllm_config.speculative_config.num_speculative_tokens
47 +
48 tokenizer = tokenizer_group.get_lora_tokenizer(None)
49 self.vocab_size = vllm_config.model_config.get_vocab_size()
50 if isinstance(tokenizer, MistralTokenizer):
@@ -118,7 +123,10 @@ class XgrammarBackend(StructuredOutputBackend):
123 f"grammar is not of valid supported types. ({request_type!s})")
124
125 return XgrammarGrammar(
121 - matcher=xgr.GrammarMatcher(ctx),
126 + matcher=xgr.GrammarMatcher(
127 + ctx,
128 + max_rollback_tokens=self.num_speculative_tokens,
129 + ),
130 vocab_size=self.vocab_size,
131 ctx=ctx,
132 )
@@ -136,7 +144,6 @@ class XgrammarGrammar(StructuredOutputGrammar):
144 # supporting different backends, in the future.
145 # For now, just xgrammar.
146 #
139 - # TODO: support max_rollback_tokens
147 # https://xgrammar.mlc.ai/docs/api/python/index.html#xgrammar.GrammarMatcher.find_jump_forward_string
148 # for jump-forward decoding
149
@@ -163,6 +170,27 @@ class XgrammarGrammar(StructuredOutputGrammar):
170 self.num_processed_tokens += 1
171 return True
172
173 + def validate_tokens(self, tokens: list[int]) -> list[int]:
174 + """Checks if the list of tokens are accepted by the FSM in sequence.
175 + Will not advance the FSM.
176 +
177 + Returns the prefix list of tokens that are accepted by the FSM.
178 + """
179 + accepted_tokens = []
180 + for token in tokens:
181 + if self.matcher.accept_token(token):
182 + accepted_tokens.append(token)
183 + else:
184 + break
185 + if len(accepted_tokens) > 0:
186 + # Rollback the FSM to the initial state
187 + self.matcher.rollback(len(accepted_tokens))
188 + return accepted_tokens
189 +
190 + def rollback(self, num_tokens: int) -> None:
191 + self.matcher.rollback(num_tokens)
192 + self.num_processed_tokens -= num_tokens
193 +
194 def fill_bitmask(self, bitmask: torch.Tensor, idx: int) -> None:
195 self.matcher.fill_next_token_bitmask(bitmask, idx)
196
vllm/v1/worker/gpu_model_runner.py
+41 -29
@@ -957,46 +957,58 @@ class GPUModelRunner(LoRAModelRunnerMixin):
957 scheduler_output: "SchedulerOutput",
958 logits: torch.Tensor,
959 ):
960 - # Serialization of np.ndarray is much more efficient than a tensor,
961 - # so we receive it in that format.
960 grammar_bitmask = scheduler_output.grammar_bitmask
961 if grammar_bitmask is None:
962 return
963
966 - # We receive the structured output bitmask from the scheduler, but the
967 - # indices of the requests in the batch may not match the indices of
968 - # the bitmask since the scheduler doesn't know how the gpu runner is
969 - # ordering the requests in the batch. We need to sort the bitmask to
970 - # match the order of the requests used here.
964 + # We receive the structured output bitmask from the scheduler,
965 + # compacted to contain bitmasks only for structured output requests.
966 + # The order of the requests in the bitmask is not guaranteed to be the
967 + # same as the order of the requests in the gpu runner's batch. We need
968 + # to sort the bitmask to match the order of the requests used here.
969 +
970 + # Get the batch indices of the structured output requests.
971 + # Keep track of the number of speculative tokens scheduled for every
972 + # request in the batch, as the logit indices are offset by this amount.
973 struct_out_req_batch_indices: dict[str, int] = {}
972 - indices_match = True
973 - for req_id in self.input_batch.req_ids:
974 - mask_index = scheduler_output.structured_output_request_ids.get(
975 - req_id)
976 - if mask_index is None:
977 - # not a structured output request
978 - continue
979 - batch_index = self.input_batch.req_id_to_index[req_id]
980 - if batch_index != mask_index:
981 - indices_match = False
982 - struct_out_req_batch_indices[req_id] = batch_index
983 -
984 - if not indices_match:
985 - # Sort the bitmask to match the order of the requests
986 - sorted_bitmask = np.zeros_like(grammar_bitmask)
987 - for req_id, batch_index in struct_out_req_batch_indices.items():
988 - orig_index = scheduler_output.structured_output_request_ids[
989 - req_id]
990 - sorted_bitmask[batch_index] = grammar_bitmask[orig_index]
991 - grammar_bitmask = sorted_bitmask
974 + cumulative_offset = 0
975 + seq = sorted(self.input_batch.req_id_to_index.items(),
976 + key=lambda x: x[1])
977 + for req_id, batch_index in seq:
978 + logit_index = batch_index + cumulative_offset
979 + cumulative_offset += len(
980 + scheduler_output.scheduled_spec_decode_tokens.get(req_id, []))
981 + if req_id in scheduler_output.structured_output_request_ids:
982 + struct_out_req_batch_indices[req_id] = logit_index
983 +
984 + out_indices = []
985 +
986 + # Reorder the bitmask to match the order of the requests in the batch.
987 + sorted_bitmask = np.zeros_like(grammar_bitmask,
988 + shape=(logits.shape[0],
989 + grammar_bitmask.shape[1]))
990 + cumulative_index = 0
991 + seq = sorted(scheduler_output.structured_output_request_ids.items(),
992 + key=lambda x: x[1])
993 + for req_id, _ in seq:
994 + logit_index = struct_out_req_batch_indices[req_id]
995 + num_spec_tokens = len(
996 + scheduler_output.scheduled_spec_decode_tokens.get(req_id, []))
997 + for i in range(1 + num_spec_tokens):
998 + sorted_bitmask[logit_index + i] = \
999 + grammar_bitmask[cumulative_index + i]
1000 + out_indices.append(logit_index + i)
1001 + cumulative_index += 1 + num_spec_tokens
1002 + grammar_bitmask = sorted_bitmask
1003
1004 + # Serialization of np.ndarray is much more efficient than a tensor,
1005 + # so we receive it in that format.
1006 grammar_bitmask = torch.from_numpy(grammar_bitmask)
1007
995 - # TODO: compatibility with spec decode
1008 xgr.apply_token_bitmask_inplace(
1009 logits,
1010 grammar_bitmask.to(self.device, non_blocking=True),
999 - indices=list(struct_out_req_batch_indices.values()),
1011 + indices=out_indices,
1012 )
1013
1014 @torch.inference_mode()