@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 025a32f9ed
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, 07:45 UTC

[Model Runner V2] Remove async barrier (#32083)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Jan 11, 2026 at 20:24 UTC 025a32f9ed53b69c90be8a8883f5c9d880880d8a
13 files changed +590 -462
vllm/v1/worker/gpu/block_table.py
+18 -34
@@ -6,9 +6,8 @@ import torch
6
7 from vllm.triton_utils import tl, triton
8 from vllm.utils.math_utils import cdiv
9 -from vllm.utils.platform_utils import is_uva_available
10 -from vllm.utils.torch_utils import get_cuda_view_from_cpu_tensor
9 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
10 +from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
11
12
13 class BlockTables:
@@ -26,19 +25,16 @@ class BlockTables:
25 self.max_model_len = max_model_len
26 self.device = device
27
29 - if not is_uva_available():
30 - raise RuntimeError("UVA is not available")
31 -
28 self.num_kv_cache_groups = len(self.block_sizes)
29 # num_kv_cache_groups x [max_num_reqs, max_num_blocks]
34 - self.block_tables: list[UvaBuffer] = []
30 + self.block_tables: list[StagedWriteTensor] = []
31 for i in range(self.num_kv_cache_groups):
32 block_size = self.block_sizes[i]
33 max_num_blocks = cdiv(self.max_model_len, block_size)
38 - block_table = UvaBuffer(
39 - self.max_num_reqs,
40 - max_num_blocks,
34 + block_table = StagedWriteTensor(
35 + (self.max_num_reqs, max_num_blocks),
36 dtype=torch.int32,
37 + device=device,
38 )
39 self.block_tables.append(block_table)
40 self.block_table_ptrs = self._make_ptr_tensor(
@@ -53,9 +49,8 @@ class BlockTables:
49 self.block_sizes_tensor = torch.tensor(
50 self.block_sizes, dtype=torch.int32, device=self.device
51 )
56 - self.num_blocks = UvaBuffer(
57 - self.num_kv_cache_groups,
58 - self.max_num_reqs,
52 + self.num_blocks = UvaBackedTensor(
53 + (self.num_kv_cache_groups, self.max_num_reqs),
54 dtype=torch.int32,
55 )
56
@@ -75,13 +70,11 @@ class BlockTables:
70
71 def _make_ptr_tensor(self, x: Iterable[torch.Tensor]) -> torch.Tensor:
72 # NOTE(woosuk): Use uint64 instead of int64 to cover all possible addresses.
78 - ptrs_tensor_cpu = torch.tensor(
73 + return torch.tensor(
74 [t.data_ptr() for t in x],
75 dtype=torch.uint64,
81 - device="cpu",
82 - pin_memory=True,
76 + device=self.device,
77 )
84 - return ptrs_tensor_cpu.to(self.device, non_blocking=True)
78
79 def append_block_ids(
80 self,
@@ -90,19 +83,17 @@ class BlockTables:
83 overwrite: bool,
84 ) -> None:
85 for i in range(self.num_kv_cache_groups):
86 + start = self.num_blocks.np[i, req_index] if not overwrite else 0
87 block_ids = new_block_ids[i]
94 - num_new_blocks = len(block_ids)
95 - if num_new_blocks == 0:
96 - continue
88 + self.block_tables[i].stage_write(req_index, start, block_ids)
89 + self.num_blocks.np[i, req_index] = start + len(block_ids)
90
98 - # TODO(woosuk): Too many Numpy invocations. Optimize this.
99 - start = self.num_blocks.np[i, req_index] if not overwrite else 0
100 - end = start + num_new_blocks
101 - if num_new_blocks == 1:
102 - self.block_tables[i].np[req_index, start] = block_ids[0]
103 - else:
104 - self.block_tables[i].np[req_index, start:end] = block_ids
105 - self.num_blocks.np[i, req_index] = end
91 + def apply_staged_writes(self) -> None:
92 + # TODO(woosuk): This can be inefficient since it launches one kernel per
93 + # block table. Implement a kernel to handle all block tables at once.
94 + for block_table in self.block_tables:
95 + block_table.apply_write()
96 + self.num_blocks.copy_to_uva()
97
98 def gather_block_tables(
99 self,
@@ -229,10 +220,3 @@ def _load_ptr(ptr_to_ptr, elem_dtype):
220 ptr = tl.load(ptr_to_ptr)
221 ptr = tl.cast(ptr, tl.pointer_type(elem_dtype))
222 return tl.multiple_of(ptr, 16)
232 -
233 -
234 -class UvaBuffer:
235 - def __init__(self, *size, dtype: torch.dtype):
236 - self.cpu = torch.zeros(*size, dtype=dtype, device="cpu", pin_memory=True)
237 - self.np = self.cpu.numpy()
238 - self.gpu = get_cuda_view_from_cpu_tensor(self.cpu)
vllm/v1/worker/gpu/buffer_utils.py new
+218
@@ -0,0 +1,218 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Sequence
4 +
5 +import numpy as np
6 +import torch
7 +
8 +from vllm.triton_utils import tl, triton
9 +from vllm.utils.math_utils import next_power_of_2
10 +from vllm.utils.platform_utils import is_uva_available
11 +from vllm.utils.torch_utils import get_cuda_view_from_cpu_tensor
12 +
13 +
14 +class UvaBuffer:
15 + def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
16 + if not is_uva_available():
17 + raise RuntimeError("UVA is not available")
18 + self.cpu = torch.zeros(size, dtype=dtype, device="cpu", pin_memory=True)
19 + self.np = self.cpu.numpy()
20 + self.uva = get_cuda_view_from_cpu_tensor(self.cpu)
21 +
22 +
23 +class UvaBufferPool:
24 + def __init__(
25 + self,
26 + size: int | Sequence[int],
27 + dtype: torch.dtype,
28 + max_concurrency: int = 2,
29 + ):
30 + self.size = size
31 + self.dtype = dtype
32 + self.max_concurrency = max_concurrency
33 +
34 + # UVA buffers for concurrency
35 + self._uva_bufs = [UvaBuffer(size, dtype) for _ in range(max_concurrency)]
36 + # Current buffer index
37 + self._curr = 0
38 +
39 + def copy_to_uva(self, x: torch.Tensor | np.ndarray | list) -> torch.Tensor:
40 + # Round robin to the next buffer.
41 + self._curr = (self._curr + 1) % self.max_concurrency
42 + buf = self._uva_bufs[self._curr]
43 + # CPU-to-CPU copy
44 + dst = buf.cpu if isinstance(x, torch.Tensor) else buf.np
45 + n = len(x)
46 + dst[:n] = x
47 + return buf.uva[:n]
48 +
49 + def copy_to_gpu(
50 + self,
51 + x: torch.Tensor | np.ndarray,
52 + out: torch.Tensor | None = None,
53 + ) -> torch.Tensor:
54 + uva = self.copy_to_uva(x)
55 + if out is None:
56 + # CPU-to-GPU copy
57 + return uva.clone()
58 + # CPU-to-GPU copy
59 + return out.copy_(uva, non_blocking=True)
60 +
61 +
62 +class UvaBackedTensor:
63 + def __init__(
64 + self,
65 + size: int | Sequence[int],
66 + dtype: torch.dtype,
67 + max_concurrency: int = 2,
68 + ):
69 + self.dtype = dtype
70 + self.max_concurrency = max_concurrency
71 +
72 + # Source of truth
73 + self.cpu = torch.zeros(size, dtype=dtype, device="cpu", pin_memory=False)
74 + self.np = self.cpu.numpy()
75 +
76 + # Buffers for concurrency
77 + self.pool = UvaBufferPool(size, dtype, max_concurrency)
78 + self.gpu = self.pool.copy_to_uva(self.np)
79 +
80 + def copy_to_uva(self, n: int | None = None) -> torch.Tensor:
81 + # CPU-to-CPU copy
82 + self.gpu = self.pool.copy_to_uva(self.np[:n] if n is not None else self.np)
83 + return self.gpu
84 +
85 +
86 +class StagedWriteTensor:
87 + def __init__(
88 + self,
89 + size: int | Sequence[int],
90 + dtype: torch.dtype,
91 + device: torch.device,
92 + max_concurrency: int = 2,
93 + uva_instead_of_gpu: bool = False,
94 + ):
95 + if dtype not in [torch.int32, torch.int64]:
96 + raise ValueError(
97 + f"Unsupported dtype {dtype}: should be either int32 or int64"
98 + )
99 + self.num_rows = size if isinstance(size, int) else size[0]
100 + self.dtype = dtype
101 + self.max_concurrency = max_concurrency
102 +
103 + if not uva_instead_of_gpu:
104 + # Create a GPU tensor (default)
105 + self.gpu = torch.zeros(size, dtype=dtype, device=device)
106 + else:
107 + # For a large but not-frequently-accessed tensor, we can use UVA instead of
108 + # GPU to save GPU memory
109 + self._uva_buf = UvaBuffer(size, dtype)
110 + self.gpu = self._uva_buf.uva
111 +
112 + self._staged_write_indices: list[int] = []
113 + self._staged_write_starts: list[int] = []
114 + self._staged_write_contents: list[int] = []
115 + self._staged_write_cu_lens: list[int] = []
116 +
117 + self.write_indices = UvaBufferPool(
118 + self.num_rows, dtype=torch.int32, max_concurrency=max_concurrency
119 + )
120 + self.write_starts = UvaBufferPool(
121 + self.num_rows, dtype=torch.int32, max_concurrency=max_concurrency
122 + )
123 + init_size = next_power_of_2(self.num_rows)
124 + self.write_contents = UvaBufferPool(
125 + init_size, dtype=dtype, max_concurrency=max_concurrency
126 + )
127 + self.write_cu_lens = UvaBufferPool(
128 + self.num_rows, dtype=torch.int32, max_concurrency=max_concurrency
129 + )
130 +
131 + def stage_write(self, index: int, start: int, x: list[int]) -> None:
132 + assert index >= 0
133 + assert start >= 0
134 + if not x:
135 + return
136 + self._staged_write_indices.append(index)
137 + self._staged_write_starts.append(start)
138 + self._staged_write_contents.extend(x)
139 + self._staged_write_cu_lens.append(len(self._staged_write_contents))
140 +
141 + def stage_write_elem(self, index: int, x: int) -> None:
142 + assert index >= 0
143 + self._staged_write_indices.append(index)
144 + self._staged_write_starts.append(0)
145 + self._staged_write_contents.append(x)
146 + self._staged_write_cu_lens.append(len(self._staged_write_contents))
147 +
148 + def apply_write(self) -> None:
149 + n = len(self._staged_write_indices)
150 + if n == 0:
151 + return
152 +
153 + indices_uva = self.write_indices.copy_to_uva(self._staged_write_indices)
154 + starts_uva = self.write_starts.copy_to_uva(self._staged_write_starts)
155 + cu_lens_uva = self.write_cu_lens.copy_to_uva(self._staged_write_cu_lens)
156 +
157 + # Special handling for write_contents
158 + diff_len = len(self._staged_write_contents)
159 + assert isinstance(self.write_contents.size, int)
160 + if diff_len > self.write_contents.size:
161 + # Re-allocate a larger buffer for the write_contents
162 + new_size = next_power_of_2(diff_len)
163 + self.write_contents = UvaBufferPool(
164 + new_size, dtype=self.dtype, max_concurrency=self.max_concurrency
165 + )
166 + # NOTE(woosuk): Since the previous write_contents buffer is released,
167 + # we perform a synchronization here to ensure that all data transfers
168 + # involving the old buffer have finished before allocating a new one.
169 + # This prevents potential race conditions. The slight overhead is
170 + # negligible because the reallocations are infrequent in practice.
171 + torch.cuda.synchronize()
172 + contents_uva = self.write_contents.copy_to_uva(self._staged_write_contents)
173 +
174 + # Write diffs to the GPU buffer
175 + _apply_write_kernel[(n,)](
176 + self.gpu,
177 + self.gpu.stride(0),
178 + indices_uva,
179 + starts_uva,
180 + contents_uva,
181 + cu_lens_uva,
182 + BLOCK_SIZE=1024,
183 + )
184 + # Clear the staged writes
185 + self.clear_staged_writes()
186 +
187 + def clear_staged_writes(self) -> None:
188 + self._staged_write_indices.clear()
189 + self._staged_write_starts.clear()
190 + self._staged_write_contents.clear()
191 + self._staged_write_cu_lens.clear()
192 +
193 +
194 +@triton.jit
195 +def _apply_write_kernel(
196 + output_ptr,
197 + output_stride,
198 + write_indices_ptr,
199 + write_starts_ptr,
200 + write_contents_ptr,
201 + write_cu_lens_ptr,
202 + BLOCK_SIZE: tl.constexpr,
203 +):
204 + pid = tl.program_id(0)
205 + row_idx = tl.load(write_indices_ptr + pid)
206 + start_idx = tl.load(write_starts_ptr + pid)
207 +
208 + cu_start = tl.load(write_cu_lens_ptr + pid - 1) if pid > 0 else 0
209 + cu_end = tl.load(write_cu_lens_ptr + pid)
210 + content_len = cu_end - cu_start
211 +
212 + for i in range(0, content_len, BLOCK_SIZE):
213 + block = i + tl.arange(0, BLOCK_SIZE)
214 + mask = block < content_len
215 + content = tl.load(write_contents_ptr + cu_start + block, mask=mask)
216 + tl.store(
217 + output_ptr + row_idx * output_stride + start_idx + block, content, mask=mask
218 + )
vllm/v1/worker/gpu/cudagraph_utils.py
+9 -6
@@ -228,10 +228,13 @@ def prepare_inputs_to_capture(
228 kv_cache_config: KVCacheConfig,
229 ) -> dict[str, Any]:
230 num_tokens_per_req = num_tokens // num_reqs
231 - query_start_loc = input_buffers.query_start_loc
232 - query_start_loc.np[: num_reqs + 1] = np.arange(num_reqs + 1) * num_tokens_per_req
233 - query_start_loc.np[num_reqs:] = num_tokens
234 - query_start_loc.copy_to_gpu()
231 +
232 + query_start_loc_np = np.arange(num_reqs + 1, dtype=np.int32) * num_tokens_per_req
233 + query_start_loc_np[-1] = num_tokens
234 + query_start_loc_cpu = torch.from_numpy(query_start_loc_np)
235 + input_buffers.query_start_loc[: num_reqs + 1] = query_start_loc_cpu
236 + input_buffers.query_start_loc[num_reqs + 1 :] = num_tokens
237 + query_start_loc = input_buffers.query_start_loc[: num_reqs + 1]
238
239 # HACK(woosuk): For faster warmup, we set seq_lens (GPU) to num_tokens
240 # rather than max_model_len.
@@ -245,8 +248,8 @@ def prepare_inputs_to_capture(
248 attn_metadata_builders=attn_metadata_builders,
249 num_reqs=num_reqs,
250 num_tokens=num_tokens,
248 - query_start_loc_gpu=query_start_loc.gpu[: num_reqs + 1],
249 - query_start_loc_cpu=query_start_loc.cpu[: num_reqs + 1],
251 + query_start_loc_gpu=query_start_loc,
252 + query_start_loc_cpu=query_start_loc_cpu,
253 seq_lens=input_buffers.seq_lens,
254 max_seq_len=max_model_len,
255 block_tables=input_block_tables,
vllm/v1/worker/gpu/input_batch.py
+57 -25
@@ -8,8 +8,6 @@ import torch
8
9 from vllm.triton_utils import tl, triton
10 from vllm.utils import random_uuid
11 -from vllm.utils.math_utils import cdiv
12 -from vllm.v1.utils import CpuGpuBuffer
11
12
13 class InputBuffers:
@@ -21,30 +19,17 @@ class InputBuffers:
19 vocab_size: int,
20 dtype: torch.dtype,
21 device: torch.device,
24 - pin_memory: bool,
22 ):
23 self.max_num_reqs = max_num_reqs
24 self.max_num_tokens = max_num_tokens
25 self.device = device
29 - self.pin_memory = pin_memory
26
31 - self.idx_mapping = self._make_buffer(max_num_reqs, dtype=torch.int32)
27 self.input_ids = torch.zeros(max_num_tokens, dtype=torch.int32, device=device)
28 self.positions = torch.zeros(max_num_tokens, dtype=torch.int64, device=device)
34 - self.query_start_loc = self._make_buffer(max_num_reqs + 1, dtype=torch.int32)
35 - self.seq_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
36 - self.cu_num_logits = self._make_buffer(max_num_reqs + 1, dtype=torch.int32)
37 -
38 - # Structured outputs.
39 - self.bitmask_indices = self._make_buffer(max_num_reqs, dtype=torch.int32)
40 - self.grammar_bitmask = self._make_buffer(
41 - max_num_reqs, cdiv(vocab_size, 32), dtype=torch.int32
42 - )
43 -
44 - def _make_buffer(self, *args, dtype: torch.dtype) -> CpuGpuBuffer:
45 - return CpuGpuBuffer(
46 - *args, dtype=dtype, pin_memory=self.pin_memory, device=self.device
29 + self.query_start_loc = torch.zeros(
30 + max_num_reqs + 1, dtype=torch.int32, device=device
31 )
32 + self.seq_lens = torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
33
34
35 @dataclass
@@ -56,6 +41,8 @@ class InputBatch:
41 # batch_idx -> req_state_idx
42 idx_mapping: torch.Tensor
43 idx_mapping_np: np.ndarray
44 + # Identical to idx_mapping except for spec decoding.
45 + expanded_idx_mapping: torch.Tensor
46
47 # [num_reqs]
48 # batch_idx -> num_scheduled_tokens
@@ -83,6 +70,7 @@ class InputBatch:
70 logits_indices: torch.Tensor
71 # [num_reqs + 1]
72 cu_num_logits: torch.Tensor
73 + cu_num_logits_np: np.ndarray
74
75 @classmethod
76 def make_dummy(
@@ -96,33 +84,41 @@ class InputBatch:
84 req_ids = [f"req_{i}_{random_uuid()}" for i in range(num_reqs)]
85 idx_mapping_np = np.arange(num_reqs, dtype=np.int32)
86 idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
87 + expanded_idx_mapping = idx_mapping
88 num_scheduled_tokens = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)
89 num_scheduled_tokens[-1] += num_tokens % num_reqs
90 assert int(num_scheduled_tokens.sum()) == num_tokens
91
103 - input_buffers.query_start_loc.np[0] = 0
104 - input_buffers.query_start_loc.np[1 : num_reqs + 1] = np.cumsum(
105 - num_scheduled_tokens
106 - )
107 - input_buffers.query_start_loc.np[num_reqs + 1 :] = num_tokens
108 - query_start_loc_np = input_buffers.query_start_loc.np[: num_reqs + 1]
109 - query_start_loc = input_buffers.query_start_loc.copy_to_gpu()[: num_reqs + 1]
92 # seq_len equals to query_len
93 input_buffers.seq_lens[:num_reqs] = num_tokens // num_reqs
94 input_buffers.seq_lens[num_reqs - 1] += num_tokens % num_reqs
95 + # Pad for full CUDA graph mode.
96 input_buffers.seq_lens[num_reqs:] = 0
97 seq_lens = input_buffers.seq_lens[:num_reqs]
98
99 + query_start_loc_np = np.empty(num_reqs + 1, dtype=np.int32)
100 + query_start_loc_np[0] = 0
101 + np.cumsum(num_scheduled_tokens, out=query_start_loc_np[1:])
102 + input_buffers.query_start_loc[0] = 0
103 + torch.cumsum(
104 + seq_lens, dim=0, out=input_buffers.query_start_loc[1 : num_reqs + 1]
105 + )
106 + # Pad for full CUDA graph mode.
107 + input_buffers.query_start_loc[num_reqs + 1 :] = num_tokens
108 + query_start_loc = input_buffers.query_start_loc[: num_reqs + 1]
109 +
110 input_ids = input_buffers.input_ids[:num_tokens]
111 positions = input_buffers.positions[:num_tokens]
112 # attn_metadata = defaultdict(lambda: None)
113 logits_indices = query_start_loc[1:] - 1
114 cu_num_logits = torch.arange(num_reqs + 1, device=device, dtype=torch.int32)
115 + cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)
116 return cls(
117 req_ids=req_ids,
118 num_reqs=num_reqs,
119 idx_mapping=idx_mapping,
120 idx_mapping_np=idx_mapping_np,
121 + expanded_idx_mapping=expanded_idx_mapping,
122 num_scheduled_tokens=num_scheduled_tokens,
123 num_tokens=num_tokens,
124 num_tokens_after_padding=num_tokens,
@@ -135,6 +131,7 @@ class InputBatch:
131 attn_metadata=None, # type: ignore
132 logits_indices=logits_indices,
133 cu_num_logits=cu_num_logits,
134 + cu_num_logits_np=cu_num_logits_np,
135 )
136
137
@@ -473,3 +470,38 @@ def post_update(
470 query_start_loc,
471 num_warps=1,
472 )
473 +
474 +
475 +@triton.jit
476 +def _expand_idx_mapping_kernel(
477 + idx_mapping_ptr,
478 + expanded_idx_mapping_ptr,
479 + cu_num_logits_ptr,
480 + BLOCK_SIZE: tl.constexpr,
481 +):
482 + req_idx = tl.program_id(0)
483 + start_idx = tl.load(cu_num_logits_ptr + req_idx)
484 + end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
485 + num_tokens = end_idx - start_idx
486 +
487 + block = tl.arange(0, BLOCK_SIZE)
488 + mask = block < num_tokens
489 + req_state_idx = tl.load(idx_mapping_ptr + req_idx)
490 + tl.store(expanded_idx_mapping_ptr + start_idx + block, req_state_idx, mask=mask)
491 +
492 +
493 +def expand_idx_mapping(
494 + idx_mapping: torch.Tensor,
495 + total_num_logits: int,
496 + cu_num_logits: torch.Tensor,
497 + max_expand_len: int,
498 +) -> torch.Tensor:
499 + num_reqs = idx_mapping.shape[0]
500 + expanded_idx_mapping = idx_mapping.new_empty(total_num_logits)
501 + _expand_idx_mapping_kernel[(num_reqs,)](
502 + idx_mapping,
503 + expanded_idx_mapping,
504 + cu_num_logits,
505 + BLOCK_SIZE=triton.next_power_of_2(max_expand_len),
506 + )
507 + return expanded_idx_mapping
vllm/v1/worker/gpu/model_runner.py
+108 -111
@@ -15,7 +15,6 @@ from vllm.forward_context import set_forward_context
15 from vllm.logger import init_logger
16 from vllm.model_executor.model_loader import get_model_loader
17 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
18 -from vllm.utils.platform_utils import is_pin_memory_available
18 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
19 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
20 from vllm.v1.kv_cache_interface import KVCacheConfig
@@ -24,7 +23,7 @@ from vllm.v1.outputs import (
23 LogprobsTensors,
24 ModelRunnerOutput,
25 )
27 -from vllm.v1.worker.gpu.async_utils import AsyncOutput, async_barrier
26 +from vllm.v1.worker.gpu.async_utils import AsyncOutput
27 from vllm.v1.worker.gpu.attn_utils import (
28 build_attn_metadata,
29 get_kv_cache_spec,
@@ -32,6 +31,7 @@ from vllm.v1.worker.gpu.attn_utils import (
31 init_kv_cache,
32 )
33 from vllm.v1.worker.gpu.block_table import BlockTables
34 +from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool
35 from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager
36 from vllm.v1.worker.gpu.dp_utils import (
37 get_batch_metadata_across_dp,
@@ -41,22 +41,20 @@ from vllm.v1.worker.gpu.input_batch import (
41 InputBatch,
42 InputBuffers,
43 combine_sampled_and_draft_tokens,
44 + expand_idx_mapping,
45 get_num_sampled_and_rejected,
46 post_update,
47 prepare_pos_seq_lens,
48 prepare_prefill_inputs,
49 )
50 from vllm.v1.worker.gpu.sample.logprob import compute_prompt_logprobs
50 -from vllm.v1.worker.gpu.sample.metadata import (
51 - SamplingMetadata,
52 - expand_sampling_metadata,
53 -)
51 +from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
52 from vllm.v1.worker.gpu.sample.output import SamplerOutput
53 from vllm.v1.worker.gpu.sample.sampler import Sampler
54 from vllm.v1.worker.gpu.spec_decode import init_speculator
55 from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample
56 from vllm.v1.worker.gpu.states import RequestState
59 -from vllm.v1.worker.gpu.structured_outputs import apply_grammar_bitmask
57 +from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
58 from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
59 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
60
@@ -81,7 +79,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
79 self.observability_config = vllm_config.observability_config
80
81 self.device = device
84 - self.pin_memory = is_pin_memory_available()
82 self.dtype = self.model_config.dtype
83 self.kv_cache_dtype = self.dtype
84 if self.cache_config.cache_dtype != "auto":
@@ -123,7 +120,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
120 num_speculative_steps=self.num_speculative_steps,
121 vocab_size=self.vocab_size,
122 device=self.device,
126 - pin_memory=self.pin_memory,
123 )
124 self.input_buffers = InputBuffers(
125 max_num_reqs=self.max_num_reqs,
@@ -132,12 +128,21 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
128 vocab_size=self.vocab_size,
129 dtype=self.dtype,
130 device=self.device,
135 - pin_memory=self.pin_memory,
131 )
132 self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)
133
134 # CUDA graphs.
135 self.cudagraph_manager = CudaGraphManager(self.vllm_config, self.device)
136 + # Structured outputs worker.
137 + self.structured_outputs_worker = StructuredOutputsWorker(
138 + max_num_logits=self.max_num_reqs * (self.num_speculative_steps + 1),
139 + vocab_size=self.vocab_size,
140 + )
141 +
142 + # Buffers for CPU-to-GPU copies.
143 + self.tmp_idx_mapping = UvaBufferPool(self.max_num_reqs, torch.int32)
144 + self.tmp_cu_num_logits = UvaBufferPool(self.max_num_reqs + 1, torch.int32)
145 + self.tmp_query_start_loc = UvaBufferPool(self.max_num_reqs + 1, torch.int32)
146
147 def update_max_model_len(self, max_model_len: int) -> None:
148 self.max_model_len = max_model_len
@@ -228,16 +233,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
233 slot_mappings = self.block_tables.get_dummy_slot_mappings(
234 input_batch.num_tokens
235 )
231 - query_start_loc = self.input_buffers.query_start_loc
232 - query_start_loc_gpu = query_start_loc.gpu[: input_batch.num_reqs + 1]
233 - query_start_loc_cpu = query_start_loc.cpu[: input_batch.num_reqs + 1]
236 attn_metadata = build_attn_metadata(
237 attn_metadata_builders=self.attn_metadata_builders,
238 num_reqs=input_batch.num_reqs,
239 num_tokens=input_batch.num_tokens,
238 - query_start_loc_gpu=query_start_loc_gpu,
239 - query_start_loc_cpu=query_start_loc_cpu,
240 - seq_lens=self.input_buffers.seq_lens,
240 + query_start_loc_gpu=input_batch.query_start_loc,
241 + query_start_loc_cpu=torch.from_numpy(input_batch.query_start_loc_np),
242 + seq_lens=input_batch.seq_lens,
243 max_seq_len=self.max_model_len,
244 block_tables=block_tables,
245 slot_mappings=slot_mappings,
@@ -396,8 +398,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
398 self.block_tables.append_block_ids(
399 req_index, new_req_data.block_ids, overwrite=True
400 )
399 - if scheduler_output.scheduled_new_reqs:
400 - self.req_states.prefill_len.copy_to_gpu()
401
402 # Add new blocks for the existing requests.
403 cached_reqs = scheduler_output.scheduled_cached_reqs
@@ -409,6 +409,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
409 req_index, req_new_block_ids, overwrite=False
410 )
411
412 + self.req_states.apply_staged_writes()
413 + self.block_tables.apply_staged_writes()
414 +
415 def prepare_inputs(
416 self,
417 scheduler_output: SchedulerOutput,
@@ -431,19 +434,19 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
434 idx_mapping_list = [
435 self.req_states.req_id_to_index[req_id] for req_id in req_ids
436 ]
434 - idx_mapping = self.input_buffers.idx_mapping
435 - idx_mapping.np[:num_reqs] = idx_mapping_list
436 - idx_mapping_np = idx_mapping.np[:num_reqs]
437 - idx_mapping = idx_mapping.copy_to_gpu(num_reqs)
437 + idx_mapping_np = np.array(idx_mapping_list, dtype=np.int32)
438 + idx_mapping = self.tmp_idx_mapping.copy_to_gpu(idx_mapping_np)
439
440 # Get the number of draft tokens for each request.
441 if not scheduler_output.scheduled_spec_decode_tokens:
442 # No draft token scheduled (common case).
443 total_num_draft_tokens = 0
444 total_num_logits = num_reqs
445 + cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)
446 cu_num_logits = torch.arange(
447 num_reqs + 1, device=self.device, dtype=torch.int32
448 )
449 + expanded_idx_mapping = idx_mapping
450 else:
451 draft_tokens = scheduler_output.scheduled_spec_decode_tokens
452 num_draft_tokens = np.array(
@@ -456,44 +459,53 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
459 total_num_draft_tokens = int(num_draft_tokens.sum())
460 total_num_logits = num_reqs + total_num_draft_tokens
461
459 - np.cumsum(
460 - num_draft_tokens + 1,
461 - out=self.input_buffers.cu_num_logits.np[1 : num_reqs + 1],
462 + num_logits = num_draft_tokens + 1
463 + cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32)
464 + cu_num_logits_np[0] = 0
465 + np.cumsum(num_logits, out=cu_num_logits_np[1:])
466 + cu_num_logits = self.tmp_cu_num_logits.copy_to_gpu(cu_num_logits_np)
467 +
468 + expanded_idx_mapping = expand_idx_mapping(
469 + idx_mapping,
470 + total_num_logits,
471 + cu_num_logits,
472 + max_expand_len=self.num_speculative_steps + 1,
473 )
463 - cu_num_logits = self.input_buffers.cu_num_logits.copy_to_gpu(num_reqs + 1)
474
475 # Block tables: num_kv_cache_groups x [num_reqs, max_num_blocks]
476 block_tables = self.block_tables.gather_block_tables(idx_mapping)
477
478 # Get query_start_loc.
469 - np.cumsum(
470 - num_scheduled_tokens,
471 - out=self.input_buffers.query_start_loc.np[1 : num_reqs + 1],
472 - )
479 + query_start_loc_np = np.empty(self.max_num_reqs + 1, dtype=np.int32)
480 + query_start_loc_np[0] = 0
481 + np.cumsum(num_scheduled_tokens, out=query_start_loc_np[1 : num_reqs + 1])
482 # Pad for full CUDA graph mode.
483 # Some attention backends like FA3 require query_start_loc to be non-decreasing.
475 - self.input_buffers.query_start_loc.np[num_reqs + 1 :] = num_tokens
476 - self.input_buffers.query_start_loc.copy_to_gpu()
477 - query_start_loc_gpu = self.input_buffers.query_start_loc.gpu[: num_reqs + 1]
478 - query_start_loc_cpu = self.input_buffers.query_start_loc.cpu[: num_reqs + 1]
479 - query_start_loc_np = self.input_buffers.query_start_loc.np[: num_reqs + 1]
484 + query_start_loc_np[num_reqs + 1 :] = num_tokens
485 + self.tmp_query_start_loc.copy_to_gpu(
486 + query_start_loc_np,
487 + out=self.input_buffers.query_start_loc,
488 + )
489 + query_start_loc_np = query_start_loc_np[: num_reqs + 1]
490 + query_start_loc_cpu = torch.from_numpy(query_start_loc_np)
491 + query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
492
493 # Get prefill tokens.
494 prepare_prefill_inputs(
495 self.input_buffers.input_ids,
496 self.req_states.next_prefill_tokens,
497 idx_mapping,
486 - query_start_loc_gpu,
498 + query_start_loc,
499 self.req_states.prefill_token_ids.gpu,
500 self.req_states.prefill_len.gpu,
489 - self.req_states.num_computed_tokens,
501 + self.req_states.num_computed_tokens.gpu,
502 )
503
504 # Prepare positions and seq_lens.
505 prepare_pos_seq_lens(
506 idx_mapping,
495 - query_start_loc_gpu,
496 - self.req_states.num_computed_tokens,
507 + query_start_loc,
508 + self.req_states.num_computed_tokens.gpu,
509 self.input_buffers.positions,
510 self.input_buffers.seq_lens,
511 )
@@ -505,7 +517,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
517 self.input_buffers.input_ids,
518 idx_mapping,
519 self.req_states.last_sampled_tokens,
508 - query_start_loc_gpu,
520 + query_start_loc,
521 seq_lens,
522 self.req_states.prefill_len.gpu,
523 self.req_states.draft_tokens,
@@ -515,7 +527,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
527
528 # Compute slot mappings: [num_kv_cache_groups, num_tokens]
529 slot_mappings = self.block_tables.compute_slot_mappings(
518 - query_start_loc_gpu, self.input_buffers.positions[:num_tokens]
530 + query_start_loc, self.input_buffers.positions[:num_tokens]
531 )
532
533 # Layer name -> attention metadata.
@@ -523,7 +535,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
535 attn_metadata_builders=self.attn_metadata_builders,
536 num_reqs=num_reqs,
537 num_tokens=num_tokens,
526 - query_start_loc_gpu=query_start_loc_gpu,
538 + query_start_loc_gpu=query_start_loc,
539 query_start_loc_cpu=query_start_loc_cpu,
540 seq_lens=self.input_buffers.seq_lens,
541 max_seq_len=self.max_model_len,
@@ -539,11 +551,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
551 num_reqs=num_reqs,
552 idx_mapping=idx_mapping,
553 idx_mapping_np=idx_mapping_np,
554 + expanded_idx_mapping=expanded_idx_mapping,
555 num_scheduled_tokens=num_scheduled_tokens,
556 num_tokens=num_tokens,
557 num_tokens_after_padding=num_tokens_after_padding,
558 num_draft_tokens=total_num_draft_tokens,
546 - query_start_loc=query_start_loc_gpu,
559 + query_start_loc=query_start_loc,
560 query_start_loc_np=query_start_loc_np,
561 seq_lens=seq_lens,
562 input_ids=input_ids,
@@ -551,6 +564,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
564 attn_metadata=attn_metadata,
565 logits_indices=logits_indices,
566 cu_num_logits=cu_num_logits,
567 + cu_num_logits_np=cu_num_logits_np,
568 )
569
570 def sample(
@@ -564,16 +578,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
578 logits = self.model.compute_logits(sample_hidden_states)
579 if grammar_output is not None:
580 # Apply grammar bitmask to the logits in-place.
567 - # TODO(woosuk): Make compatible with spec decoding.
568 - assert input_batch.num_draft_tokens == 0
569 - with async_barrier(self.structured_outputs_event):
570 - apply_grammar_bitmask(
571 - logits,
572 - input_batch.req_ids,
573 - grammar_output.structured_output_request_ids,
574 - grammar_output.grammar_bitmask,
575 - self.input_buffers,
576 - )
581 + self.structured_outputs_worker.apply_grammar_bitmask(
582 + logits,
583 + input_batch,
584 + grammar_output.structured_output_request_ids,
585 + grammar_output.grammar_bitmask,
586 + )
587
588 # Sample tokens and compute logprobs (if needed).
589 sampler_output = self.sampler(logits, sampling_metadata)
@@ -641,8 +651,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
651 # Handle chunked prompts.
652 pos_after_step = computed_prefill + input_batch.num_scheduled_tokens
653 is_prompt_chunked = pos_after_step < prompt_lens
644 - prefill_token_ids = self.req_states.prefill_token_ids.np
645 - query_start_loc = self.input_buffers.query_start_loc.np
654 + prefill_token_ids = self.req_states.prefill_token_ids.gpu
655 + query_start_loc_np = input_batch.query_start_loc_np
656 for i, req_id in enumerate(input_batch.req_ids):
657 if not needs_prompt_logprobs[i]:
658 continue
@@ -650,10 +660,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
660 continue
661 # The prompt is chunked. Get the next prompt token.
662 req_idx = input_batch.idx_mapping_np[i]
653 - next_prompt_token = int(prefill_token_ids[req_idx, pos_after_step[i]])
654 - idx = int(query_start_loc[i + 1] - 1)
655 - # Set the next prompt token.
656 - # NOTE(woosuk): This triggers a GPU operation.
663 + idx = int(query_start_loc_np[i + 1] - 1)
664 + # NOTE(woosuk): This triggers two GPU operations.
665 + next_prompt_token = prefill_token_ids[req_idx, pos_after_step[i]]
666 token_ids[idx] = next_prompt_token
667
668 # NOTE(woosuk): We mask out logprobs for negative tokens.
@@ -669,8 +678,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
678 if not needs_prompt_logprobs[i]:
679 continue
680
672 - start_idx = query_start_loc[i]
673 - end_idx = query_start_loc[i + 1]
681 + start_idx = query_start_loc_np[i]
682 + end_idx = query_start_loc_np[i + 1]
683 assert start_idx < end_idx, (
684 f"start_idx ({start_idx}) >= end_idx ({end_idx})"
685 )
@@ -714,7 +723,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
723 # Update the number of computed tokens.
724 post_update(
725 input_batch.idx_mapping,
717 - self.req_states.num_computed_tokens,
726 + self.req_states.num_computed_tokens.gpu,
727 self.req_states.last_sampled_tokens,
728 self.req_states.output_bin_counts,
729 sampled_tokens,
@@ -825,61 +834,49 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
834 assert intermediate_tensors is None
835 if scheduler_output.total_num_scheduled_tokens == 0 and not dummy_run:
836 # No need to run the model.
828 - with async_barrier(self.input_prep_event):
829 - self.update_states(scheduler_output)
830 - return EMPTY_MODEL_RUNNER_OUTPUT
837 + self.update_states(scheduler_output)
838 + return EMPTY_MODEL_RUNNER_OUTPUT
839
832 - # NOTE: Call this before the async barrier so CPU all-reduce and
833 - # GPU execution can overlap.
840 cudagraph_mode, num_tokens_after_padding, num_tokens_across_dp = (
841 self.get_cudagraph_and_dp_padding(scheduler_output)
842 )
837 - with async_barrier(self.input_prep_event):
838 - self.update_states(scheduler_output)
839 - if num_tokens_after_padding == 0:
840 - # All DP ranks have zero tokens to run.
841 - return EMPTY_MODEL_RUNNER_OUTPUT
842 -
843 - if not dummy_run:
844 - # Common case.
845 - # Prepare all the inputs and copy to the input buffers.
846 - input_batch = self.prepare_inputs(
847 - scheduler_output,
848 - num_tokens_after_padding,
849 - )
843 + self.update_states(scheduler_output)
844 + if num_tokens_after_padding == 0:
845 + # All DP ranks have zero tokens to run.
846 + return EMPTY_MODEL_RUNNER_OUTPUT
847 +
848 + if not dummy_run:
849 + # Common case.
850 + # Prepare all the inputs and copy to the input buffers.
851 + input_batch = self.prepare_inputs(
852 + scheduler_output,
853 + num_tokens_after_padding,
854 + )
855
851 - # NOTE(woosuk): Sampling metadata should be built under the async
852 - # barrier to avoid race conditions.
853 - pos = input_batch.positions[input_batch.logits_indices]
854 - sampling_metadata = self.req_states.make_sampling_metadata(
855 - input_batch.idx_mapping, input_batch.idx_mapping_np, pos
856 - )
857 - if input_batch.num_draft_tokens > 0:
858 - sampling_metadata = expand_sampling_metadata(
859 - sampling_metadata,
860 - input_batch.cu_num_logits,
861 - max_expand_len=self.num_speculative_steps + 1,
862 - )
863 -
864 - if self.lora_config:
865 - # Activate LoRA adapters.
866 - lora_inputs = self.req_states.make_lora_inputs(
867 - input_batch.req_ids,
868 - input_batch.idx_mapping_np,
869 - input_batch.num_scheduled_tokens,
870 - )
871 - self._set_active_loras(*lora_inputs)
872 - else:
873 - # No actual tokens to run. A dummy run for DP.
874 - num_reqs = min(num_tokens_after_padding, self.max_num_reqs)
875 - input_batch = InputBatch.make_dummy(
876 - num_reqs=num_reqs,
877 - num_tokens=num_tokens_after_padding,
878 - input_buffers=self.input_buffers,
879 - device=self.device,
856 + pos = input_batch.positions[input_batch.logits_indices]
857 + sampling_metadata = self.req_states.make_sampling_metadata(
858 + input_batch.expanded_idx_mapping, input_batch.idx_mapping_np, pos
859 + )
860 +
861 + if self.lora_config:
862 + # Activate LoRA adapters.
863 + lora_inputs = self.req_states.make_lora_inputs(
864 + input_batch.req_ids,
865 + input_batch.idx_mapping_np,
866 + input_batch.num_scheduled_tokens,
867 )
881 - self.prepare_dummy_attn_metadata(input_batch)
882 - sampling_metadata = None
868 + self._set_active_loras(*lora_inputs)
869 + else:
870 + # No actual tokens to run. A dummy run for DP.
871 + num_reqs = min(num_tokens_after_padding, self.max_num_reqs)
872 + input_batch = InputBatch.make_dummy(
873 + num_reqs=num_reqs,
874 + num_tokens=num_tokens_after_padding,
875 + input_buffers=self.input_buffers,
876 + device=self.device,
877 + )
878 + self.prepare_dummy_attn_metadata(input_batch)
879 + sampling_metadata = None
880
881 # Run model.
882 if cudagraph_mode == CUDAGraphMode.FULL:
vllm/v1/worker/gpu/sample/gumbel.py
+12 -7
@@ -13,6 +13,7 @@ def _gumbel_sample_kernel(
13 local_max_stride,
14 logits_ptr,
15 logits_stride,
16 + idx_mapping_ptr,
17 seeds_ptr,
18 pos_ptr,
19 temp_ptr,
@@ -20,22 +21,24 @@ def _gumbel_sample_kernel(
21 BLOCK_SIZE: tl.constexpr,
22 APPLY_TEMPERATURE: tl.constexpr,
23 ):
23 - req_idx = tl.program_id(0)
24 + batch_idx = tl.program_id(0)
25 + req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
26 +
27 block_idx = tl.program_id(1)
28 block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
29 mask = block < vocab_size
30 logits = tl.load(
28 - logits_ptr + req_idx * logits_stride + block,
31 + logits_ptr + batch_idx * logits_stride + block,
32 mask=mask,
33 other=float("-inf"),
34 )
35 logits = logits.to(tl.float32)
36
34 - temp = tl.load(temp_ptr + req_idx).to(tl.float32)
37 + temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
38 if temp != 0.0:
39 # Calculate the seed for gumbel noise.
37 - seed = tl.load(seeds_ptr + req_idx)
38 - pos = tl.load(pos_ptr + req_idx)
40 + seed = tl.load(seeds_ptr + req_state_idx)
41 + pos = tl.load(pos_ptr + batch_idx)
42 gumbel_seed = tl.randint(seed, pos)
43
44 # Generate gumbel noise.
@@ -55,12 +58,13 @@ def _gumbel_sample_kernel(
58 idx = tl.argmax(logits, axis=0)
59 token_id = block_idx * BLOCK_SIZE + idx
60 value = tl.max(logits, axis=0)
58 - tl.store(local_argmax_ptr + req_idx * local_argmax_stride + block_idx, token_id)
59 - tl.store(local_max_ptr + req_idx * local_max_stride + block_idx, value)
61 + tl.store(local_argmax_ptr + batch_idx * local_argmax_stride + block_idx, token_id)
62 + tl.store(local_max_ptr + batch_idx * local_max_stride + block_idx, value)
63
64
65 def gumbel_sample(
66 logits: torch.Tensor, # [num_reqs, vocab_size]
67 + idx_mapping: torch.Tensor, # [num_reqs]
68 temperature: torch.Tensor, # [num_reqs]
69 seed: torch.Tensor, # [num_reqs]
70 pos: torch.Tensor, # [num_reqs]
@@ -88,6 +92,7 @@ def gumbel_sample(
92 local_max.stride(0),
93 logits,
94 logits.stride(0),
95 + idx_mapping,
96 seed,
97 pos,
98 temperature,
vllm/v1/worker/gpu/sample/metadata.py
+14 -127
@@ -4,20 +4,23 @@ from dataclasses import dataclass
4
5 import torch
6
7 -from vllm.triton_utils import tl, triton
8 -
7
8 @dataclass
9 class SamplingMetadata:
10 + idx_mapping: torch.Tensor
11 +
12 temperature: torch.Tensor
13
14 top_p: torch.Tensor | None
15 top_k: torch.Tensor | None
16 min_p: torch.Tensor | None
17
18 + # For penalties
19 repetition_penalty: torch.Tensor
20 frequency_penalty: torch.Tensor
21 presence_penalty: torch.Tensor
22 + prompt_bin_mask: torch.Tensor
23 + output_bin_counts: torch.Tensor
24
25 seeds: torch.Tensor
26 pos: torch.Tensor
@@ -25,11 +28,6 @@ class SamplingMetadata:
28 # None means no logprobs, 0 means sampled token logprobs only
29 max_num_logprobs: int | None
30
28 - # For penalties
29 - idx_mapping: torch.Tensor
30 - prompt_bin_mask: torch.Tensor
31 - output_bin_counts: torch.Tensor
32 -
31 @classmethod
32 def make_dummy(
33 cls,
@@ -37,6 +35,8 @@ class SamplingMetadata:
35 device: torch.device,
36 ) -> "SamplingMetadata":
37 assert num_reqs > 0
38 + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
39 +
40 temperature = torch.zeros(num_reqs, dtype=torch.float32, device=device)
41 temperature[0] = 0.5
42 # TODO(woosuk): Use top-p and top-k for dummy sampler.
@@ -51,18 +51,19 @@ class SamplingMetadata:
51 repetition_penalty = torch.ones(num_reqs, dtype=torch.float32, device=device)
52 frequency_penalty = torch.zeros(num_reqs, dtype=torch.float32, device=device)
53 presence_penalty = torch.zeros(num_reqs, dtype=torch.float32, device=device)
54 - seeds = torch.zeros(num_reqs, dtype=torch.int64, device=device)
55 - pos = torch.zeros(num_reqs, dtype=torch.int64, device=device)
56 - max_num_logprobs = 20
54
58 - idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
55 # NOTE(woosuk): These are placeholder tensors to avoid None checks in the
56 # penalties kernel. We use 2 instead of 1 as vocab_size to avoid Triton
57 # specialization and re-compilation at runtime.
58 prompt_bin_mask = torch.zeros(num_reqs, 2, dtype=torch.int32, device=device)
59 output_bin_counts = torch.zeros(num_reqs, 2, dtype=torch.int32, device=device)
60
61 + seeds = torch.zeros(num_reqs, dtype=torch.int64, device=device)
62 + pos = torch.zeros(num_reqs, dtype=torch.int64, device=device)
63 + max_num_logprobs = 20
64 +
65 return cls(
66 + idx_mapping=idx_mapping,
67 temperature=temperature,
68 top_p=top_p,
69 top_k=top_k,
@@ -70,123 +71,9 @@ class SamplingMetadata:
71 repetition_penalty=repetition_penalty,
72 frequency_penalty=frequency_penalty,
73 presence_penalty=presence_penalty,
74 + prompt_bin_mask=prompt_bin_mask,
75 + output_bin_counts=output_bin_counts,
76 seeds=seeds,
77 pos=pos,
78 max_num_logprobs=max_num_logprobs,
76 - idx_mapping=idx_mapping,
77 - prompt_bin_mask=prompt_bin_mask,
78 - output_bin_counts=output_bin_counts,
79 )
80 -
81 -
82 -# NOTE(woosuk): Re-compilation can happen at runtime since top_p and top_k can be None.
83 -@triton.jit
84 -def _expand_sampling_metadata_kernel(
85 - temp_ptr,
86 - expanded_temp_ptr,
87 - top_p_ptr,
88 - expanded_top_p_ptr,
89 - top_k_ptr,
90 - expanded_top_k_ptr,
91 - min_p_ptr,
92 - expanded_min_p_ptr,
93 - rep_penalty_ptr,
94 - expanded_rep_penalty_ptr,
95 - freq_penalty_ptr,
96 - expanded_freq_penalty_ptr,
97 - pres_penalty_ptr,
98 - expanded_pres_penalty_ptr,
99 - seeds_ptr,
100 - expanded_seeds_ptr,
101 - cu_num_logits_ptr,
102 - BLOCK_SIZE: tl.constexpr,
103 -):
104 - req_idx = tl.program_id(0)
105 - start_idx = tl.load(cu_num_logits_ptr + req_idx)
106 - end_idx = tl.load(cu_num_logits_ptr + req_idx + 1)
107 - num_tokens = end_idx - start_idx
108 -
109 - block = tl.arange(0, BLOCK_SIZE)
110 - mask = block < num_tokens
111 -
112 - temp = tl.load(temp_ptr + req_idx)
113 - tl.store(expanded_temp_ptr + start_idx + block, temp, mask=mask)
114 -
115 - if top_p_ptr is not None:
116 - top_p = tl.load(top_p_ptr + req_idx)
117 - tl.store(expanded_top_p_ptr + start_idx + block, top_p, mask=mask)
118 -
119 - if top_k_ptr is not None:
120 - top_k = tl.load(top_k_ptr + req_idx)
121 - tl.store(expanded_top_k_ptr + start_idx + block, top_k, mask=mask)
122 -
123 - if min_p_ptr is not None:
124 - min_p = tl.load(min_p_ptr + req_idx)
125 - tl.store(expanded_min_p_ptr + start_idx + block, min_p, mask=mask)
126 -
127 - rep_penalty = tl.load(rep_penalty_ptr + req_idx)
128 - tl.store(expanded_rep_penalty_ptr + start_idx + block, rep_penalty, mask=mask)
129 -
130 - freq_penalty = tl.load(freq_penalty_ptr + req_idx)
131 - tl.store(expanded_freq_penalty_ptr + start_idx + block, freq_penalty, mask=mask)
132 -
133 - pres_penalty = tl.load(pres_penalty_ptr + req_idx)
134 - tl.store(expanded_pres_penalty_ptr + start_idx + block, pres_penalty, mask=mask)
135 -
136 - seed = tl.load(seeds_ptr + req_idx)
137 - tl.store(expanded_seeds_ptr + start_idx + block, seed, mask=mask)
138 -
139 -
140 -def expand_sampling_metadata(
141 - sampling_metadata: SamplingMetadata,
142 - cu_num_logits: torch.Tensor,
143 - max_expand_len: int,
144 -) -> SamplingMetadata:
145 - total_num_logits = sampling_metadata.pos.shape[0]
146 - create_empty = lambda x: x.new_empty(total_num_logits) if x is not None else None
147 - expanded_temp = create_empty(sampling_metadata.temperature)
148 - expanded_top_p = create_empty(sampling_metadata.top_p)
149 - expanded_top_k = create_empty(sampling_metadata.top_k)
150 - expanded_min_p = create_empty(sampling_metadata.min_p)
151 - expanded_repetition_penalty = create_empty(sampling_metadata.repetition_penalty)
152 - expanded_frequency_penalty = create_empty(sampling_metadata.frequency_penalty)
153 - expanded_presence_penalty = create_empty(sampling_metadata.presence_penalty)
154 - expanded_seeds = create_empty(sampling_metadata.seeds)
155 -
156 - num_reqs = cu_num_logits.shape[0] - 1
157 - _expand_sampling_metadata_kernel[(num_reqs,)](
158 - sampling_metadata.temperature,
159 - expanded_temp,
160 - sampling_metadata.top_p,
161 - expanded_top_p,
162 - sampling_metadata.top_k,
163 - expanded_top_k,
164 - sampling_metadata.min_p,
165 - expanded_min_p,
166 - sampling_metadata.repetition_penalty,
167 - expanded_repetition_penalty,
168 - sampling_metadata.frequency_penalty,
169 - expanded_frequency_penalty,
170 - sampling_metadata.presence_penalty,
171 - expanded_presence_penalty,
172 - sampling_metadata.seeds,
173 - expanded_seeds,
174 - cu_num_logits,
175 - BLOCK_SIZE=triton.next_power_of_2(max_expand_len),
176 - )
177 - return SamplingMetadata(
178 - temperature=expanded_temp,
179 - top_p=expanded_top_p,
180 - top_k=expanded_top_k,
181 - min_p=expanded_min_p,
182 - seeds=expanded_seeds,
183 - repetition_penalty=expanded_repetition_penalty,
184 - frequency_penalty=expanded_frequency_penalty,
185 - presence_penalty=expanded_presence_penalty,
186 - pos=sampling_metadata.pos,
187 - max_num_logprobs=sampling_metadata.max_num_logprobs,
188 - # TODO(woosuk): Support penalties with spec decoding.
189 - idx_mapping=sampling_metadata.idx_mapping,
190 - prompt_bin_mask=sampling_metadata.prompt_bin_mask,
191 - output_bin_counts=sampling_metadata.output_bin_counts,
192 - )
vllm/v1/worker/gpu/sample/min_p.py
+9 -2
@@ -9,12 +9,14 @@ from vllm.triton_utils import tl, triton
9 def _min_p_kernel(
10 logits_ptr,
11 logits_stride,
12 + idx_mapping_ptr,
13 min_p_ptr,
14 vocab_size,
15 BLOCK_SIZE: tl.constexpr,
16 ):
17 req_idx = tl.program_id(0)
17 - min_p = tl.load(min_p_ptr + req_idx).to(tl.float32)
18 + req_state_idx = tl.load(idx_mapping_ptr + req_idx)
19 + min_p = tl.load(min_p_ptr + req_state_idx).to(tl.float32)
20 if min_p == 0.0:
21 return
22
@@ -39,12 +41,17 @@ def _min_p_kernel(
41 tl.store(logits_ptr + req_idx * logits_stride + block, logits, mask=mask)
42
43
42 -def apply_min_p(logits: torch.Tensor, min_p: torch.Tensor) -> None:
44 +def apply_min_p(
45 + logits: torch.Tensor,
46 + idx_mapping: torch.Tensor,
47 + min_p: torch.Tensor,
48 +) -> None:
49 num_reqs, vocab_size = logits.shape
50 BLOCK_SIZE = 1024
51 _min_p_kernel[(num_reqs,)](
52 logits,
53 logits.stride(0),
54 + idx_mapping,
55 min_p,
56 vocab_size,
57 BLOCK_SIZE=BLOCK_SIZE,
vllm/v1/worker/gpu/sample/penalties.py
+7 -7
@@ -10,11 +10,11 @@ from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
10 def _penalties_and_temperature_kernel(
11 logits_ptr,
12 logits_stride,
13 + idx_mapping_ptr,
14 repetition_penalty_ptr,
15 frequency_penalty_ptr,
16 presence_penalty_ptr,
17 temperature_ptr,
17 - idx_mapping_ptr,
18 prompt_bin_mask_ptr,
19 prompt_bin_mask_stride,
20 output_bin_counts_ptr,
@@ -23,10 +23,11 @@ def _penalties_and_temperature_kernel(
23 BLOCK_SIZE: tl.constexpr,
24 ):
25 batch_idx = tl.program_id(0)
26 - rep_penalty = tl.load(repetition_penalty_ptr + batch_idx)
27 - freq_penalty = tl.load(frequency_penalty_ptr + batch_idx)
28 - pres_penalty = tl.load(presence_penalty_ptr + batch_idx)
29 - temperature = tl.load(temperature_ptr + batch_idx)
26 + req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
27 + rep_penalty = tl.load(repetition_penalty_ptr + req_state_idx)
28 + freq_penalty = tl.load(frequency_penalty_ptr + req_state_idx)
29 + pres_penalty = tl.load(presence_penalty_ptr + req_state_idx)
30 + temperature = tl.load(temperature_ptr + req_state_idx)
31 temperature = tl.where(temperature == 0.0, 1.0, temperature)
32
33 use_rep_penalty = rep_penalty != 1.0
@@ -45,7 +46,6 @@ def _penalties_and_temperature_kernel(
46 logits = logits.to(tl.float32)
47
48 if use_penalty:
48 - req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
49 output_bin_counts = tl.load(
50 output_bin_counts_ptr + req_state_idx * output_bin_counts_stride + block,
51 mask=mask,
@@ -92,11 +92,11 @@ def apply_penalties_and_temperature(
92 _penalties_and_temperature_kernel[(num_reqs, num_blocks)](
93 logits,
94 logits.stride(0),
95 + sampling_metadata.idx_mapping,
96 sampling_metadata.repetition_penalty,
97 sampling_metadata.frequency_penalty,
98 sampling_metadata.presence_penalty,
99 sampling_metadata.temperature,
99 - sampling_metadata.idx_mapping,
100 sampling_metadata.prompt_bin_mask,
101 sampling_metadata.prompt_bin_mask.stride(0),
102 sampling_metadata.output_bin_counts,
vllm/v1/worker/gpu/sample/sampler.py
+2 -1
@@ -71,7 +71,7 @@ class Sampler:
71 apply_penalties_and_temperature(logits, sampling_metadata)
72 # Apply min_p in place.
73 if sampling_metadata.min_p is not None:
74 - apply_min_p(logits, sampling_metadata.min_p)
74 + apply_min_p(logits, sampling_metadata.idx_mapping, sampling_metadata.min_p)
75 # Apply top_k and/or top_p. This might return a new tensor.
76 logits = apply_top_k_top_p(
77 logits, sampling_metadata.top_k, sampling_metadata.top_p
@@ -79,6 +79,7 @@ class Sampler:
79
80 sampled = gumbel_sample(
81 logits,
82 + sampling_metadata.idx_mapping,
83 sampling_metadata.temperature,
84 sampling_metadata.seeds,
85 sampling_metadata.pos,
vllm/v1/worker/gpu/spec_decode/eagle.py
+27 -23
@@ -2,7 +2,6 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from typing import Any
4
5 -import numpy as np
5 import torch
6 import torch.nn as nn
7
@@ -12,7 +11,6 @@ from vllm.forward_context import set_forward_context
11 from vllm.logger import init_logger
12 from vllm.model_executor.model_loader import get_model
13 from vllm.triton_utils import tl, triton
15 -from vllm.utils.platform_utils import is_pin_memory_available
14 from vllm.v1.attention.backends.utils import AttentionMetadataBuilder
15 from vllm.v1.kv_cache_interface import KVCacheConfig
16 from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
@@ -46,7 +44,6 @@ class EagleSpeculator:
44 self.hidden_size = self.draft_model_config.get_hidden_size()
45 self.inputs_embeds_size = self.draft_model_config.get_inputs_embeds_size()
46 self.vocab_size = self.draft_model_config.get_vocab_size()
49 - self.pin_memory = is_pin_memory_available()
47 self.dtype = vllm_config.model_config.dtype
48
49 self.input_buffers = InputBuffers(
@@ -56,7 +53,6 @@ class EagleSpeculator:
53 vocab_size=self.vocab_size,
54 dtype=self.dtype,
55 device=device,
59 - pin_memory=self.pin_memory,
56 )
57 self.hidden_states = torch.zeros(
58 self.max_num_tokens,
@@ -64,6 +60,11 @@ class EagleSpeculator:
60 dtype=self.dtype,
61 device=device,
62 )
63 + self.idx_mapping = torch.zeros(
64 + self.max_num_reqs,
65 + dtype=torch.int32,
66 + device=device,
67 + )
68 self.temperature = torch.zeros(
69 self.max_num_reqs,
70 dtype=torch.float32,
@@ -140,7 +141,7 @@ class EagleSpeculator:
141 num_tokens_across_dp: torch.Tensor | None,
142 ) -> None:
143 pos = self.input_buffers.positions[:num_reqs]
143 - query_start_loc = self.input_buffers.query_start_loc.gpu[: num_reqs + 1]
144 + query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
145 for step in range(1, self.num_speculative_steps):
146 # Run the eagle model.
147 last_hidden_states, hidden_states = self.run_model(
@@ -152,8 +153,9 @@ class EagleSpeculator:
153 # used for draft and target sampling.
154 draft_tokens = gumbel_sample(
155 logits,
155 - self.temperature[:num_reqs],
156 - self.seeds[:num_reqs],
156 + self.idx_mapping[:num_reqs],
157 + self.temperature,
158 + self.seeds,
159 pos + 1,
160 apply_temperature=True,
161 )
@@ -237,23 +239,27 @@ class EagleSpeculator:
239 logits = self.model.compute_logits(sample_hidden_states)
240
241 num_reqs = input_batch.num_reqs
240 - cu_num_logits = input_batch.cu_num_logits[:num_reqs]
242 # NOTE(woosuk): For draft sampling, we only consider the temperature
243 # and ignore the other sampling parameters such as top_k and top_p,
244 # for simplicity and performance.
245 # While this may slightly degrade the acceptance rate, it does not
246 # affect the output distribution after rejection sampling.
246 - temperature = self.temperature[:num_reqs]
247 - seeds = self.seeds[:num_reqs]
248 - pos = self.input_buffers.positions[:num_reqs]
247 + idx_mapping = self.idx_mapping[:num_reqs]
248 + idx_mapping.copy_(input_batch.idx_mapping)
249 + self.temperature.copy_(sampling_metadata.temperature)
250 + self.seeds.copy_(sampling_metadata.seeds)
251 # Gather the values and copy them to the pre-allocated buffers.
250 - torch.gather(sampling_metadata.temperature, 0, cu_num_logits, out=temperature)
251 - torch.gather(sampling_metadata.seeds, 0, cu_num_logits, out=seeds)
252 + pos = self.input_buffers.positions[:num_reqs]
253 torch.gather(input_batch.positions, 0, last_token_indices, out=pos)
254 # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
255 # used for draft and target sampling.
256 draft_tokens = gumbel_sample(
256 - logits, temperature, seeds, pos + 1, apply_temperature=True
257 + logits,
258 + idx_mapping,
259 + self.temperature,
260 + self.seeds,
261 + pos + 1,
262 + apply_temperature=True,
263 )
264 if self.num_speculative_steps == 1:
265 # Early exit.
@@ -273,11 +279,8 @@ class EagleSpeculator:
279 self.max_model_len,
280 self.max_num_reqs,
281 )
276 - query_start_loc = self.input_buffers.query_start_loc
277 - query_start_loc_gpu = query_start_loc.gpu[: num_reqs + 1]
278 - slot_mappings = self.block_tables.compute_slot_mappings(
279 - query_start_loc_gpu, pos
280 - )
282 + query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
283 + slot_mappings = self.block_tables.compute_slot_mappings(query_start_loc, pos)
284
285 cudagraph_size = self.cudagraph_manager.get_cudagraph_size(num_reqs)
286 if cudagraph_size is not None:
@@ -286,8 +289,9 @@ class EagleSpeculator:
289 return self.draft_tokens[:num_reqs]
290
291 # Run eager mode.
289 - query_start_loc.np[: num_reqs + 1] = np.arange(num_reqs + 1)
290 - query_start_loc_cpu = query_start_loc.cpu[: num_reqs + 1]
292 + query_start_loc_cpu = torch.arange(
293 + num_reqs + 1, dtype=torch.int32, device="cpu"
294 + )
295 block_tables = [x[:num_reqs] for x in self.block_tables.input_block_tables]
296
297 # FIXME(woosuk): This is UNSAFE!!
@@ -295,7 +299,7 @@ class EagleSpeculator:
299 attn_metadata_builders=self.attn_metadata_builders,
300 num_reqs=num_reqs,
301 num_tokens=num_reqs,
298 - query_start_loc_gpu=query_start_loc_gpu,
302 + query_start_loc_gpu=query_start_loc,
303 query_start_loc_cpu=query_start_loc_cpu,
304 seq_lens=self.input_buffers.seq_lens[:num_reqs],
305 max_seq_len=self.max_model_len,
@@ -484,7 +488,7 @@ def prepare_eagle_decode(
488 input_buffers.positions,
489 input_hidden_states,
490 input_hidden_states.stride(0),
487 - input_buffers.query_start_loc.gpu,
491 + input_buffers.query_start_loc,
492 input_buffers.seq_lens,
493 hidden_size,
494 max_model_len,
vllm/v1/worker/gpu/states.py
+51 -85
@@ -8,10 +8,8 @@ import torch
8 from vllm.lora.request import LoRARequest
9 from vllm.sampling_params import SamplingParams
10 from vllm.utils.math_utils import cdiv
11 -from vllm.utils.platform_utils import is_uva_available
12 -from vllm.utils.torch_utils import get_cuda_view_from_cpu_tensor
11 from vllm.v1.outputs import LogprobsTensors
14 -from vllm.v1.utils import CpuGpuBuffer
12 +from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
13 from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
14 from vllm.v1.worker.gpu.sample.penalties import bincount
15
@@ -29,7 +27,6 @@ class RequestState:
27 num_speculative_steps: int,
28 vocab_size: int,
29 device: torch.device,
32 - pin_memory: bool,
30 ):
31 self.max_num_reqs = max_num_reqs
32 self.max_model_len = max_model_len
@@ -37,7 +34,6 @@ class RequestState:
34 self.num_speculative_steps = num_speculative_steps
35 self.vocab_size = vocab_size
36 self.device = device
40 - self.pin_memory = pin_memory
37
38 self.req_id_to_index: dict[str, int] = {}
39 self.index_to_req_id: dict[int, str] = {}
@@ -47,16 +43,18 @@ class RequestState:
43 self.prompt_len = np.zeros(self.max_num_reqs, dtype=np.int32)
44 # NOTE(woosuk): This tensor can be extremely large (e.g., several GBs)
45 # depending on the configured max_num_reqs and max_model_len.
50 - self.prefill_token_ids = UvaBuffer(
51 - self.max_num_reqs, self.max_model_len, dtype=torch.int32
46 + # To save GPU memory, we use UVA instead of GPU for this tensor.
47 + self.prefill_token_ids = StagedWriteTensor(
48 + (self.max_num_reqs, self.max_model_len),
49 + dtype=torch.int32,
50 + device=device,
51 + uva_instead_of_gpu=True,
52 )
53 - # NOTE(woosuk): We don't use UVA for prefill_len because its GPU view
54 - # can be used outside of update_states and prepare_inputs.
55 - # Without async barrier, using UVA can cause race conditions.
56 - self.prefill_len = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
53 + self.prefill_len = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
54 +
55 # Number of computed tokens.
56 self.num_computed_prefill_tokens = np.zeros(self.max_num_reqs, dtype=np.int32)
59 - self.num_computed_tokens = torch.zeros(
57 + self.num_computed_tokens = StagedWriteTensor(
58 self.max_num_reqs, dtype=torch.int32, device=device
59 )
60
@@ -84,14 +82,16 @@ class RequestState:
82 self.lora_ids.fill(NO_LORA_ID)
83
84 # Sampling parameters.
87 - self.temperature = self._make_param(self.max_num_reqs, torch.float32)
88 - self.top_p = self._make_param(self.max_num_reqs, torch.float32)
89 - self.top_k = self._make_param(self.max_num_reqs, torch.int32)
90 - self.min_p = self._make_param(self.max_num_reqs, torch.float32)
91 - self.repetition_penalty = self._make_param(self.max_num_reqs, torch.float32)
92 - self.frequency_penalty = self._make_param(self.max_num_reqs, torch.float32)
93 - self.presence_penalty = self._make_param(self.max_num_reqs, torch.float32)
94 - self.seeds = self._make_param(self.max_num_reqs, torch.int64)
85 + self.temperature = UvaBackedTensor(self.max_num_reqs, dtype=torch.float32)
86 + self.top_p = UvaBackedTensor(self.max_num_reqs, dtype=torch.float32)
87 + self.top_k = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
88 + self.min_p = UvaBackedTensor(self.max_num_reqs, dtype=torch.float32)
89 + self.repetition_penalty = UvaBackedTensor(
90 + self.max_num_reqs, dtype=torch.float32
91 + )
92 + self.frequency_penalty = UvaBackedTensor(self.max_num_reqs, dtype=torch.float32)
93 + self.presence_penalty = UvaBackedTensor(self.max_num_reqs, dtype=torch.float32)
94 + self.seeds = UvaBackedTensor(self.max_num_reqs, dtype=torch.int64)
95
96 self.num_logprobs = np.empty(self.max_num_reqs, dtype=np.int32)
97 # -1 means no logprobs are requested.
@@ -111,13 +111,7 @@ class RequestState:
111 self.max_num_reqs, self.vocab_size, dtype=torch.int32, device=self.device
112 )
113
114 - def _make_param(self, size: int, dtype: torch.dtype) -> "Param":
115 - return Param(size, dtype=dtype, device=self.device, pin_memory=self.pin_memory)
116 -
117 - def _make_buffer(self, size: int, dtype: torch.dtype) -> CpuGpuBuffer:
118 - return CpuGpuBuffer(
119 - size, dtype=dtype, device=self.device, pin_memory=self.pin_memory
120 - )
114 + self._penalties_reqs: list[int] = []
115
116 @property
117 def num_reqs(self) -> int:
@@ -144,12 +138,9 @@ class RequestState:
138 f"prefill_len {prefill_len} < prompt_len {prompt_len}"
139 )
140 self.prefill_len.np[req_idx] = prefill_len
147 - self.prefill_token_ids.np[req_idx, :prefill_len] = prefill_token_ids
148 -
141 + self.prefill_token_ids.stage_write(req_idx, 0, prefill_token_ids)
142 self.num_computed_prefill_tokens[req_idx] = num_computed_tokens
150 - # FIXME(woosuk): This triggers a GPU operation whenever adding a new request.
151 - # Optimize this.
152 - self.num_computed_tokens[req_idx] = num_computed_tokens
143 + self.num_computed_tokens.stage_write_elem(req_idx, num_computed_tokens)
144
145 if lora_request is not None:
146 self.lora_ids[req_idx] = lora_request.lora_int_id
@@ -169,13 +160,7 @@ class RequestState:
160 self.presence_penalty.np[req_idx] = sampling_params.presence_penalty
161
162 if use_penalty(sampling_params):
172 - bincount(
173 - self.prefill_token_ids.gpu[req_idx],
174 - prefill_len,
175 - prompt_len,
176 - self.prompt_bin_mask[req_idx],
177 - self.output_bin_counts[req_idx],
178 - )
163 + self._penalties_reqs.append(req_idx)
164
165 if sampling_params.seed is not None:
166 seed = sampling_params.seed
@@ -193,6 +178,22 @@ class RequestState:
178 needs_prompt_logprobs = sampling_params.prompt_logprobs is not None
179 self.needs_prompt_logprobs[req_idx] = needs_prompt_logprobs
180
181 + def apply_staged_writes(self) -> None:
182 + self.prefill_len.copy_to_uva()
183 + self.prefill_token_ids.apply_write()
184 + self.num_computed_tokens.apply_write()
185 +
186 + # TODO(woosuk): Optimize this.
187 + for req_idx in self._penalties_reqs:
188 + bincount(
189 + self.prefill_token_ids.gpu[req_idx],
190 + int(self.prefill_len.np[req_idx]),
191 + int(self.prompt_len[req_idx]),
192 + self.prompt_bin_mask[req_idx],
193 + self.output_bin_counts[req_idx],
194 + )
195 + self._penalties_reqs.clear()
196 +
197 def remove_request(self, req_id: str) -> None:
198 self.extra_data.pop(req_id, None)
199 req_idx = self.req_id_to_index.pop(req_id, None)
@@ -208,30 +209,25 @@ class RequestState:
209 idx_mapping_np: np.ndarray,
210 pos: torch.Tensor,
211 ) -> SamplingMetadata:
211 - temperature = self.temperature.np[idx_mapping_np]
212 - temperature = self.temperature.copy_np_to_gpu(temperature)
212 + temperature = self.temperature.copy_to_uva()
213
214 top_p = self.top_p.np[idx_mapping_np]
215 no_top_p = np.all(top_p == 1.0)
216 - top_p = self.top_p.copy_np_to_gpu(top_p) if not no_top_p else None
216 + top_p = self.top_p.copy_to_uva()[idx_mapping] if not no_top_p else None
217
218 top_k = self.top_k.np[idx_mapping_np]
219 no_top_k = np.all(top_k == self.vocab_size)
220 - top_k = self.top_k.copy_np_to_gpu(top_k) if not no_top_k else None
220 + top_k = self.top_k.copy_to_uva()[idx_mapping] if not no_top_k else None
221
222 min_p = self.min_p.np[idx_mapping_np]
223 no_min_p = np.all(min_p == 0.0)
224 - min_p = self.min_p.copy_np_to_gpu(min_p) if not no_min_p else None
224 + min_p = self.min_p.copy_to_uva() if not no_min_p else None
225
226 - rep_penalty = self.repetition_penalty.np[idx_mapping_np]
227 - rep_penalty = self.repetition_penalty.copy_np_to_gpu(rep_penalty)
228 - freq_penalty = self.frequency_penalty.np[idx_mapping_np]
229 - freq_penalty = self.frequency_penalty.copy_np_to_gpu(freq_penalty)
230 - pres_penalty = self.presence_penalty.np[idx_mapping_np]
231 - pres_penalty = self.presence_penalty.copy_np_to_gpu(pres_penalty)
226 + rep_penalty = self.repetition_penalty.copy_to_uva()
227 + freq_penalty = self.frequency_penalty.copy_to_uva()
228 + pres_penalty = self.presence_penalty.copy_to_uva()
229
233 - seeds = self.seeds.np[idx_mapping_np]
234 - seeds = self.seeds.copy_np_to_gpu(seeds)
230 + seeds = self.seeds.copy_to_uva()
231
232 num_logprobs = self.num_logprobs[idx_mapping_np]
233 max_num_logprobs: int | None = int(np.max(num_logprobs))
@@ -239,6 +235,7 @@ class RequestState:
235 max_num_logprobs = None
236
237 return SamplingMetadata(
238 + idx_mapping=idx_mapping,
239 temperature=temperature,
240 top_p=top_p,
241 top_k=top_k,
@@ -246,12 +243,11 @@ class RequestState:
243 repetition_penalty=rep_penalty,
244 frequency_penalty=freq_penalty,
245 presence_penalty=pres_penalty,
246 + prompt_bin_mask=self.prompt_bin_mask,
247 + output_bin_counts=self.output_bin_counts,
248 seeds=seeds,
249 pos=pos,
250 max_num_logprobs=max_num_logprobs,
252 - idx_mapping=idx_mapping,
253 - prompt_bin_mask=self.prompt_bin_mask,
254 - output_bin_counts=self.output_bin_counts,
251 )
252
253 def make_lora_inputs(
@@ -272,42 +268,12 @@ class RequestState:
268 return prompt_lora_mapping, token_lora_mapping, active_lora_requests
269
270
275 -class Param:
276 - def __init__(
277 - self,
278 - size: int,
279 - dtype: torch.dtype,
280 - device: torch.device,
281 - pin_memory: bool,
282 - ):
283 - self.buffer = CpuGpuBuffer(
284 - size,
285 - dtype=dtype,
286 - device=device,
287 - pin_memory=pin_memory,
288 - )
289 - self.np = np.zeros_like(self.buffer.np)
290 -
291 - def copy_np_to_gpu(self, x: np.ndarray) -> torch.Tensor:
292 - n = x.shape[0]
293 - self.buffer.np[:n] = x
294 - return self.buffer.copy_to_gpu(n)
295 -
296 -
271 @dataclass
272 class ExtraData:
273 lora_request: LoRARequest | None
274 in_progress_prompt_logprobs: list[LogprobsTensors] = field(default_factory=list)
275
276
303 -class UvaBuffer:
304 - def __init__(self, *size: int | torch.SymInt, dtype: torch.dtype):
305 - assert is_uva_available()
306 - self.cpu = torch.zeros(*size, dtype=dtype, device="cpu", pin_memory=True)
307 - self.np = self.cpu.numpy()
308 - self.gpu = get_cuda_view_from_cpu_tensor(self.cpu)
309 -
310 -
277 def use_penalty(sampling_params: SamplingParams) -> bool:
278 return (
279 sampling_params.repetition_penalty != 1.0
vllm/v1/worker/gpu/structured_outputs.py
+58 -34
@@ -4,38 +4,65 @@ import numpy as np
4 import torch
5
6 from vllm.triton_utils import tl, triton
7 -from vllm.v1.worker.gpu.input_batch import InputBuffers
7 +from vllm.utils.math_utils import cdiv
8 +from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool
9 +from vllm.v1.worker.gpu.input_batch import InputBatch
10
11
10 -def apply_grammar_bitmask(
11 - logits: torch.Tensor,
12 - req_ids: list[str],
13 - grammar_req_ids: list[str],
14 - grammar_bitmask: np.ndarray,
15 - input_buffers: InputBuffers,
16 -) -> None:
17 - input_buffers.grammar_bitmask.np[: grammar_bitmask.shape[0]] = grammar_bitmask
18 - input_buffers.grammar_bitmask.copy_to_gpu(grammar_bitmask.shape[0])
12 +class StructuredOutputsWorker:
13 + def __init__(
14 + self,
15 + max_num_logits: int,
16 + vocab_size: int,
17 + ):
18 + # NOTE(woosuk): Here, we use UvaBufferPool instead of UvaBackedTensor
19 + # to save a unnecessary CPU-to-CPU copy.
20 + self.logits_indices = UvaBufferPool(max_num_logits, torch.int32)
21 + self.grammar_bitmask = UvaBufferPool(
22 + (max_num_logits, cdiv(vocab_size, 32)), torch.int32
23 + )
24
20 - batch_size = logits.shape[0]
21 - grammar_req_id_to_idx = {req_id: i for i, req_id in enumerate(grammar_req_ids)}
22 - # logits -> bitmask mapping
23 - mapping = [grammar_req_id_to_idx.get(req_id, -1) for req_id in req_ids]
24 - input_buffers.bitmask_indices.np[:batch_size] = mapping
25 - input_buffers.bitmask_indices.copy_to_gpu(batch_size)
25 + def apply_grammar_bitmask(
26 + self,
27 + logits: torch.Tensor,
28 + input_batch: InputBatch,
29 + grammar_req_ids: list[str],
30 + grammar_bitmask: np.ndarray,
31 + ) -> None:
32 + if not grammar_req_ids:
33 + return
34
27 - vocab_size = logits.shape[-1]
28 - BLOCK_SIZE = 8192
29 - grid = (batch_size, triton.cdiv(vocab_size, BLOCK_SIZE))
30 - _apply_grammar_bitmask_kernel[grid](
31 - logits,
32 - logits.stride(0),
33 - input_buffers.grammar_bitmask.gpu,
34 - input_buffers.grammar_bitmask.gpu.stride(0),
35 - input_buffers.bitmask_indices.gpu,
36 - vocab_size,
37 - BLOCK_SIZE=BLOCK_SIZE,
38 - )
35 + # Construct bitmask -> logits mapping
36 + mapping: list[int] = []
37 + req_ids = input_batch.req_ids
38 + cu_num_logits = input_batch.cu_num_logits_np.tolist()
39 + req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)}
40 + for grammar_req_id in grammar_req_ids:
41 + req_idx = req_id_to_idx[grammar_req_id]
42 + logits_start_idx = cu_num_logits[req_idx]
43 + logits_end_idx = cu_num_logits[req_idx + 1]
44 + mapping.extend(range(logits_start_idx, logits_end_idx))
45 + # Copy the mapping.
46 + mapping_np = np.array(mapping, dtype=np.int32)
47 + logits_indices = self.logits_indices.copy_to_uva(mapping_np)
48 +
49 + # Copy the bitmask.
50 + bitmask = self.grammar_bitmask.copy_to_uva(grammar_bitmask)
51 +
52 + num_masks = bitmask.shape[0]
53 + assert num_masks == len(mapping)
54 + vocab_size = logits.shape[-1]
55 + BLOCK_SIZE = 8192
56 + grid = (num_masks, triton.cdiv(vocab_size, BLOCK_SIZE))
57 + _apply_grammar_bitmask_kernel[grid](
58 + logits,
59 + logits.stride(0),
60 + logits_indices,
61 + bitmask,
62 + bitmask.stride(0),
63 + vocab_size,
64 + BLOCK_SIZE=BLOCK_SIZE,
65 + )
66
67
68 # Adapted from
@@ -44,17 +71,14 @@ def apply_grammar_bitmask(
71 def _apply_grammar_bitmask_kernel(
72 logits_ptr,
73 logits_stride,
74 + logits_indices_ptr,
75 bitmask_ptr,
76 bitmask_stride,
49 - bitmask_indices_ptr,
77 vocab_size,
78 BLOCK_SIZE: tl.constexpr,
79 ):
53 - logits_idx = tl.program_id(0)
54 - bitmask_idx = tl.load(bitmask_indices_ptr + logits_idx)
55 - if bitmask_idx == -1:
56 - # No bitmask to apply.
57 - return
80 + bitmask_idx = tl.program_id(0)
81 + logits_idx = tl.load(logits_indices_ptr + bitmask_idx)
82
83 # Load the bitmask.
84 block_id = tl.program_id(1)