| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """In-place interleaved RoPE on the Q/K slices of a packed QKV buffer.""" |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.triton_utils import tl, triton |
| 8 | |
| 9 | |
| 10 | # seqlen is not specialized: image sizes vary per request. |
| 11 | @triton.jit(do_not_specialize=["seqlen"]) |
| 12 | def _packed_qk_rope_kernel( |
| 13 | qk_ptr, # (seqlen, 2 * nheads, headdim) strided view of the packed QKV |
| 14 | freqs_ptr, # (seqlen, rotary_dim) fp32 re/im view of the complex freqs |
| 15 | seqlen, |
| 16 | qk_row_stride, |
| 17 | qk_head_stride, |
| 18 | rotary_dim: tl.constexpr, |
| 19 | BLOCK_M: tl.constexpr, |
| 20 | ): |
| 21 | BLOCK_K: tl.constexpr = triton.next_power_of_2(rotary_dim) |
| 22 | pid_m = tl.program_id(0) |
| 23 | pid_h = tl.program_id(1) |
| 24 | rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) |
| 25 | rk = tl.arange(0, BLOCK_K) |
| 26 | mask_m = rm < seqlen |
| 27 | |
| 28 | base = qk_ptr + pid_h * qk_head_stride + rm[:, None] * qk_row_stride |
| 29 | mask_x = mask_m[:, None] & (rk < rotary_dim)[None, :] |
| 30 | x = tl.load(base + rk[None, :], mask=mask_x, other=0.0).to(tl.float32) |
| 31 | f = tl.load(freqs_ptr + rm[:, None] * rotary_dim + rk[None, :], mask=mask_x) |
| 32 | cos, sin = tl.split(tl.reshape(f, (BLOCK_M, BLOCK_K // 2, 2))) |
| 33 | |
| 34 | x0, x1 = tl.split(tl.reshape(x, (BLOCK_M, BLOCK_K // 2, 2))) |
| 35 | o0 = x0 * cos - x1 * sin |
| 36 | o1 = x1 * cos + x0 * sin |
| 37 | tl.store(base + rk[None, :], tl.interleave(o0, o1), mask=mask_x) |
| 38 | |
| 39 | |
| 40 | def packed_qk_rope_(xqkv: torch.Tensor, freqs_cis: torch.Tensor) -> None: |
| 41 | """Rotate Q and K in place inside a packed QKV buffer. |
| 42 | |
| 43 | Args: |
| 44 | xqkv: contiguous (seqlen, 3, nheads, headdim); only the Q and K |
| 45 | slices are rotated. |
| 46 | freqs_cis: contiguous (seqlen, headdim // 2) complex64 rotary freqs, |
| 47 | read through its interleaved re/im fp32 view. |
| 48 | |
| 49 | Matches ``ApplyRotaryEmb(enable_fp32_compute=True)`` applied to Q and K |
| 50 | separately to within one ULP, but in one kernel with ~5x less memory |
| 51 | traffic. Both compute in fp32; the gap is which product the compiler |
| 52 | keeps exact inside the fused multiply-add, which differs per backend. |
| 53 | Requires triton: callers must check HAS_TRITON and use the unfused path |
| 54 | otherwise. Contract violations raise. |
| 55 | |
| 56 | """ |
| 57 | assert xqkv.ndim == 4 and xqkv.size(1) == 3, xqkv.shape |
| 58 | assert xqkv.is_contiguous(), xqkv.shape |
| 59 | assert freqs_cis.ndim == 2 and freqs_cis.is_contiguous(), freqs_cis.shape |
| 60 | assert freqs_cis.dtype == torch.complex64, freqs_cis.dtype |
| 61 | |
| 62 | seq_length, _, num_heads, head_dim = xqkv.shape |
| 63 | rotary_dim = freqs_cis.size(-1) * 2 |
| 64 | assert freqs_cis.size(0) == seq_length, (freqs_cis.shape, xqkv.shape) |
| 65 | assert rotary_dim == head_dim, (freqs_cis.shape, xqkv.shape) # full rotation |
| 66 | qk = xqkv.as_strided( |
| 67 | (seq_length, 2 * num_heads, head_dim), |
| 68 | (3 * num_heads * head_dim, head_dim, 1), |
| 69 | ) |
| 70 | |
| 71 | BLOCK_M = 4 |
| 72 | # seqlen on grid.x: HIP caps gridDim.y at 65535. |
| 73 | grid = (triton.cdiv(seq_length, BLOCK_M), 2 * num_heads) |
| 74 | _packed_qk_rope_kernel[grid]( |
| 75 | qk, |
| 76 | freqs_cis.view(torch.float32), |
| 77 | seq_length, |
| 78 | qk.stride(0), |
| 79 | qk.stride(1), |
| 80 | rotary_dim=rotary_dim, |
| 81 | BLOCK_M=BLOCK_M, |
| 82 | num_warps=2 if rotary_dim <= 64 else 4, |
| 83 | ) |