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 3, 21:30 UTC
[Perf][1/n] Eliminate various GPU<->CPU syncs (#41429)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Nick Hill committed
May 11, 2026 at 13:36 UTC
bbee5329880bfe45a8b0cf7753d17501760e62c9
7 files changed
+49
-33
vllm/model_executor/layers/mamba/mamba_mixer.py
+1
@@ -362,6 +362,7 @@ class MambaMixer(MambaBase, PluggableLayer):
362
initial_state_idx=block_idx_last_computed_token_p,
363
num_computed_tokens=num_computed_tokens_p,
364
block_size_to_align=mamba_block_size,
365
+ metadata=attn_metadata,
366
)
367
# 3. State Space Model sequence transformations.
368
discrete_time_step_p, B_p, C_p = self._ssm_transform(
vllm/v1/sample/ops/bad_words.py
+2
-1
@@ -23,7 +23,8 @@ def _apply_bad_words_single_batch(
23
assert len(actual_prefix) == len(expected_prefix)
24
25
if actual_prefix == expected_prefix:
26
- logits[last_token_id] = _SMALLEST_LOGIT
26
+ # Assign to slice to avoid cpu->gpu sync.
27
+ logits[last_token_id : last_token_id + 1] = _SMALLEST_LOGIT
28
29
30
def apply_bad_words(
vllm/v1/sample/sampler.py
+17
-14
@@ -175,27 +175,30 @@ class Sampler(nn.Module):
175
176
# Find max number of tokens across all requests
177
max_num_tokens = max(len(tids) for tids in logprob_token_ids.values())
178
+ pin = self.pin_memory
179
179
- # Create padded token_ids tensor: [batch_size, max_num_tokens + 1]
180
- # +1 for sampled token in first position
181
- token_ids_tensor = torch.zeros(
182
- batch_size, max_num_tokens + 1, dtype=torch.int64, device=device
180
+ # Build the padded token_ids and valid_mask matrices on pinned CPU,
181
+ # then upload non-blocking.
182
+ token_ids_cpu = torch.zeros(
183
+ batch_size, max_num_tokens + 1, dtype=torch.int64, pin_memory=pin
184
)
184
- token_ids_tensor[:, 0] = sampled # First column is sampled token
185
-
185
# Create mask for valid positions (True = valid, False = padded)
187
- valid_mask = torch.zeros(
188
- batch_size, max_num_tokens + 1, dtype=torch.bool, device=device
186
+ valid_mask_cpu = torch.zeros(
187
+ batch_size, max_num_tokens + 1, dtype=torch.bool, pin_memory=pin
188
)
190
- valid_mask[:, 0] = True # Sampled token is always valid
191
-
192
- # Fill in token IDs for each request
189
+ valid_mask_cpu[:, 0] = True # Sampled token is always valid
190
for req_idx, token_ids in logprob_token_ids.items():
191
num_tokens = len(token_ids)
195
- token_ids_tensor[req_idx, 1 : num_tokens + 1] = torch.tensor(
196
- token_ids, dtype=torch.int64, device=device
192
+ token_ids_cpu[req_idx, 1 : num_tokens + 1] = torch.as_tensor(
193
+ token_ids, dtype=torch.int64
194
)
198
- valid_mask[req_idx, 1 : num_tokens + 1] = True
195
+ valid_mask_cpu[req_idx, 1 : num_tokens + 1] = True
196
+
197
+ token_ids_tensor = token_ids_cpu.to(device, non_blocking=True)
198
+ valid_mask = valid_mask_cpu.to(device, non_blocking=True)
199
+ # Sampled token in column 0 — fill on-device from the sampled GPU
200
+ # tensor so we don't need to D2H + re-upload.
201
+ token_ids_tensor[:, 0] = sampled
202
203
# Compute logprobs using the fused Triton kernel (log_softmax + gather)
204
logprobs = compute_token_logprobs(logits, token_ids_tensor)
vllm/v1/spec_decode/ngram_proposer_gpu.py
+3
-2
@@ -18,6 +18,7 @@ from vllm.config import (
18
VllmConfig,
19
)
20
from vllm.forward_context import set_forward_context
21
+from vllm.utils.torch_utils import async_tensor_h2d
22
from vllm.v1.core.sched.output import SchedulerOutput
23
from vllm.v1.utils import record_function_or_nullcontext
24
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
@@ -569,8 +570,8 @@ def update_ngram_gpu_tensors_incremental(
570
reorder_dst.append(curr_idx)
571
572
if reorder_src:
572
- src_tensor = torch.tensor(reorder_src, dtype=torch.long, device=device)
573
- dst_tensor = torch.tensor(reorder_dst, dtype=torch.long, device=device)
573
+ src_tensor = async_tensor_h2d(reorder_src, dtype=torch.long, device=device)
574
+ dst_tensor = async_tensor_h2d(reorder_dst, dtype=torch.long, device=device)
575
576
temp_token_ids = token_ids_gpu_tensor[src_tensor].clone()
577
temp_num_tokens = num_tokens_no_spec_gpu[src_tensor].clone()
vllm/v1/worker/dp_utils.py
+7
-5
@@ -43,11 +43,13 @@ def _run_ar(
43
dp_size = parallel_config.data_parallel_size
44
dp_rank = parallel_config.data_parallel_rank
45
device, group = _get_device_and_group(parallel_config)
46
- tensor = torch.zeros(4, dp_size, device=device, dtype=torch.int32)
47
- tensor[0][dp_rank] = orig_num_tokens_per_ubatch
48
- tensor[1][dp_rank] = padded_num_tokens_per_ubatch
49
- tensor[2][dp_rank] = 1 if should_ubatch else 0
50
- tensor[3][dp_rank] = cudagraph_mode
46
+ # Populate this rank's contribution on CPU to reduce GPU syncs.
47
+ tensor_cpu = torch.zeros(4, dp_size, dtype=torch.int32)
48
+ tensor_cpu[0][dp_rank] = orig_num_tokens_per_ubatch
49
+ tensor_cpu[1][dp_rank] = padded_num_tokens_per_ubatch
50
+ tensor_cpu[2][dp_rank] = 1 if should_ubatch else 0
51
+ tensor_cpu[3][dp_rank] = cudagraph_mode
52
+ tensor = tensor_cpu.to(device, non_blocking=True)
53
dist.all_reduce(tensor, group=group)
54
return tensor
55
vllm/v1/worker/gpu/sample/penalties.py
+4
-2
@@ -283,8 +283,10 @@ def bincount(
283
output_bin_counts: torch.Tensor,
284
max_prefill_len: int,
285
) -> None:
286
- prompt_bin_mask[expanded_idx_mapping] = 0
287
- output_bin_counts[expanded_idx_mapping] = 0
286
+ # Use index_fill_ instead of `tensor[idx] = 0` to avoid sync.
287
+ idx_long = expanded_idx_mapping.long()
288
+ prompt_bin_mask.index_fill_(0, idx_long, 0)
289
+ output_bin_counts.index_fill_(0, idx_long, 0)
290
num_tokens = expanded_idx_mapping.shape[0]
291
BLOCK_SIZE = 1024
292
num_blocks = triton.cdiv(max_prefill_len, BLOCK_SIZE)
vllm/v1/worker/gpu_model_runner.py
+15
-9
@@ -995,16 +995,23 @@ class GPUModelRunner(
995
if len(token_type_id_requests) == 0:
996
return model_kwargs
997
998
- seq_lens = self.seq_lens[:num_reqs]
998
+ # Build ids on CPU using the CPU-resident upper bound for seq_lens;
999
+ # `torch.arange(seq_lens[i])` with a GPU scalar would force a sync.
1000
+ seq_lens_cpu = self.optimistic_seq_lens_cpu[:num_reqs].tolist()
1001
token_type_ids = []
1002
1003
for i in range(num_reqs):
1002
- pos = token_type_id_requests.get(i, seq_lens[i])
1003
- ids = (torch.arange(seq_lens[i]) >= pos).int()
1004
+ seq_len_i = seq_lens_cpu[i]
1005
+ pos = token_type_id_requests.get(i, seq_len_i)
1006
+ ids = (torch.arange(seq_len_i) >= pos).int()
1007
token_type_ids.append(ids)
1008
1006
- model_kwargs["token_type_ids"] = torch.concat(token_type_ids).to(
1007
- device=self.device
1009
+ token_type_ids_cpu = torch.empty(
1010
+ sum(seq_lens_cpu), dtype=torch.int32, pin_memory=self.pin_memory
1011
+ )
1012
+ torch.cat(token_type_ids, out=token_type_ids_cpu)
1013
+ model_kwargs["token_type_ids"] = token_type_ids_cpu.to(
1014
+ device=self.device, non_blocking=True
1015
)
1016
return model_kwargs
1017
@@ -2715,10 +2722,9 @@ class GPUModelRunner(
2722
# There might have leftover indices in logits_indices[num_logits:]
2723
# from previous iterations, whose values may be greater than the
2724
# batch size in the current iteration. To ensure indices are always
2718
- # valid, we fill the padded indices with the last index.
2719
- self.kv_sharing_fast_prefill_logits_indices[num_logits:].fill_(
2720
- logits_indices[-1].item()
2721
- )
2725
+ # valid, fill the padded indices with the last index. Broadcast the
2726
+ # scalar GPU-side to avoid a D2H sync on `.item()`.
2727
+ self.kv_sharing_fast_prefill_logits_indices[num_logits:] = logits_indices[-1]
2728
# Dispatch for the decoder portion of the model.
2729
_, batch_desc = self.cudagraph_dispatcher.dispatch(
2730
num_logits, invalid_modes={CUDAGraphMode.FULL}