| 209 | |
| 210 | @flags.compile_wrap |
| 211 | def _apply_rotary_emb_inplace(x, theta, conj): |
| 212 | dtype = reduce(torch.promote_types, (x.dtype, theta.dtype, torch.float32)) |
| 213 | d = theta.shape[-1] |
| 214 | assert d * 2 <= x.shape[-1] |
| 215 | x1, x2 = x[..., :d], x[..., d : d * 2] |
| 216 | x1_, x2_, theta = x1.to(dtype), x2.to(dtype), theta.to(dtype) |
| 217 | cos, sin = torch.cos(theta), torch.sin(theta) |
| 218 | sin = -sin if conj else sin |
| 219 | y1 = x1_ * cos - x2_ * sin |
| 220 | y2 = x2_ * cos + x1_ * sin |
| 221 | x1.copy_(y1) |
| 222 | x2.copy_(y2) |
| 223 | |
| 224 | |
| 225 | class ApplyRotaryEmbeddingInplace(torch.autograd.Function): |