[ModelRunner V2] Misc minor simplifications and optimizations (vllm-project#33467)

Signed-off-by: Nick Hill <[email protected]>

Author: njhill

vllm/v1/worker/gpu/async_utils.py

@@ -4,11 +4,7 @@
 import numpy as np
 import torch
 
-from vllm.v1.outputs import (
-    AsyncModelRunnerOutput,
-    LogprobsTensors,
-    ModelRunnerOutput,
-)
+from vllm.v1.outputs import AsyncModelRunnerOutput, LogprobsTensors, ModelRunnerOutput
 from vllm.v1.worker.gpu.sample.output import SamplerOutput
 
 

vllm/v1/worker/gpu/attn_utils.py

@@ -32,9 +32,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:
 
 
 def init_attn_backend(
-    kv_cache_config: KVCacheConfig,
-    vllm_config: VllmConfig,
-    device: torch.device,
+    kv_cache_config: KVCacheConfig, vllm_config: VllmConfig, device: torch.device
 ):
     attn_backends: dict[str, type[AttentionBackend]] = {}
     attn_metadata_builders: list[AttentionMetadataBuilder] = []
@@ -50,10 +48,7 @@ def init_attn_backend(
             attn_backends[layer_name] = attn_backend
 
         attn_metadata_builder = attn_backend.get_builder_cls()(
-            kv_cache_group_spec.kv_cache_spec,
-            layer_names,
-            vllm_config,
-            device,
+            kv_cache_group_spec.kv_cache_spec, layer_names, vllm_config, device
         )
         attn_metadata_builders.append(attn_metadata_builder)  # type: ignore
 
@@ -65,10 +60,7 @@ def init_attn_backend(
     return attn_backends, attn_metadata_builders
 
 
-def _allocate_kv_cache(
-    kv_cache_config: KVCacheConfig,
-    device: torch.device,
-):
+def _allocate_kv_cache(kv_cache_config: KVCacheConfig, device: torch.device):
     kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
     for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
         tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device)
@@ -141,12 +133,11 @@ def init_kv_cache(
 
 
 def build_slot_mappings_by_layer(
-    slot_mappings: torch.Tensor,
-    kv_cache_config: KVCacheConfig,
+    slot_mappings: torch.Tensor, kv_cache_config: KVCacheConfig
 ) -> dict[str, torch.Tensor]:
     slot_mappings_by_layer: dict[str, torch.Tensor] = {}
-    for i, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
-        slot_mapping = slot_mappings[i]
+    kv_cache_groups = kv_cache_config.kv_cache_groups
+    for slot_mapping, kv_cache_group in zip(slot_mappings, kv_cache_groups):
         for layer_name in kv_cache_group.layer_names:
             slot_mappings_by_layer[layer_name] = slot_mapping
     return slot_mappings_by_layer
@@ -188,8 +179,7 @@ def build_attn_metadata(
 
         attn_metadata_builder = attn_metadata_builders[i]
         metadata = attn_metadata_builder.build(
-            common_prefix_len=0,
-            common_attn_metadata=common_attn_metadata,
+            common_prefix_len=0, common_attn_metadata=common_attn_metadata
         )
         for layer_name in kv_cache_spec.layer_names:
             attn_metadata[layer_name] = metadata

vllm/v1/worker/gpu/block_table.py

@@ -71,9 +71,7 @@ def __init__(
     def _make_ptr_tensor(self, x: Iterable[torch.Tensor]) -> torch.Tensor:
         # NOTE(woosuk): Use uint64 instead of int64 to cover all possible addresses.
         return torch.tensor(
-            [t.data_ptr() for t in x],
-            dtype=torch.uint64,
-            device=self.device,
+            [t.data_ptr() for t in x], dtype=torch.uint64, device=self.device
         )
 
     def append_block_ids(
@@ -96,8 +94,7 @@ def apply_staged_writes(self) -> None:
         self.num_blocks.copy_to_uva()
 
     def gather_block_tables(
-        self,
-        idx_mapping: torch.Tensor,
+        self, idx_mapping: torch.Tensor
     ) -> tuple[torch.Tensor, ...]:
         num_reqs = idx_mapping.shape[0]
         _gather_block_tables_kernel[(self.num_kv_cache_groups, num_reqs)](

vllm/v1/worker/gpu/buffer_utils.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Iterable, Sequence
+from functools import partial
 
 import numpy as np
 import torch
@@ -81,10 +82,7 @@ def copy_to_gpu(
 
 class UvaBackedTensor:
     def __init__(
-        self,
-        size: int | Sequence[int],
-        dtype: torch.dtype,
-        max_concurrency: int = 2,
+        self, size: int | Sequence[int], dtype: torch.dtype, max_concurrency: int = 2
     ):
         self.dtype = dtype
         self.max_concurrency = max_concurrency
@@ -135,25 +133,16 @@ def __init__(
         self._staged_write_contents: list[int | float] = []
         self._staged_write_cu_lens: list[int] = []
 
-        self.write_indices = UvaBufferPool(
-            self.num_rows, dtype=torch.int32, max_concurrency=max_concurrency
-        )
-        self.write_starts = UvaBufferPool(
-            self.num_rows, dtype=torch.int32, max_concurrency=max_concurrency
-        )
+        new_buffer = partial(UvaBufferPool, max_concurrency=max_concurrency)
+
+        self.write_indices = new_buffer(self.num_rows, dtype=torch.int32)
+        self.write_starts = new_buffer(self.num_rows, dtype=torch.int32)
         init_size = next_power_of_2(self.num_rows)
-        self.write_contents = UvaBufferPool(
-            init_size, dtype=dtype, max_concurrency=max_concurrency
-        )
-        self.write_cu_lens = UvaBufferPool(
-            self.num_rows, dtype=torch.int32, max_concurrency=max_concurrency
-        )
+        self.write_contents = new_buffer(init_size, dtype=dtype)
+        self.write_cu_lens = new_buffer(self.num_rows, dtype=torch.int32)
 
     def stage_write(
-        self,
-        index: int,
-        start: int,
-        x: Iterable[int] | Iterable[float],
+        self, index: int, start: int, x: Iterable[int] | Iterable[float]
     ) -> None:
         assert index >= 0
         assert start >= 0

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -24,12 +24,7 @@
 
 
 class CudaGraphManager:
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        uses_mrope: bool,
-        device: torch.device,
-    ):
+    def __init__(self, vllm_config: VllmConfig, uses_mrope: bool, device: torch.device):
         self.vllm_config = vllm_config
         self.scheduler_config = vllm_config.scheduler_config
         self.uses_mrope = uses_mrope
@@ -41,11 +36,7 @@ def __init__(
         self.dp_size = vllm_config.parallel_config.data_parallel_size
         self.compilation_config = vllm_config.compilation_config
         assert self.compilation_config is not None
-        self.cudagraph_mode: CUDAGraphMode
-        if self.compilation_config.cudagraph_mode is None:
-            self.cudagraph_mode = CUDAGraphMode.NONE
-        else:
-            self.cudagraph_mode = self.compilation_config.cudagraph_mode
+        self.cudagraph_mode = self.compilation_config.cudagraph_mode
         self.cudagraph_sizes = get_cudagraph_sizes(
             self.compilation_config.cudagraph_capture_sizes,
             self.max_num_reqs,

vllm/v1/worker/gpu/dp_utils.py

@@ -13,10 +13,7 @@ def make_num_tokens_across_dp(dp_size: int, num_tokens: int) -> torch.Tensor | N
 
 
 def get_batch_metadata_across_dp(
-    num_tokens: int,
-    cudagraph_size: int,
-    dp_size: int,
-    dp_rank: int,
+    num_tokens: int, cudagraph_size: int, dp_size: int, dp_rank: int
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert dp_size > 1
     # Use CPU group to avoid CPU-GPU synchronization.
@@ -29,10 +26,7 @@ def get_batch_metadata_across_dp(
 
 
 def get_cudagraph_and_dp_padding(
-    num_tokens: int,
-    cudagraph_size: int | None,
-    dp_size: int,
-    dp_rank: int,
+    num_tokens: int, cudagraph_size: int | None, dp_size: int, dp_rank: int
 ) -> tuple[bool, int, torch.Tensor | None]:
     if dp_size == 1:
         if cudagraph_size is not None:

vllm/v1/worker/gpu/kv_connector.py

@@ -65,10 +65,10 @@ def pre_forward(self, scheduler_output: "SchedulerOutput") -> None:
 
         if scheduler_output.preempted_req_ids:
             self.kv_connector.handle_preemptions(scheduler_output.preempted_req_ids)
-        assert scheduler_output.kv_connector_metadata is not None
-        self.kv_connector.bind_connector_metadata(
-            scheduler_output.kv_connector_metadata
-        )
+        kv_connector_metadata = scheduler_output.kv_connector_metadata
+        assert kv_connector_metadata is not None
+        self.kv_connector.bind_connector_metadata(kv_connector_metadata)
+
         # TODO: sort out KV Connectors' use of forward_context
         if is_forward_context_available():
             self.kv_connector.start_load_kv(get_forward_context())

vllm/v1/worker/gpu/lora_utils.py

@@ -15,10 +15,7 @@ def __init__(self, max_num_reqs: int):
         self.lora_requests: dict[str, LoRARequest] = {}
 
     def add_request(
-        self,
-        req_id: str,
-        req_index: int,
-        lora_request: LoRARequest | None,
+        self, req_id: str, req_index: int, lora_request: LoRARequest | None
     ) -> None:
         if lora_request is not None:
             self.lora_requests[req_id] = lora_request
@@ -41,7 +38,7 @@ def make_lora_inputs(
 
         active_lora_requests: set[LoRARequest] = set()
         for req_id in req_ids:
-            lora_request = self.lora_requests.get(req_id, None)
+            lora_request = self.lora_requests.get(req_id)
             if lora_request is not None:
                 active_lora_requests.add(lora_request)
         return prompt_lora_mapping, token_lora_mapping, active_lora_requests

vllm/v1/worker/gpu/mm/encoder_runner.py

@@ -23,10 +23,7 @@ def __init__(
         self.device = device
 
         self.inputs_embeds = torch.zeros(
-            max_num_tokens,
-            hidden_size,
-            dtype=dtype,
-            device=device,
+            max_num_tokens, hidden_size, dtype=dtype, device=device
         )
         self.req_id_to_mm_features: dict[str, list[MultiModalFeatureSpec]] = {}
         self.encoder_cache: dict[str, torch.Tensor] = {}
@@ -57,8 +54,7 @@ def remove_request(self, req_id: str) -> None:
         self.req_id_to_mm_features.pop(req_id, None)
 
     def prepare_mm_inputs(
-        self,
-        scheduled_encoder_inputs: dict[str, list[int]],
+        self, scheduled_encoder_inputs: dict[str, list[int]]
     ) -> tuple[list[str], list[tuple[str, MultiModalKwargsItem]]]:
         mm_hashes: list[str] = []
         mm_kwargs: list[tuple[str, MultiModalKwargsItem]] = []
@@ -85,20 +81,16 @@ def execute_mm_encoder(
 
         encoder_outputs: list[torch.Tensor] = []
         for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(
-            mm_kwargs,
-            device=self.device,
-            pin_memory=False,
+            mm_kwargs, device=self.device, pin_memory=False
         ):
             curr_group_outputs = model.embed_multimodal(**mm_kwargs_group)
             sanity_check_mm_encoder_outputs(
-                curr_group_outputs,
-                expected_num_items=num_items,
+                curr_group_outputs, expected_num_items=num_items
             )
             encoder_outputs.extend(curr_group_outputs)
 
         # Cache the encoder outputs by mm_hash
-        for mm_hash, output in zip(mm_hashes, encoder_outputs):
-            self.encoder_cache[mm_hash] = output
+        self.encoder_cache.update(zip(mm_hashes, encoder_outputs))
         return encoder_outputs
 
     def gather_mm_embeddings(
@@ -115,20 +107,15 @@ def gather_mm_embeddings(
         if all_decode:
             # All decode requests, so no need to gather any embeddings.
             return [], torch.zeros(
-                total_num_scheduled_tokens,
-                dtype=torch.bool,
-                device=self.device,
+                total_num_scheduled_tokens, dtype=torch.bool, device=self.device
             )
 
         query_start = computed_prefill_lens.tolist()
         query_end = (computed_prefill_lens + num_scheduled_tokens).tolist()
 
         mm_embeds: list[torch.Tensor] = []
         is_mm_embed = torch.zeros(
-            total_num_scheduled_tokens,
-            dtype=torch.bool,
-            device="cpu",
-            pin_memory=True,
+            total_num_scheduled_tokens, dtype=torch.bool, device="cpu", pin_memory=True
         )
         for i, req_id in enumerate(req_ids):
             if not is_prefilling[i]:
@@ -189,9 +176,7 @@ def get_inputs_embeds(
         is_mm_embed: torch.Tensor,
     ) -> torch.Tensor:
         x = model.embed_input_ids(
-            input_ids,
-            multimodal_embeddings=mm_embeds,
-            is_multimodal=is_mm_embed,
+            input_ids, multimodal_embeddings=mm_embeds, is_multimodal=is_mm_embed
         )
         # Copy to the pre-allocated buffer for CUDA graphs.
         self.inputs_embeds[: x.shape[0]] = x

vllm/v1/worker/gpu/mm/mrope_utils.py

@@ -51,10 +51,7 @@ def init_prefill_mrope_positions(
         mm_features: list,
     ) -> None:
         prefill_mrope_positions, prefill_mrope_delta = (
-            mrope_model.get_mrope_input_positions(
-                prefill_token_ids,
-                mm_features,
-            )
+            mrope_model.get_mrope_input_positions(prefill_token_ids, mm_features)
         )
         for i in range(3):
             pos = prefill_mrope_positions[i].tolist()