Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / rotary_embedding.py: 21%
100 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""NPU-accelerated Transformers-compatible rotary position embeddings."""
17from __future__ import annotations
19import numpy as np
20import torch # pylint: disable=forbidden-backend-import
21import torch_npu
24def apply_rotary_pos_emb(
25 q: torch.Tensor,
26 k: torch.Tensor,
27 cos: torch.Tensor,
28 sin: torch.Tensor,
29 unsqueeze_dim: int = 1,
30) -> tuple[torch.Tensor, torch.Tensor]:
31 """Apply NPU RoPE using the Transformers ``apply_rotary_pos_emb`` contract.
33 Args:
34 q: Query tensor.
35 k: Key tensor.
36 cos: Cosine part of the rotary embedding.
37 sin: Sine part of the rotary embedding.
38 unsqueeze_dim: Dimension used to make ``cos`` and ``sin`` broadcastable
39 to ``q`` and ``k``.
41 Returns:
42 Rotated query and key tensors with the same shapes as the inputs.
44 Raises:
45 ValueError: If the rotary dimension does not cover the complete Q/K
46 head dimension supported by the NPU operator interface.
47 """
48 if cos.ndim == 2:
49 cos = cos.unsqueeze(0)
50 sin = sin.unsqueeze(0)
51 if cos.shape[0] == 1 and q.shape[0] > 1:
52 cos = cos.expand(q.shape[0], -1, -1)
53 sin = sin.expand(q.shape[0], -1, -1)
54 cos = cos.unsqueeze(unsqueeze_dim)
55 sin = sin.unsqueeze(unsqueeze_dim)
56 rotary_dim = cos.shape[-1]
57 if rotary_dim != q.shape[-1] or rotary_dim != k.shape[-1]:
58 raise ValueError(
59 "NPU RoPE requires the rotary dimension to equal the Q/K head dimension"
60 )
61 q_embed = torch_npu.npu_rotary_mul(q.clone(), cos, sin, rotary_mode="half")
62 k_embed = torch_npu.npu_rotary_mul(k.clone(), cos, sin, rotary_mode="half")
63 return q_embed, k_embed
66def _apply_interleaved_rope(
67 tensor: torch.Tensor,
68 pass_through: torch.Tensor,
69 cos: torch.Tensor,
70 sin: torch.Tensor,
71 unsqueeze_dim: int,
72) -> torch.Tensor:
73 """Apply interleaved RoPE and append dimensions that bypass rotation."""
74 sequence_first = (
75 tensor.permute(2, 0, 1, 3)
76 if unsqueeze_dim == 1
77 else tensor.permute(1, 0, 2, 3)
78 )
79 seq_length, batch_size, num_heads, head_dim = sequence_first.shape
80 if batch_size > 1 and seq_length > 1:
81 rotated = torch_npu.npu_rotary_mul(
82 sequence_first.reshape(batch_size * seq_length, 1, num_heads, head_dim),
83 cos.reshape(batch_size * seq_length, 1, 1, head_dim),
84 sin.reshape(batch_size * seq_length, 1, 1, head_dim),
85 rotary_mode="interleave",
86 ).reshape(seq_length, batch_size, num_heads, head_dim)
87 else:
88 rotated = torch_npu.npu_rotary_mul(
89 sequence_first.clone(),
90 cos.unsqueeze(-2),
91 sin.unsqueeze(-2),
92 rotary_mode="interleave",
93 )
94 rotated = rotated.permute(1, 2, 0, 3)
95 rotated = torch.cat((rotated[..., 0::2], rotated[..., 1::2]), dim=-1)
96 if unsqueeze_dim == 2:
97 rotated = rotated.transpose(1, 2)
98 return torch.cat((rotated, pass_through), dim=-1)
101def apply_rotary_pos_emb_interleave(
102 q: torch.Tensor,
103 k: torch.Tensor,
104 cos: torch.Tensor,
105 sin: torch.Tensor,
106 position_ids: torch.Tensor | None = None,
107 unsqueeze_dim: int = 1,
108) -> tuple[torch.Tensor, torch.Tensor]:
109 """Apply NPU RoPE using Transformers' interleaved RoPE contract.
111 Args:
112 q: Interleaved-pair query tensor.
113 k: Interleaved-pair key tensor.
114 cos: Transformers ``cat(freqs, freqs)`` cosine tensor.
115 sin: Sine tensor with the same shape and dtype as ``cos``.
116 position_ids: Accepted for Transformers API compatibility.
117 unsqueeze_dim: Dimension used to make ``cos`` and ``sin`` broadcastable
118 to ``q`` and ``k``.
120 Returns:
121 Rotated query and key tensors with the same shapes as the inputs.
123 Raises:
124 ValueError: If ``unsqueeze_dim`` or the rotary dimension is unsupported
125 by the NPU interleave kernel.
126 """
127 del position_ids
128 if unsqueeze_dim not in (1, 2):
129 raise ValueError("NPU interleaved RoPE supports unsqueeze_dim 1 or 2")
130 if cos.ndim == 2:
131 cos = cos.unsqueeze(0)
132 sin = sin.unsqueeze(0)
133 rotary_dim = cos.shape[-1]
134 if rotary_dim % 2:
135 raise ValueError("interleaved RoPE requires an even rotary dimension")
136 if rotary_dim > q.shape[-1] or rotary_dim > k.shape[-1]:
137 raise ValueError("RoPE frequencies cannot exceed the Q/K head dimension")
138 q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
139 k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
141 # Transformers supplies cat(freqs, freqs). The NPU interleave mode expects
142 # each frequency next to the corresponding real/imaginary input pair.
143 cos = cos[..., : rotary_dim // 2].repeat_interleave(2, dim=-1)
144 sin = sin[..., : rotary_dim // 2].repeat_interleave(2, dim=-1)
145 if cos.shape[0] == 1 and q.shape[0] > 1:
146 cos = cos.expand(q.shape[0], -1, -1)
147 sin = sin.expand(q.shape[0], -1, -1)
148 cos = cos.transpose(0, 1)
149 sin = sin.transpose(0, 1)
151 return (
152 _apply_interleaved_rope(q_rot, q_pass, cos, sin, unsqueeze_dim),
153 _apply_interleaved_rope(k_rot, k_pass, cos, sin, unsqueeze_dim),
154 )
157def _rotate_half(
158 x: torch.Tensor,
159 rotary_interleaved: bool = False,
160) -> torch.Tensor:
161 """Rotate the last dimension of ``x`` for rotary position embedding.
163 Args:
164 x: Input tensor whose last dimension is the rotary dimension.
165 rotary_interleaved: If ``True``, rotate adjacent even/odd pairs
166 (interleaved layout). If ``False``, split the last dimension
167 into two halves and rotate across them (half layout).
169 Returns:
170 A tensor with the same shape and dtype as ``x``.
171 """
172 if not rotary_interleaved:
173 x1, x2 = torch.chunk(x, 2, dim=-1)
174 return torch.cat((-x2, x1), dim=-1)
176 dim = x.shape[-1]
177 index1 = np.ones(dim)
178 index1[::2] = 0
179 index2 = np.zeros(dim)
180 index2[::2] = -1
181 rotation_matrix = np.eye(dim, k=1) * index1 + np.eye(dim, k=-1) * index2
182 rotation_matrix = torch.from_numpy(rotation_matrix[None, None, :, :]).to(x.dtype).to(x.device)
183 return torch.matmul(x, rotation_matrix)
186def apply_rotary_pos_emb_single(
187 t: torch.Tensor,
188 cos: torch.Tensor,
189 sin: torch.Tensor,
190 rotary_interleaved: bool = False,
191 use_fused_rotary_pos_emb: bool = False,
192) -> torch.Tensor:
193 """Apply interleaved RoPE and append dimensions that bypass rotation.
195 Args:
196 tensor: Tensor slice to rotate (the rotary part of Q or K).
197 pass_through: Tensor slice that bypasses rotation, concatenated
198 back after the rotated part.
199 cos: Cosine tensor for the interleaved rotary mode.
200 sin: Sine tensor for the interleaved rotary mode.
201 unsqueeze_dim: Dimension used to make ``cos`` and ``sin``
202 broadcastable to ``tensor``. Must be 1 or 2.
204 Returns:
205 A tensor of shape ``..., rotary_dim + pass_through_dim`` formed by
206 concatenating the rotated part and ``pass_through`` along the last
207 dimension.
208 """
209 rot_dim = cos.shape[-1]
210 t_dim = t.shape[-1]
211 t_pass = None
212 if rot_dim != t_dim:
213 t, t_pass = t[..., :rot_dim], t[..., rot_dim:]
215 cos = cos.to(t.dtype)
216 sin = sin.to(t.dtype)
218 while cos.ndim < t.ndim:
219 cos = cos.unsqueeze(-2)
220 sin = sin.unsqueeze(-2)
222 if use_fused_rotary_pos_emb and t.device.type != "cpu":
223 rotary_mode = "interleave" if rotary_interleaved else "half"
224 sequence_length, batch_size, num_heads, head_dim = t.shape
225 if rotary_interleaved and batch_size > 1 and sequence_length > 1:
226 t = torch_npu.npu_rotary_mul(
227 t.reshape(batch_size * sequence_length, 1, num_heads, head_dim),
228 cos.reshape(batch_size * sequence_length, 1, cos.shape[-2], cos.shape[-1]),
229 sin.reshape(batch_size * sequence_length, 1, sin.shape[-2], sin.shape[-1]),
230 rotary_mode=rotary_mode,
231 ).reshape(sequence_length, batch_size, num_heads, head_dim)
232 else:
233 t = torch_npu.npu_rotary_mul(t.clone(), cos, sin, rotary_mode=rotary_mode)
234 else:
235 t = (t * cos) + (_rotate_half(t, rotary_interleaved) * sin)
237 if rot_dim != t_dim:
238 return torch.cat((t, t_pass), dim=-1)
239 return t
242def apply_rotary_pos_emb_single_freqs(
243 t: torch.Tensor,
244 freqs: torch.Tensor,
245 rotary_interleaved: bool = False,
246) -> torch.Tensor:
247 """Apply rotary embeddings from raw frequency values.
249 Args:
250 t: Input tensor to rotate.
251 freqs: Frequency tensor whose last dimension is the rotary
252 dimension. ``cos(freqs)`` and ``sin(freqs)`` are used
253 internally.
254 rotary_interleaved: If ``True``, use interleaved rotary mode;
255 otherwise use half rotary mode.
257 Returns:
258 A tensor with the same shape as ``t``. If ``freqs.shape[-1]`` is
259 smaller than ``t.shape[-1]``, the remaining dimensions of ``t``
260 are passed through unchanged and concatenated back.
261 """
262 rot_dim = freqs.shape[-1]
263 t_dim = t.shape[-1]
264 t_pass = None
265 if rot_dim != t_dim:
266 t, t_pass = t[..., :rot_dim], t[..., rot_dim:]
268 cos_ = torch.cos(freqs).to(t.dtype)
269 sin_ = torch.sin(freqs).to(t.dtype)
271 rotary_mode = "interleave" if rotary_interleaved else "half"
272 t_copy = t.clone()
273 t = torch_npu.npu_rotary_mul(t_copy, cos_, sin_, rotary_mode=rotary_mode)
274 t_copy.untyped_storage().resize_(0)
276 if t_pass is not None:
277 return torch.cat((t, t_pass), dim=-1)
278 return t