@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 185 lines 6.14 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 #
4 # Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
5 # Adapted from vllm/model_executor/layers/rotary_embedding/__init__.py
6 # Copyright 2023 The vLLM team.
7 #
8 # This file is a part of the vllm-ascend project.
9 #
10 # Licensed under the Apache License, Version 2.0 (the "License");
11 # you may not use this file except in compliance with the License.
12 # You may obtain a copy of the License at
13 #
14 # http://www.apache.org/licenses/LICENSE-2.0
15 #
16 # Unless required by applicable law or agreed to in writing, software
17 # distributed under the License is distributed on an "AS IS" BASIS,
18 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
19 # See the License for the specific language governing permissions and
20 # limitations under the License.
21
22 import torch
23
24 from .mrope import MRotaryEmbedding
25
26
27 # MRotaryEmbedding with interleaved
28 class MRotaryEmbeddingInterleaved(MRotaryEmbedding):
29 """Rotary Embedding with Multimodal Sections and Interleaved Support."""
30
31 def __init__(
32 self,
33 head_size: int,
34 rotary_dim: int,
35 max_position_embeddings: int,
36 base: float,
37 is_neox_style: bool,
38 dtype: torch.dtype,
39 mrope_section: list[int],
40 mrope_interleaved: bool = True,
41 ) -> None:
42 # Enlarge max_position_embeddings for video inputs
43 self.cache_max_position_num = max_position_embeddings
44 super().__init__(
45 head_size,
46 rotary_dim,
47 self.cache_max_position_num,
48 base,
49 is_neox_style,
50 dtype,
51 )
52
53 self.mrope_section = mrope_section
54 self.mrope_interleaved = mrope_interleaved
55
56 if self.mrope_section is None:
57 raise ValueError("mrope_section cannot be None.")
58 if sum(self.mrope_section) != rotary_dim // 2:
59 raise ValueError("Sum of mrope_section must equal rotary_dim // 2.")
60 if not self.mrope_interleaved:
61 raise ValueError(
62 "mrope_interleaved must be True when mrope_section is provided."
63 )
64
65 # Generate interleaved indices
66 if len(mrope_section) == 2:
67 h_num, w_num = mrope_section[0], mrope_section[1]
68 mrope_dim = self.get_mrope_interleaved_id_list(h_num, w_num, 0)
69 elif len(mrope_section) == 3:
70 t_num, h_num, w_num = mrope_section[0], mrope_section[1], mrope_section[2]
71 mrope_dim = self.get_mrope_interleaved_id_list(
72 t_num, h_num, w_num, force_last=True
73 )
74 else:
75 raise AssertionError(
76 "Cannot support the length of mrope section is not 2 or 3."
77 )
78
79 mrope_dim = mrope_dim * 2
80 self.mrope_dim = mrope_dim
81
82 self.layer_cache = None
83
84 def _rebuild_pos_emb(
85 self,
86 positions: torch.Tensor,
87 ) -> tuple[torch.Tensor, torch.Tensor]:
88 """Interleave the rotary embedding."""
89 cos_sin = self.cos_sin_cache[positions]
90 mrope_section_3d = [1] * len(self.mrope_dim)
91 mrope_dim = self.mrope_dim
92 cos_sin = torch.cat(
93 [
94 m[mrope_dim[i]]
95 for i, m in enumerate(cos_sin.split(mrope_section_3d, dim=-1))
96 ],
97 dim=-1,
98 )
99 return cos_sin, torch.arange(cos_sin.shape[0], device=positions.device)
100
101 def forward(
102 self,
103 positions: torch.Tensor,
104 query: torch.Tensor,
105 key: torch.Tensor | None = None,
106 ) -> tuple[torch.Tensor, torch.Tensor | None]:
107 """Forward pass with interleaved rotary embedding."""
108 cos_sin, positions = self._rebuild_pos_emb(positions)
109 cos, sin = cos_sin.chunk(2, dim=-1)
110
111 query_shape = query.shape
112 positions = positions.flatten()
113 num_tokens = positions.shape[0]
114 query = query.view(num_tokens, -1, self.head_size)
115 query_rot = query[..., : self.rotary_dim]
116 query_pass = query[..., self.rotary_dim :]
117 query_rot = self.apply_rotary_emb.forward_native(
118 query_rot,
119 cos,
120 sin,
121 )
122 query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
123
124 # key may be None in some cases, e.g. cross-layer KV sharing
125 if key is not None:
126 key_shape = key.shape
127 key = key.view(num_tokens, -1, self.head_size)
128 key_rot = key[..., : self.rotary_dim]
129 key_pass = key[..., self.rotary_dim :]
130 key_rot = self.apply_rotary_emb.forward_native(
131 key_rot,
132 cos,
133 sin,
134 )
135 key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
136 return query, key
137
138 @staticmethod
139 def get_mrope_interleaved_id_list(
140 a: int, b: int, c: int, force_last: bool = False
141 ) -> list[int]:
142 """Generate an interleaved list of indices for multi-modal rotary embedding.
143
144 Args:
145 a: Number of indices for first modality
146 b: Number of indices for second modality
147 c: Number of indices for third modality
148 force_last: Whether to force the last element to be from the first modality
149
150 Returns:
151 List of interleaved indices
152
153 """
154 if force_last:
155 a -= 1
156
157 counts = {0: a, 1: b, 2: c}
158 placed = {k: 0 for k in counts}
159 rem = counts.copy()
160 seq: list[int] = []
161 last = None
162
163 total = a + b + c
164 for _ in range(total):
165 # Candidates: remaining > 0 and ≠ last
166 cands = [k for k in rem if rem[k] > 0 and k != last]
167 if not cands:
168 # If only last remains, relax the condition
169 cands = [k for k in rem if rem[k] > 0]
170
171 # Select the rarest candidate
172 try:
173 best = min(cands, key=lambda k: (placed[k] / counts[k], k))
174 except KeyError:
175 best = 0
176
177 seq.append(best)
178 placed[best] += 1
179 rem[best] -= 1
180 last = best
181
182 if force_last:
183 seq.append(0)
184
185 return seq