@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 152 lines 5.79 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import torch
4
5 from vllm.triton_utils import tl, triton
6 from vllm.utils.torch_utils import async_tensor_h2d
7 from vllm.v1.core.sched.output import NewRequestData
8 from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
9 from vllm.v1.worker.gpu.input_batch import InputBatch
10
11 TOKEN_BLOCK = 16
12
13
14 class PromptEmbedsState:
15 """GPU-side state for user-provided prompt embeddings.
16
17 Each request's embeddings are copied to the GPU once at `add_request`,
18 off the per-step hot path. A per-request pointer table (UVA) then lets a
19 single triton kernel overlay all scheduled prompt-embeds rows onto
20 `inputs_embeds` each step, with no python loops or per-request H2D copies.
21 """
22
23 def __init__(
24 self,
25 max_num_reqs: int,
26 hidden_size: int,
27 dtype: torch.dtype,
28 device: torch.device,
29 ):
30 self.hidden_size = hidden_size
31 self.dtype = dtype
32 self.device = device
33
34 # req_id -> (embeds, is_token_ids mask or None). Holds the references
35 # that keep the pointer table below valid.
36 self.gpu_tensors: dict[str, tuple[torch.Tensor, torch.Tensor | None]] = {}
37
38 # Indexed by req_state index. Stale entries after removal are
39 # harmless: add_request rewrites all fields for every index it claims.
40 self.embeds_ptrs = UvaBackedTensor(max_num_reqs, dtype=torch.uint64)
41 self.mask_ptrs = UvaBackedTensor(max_num_reqs, dtype=torch.uint64)
42 self.embeds_lens = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
43
44 def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
45 prompt_embeds = new_req_data.prompt_embeds
46 if prompt_embeds is None:
47 self.gpu_tensors.pop(new_req_data.req_id, None)
48 self.embeds_lens.np[req_index] = 0
49 return
50
51 embeds = async_tensor_h2d(prompt_embeds, device=self.device, dtype=self.dtype)
52 embeds = embeds.contiguous()
53 is_token_ids = new_req_data.prompt_is_token_ids
54 mask = None
55 if is_token_ids is not None:
56 mask = async_tensor_h2d(is_token_ids, device=self.device, dtype=torch.uint8)
57 self.gpu_tensors[new_req_data.req_id] = (embeds, mask)
58 self.embeds_ptrs.np[req_index] = embeds.data_ptr()
59 self.mask_ptrs.np[req_index] = 0 if mask is None else mask.data_ptr()
60 self.embeds_lens.np[req_index] = embeds.shape[0]
61
62 def remove_request(self, req_id: str) -> None:
63 self.gpu_tensors.pop(req_id, None)
64
65 def apply_staged_writes(self) -> None:
66 self.embeds_ptrs.copy_to_uva()
67 self.mask_ptrs.copy_to_uva()
68 self.embeds_lens.copy_to_uva()
69
70 def apply(
71 self,
72 input_batch: InputBatch,
73 num_computed_tokens: torch.Tensor,
74 inputs_embeds: torch.Tensor,
75 ) -> None:
76 """Overlay prompt embeddings onto `inputs_embeds` for the batch."""
77 if not self.gpu_tensors:
78 return
79 # The kernel reinterprets raw source pointers as inputs_embeds' dtype.
80 assert inputs_embeds.dtype == self.dtype
81 num_reqs = input_batch.num_reqs
82 max_query_len = int(input_batch.num_scheduled_tokens.max())
83 grid = (num_reqs, triton.cdiv(max_query_len, TOKEN_BLOCK))
84 _apply_prompt_embeds_kernel[grid](
85 inputs_embeds,
86 inputs_embeds.stride(0),
87 self.embeds_ptrs.gpu,
88 self.mask_ptrs.gpu,
89 self.embeds_lens.gpu,
90 input_batch.idx_mapping,
91 input_batch.query_start_loc,
92 num_computed_tokens,
93 self.hidden_size,
94 TOKEN_BLOCK=TOKEN_BLOCK,
95 BLOCK_SIZE=1024,
96 )
97
98
99 @triton.jit
100 def _apply_prompt_embeds_kernel(
101 inputs_embeds_ptr,
102 inputs_embeds_stride,
103 embeds_ptrs_ptr, # uint64 [max_num_reqs], device pointers (0-len = unused)
104 mask_ptrs_ptr, # uint64 [max_num_reqs], 0 = no is-token-ids mask
105 embeds_lens_ptr, # int32 [max_num_reqs]
106 idx_mapping_ptr,
107 query_start_loc_ptr,
108 num_computed_tokens_ptr,
109 hidden_size: tl.constexpr,
110 TOKEN_BLOCK: tl.constexpr,
111 BLOCK_SIZE: tl.constexpr,
112 ):
113 batch_idx = tl.program_id(0)
114 req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
115 embeds_len = tl.load(embeds_lens_ptr + req_state_idx)
116 num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
117 if num_computed >= embeds_len:
118 # No prompt embeds for this request, or they are fully consumed.
119 return
120
121 query_start = tl.load(query_start_loc_ptr + batch_idx)
122 query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
123 num_rows = tl.minimum(query_end - query_start, embeds_len - num_computed)
124
125 t_start = tl.program_id(1) * TOKEN_BLOCK
126 if t_start >= num_rows:
127 return
128
129 src_ptr = tl.load(embeds_ptrs_ptr + req_state_idx).to(
130 tl.pointer_type(inputs_embeds_ptr.dtype.element_ty)
131 )
132 mask_int = tl.load(mask_ptrs_ptr + req_state_idx)
133 mask_ptr = mask_int.to(tl.pointer_type(tl.int8))
134
135 for t_offset in tl.static_range(TOKEN_BLOCK):
136 t = t_start + t_offset
137 if t < num_rows:
138 src_row = (num_computed + t).to(tl.int64)
139 is_token_id = 0
140 if mask_int != 0:
141 is_token_id = tl.load(mask_ptr + src_row).to(tl.int32)
142 if is_token_id == 0:
143 dst_row = (query_start + t).to(tl.int64)
144 for h in tl.range(0, hidden_size, BLOCK_SIZE):
145 offs = h + tl.arange(0, BLOCK_SIZE)
146 h_mask = offs < hidden_size
147 row = tl.load(src_ptr + src_row * hidden_size + offs, mask=h_mask)
148 tl.store(
149 inputs_embeds_ptr + dst_row * inputs_embeds_stride + offs,
150 row,
151 mask=h_mask,
152 )