@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 199 lines 7.22 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import numpy as np
4 import torch
5
6 from vllm.triton_utils import tl, triton
7 from vllm.utils.math_utils import cdiv
8 from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d
9 from vllm.v1.worker.gpu.input_batch import InputBatch
10
11
12 def _build_grammar_mapping(
13 req_ids: list[str],
14 grammar_req_ids: list[str],
15 cu_num_logits_np: np.ndarray,
16 num_draft_tokens_per_req: np.ndarray | None,
17 num_bonus_tokens: int,
18 mask_stride: int,
19 ) -> list[int]:
20 mapping: list[int] = []
21 req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)}
22 for grammar_req_id in grammar_req_ids:
23 req_idx = req_id_to_idx[grammar_req_id]
24 if num_draft_tokens_per_req is None:
25 num_positions = int(
26 cu_num_logits_np[req_idx + 1] - cu_num_logits_np[req_idx]
27 )
28 else:
29 # Grammar masks follow the scheduled layout even when adaptive
30 # verification compacts the actual CPU logit offsets to bonus-only.
31 num_positions = int(num_draft_tokens_per_req[req_idx]) + num_bonus_tokens
32 mapping.extend(
33 req_idx * mask_stride + position for position in range(num_positions)
34 )
35 return mapping
36
37
38 def grammar_invalid_drafts(
39 input_batch: InputBatch,
40 grammar_req_ids: list[str],
41 num_acceptable_drafts: list[int] | None,
42 ) -> torch.Tensor | None:
43 """Mask over logit rows whose draft was verified against a permissive row.
44
45 Draft i sits at local position i + 1, and drafts from `num_acceptable_drafts`
46 on met `_full_mask`. Local positions come from the device `cu_num_logits`, so
47 this stays right when adaptive verification trims drafts on device.
48 Returns None when there is nothing to invalidate.
49 """
50 if not grammar_req_ids or input_batch.num_draft_tokens == 0:
51 return None
52 req_id_to_idx = {req_id: i for i, req_id in enumerate(input_batch.req_ids)}
53 limit = np.full(input_batch.num_reqs, np.iinfo(np.int32).max, dtype=np.int32)
54 for i, req_id in enumerate(grammar_req_ids):
55 req_idx = req_id_to_idx.get(req_id)
56 if req_idx is not None:
57 # None (an older scheduler, or warmup): invalidate the whole window.
58 limit[req_idx] = (
59 num_acceptable_drafts[i] if num_acceptable_drafts is not None else 0
60 )
61 # Adaptive verification only trims drafts, so this bound holds either way.
62 num_drafts = input_batch.num_draft_tokens_per_req
63 if num_drafts is not None and (limit >= num_drafts).all():
64 return None
65 local_pos = input_batch.expanded_local_pos
66 cu_num_logits = input_batch.cu_num_logits
67 row_limit = torch.repeat_interleave(
68 async_tensor_h2d(limit, device=local_pos.device),
69 cu_num_logits[1:] - cu_num_logits[:-1],
70 output_size=local_pos.shape[0],
71 )
72 return local_pos > row_limit
73
74
75 class StructuredOutputsWorker:
76 def __init__(
77 self,
78 max_num_logits: int,
79 vocab_size: int,
80 device: torch.device,
81 mask_stride: int,
82 num_bonus_tokens: int,
83 ):
84 self.logits_indices = torch.zeros(
85 max_num_logits, dtype=torch.int32, device=device
86 )
87 self.grammar_bitmask = torch.zeros(
88 (max_num_logits, cdiv(vocab_size, 32)), dtype=torch.int32, device=device
89 )
90 self.device = device
91 self.copy_stream = torch.cuda.Stream()
92 self.mask_stride = mask_stride
93 self.num_bonus_tokens = num_bonus_tokens
94
95 def apply_grammar_bitmask(
96 self,
97 logits: torch.Tensor,
98 input_batch: InputBatch,
99 grammar_req_ids: list[str],
100 grammar_bitmask: np.ndarray,
101 ) -> None:
102 if not grammar_req_ids:
103 return
104
105 # Asynchronously copy the bitmask to GPU.
106 with torch.cuda.stream(self.copy_stream):
107 bitmask = async_tensor_h2d(
108 grammar_bitmask, out=self.grammar_bitmask[: grammar_bitmask.shape[0]]
109 )
110
111 # Construct bitmask -> logits mapping
112 # Key by (request, position) rather than absolute logit index:
113 # adaptive verification finalizes per-request logit offsets on
114 # device, so the kernel resolves them from the GPU cu_num_logits.
115 mapping = _build_grammar_mapping(
116 input_batch.req_ids,
117 grammar_req_ids,
118 input_batch.cu_num_logits_np,
119 input_batch.num_draft_tokens_per_req,
120 self.num_bonus_tokens,
121 self.mask_stride,
122 )
123
124 # Asynchronously copy the mapping to GPU.
125 with torch.cuda.stream(self.copy_stream):
126 logits_indices = torch.tensor(
127 mapping, dtype=torch.int32, device="cpu", pin_memory=PIN_MEMORY
128 )
129 logits_indices = self.logits_indices[: len(mapping)].copy_(
130 logits_indices, non_blocking=True
131 )
132
133 # Ensure all async copies are complete before launching the kernel.
134 current_stream = torch.cuda.current_stream()
135 current_stream.wait_stream(self.copy_stream)
136
137 num_masks = bitmask.shape[0]
138 assert num_masks == len(mapping)
139 vocab_size = logits.shape[-1]
140 BLOCK_SIZE = 8192
141 grid = (num_masks, triton.cdiv(vocab_size, BLOCK_SIZE))
142 _apply_grammar_bitmask_kernel[grid](
143 logits,
144 logits.stride(0),
145 logits_indices,
146 input_batch.cu_num_logits,
147 bitmask,
148 bitmask.stride(0),
149 vocab_size,
150 MASK_STRIDE=self.mask_stride,
151 BLOCK_SIZE=BLOCK_SIZE,
152 )
153
154 # Ensure the copy stream waits for the device tensors to finish being used
155 # before it re-uses or deallocates them
156 self.copy_stream.wait_stream(current_stream)
157
158
159 # Adapted from
160 # https://github.com/mlc-ai/xgrammar/blob/main/python/xgrammar/kernels/apply_token_bitmask_inplace_triton.py
161 @triton.jit
162 def _apply_grammar_bitmask_kernel(
163 logits_ptr,
164 logits_stride,
165 logits_indices_ptr,
166 cu_num_logits_ptr,
167 bitmask_ptr,
168 bitmask_stride,
169 vocab_size,
170 MASK_STRIDE: tl.constexpr,
171 BLOCK_SIZE: tl.constexpr,
172 ):
173 bitmask_idx = tl.program_id(0)
174 mapping_idx = tl.load(logits_indices_ptr + bitmask_idx)
175 req_idx = mapping_idx // MASK_STRIDE
176 position_idx = mapping_idx % MASK_STRIDE
177 logits_idx = tl.load(cu_num_logits_ptr + req_idx)
178 num_req_logits = tl.load(cu_num_logits_ptr + req_idx + 1) - logits_idx
179 logits_idx += position_idx
180 position_is_active = position_idx < num_req_logits
181
182 # Load the bitmask.
183 block_id = tl.program_id(1)
184 bitmask_offset = (block_id * BLOCK_SIZE) // 32 + tl.arange(0, BLOCK_SIZE // 32)
185 packed_bitmask = tl.load(
186 bitmask_ptr + bitmask_idx * bitmask_stride + bitmask_offset,
187 mask=bitmask_offset < bitmask_stride,
188 )
189 # Unpack the bitmask.
190 bitmask = ((packed_bitmask[:, None] >> (tl.arange(0, 32)[None, :])) & 1) == 0
191 bitmask = bitmask.reshape(BLOCK_SIZE)
192
193 # Apply the bitmask to the logits.
194 block_offset = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
195 tl.store(
196 logits_ptr + logits_idx * logits_stride + block_offset,
197 -float("inf"),
198 mask=position_is_active & bitmask & (block_offset < vocab_size),
199 )