@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 603fbbbce0

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

[Misc] Misc code cleanup/simplification (#23304)

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

Nick Hill committed Aug 21, 2025 at 10:22 UTC 603fbbbce0d163a3dd1fec6c5b56cf64c7021691

4 files changed +52 -58

vllm/v1/sample/sampler.py

+1 -1

@@ -91,7 +91,7 @@ class Sampler(nn.Module):

91 logits = self.apply_bad_words(logits, sampling_metadata)

93 # Apply logits processors which can impact greedy sampling

94 - for processor in (sampling_metadata.logitsprocs.non_argmax_invariant):

94 + for processor in sampling_metadata.logitsprocs.non_argmax_invariant:

95 logits = processor.apply(logits)

97 # Apply penalties (e.g., min_tokens, freq_penalties).

vllm/v1/worker/gpu_input_batch.py

+5 -4

@@ -442,10 +442,11 @@ class InputBatch:

442 # LoRA

443 lora_id = self.request_lora_mapping[req_index]

444 if lora_id != 0:

445 - self.lora_id_to_request_ids[lora_id].discard(req_id)

446 - if len(self.lora_id_to_request_ids[lora_id]) == 0:

447 - self.lora_id_to_request_ids.pop(lora_id)

448 - self.lora_id_to_lora_request.pop(lora_id)

445 + lora_req_ids = self.lora_id_to_request_ids[lora_id]

446 + lora_req_ids.discard(req_id)

447 + if not lora_req_ids:

448 + del self.lora_id_to_request_ids[lora_id]

449 + del self.lora_id_to_lora_request[lora_id]

450 self.request_lora_mapping[req_index] = 0

451

452 self.has_allowed_token_ids.discard(req_id)

vllm/v1/worker/gpu_model_runner.py

+45 -52

@@ -358,6 +358,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

358 if num_pooling_reqs == 0:

359 return model_kwargs

360

361 + # This does nontrivial work.

362 pooling_params = self.input_batch.pooling_metadata.pooling_params

363

364 assert num_pooling_reqs == num_reqs

@@ -465,7 +466,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

466 for req_id in unscheduled_req_ids:

467 self.input_batch.remove_request(req_id)

468

468 - req_ids_to_add: list[str] = []

469 + reqs_to_add: list[CachedRequestState] = []

470 # Add new requests to the cached states.

471 for new_req_data in scheduler_output.scheduled_new_reqs:

472 req_id = new_req_data.req_id

@@ -480,14 +481,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

481 generator = None

482

483 if pooling_params:

483 - assert (task := pooling_params.task) is not None, (

484 - "You did not set `task` in the API")

484 + task = pooling_params.task

485 + assert task is not None, "You did not set `task` in the API"

486

487 model = cast(VllmModelForPooling, self.get_model())

488 to_update = model.pooler.get_pooling_updates(task)

489 to_update.apply(pooling_params)

490

490 - self.requests[req_id] = CachedRequestState(

491 + req_state = CachedRequestState(

492 req_id=req_id,

493 prompt_token_ids=new_req_data.prompt_token_ids,

494 mm_kwargs=new_req_data.mm_kwargs,

@@ -501,6 +502,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

502 lora_request=new_req_data.lora_request,

503 )

504

505 + self.requests[req_id] = req_state

506 +

507 # Only relevant for models using M-RoPE (e.g, Qwen2-VL)

508 if self.uses_mrope:

509 image_grid_thw = []

@@ -508,29 +511,25 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

511 second_per_grid_ts = []

512 audio_feature_lengths = []

513 use_audio_in_video = False

511 - for mm_item in self.requests[req_id].mm_kwargs:

514 + for mm_item in req_state.mm_kwargs:

515 mm_input = mm_item.get_data()

513 - if mm_input.get("image_grid_thw") is not None:

514 - image_grid_thw.append(

515 - mm_input["image_grid_thw"].tolist())

516 - if mm_input.get("video_grid_thw") is not None:

517 - video_grid_thw.append(

518 - mm_input["video_grid_thw"].tolist())

519 - if mm_input.get("second_per_grid_ts") is not None:

520 - second_per_grid_ts.append(

521 - mm_input["second_per_grid_ts"])

522 - if mm_input.get("audio_feature_lengths") is not None:

523 - audio_feature_lengths.append(

524 - mm_input["audio_feature_lengths"])

516 + if (t := mm_input.get("image_grid_thw")) is not None:

517 + image_grid_thw.append(t.tolist())

518 + if (t := mm_input.get("video_grid_thw")) is not None:

519 + video_grid_thw.append(t.tolist())

520 + if (t := mm_input.get("second_per_grid_ts")) is not None:

521 + second_per_grid_ts.append(t)

522 + if (t :=

523 + mm_input.get("audio_feature_lengths")) is not None:

524 + audio_feature_lengths.append(t)

525 if mm_input.get("use_audio_in_video") is True:

526 use_audio_in_video = True

527

528 hf_config = self.model_config.hf_config

529

530 - self.requests[req_id].mrope_positions, \

531 - self.requests[req_id].mrope_position_delta = \

530 + req_state.mrope_positions, req_state.mrope_position_delta = \

531 MRotaryEmbedding.get_input_positions_tensor(

533 - self.requests[req_id].prompt_token_ids,

532 + req_state.prompt_token_ids,

533 hf_config=hf_config,

534 image_grid_thw=image_grid_thw,

535 video_grid_thw=video_grid_thw,

@@ -539,7 +538,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

538 use_audio_in_video=use_audio_in_video,

539 )

540

542 - req_ids_to_add.append(req_id)

541 + reqs_to_add.append(req_state)

542

543 # Update the states of the running/resumed requests.

544 is_last_rank = get_pp_group().is_last_rank

@@ -587,7 +586,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

586 # The request is not in the persistent batch.

587 # The request was either preempted and resumed later, or was not

588 # scheduled in the previous step and needs to be added again.

590 - req_ids_to_add.append(req_id)

589 + reqs_to_add.append(req_state)

590 continue

591

592 # Update the persistent batch.

@@ -624,9 +623,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

623

624 # Add the new or resumed requests to the persistent batch.

625 # The smaller empty indices are filled first.

627 - for req_id in req_ids_to_add:

628 - req_state = self.requests[req_id]

629 - self.input_batch.add_request(req_state)

626 + for request in reqs_to_add:

627 + self.input_batch.add_request(request)

628

629 # Condense the batched states if there are gaps left by removed requests

630 self.input_batch.condense()

@@ -639,38 +637,32 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

637 self,

638 scheduler_output: "SchedulerOutput",

639 ) -> BatchedTensorInputs:

642 - if self.is_multimodal_raw_input_supported: # noqa: SIM102

643 - if scheduler_output:

644 - mm_kwargs = list[MultiModalKwargsItem]()

645 - for req in scheduler_output.scheduled_new_reqs:

646 - req_mm_kwargs = req.mm_kwargs

647 - if not isinstance(req_mm_kwargs, list):

648 - req_mm_kwargs = list(req_mm_kwargs)

649 - mm_kwargs.extend(req_mm_kwargs)

650 -

651 - # Input all modalities at once

652 - mm_kwargs_combined: BatchedTensorInputs = {}

653 - for _, _, mm_kwargs_group in group_mm_kwargs_by_modality(

654 - mm_kwargs,

655 - device=self.device,

656 - pin_memory=self.pin_memory,

657 - ):

658 - mm_kwargs_combined.update(mm_kwargs_group)

659 -

660 - return mm_kwargs_combined

640 + if not self.is_multimodal_raw_input_supported or not scheduler_output: # noqa: SIM102

641 + return {}

642

662 - return {}

643 + mm_kwargs = list[MultiModalKwargsItem]()

644 + for req in scheduler_output.scheduled_new_reqs:

645 + mm_kwargs.extend(req.mm_kwargs)

646

664 - def _dummy_mm_kwargs(self, num_seqs: int) -> BatchedTensorInputs:

665 - if self.is_multimodal_raw_input_supported:

666 - mm_budget = self.mm_budget

667 - assert mm_budget is not None

647 + # Input all modalities at once

648 + mm_kwargs_combined: BatchedTensorInputs = {}

649 + for _, _, mm_kwargs_group in group_mm_kwargs_by_modality(

650 + mm_kwargs,

651 + device=self.device,

652 + pin_memory=self.pin_memory,

653 + ):

654 + mm_kwargs_combined.update(mm_kwargs_group)

655

669 - dummy_modality = mm_budget.get_modality_with_max_tokens()

656 + return mm_kwargs_combined

657

671 - return self._get_mm_dummy_batch(dummy_modality, num_seqs)

658 + def _dummy_mm_kwargs(self, num_seqs: int) -> BatchedTensorInputs:

659 + if not self.is_multimodal_raw_input_supported:

660 + return {}

661 + mm_budget = self.mm_budget

662 + assert mm_budget is not None

663

673 - return {}

664 + dummy_modality = mm_budget.get_modality_with_max_tokens()

665 + return self._get_mm_dummy_batch(dummy_modality, num_seqs)

666

667 def _get_cumsum_and_arange(

668 self,

@@ -1612,6 +1604,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1604 batch_descriptor=batch_descriptor,

1605 ), self.maybe_get_kv_connector_output(

1606 scheduler_output) as kv_connector_output:

1607 +

1608 model_output = self.model(

1609 input_ids=input_ids,

1610 positions=positions,

vllm/worker/worker_base.py

+1 -1

@@ -544,7 +544,7 @@ class WorkerWrapperBase:

544 Arguments are passed to the worker class constructor.

545 """

546 kwargs = all_kwargs[self.rpc_rank]

547 - self.vllm_config = kwargs.get("vllm_config", None)

547 + self.vllm_config = kwargs.get("vllm_config")

548 assert self.vllm_config is not None, (

549 "vllm_config is required to initialize the worker")

550 enable_trace_function_call_for_thread(self.vllm_config)