@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 192 lines 6.18 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from typing import cast
4
5 import torch
6 import torch.nn as nn
7
8 from vllm.config import ModelConfig
9 from vllm.model_executor.models.interfaces import SupportsMRoPE
10 from vllm.triton_utils import tl, triton
11 from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
12
13
14 class RopeState:
15 """State for multi-dimensional (M-RoPE) positions.
16
17 `num_dims` is the number of position channels the model consumes, which
18 the model config derives from its M-RoPE sections.
19
20 NOTE: `positions` is implemented with one additional dummy position on
21 purpose to make it non-contiguous so that it can work with torch compile.
22 See detailed explanation in
23 https://github.com/vllm-project/vllm/pull/12128#discussion_r1926431923
24
25 NOTE: For text-only inputs, each dimension has identical position IDs,
26 making M-RoPE functionally equivalent to 1D-RoPE.
27 See page 5 of https://arxiv.org/abs/2409.12191
28 """
29
30 def __init__(
31 self,
32 num_dims: int,
33 max_num_reqs: int,
34 max_num_tokens: int,
35 max_model_len: int,
36 device: torch.device,
37 ):
38 self.num_dims = num_dims
39 self.max_num_reqs = max_num_reqs
40 self.max_num_tokens = max_num_tokens
41 self.max_model_len = max_model_len
42 self.device = device
43
44 # NOTE(woosuk): This tensor can be extremely large (e.g., several GBs)
45 # wasting a lot of CPU memory.
46 self.prefill_positions = StagedWriteTensor(
47 (max_num_reqs * num_dims, max_model_len),
48 dtype=torch.int32,
49 device=device,
50 uva_instead_of_gpu=True,
51 )
52 self.positions = torch.zeros(
53 (num_dims, max_num_tokens + 1), dtype=torch.int64, device=device
54 )
55
56 self.prefill_delta = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
57
58 def init_prefill_positions(
59 self,
60 req_idx: int,
61 model: nn.Module,
62 prefill_token_ids: list[int],
63 mm_features: list,
64 ) -> None:
65 mrope_model = cast(SupportsMRoPE, model)
66 prefill_positions, delta = mrope_model.get_mrope_input_positions(
67 prefill_token_ids, mm_features
68 )
69 self.prefill_delta.np[req_idx] = delta
70
71 for i in range(self.num_dims):
72 pos = prefill_positions[i].tolist()
73 self.prefill_positions.stage_write(self.num_dims * req_idx + i, 0, pos)
74
75 def apply_staged_writes(self) -> None:
76 self.prefill_positions.apply_write()
77 self.prefill_delta.copy_to_uva()
78
79 def get_positions(self, num_tokens: int) -> torch.Tensor:
80 return self.positions[:, :num_tokens]
81
82 def read_prefill_positions(self, req_idx: int, length: int) -> torch.Tensor:
83 """Return staged per-request prefill positions as [num_dims, length]."""
84 base = self.num_dims * req_idx
85 return self.prefill_positions.gpu[base : base + self.num_dims, :length]
86
87 def update_prefill_positions(
88 self, req_idx: int, positions: torch.Tensor, delta: int
89 ) -> None:
90 """Overwrite a request's staged prefill positions with recomputed values."""
91 base = self.num_dims * req_idx
92 length = positions.shape[1]
93 self.prefill_positions.gpu[base : base + self.num_dims, :length].copy_(
94 positions
95 )
96 self.prefill_delta.np[req_idx] = delta
97
98 def prepare_positions(
99 self,
100 idx_mapping: torch.Tensor,
101 query_start_loc: torch.Tensor,
102 prefill_lens: torch.Tensor,
103 num_computed_tokens: torch.Tensor,
104 ) -> None:
105 num_reqs = idx_mapping.shape[0]
106 _prepare_rope_positions_kernel[(num_reqs,)](
107 self.positions,
108 self.positions.stride(0),
109 self.prefill_positions.gpu,
110 self.num_dims * self.max_model_len,
111 self.max_model_len,
112 self.prefill_delta.gpu,
113 idx_mapping,
114 query_start_loc,
115 prefill_lens,
116 num_computed_tokens,
117 BLOCK_SIZE=1024,
118 NUM_DIMS=self.num_dims,
119 )
120
121
122 def get_rope_state(
123 model_config: ModelConfig,
124 model: nn.Module,
125 max_num_reqs: int,
126 max_num_tokens: int,
127 max_model_len: int,
128 device: torch.device,
129 ) -> RopeState | None:
130 """Create a RopeState if the model uses multi-dimensional RoPE."""
131 if not model_config.uses_mrope:
132 return None
133
134 assert isinstance(model, SupportsMRoPE)
135 return RopeState(
136 num_dims=model_config.mrope_num_dims,
137 max_num_reqs=max_num_reqs,
138 max_num_tokens=max_num_tokens,
139 max_model_len=max_model_len,
140 device=device,
141 )
142
143
144 @triton.jit
145 def _prepare_rope_positions_kernel(
146 positions_ptr,
147 positions_stride,
148 prefill_positions_ptr,
149 prefill_positions_stride0,
150 prefill_positions_stride1,
151 prefill_delta_ptr,
152 idx_mapping_ptr,
153 query_start_loc_ptr,
154 prefill_lens_ptr,
155 num_computed_tokens_ptr,
156 BLOCK_SIZE: tl.constexpr,
157 NUM_DIMS: tl.constexpr,
158 ):
159 batch_idx = tl.program_id(0)
160 req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
161
162 prefill_len = tl.load(prefill_lens_ptr + req_state_idx)
163 num_computed = tl.load(num_computed_tokens_ptr + req_state_idx)
164 is_prefill = num_computed < prefill_len
165
166 query_start = tl.load(query_start_loc_ptr + batch_idx)
167 query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
168 query_len = query_end - query_start
169
170 delta = tl.load(prefill_delta_ptr + req_state_idx)
171
172 for i in range(0, query_len, BLOCK_SIZE):
173 block = i + tl.arange(0, BLOCK_SIZE)
174 mask = block < query_len
175 orig_pos = num_computed + block
176
177 for j in tl.static_range(NUM_DIMS):
178 if is_prefill:
179 pos = tl.load(
180 prefill_positions_ptr
181 + req_state_idx * prefill_positions_stride0
182 + j * prefill_positions_stride1
183 + orig_pos,
184 mask=mask,
185 )
186 else:
187 pos = orig_pos + delta
188 tl.store(
189 positions_ptr + j * positions_stride + query_start + block,
190 pos,
191 mask=mask,
192 )