karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

Revert "remove ring buffer"

Browse Source 
This reverts commit 48a772a0d4.
This commit is contained in:
yewentao256
2026-03-02 20:59:14 +00:00
parent d526f9d91f
commit e2ded9d884
2 changed files with 51 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/v1/worker/gpu/buffer_utils.py
+39 -10
View File
@@ -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):
vllm/v1/worker/gpu/model_runner.py
+12 -6
View File
@@ -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]