[KimiViT][Perf] Fuse per-layer QK RoPE into one in-place kernel (#58651)
Thien Tran committed
Sep 28, 2026 at 09:18 UTC
39d49539de86a9fe6bd556a7321474b53d0f7a63
3 files changed
+140
-7
tests/kernels/core/test_apply_rotary_emb.py
+43
@@ -225,3 +225,46 @@ def test_apply_rotary_emb_xpu_matches_native(
225
expected = op.forward_native(x, cos, sin)
226
actual = op.forward_xpu(x, cos, sin)
227
torch.testing.assert_close(actual, expected, atol=0, rtol=0)
228
+
229
+
230
+@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="CUDA/ROCm only test.")
231
+@pytest.mark.parametrize("num_tokens", [257, 2304])
232
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
233
+def test_packed_qk_rope_correctness(
234
+ num_tokens: int, dtype: torch.dtype, default_vllm_config
235
+):
236
+ """packed_qk_rope_ must be bitwise identical to the per-tensor
237
+ ApplyRotaryEmb path with enable_fp32_compute=True, and must leave the V
238
+ slice untouched."""
239
+ from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb
240
+ from vllm.model_executor.layers.rotary_embedding.packed_qk_rope import (
241
+ packed_qk_rope_,
242
+ )
243
+
244
+ num_heads, head_dim = 12, 128
245
+ rng = torch.Generator(device="cuda").manual_seed(0)
246
+
247
+ angles = (
248
+ torch.rand(num_tokens, head_dim // 2, device="cuda", generator=rng) * torch.pi
249
+ )
250
+ freqs_cis = torch.polar(torch.ones_like(angles), angles)
251
+ cos = freqs_cis.real.contiguous()
252
+ sin = freqs_cis.imag.contiguous()
253
+
254
+ xqkv = torch.randn(
255
+ num_tokens, 3, num_heads, head_dim, dtype=dtype, device="cuda", generator=rng
256
+ )
257
+ xqkv_ref = xqkv.clone()
258
+ xq, xk, xv = torch.unbind(xqkv_ref, dim=-3)
259
+
260
+ op = ApplyRotaryEmb(
261
+ enforce_enable=True, is_neox_style=False, enable_fp32_compute=True
262
+ )
263
+ # op() dispatches to forward_cuda / forward_hip per platform.
264
+ xq_ref = op(xq, cos, sin)
265
+ xk_ref = op(xk, cos, sin)
266
+
267
+ packed_qk_rope_(xqkv, freqs_cis)
268
+ torch.testing.assert_close(xqkv[:, 0], xq_ref, atol=0, rtol=0)
269
+ torch.testing.assert_close(xqkv[:, 1], xk_ref, atol=0, rtol=0)
270
+ torch.testing.assert_close(xqkv[:, 2], xv, atol=0, rtol=0)
vllm/model_executor/layers/rotary_embedding/packed_qk_rope.py
new
+81
@@ -0,0 +1,81 @@
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
+ Bitwise identical to ``ApplyRotaryEmb(enable_fp32_compute=True)`` applied
50
+ to Q and K separately, but in one kernel with ~5x less memory traffic.
51
+ Requires triton: callers must check HAS_TRITON and use the unfused path
52
+ otherwise. Contract violations raise.
53
+
54
+ """
55
+ assert xqkv.ndim == 4 and xqkv.size(1) == 3, xqkv.shape
56
+ assert xqkv.is_contiguous(), xqkv.shape
57
+ assert freqs_cis.ndim == 2 and freqs_cis.is_contiguous(), freqs_cis.shape
58
+ assert freqs_cis.dtype == torch.complex64, freqs_cis.dtype
59
+
60
+ seq_length, _, num_heads, head_dim = xqkv.shape
61
+ rotary_dim = freqs_cis.size(-1) * 2
62
+ assert freqs_cis.size(0) == seq_length, (freqs_cis.shape, xqkv.shape)
63
+ assert rotary_dim == head_dim, (freqs_cis.shape, xqkv.shape) # full rotation
64
+ qk = xqkv.as_strided(
65
+ (seq_length, 2 * num_heads, head_dim),
66
+ (3 * num_heads * head_dim, head_dim, 1),
67
+ )
68
+
69
+ BLOCK_M = 4
70
+ # seqlen on grid.x: HIP caps gridDim.y at 65535.
71
+ grid = (triton.cdiv(seq_length, BLOCK_M), 2 * num_heads)
72
+ _packed_qk_rope_kernel[grid](
73
+ qk,
74
+ freqs_cis.view(torch.float32),
75
+ seq_length,
76
+ qk.stride(0),
77
+ qk.stride(1),
78
+ rotary_dim=rotary_dim,
79
+ BLOCK_M=BLOCK_M,
80
+ num_warps=2 if rotary_dim <= 64 else 4,
81
+ )
vllm/model_executor/models/kimi_k25_vit.py
+16
-7
@@ -30,6 +30,9 @@ from vllm.model_executor.layers.linear import (
30
)
31
from vllm.model_executor.layers.quantization import QuantizationConfig
32
from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb
33
+from vllm.model_executor.layers.rotary_embedding.packed_qk_rope import (
34
+ packed_qk_rope_,
35
+)
36
from vllm.model_executor.models.utils import maybe_prefix
37
from vllm.model_executor.models.vision import (
38
is_vit_use_data_parallel,
@@ -37,6 +40,7 @@ from vllm.model_executor.models.vision import (
40
)
41
from vllm.platforms import current_platform
42
from vllm.transformers_utils.configs.kimi_k25 import KimiK25VisionConfig
43
+from vllm.triton_utils import HAS_TRITON
44
from vllm.utils.torch_utils import async_tensor_h2d
45
46
logger = init_logger(__name__)
@@ -451,14 +455,19 @@ class MoonViTEncoderLayer(nn.Module):
455
)
456
# xqkv: (seqlen, 3, nheads, headdim)
457
xqkv = xqkv.view(*qkv_shape)
454
- xq, xk, xv = torch.unbind(xqkv, dim=-3)
458
456
- _apply_rope_input_validation(xq, rope_freqs_cis)
457
- _apply_rope_input_validation(xk, rope_freqs_cis)
458
- rope_cos = rope_freqs_cis.real.contiguous()
459
- rope_sin = rope_freqs_cis.imag.contiguous()
460
- xq = self.apply_rotary_emb(xq, rope_cos, rope_sin)
461
- xk = self.apply_rotary_emb(xk, rope_cos, rope_sin)
459
+ # xq/xk alias xqkv, so the in-place rotation below is visible through
460
+ # them; only the fallback needs to rebind.
461
+ xq, xk, xv = torch.unbind(xqkv, dim=-3)
462
+ if HAS_TRITON:
463
+ packed_qk_rope_(xqkv, rope_freqs_cis)
464
+ else:
465
+ _apply_rope_input_validation(xq, rope_freqs_cis)
466
+ _apply_rope_input_validation(xk, rope_freqs_cis)
467
+ rope_cos = rope_freqs_cis.real.contiguous()
468
+ rope_sin = rope_freqs_cis.imag.contiguous()
469
+ xq = self.apply_rotary_emb(xq, rope_cos, rope_sin)
470
+ xk = self.apply_rotary_emb(xk, rope_cos, rope_sin)
471
472
if max_seqlen is None:
473
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()