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 14, 05:30 UTC
[Structured Outputs] Refactor bitmask construction into get_grammar_bitmask (#23361)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Aug 21, 2025 at 13:53 UTC
800349c2a50ac1f823d3387cb0f44d0bd6d470d1
1 file changed
+33
-22
vllm/v1/core/sched/scheduler.py
+33
-22
@@ -177,14 +177,6 @@ class Scheduler(SchedulerInterface):
177
scheduled_running_reqs: list[Request] = []
178
preempted_reqs: list[Request] = []
179
180
- # NOTE: structured_output_request_ids maps
181
- # a request's (request that uses structured output)
182
- # request_id to the running request index.
183
- # This will helps us determine to slice the grammar bitmask
184
- # and only applies valid mask for requests that
185
- # uses structured decoding.
186
- structured_output_request_ids: dict[str, int] = {}
187
-
180
req_to_new_blocks: dict[str, KVCacheBlocks] = {}
181
num_scheduled_tokens: dict[str, int] = {}
182
token_budget = self.max_num_scheduled_tokens
@@ -282,12 +274,6 @@ class Scheduler(SchedulerInterface):
274
275
# Schedule the request.
276
scheduled_running_reqs.append(request)
285
- if request.use_structured_output:
286
- # PERF: in case of chunked prefill,
287
- # request might not include any new tokens.
288
- # Therefore, we might introduce some additional
289
- # cycle to fill in the bitmask, which could be a big no-op.
290
- structured_output_request_ids[request.request_id] = req_index
277
req_to_new_blocks[request.request_id] = new_blocks
278
num_scheduled_tokens[request.request_id] = num_new_tokens
279
token_budget -= num_new_tokens
@@ -477,9 +463,6 @@ class Scheduler(SchedulerInterface):
463
request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
464
continue
465
480
- if request.use_structured_output:
481
- structured_output_request_ids[request.request_id] = (
482
- req_index)
466
req_index += 1
467
self.running.append(request)
468
if self.log_stats:
@@ -538,11 +521,6 @@ class Scheduler(SchedulerInterface):
521
self.kv_cache_manager.get_num_common_prefix_blocks(
522
any_request, len(self.running)))
523
541
- grammar_bitmask = self.structured_output_manager.grammar_bitmask(
542
- self.requests,
543
- structured_output_request_ids,
544
- scheduled_spec_decode_tokens,
545
- )
524
# Construct the scheduler output.
525
new_reqs_data = [
526
NewRequestData.from_request(
@@ -556,6 +534,9 @@ class Scheduler(SchedulerInterface):
534
scheduled_spec_decode_tokens,
535
req_to_new_blocks,
536
)
537
+ structured_output_request_ids, grammar_bitmask = (
538
+ self.get_grammar_bitmask(self.running,
539
+ scheduled_spec_decode_tokens))
540
scheduler_output = SchedulerOutput(
541
scheduled_new_reqs=new_reqs_data,
542
scheduled_cached_reqs=cached_reqs_data,
@@ -753,6 +734,36 @@ class Scheduler(SchedulerInterface):
734
encoder_inputs_to_schedule.append(i)
735
return encoder_inputs_to_schedule, num_new_tokens, encoder_budget
736
737
+ def get_grammar_bitmask(
738
+ self,
739
+ requests: list[Request],
740
+ scheduled_spec_decode_tokens: dict[str, list[int]],
741
+ ):
742
+ # NOTE: structured_output_request_ids maps
743
+ # a request's (request that uses structured output)
744
+ # request_id to its index in the batch.
745
+ # This will helps us determine to slice the grammar bitmask
746
+ # and only applies valid mask for requests that
747
+ # uses structured decoding.
748
+ structured_output_request_ids: dict[str, int] = {}
749
+ for i, req in enumerate(requests):
750
+ if req.use_structured_output:
751
+ # PERF: in case of chunked prefill,
752
+ # request might not include any new tokens.
753
+ # Therefore, we might introduce some additional
754
+ # cycle to fill in the bitmask, which could be a big no-op.
755
+ structured_output_request_ids[req.request_id] = i
756
+
757
+ if not structured_output_request_ids:
758
+ bitmask = None
759
+ else:
760
+ bitmask = self.structured_output_manager.grammar_bitmask(
761
+ self.requests,
762
+ structured_output_request_ids,
763
+ scheduled_spec_decode_tokens,
764
+ )
765
+ return structured_output_request_ids, bitmask
766
+
767
def update_from_output(
768
self,
769
scheduler_output: SchedulerOutput,