[Kernel][Kimi] fused vision q/k roper kernel (#50400)
Signed-off-by: rongfu.leng <lenronfu@gmail.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
rongfu.leng committed
Aug 20, 2026 at 17:15 UTC
c8de519917ce549f72132952116185e38b37c95d
1 file changed
+14
-26
vllm/model_executor/models/kimi_k25_vit.py
+14
-26
@@ -29,6 +29,7 @@ from vllm.model_executor.layers.linear import (
29
RowParallelLinear,
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.models.utils import maybe_prefix
34
from vllm.model_executor.models.vision import (
35
is_vit_use_data_parallel,
@@ -79,29 +80,6 @@ def get_rope_shape(org, interpolation_mode, shape):
80
)
81
82
82
-def apply_rope(
83
- xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor
84
-) -> tuple[torch.Tensor, torch.Tensor]:
85
- """
86
- Args: (The leading dimensions of all inputs should be the same)
87
- xq: query, tensor of shape (..., num_heads, head_dim)
88
- xk: key, tensor of shape (..., num_heads, head_dim)
89
- freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64.
90
- Returns:
91
- xq_out, xk_out: tensors of shape (..., num_heads, head_dim)
92
- """
93
- _apply_rope_input_validation(xq, freqs_cis)
94
- _apply_rope_input_validation(xk, freqs_cis)
95
-
96
- freqs_cis = freqs_cis.unsqueeze(-2) # ..., 1, head_dim/2
97
- # ..., num_heads, head_dim/2
98
- xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2))
99
- xk_ = torch.view_as_complex(xk.float().view(*xq.shape[:-1], -1, 2))
100
- xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(-2) # ..., num_heads, head_dim
101
- xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(-2) # ..., num_heads, head_dim
102
- return xq_out.type_as(xq), xk_out.type_as(xk)
103
-
104
-
83
def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
84
"""Generate 1D sincos positional embedding from grid positions."""
85
assert embed_dim % 2 == 0
@@ -451,12 +429,17 @@ class MoonViTEncoderLayer(nn.Module):
429
scale=self.hidden_size_per_attention_head**-0.5,
430
prefix=f"{prefix}.attn",
431
)
432
+ self.apply_rotary_emb = ApplyRotaryEmb(
433
+ enforce_enable=True,
434
+ is_neox_style=False,
435
+ enable_fp32_compute=True,
436
+ )
437
438
def attention_qkvpacked(
439
self,
440
x: torch.Tensor,
441
cu_seqlens: torch.Tensor,
459
- rope_freqs_cis: torch.Tensor | None = None,
442
+ rope_freqs_cis: torch.Tensor,
443
max_seqlen: torch.Tensor | None = None,
444
sequence_lengths: torch.Tensor | None = None,
445
):
@@ -478,7 +461,12 @@ class MoonViTEncoderLayer(nn.Module):
461
xqkv = xqkv.view(*qkv_shape)
462
xq, xk, xv = torch.unbind(xqkv, dim=-3)
463
481
- xq, xk = apply_rope(xq, xk, rope_freqs_cis)
464
+ _apply_rope_input_validation(xq, rope_freqs_cis)
465
+ _apply_rope_input_validation(xk, rope_freqs_cis)
466
+ rope_cos = rope_freqs_cis.real.contiguous()
467
+ rope_sin = rope_freqs_cis.imag.contiguous()
468
+ xq = self.apply_rotary_emb(xq, rope_cos, rope_sin)
469
+ xk = self.apply_rotary_emb(xk, rope_cos, rope_sin)
470
471
if max_seqlen is None:
472
max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
@@ -502,7 +490,7 @@ class MoonViTEncoderLayer(nn.Module):
490
self,
491
hidden_states: torch.Tensor,
492
cu_seqlens: torch.Tensor,
505
- rope_freqs_cis: torch.Tensor | None = None,
493
+ rope_freqs_cis: torch.Tensor,
494
max_seqlen: torch.Tensor | None = None,
495
sequence_lengths: torch.Tensor | None = None,
496
):