@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4aed506b65

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 14, 05:30 UTC

[Core] Streamline some structured output related code (#26737)

Signed-off-by: Nick Hill <nhill@redhat.com>

Nick Hill committed Oct 14, 2025 at 16:27 UTC 4aed506b6538ec4f284c480bf4449e9dc5f72054

13 files changed +119 -136

tests/v1/core/test_scheduler.py

+8 -10

@@ -30,7 +30,6 @@ from vllm.v1.kv_cache_interface import (

30 from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput

31 from vllm.v1.request import Request, RequestStatus

32 from vllm.v1.structured_output import StructuredOutputManager

33 -from vllm.v1.structured_output.request import StructuredOutputRequest

34 from .utils import EOS_TOKEN_ID, create_requests, create_scheduler

@@ -335,10 +334,10 @@ def test_stop_via_update_from_output():

334 requests[0].request_id: [],

335 requests[1].request_id: [10],

336 },

338 - num_common_prefix_blocks=0,

337 + num_common_prefix_blocks=[],

338 finished_req_ids=set(),

339 free_encoder_mm_hashes=[],

341 - structured_output_request_ids={},

340 + structured_output_request_ids=[],

341 grammar_bitmask=None,

342 )

343

@@ -383,10 +382,10 @@ def test_stop_via_update_from_output():

382 requests[0].request_id: [10, 42],

383 requests[1].request_id: [13],

384 },

386 - num_common_prefix_blocks=0,

385 + num_common_prefix_blocks=[],

386 finished_req_ids=set(),

387 free_encoder_mm_hashes=[],

389 - structured_output_request_ids={},

388 + structured_output_request_ids=[],

389 grammar_bitmask=None,

390 )

391

@@ -429,10 +428,10 @@ def test_stop_via_update_from_output():

428 requests[0].request_id: [10, 11],

429 requests[1].request_id: [],

430 },

432 - num_common_prefix_blocks=0,

431 + num_common_prefix_blocks=[],

432 finished_req_ids=set(),

433 free_encoder_mm_hashes=[],

435 - structured_output_request_ids={},

434 + structured_output_request_ids=[],

435 grammar_bitmask=None,

436 )

437

@@ -470,10 +469,10 @@ def test_stop_via_update_from_output():

469 total_num_scheduled_tokens=3,

470 scheduled_encoder_inputs={},

471 scheduled_spec_decode_tokens={requests[0].request_id: [EOS_TOKEN_ID, 10]},

473 - num_common_prefix_blocks=0,

472 + num_common_prefix_blocks=[],

473 finished_req_ids=set(),

474 free_encoder_mm_hashes=[],

476 - structured_output_request_ids={},

475 + structured_output_request_ids=[],

476 grammar_bitmask=None,

477 )

478

@@ -1941,7 +1940,6 @@ def test_schedule_skip_tokenizer_init_structured_output_request():

1940 sampling_params=sampling_params,

1941 pooling_params=None,

1942 eos_token_id=EOS_TOKEN_ID,

1944 - structured_output_request=StructuredOutputRequest(sampling_params),

1943 )

1944 scheduler.add_request(request)

1945 output = scheduler.schedule()

tests/v1/kv_connector/unit/test_kv_connector_lifecyle.py

+1 -1

@@ -26,7 +26,7 @@ def _make_empty_scheduler_output():

26 num_common_prefix_blocks=[],

27 finished_req_ids=set(),

28 free_encoder_mm_hashes=[],

29 - structured_output_request_ids={},

29 + structured_output_request_ids=[],

30 grammar_bitmask=None,

31 kv_connector_metadata=SharedStorageConnectorMetadata(),

32 )

tests/v1/tpu/worker/test_tpu_model_runner.py

+12 -12

@@ -89,10 +89,10 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:

89 total_num_scheduled_tokens=total_num_scheduled_tokens,

90 scheduled_spec_decode_tokens={},

91 scheduled_encoder_inputs={},

92 - num_common_prefix_blocks=0,

92 + num_common_prefix_blocks=[],

93 finished_req_ids=set(),

94 free_encoder_mm_hashes=[],

95 - structured_output_request_ids={},

95 + structured_output_request_ids=[],

96 grammar_bitmask=None,

97 )

@@ -168,10 +168,10 @@ def test_update_states_request_finished(model_runner):

168 total_num_scheduled_tokens=0,

169 scheduled_spec_decode_tokens={},

170 scheduled_encoder_inputs={},

171 - num_common_prefix_blocks=0,

171 + num_common_prefix_blocks=[],

172 finished_req_ids={req_id},

173 free_encoder_mm_hashes=[],

174 - structured_output_request_ids={},

174 + structured_output_request_ids=[],

175 grammar_bitmask=None,

176 )

177

@@ -198,10 +198,10 @@ def test_update_states_request_resumed(model_runner):

198 total_num_scheduled_tokens=0,

199 scheduled_spec_decode_tokens={},

200 scheduled_encoder_inputs={},

201 - num_common_prefix_blocks=0,

201 + num_common_prefix_blocks=[],

202 finished_req_ids=set(),

203 free_encoder_mm_hashes=[],

204 - structured_output_request_ids={},

204 + structured_output_request_ids=[],

205 grammar_bitmask=None,

206 )

207

@@ -225,10 +225,10 @@ def test_update_states_request_resumed(model_runner):

225 total_num_scheduled_tokens=1,

226 scheduled_spec_decode_tokens={},

227 scheduled_encoder_inputs={},

228 - num_common_prefix_blocks=0,

228 + num_common_prefix_blocks=[],

229 finished_req_ids=set(),

230 free_encoder_mm_hashes=[],

231 - structured_output_request_ids={},

231 + structured_output_request_ids=[],

232 grammar_bitmask=None,

233 )

234

@@ -256,10 +256,10 @@ def test_update_states_no_changes(model_runner):

256 total_num_scheduled_tokens=1,

257 scheduled_spec_decode_tokens={},

258 scheduled_encoder_inputs={},

259 - num_common_prefix_blocks=0,

259 + num_common_prefix_blocks=[],

260 finished_req_ids=set(),

261 free_encoder_mm_hashes=[],

262 - structured_output_request_ids={},

262 + structured_output_request_ids=[],

263 grammar_bitmask=None,

264 )

265

@@ -291,10 +291,10 @@ def test_update_states_request_unscheduled(model_runner):

291 total_num_scheduled_tokens=1,

292 scheduled_spec_decode_tokens={},

293 scheduled_encoder_inputs={},

294 - num_common_prefix_blocks=0,

294 + num_common_prefix_blocks=[],

295 finished_req_ids=set(),

296 free_encoder_mm_hashes=[],

297 - structured_output_request_ids={},

297 + structured_output_request_ids=[],

298 grammar_bitmask=None,

299 )

300

tests/v1/worker/test_gpu_model_runner.py

+12 -12

@@ -146,10 +146,10 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:

146 total_num_scheduled_tokens=total_num_scheduled_tokens,

147 scheduled_spec_decode_tokens={},

148 scheduled_encoder_inputs={},

149 - num_common_prefix_blocks=0,

149 + num_common_prefix_blocks=[],

150 finished_req_ids=set(),

151 free_encoder_mm_hashes=[],

152 - structured_output_request_ids={},

152 + structured_output_request_ids=[],

153 grammar_bitmask=None,

154 )

155

@@ -212,10 +212,10 @@ def test_update_states_request_finished(model_runner, dist_init):

212 total_num_scheduled_tokens=0,

213 scheduled_spec_decode_tokens={},

214 scheduled_encoder_inputs={},

215 - num_common_prefix_blocks=0,

215 + num_common_prefix_blocks=[],

216 finished_req_ids={req_id},

217 free_encoder_mm_hashes=[],

218 - structured_output_request_ids={},

218 + structured_output_request_ids=[],

219 grammar_bitmask=None,

220 )

221

@@ -244,10 +244,10 @@ def test_update_states_request_resumed(model_runner, dist_init):

244 total_num_scheduled_tokens=0,

245 scheduled_spec_decode_tokens={},

246 scheduled_encoder_inputs={},

247 - num_common_prefix_blocks=0,

247 + num_common_prefix_blocks=[],

248 finished_req_ids=set(),

249 free_encoder_mm_hashes=[],

250 - structured_output_request_ids={},

250 + structured_output_request_ids=[],

251 grammar_bitmask=None,

252 )

253

@@ -273,10 +273,10 @@ def test_update_states_request_resumed(model_runner, dist_init):

273 total_num_scheduled_tokens=1,

274 scheduled_spec_decode_tokens={},

275 scheduled_encoder_inputs={},

276 - num_common_prefix_blocks=0,

276 + num_common_prefix_blocks=[],

277 finished_req_ids=set(),

278 free_encoder_mm_hashes=[],

279 - structured_output_request_ids={},

279 + structured_output_request_ids=[],

280 grammar_bitmask=None,

281 )

282

@@ -366,10 +366,10 @@ def test_update_states_no_changes(model_runner, dist_init):

366 total_num_scheduled_tokens=1,

367 scheduled_spec_decode_tokens={},

368 scheduled_encoder_inputs={},

369 - num_common_prefix_blocks=0,

369 + num_common_prefix_blocks=[],

370 finished_req_ids=set(),

371 free_encoder_mm_hashes=[],

372 - structured_output_request_ids={},

372 + structured_output_request_ids=[],

373 grammar_bitmask=None,

374 )

375

@@ -403,10 +403,10 @@ def test_update_states_request_unscheduled(model_runner, dist_init):

403 total_num_scheduled_tokens=1,

404 scheduled_spec_decode_tokens={},

405 scheduled_encoder_inputs={},

406 - num_common_prefix_blocks=0,

406 + num_common_prefix_blocks=[],

407 finished_req_ids=set(),

408 free_encoder_mm_hashes=[],

409 - structured_output_request_ids={},

409 + structured_output_request_ids=[],

410 grammar_bitmask=None,

411 )

412

vllm/v1/core/sched/output.py

+2 -3

@@ -165,9 +165,8 @@ class SchedulerOutput:

165 # freed from the encoder cache.

166 free_encoder_mm_hashes: list[str]

167

168 - # Dict of request ids to their index within the batch

169 - # for filling the next token bitmask

170 - structured_output_request_ids: dict[str, int]

168 + # ids of structured outputs requests included in the bitmask, in order.

169 + structured_output_request_ids: list[str]

170 # the bitmask for the whole batch

171 grammar_bitmask: "npt.NDArray[np.int32] | None"

172

vllm/v1/core/sched/scheduler.py

+30 -35

@@ -5,7 +5,7 @@ import itertools

5 import time

6 from collections import defaultdict

7 from collections.abc import Iterable

8 -from typing import Any

8 +from typing import TYPE_CHECKING, Any

10 from vllm.config import VllmConfig

11 from vllm.distributed.kv_events import EventPublisherFactory, KVEventBatch

@@ -34,6 +34,10 @@ from vllm.v1.request import Request, RequestStatus

34 from vllm.v1.spec_decode.metrics import SpecDecodingStats

35 from vllm.v1.structured_output import StructuredOutputManager

37 +if TYPE_CHECKING:

38 + import numpy as np

39 + import numpy.typing as npt

40 +

41 logger = init_logger(__name__)

@@ -608,11 +612,8 @@ class Scheduler(SchedulerInterface):

612 scheduled_spec_decode_tokens,

613 req_to_new_blocks,

614 )

611 - scheduled_requests = (

612 - scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs

613 - )

615 structured_output_request_ids, grammar_bitmask = self.get_grammar_bitmask(

615 - scheduled_requests, scheduled_spec_decode_tokens

616 + num_scheduled_tokens.keys(), scheduled_spec_decode_tokens

617 )

618 scheduler_output = SchedulerOutput(

619 scheduled_new_reqs=new_reqs_data,

@@ -876,32 +877,28 @@ class Scheduler(SchedulerInterface):

877

878 def get_grammar_bitmask(

879 self,

879 - requests: list[Request],

880 + scheduled_request_ids: Iterable[str],

881 scheduled_spec_decode_tokens: dict[str, list[int]],

881 - ):

882 - # NOTE: structured_output_request_ids maps

883 - # a request's (request that uses structured output)

884 - # request_id to its index in the batch.

885 - # This will help us determine to slice the grammar bitmask

886 - # and only applies valid mask for requests that

887 - # uses structured decoding.

888 - structured_output_request_ids: dict[str, int] = {}

889 - for i, req in enumerate(requests):

890 - if req.use_structured_output:

891 - # PERF: in case of chunked prefill,

892 - # request might not include any new tokens.

893 - # Therefore, we might introduce some additional

894 - # cycle to fill in the bitmask, which could be a big no-op.

895 - structured_output_request_ids[req.request_id] = i

896 -

882 + ) -> tuple[list[str], "npt.NDArray[np.int32] | None"]:

883 + # Collect list of scheduled request ids that use structured output.

884 + # The corresponding rows of the bitmask will be in this order.

885 + # PERF: in case of chunked prefill,

886 + # request might not include any new tokens.

887 + # Therefore, we might introduce some additional

888 + # cycle to fill in the bitmask, which could be a big no-op.

889 + structured_output_request_ids = [

890 + req_id

891 + for req_id in scheduled_request_ids

892 + if (req := self.requests.get(req_id)) and req.use_structured_output

893 + ]

894 if not structured_output_request_ids:

898 - bitmask = None

899 - else:

900 - bitmask = self.structured_output_manager.grammar_bitmask(

901 - self.requests,

902 - structured_output_request_ids,

903 - scheduled_spec_decode_tokens,

904 - )

895 + return structured_output_request_ids, None

896 +

897 + bitmask = self.structured_output_manager.grammar_bitmask(

898 + self.requests,

899 + structured_output_request_ids,

900 + scheduled_spec_decode_tokens,

901 + )

902 return structured_output_request_ids, bitmask

903

904 def update_from_output(

@@ -1013,12 +1010,10 @@ class Scheduler(SchedulerInterface):

1010 new_logprobs = logprobs.slice(req_index, req_index + 1)

1011

1012 if new_token_ids and self.structured_output_manager.should_advance(request):

1016 - # NOTE: structured_output_request

1017 - # should not be None if use_structured_output, we have

1018 - # checked above, so safe to ignore type warning

1019 - request.structured_output_request.grammar.accept_tokens( # type: ignore[union-attr]

1020 - req_id, new_token_ids

1021 - )

1013 + struct_output_request = request.structured_output_request

1014 + assert struct_output_request is not None

1015 + assert struct_output_request.grammar is not None

1016 + struct_output_request.grammar.accept_tokens(req_id, new_token_ids)

1017

1018 if num_nans_in_logits is not None and req_id in num_nans_in_logits:

1019 request.num_nans_in_logits = num_nans_in_logits[req_id]

vllm/v1/request.py

+8 -10

@@ -40,7 +40,6 @@ class Request:

40 prompt_embeds: torch.Tensor | None = None,

41 mm_features: list[MultiModalFeatureSpec] | None = None,

42 lora_request: Optional["LoRARequest"] = None,

43 - structured_output_request: Optional["StructuredOutputRequest"] = None,

43 cache_salt: str | None = None,

44 priority: int = 0,

45 trace_headers: Mapping[str, str] | None = None,

@@ -54,11 +53,12 @@ class Request:

53 # Because of LoRA, the eos token id can be different for each request.

54 self.eos_token_id = eos_token_id

55 self.lora_request = lora_request

57 - self.structured_output_request = structured_output_request

56 + self.structured_output_request = StructuredOutputRequest.from_sampling_params(

57 + sampling_params

58 + )

59 self.arrival_time = arrival_time if arrival_time is not None else time.time()

61 self.status = RequestStatus.WAITING

61 - self.use_structured_output = False

62 self.events: list[EngineCoreEvent] = []

63 self.stop_reason: int | str | None = None

@@ -72,9 +72,8 @@ class Request:

72 # Generative models.

73 assert sampling_params.max_tokens is not None

74 self.max_tokens = sampling_params.max_tokens

75 - if sampling_params.structured_outputs is not None:

75 + if self.structured_output_request is not None:

76 self.status = RequestStatus.WAITING_FOR_FSM

77 - self.use_structured_output = True

78 if sampling_params.extra_args is not None:

79 self.kv_transfer_params = sampling_params.extra_args.get(

@@ -145,11 +144,6 @@ class Request:

144 eos_token_id=request.eos_token_id,

145 arrival_time=request.arrival_time,

146 lora_request=request.lora_request,

148 - structured_output_request=StructuredOutputRequest(

149 - sampling_params=request.sampling_params

150 - )

151 - if request.sampling_params

152 - else None,

147 cache_salt=request.cache_salt,

148 priority=request.priority,

149 trace_headers=request.trace_headers,

@@ -170,6 +164,10 @@ class Request:

164 if self.get_hash_new_full_blocks is not None:

165 self.block_hashes.extend(self.get_hash_new_full_blocks())

166

167 + @property

168 + def use_structured_output(self) -> bool:

169 + return self.structured_output_request is not None

170 +

171 @property

172 def is_output_corrupted(self) -> bool:

173 return self.num_nans_in_logits > 0

vllm/v1/structured_output/__init__.py

+15 -17

@@ -167,7 +167,7 @@ class StructuredOutputManager:

167 def grammar_bitmask(

168 self,

169 requests: dict[str, Request],

170 - structured_output_request_ids: dict[str, int],

170 + structured_output_request_ids: list[str],

171 scheduled_spec_decode_tokens: dict[str, list[int]],

172 ) -> "npt.NDArray[np.int32] | None":

173 # Prepare the structured output bitmask for this batch.

@@ -196,17 +196,16 @@ class StructuredOutputManager:

196 # masks for each request, one for each possible bonus token position.

197 # These are stored inline in the tensor and unpacked by the gpu runner.

198 cumulative_index = 0

199 - ordered_seq = sorted(structured_output_request_ids.items(), key=lambda x: x[1])

199

200 # Optimized parallel filling of bitmasks for

201 # non-spec, large-batch-size cases

202 if (

204 - len(ordered_seq) > self.fill_bitmask_parallel_threshold

203 + len(structured_output_request_ids) > self.fill_bitmask_parallel_threshold

204 and max_num_spec_tokens == 0

205 ):

206 promises = []

207 batch = []

209 - for req_id, _ in ordered_seq:

208 + for req_id in structured_output_request_ids:

209 request = requests[req_id]

210 structured_output_request = request.structured_output_request

211 if TYPE_CHECKING:

@@ -230,7 +229,7 @@ class StructuredOutputManager:

229 promise.result()

230 else:

231 # Fallback to serial filling of bitmasks for small-batch-size cases

233 - for req_id, _ in ordered_seq:

232 + for req_id in structured_output_request_ids:

233 request = requests[req_id]

234 structured_output_request = request.structured_output_request

235

@@ -295,21 +294,20 @@ class StructuredOutputManager:

294 assert request.structured_output_request.grammar is not None

295 # by default, we should always advance

296 # for cases that don't use thinking mode.

298 - if self.reasoner is not None:

299 - structured_req = request.structured_output_request

297 + if self.reasoner is None:

298 + return True

299

301 - if structured_req.reasoning_ended:

302 - return True

300 + structured_req = request.structured_output_request

301 + if structured_req.reasoning_ended:

302 + return True

303

304 - # Check if reasoning ends in *this* step

305 - if self.reasoner.is_reasoning_end(request.all_token_ids):

306 - # Reasoning just ended, so we shouldn't advance til

307 - # next pass

308 - structured_req.reasoning_ended = True

304 + # Check if reasoning ends in *this* step

305 + if self.reasoner.is_reasoning_end(request.all_token_ids):

306 + # Reasoning just ended, so we shouldn't advance til

307 + # next pass

308 + structured_req.reasoning_ended = True

309

310 - return False

311 - else:

312 - return True

310 + return False

311

312 def clear_backend(self) -> None:

313 if self.backend is not None:

vllm/v1/structured_output/backend_guidance.py

+1 -1

@@ -252,7 +252,7 @@ def serialize_guidance_grammar(

252 def validate_guidance_grammar(

253 sampling_params: SamplingParams, tokenizer: llguidance.LLTokenizer | None = None

254 ) -> None:

255 - tp, grm = get_structured_output_key(sampling_params)

255 + tp, grm = get_structured_output_key(sampling_params.structured_outputs)

256 guidance_grm = serialize_guidance_grammar(tp, grm)

257 err = llguidance.LLMatcher.validate_grammar(guidance_grm, tokenizer)

258 if err:

vllm/v1/structured_output/request.py

+25 -19

@@ -7,7 +7,7 @@ from concurrent.futures import Future

7 from concurrent.futures._base import TimeoutError

8 from typing import cast

10 -from vllm.sampling_params import SamplingParams

10 +from vllm.sampling_params import SamplingParams, StructuredOutputsParams

11 from vllm.v1.structured_output.backend_types import (

12 StructuredOutputGrammar,

13 StructuredOutputKey,

@@ -17,10 +17,19 @@ from vllm.v1.structured_output.backend_types import (

18 @dataclasses.dataclass

19 class StructuredOutputRequest:

20 - sampling_params: SamplingParams

20 + params: StructuredOutputsParams

21 _grammar: Future[StructuredOutputGrammar] | StructuredOutputGrammar | None = None

22 reasoning_ended: bool | None = None

24 + @staticmethod

25 + def from_sampling_params(

26 + sampling_params: SamplingParams | None,

27 + ) -> "StructuredOutputRequest | None":

28 + if sampling_params is None:

29 + return None

30 + params = sampling_params.structured_outputs

31 + return StructuredOutputRequest(params=params) if params else None

32 +

33 def _check_grammar_completion(self) -> bool:

34 # NOTE: We have to lazy import to gate circular imports

35 from vllm.v1.request import RequestStatus

@@ -53,31 +62,28 @@ class StructuredOutputRequest:

63 @functools.cached_property

64 def structured_output_key(self) -> StructuredOutputKey:

56 - return get_structured_output_key(self.sampling_params)

65 + return get_structured_output_key(self.params)

59 -def get_structured_output_key(sampling_params: SamplingParams) -> StructuredOutputKey:

60 - params = sampling_params.structured_outputs

61 - assert params is not None, "params can't be None."

68 +def get_structured_output_key(params: StructuredOutputsParams) -> StructuredOutputKey:

69 if params.json is not None:

70 if not isinstance(params.json, str):

71 json_str = json.dumps(params.json)

72 else:

73 json_str = params.json

67 - return (StructuredOutputOptions.JSON, json_str)

68 - elif params.json_object:

69 - return (StructuredOutputOptions.JSON_OBJECT, "")

70 - elif params.regex is not None:

71 - return (StructuredOutputOptions.REGEX, params.regex)

72 - elif params.choice is not None:

74 + return StructuredOutputOptions.JSON, json_str

75 + if params.json_object:

76 + return StructuredOutputOptions.JSON_OBJECT, ""

77 + if params.regex is not None:

78 + return StructuredOutputOptions.REGEX, params.regex

79 + if params.choice is not None:

80 if not isinstance(params.choice, str):

81 json_str = json.dumps(params.choice)

82 else:

83 json_str = params.choice

77 - return (StructuredOutputOptions.CHOICE, json_str)

78 - elif params.grammar is not None:

79 - return (StructuredOutputOptions.GRAMMAR, params.grammar)

80 - elif params.structural_tag is not None:

81 - return (StructuredOutputOptions.STRUCTURAL_TAG, params.structural_tag)

82 - else:

83 - raise ValueError("No valid structured output parameter found")

84 + return StructuredOutputOptions.CHOICE, json_str

85 + if params.grammar is not None:

86 + return StructuredOutputOptions.GRAMMAR, params.grammar

87 + if params.structural_tag is not None:

88 + return StructuredOutputOptions.STRUCTURAL_TAG, params.structural_tag

89 + raise ValueError("No valid structured output parameter found")

vllm/v1/structured_output/utils.py

+2 -7

@@ -47,7 +47,6 @@ def apply_grammar_bitmask(

47 scheduler_output: SchedulerOutput,

48 input_batch: InputBatch,

49 logits: torch.Tensor,

50 - device: torch.device,

50 ) -> None:

51 """

52 Apply grammar bitmask to output logits of the model with xgrammar function.

@@ -56,7 +55,6 @@ def apply_grammar_bitmask(

55 scheduler_output (SchedulerOutput): The result of engine scheduling.

56 input_batch (InputBatch): The input of model runner.

57 logits (torch.Tensor): The output logits of model forward.

59 - device (torch.device): The device that model runner running on.

58 """

59 grammar_bitmask = scheduler_output.grammar_bitmask

60 if grammar_bitmask is None:

@@ -91,10 +89,7 @@ def apply_grammar_bitmask(

89 dtype=grammar_bitmask.dtype,

90 )

91 cumulative_index = 0

94 - seq = sorted(

95 - scheduler_output.structured_output_request_ids.items(), key=lambda x: x[1]

96 - )

97 - for req_id, _ in seq:

92 + for req_id in scheduler_output.structured_output_request_ids:

93 num_spec_tokens = len(

94 scheduler_output.scheduled_spec_decode_tokens.get(req_id, [])

95 )

@@ -117,7 +112,7 @@ def apply_grammar_bitmask(

112

113 xgr.apply_token_bitmask_inplace(

114 logits,

120 - grammar_bitmask.to(device, non_blocking=True),

115 + grammar_bitmask.to(logits.device, non_blocking=True),

116 indices=out_indices if not skip_out_indices else None,

117 )

118

vllm/v1/worker/gpu_model_runner.py

+2 -4

@@ -2568,10 +2568,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2568 logits = model_output_broadcast_data["logits"]

2569

2570 # Apply structured output bitmasks if present

2571 - if scheduler_output.grammar_bitmask is not None:

2572 - apply_grammar_bitmask(

2573 - scheduler_output, self.input_batch, logits, self.device

2574 - )

2571 + if scheduler_output.structured_output_request_ids:

2572 + apply_grammar_bitmask(scheduler_output, self.input_batch, logits)

2573

2574 with record_function_or_nullcontext("Sample"):

2575 sampler_output = self._sample(logits, spec_decode_metadata)

vllm/v1/worker/tpu_model_runner.py

+1 -5

@@ -1963,12 +1963,8 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1963 self.grammar_bitmask_cpu.zero_()

1964 self.require_structured_out_cpu.zero_()

1965

1966 - sorted_struct_requests = sorted(

1967 - scheduler_output.structured_output_request_ids.items(),

1968 - key=lambda item: item[1],

1969 - )

1966 cumulative_mask_idx = 0

1971 - for req_id, _ in sorted_struct_requests:

1967 + for req_id in scheduler_output.structured_output_request_ids:

1968 if req_id not in self.input_batch.req_id_to_index:

1969 continue

1970 batch_index = self.input_batch.req_id_to_index[req_id]