Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / rotary_embedding.py: 23%
100 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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 # Transpose pairs instead of strided slices to avoid scattered writes in backward.
96 rotated = rotated.unflatten(-1, (-1, 2)).transpose(-1, -2).flatten(-2)
97 if unsqueeze_dim == 2:
98 rotated = rotated.transpose(1, 2)
99 return torch.cat((rotated, pass_through), dim=-1) if pass_through.shape[-1] else rotated
102def apply_rotary_pos_emb_interleave(
103 q: torch.Tensor,
104 k: torch.Tensor,
105 cos: torch.Tensor,
106 sin: torch.Tensor,
107 position_ids: torch.Tensor | None = None,
108 unsqueeze_dim: int = 1,
109) -> tuple[torch.Tensor, torch.Tensor]:
110 """Apply NPU RoPE using Transformers' interleaved RoPE contract.
112 Args:
113 q: Interleaved-pair query tensor.
114 k: Interleaved-pair key tensor.
115 cos: Transformers ``cat(freqs, freqs)`` cosine tensor.
116 sin: Sine tensor with the same shape and dtype as ``cos``.
117 position_ids: Accepted for Transformers API compatibility.
118 unsqueeze_dim: Dimension used to make ``cos`` and ``sin`` broadcastable
119 to ``q`` and ``k``.
121 Returns:
122 Rotated query and key tensors with the same shapes as the inputs.
124 Raises:
125 ValueError: If ``unsqueeze_dim`` or the rotary dimension is unsupported
126 by the NPU interleave kernel.
127 """
128 del position_ids
129 if unsqueeze_dim not in (1, 2):
130 raise ValueError("NPU interleaved RoPE supports unsqueeze_dim 1 or 2")
131 if cos.ndim == 2:
132 cos = cos.unsqueeze(0)
133 sin = sin.unsqueeze(0)
134 rotary_dim = cos.shape[-1]
135 if rotary_dim % 2:
136 raise ValueError("interleaved RoPE requires an even rotary dimension")
137 if rotary_dim > q.shape[-1] or rotary_dim > k.shape[-1]:
138 raise ValueError("RoPE frequencies cannot exceed the Q/K head dimension")
139 q_rot, q_pass = q.split((rotary_dim, q.shape[-1] - rotary_dim), dim=-1)
140 k_rot, k_pass = k.split((rotary_dim, k.shape[-1] - rotary_dim), dim=-1)
142 # Transformers supplies cat(freqs, freqs). The NPU interleave mode expects
143 # each frequency next to the corresponding real/imaginary input pair.
144 cos = cos[..., : rotary_dim // 2].repeat_interleave(2, dim=-1)
145 sin = sin[..., : rotary_dim // 2].repeat_interleave(2, dim=-1)
146 if cos.shape[0] == 1 and q.shape[0] > 1:
147 cos = cos.expand(q.shape[0], -1, -1)
148 sin = sin.expand(q.shape[0], -1, -1)
149 cos = cos.transpose(0, 1)
150 sin = sin.transpose(0, 1)
152 return (
153 _apply_interleaved_rope(q_rot, q_pass, cos, sin, unsqueeze_dim),
154 _apply_interleaved_rope(k_rot, k_pass, cos, sin, unsqueeze_dim),
155 )
158def _rotate_half(
159 x: torch.Tensor,
160 rotary_interleaved: bool = False,
161) -> torch.Tensor:
162 """Rotate the last dimension of ``x`` for rotary position embedding.
164 Args:
165 x: Input tensor whose last dimension is the rotary dimension.
166 rotary_interleaved: If ``True``, rotate adjacent even/odd pairs
167 (interleaved layout). If ``False``, split the last dimension
168 into two halves and rotate across them (half layout).
170 Returns:
171 A tensor with the same shape and dtype as ``x``.
172 """
173 if not rotary_interleaved:
174 x1, x2 = torch.chunk(x, 2, dim=-1)
175 return torch.cat((-x2, x1), dim=-1)
177 dim = x.shape[-1]
178 index1 = np.ones(dim)
179 index1[::2] = 0
180 index2 = np.zeros(dim)
181 index2[::2] = -1
182 rotation_matrix = np.eye(dim, k=1) * index1 + np.eye(dim, k=-1) * index2
183 rotation_matrix = torch.from_numpy(rotation_matrix[None, None, :, :]).to(x.dtype).to(x.device)
184 return torch.matmul(x, rotation_matrix)
187def apply_rotary_pos_emb_single(
188 t: torch.Tensor,
189 cos: torch.Tensor,
190 sin: torch.Tensor,
191 rotary_interleaved: bool = False,
192 use_fused_rotary_pos_emb: bool = False,
193) -> torch.Tensor:
194 """Apply interleaved RoPE and append dimensions that bypass rotation.
196 Args:
197 tensor: Tensor slice to rotate (the rotary part of Q or K).
198 pass_through: Tensor slice that bypasses rotation, concatenated
199 back after the rotated part.
200 cos: Cosine tensor for the interleaved rotary mode.
201 sin: Sine tensor for the interleaved rotary mode.
202 unsqueeze_dim: Dimension used to make ``cos`` and ``sin``
203 broadcastable to ``tensor``. Must be 1 or 2.
205 Returns:
206 A tensor of shape ``..., rotary_dim + pass_through_dim`` formed by
207 concatenating the rotated part and ``pass_through`` along the last
208 dimension.
209 """
210 rot_dim = cos.shape[-1]
211 t_dim = t.shape[-1]
212 t_pass = None
213 if rot_dim != t_dim:
214 t, t_pass = t[..., :rot_dim], t[..., rot_dim:]
216 cos = cos.to(t.dtype)
217 sin = sin.to(t.dtype)
219 while cos.ndim < t.ndim:
220 cos = cos.unsqueeze(-2)
221 sin = sin.unsqueeze(-2)
223 if use_fused_rotary_pos_emb and t.device.type != "cpu":
224 rotary_mode = "interleave" if rotary_interleaved else "half"
225 sequence_length, batch_size, num_heads, head_dim = t.shape
226 if rotary_interleaved and batch_size > 1 and sequence_length > 1:
227 t = torch_npu.npu_rotary_mul(
228 t.reshape(batch_size * sequence_length, 1, num_heads, head_dim),
229 cos.reshape(batch_size * sequence_length, 1, cos.shape[-2], cos.shape[-1]),
230 sin.reshape(batch_size * sequence_length, 1, sin.shape[-2], sin.shape[-1]),
231 rotary_mode=rotary_mode,
232 ).reshape(sequence_length, batch_size, num_heads, head_dim)
233 else:
234 t = torch_npu.npu_rotary_mul(t.clone(), cos, sin, rotary_mode=rotary_mode)
235 else:
236 t = (t * cos) + (_rotate_half(t, rotary_interleaved) * sin)
238 if rot_dim != t_dim:
239 return torch.cat((t, t_pass), dim=-1)
240 return t
243def apply_rotary_pos_emb_single_freqs(
244 t: torch.Tensor,
245 freqs: torch.Tensor,
246 rotary_interleaved: bool = False,
247) -> torch.Tensor:
248 """Apply rotary embeddings from raw frequency values.
250 Args:
251 t: Input tensor to rotate.
252 freqs: Frequency tensor whose last dimension is the rotary
253 dimension. ``cos(freqs)`` and ``sin(freqs)`` are used
254 internally.
255 rotary_interleaved: If ``True``, use interleaved rotary mode;
256 otherwise use half rotary mode.
258 Returns:
259 A tensor with the same shape as ``t``. If ``freqs.shape[-1]`` is
260 smaller than ``t.shape[-1]``, the remaining dimensions of ``t``
261 are passed through unchanged and concatenated back.
262 """
263 rot_dim = freqs.shape[-1]
264 t_dim = t.shape[-1]
265 t_pass = None
266 if rot_dim != t_dim:
267 t, t_pass = t[..., :rot_dim], t[..., rot_dim:]
269 cos_ = torch.cos(freqs).to(t.dtype)
270 sin_ = torch.sin(freqs).to(t.dtype)
272 rotary_mode = "interleave" if rotary_interleaved else "half"
273 t_copy = t.clone()
274 t = torch_npu.npu_rotary_mul(t_copy, cos_, sin_, rotary_mode=rotary_mode)
275 t_copy.untyped_storage().resize_(0)
277 if t_pass is not None:
278 return torch.cat((t, t_pass), dim=-1)
279 return t