A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 09:00 UTC
Fuse RoPE and MLA KV-cache write (#25774)
Signed-off-by: Patryk Saffer <patryk.saffer99@gmail.com> Signed-off-by: PatrykSaffer <patryk.saffer@mistral.ai> Co-authored-by: Patryk Saffer <patryk.saffer99@gmail.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
PatrykSaffer committed
Jan 10, 2026 at 04:18 UTC
80fead8bf6508b7907a565d05cb0a6aed9101a3d
6 files changed
+490
CMakeLists.txt
+1
@@ -282,6 +282,7 @@ endif()
282
set(VLLM_EXT_SRC
283
"csrc/mamba/mamba_ssm/selective_scan_fwd.cu"
284
"csrc/cache_kernels.cu"
285
+ "csrc/cache_kernels_fused.cu"
286
"csrc/attention/paged_attention_v1.cu"
287
"csrc/attention/paged_attention_v2.cu"
288
"csrc/attention/merge_attn_states.cu"
csrc/cache.h
+7
@@ -27,6 +27,13 @@ void concat_and_cache_mla(torch::Tensor& kv_c, torch::Tensor& k_pe,
27
const std::string& kv_cache_dtype,
28
torch::Tensor& scale);
29
30
+// NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
31
+void concat_and_cache_mla_rope_fused(
32
+ torch::Tensor& positions, torch::Tensor& q_pe, torch::Tensor& k_pe,
33
+ torch::Tensor& kv_c, torch::Tensor& rope_cos_sin_cache, bool rope_is_neox,
34
+ torch::Tensor& kv_cache_slot_mapping, torch::Tensor& kv_cache,
35
+ const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale);
36
+
37
// Just for unittest
38
void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,
39
const double scale, const std::string& kv_cache_dtype);
csrc/cache_kernels_fused.cu
new
+279
@@ -0,0 +1,279 @@
1
+#include <torch/all.h>
2
+#include <ATen/cuda/CUDAContext.h>
3
+#include <c10/cuda/CUDAGuard.h>
4
+
5
+#include "cuda_compat.h"
6
+#include "dispatch_utils.h"
7
+
8
+#include "quantization/w8a8/fp8/common.cuh"
9
+#ifdef USE_ROCM
10
+ #include "quantization/w8a8/fp8/amd/quant_utils.cuh"
11
+#else
12
+ #include "quantization/w8a8/fp8/nvidia/quant_utils.cuh"
13
+#endif
14
+
15
+#ifdef USE_ROCM
16
+ #include <hip/hip_bf16.h>
17
+typedef __hip_bfloat16 __nv_bfloat16;
18
+#endif
19
+
20
+namespace vllm {
21
+
22
+// NOTE Be EXTRA careful with raw_kv_scalar_t, for __half and __nv_bfloat16 it's
23
+// using u16 as the backing type.
24
+template <typename qk_t, bool IS_NEOX, typename raw_kv_scalar_t,
25
+ typename cache_t, Fp8KVCacheDataType kv_dt>
26
+__global__ void concat_and_cache_mla_rope_fused_kernel(
27
+ const int64_t* __restrict__ positions, // [num_tokens]
28
+ qk_t* __restrict__ q_pe, // [num_tokens, num_q_heads, rot_dim]
29
+ qk_t* __restrict__ k_pe, // [num_tokens, rot_dim]
30
+ const qk_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank]
31
+ const qk_t* __restrict__ rope_cos_sin_cache, // [max_position, 2,
32
+ // rot_dim // 2]
33
+ const int rot_dim, const int64_t q_pe_stride_token,
34
+ const int64_t q_pe_stride_head, const int64_t k_pe_stride,
35
+ const int64_t kv_c_stride, const int num_q_heads,
36
+ cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank +
37
+ // rot_dim)]
38
+ const int64_t* __restrict__ kv_cache_slot_mapping, // [num_tokens]
39
+ const int block_stride, const int entry_stride, const int kv_lora_rank,
40
+ const int block_size, const float* kv_cache_quant_scale) {
41
+ // Each thread block is responsible for one token.
42
+ const int64_t token_idx = blockIdx.x;
43
+ const int64_t pos = positions[token_idx];
44
+
45
+ const qk_t* cos_sin_ptr = rope_cos_sin_cache + pos * rot_dim;
46
+
47
+ const int embed_dim = rot_dim / 2;
48
+
49
+ // Q ROPE
50
+ const int nq = num_q_heads * embed_dim;
51
+ for (int i = threadIdx.x; i < nq; i += blockDim.x) {
52
+ int head_idx = i / embed_dim;
53
+ int pair_idx = i % embed_dim;
54
+
55
+ // NOTE: Would be nice to have interleaved sin/cos so we could just load
56
+ // both at the same time.
57
+ qk_t cos = VLLM_LDG(cos_sin_ptr + pair_idx);
58
+ qk_t sin = VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim);
59
+
60
+ qk_t* q_pe_head_ptr =
61
+ q_pe + token_idx * q_pe_stride_token + head_idx * q_pe_stride_head;
62
+
63
+ int pair_idx_x, pair_idx_y;
64
+ if constexpr (IS_NEOX) {
65
+ // GPT-NeoX style rotary embedding.
66
+ pair_idx_x = pair_idx;
67
+ pair_idx_y = embed_dim + pair_idx;
68
+ } else {
69
+ // GPT-J style rotary embedding.
70
+ pair_idx_x = pair_idx * 2;
71
+ pair_idx_y = pair_idx * 2 + 1;
72
+ }
73
+
74
+ qk_t x_src = q_pe_head_ptr[pair_idx_x];
75
+ qk_t y_src = q_pe_head_ptr[pair_idx_y];
76
+
77
+ qk_t x_dst = x_src * cos - y_src * sin;
78
+ qk_t y_dst = y_src * cos + x_src * sin;
79
+
80
+ q_pe_head_ptr[pair_idx_x] = x_dst;
81
+ q_pe_head_ptr[pair_idx_y] = y_dst;
82
+ }
83
+
84
+ const int64_t slot_idx = kv_cache_slot_mapping[token_idx];
85
+ const int64_t block_idx = slot_idx / block_size;
86
+ const int64_t entry_idx = slot_idx % block_size;
87
+
88
+ // NOTE: slot_idx can be -1 if the token is padded
89
+ if (slot_idx < 0) {
90
+ return;
91
+ }
92
+
93
+ // K with 1 HEAD
94
+ for (int i = threadIdx.x; i < embed_dim; i += blockDim.x) {
95
+ int pair_idx = i;
96
+
97
+ qk_t cos = VLLM_LDG(cos_sin_ptr + pair_idx);
98
+ qk_t sin = VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim);
99
+
100
+ qk_t* k_pe_head_ptr = k_pe + token_idx * k_pe_stride;
101
+
102
+ int pair_idx_x, pair_idx_y;
103
+ if constexpr (IS_NEOX) {
104
+ // GPT-NeoX style rotary embedding.
105
+ pair_idx_x = pair_idx;
106
+ pair_idx_y = embed_dim + pair_idx;
107
+ } else {
108
+ // GPT-J style rotary embedding.
109
+ pair_idx_x = pair_idx * 2;
110
+ pair_idx_y = pair_idx * 2 + 1;
111
+ }
112
+
113
+ qk_t x_src = k_pe_head_ptr[pair_idx_x];
114
+ qk_t y_src = k_pe_head_ptr[pair_idx_y];
115
+
116
+ qk_t x_dst = x_src * cos - y_src * sin;
117
+ qk_t y_dst = y_src * cos + x_src * sin;
118
+
119
+ k_pe_head_ptr[pair_idx_x] = x_dst;
120
+ k_pe_head_ptr[pair_idx_y] = y_dst;
121
+
122
+ // NOTE Why is this monster necessary?
123
+ // When K is of type float16, the actual template replacement for
124
+ // raw_kv_scalar_t with be u16. That's why it's used at the last moment
125
+ // otherwise CUDA ALU would break.
126
+ const raw_kv_scalar_t raw_x_value =
127
+ *reinterpret_cast<const raw_kv_scalar_t*>(&x_dst);
128
+ const raw_kv_scalar_t raw_y_value =
129
+ *reinterpret_cast<const raw_kv_scalar_t*>(&y_dst);
130
+
131
+ cache_t* kv_cache_ptr = kv_cache + block_idx * block_stride +
132
+ entry_idx * entry_stride + kv_lora_rank;
133
+
134
+ // MLA Cache Store
135
+ if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
136
+ kv_cache_ptr[pair_idx_x] = raw_x_value;
137
+ kv_cache_ptr[pair_idx_y] = raw_y_value;
138
+ } else {
139
+ kv_cache_ptr[pair_idx_x] =
140
+ fp8::scaled_convert<cache_t, raw_kv_scalar_t, kv_dt>(
141
+ raw_x_value, *kv_cache_quant_scale);
142
+ kv_cache_ptr[pair_idx_y] =
143
+ fp8::scaled_convert<cache_t, raw_kv_scalar_t, kv_dt>(
144
+ raw_y_value, *kv_cache_quant_scale);
145
+ }
146
+ }
147
+
148
+ // NOPE
149
+ for (int i = threadIdx.x; i < kv_lora_rank; i += blockDim.x) {
150
+ const qk_t* src_ptr = kv_c + token_idx * kv_c_stride + i;
151
+ const raw_kv_scalar_t src_value =
152
+ *reinterpret_cast<const raw_kv_scalar_t*>(src_ptr);
153
+
154
+ cache_t* kv_cache_ptr =
155
+ kv_cache + block_idx * block_stride + entry_idx * entry_stride;
156
+
157
+ if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) {
158
+ kv_cache_ptr[i] = src_value;
159
+ } else {
160
+ kv_cache_ptr[i] = fp8::scaled_convert<cache_t, raw_kv_scalar_t, kv_dt>(
161
+ src_value, *kv_cache_quant_scale);
162
+ }
163
+ }
164
+}
165
+
166
+} // namespace vllm
167
+
168
+#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \
169
+ do { \
170
+ VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] { \
171
+ using qk_t = scalar_t; \
172
+ if (rope_is_neox) { \
173
+ vllm::concat_and_cache_mla_rope_fused_kernel<qk_t, true, RAW_KV_T, \
174
+ CACHE_T, KV_DTYPE> \
175
+ <<<grid, block, 0, stream>>>( \
176
+ positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
177
+ k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
178
+ rope_cos_sin_cache.data_ptr<qk_t>(), rot_dim, \
179
+ q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \
180
+ num_q_heads, reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
181
+ kv_cache_slot_mapping.data_ptr<int64_t>(), block_stride, \
182
+ entry_stride, kv_lora_rank, block_size, \
183
+ kv_cache_quant_scale.data_ptr<float>()); \
184
+ } else { \
185
+ vllm::concat_and_cache_mla_rope_fused_kernel<qk_t, false, RAW_KV_T, \
186
+ CACHE_T, KV_DTYPE> \
187
+ <<<grid, block, 0, stream>>>( \
188
+ positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
189
+ k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
190
+ rope_cos_sin_cache.data_ptr<qk_t>(), rot_dim, \
191
+ q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \
192
+ num_q_heads, reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
193
+ kv_cache_slot_mapping.data_ptr<int64_t>(), block_stride, \
194
+ entry_stride, kv_lora_rank, block_size, \
195
+ kv_cache_quant_scale.data_ptr<float>()); \
196
+ } \
197
+ }); \
198
+ } while (false)
199
+
200
+// Executes RoPE on q_pe and k_pe, then writes k_pe and kv_c in the kv cache.
201
+// q_pe and k_pe are modified in place.
202
+// Replaces DeepseekScalingRotaryEmbedding.self.rotary_emb and
203
+// concat_and_cache_mla.
204
+void concat_and_cache_mla_rope_fused(
205
+ torch::Tensor& positions, // [num_tokens]
206
+ torch::Tensor& q_pe, // [num_tokens, num_q_heads, rot_dim]
207
+ torch::Tensor& k_pe, // [num_tokens, rot_dim]
208
+ torch::Tensor& kv_c, // [num_tokens, kv_lora_rank]
209
+ torch::Tensor& rope_cos_sin_cache, // [max_position, rot_dim]
210
+ bool rope_is_neox,
211
+ torch::Tensor&
212
+ kv_cache_slot_mapping, // [num_tokens] or [num_actual_tokens]
213
+ torch::Tensor&
214
+ kv_cache, // [num_blocks, block_size, (kv_lora_rank + rot_dim)]
215
+ const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale) {
216
+ const int64_t num_tokens = q_pe.size(0);
217
+
218
+ const int num_q_heads = q_pe.size(1);
219
+ const int rot_dim = q_pe.size(2);
220
+ const int kv_lora_rank = kv_c.size(1);
221
+
222
+ TORCH_CHECK(positions.size(0) >=
223
+ num_tokens); // CUDA Graphs might pad this for us
224
+ TORCH_CHECK_EQ(positions.dim(), 1);
225
+ TORCH_CHECK_EQ(positions.scalar_type(), c10::ScalarType::Long);
226
+
227
+ TORCH_CHECK_EQ(q_pe.size(0), num_tokens);
228
+ TORCH_CHECK_EQ(q_pe.size(1), num_q_heads);
229
+ TORCH_CHECK_EQ(q_pe.size(2), rot_dim);
230
+ TORCH_CHECK_EQ(q_pe.dim(), 3);
231
+
232
+ TORCH_CHECK_EQ(k_pe.size(0), num_tokens);
233
+ TORCH_CHECK_EQ(k_pe.size(1), rot_dim);
234
+ TORCH_CHECK_EQ(k_pe.dim(), 2);
235
+ TORCH_CHECK_EQ(k_pe.scalar_type(), q_pe.scalar_type());
236
+
237
+ TORCH_CHECK_EQ(kv_c.size(0), num_tokens);
238
+ TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank);
239
+ TORCH_CHECK_EQ(kv_c.dim(), 2);
240
+ TORCH_CHECK_EQ(kv_c.scalar_type(), q_pe.scalar_type());
241
+ TORCH_CHECK_EQ(kv_c.dtype(), q_pe.dtype());
242
+
243
+ TORCH_CHECK_EQ(rope_cos_sin_cache.size(1), rot_dim);
244
+ TORCH_CHECK_EQ(rope_cos_sin_cache.scalar_type(), q_pe.scalar_type());
245
+
246
+ TORCH_CHECK_EQ(kv_cache_slot_mapping.size(0), num_tokens);
247
+ TORCH_CHECK_EQ(kv_cache_slot_mapping.scalar_type(), c10::ScalarType::Long);
248
+
249
+ TORCH_CHECK_EQ(kv_cache.size(2), kv_lora_rank + rot_dim);
250
+ TORCH_CHECK_EQ(kv_cache.dim(), 3);
251
+
252
+ TORCH_CHECK_EQ(kv_cache_quant_scale.numel(), 1);
253
+ TORCH_CHECK_EQ(kv_cache_quant_scale.scalar_type(), c10::ScalarType::Float);
254
+
255
+ int64_t q_pe_stride_token = q_pe.stride(0);
256
+ int64_t q_pe_stride_head = q_pe.stride(1);
257
+
258
+ int64_t k_pe_stride = k_pe.stride(0);
259
+ int64_t kv_c_stride = kv_c.stride(0);
260
+
261
+ int block_size = kv_cache.size(1);
262
+
263
+ int block_stride = kv_cache.stride(0);
264
+ int entry_stride = kv_cache.stride(1);
265
+
266
+ int rope_block_size = std::min(num_q_heads * rot_dim / 2, 512);
267
+ int mla_block_size = kv_lora_rank;
268
+ int thread_block_size =
269
+ std::min(std::max(rope_block_size, mla_block_size), 512);
270
+
271
+ dim3 grid(num_tokens, 1, 1);
272
+ dim3 block(thread_block_size, 1, 1);
273
+
274
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(positions));
275
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
276
+
277
+ DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype,
278
+ CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED);
279
+}
csrc/torch_bindings.cpp
+16
@@ -737,6 +737,22 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
737
" Tensor scale) -> ()");
738
cache_ops.impl("concat_and_cache_mla", torch::kCUDA, &concat_and_cache_mla);
739
740
+ // Rotate Q and K, then write to kv cache for MLA
741
+ cache_ops.def(
742
+ "concat_and_cache_mla_rope_fused("
743
+ " Tensor positions,"
744
+ " Tensor! q_pe,"
745
+ " Tensor! k_pe,"
746
+ " Tensor kv_c,"
747
+ " Tensor cos_sin_cache,"
748
+ " bool is_neox,"
749
+ " Tensor slot_mapping,"
750
+ " Tensor! kv_cache,"
751
+ " str kv_cache_dtype,"
752
+ " Tensor kv_cache_scale) -> ()");
753
+ cache_ops.impl("concat_and_cache_mla_rope_fused", torch::kCUDA,
754
+ &concat_and_cache_mla_rope_fused);
755
+
756
// Convert the key and value cache to fp8 data type.
757
cache_ops.def(
758
"convert_fp8(Tensor! dst_cache, Tensor src_cache, float scale, "
tests/kernels/core/test_rotary_embedding_mla_cache_fused.py
new
+161
@@ -0,0 +1,161 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Tests for fused MLA KV-cache write and RoPE fused kernel
5
+"""
6
+
7
+import random
8
+
9
+import pytest
10
+import torch
11
+
12
+from tests.kernels.allclose_default import get_default_atol, get_default_rtol
13
+from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
14
+from vllm import _custom_ops as ops
15
+from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
16
+from vllm.platforms import current_platform
17
+
18
+
19
+@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16, torch.float])
20
+@pytest.mark.parametrize("is_neox_style", [False, True])
21
+@pytest.mark.parametrize("seq_len", [11, 42])
22
+@pytest.mark.parametrize("qk_rope_head_dim", [64, 128])
23
+@pytest.mark.parametrize("num_q_heads", [128])
24
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
25
+@pytest.mark.parametrize("kv_lora_rank", [512])
26
+@pytest.mark.parametrize("num_blocks", [64])
27
+@pytest.mark.parametrize("block_size", [16, 64, 256])
28
+@pytest.mark.parametrize("seed", [0])
29
+@pytest.mark.parametrize(
30
+ "device", [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
31
+)
32
+@torch.inference_mode()
33
+def test_concat_and_cache_mla_rope_fused(
34
+ dtype: torch.dtype,
35
+ is_neox_style: bool,
36
+ seq_len: int,
37
+ qk_rope_head_dim: int,
38
+ num_q_heads: int,
39
+ kv_cache_dtype: str,
40
+ kv_lora_rank: int,
41
+ num_blocks: int,
42
+ block_size: int,
43
+ seed: int,
44
+ device: str,
45
+ max_position: int = 8192,
46
+ base: float = 10000,
47
+) -> None:
48
+ current_platform.seed_everything(seed)
49
+ torch.set_default_device(device)
50
+
51
+ rope = RotaryEmbedding(
52
+ qk_rope_head_dim,
53
+ qk_rope_head_dim,
54
+ max_position,
55
+ base,
56
+ is_neox_style,
57
+ torch.float32,
58
+ )
59
+
60
+ rope = rope.to(dtype=dtype, device=torch.get_default_device())
61
+
62
+ positions = torch.randint(0, max_position, (seq_len,))
63
+
64
+ query = torch.randn(seq_len, num_q_heads, qk_rope_head_dim, dtype=dtype)
65
+ key = torch.randn(seq_len, 1, qk_rope_head_dim + kv_lora_rank, dtype=dtype)
66
+
67
+ k_pe = torch.flatten(key[..., :qk_rope_head_dim], start_dim=1).to(device=device)
68
+ kv_c = torch.flatten(key[..., qk_rope_head_dim:], start_dim=1).to(device=device)
69
+
70
+ # NOTE(woosuk): The reference implementation should be executed first
71
+ # because the custom kernel is in-place.
72
+ ref_q_pe, ref_k_pe = rope.forward_native(positions, query, k_pe)
73
+ assert ref_k_pe is not None
74
+
75
+ ref_k_pe = torch.flatten(ref_k_pe, start_dim=1).to(device=device)
76
+ ref_k_rope = ref_k_pe[..., :qk_rope_head_dim]
77
+
78
+ total_available_slots = num_blocks * block_size
79
+ total_needed_slots = seq_len
80
+ assert total_available_slots >= total_needed_slots, "Not enough kv slots!"
81
+
82
+ slot_mapping_lst = random.sample(range(total_available_slots), total_needed_slots)
83
+ slot_mapping = torch.tensor(slot_mapping_lst, dtype=torch.long, device=device)
84
+
85
+ entry_size = kv_lora_rank + qk_rope_head_dim
86
+
87
+ kv_cache_scale = torch.tensor([0.1], dtype=torch.float32, device=device)
88
+
89
+ kv_cache = torch.zeros(
90
+ num_blocks,
91
+ block_size,
92
+ entry_size,
93
+ dtype=torch.uint8 if kv_cache_dtype == "fp8" else dtype,
94
+ device=device,
95
+ )
96
+
97
+ ref_temp = torch.zeros(*kv_cache.shape, dtype=dtype, device=device)
98
+
99
+ for i in range(seq_len):
100
+ slot = slot_mapping[i].item()
101
+ block_idx = slot // block_size
102
+ block_offset = slot % block_size
103
+ ref_temp[block_idx, block_offset] = torch.cat((kv_c[i], ref_k_rope[i]), -1)
104
+
105
+ if kv_cache_dtype == "fp8":
106
+ ref_kv_cache = torch.empty_like(ref_temp, dtype=kv_cache.dtype)
107
+ ops.convert_fp8(
108
+ ref_kv_cache, ref_temp, kv_cache_scale.item(), kv_dtype=kv_cache_dtype
109
+ )
110
+ else:
111
+ ref_kv_cache = ref_temp
112
+
113
+ opcheck(
114
+ torch.ops._C_cache_ops.concat_and_cache_mla_rope_fused,
115
+ (
116
+ positions,
117
+ query,
118
+ k_pe,
119
+ kv_c,
120
+ rope.cos_sin_cache,
121
+ is_neox_style,
122
+ slot_mapping,
123
+ kv_cache,
124
+ kv_cache_dtype,
125
+ kv_cache_scale,
126
+ ),
127
+ test_utils=DEFAULT_OPCHECK_TEST_UTILS,
128
+ )
129
+
130
+ ops.concat_and_cache_mla_rope_fused(
131
+ positions,
132
+ query,
133
+ k_pe,
134
+ kv_c,
135
+ rope.cos_sin_cache,
136
+ is_neox_style,
137
+ slot_mapping,
138
+ kv_cache,
139
+ kv_cache_dtype,
140
+ kv_cache_scale,
141
+ )
142
+
143
+ if kv_cache_dtype == "fp8":
144
+ result_temp = torch.empty_like(kv_cache, dtype=torch.float16)
145
+ ops.convert_fp8(
146
+ result_temp,
147
+ kv_cache.contiguous(),
148
+ kv_cache_scale.item(),
149
+ kv_dtype=kv_cache_dtype,
150
+ )
151
+ expected_temp = torch.empty_like(ref_kv_cache, dtype=torch.float16)
152
+ ops.convert_fp8(
153
+ expected_temp, ref_kv_cache, kv_cache_scale.item(), kv_dtype=kv_cache_dtype
154
+ )
155
+ torch.testing.assert_close(result_temp, expected_temp, atol=0.001, rtol=0.1)
156
+ else:
157
+ torch.testing.assert_close(kv_cache, ref_kv_cache)
158
+
159
+ torch.testing.assert_close(
160
+ query, ref_q_pe, atol=get_default_atol(query), rtol=get_default_rtol(query)
161
+ )
vllm/_custom_ops.py
+26
@@ -2418,6 +2418,32 @@ def concat_and_cache_mla(
2418
)
2419
2420
2421
+def concat_and_cache_mla_rope_fused(
2422
+ positions: torch.Tensor,
2423
+ q_pe: torch.Tensor,
2424
+ k_pe: torch.Tensor,
2425
+ kv_c: torch.Tensor,
2426
+ cos_sin_cache: torch.Tensor,
2427
+ is_neox: bool,
2428
+ slot_mapping: torch.Tensor,
2429
+ kv_cache: torch.Tensor,
2430
+ kv_cache_dtype: str,
2431
+ kv_cache_scale: torch.Tensor,
2432
+) -> None:
2433
+ torch.ops._C_cache_ops.concat_and_cache_mla_rope_fused(
2434
+ positions,
2435
+ q_pe,
2436
+ k_pe,
2437
+ kv_c,
2438
+ cos_sin_cache,
2439
+ is_neox,
2440
+ slot_mapping,
2441
+ kv_cache,
2442
+ kv_cache_dtype,
2443
+ kv_cache_scale,
2444
+ )
2445
+
2446
+
2447
def swap_blocks(
2448
src: torch.Tensor, dst: torch.Tensor, block_mapping: torch.Tensor
2449
) -> None: