@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0ff902f3b4

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Refactor] Refactor persistent buffers with CpuGpuBuffer (#23515)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Woosuk Kwon committed Aug 25, 2025 at 08:44 UTC 0ff902f3b48376aadfe4e753ce0872b49c8f9f7e

2 files changed +99 -103

vllm/v1/worker/gpu_model_runner.py

+70 -103

@@ -83,8 +83,9 @@ from vllm.v1.worker.kv_connector_model_runner_mixin import (

83 KVConnectorModelRunnerMixin, KVConnectorOutput)

84 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

86 -from .utils import (AttentionGroup, MultiModalBudget, bind_kv_cache,

87 - gather_mm_placeholders, initialize_kv_cache_for_kv_sharing,

86 +from .utils import (AttentionGroup, CpuGpuBuffer, MultiModalBudget,

87 + bind_kv_cache, gather_mm_placeholders,

88 + initialize_kv_cache_for_kv_sharing,

89 sanity_check_mm_encoder_outputs, scatter_mm_placeholders)

91 if TYPE_CHECKING:

@@ -149,6 +150,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

150 parallel_config)

151 self.hidden_size = model_config.get_hidden_size()

152 self.attention_chunk_size = model_config.attention_chunk_size

153 + # Only relevant for models using ALiBi (e.g, MPT)

154 + self.use_alibi = check_use_alibi(model_config)

155

156 self.cascade_attn_enabled = not self.model_config.disable_cascade_attn

157

@@ -242,21 +245,17 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

245 self._init_device_properties()

246

247 # Persistent buffers for CUDA graphs.

245 - self.input_ids = torch.zeros(self.max_num_tokens,

246 - dtype=torch.int32,

247 - device=self.device)

248 - self.positions = torch.zeros(self.max_num_tokens,

249 - dtype=torch.int64,

250 - device=self.device)

251 - self.query_start_loc = torch.zeros(self.max_num_reqs + 1,

252 - dtype=torch.int32,

253 - device=self.device)

254 - self.seq_lens = torch.zeros(self.max_num_reqs,

255 - dtype=torch.int32,

256 - device=self.device)

257 -

258 - # None in the first PP rank. The rest are set after load_model.

259 - self.intermediate_tensors: Optional[IntermediateTensors] = None

248 + self.input_ids = self._make_buffer(self.max_num_tokens,

249 + dtype=torch.int32)

250 + self.positions = self._make_buffer(self.max_num_tokens,

251 + dtype=torch.int64)

252 + self.query_start_loc = self._make_buffer(self.max_num_reqs + 1,

253 + dtype=torch.int32)

254 + self.seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)

255 + self.inputs_embeds = torch.zeros(

256 + (self.max_num_tokens, self.hidden_size),

257 + dtype=self.dtype,

258 + device=self.device)

259

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

261 if self.uses_mrope:

@@ -270,23 +269,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

269 # identical position IDs, making M-RoPE functionally equivalent to

270 # 1D-RoPE.

271 # See page 5 of https://arxiv.org/abs/2409.12191

273 - self.mrope_positions = torch.zeros((3, self.max_num_tokens + 1),

274 - dtype=torch.int64,

275 - device=self.device)

276 - self.mrope_positions_cpu = torch.zeros(

277 - (3, self.max_num_tokens + 1),

278 - dtype=torch.int64,

279 - device="cpu",

280 - pin_memory=self.pin_memory)

281 - self.mrope_positions_np = self.mrope_positions_cpu.numpy()

282 -

283 - # Only relevant for models using ALiBi (e.g, MPT)

284 - self.use_alibi = check_use_alibi(model_config)

272 + self.mrope_positions = self._make_buffer(

273 + (3, self.max_num_tokens + 1), dtype=torch.int64)

274

286 - self.inputs_embeds = torch.zeros(

287 - (self.max_num_tokens, self.hidden_size),

288 - dtype=self.dtype,

289 - device=self.device)

275 + # None in the first PP rank. The rest are set after load_model.

276 + self.intermediate_tensors: Optional[IntermediateTensors] = None

277

278 # OPTIMIZATION: Cache the tensors rather than creating them every step.

279 # Keep in int64 to avoid overflow with long context

@@ -294,28 +281,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

281 self.max_model_len,

282 self.max_num_tokens),

283 dtype=np.int64)

297 - # NOTE(woosuk): These tensors are "stateless", i.e., they are literally

298 - # a faster version of creating a new tensor every time. Thus, we should

299 - # not make any assumptions about the values in these tensors.

300 - self.input_ids_cpu = torch.zeros(self.max_num_tokens,

301 - dtype=torch.int32,

302 - device="cpu",

303 - pin_memory=self.pin_memory)

304 - self.positions_cpu = torch.zeros(self.max_num_tokens,

305 - dtype=torch.int64,

306 - device="cpu",

307 - pin_memory=self.pin_memory)

308 - self.positions_np = self.positions_cpu.numpy()

309 - self.query_start_loc_cpu = torch.zeros(self.max_num_reqs + 1,

310 - dtype=torch.int32,

311 - device="cpu",

312 - pin_memory=self.pin_memory)

313 - self.query_start_loc_np = self.query_start_loc_cpu.numpy()

314 - self.seq_lens_cpu = torch.zeros(self.max_num_reqs,

315 - dtype=torch.int32,

316 - device="cpu",

317 - pin_memory=self.pin_memory)

318 - self.seq_lens_np = self.seq_lens_cpu.numpy()

284

285 # Layer pairings for cross-layer KV sharing.

286 # If an Attention layer `layer_name` is in the keys of this dict, it

@@ -352,6 +317,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

317 self._draft_token_ids: Optional[Union[list[list[int]],

318 torch.Tensor]] = None

319

320 + def _make_buffer(self, *args, dtype: torch.dtype) -> CpuGpuBuffer:

321 + return CpuGpuBuffer(*args,

322 + dtype=dtype,

323 + device=self.device,

324 + pin_memory=self.pin_memory)

325 +

326 def _init_model_kwargs(self, num_tokens: int):

327 model_kwargs = dict[str, Any]()

328 num_reqs = self.input_batch.num_reqs

@@ -376,7 +347,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

347 if len(token_type_id_requests) == 0:

348 return model_kwargs

349

379 - seq_lens = self.seq_lens[:num_reqs]

350 + seq_lens = self.seq_lens.gpu[:num_reqs]

351 token_type_ids = []

352

353 for i in range(num_reqs):

@@ -719,7 +690,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

690 num_scheduled_tokens)

691

692 # Get positions.

722 - positions_np = self.positions_np[:total_num_scheduled_tokens]

693 + positions_np = self.positions.np[:total_num_scheduled_tokens]

694 np.add(self.input_batch.num_computed_tokens_cpu[req_indices],

695 arange,

696 out=positions_np)

@@ -742,7 +713,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

713 torch.index_select(self.input_batch.token_ids_cpu_tensor.flatten(),

714 0,

715 torch.from_numpy(token_indices),

745 - out=self.input_ids_cpu[:total_num_scheduled_tokens])

716 + out=self.input_ids.cpu[:total_num_scheduled_tokens])

717

718 self.input_batch.block_table.compute_slot_mapping(

719 req_indices, positions_np)

@@ -750,36 +721,33 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

721 total_num_scheduled_tokens)

722

723 # Prepare the attention metadata.

753 - self.query_start_loc_np[0] = 0

754 - self.query_start_loc_np[1:num_reqs + 1] = cu_num_tokens

724 + self.query_start_loc.np[0] = 0

725 + self.query_start_loc.np[1:num_reqs + 1] = cu_num_tokens

726 # Note: pad query_start_loc to be non-decreasing, as kernels

727 # like FlashAttention requires that

757 - self.query_start_loc_np[num_reqs + 1:].fill(cu_num_tokens[-1])

758 - self.query_start_loc.copy_(self.query_start_loc_cpu, non_blocking=True)

759 - query_start_loc = self.query_start_loc[:num_reqs + 1]

728 + self.query_start_loc.np[num_reqs + 1:].fill(cu_num_tokens[-1])

729 + self.query_start_loc.copy_to_gpu()

730 + query_start_loc = self.query_start_loc.gpu[:num_reqs + 1]

731

761 - self.seq_lens_np[:num_reqs] = (

732 + self.seq_lens.np[:num_reqs] = (

733 self.input_batch.num_computed_tokens_cpu[:num_reqs] +

734 num_scheduled_tokens)

735 # Fill unused with 0 for full cuda graph mode.

765 - self.seq_lens_np[num_reqs:].fill(0)

766 - self.seq_lens.copy_(self.seq_lens_cpu, non_blocking=True)

767 - seq_lens = self.seq_lens[:num_reqs]

768 - max_seq_len = self.seq_lens_np[:num_reqs].max().item()

736 + self.seq_lens.np[num_reqs:].fill(0)

737 + self.seq_lens.copy_to_gpu()

738 + seq_lens = self.seq_lens.gpu[:num_reqs]

739 + max_seq_len = self.seq_lens.np[:num_reqs].max().item()

740

741 # Copy the tensors to the GPU.

771 - self.input_ids[:total_num_scheduled_tokens].copy_(

772 - self.input_ids_cpu[:total_num_scheduled_tokens], non_blocking=True)

742 + self.input_ids.copy_to_gpu(total_num_scheduled_tokens)

743 if self.uses_mrope:

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

775 - self.mrope_positions[:, :total_num_scheduled_tokens].copy_(

776 - self.mrope_positions_cpu[:, :total_num_scheduled_tokens],

745 + self.mrope_positions.gpu[:, :total_num_scheduled_tokens].copy_(

746 + self.mrope_positions.cpu[:, :total_num_scheduled_tokens],

747 non_blocking=True)

748 else:

749 # Common case (1D positions)

780 - self.positions[:total_num_scheduled_tokens].copy_(

781 - self.positions_cpu[:total_num_scheduled_tokens],

782 - non_blocking=True)

750 + self.positions.copy_to_gpu(total_num_scheduled_tokens)

751

752 use_spec_decode = len(

753 scheduler_output.scheduled_spec_decode_tokens) > 0

@@ -833,8 +801,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

801 attn_metadata: dict[str, Any] = {}

802

803 # Used in the below loop.

836 - query_start_loc_cpu = self.query_start_loc_cpu[:num_reqs + 1]

837 - seq_lens_cpu = self.seq_lens_cpu[:num_reqs]

804 + query_start_loc_cpu = self.query_start_loc.cpu[:num_reqs + 1]

805 + seq_lens_cpu = self.seq_lens.cpu[:num_reqs]

806 num_computed_tokens_cpu = (

807 self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs])

808 spec_decode_common_attn_metadata = None

@@ -1065,9 +1033,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1033 src_start = num_computed_tokens

1034 src_end = num_computed_tokens + prompt_part_len

1035

1068 - self.mrope_positions_cpu[:, dst_start:dst_end] = \

1069 - req.mrope_positions[:,src_start:src_end]

1070 -

1036 + self.mrope_positions.cpu[:, dst_start:dst_end] = (

1037 + req.mrope_positions[:, src_start:src_end])

1038 mrope_pos_ptr += prompt_part_len

1039

1040 if completion_part_len > 0:

@@ -1076,7 +1043,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1043 dst_end = mrope_pos_ptr + completion_part_len

1044

1045 MRotaryEmbedding.get_next_input_positions_tensor(

1079 - out=self.mrope_positions_np,

1046 + out=self.mrope_positions.np,

1047 out_offset=dst_start,

1048 mrope_position_delta=req.mrope_position_delta,

1049 context_len=num_computed_tokens + prompt_part_len,

@@ -1140,7 +1107,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1107

1108 # Compute the draft token ids.

1109 # draft_token_indices: [ 1, 2, 3, 105, 106, 208]

1143 - draft_token_ids = self.input_ids[logits_indices]

1110 + draft_token_ids = self.input_ids.gpu[logits_indices]

1111 draft_token_ids = draft_token_ids[target_logits_indices + 1]

1112

1113 metadata = SpecDecodeMetadata(

@@ -1471,7 +1438,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1438 pooling_metadata = self.input_batch.pooling_metadata

1439 pooling_metadata.build_pooling_cursor(num_scheduled_tokens_np.tolist(),

1440 device=hidden_states.device)

1474 - seq_lens_cpu = self.seq_lens_cpu[:self.input_batch.num_reqs]

1441 + seq_lens_cpu = self.seq_lens.cpu[:self.input_batch.num_reqs]

1442

1443 # Pooling models D2H & synchronize occurs in pooler.py:build_output

1444 raw_pooler_output = self.model.pooler(

@@ -1550,7 +1517,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1517 # embeddings), we always use embeddings (rather than token ids)

1518 # as input to the multimodal model, even when the input is text.

1519 inputs_embeds_scheduled = self.model.get_input_embeddings(

1553 - input_ids=self.input_ids[:num_scheduled_tokens],

1520 + input_ids=self.input_ids.gpu[:num_scheduled_tokens],

1521 multimodal_embeddings=mm_embeds or None,

1522 )

1523

@@ -1569,13 +1536,13 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1536 # While it is possible to use embeddings as input just like the

1537 # multimodal models, it is not desirable for performance since

1538 # then the embedding layer is not included in the CUDA graph.

1572 - input_ids = self.input_ids[:num_input_tokens]

1539 + input_ids = self.input_ids.gpu[:num_input_tokens]

1540 inputs_embeds = None

1541 model_kwargs = self._init_model_kwargs(num_input_tokens)

1542 if self.uses_mrope:

1576 - positions = self.mrope_positions[:, :num_input_tokens]

1543 + positions = self.mrope_positions.gpu[:, :num_input_tokens]

1544 else:

1578 - positions = self.positions[:num_input_tokens]

1545 + positions = self.positions.gpu[:num_input_tokens]

1546

1547 if get_pp_group().is_first_rank:

1548 intermediate_tensors = None

@@ -1857,9 +1824,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1824

1825 if spec_decode_metadata is None:

1826 # input_ids can be None for multimodal models.

1860 - target_token_ids = self.input_ids[:num_scheduled_tokens]

1827 + target_token_ids = self.input_ids.gpu[:num_scheduled_tokens]

1828 # TODO(woosuk): Support M-RoPE.

1862 - target_positions = self.positions[:num_scheduled_tokens]

1829 + target_positions = self.positions.gpu[:num_scheduled_tokens]

1830 if self.use_aux_hidden_state_outputs:

1831 target_hidden_states = torch.cat(

1832 [h[:num_scheduled_tokens] for h in aux_hidden_states],

@@ -1879,9 +1846,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1846 self.drafter.prepare_inputs(

1847 common_attn_metadata, num_rejected_tokens_cpu)

1848

1882 - target_token_ids = self.input_ids[token_indices]

1849 + target_token_ids = self.input_ids.gpu[token_indices]

1850 # TODO(woosuk): Support M-RoPE.

1884 - target_positions = self.positions[token_indices]

1851 + target_positions = self.positions.gpu[token_indices]

1852 if self.use_aux_hidden_state_outputs:

1853 target_hidden_states = torch.cat(

1854 [h[token_indices] for h in aux_hidden_states], dim=-1)

@@ -2123,7 +2090,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2090 # If this is a partial request (i.e. chunked prefill),

2091 # then there is prompt logprob generated for each index.

2092 req_idx = self.input_batch.req_id_to_index[req_id]

2126 - offset = self.query_start_loc_np[req_idx].item()

2093 + offset = self.query_start_loc.np[req_idx].item()

2094 prompt_hidden_states = hidden_states[offset:offset + num_logits]

2095 logits = self.model.compute_logits(prompt_hidden_states, None)

2096

@@ -2196,7 +2163,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2163 @functools.cache

2164 def rand_input_ids() -> torch.Tensor:

2165 return torch.randint_like(

2199 - self.input_ids,

2166 + self.input_ids.gpu,

2167 low=0,

2168 high=self.model_config.get_vocab_size(),

2169 dtype=input_ids.dtype)

@@ -2313,18 +2280,18 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2280 attn_metadata = {}

2281

2282 # Make sure max_model_len is used at the graph capture time.

2316 - self.seq_lens_np[:num_reqs] = self.max_model_len

2317 - self.seq_lens_np[num_reqs:] = 0

2318 - self.seq_lens.copy_(self.seq_lens_cpu, non_blocking=True)

2283 + self.seq_lens.np[:num_reqs] = self.max_model_len

2284 + self.seq_lens.np[num_reqs:] = 0

2285 + self.seq_lens.copy_to_gpu()

2286

2287 for kv_cache_group_id, kv_cache_group_spec in enumerate(

2288 self.kv_cache_config.kv_cache_groups):

2289 common_attn_metadata = CommonAttentionMetadata(

2323 - query_start_loc=self.query_start_loc[:num_reqs + 1],

2324 - query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs +

2290 + query_start_loc=self.query_start_loc.gpu[:num_reqs + 1],

2291 + query_start_loc_cpu=self.query_start_loc.cpu[:num_reqs +

2292 1],

2326 - seq_lens=self.seq_lens[:num_reqs],

2327 - seq_lens_cpu=self.seq_lens_cpu[:num_reqs],

2293 + seq_lens=self.seq_lens.gpu[:num_reqs],

2294 + seq_lens_cpu=self.seq_lens.cpu[:num_reqs],

2295 num_computed_tokens_cpu=self.input_batch.

2296 num_computed_tokens_cpu_tensor[:num_reqs],

2297 num_reqs=num_reqs,

@@ -2353,14 +2320,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2320 **self._dummy_mm_kwargs(num_reqs),

2321 }

2322 else:

2356 - input_ids = self.input_ids[:num_tokens]

2323 + input_ids = self.input_ids.gpu[:num_tokens]

2324 inputs_embeds = None

2325 model_kwargs = self._init_model_kwargs(num_tokens)

2326

2327 if self.uses_mrope:

2361 - positions = self.mrope_positions[:, :num_tokens]

2328 + positions = self.mrope_positions.gpu[:, :num_tokens]

2329 else:

2363 - positions = self.positions[:num_tokens]

2330 + positions = self.positions.gpu[:num_tokens]

2331

2332 if get_pp_group().is_first_rank:

2333 intermediate_tensors = None

vllm/v1/worker/utils.py

+29

@@ -298,3 +298,32 @@ def bind_kv_cache(

298 for layer_name, kv_cache in kv_caches.items():

299 # NOTE: Use list because of v0 PP virtual engine.

300 forward_context[layer_name].kv_cache = [kv_cache]

301 +

302 +

303 +class CpuGpuBuffer:

304 +

305 + def __init__(

306 + self,

307 + *args,

308 + dtype: torch.dtype,

309 + device: torch.device,

310 + pin_memory: bool,

311 + ):

312 + self.cpu = torch.zeros(*args,

313 + dtype=dtype,

314 + device="cpu",

315 + pin_memory=pin_memory)

316 + self.np = self.cpu.numpy()

317 + self.gpu = self.cpu.to(device)

318 +

319 + def copy_to_gpu(self, n: Optional[int] = None) -> None:

320 + if n is None:

321 + return self.gpu.copy_(self.cpu, non_blocking=True)

322 + return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True)

323 +

324 + def copy_to_cpu(self, n: Optional[int] = None) -> None:

325 + """NOTE: Because this method is non-blocking, explicit synchronization

326 + is needed to ensure the data is copied to CPU."""

327 + if n is None:

328 + return self.cpu.copy_(self.gpu, non_blocking=True)

329 + return self.cpu[:n].copy_(self.gpu[:n], non_blocking=True)