@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 43e3f8e4a9

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 10, 19:00 UTC

[Misc] Various code simplifications (#31666)

Signed-off-by: njhill <nickhill123@gmail.com>

Nick Hill committed Jan 4, 2026 at 18:35 UTC 43e3f8e4a9f08597847a3844f419d1efe22cc2d8

7 files changed +66 -135

vllm/v1/core/sched/async_scheduler.py

+2 -7

@@ -10,10 +10,7 @@ logger = init_logger(__name__)

12 class AsyncScheduler(Scheduler):

13 - def _update_after_schedule(

14 - self,

15 - scheduler_output: SchedulerOutput,

16 - ) -> None:

13 + def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:

14 super()._update_after_schedule(scheduler_output)

15 pending_structured_output_tokens = False

16 spec_decode_tokens = scheduler_output.scheduled_spec_decode_tokens

@@ -41,9 +38,7 @@ class AsyncScheduler(Scheduler):

38 )

40 def _update_request_with_output(

44 - self,

45 - request: Request,

46 - new_token_ids: list[int],

41 + self, request: Request, new_token_ids: list[int]

42 ) -> tuple[list[int], bool]:

43 if request.discard_latest_async_tokens:

44 # If the request is force preempted in reset_prefix_cache, we

vllm/v1/core/sched/interface.py

+1 -4

@@ -85,10 +85,7 @@ class SchedulerInterface(ABC):

85 raise NotImplementedError

87 @abstractmethod

88 - def update_draft_token_ids(

89 - self,

90 - draft_token_ids: "DraftTokenIds",

91 - ) -> None:

88 + def update_draft_token_ids(self, draft_token_ids: "DraftTokenIds") -> None:

89 """Update the draft token ids for the scheduled requests."""

90 raise NotImplementedError

vllm/v1/core/sched/scheduler.py

+7 -23

@@ -762,11 +762,7 @@ class Scheduler(SchedulerInterface):

762 self._update_after_schedule(scheduler_output)

763 return scheduler_output

764

765 - def _preempt_request(

766 - self,

767 - request: Request,

768 - timestamp: float,

769 - ) -> None:

765 + def _preempt_request(self, request: Request, timestamp: float) -> None:

766 """Preempt a request and put it back to the waiting queue.

767

768 NOTE: The request should be popped from the running queue outside of this

@@ -786,10 +782,7 @@ class Scheduler(SchedulerInterface):

782 # Put the request back to the waiting queue.

783 self.waiting.prepend_request(request)

784

789 - def _update_after_schedule(

790 - self,

791 - scheduler_output: SchedulerOutput,

792 - ) -> None:

785 + def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:

786 # Advance the number of computed tokens for the request AFTER

787 # the request is scheduled.

788 # 1. The scheduler_output of the current step has to include the

@@ -1006,8 +999,7 @@ class Scheduler(SchedulerInterface):

999 )

1000 curr_embeds_start, curr_embeds_end = (

1001 mm_feature.mm_position.get_embeds_indices_in_range(

1009 - start_idx_rel,

1010 - end_idx_rel,

1002 + start_idx_rel, end_idx_rel

1003 )

1004 )

1005 # There's no embeddings in the current range of encoder placeholder tokens

@@ -1034,8 +1026,7 @@ class Scheduler(SchedulerInterface):

1026 )

1027

1028 def get_grammar_bitmask(

1037 - self,

1038 - scheduler_output: SchedulerOutput,

1029 + self, scheduler_output: SchedulerOutput

1030 ) -> GrammarOutput | None:

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

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

@@ -1285,9 +1276,7 @@ class Scheduler(SchedulerInterface):

1276 return engine_core_outputs

1277

1278 def _update_request_with_output(

1288 - self,

1289 - request: Request,

1290 - new_token_ids: list[int],

1279 + self, request: Request, new_token_ids: list[int]

1280 ) -> tuple[list[int], bool]:

1281 # Append generated tokens and check for stop. Note that if

1282 # a request is still being prefilled, we expect the model runner

@@ -1328,10 +1317,7 @@ class Scheduler(SchedulerInterface):

1317 # in the decoder's KV cache.

1318 self.encoder_cache_manager.free_encoder_input(request, input_id)

1319

1331 - def update_draft_token_ids(

1332 - self,

1333 - draft_token_ids: DraftTokenIds,

1334 - ) -> None:

1320 + def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:

1321 for req_id, spec_token_ids in zip(

1322 draft_token_ids.req_ids,

1323 draft_token_ids.draft_token_ids,

@@ -1361,9 +1347,7 @@ class Scheduler(SchedulerInterface):

1347 request.record_event(EngineCoreEventType.QUEUED)

1348

1349 def finish_requests(

1364 - self,

1365 - request_ids: str | Iterable[str],

1366 - finished_status: RequestStatus,

1350 + self, request_ids: str | Iterable[str], finished_status: RequestStatus

1351 ) -> None:

1352 """Handles the finish signal from outside the scheduler.

1353

vllm/v1/spec_decode/eagle.py

+10 -23

@@ -204,10 +204,7 @@ class EagleProposer:

204 )

205 # Precompute draft position offsets in flattened tree.

206 self.tree_draft_pos_offsets = torch.arange(

207 - 1,

208 - len(self.tree_choices) + 1,

209 - device=device,

210 - dtype=torch.int32,

207 + 1, len(self.tree_choices) + 1, device=device, dtype=torch.int32

208 ).repeat(max_batch_size, 1)

209

210 def _get_positions(self, num_tokens: int):

@@ -287,8 +284,7 @@ class EagleProposer:

284 per_layer_attn_metadata[layer_name] = draft_indexer_metadata

285

286 num_tokens_dp_padded, num_tokens_across_dp = self._pad_batch_across_dp(

290 - num_tokens_unpadded=num_tokens,

291 - num_tokens_padded=num_tokens,

287 + num_tokens_unpadded=num_tokens, num_tokens_padded=num_tokens

288 )

289

290 cudagraph_runtime_mode = CUDAGraphMode.NONE

@@ -391,8 +387,7 @@ class EagleProposer:

387 draft_token_ids_list = [draft_token_ids]

388

389 batch_size_dp_padded, batch_size_across_dp = self._pad_batch_across_dp(

394 - num_tokens_unpadded=batch_size,

395 - num_tokens_padded=batch_size,

390 + num_tokens_unpadded=batch_size, num_tokens_padded=batch_size

391 )

392

393 if (

@@ -610,10 +605,8 @@ class EagleProposer:

605 assert discard_request_mask.dtype == torch.bool

606 assert backup_tokens_gpu.dtype == torch.int32

607

613 - next_token_ids = torch.empty((batch_size,), dtype=torch.int32, device=device)

614 - valid_sampled_tokens_count = torch.empty(

615 - (batch_size,), dtype=torch.int32, device=device

616 - )

608 + next_token_ids = torch.empty(batch_size, dtype=torch.int32, device=device)

609 + valid_sampled_tokens_count = next_token_ids.new_empty(batch_size)

610

611 # Kernel grid: one program per request (row)

612 grid = (batch_size,)

@@ -782,8 +775,7 @@ class EagleProposer:

775 max_query_len=query_len,

776 )

777 attn_metadata = tree_attn_metadata_builder.build_for_drafting(

785 - common_attn_metadata=common_attn_metadata,

786 - draft_index=level + 1,

778 + common_attn_metadata=common_attn_metadata, draft_index=level + 1

779 )

780

781 # Apply new attention metadata to all layers.

@@ -1161,8 +1153,8 @@ class EagleProposer:

1153 def dummy_run(

1154 self,

1155 num_tokens: int,

1164 - use_cudagraphs=True,

1165 - is_graph_capturing=False,

1156 + use_cudagraphs: bool = True,

1157 + is_graph_capturing: bool = False,

1158 ) -> None:

1159 # Determine if CUDA graphs should be used for this run.

1160 cudagraphs_enabled = use_cudagraphs and self.use_cuda_graph

@@ -1174,8 +1166,7 @@ class EagleProposer:

1166 ):

1167 if fwd_idx <= 1:

1168 num_tokens_dp_padded, num_tokens_across_dp = self._pad_batch_across_dp(

1177 - num_tokens_unpadded=num_tokens,

1178 - num_tokens_padded=num_tokens,

1169 + num_tokens_unpadded=num_tokens, num_tokens_padded=num_tokens

1170 )

1171 if (

1172 cudagraphs_enabled

@@ -1342,9 +1333,5 @@ def compute_probs_and_sample_next_token(

1333 next_token_ids = probs.div(q).argmax(dim=-1).view(-1)

1334 if not sampling_metadata.all_random:

1335 greedy_token_ids = probs.argmax(dim=-1)

1345 - next_token_ids = torch.where(

1346 - is_greedy,

1347 - greedy_token_ids,

1348 - next_token_ids,

1349 - )

1336 + next_token_ids = torch.where(is_greedy, greedy_token_ids, next_token_ids)

1337 return next_token_ids, probs

vllm/v1/structured_output/__init__.py

+18 -37

@@ -28,8 +28,6 @@ if TYPE_CHECKING:

28 else:

29 torch = LazyLoader("torch", globals(), "torch")

31 - ReasoningParser = object

32 - Request = object

32 logger = init_logger(__name__)

@@ -98,7 +96,7 @@ class StructuredOutputManager:

96 self.vllm_config.structured_outputs_config.enable_in_reasoning

97 )

101 - def grammar_init(self, request: Request) -> None:

99 + def grammar_init(self, request: "Request") -> None:

100 if request.structured_output_request is None:

101 return

102

@@ -156,10 +154,7 @@ class StructuredOutputManager:

154 grammar = self._create_grammar(request) # type: ignore[assignment]

155 request.structured_output_request.grammar = grammar # type: ignore[assignment]

156

159 - def _create_grammar(

160 - self,

161 - request: Request,

162 - ) -> StructuredOutputGrammar:

157 + def _create_grammar(self, request: "Request") -> StructuredOutputGrammar:

158 key = request.structured_output_request.structured_output_key # type: ignore[union-attr]

159

160 # Note that the request was validated in the engine core client,

@@ -173,8 +168,7 @@ class StructuredOutputManager:

168 return self.backend.compile_grammar(request_type, grammar_spec)

169

170 def _fill_bitmasks(

176 - self,

177 - batch: Iterable[tuple[StructuredOutputGrammar, int, bool]],

171 + self, batch: Iterable[tuple[StructuredOutputGrammar, int, bool]]

172 ) -> None:

173 assert self._grammar_bitmask is not None

174 for grammar, index, apply_bitmask in batch:

@@ -187,14 +181,13 @@ class StructuredOutputManager:

181 self._grammar_bitmask[index].fill_(self._full_mask)

182

183 def _async_submit_fill_bitmask(

190 - self,

191 - batch: list[tuple[StructuredOutputGrammar, int, bool]],

184 + self, batch: list[tuple[StructuredOutputGrammar, int, bool]]

185 ) -> Future:

186 return self.executor_for_fillmask.submit(self._fill_bitmasks, batch)

187

188 def grammar_bitmask(

189 self,

197 - requests: dict[str, Request],

190 + requests: dict[str, "Request"],

191 structured_output_request_ids: list[str],

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

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

@@ -239,11 +232,10 @@ class StructuredOutputManager:

232 if TYPE_CHECKING:

233 assert structured_output_request is not None

234 assert structured_output_request.grammar is not None

235 + grammar = structured_output_request.grammar

236

237 apply_bitmask = self.should_fill_bitmask(request)

244 - batch.append(

245 - (structured_output_request.grammar, cumulative_index, apply_bitmask)

246 - )

238 + batch.append((grammar, cumulative_index, apply_bitmask))

239 if len(batch) == self.fill_bitmask_parallel_batch_size:

240 promises.append(self._async_submit_fill_bitmask(batch))

241 batch = []

@@ -264,34 +256,23 @@ class StructuredOutputManager:

256 if TYPE_CHECKING:

257 assert structured_output_request is not None

258 assert structured_output_request.grammar is not None

259 + grammar = structured_output_request.grammar

260 apply_bitmask = self.should_fill_bitmask(request)

261

262 state_advancements = 0

263 req_tokens = scheduled_spec_decode_tokens.get(req_id, ())

271 - for token in itertools.chain(req_tokens, (None,)):

272 - self._fill_bitmasks(

273 - (

274 - (

275 - structured_output_request.grammar,

276 - cumulative_index,

277 - apply_bitmask,

278 - ),

279 - )

280 - )

281 -

282 - if (

283 - apply_bitmask

284 - and token is not None

285 - and not structured_output_request.grammar.is_terminated()

286 - ):

287 - accepted = structured_output_request.grammar.accept_tokens(

288 - req_id, [token]

289 - )

264 + for token in itertools.chain(req_tokens, (-1,)):

265 + self._fill_bitmasks(((grammar, cumulative_index, apply_bitmask),))

266 + if token == -1:

267 + # Stop advancing the grammar once we hit a padding token.

268 + apply_bitmask = False

269 + if apply_bitmask and not grammar.is_terminated():

270 + accepted = grammar.accept_tokens(req_id, [token])

271 assert accepted, (token, req_id, scheduled_spec_decode_tokens)

272 state_advancements += 1

273 cumulative_index += 1

274 if state_advancements > 0:

294 - structured_output_request.grammar.rollback(state_advancements)

275 + grammar.rollback(state_advancements)

276

277 bitmask_tensor = self._grammar_bitmask

278 if cumulative_index < bitmask_tensor.shape[0]:

@@ -302,7 +283,7 @@ class StructuredOutputManager:

283 # and deserialization when sending this to the GPU workers.

284 return bitmask_tensor.numpy()

285

305 - def should_fill_bitmask(self, request: Request) -> bool:

286 + def should_fill_bitmask(self, request: "Request") -> bool:

287 # NOTE (Hanchen) if enable_in_reasoning is True, it means that

288 # the model needs to be constrained in reasoning. So we should always

289 # enable the bitmask filling.

@@ -318,7 +299,7 @@ class StructuredOutputManager:

299 return request.structured_output_request.reasoning_ended

300 return True

301

321 - def should_advance(self, request: Request) -> bool:

302 + def should_advance(self, request: "Request") -> bool:

303 if not request.use_structured_output:

304 return False

305

vllm/v1/structured_output/utils.py

+24 -36

@@ -5,6 +5,7 @@ from __future__ import annotations

5 import hashlib

6 import importlib.metadata

7 import os

8 +import tempfile

9 from typing import TYPE_CHECKING

11 import numpy as np

@@ -34,9 +35,6 @@ else:

35 "convert_slow_tokenizer", globals(), "transformers.convert_slow_tokenizer"

36 )

37 - TokenizerLike = object

38 - SchedulerOutput = object

39 - InputBatch = object

39 logger = init_logger(__name__)

@@ -72,13 +70,12 @@ def apply_grammar_bitmask(

70 # request in the batch, as the logit indices are offset by this amount.

71 struct_out_req_batch_indices: dict[str, int] = {}

72 cumulative_offset = 0

75 - seq = sorted(input_batch.req_id_to_index.items(), key=lambda x: x[1])

76 - for req_id, batch_index in seq:

73 + spec_tokens = scheduler_output.scheduled_spec_decode_tokens

74 + struct_out_req_ids = set(grammar_output.structured_output_request_ids)

75 + for batch_index, req_id in enumerate(input_batch.req_ids):

76 logit_index = batch_index + cumulative_offset

78 - cumulative_offset += len(

79 - scheduler_output.scheduled_spec_decode_tokens.get(req_id, [])

80 - )

81 - if req_id in grammar_output.structured_output_request_ids:

77 + cumulative_offset += len(spec_tokens.get(req_id, ()))

78 + if req_id in struct_out_req_ids:

79 struct_out_req_batch_indices[req_id] = logit_index

81 out_indices = []

@@ -91,14 +88,12 @@ def apply_grammar_bitmask(

88 )

89 cumulative_index = 0

90 for req_id in grammar_output.structured_output_request_ids:

94 - num_spec_tokens = len(

95 - scheduler_output.scheduled_spec_decode_tokens.get(req_id, [])

96 - )

97 - if req_id in struct_out_req_batch_indices:

98 - logit_index = struct_out_req_batch_indices[req_id]

91 + num_spec_tokens = len(spec_tokens.get(req_id, ()))

92 + if (logit_idx := struct_out_req_batch_indices.get(req_id)) is not None:

93 for i in range(1 + num_spec_tokens):

100 - sorted_bitmask[logit_index + i] = grammar_bitmask[cumulative_index + i]

101 - out_indices.append(logit_index + i)

94 + bitmask_index = logit_idx + i

95 + sorted_bitmask[bitmask_index] = grammar_bitmask[cumulative_index + i]

96 + out_indices.append(bitmask_index)

97 cumulative_index += 1 + num_spec_tokens

99 # Copy async to device as tensor.

@@ -149,21 +144,19 @@ def get_outlines_cache_path() -> str:

144 if outlines_cache_dir:

145 # OUTLINES_CACHE_DIR takes precedence

146 return outlines_cache_dir

152 - elif xdg_cache_home:

147 + if xdg_cache_home:

148 return os.path.join(xdg_cache_home, ".cache", "outlines")

149 # If homedir is "/", we may be inside a container, and thus writing to

150 # root would be problematic, so we fall back to using a tempfile.

151 # Also validate the path exists, since os.path.expanduser does

152 # not guarantee existence.

158 - elif os.path.isdir(home_dir) and home_dir != "/":

153 + if os.path.isdir(home_dir) and home_dir != "/":

154 # Default Unix fallback: ~/.cache/outlines

155 return os.path.join(home_dir, ".cache", "outlines")

161 - else:

162 - import tempfile

156

164 - # home_dir may be / inside a docker container without existing user

165 - tempdir = tempfile.gettempdir()

166 - return os.path.join(tempdir, ".cache", "outlines")

157 + # home_dir may be / inside a docker container without existing user

158 + tempdir = tempfile.gettempdir()

159 + return os.path.join(tempdir, ".cache", "outlines")

160

161

162 def get_outlines_cache():

@@ -184,8 +177,8 @@ def get_outlines_cache():

177 cache.clear()

178 cache.set("__version__", outlines_version)

179 return cache

187 - else:

188 - return LRUCache(maxsize=128)

180 +

181 + return LRUCache(maxsize=128)

182

183

184 re_llama_byte_token = re.compile(r"^<0x[0-9A-F]{2}>$")

@@ -193,8 +186,7 @@ re_replacement_seq = re.compile(r"^.{0,6}�+.{0,6}$")

186

187

188 def _reduced_vocabulary(

196 - tokenizer: TokenizerLike,

197 - eos_token_id: int,

189 + tokenizer: TokenizerLike, eos_token_id: int

190 ) -> dict[bytes, list[int]]:

191 """Create a map from vocabulary tokens to lists of equivalent token ids.

192

@@ -267,17 +259,13 @@ def get_outlines_vocabulary(tokenizer: TokenizerLike) -> oc.Vocabulary:

259 return tokenizer._outlines_vocabulary # type: ignore

260

261 try:

270 - if (

271 - hasattr(

272 - tokenizer,

273 - "eos_token_id",

274 - )

275 - and tokenizer.eos_token_id is not None

276 - ):

262 + if hasattr(tokenizer, "eos_token_id") and tokenizer.eos_token_id is not None:

263 eos_token_id = tokenizer.eos_token_id

264 else:

265 raise ValueError(

280 - f"Error during structured outputs setup for outlines: Tokenizer ({type(tokenizer)}) has no `eos_token_id` property, but `eos_token_id` is required for structured outputs to work properly." # noqa: E501

266 + "Error during structured outputs setup for outlines: Tokenizer "

267 + f"({type(tokenizer)}) has no `eos_token_id` property, but "

268 + "`eos_token_id` is required for structured outputs to work properly."

269 )

270

271 reduced_vocab = _reduced_vocabulary(

@@ -290,7 +278,7 @@ def get_outlines_vocabulary(tokenizer: TokenizerLike) -> oc.Vocabulary:

278 return vocabulary

279 except AttributeError as e:

280 raise ValueError(

293 - f"Cannot get the vocabulary of the tokenizer "

281 + "Cannot get the vocabulary of the tokenizer "

282 f"({type(tokenizer)}). The tokenizer should have a "

283 "get_vocab method."

284 ) from e

vllm/v1/worker/gpu_model_runner.py

+4 -5

@@ -3564,14 +3564,13 @@ class GPUModelRunner(

3564 def _get_valid_sampled_token_count(self) -> list[int]:

3565 # Wait until valid_sampled_tokens_count is copied to cpu,

3566 prev_sampled_token_ids = self.input_batch.prev_sampled_token_ids

3567 - if (

3568 - self.valid_sampled_token_count_event is None

3569 - or prev_sampled_token_ids is None

3570 - ):

3567 + sampled_count_event = self.valid_sampled_token_count_event

3568 + if sampled_count_event is None or prev_sampled_token_ids is None:

3569 return []

3570

3571 counts_cpu = self.valid_sampled_token_count_cpu

3574 - self.valid_sampled_token_count_event.synchronize()

3572 + assert counts_cpu is not None

3573 + sampled_count_event.synchronize()

3574 return counts_cpu[: prev_sampled_token_ids.shape[0]].tolist()

3575

3576 def propose_draft_token_ids(