@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 68c09efc37
Fork 0

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 13, 12:00 UTC

[Kernel][Perf] fuse QK Norm and RoPE into one cuda kernel for Qwen Model (#27165)

Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>

zhrrr committed Nov 12, 2025 at 01:00 UTC 68c09efc37e87032640cf8db571eaf486bd744ac
16 files changed +1234 -29
.buildkite/test-pipeline.yaml
+1
@@ -451,6 +451,7 @@ steps:
451 - pytest -v -s compile/test_decorator.py
452 - pytest -v -s compile/test_noop_elimination.py
453 - pytest -v -s compile/test_aot_compile.py
454 + - pytest -v -s compile/test_qk_norm_rope_fusion.py
455
456 - label: PyTorch Fullgraph Smoke Test # 15min
457 timeout_in_minutes: 30
CMakeLists.txt
+1
@@ -265,6 +265,7 @@ set(VLLM_EXT_SRC
265 "csrc/pos_encoding_kernels.cu"
266 "csrc/activation_kernels.cu"
267 "csrc/layernorm_kernels.cu"
268 + "csrc/fused_qknorm_rope_kernel.cu"
269 "csrc/layernorm_quant_kernels.cu"
270 "csrc/sampler.cu"
271 "csrc/cuda_view.cu"
csrc/fused_qknorm_rope_kernel.cu new
+418
@@ -0,0 +1,418 @@
1 +/*
2 + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
3 + *
4 + * Licensed under the Apache License, Version 2.0 (the "License");
5 + * you may not use this file except in compliance with the License.
6 + * You may obtain a copy of the License at
7 + *
8 + * http://www.apache.org/licenses/LICENSE-2.0
9 + *
10 + * Unless required by applicable law or agreed to in writing, software
11 + * distributed under the License is distributed on an "AS IS" BASIS,
12 + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 + * See the License for the specific language governing permissions and
14 + * limitations under the License.
15 + */
16 +
17 +#include <cmath>
18 +#include <cuda_runtime.h>
19 +#include <type_traits>
20 +
21 +#include <torch/cuda.h>
22 +#include <c10/cuda/CUDAGuard.h>
23 +
24 +#include "cuda_compat.h"
25 +#include "dispatch_utils.h"
26 +#include "type_convert.cuh"
27 +
28 +#define CHECK_TYPE(x, st) \
29 + TORCH_CHECK(x.scalar_type() == st, #x " dtype is ", x.scalar_type(), \
30 + ", while ", st, " is expected")
31 +#define CHECK_TH_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
32 +#define CHECK_CONTIGUOUS(x) \
33 + TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
34 +#define CHECK_INPUT(x) \
35 + CHECK_TH_CUDA(x); \
36 + CHECK_CONTIGUOUS(x)
37 +
38 +#define FINAL_MASK 0xffffffff
39 +
40 +// TODO: suport for AMD ROCM platform
41 +#ifndef USE_ROCM
42 +namespace tensorrt_llm::common {
43 +template <typename T, int num>
44 +struct packed_as;
45 +// Specialization for packed_as used in this kernel.
46 +template <>
47 +struct packed_as<uint, 1> {
48 + using type = uint;
49 +};
50 +
51 +template <>
52 +struct packed_as<uint, 2> {
53 + using type = uint2;
54 +};
55 +
56 +template <>
57 +struct packed_as<uint, 4> {
58 + using type = uint4;
59 +};
60 +
61 +template <typename T>
62 +__inline__ __device__ T warpReduceSum(T val) {
63 + #pragma unroll
64 + for (int mask = 16; mask > 0; mask >>= 1)
65 + val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
66 + return val;
67 +}
68 +
69 +template <typename T>
70 +inline __device__ __host__ T divUp(T m, T n) {
71 + return (m + n - 1) / n;
72 +}
73 +
74 +} // namespace tensorrt_llm::common
75 +
76 +namespace tensorrt_llm::kernels {
77 +// NOTE(zhuhaoran): This kernel is adapted from TensorRT-LLM implementation,
78 +// with added support for passing the cos_sin_cache as an input.
79 +// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu
80 +
81 +// Perform per-head QK Norm and RoPE in a single kernel.
82 +// scalar_t_in: data type of QKV and RMSNorm weights
83 +// scalar_t_cache: data type of cos/sin cache
84 +// head_dim: the dimension of each head
85 +// interleave: interleave=!is_neox.
86 +template <typename scalar_t_in, typename scalar_t_cache, int head_dim,
87 + bool interleave>
88 +__global__ void fusedQKNormRopeKernel(
89 + void* qkv_void, // Combined QKV tensor
90 + int const num_heads_q, // Number of query heads
91 + int const num_heads_k, // Number of key heads
92 + int const num_heads_v, // Number of value heads
93 + float const eps, // Epsilon for RMS normalization
94 + void const* q_weight_void, // RMSNorm weights for query
95 + void const* k_weight_void, // RMSNorm weights for key
96 + void const* cos_sin_cache_void, // Pre-computed cos/sin cache
97 + int64_t const* position_ids, // Position IDs for RoPE
98 + int const num_tokens // Number of tokens
99 +) {
100 + #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800
101 + if constexpr ((std::is_same_v<scalar_t_in, c10::BFloat16>) ||
102 + std::is_same_v<scalar_t_cache, c10::BFloat16>) {
103 + return;
104 + } else {
105 + #endif
106 +
107 + using Converter = vllm::_typeConvert<scalar_t_in>;
108 + static_assert(Converter::exists,
109 + "Input QKV data type is not supported for this CUDA "
110 + "architecture or toolkit version.");
111 + using T_in = typename Converter::hip_type;
112 + using T2_in = typename Converter::packed_hip_type;
113 +
114 + using CacheConverter = vllm::_typeConvert<scalar_t_cache>;
115 + static_assert(CacheConverter::exists,
116 + "Cache data type is not supported for this CUDA architecture "
117 + "or toolkit version.");
118 + using T_cache = typename CacheConverter::hip_type;
119 +
120 + T_in* qkv = reinterpret_cast<T_in*>(qkv_void);
121 + T_in const* q_weight = reinterpret_cast<T_in const*>(q_weight_void);
122 + T_in const* k_weight = reinterpret_cast<T_in const*>(k_weight_void);
123 + T_cache const* cos_sin_cache =
124 + reinterpret_cast<T_cache const*>(cos_sin_cache_void);
125 +
126 + int const warpsPerBlock = blockDim.x / 32;
127 + int const warpId = threadIdx.x / 32;
128 + int const laneId = threadIdx.x % 32;
129 +
130 + // Calculate global warp index to determine which head/token this warp
131 + // processes
132 + int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
133 +
134 + // Total number of attention heads (Q and K)
135 + int const total_qk_heads = num_heads_q + num_heads_k;
136 +
137 + // Determine which token and head type (Q or K) this warp processes
138 + int const tokenIdx = globalWarpIdx / total_qk_heads;
139 + int const localHeadIdx = globalWarpIdx % total_qk_heads;
140 +
141 + // Skip if this warp is assigned beyond the number of tokens
142 + if (tokenIdx >= num_tokens) return;
143 +
144 + bool const isQ = localHeadIdx < num_heads_q;
145 + int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q;
146 +
147 + int const num_heads = num_heads_q + num_heads_k + num_heads_v;
148 +
149 + static_assert(head_dim % (32 * 2) == 0,
150 + "head_dim must be divisible by 64 (each warp processes one "
151 + "head, and each thread gets even number of "
152 + "elements)");
153 + constexpr int numElemsPerThread = head_dim / 32;
154 + float elements[numElemsPerThread];
155 + constexpr int elemSizeBytes = numElemsPerThread * sizeof(__nv_bfloat16);
156 + static_assert(elemSizeBytes % 4 == 0,
157 + "numSizeBytes must be a multiple of 4");
158 + constexpr int vecSize =
159 + elemSizeBytes /
160 + 4; // Use packed_as<uint, vecSize> to perform loading/saving.
161 + using vec_T = typename tensorrt_llm::common::packed_as<uint, vecSize>::type;
162 +
163 + int offsetWarp; // Offset for the warp
164 + if (isQ) {
165 + // Q segment: token offset + head offset within Q segment
166 + offsetWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim;
167 + } else {
168 + // K segment: token offset + entire Q segment + head offset within K
169 + // segment
170 + offsetWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim +
171 + headIdx * head_dim;
172 + }
173 + int offsetThread = offsetWarp + laneId * numElemsPerThread;
174 +
175 + // Sum of squares for RMSNorm
176 + float sumOfSquares = 0.0f;
177 +
178 + // Load.
179 + {
180 + vec_T vec = *reinterpret_cast<vec_T const*>(&qkv[offsetThread]);
181 + constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
182 + #pragma unroll
183 + for (int i = 0; i < num_packed_elems; i++) {
184 + // Interpret the generic vector chunk as the specific packed type
185 + T2_in packed_val = *(reinterpret_cast<T2_in*>(&vec) + i);
186 + // Convert to float2 for computation
187 + float2 vals = Converter::convert(packed_val);
188 + sumOfSquares += vals.x * vals.x;
189 + sumOfSquares += vals.y * vals.y;
190 +
191 + elements[2 * i] = vals.x;
192 + elements[2 * i + 1] = vals.y;
193 + }
194 + }
195 +
196 + // Reduce sum across warp using the utility function
197 + sumOfSquares = tensorrt_llm::common::warpReduceSum(sumOfSquares);
198 +
199 + // Compute RMS normalization factor
200 + float rms_rcp = rsqrtf(sumOfSquares / static_cast<float>(head_dim) + eps);
201 +
202 + // Normalize elements
203 + #pragma unroll
204 + for (int i = 0; i < numElemsPerThread; i++) {
205 + int dim = laneId * numElemsPerThread + i;
206 + float weight = isQ ? Converter::convert(q_weight[dim])
207 + : Converter::convert(k_weight[dim]);
208 + elements[i] *= rms_rcp * weight;
209 + }
210 +
211 + // Apply RoPE to normalized elements
212 + float elements2[numElemsPerThread]; // Additional buffer required for RoPE.
213 +
214 + int64_t pos_id = position_ids[tokenIdx];
215 +
216 + // Calculate cache pointer for this position - similar to
217 + // pos_encoding_kernels.cu
218 + T_cache const* cache_ptr = cos_sin_cache + pos_id * head_dim;
219 + int const embed_dim = head_dim / 2;
220 + T_cache const* cos_ptr = cache_ptr;
221 + T_cache const* sin_ptr = cache_ptr + embed_dim;
222 +
223 + if constexpr (interleave) {
224 + // Perform interleaving. Use pre-computed cos/sin values.
225 + #pragma unroll
226 + for (int i = 0; i < numElemsPerThread / 2; ++i) {
227 + int const idx0 = 2 * i;
228 + int const idx1 = 2 * i + 1;
229 +
230 + float const val0 = elements[idx0];
231 + float const val1 = elements[idx1];
232 +
233 + int const dim_idx = laneId * numElemsPerThread + idx0;
234 + int const half_dim = dim_idx / 2;
235 + float const cos_val =
236 + CacheConverter::convert(VLLM_LDG(cos_ptr + half_dim));
237 + float const sin_val =
238 + CacheConverter::convert(VLLM_LDG(sin_ptr + half_dim));
239 +
240 + elements[idx0] = val0 * cos_val - val1 * sin_val;
241 + elements[idx1] = val0 * sin_val + val1 * cos_val;
242 + }
243 + } else {
244 + // Before data exchange with in warp, we need to sync.
245 + __syncwarp();
246 + // Get the data from the other half of the warp. Use pre-computed cos/sin
247 + // values.
248 + #pragma unroll
249 + for (int i = 0; i < numElemsPerThread; i++) {
250 + elements2[i] = __shfl_xor_sync(0xffffffff, elements[i], 16);
251 + if (laneId < 16) {
252 + elements2[i] = -elements2[i];
253 + }
254 +
255 + int dim_idx = laneId * numElemsPerThread + i;
256 + dim_idx = (dim_idx * 2) % head_dim;
257 + int half_dim = dim_idx / 2;
258 + // Use pre-computed cos/sin from cache
259 + float cos_val = CacheConverter::convert(VLLM_LDG(cos_ptr + half_dim));
260 + float sin_val = CacheConverter::convert(VLLM_LDG(sin_ptr + half_dim));
261 +
262 + elements[i] = elements[i] * cos_val + elements2[i] * sin_val;
263 + }
264 + // __shfl_xor_sync does not provide memfence. Need to sync again.
265 + __syncwarp();
266 + }
267 +
268 + // Store.
269 + {
270 + vec_T vec;
271 + constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
272 + #pragma unroll
273 + for (int i = 0; i < num_packed_elems; i++) {
274 + // Convert from float2 back to the specific packed type
275 + T2_in packed_val = Converter::convert(
276 + make_float2(elements[2 * i], elements[2 * i + 1]));
277 + // Place it into the generic vector
278 + *(reinterpret_cast<T2_in*>(&vec) + i) = packed_val;
279 + }
280 + *reinterpret_cast<vec_T*>(&qkv[offsetThread]) = vec;
281 + }
282 +
283 + #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800
284 + }
285 + #endif
286 +}
287 +
288 + // Borrowed from
289 + // https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568
290 + #define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \
291 + if (interleave) { \
292 + const bool INTERLEAVE = true; \
293 + __VA_ARGS__ \
294 + } else { \
295 + const bool INTERLEAVE = false; \
296 + __VA_ARGS__ \
297 + }
298 +
299 +template <typename scalar_t_in, typename scalar_t_cache>
300 +void launchFusedQKNormRope(void* qkv, int const num_tokens,
301 + int const num_heads_q, int const num_heads_k,
302 + int const num_heads_v, int const head_dim,
303 + float const eps, void const* q_weight,
304 + void const* k_weight, void const* cos_sin_cache,
305 + bool const interleave, int64_t const* position_ids,
306 + cudaStream_t stream) {
307 + constexpr int blockSize = 256;
308 +
309 + int const warpsPerBlock = blockSize / 32;
310 + int const totalQKHeads = num_heads_q + num_heads_k;
311 + int const totalWarps = num_tokens * totalQKHeads;
312 +
313 + int const gridSize = common::divUp(totalWarps, warpsPerBlock);
314 + dim3 gridDim(gridSize);
315 + dim3 blockDim(blockSize);
316 +
317 + switch (head_dim) {
318 + case 64:
319 + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
320 + fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 64, INTERLEAVE>
321 + <<<gridDim, blockDim, 0, stream>>>(
322 + qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight,
323 + k_weight, cos_sin_cache, position_ids, num_tokens);
324 + });
325 + break;
326 + case 128:
327 + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
328 + fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 128, INTERLEAVE>
329 + <<<gridDim, blockDim, 0, stream>>>(
330 + qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight,
331 + k_weight, cos_sin_cache, position_ids, num_tokens);
332 + });
333 + break;
334 + case 256:
335 + DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
336 + fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 256, INTERLEAVE>
337 + <<<gridDim, blockDim, 0, stream>>>(
338 + qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight,
339 + k_weight, cos_sin_cache, position_ids, num_tokens);
340 + });
341 + break;
342 + default:
343 + TORCH_CHECK(false,
344 + "Unsupported head dimension for fusedQKNormRope: ", head_dim);
345 + }
346 +}
347 +} // namespace tensorrt_llm::kernels
348 +
349 +void fused_qk_norm_rope(
350 + torch::Tensor& qkv, // Combined QKV tensor [num_tokens,
351 + // (num_heads_q+num_heads_k+num_heads_v)*head_dim]
352 + int64_t num_heads_q, // Number of query heads
353 + int64_t num_heads_k, // Number of key heads
354 + int64_t num_heads_v, // Number of value heads
355 + int64_t head_dim, // Dimension per head
356 + double eps, // Epsilon for RMS normalization
357 + torch::Tensor& q_weight, // RMSNorm weights for query [head_dim]
358 + torch::Tensor& k_weight, // RMSNorm weights for key [head_dim]
359 + torch::Tensor& cos_sin_cache, // Cos/sin cache [max_position, head_dim]
360 + bool is_neox, // Whether RoPE is applied in Neox style
361 + torch::Tensor& position_ids // Position IDs for RoPE [num_tokens]
362 +) {
363 + // Input validation
364 + CHECK_INPUT(qkv);
365 + CHECK_INPUT(position_ids);
366 + CHECK_INPUT(q_weight);
367 + CHECK_INPUT(k_weight);
368 + CHECK_INPUT(cos_sin_cache);
369 + CHECK_TYPE(position_ids, torch::kInt64);
370 +
371 + TORCH_CHECK(qkv.dim() == 2,
372 + "QKV tensor must be 2D: [num_tokens, "
373 + "(num_heads_q+num_heads_k+num_heads_v)*head_dim]");
374 + TORCH_CHECK(position_ids.dim() == 1, "Position IDs must be 1D: [num_tokens]");
375 + TORCH_CHECK(q_weight.dim() == 1, "Query weights must be 1D: [head_dim]");
376 + TORCH_CHECK(k_weight.dim() == 1, "Key weights must be 1D: [head_dim]");
377 + TORCH_CHECK(cos_sin_cache.dim() == 2,
378 + "Cos/sin cache must be 2D: [max_position, head_dim]");
379 + TORCH_CHECK(q_weight.size(0) == head_dim,
380 + "Query weights size must match head dimension");
381 + TORCH_CHECK(k_weight.size(0) == head_dim,
382 + "Key weights size must match head dimension");
383 + TORCH_CHECK(cos_sin_cache.size(1) == head_dim,
384 + "Cos/sin cache dimension must match head_dim");
385 + TORCH_CHECK(qkv.scalar_type() == q_weight.scalar_type() &&
386 + qkv.scalar_type() == k_weight.scalar_type(),
387 + "qkv, q_weight and k_weight must have the same dtype");
388 +
389 + int64_t num_tokens = qkv.size(0);
390 + TORCH_CHECK(position_ids.size(0) == num_tokens,
391 + "Number of tokens in position_ids must match QKV");
392 +
393 + int64_t total_heads = num_heads_q + num_heads_k + num_heads_v;
394 + TORCH_CHECK(
395 + qkv.size(1) == total_heads * head_dim,
396 + "QKV tensor size must match total number of heads and head dimension");
397 +
398 + auto stream = at::cuda::getCurrentCUDAStream(qkv.get_device());
399 +
400 + VLLM_DISPATCH_HALF_TYPES(qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
401 + using qkv_scalar_t = scalar_t;
402 + VLLM_DISPATCH_FLOATING_TYPES(
403 + cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
404 + using cache_scalar_t = scalar_t;
405 + tensorrt_llm::kernels::launchFusedQKNormRope<qkv_scalar_t,
406 + cache_scalar_t>(
407 + qkv.data_ptr(), static_cast<int>(num_tokens),
408 + static_cast<int>(num_heads_q), static_cast<int>(num_heads_k),
409 + static_cast<int>(num_heads_v), static_cast<int>(head_dim),
410 + static_cast<float>(eps), q_weight.data_ptr(), k_weight.data_ptr(),
411 + cos_sin_cache.data_ptr(), !is_neox,
412 + reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
413 + stream);
414 + });
415 + });
416 +}
417 +
418 +#endif // not USE_ROCM
\ No newline at end of file
csrc/ops.h
+6
@@ -92,6 +92,12 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
92 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
93 torch::Tensor& weight, double epsilon);
94
95 +void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q,
96 + int64_t num_heads_k, int64_t num_heads_v,
97 + int64_t head_dim, double eps, torch::Tensor& q_weight,
98 + torch::Tensor& k_weight, torch::Tensor& cos_sin_cache,
99 + bool is_neox, torch::Tensor& position_ids);
100 +
101 void apply_repetition_penalties_(torch::Tensor& logits,
102 const torch::Tensor& prompt_mask,
103 const torch::Tensor& output_mask,
csrc/torch_bindings.cpp
+10
@@ -175,6 +175,16 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
175 "float epsilon) -> ()");
176 ops.impl("fused_add_rms_norm", torch::kCUDA, &fused_add_rms_norm);
177
178 +#ifndef USE_ROCM
179 + // Function for fused QK Norm and RoPE
180 + ops.def(
181 + "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, "
182 + "int num_heads_k, int num_heads_v, int head_dim, float eps, "
183 + "Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, "
184 + "bool is_neox, Tensor position_ids) -> ()");
185 + ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
186 +#endif
187 +
188 // Apply repetition penalties to logits in-place
189 ops.def(
190 "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
csrc/type_convert.cuh
+44 -16
@@ -29,6 +29,22 @@ struct _typeConvert {
29 static constexpr bool exists = false;
30 };
31
32 +template <>
33 +struct _typeConvert<float> {
34 + static constexpr bool exists = true;
35 + using hip_type = float;
36 + using packed_hip_type = float2;
37 + using packed_hip_type4 = float4; // For 128-bit vectorization
38 +
39 + __device__ static __forceinline__ float convert(hip_type x) { return x; }
40 + __device__ static __forceinline__ float2 convert(packed_hip_type x) {
41 + return x;
42 + }
43 + __device__ static __forceinline__ float4 convert(packed_hip_type4 x) {
44 + return x;
45 + }
46 +};
47 +
48 #if defined(USE_ROCM) || (defined(CUDA_VERSION) && (CUDA_VERSION >= 12000))
49 // CUDA < 12.0 runs into issues with packed type conversion
50 template <>
@@ -37,14 +53,16 @@ struct _typeConvert<c10::Half> {
53 using hip_type = __half;
54 using packed_hip_type = __half2;
55
40 - __device__ static inline float convert(hip_type x) { return __half2float(x); }
41 - __device__ static inline float2 convert(packed_hip_type x) {
56 + __device__ static __forceinline__ float convert(hip_type x) {
57 + return __half2float(x);
58 + }
59 + __device__ static __forceinline__ float2 convert(packed_hip_type x) {
60 return __half22float2(x);
61 }
44 - __device__ static inline hip_type convert(float x) {
62 + __device__ static __forceinline__ hip_type convert(float x) {
63 return __float2half_rn(x);
64 }
47 - __device__ static inline packed_hip_type convert(float2 x) {
65 + __device__ static __forceinline__ packed_hip_type convert(float2 x) {
66 return __float22half2_rn(x);
67 }
68 };
@@ -58,16 +76,16 @@ struct _typeConvert<c10::BFloat16> {
76 using hip_type = __nv_bfloat16;
77 using packed_hip_type = __nv_bfloat162;
78
61 - __device__ static inline float convert(hip_type x) {
79 + __device__ static __forceinline__ float convert(hip_type x) {
80 return __bfloat162float(x);
81 }
64 - __device__ static inline float2 convert(packed_hip_type x) {
82 + __device__ static __forceinline__ float2 convert(packed_hip_type x) {
83 return __bfloat1622float2(x);
84 }
67 - __device__ static inline hip_type convert(float x) {
85 + __device__ static __forceinline__ hip_type convert(float x) {
86 return __float2bfloat16(x);
87 }
70 - __device__ static inline packed_hip_type convert(float2 x) {
88 + __device__ static __forceinline__ packed_hip_type convert(float2 x) {
89 return __float22bfloat162_rn(x);
90 }
91 };
@@ -95,10 +113,15 @@ struct alignas(16) _f16Vec {
113 if constexpr (width % 2 == 0) {
114 #pragma unroll
115 for (int i = 0; i < width; i += 2) {
98 - T2 temp{data[i], data[i + 1]};
99 - temp += T2{other.data[i], other.data[i + 1]};
100 - data[i] = temp.x;
101 - data[i + 1] = temp.y;
116 + if constexpr (std::is_same_v<T2, float2>) {
117 + data[i] += other.data[i];
118 + data[i + 1] += other.data[i + 1];
119 + } else {
120 + T2 temp{data[i], data[i + 1]};
121 + temp += T2{other.data[i], other.data[i + 1]};
122 + data[i] = temp.x;
123 + data[i + 1] = temp.y;
124 + }
125 }
126 } else {
127 #pragma unroll
@@ -111,10 +134,15 @@ struct alignas(16) _f16Vec {
134 if constexpr (width % 2 == 0) {
135 #pragma unroll
136 for (int i = 0; i < width; i += 2) {
114 - T2 temp{data[i], data[i + 1]};
115 - temp *= T2{other.data[i], other.data[i + 1]};
116 - data[i] = temp.x;
117 - data[i + 1] = temp.y;
137 + if constexpr (std::is_same_v<T2, float2>) {
138 + data[i] *= other.data[i];
139 + data[i + 1] *= other.data[i + 1];
140 + } else {
141 + T2 temp{data[i], data[i + 1]};
142 + temp *= T2{other.data[i], other.data[i + 1]};
143 + data[i] = temp.x;
144 + data[i + 1] = temp.y;
145 + }
146 }
147 } else {
148 #pragma unroll
tests/compile/test_qk_norm_rope_fusion.py new
+195
@@ -0,0 +1,195 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from tests.compile.backend import TestBackend
8 +from vllm.attention import Attention, AttentionType
9 +from vllm.compilation.matcher_utils import FLASHINFER_ROTARY_OP, RMS_OP, ROTARY_OP
10 +from vllm.compilation.noop_elimination import NoOpEliminationPass
11 +from vllm.compilation.post_cleanup import PostCleanupPass
12 +from vllm.compilation.qk_norm_rope_fusion import (
13 + FUSED_QK_ROPE_OP,
14 + QKNormRoPEFusionPass,
15 +)
16 +from vllm.config import (
17 + CompilationConfig,
18 + CompilationMode,
19 + ModelConfig,
20 + PassConfig,
21 + VllmConfig,
22 + set_current_vllm_config,
23 +)
24 +from vllm.model_executor.layers.layernorm import RMSNorm
25 +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
26 +from vllm.platforms import current_platform
27 +
28 +RSQRT_OP = torch.ops.aten.rsqrt.default
29 +INDEX_SELECT_OP = torch.ops.aten.index.Tensor
30 +
31 +
32 +class QKNormRoPETestModel(torch.nn.Module):
33 + def __init__(
34 + self,
35 + *,
36 + num_heads: int,
37 + num_kv_heads: int,
38 + head_dim: int,
39 + eps: float,
40 + is_neox: bool,
41 + vllm_config: VllmConfig,
42 + dtype: torch.dtype,
43 + prefix: str = "model.layers.0.self_attn.attn",
44 + ) -> None:
45 + super().__init__()
46 + self.num_heads = num_heads
47 + self.num_kv_heads = num_kv_heads
48 + self.head_dim = head_dim
49 + self.q_size = num_heads * head_dim
50 + self.kv_size = num_kv_heads * head_dim
51 + self.rotary_dim = head_dim
52 + self.eps = eps
53 + self.dtype = dtype
54 +
55 + # Register layer metadata for the fusion pass via Attention.
56 + self.attn = Attention(
57 + num_heads=self.num_heads,
58 + head_size=self.head_dim,
59 + scale=1.0 / self.head_dim**0.5,
60 + num_kv_heads=self.num_kv_heads,
61 + cache_config=vllm_config.cache_config,
62 + prefix=prefix,
63 + attn_type=AttentionType.DECODER,
64 + )
65 +
66 + self.q_norm = RMSNorm(self.head_dim, eps=self.eps)
67 + self.k_norm = RMSNorm(self.head_dim, eps=self.eps)
68 + self.rotary_emb = RotaryEmbedding(
69 + self.head_dim,
70 + rotary_dim=self.rotary_dim,
71 + max_position_embeddings=4096,
72 + base=10000,
73 + is_neox_style=is_neox,
74 + dtype=self.dtype,
75 + )
76 + self.enable_rms_norm_custom_op = self.q_norm.enabled()
77 + self.enable_rope_custom_op = self.rotary_emb.enabled()
78 +
79 + def forward(self, qkv: torch.Tensor, positions: torch.Tensor):
80 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
81 + q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim)
82 + q_by_head = self.q_norm(q_by_head)
83 + q = q_by_head.view(q.shape)
84 + k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim)
85 + k_by_head = self.k_norm(k_by_head)
86 + k = k_by_head.view(k.shape)
87 + q, k = self.rotary_emb(positions, q, k)
88 + return q, k, v
89 +
90 + def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
91 + ops = []
92 + if self.enable_rms_norm_custom_op:
93 + ops.append(RMS_OP)
94 + else:
95 + ops.append(RSQRT_OP)
96 +
97 + if self.enable_rope_custom_op:
98 + if self.rotary_emb.use_flashinfer:
99 + ops.append(FLASHINFER_ROTARY_OP)
100 + else:
101 + ops.append(ROTARY_OP)
102 + else:
103 + ops.append(INDEX_SELECT_OP)
104 + return ops
105 +
106 + def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
107 + return [FUSED_QK_ROPE_OP]
108 +
109 +
110 +@pytest.mark.parametrize("eps", [1e-5, 1e-6])
111 +@pytest.mark.parametrize("is_neox", [True, False])
112 +@pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
113 +@pytest.mark.parametrize("enable_rope_custom_op", [True])
114 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
115 +@pytest.mark.skipif(
116 + not current_platform.is_cuda(),
117 + reason="Only test on cuda platform",
118 +)
119 +def test_qk_norm_rope_fusion(
120 + eps, is_neox, enable_rms_norm_custom_op, enable_rope_custom_op, dtype
121 +):
122 + if not hasattr(torch.ops._C, "fused_qk_norm_rope"):
123 + pytest.skip("fused_qk_norm_rope custom op not available")
124 +
125 + torch.set_default_device("cuda")
126 + torch.set_default_dtype(dtype)
127 + torch.manual_seed(0)
128 +
129 + custom_ops: list[str] = []
130 + if enable_rms_norm_custom_op:
131 + custom_ops.append("+rms_norm")
132 + if enable_rope_custom_op:
133 + custom_ops.append("+rotary_embedding")
134 +
135 + vllm_config = VllmConfig(
136 + model_config=ModelConfig(dtype=dtype),
137 + compilation_config=CompilationConfig(
138 + mode=CompilationMode.VLLM_COMPILE,
139 + custom_ops=custom_ops,
140 + pass_config=PassConfig(
141 + enable_qk_norm_rope_fusion=True,
142 + enable_noop=True,
143 + ),
144 + ),
145 + )
146 +
147 + num_heads, num_kv_heads, head_dim = 16, 4, 128
148 + T = 5
149 +
150 + with set_current_vllm_config(vllm_config):
151 + model = QKNormRoPETestModel(
152 + num_heads=num_heads,
153 + num_kv_heads=num_kv_heads,
154 + head_dim=head_dim,
155 + eps=eps,
156 + is_neox=is_neox,
157 + vllm_config=vllm_config,
158 + dtype=dtype,
159 + )
160 +
161 + noop_pass = NoOpEliminationPass(vllm_config)
162 + fusion_pass = QKNormRoPEFusionPass(vllm_config)
163 + cleanup_pass = PostCleanupPass(vllm_config)
164 +
165 + backend = TestBackend(noop_pass, fusion_pass, cleanup_pass)
166 + backend_baseline = TestBackend(noop_pass, cleanup_pass)
167 +
168 + qkv = torch.randn(T, model.q_size + 2 * model.kv_size)
169 + pos = torch.arange(T, dtype=torch.long, device=qkv.device)
170 + qkv_unfused = qkv.clone()
171 + pos_unfused = pos.clone()
172 +
173 + torch._dynamo.mark_dynamic(qkv, 0)
174 + torch._dynamo.mark_dynamic(pos, 0)
175 + model_fused = torch.compile(model, backend=backend)
176 + q_fused, k_fused, v_fused = model_fused(qkv, pos)
177 +
178 + torch._dynamo.mark_dynamic(qkv_unfused, 0)
179 + torch._dynamo.mark_dynamic(pos_unfused, 0)
180 + model_unfused = torch.compile(model, backend=backend_baseline)
181 + q_unfused, k_unfused, v_unfused = model_unfused(qkv_unfused, pos_unfused)
182 +
183 + if dtype == torch.float16:
184 + ATOL, RTOL = (2e-3, 2e-3)
185 + else:
186 + ATOL, RTOL = (1e-2, 1e-2)
187 +
188 + torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
189 + torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
190 + torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
191 +
192 + assert fusion_pass.matched_count == 1
193 +
194 + backend.check_before_ops(model.ops_in_model_before())
195 + backend.check_after_ops(model.ops_in_model_after())
tests/kernels/core/test_fused_qk_norm_rope.py new
+141
@@ -0,0 +1,141 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from tests.kernels.utils import opcheck
8 +from vllm.model_executor.layers.layernorm import RMSNorm
9 +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
10 +from vllm.platforms import current_platform
11 +
12 +DTYPES = [torch.bfloat16, torch.float16]
13 +IS_NEOX = [True, False]
14 +EPS_VALUES = [1e-5, 1e-6]
15 +SEEDS = [13]
16 +CUDA_DEVICES = ["cuda:0"]
17 +
18 +
19 +def _apply_qk_norm_rope(
20 + qkv: torch.Tensor,
21 + positions: torch.Tensor,
22 + q_norm: RMSNorm,
23 + k_norm: RMSNorm,
24 + rope: RotaryEmbedding,
25 + num_heads_q: int,
26 + num_heads_kv: int,
27 + head_dim: int,
28 +) -> torch.Tensor:
29 + q_size = num_heads_q * head_dim
30 + kv_size = num_heads_kv * head_dim
31 +
32 + q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
33 +
34 + q_by_head = q.view(*q.shape[:-1], q.shape[-1] // head_dim, head_dim)
35 + q_by_head = q_norm.forward_native(q_by_head)
36 + q = q_by_head.view(q.shape)
37 +
38 + k_by_head = k.view(*k.shape[:-1], k.shape[-1] // head_dim, head_dim)
39 + k_by_head = k_norm.forward_native(k_by_head)
40 + k = k_by_head.view(k.shape)
41 +
42 + q, k = rope.forward_native(positions, q, k)
43 + return torch.cat([q, k, v], dim=-1)
44 +
45 +
46 +@pytest.mark.skipif(
47 + not current_platform.is_cuda(),
48 + reason="fused_qk_norm_rope custom op requires cuda platform",
49 +)
50 +@pytest.mark.parametrize("device", CUDA_DEVICES)
51 +@pytest.mark.parametrize("dtype", DTYPES)
52 +@pytest.mark.parametrize("is_neox", IS_NEOX)
53 +@pytest.mark.parametrize("eps", EPS_VALUES)
54 +@pytest.mark.parametrize("seed", SEEDS)
55 +@torch.inference_mode()
56 +def test_fused_qk_norm_rope_matches_reference(
57 + device: str,
58 + dtype: torch.dtype,
59 + is_neox: bool,
60 + eps: float,
61 + seed: int,
62 +):
63 + torch.set_default_device(device)
64 + current_platform.seed_everything(seed)
65 + num_heads, num_kv_heads, head_dim = 16, 4, 128
66 + num_tokens = 4
67 +
68 + total_dim = (num_heads + 2 * num_kv_heads) * head_dim
69 + qkv_base = torch.randn(num_tokens, total_dim, dtype=dtype, device=device)
70 + qkv_fused = qkv_base.clone()
71 + positions = torch.arange(num_tokens, dtype=torch.long, device=device)
72 +
73 + q_norm = RMSNorm(head_dim, eps=eps).to(device=device, dtype=dtype)
74 + k_norm = RMSNorm(head_dim, eps=eps).to(device=device, dtype=dtype)
75 + q_norm.weight.data.normal_(mean=1.0, std=0.1)
76 + k_norm.weight.data.normal_(mean=1.0, std=0.1)
77 + q_weight = q_norm.weight.data
78 + k_weight = k_norm.weight.data
79 +
80 + rope = RotaryEmbedding(
81 + head_size=head_dim,
82 + rotary_dim=head_dim,
83 + max_position_embeddings=4096,
84 + base=10000.0,
85 + is_neox_style=is_neox,
86 + dtype=dtype,
87 + ).to(device)
88 +
89 + ref_result = _apply_qk_norm_rope(
90 + qkv=qkv_base,
91 + positions=positions,
92 + q_norm=q_norm,
93 + k_norm=k_norm,
94 + rope=rope,
95 + num_heads_q=num_heads,
96 + num_heads_kv=num_kv_heads,
97 + head_dim=head_dim,
98 + )
99 +
100 + opcheck(
101 + torch.ops._C.fused_qk_norm_rope,
102 + (
103 + qkv_fused.clone(),
104 + num_heads,
105 + num_kv_heads,
106 + num_kv_heads,
107 + head_dim,
108 + eps,
109 + q_weight,
110 + k_weight,
111 + rope.cos_sin_cache,
112 + is_neox,
113 + positions.view(-1),
114 + ),
115 + )
116 +
117 + torch.ops._C.fused_qk_norm_rope(
118 + qkv_fused,
119 + num_heads,
120 + num_kv_heads,
121 + num_kv_heads,
122 + head_dim,
123 + eps,
124 + q_weight,
125 + k_weight,
126 + rope.cos_sin_cache,
127 + is_neox,
128 + positions.view(-1),
129 + )
130 +
131 + if dtype == torch.float16:
132 + ATOL, RTOL = (2e-3, 2e-3)
133 + else:
134 + ATOL, RTOL = (1e-2, 1e-2)
135 +
136 + torch.testing.assert_close(
137 + qkv_fused,
138 + ref_result,
139 + atol=ATOL,
140 + rtol=RTOL,
141 + )
vllm/_custom_ops.py
+29
@@ -329,6 +329,7 @@ def rms_norm(
329 out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, epsilon: float
330 ) -> None:
331 # TODO: Remove this contiguous call when the kernel is updated to support non-contiguous input
332 + # If removed, also need to remove contiguous in MatcherRMSNorm
333 input_contiguous = input.contiguous()
334 torch.ops._C.rms_norm(out, input_contiguous, weight, epsilon)
335
@@ -339,6 +340,34 @@ def fused_add_rms_norm(
340 torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon)
341
342
343 +def fused_qk_norm_rope(
344 + qkv: torch.Tensor,
345 + num_heads_q: int,
346 + num_heads_k: int,
347 + num_heads_v: int,
348 + head_dim: int,
349 + eps: float,
350 + q_weight: torch.Tensor,
351 + k_weight: torch.Tensor,
352 + cos_sin_cache: torch.Tensor,
353 + is_neox: bool,
354 + position_ids: torch.Tensor,
355 +) -> None:
356 + torch.ops._C.fused_qk_norm_rope(
357 + qkv,
358 + num_heads_q,
359 + num_heads_k,
360 + num_heads_v,
361 + head_dim,
362 + eps,
363 + q_weight,
364 + k_weight,
365 + cos_sin_cache,
366 + is_neox,
367 + position_ids,
368 + )
369 +
370 +
371 def apply_repetition_penalties_torch(
372 logits: torch.Tensor,
373 prompt_mask: torch.Tensor,
vllm/compilation/fix_functionalization.py
+17
@@ -132,6 +132,23 @@ class FixFunctionalizationPass(VllmInductorPass):
132 "input_global_scale",
133 ),
134 )
135 + # Defunctionalize fused_qk_norm_rope to remove higher-order wrapper.
136 + elif at_target == torch.ops._C.fused_qk_norm_rope.default:
137 + mutated_args = {1: "qkv"}
138 + args = (
139 + "qkv",
140 + "num_heads_q",
141 + "num_heads_k",
142 + "num_heads_v",
143 + "head_dim",
144 + "eps",
145 + "q_weight",
146 + "k_weight",
147 + "cos_sin_cache",
148 + "is_neox",
149 + "position_ids",
150 + )
151 + self.defunctionalize(graph, node, mutated_args=mutated_args, args=args)
152 else:
153 continue # skip the count
154
vllm/compilation/fusion.py
+4
@@ -44,6 +44,10 @@ def empty_i32(*args, **kwargs):
44 return torch.empty(*args, **kwargs, dtype=torch.int32, device="cuda")
45
46
47 +def empty_i64(*args, **kwargs):
48 + return torch.empty(*args, **kwargs, dtype=torch.int64, device="cuda")
49 +
50 +
51 RMS_OP = torch.ops._C.rms_norm.default
52 RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
53
vllm/compilation/matcher_utils.py
+80 -1
@@ -18,10 +18,13 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
18 kFp8StaticTensorSym,
19 kNvfp4Quant,
20 )
21 +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
22 from vllm.platforms import current_platform
23
24 RMS_OP = torch.ops._C.rms_norm.default
25 RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
26 +ROTARY_OP = torch.ops._C.rotary_embedding.default
27 +FLASHINFER_ROTARY_OP = torch.ops.vllm.flashinfer_rotary_embedding.default
28
29 QUANT_OPS: dict[QuantKey, OpOverload] = {
30 kFp8StaticTensorSym: torch.ops._C.static_scaled_fp8_quant.default, # noqa: E501
@@ -58,6 +61,9 @@ class MatcherCustomOp(ABC):
61 def empty(self, *args, **kws):
62 return torch.empty(*args, dtype=self.model_dtype, device=self.device, **kws)
63
64 + def empty_int64(self, *args, **kws):
65 + return torch.empty(*args, dtype=torch.int64, device=self.device, **kws)
66 +
67 def empty_f32(self, *args, **kws):
68 return torch.empty(*args, dtype=torch.float32, device=self.device, **kws)
69
@@ -66,6 +72,77 @@ class MatcherCustomOp(ABC):
72 raise NotImplementedError
73
74
75 +class MatcherRotaryEmbedding(MatcherCustomOp):
76 + def __init__(
77 + self,
78 + is_neox: bool,
79 + head_size: int,
80 + num_heads: int,
81 + num_kv_heads: int,
82 + use_flashinfer: bool = False,
83 + enabled: bool | None = None,
84 + ) -> None:
85 + if enabled is None:
86 + enabled = RotaryEmbedding.enabled()
87 +
88 + super().__init__(enabled)
89 + self.is_neox = is_neox
90 + self.head_size = head_size
91 + self.num_heads = num_heads
92 + self.num_kv_heads = num_kv_heads
93 + self.q_size = self.num_heads * self.head_size
94 + self.kv_size = self.num_kv_heads * self.head_size
95 + self.rotary_dim = head_size
96 + if use_flashinfer:
97 + self.rotary_op = FLASHINFER_ROTARY_OP
98 + else:
99 + self.rotary_op = ROTARY_OP
100 +
101 + def inputs(self) -> list[torch.Tensor]:
102 + positions = self.empty_int64(5)
103 + query = self.empty(5, self.q_size)
104 + key = self.empty(5, self.kv_size)
105 + cos_sin_cache = self.empty(4096, self.rotary_dim)
106 + return [positions, query, key, cos_sin_cache]
107 +
108 + def forward_custom(
109 + self,
110 + positions: torch.Tensor,
111 + query: torch.Tensor,
112 + key: torch.Tensor | None,
113 + cos_sin_cache: torch.Tensor,
114 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
115 + result = auto_functionalized(
116 + self.rotary_op,
117 + positions=positions,
118 + query=query,
119 + key=key,
120 + head_size=self.head_size,
121 + cos_sin_cache=cos_sin_cache,
122 + is_neox=self.is_neox,
123 + )
124 + query_out = result[1]
125 + key_out = result[2] if len(result) > 2 else None
126 + return query_out, key_out
127 +
128 + def forward_native(
129 + self,
130 + positions: torch.Tensor,
131 + query: torch.Tensor,
132 + key: torch.Tensor | None,
133 + cos_sin_cache: torch.Tensor,
134 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
135 + return RotaryEmbedding.forward_static(
136 + positions,
137 + query,
138 + key,
139 + self.head_size,
140 + self.rotary_dim,
141 + cos_sin_cache,
142 + self.is_neox,
143 + )
144 +
145 +
146 class MatcherRMSNorm(MatcherCustomOp):
147 def __init__(self, epsilon: float, enabled: bool | None = None):
148 if enabled is None:
@@ -85,10 +162,12 @@ class MatcherRMSNorm(MatcherCustomOp):
162 weight: torch.Tensor,
163 ) -> torch.Tensor:
164 result = torch.empty_like(input)
165 + # TODO: support non-contiguous input for RMSNorm and remove this
166 + input_contiguous = input.contiguous()
167 _, result = auto_functionalized(
168 RMS_OP,
169 result=result,
91 - input=input,
170 + input=input_contiguous,
171 weight=weight,
172 epsilon=self.epsilon,
173 )
vllm/compilation/pass_manager.py
+4
@@ -17,6 +17,7 @@ if current_platform.is_cuda_alike():
17 from .activation_quant_fusion import ActivationQuantFusionPass
18 from .fusion import RMSNormQuantFusionPass
19 from .fusion_attn import AttnFusionPass
20 + from .qk_norm_rope_fusion import QKNormRoPEFusionPass
21
22 if current_platform.is_cuda():
23 from .collective_fusion import AllReduceFusionPass, AsyncTPPass
@@ -109,6 +110,9 @@ class PostGradPassManager(CustomGraphPass):
110 if self.pass_config.enable_attn_fusion:
111 self.passes += [AttnFusionPass(config)]
112
113 + if self.pass_config.enable_qk_norm_rope_fusion:
114 + self.passes += [QKNormRoPEFusionPass(config)]
115 +
116 # needs a functional graph
117 self.post_cleanup = PostCleanupPass(config)
118 self.fix_functionalization = FixFunctionalizationPass(config)
vllm/compilation/qk_norm_rope_fusion.py new
+238
@@ -0,0 +1,238 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Callable
5 +
6 +import torch
7 +import torch._inductor.pattern_matcher as pm
8 +from torch import fx
9 +from torch._higher_order_ops.auto_functionalize import auto_functionalized
10 +from torch._inductor.pattern_matcher import PatternMatcherPass
11 +
12 +from vllm.attention import Attention
13 +from vllm.config import VllmConfig, get_layers_from_vllm_config
14 +from vllm.logger import init_logger
15 +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
16 +
17 +from .fusion import empty_bf16, empty_fp32, empty_i64
18 +from .inductor_pass import enable_fake_mode
19 +from .matcher_utils import MatcherRMSNorm, MatcherRotaryEmbedding
20 +from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
21 +
22 +logger = init_logger(__name__)
23 +
24 +FUSED_QK_ROPE_OP = torch.ops._C.fused_qk_norm_rope.default
25 +
26 +
27 +class QkNormRopePattern:
28 + """
29 + Match the unfused sequence in attention blocks and replace with the fused op.
30 +
31 + Unfused (conceptually):
32 + q, k, v = split(qkv, [qsz, kvsz, kvsz], -1)
33 + qh = reshape(q, [-1, num_heads, head_dim])
34 + kh = reshape(k, [-1, num_kv_heads, head_dim])
35 + qn = rms_norm(qh, q_weight, eps)
36 + kn = rms_norm(kh, k_weight, eps)
37 + qf = reshape(qn, [-1, num_heads * head_dim])
38 + kf = reshape(kn, [-1, num_kv_heads * head_dim])
39 + qf, kf = rotary_embedding(positions, qf, kf, head_dim, cos_sin_cache, is_neox)
40 + return qf, kf, v
41 +
42 + Fused replacement:
43 + fused_qk_norm_rope(qkv, num_heads, num_kv_heads, num_kv_heads, head_dim,
44 + eps, q_weight, k_weight, cos_sin_cache, is_neox,
45 + positions.view(-1))
46 + return split(qkv, [qsz, kvsz, kvsz], -1)
47 + """
48 +
49 + def __init__(
50 + self,
51 + head_dim: int,
52 + num_heads: int,
53 + num_kv_heads: int,
54 + eps: float,
55 + is_neox: bool,
56 + rope_flashinfer: bool = False,
57 + ) -> None:
58 + self.num_heads = num_heads
59 + self.num_kv_heads = num_kv_heads
60 + self.head_dim = head_dim
61 + self.q_size = self.num_heads * self.head_dim
62 + self.kv_size = self.num_kv_heads * self.head_dim
63 + self.eps = eps
64 + self.rmsnorm_matcher = MatcherRMSNorm(eps)
65 + self.is_neox = is_neox
66 + self.rope_flashinfer = rope_flashinfer
67 + self.rope_matcher = MatcherRotaryEmbedding(
68 + is_neox=is_neox,
69 + head_size=self.head_dim,
70 + num_heads=self.num_heads,
71 + num_kv_heads=self.num_kv_heads,
72 + use_flashinfer=self.rope_flashinfer,
73 + )
74 +
75 + def get_inputs(self):
76 + # Sample inputs to help pattern tracing
77 + T = 5
78 + qkv = empty_bf16(T, self.q_size + 2 * self.kv_size)
79 + positions = empty_i64(T)
80 + q_weight = empty_bf16(1, self.head_dim)
81 + k_weight = empty_bf16(1, self.head_dim)
82 + if self.rope_flashinfer:
83 + cos_sin_cache = empty_fp32(4096, self.head_dim)
84 + else:
85 + cos_sin_cache = empty_bf16(4096, self.head_dim)
86 + return [
87 + qkv,
88 + positions,
89 + q_weight,
90 + k_weight,
91 + cos_sin_cache,
92 + ]
93 +
94 + @staticmethod
95 + def wrap_trace_fn(trace_fn, *process_fx_fns: Callable[[fx.GraphModule], None]):
96 + def wrapped(*args, **kwargs):
97 + gm = trace_fn(*args, **kwargs)
98 + for process_fx in process_fx_fns:
99 + process_fx(gm)
100 +
101 + return gm
102 +
103 + return wrapped
104 +
105 + @staticmethod
106 + def fx_view_to_reshape(gm: torch.fx.GraphModule):
107 + from torch._inductor.fx_passes.post_grad import view_to_reshape
108 +
109 + view_to_reshape(gm)
110 +
111 + def register(self, pm_pass: PatternMatcherPass):
112 + def pattern(
113 + qkv: torch.Tensor,
114 + positions: torch.Tensor,
115 + q_weight: torch.Tensor,
116 + k_weight: torch.Tensor,
117 + cos_sin_cache: torch.Tensor,
118 + ):
119 + # split qkv -> q,k,v
120 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
121 +
122 + # Q path: view -> RMS -> view back to q.shape
123 + q_by_head = q.view(
124 + *q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim
125 + )
126 + q_normed_by_head = self.rmsnorm_matcher(q_by_head, q_weight)
127 + q_flat = q_normed_by_head.view(q.shape)
128 +
129 + # K path: view -> RMS -> view back to k.shape
130 + k_by_head = k.view(
131 + *k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim
132 + )
133 + k_normed_by_head = self.rmsnorm_matcher(k_by_head, k_weight)
134 + k_flat = k_normed_by_head.view(k.shape)
135 +
136 + # RoPE: apply to flattened q/k
137 + q_rope, k_rope = self.rope_matcher(positions, q_flat, k_flat, cos_sin_cache)
138 + return q_rope, k_rope, v
139 +
140 + def replacement(
141 + qkv: torch.Tensor,
142 + positions: torch.Tensor,
143 + q_weight: torch.Tensor,
144 + k_weight: torch.Tensor,
145 + cos_sin_cache: torch.Tensor,
146 + ):
147 + # Run fused qk_norm_rope op
148 + result = auto_functionalized(
149 + FUSED_QK_ROPE_OP,
150 + qkv=qkv,
151 + num_heads_q=self.num_heads,
152 + num_heads_k=self.num_kv_heads,
153 + num_heads_v=self.num_kv_heads,
154 + head_dim=self.head_dim,
155 + eps=self.eps,
156 + q_weight=q_weight,
157 + k_weight=k_weight,
158 + cos_sin_cache=cos_sin_cache,
159 + is_neox=self.is_neox,
160 + position_ids=positions.view(-1),
161 + )
162 + result_qkv = result[1]
163 +
164 + # Split back to q,k,v and return
165 + return result_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
166 +
167 + # NOTE: use fx_view_to_reshape to unify view/reshape to simplify
168 + # pattern and increase matching opportunities
169 + pm.register_replacement(
170 + pattern,
171 + replacement,
172 + self.get_inputs(),
173 + QkNormRopePattern.wrap_trace_fn(
174 + pm.fwd_only,
175 + QkNormRopePattern.fx_view_to_reshape,
176 + ),
177 + pm_pass,
178 + )
179 +
180 +
181 +class QKNormRoPEFusionPass(VllmPatternMatcherPass):
182 + """Fuse Q/K RMSNorm + RoPE into fused_qk_norm_rope when the custom op exists."""
183 +
184 + @enable_fake_mode
185 + def __init__(self, config: VllmConfig):
186 + super().__init__(config)
187 + self.patterns: PatternMatcherPass = PatternMatcherPass(
188 + pass_name="qk_norm_rope_fusion_pass"
189 + )
190 +
191 + dtype = config.model_config.dtype
192 + if dtype not in (torch.bfloat16, torch.float16):
193 + logger.warning_once(
194 + "QK Norm+RoPE fusion not enabled: unsupported dtype %s", dtype
195 + )
196 + return
197 +
198 + # use one attn layer to get meta (such as head_dim) for QkNormRopePattern
199 + attn_layers: dict[str, Attention] = get_layers_from_vllm_config(
200 + config, Attention
201 + )
202 + if len(attn_layers) == 0:
203 + logger.warning_once(
204 + "QK Norm+RoPE fusion enabled, but no Attention layers were discovered."
205 + )
206 + return
207 + layer = next(iter(attn_layers.values()))
208 +
209 + for epsilon in [1e-5, 1e-6]:
210 + for neox in [True, False]:
211 + if RotaryEmbedding.enabled():
212 + for rope_flashinfer in [False, True]:
213 + QkNormRopePattern(
214 + head_dim=layer.head_size,
215 + num_heads=layer.num_heads,
216 + num_kv_heads=layer.num_kv_heads,
217 + eps=epsilon,
218 + is_neox=neox,
219 + rope_flashinfer=rope_flashinfer,
220 + ).register(self.patterns)
221 + else:
222 + QkNormRopePattern(
223 + head_dim=layer.head_size,
224 + num_heads=layer.num_heads,
225 + num_kv_heads=layer.num_kv_heads,
226 + eps=epsilon,
227 + is_neox=neox,
228 + ).register(self.patterns)
229 +
230 + self.dump_patterns(config, self.patterns)
231 +
232 + @VllmInductorPass.time_and_log
233 + def __call__(self, graph: fx.Graph) -> None:
234 + self.matched_count = self.patterns.apply(graph)
235 + logger.debug("Fused QK Norm+RoPE on %s sites", self.matched_count)
236 +
237 + def uuid(self):
238 + return VllmInductorPass.hash_source(self, QkNormRopePattern)
vllm/config/compilation.py
+13
@@ -129,6 +129,8 @@ class PassConfig:
129 8: 1, # 1MB
130 },
131 }, where key is the device capability"""
132 + enable_qk_norm_rope_fusion: bool = False
133 + """Whether to enable the fused Q/K RMSNorm + RoPE pass."""
134
135 # TODO(luka) better pass enabling system.
136
@@ -182,6 +184,12 @@ class PassConfig:
184 "Fusion enabled but reshape elimination disabled. "
185 "Allreduce + rms norm + quant (fp8) fusion might not work"
186 )
187 + if self.enable_qk_norm_rope_fusion and not current_platform.is_cuda():
188 + logger.warning_once(
189 + "QK Norm + RoPE fusion enabled but the current platform is not "
190 + "CUDA. The fusion will be disabled."
191 + )
192 + self.enable_qk_norm_rope_fusion = False
193
194
195 @config
@@ -640,6 +648,11 @@ class CompilationConfig:
648 if isinstance(self.pass_config, dict):
649 self.pass_config = PassConfig(**self.pass_config)
650
651 + if self.pass_config.enable_qk_norm_rope_fusion:
652 + # TODO(zhuhaoran): support rope native forward match and remove this.
653 + # Linked issue: https://github.com/vllm-project/vllm/issues/28042
654 + self.custom_ops.append("+rotary_embedding")
655 +
656 if (
657 is_torch_equal_or_newer("2.9.0.dev")
658 and "combo_kernels" not in self.inductor_compile_config
vllm/model_executor/layers/rotary_embedding/base.py
+33 -12
@@ -98,35 +98,56 @@ class RotaryEmbedding(RotaryEmbeddingBase):
98 head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
99 )
100
101 - def forward_native(
102 - self,
101 + @staticmethod
102 + def forward_static(
103 positions: torch.Tensor,
104 query: torch.Tensor,
105 - key: torch.Tensor | None = None,
105 + key: torch.Tensor | None,
106 + head_size: int,
107 + rotary_dim: int,
108 + cos_sin_cache: torch.Tensor,
109 + is_neox_style: bool,
110 ) -> tuple[torch.Tensor, torch.Tensor | None]:
111 """A PyTorch-native implementation of forward()."""
112 positions = positions.flatten()
113 num_tokens = positions.shape[0]
110 - cos_sin = self.cos_sin_cache.index_select(0, positions)
114 + cos_sin = cos_sin_cache.index_select(0, positions)
115 cos, sin = cos_sin.chunk(2, dim=-1)
116
117 query_shape = query.shape
114 - query = query.view(num_tokens, -1, self.head_size)
115 - query_rot = query[..., : self.rotary_dim]
116 - query_pass = query[..., self.rotary_dim :]
117 - query_rot = apply_rotary_emb_torch(query_rot, cos, sin, self.is_neox_style)
118 + query = query.view(num_tokens, -1, head_size)
119 + query_rot = query[..., :rotary_dim]
120 + query_pass = query[..., rotary_dim:]
121 + query_rot = apply_rotary_emb_torch(query_rot, cos, sin, is_neox_style)
122 query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
123
124 # key may be None in some cases, e.g. cross-layer KV sharing
125 if key is not None:
126 key_shape = key.shape
123 - key = key.view(num_tokens, -1, self.head_size)
124 - key_rot = key[..., : self.rotary_dim]
125 - key_pass = key[..., self.rotary_dim :]
126 - key_rot = apply_rotary_emb_torch(key_rot, cos, sin, self.is_neox_style)
127 + key = key.view(num_tokens, -1, head_size)
128 + key_rot = key[..., :rotary_dim]
129 + key_pass = key[..., rotary_dim:]
130 + key_rot = apply_rotary_emb_torch(key_rot, cos, sin, is_neox_style)
131 key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
132 return query, key
133
134 + def forward_native(
135 + self,
136 + positions: torch.Tensor,
137 + query: torch.Tensor,
138 + key: torch.Tensor | None = None,
139 + ) -> tuple[torch.Tensor, torch.Tensor | None]:
140 + """A PyTorch-native implementation of forward()."""
141 + return self.forward_static(
142 + positions,
143 + query,
144 + key,
145 + self.head_size,
146 + self.rotary_dim,
147 + self.cos_sin_cache,
148 + self.is_neox_style,
149 + )
150 +
151 def forward_cuda(
152 self,
153 positions: torch.Tensor,