diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index 2733fd396e6..1dd85efca4f 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/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 dataclasses import dataclass from functools import partial import numpy as np @@ -35,18 +36,34 @@ def async_copy_to_gpu( return out.copy_(tmp, non_blocking=True) +@dataclass +class PrepareInputsRingSlot: + idx_mapping_cpu: torch.Tensor + query_start_loc_cpu: torch.Tensor + cu_num_logits_cpu: torch.Tensor + event: torch.cuda.Event + in_use: bool = False + + class PrepareInputsBuffers: - def __init__(self, max_num_reqs: int, device: torch.device): + def __init__(self, ring_size: int, max_num_reqs: int, device: torch.device): self.device = device - self.idx_mapping_cpu = torch.empty( - max_num_reqs, dtype=torch.int32, pin_memory=True - ) - self.query_start_loc_cpu = torch.empty( - max_num_reqs + 1, dtype=torch.int32, pin_memory=True - ) - self.cu_num_logits_cpu = torch.empty( - max_num_reqs + 1, dtype=torch.int32, pin_memory=True - ) + self.ring_slots = [ + PrepareInputsRingSlot( + idx_mapping_cpu=torch.empty( + max_num_reqs, dtype=torch.int32, pin_memory=True + ), + query_start_loc_cpu=torch.empty( + max_num_reqs + 1, dtype=torch.int32, pin_memory=True + ), + cu_num_logits_cpu=torch.empty( + max_num_reqs + 1, dtype=torch.int32, pin_memory=True + ), + event=torch.cuda.Event(), + ) + for _ in range(ring_size) + ] + self.ring_slot_idx = -1 self.idx_mapping_gpu = torch.empty( max_num_reqs, dtype=torch.int32, device=device ) @@ -61,6 +78,18 @@ class PrepareInputsBuffers: max_num_reqs, dtype=torch.int32, device=device ) + def acquire_ring_slot(self) -> PrepareInputsRingSlot: + slot_idx = (self.ring_slot_idx + 1) % len(self.ring_slots) + self.ring_slot_idx = slot_idx + slot = self.ring_slots[slot_idx] + if slot.in_use and not slot.event.query(): + slot.event.synchronize() + return slot + + def mark_ring_slot_inflight(self, slot: PrepareInputsRingSlot) -> None: + slot.event.record(torch.cuda.current_stream(self.device)) + slot.in_use = True + class UvaBuffer: def __init__(self, size: int | Sequence[int], dtype: torch.dtype): diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 265353e01f8..0ecf442ce87 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -190,7 +190,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): max_num_tokens=self.max_num_tokens, device=self.device, ) + # ring-buffered staging for prepare_inputs self._prepare_inputs_buffers = PrepareInputsBuffers( + ring_size=max(self.pp_size, 2), max_num_reqs=self.max_num_reqs, device=self.device, ) @@ -557,6 +559,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): num_tokens_per_req = scheduler_output.num_scheduled_tokens num_reqs = len(num_tokens_per_req) prep_buffers = self._prepare_inputs_buffers + slot = prep_buffers.acquire_ring_slot() # Decode first, then prefill. # batch_idx -> req_id @@ -566,10 +569,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): idx_mapping_iter = map(self.req_states.req_id_to_index.get, req_ids) idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.int32, count=num_reqs) - idx_mapping_cpu = prep_buffers.idx_mapping_cpu[:num_reqs] + idx_mapping_cpu = slot.idx_mapping_cpu[:num_reqs] np.copyto(idx_mapping_cpu.numpy(), idx_mapping_np, casting="no") idx_mapping = prep_buffers.idx_mapping_gpu[:num_reqs] - idx_mapping.copy_(idx_mapping_cpu, non_blocking=False) + idx_mapping.copy_(idx_mapping_cpu, non_blocking=True) # Get the number of draft tokens for each request. draft_tokens = scheduler_output.scheduled_spec_decode_tokens @@ -590,12 +593,13 @@ class GPUModelRunner(LoRAModelRunnerMixin): total_num_logits = num_reqs + total_num_draft_tokens num_logits = num_draft_tokens + 1 - cu_num_logits_cpu = prep_buffers.cu_num_logits_cpu[: num_reqs + 1] + cu_num_logits_cpu = slot.cu_num_logits_cpu[: num_reqs + 1] cu_num_logits_cpu_np = cu_num_logits_cpu.numpy() cu_num_logits_cpu_np[0] = 0 np.cumsum(num_logits, out=cu_num_logits_cpu_np[1:]) cu_num_logits = prep_buffers.cu_num_logits_gpu[: num_reqs + 1] - cu_num_logits.copy_(cu_num_logits_cpu, non_blocking=False) + cu_num_logits.copy_(cu_num_logits_cpu, non_blocking=True) + # keep an independent CPU snapshot because ring slots are reused. cu_num_logits_np = cu_num_logits_cpu_np.copy() max_expand_len = self.num_speculative_steps + 1 @@ -604,7 +608,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) # Get query_start_loc. - query_start_loc_cpu_full = prep_buffers.query_start_loc_cpu + query_start_loc_cpu_full = slot.query_start_loc_cpu query_start_loc_np_full = query_start_loc_cpu_full.numpy() query_start_loc_np_full[0] = 0 np.cumsum(num_scheduled_tokens, out=query_start_loc_np_full[1 : num_reqs + 1]) @@ -612,9 +616,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Some attention backends like FA3 require query_start_loc to be non-decreasing. query_start_loc_np_full[num_reqs + 1 :] = num_tokens self.input_buffers.query_start_loc.copy_( - query_start_loc_cpu_full, non_blocking=False + query_start_loc_cpu_full, non_blocking=True ) + prep_buffers.mark_ring_slot_inflight(slot) + # keep an independent CPU snapshot because ring slots are reused. query_start_loc_np = query_start_loc_np_full[: num_reqs + 1].copy() query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]