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 9, 12:30 UTC
[Model Runner V2] Use a different stream for grammar bitmask h2d copy (#33059)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Co-authored-by: Nick Hill <nhill@redhat.com>
Woosuk Kwon committed
Jan 27, 2026 at 16:37 UTC
46ec6d71c71c50d088a03e0b04e6f49cf8637c00
2 files changed
+32
-11
vllm/v1/worker/gpu/model_runner.py
+1
@@ -168,6 +168,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
168
self.structured_outputs_worker = StructuredOutputsWorker(
169
max_num_logits=self.max_num_reqs * (self.num_speculative_steps + 1),
170
vocab_size=self.vocab_size,
171
+ device=self.device,
172
)
173
# LoRA-related workers.
174
self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
vllm/v1/worker/gpu/structured_outputs.py
+31
-11
@@ -5,7 +5,7 @@ import torch
5
6
from vllm.triton_utils import tl, triton
7
from vllm.utils.math_utils import cdiv
8
-from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool
8
+from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
9
from vllm.v1.worker.gpu.input_batch import InputBatch
10
11
@@ -14,13 +14,16 @@ class StructuredOutputsWorker:
14
self,
15
max_num_logits: int,
16
vocab_size: int,
17
+ device: torch.device,
18
):
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
19
+ self.logits_indices = torch.zeros(
20
+ max_num_logits, dtype=torch.int32, device=device
21
)
22
+ self.grammar_bitmask = torch.zeros(
23
+ (max_num_logits, cdiv(vocab_size, 32)), dtype=torch.int32, device=device
24
+ )
25
+ self.device = device
26
+ self.copy_stream = torch.cuda.Stream()
27
28
def apply_grammar_bitmask(
29
self,
@@ -32,6 +35,12 @@ class StructuredOutputsWorker:
35
if not grammar_req_ids:
36
return
37
38
+ # Asynchronously copy the bitmask to GPU.
39
+ with torch.cuda.stream(self.copy_stream):
40
+ bitmask = async_copy_to_gpu(
41
+ grammar_bitmask, out=self.grammar_bitmask[: grammar_bitmask.shape[0]]
42
+ )
43
+
44
# Construct bitmask -> logits mapping
45
mapping: list[int] = []
46
req_ids = input_batch.req_ids
@@ -42,12 +51,19 @@ class StructuredOutputsWorker:
51
logits_start_idx = cu_num_logits[req_idx]
52
logits_end_idx = cu_num_logits[req_idx + 1]
53
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)
54
49
- # Copy the bitmask.
50
- bitmask = self.grammar_bitmask.copy_to_uva(grammar_bitmask)
55
+ # Asynchronously copy the mapping to GPU.
56
+ with torch.cuda.stream(self.copy_stream):
57
+ logits_indices = torch.tensor(
58
+ mapping, dtype=torch.int32, device="cpu", pin_memory=True
59
+ )
60
+ logits_indices = self.logits_indices[: len(mapping)].copy_(
61
+ logits_indices, non_blocking=True
62
+ )
63
+
64
+ # Ensure all async copies are complete before launching the kernel.
65
+ current_stream = torch.cuda.current_stream()
66
+ current_stream.wait_stream(self.copy_stream)
67
68
num_masks = bitmask.shape[0]
69
assert num_masks == len(mapping)
@@ -64,6 +80,10 @@ class StructuredOutputsWorker:
80
BLOCK_SIZE=BLOCK_SIZE,
81
)
82
83
+ # Ensure the copy stream waits for the device tensors to finish being used
84
+ # before it re-uses or deallocates them
85
+ self.copy_stream.wait_stream(current_stream)
86
+
87
88
# Adapted from
89
# https://github.com/mlc-ai/xgrammar/blob/main/python/xgrammar/kernels/apply_token_bitmask_inplace_triton.py