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 11, 03:45 UTC
[Kernel] Enable fused_qknorm_rope_kernel supports partial rope (#30821)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Dec 22, 2025 at 10:39 UTC
097978a15dc3757e6fdbbae6ad752b97691581bd
2 files changed
+67
-55
csrc/fused_qknorm_rope_kernel.cu
+62
-53
@@ -107,7 +107,8 @@ __global__ void fusedQKNormRopeKernel(
107
void const* k_weight_void, // RMSNorm weights for key
108
void const* cos_sin_cache_void, // Pre-computed cos/sin cache
109
int64_t const* position_ids, // Position IDs for RoPE
110
- int const num_tokens // Number of tokens
110
+ int const num_tokens, // Number of tokens
111
+ int const rotary_dim // Dimension for RoPE
112
) {
113
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
114
if constexpr ((std::is_same_v<scalar_t_in, c10::BFloat16>) ||
@@ -227,56 +228,59 @@ __global__ void fusedQKNormRopeKernel(
228
229
// Calculate cache pointer for this position - similar to
230
// pos_encoding_kernels.cu
230
- T_cache const* cache_ptr = cos_sin_cache + pos_id * head_dim;
231
- int const embed_dim = head_dim / 2;
231
+ T_cache const* cache_ptr = cos_sin_cache + pos_id * rotary_dim;
232
+ int const embed_dim = rotary_dim / 2;
233
T_cache const* cos_ptr = cache_ptr;
234
T_cache const* sin_ptr = cache_ptr + embed_dim;
234
-
235
- if constexpr (interleave) {
236
- // Perform interleaving. Use pre-computed cos/sin values.
235
+ int const rotary_lanes = rotary_dim / numElemsPerThread; // rotary range
236
+ if (laneId < rotary_lanes) {
237
+ if constexpr (interleave) {
238
+ // Perform interleaving. Use pre-computed cos/sin values.
239
#pragma unroll
238
- for (int i = 0; i < numElemsPerThread / 2; ++i) {
239
- int const idx0 = 2 * i;
240
- int const idx1 = 2 * i + 1;
241
-
242
- float const val0 = elements[idx0];
243
- float const val1 = elements[idx1];
244
-
245
- int const dim_idx = laneId * numElemsPerThread + idx0;
246
- int const half_dim = dim_idx / 2;
247
- float const cos_val =
248
- CacheConverter::convert(VLLM_LDG(cos_ptr + half_dim));
249
- float const sin_val =
250
- CacheConverter::convert(VLLM_LDG(sin_ptr + half_dim));
251
-
252
- elements[idx0] = val0 * cos_val - val1 * sin_val;
253
- elements[idx1] = val0 * sin_val + val1 * cos_val;
254
- }
255
- } else {
256
- // Before data exchange with in warp, we need to sync.
257
- __syncwarp();
258
- // Get the data from the other half of the warp. Use pre-computed cos/sin
259
- // values.
260
-#pragma unroll
261
- for (int i = 0; i < numElemsPerThread; i++) {
262
- elements2[i] = __shfl_xor_sync(FINAL_MASK, elements[i], 16);
263
- if (laneId < 16) {
264
- elements2[i] = -elements2[i];
240
+ for (int i = 0; i < numElemsPerThread / 2; ++i) {
241
+ int const idx0 = 2 * i;
242
+ int const idx1 = 2 * i + 1;
243
+ // Global dimension index in the head
244
+ int const dim_idx = laneId * numElemsPerThread + idx0;
245
+
246
+ float const val0 = elements[idx0];
247
+ float const val1 = elements[idx1];
248
+
249
+ int const half_dim = dim_idx / 2;
250
+ float const cos_val =
251
+ CacheConverter::convert(VLLM_LDG(cos_ptr + half_dim));
252
+ float const sin_val =
253
+ CacheConverter::convert(VLLM_LDG(sin_ptr + half_dim));
254
+
255
+ elements[idx0] = val0 * cos_val - val1 * sin_val;
256
+ elements[idx1] = val0 * sin_val + val1 * cos_val;
257
}
258
+ } else {
259
+ // Before data exchange with in warp, we need to sync.
260
+ __syncwarp();
261
+ int pairOffset = (rotary_dim / 2) / numElemsPerThread;
262
+ // Get the data from the other half of the warp. Use pre-computed
263
+ // cos/sin values.
264
+#pragma unroll
265
+ for (int i = 0; i < numElemsPerThread; i++) {
266
+ elements2[i] = __shfl_xor_sync(FINAL_MASK, elements[i], pairOffset);
267
267
- int dim_idx = laneId * numElemsPerThread + i;
268
- dim_idx = (dim_idx * 2) % head_dim;
269
- int half_dim = dim_idx / 2;
270
- // Use pre-computed cos/sin from cache
271
- float cos_val = CacheConverter::convert(VLLM_LDG(cos_ptr + half_dim));
272
- float sin_val = CacheConverter::convert(VLLM_LDG(sin_ptr + half_dim));
268
+ if (laneId < pairOffset) {
269
+ elements2[i] = -elements2[i];
270
+ }
271
+ int dim_idx = laneId * numElemsPerThread + i;
272
274
- elements[i] = elements[i] * cos_val + elements2[i] * sin_val;
273
+ dim_idx = (dim_idx * 2) % rotary_dim;
274
+ int half_dim = dim_idx / 2;
275
+ float cos_val = CacheConverter::convert(VLLM_LDG(cos_ptr + half_dim));
276
+ float sin_val = CacheConverter::convert(VLLM_LDG(sin_ptr + half_dim));
277
+
278
+ elements[i] = elements[i] * cos_val + elements2[i] * sin_val;
279
+ }
280
+ // __shfl_xor_sync does not provide memfence. Need to sync again.
281
+ __syncwarp();
282
}
276
- // __shfl_xor_sync does not provide memfence. Need to sync again.
277
- __syncwarp();
283
}
279
-
284
// Store.
285
{
286
vec_T vec;
@@ -312,10 +316,10 @@ template <typename scalar_t_in, typename scalar_t_cache>
316
void launchFusedQKNormRope(void* qkv, int const num_tokens,
317
int const num_heads_q, int const num_heads_k,
318
int const num_heads_v, int const head_dim,
315
- float const eps, void const* q_weight,
316
- void const* k_weight, void const* cos_sin_cache,
317
- bool const interleave, int64_t const* position_ids,
318
- cudaStream_t stream) {
319
+ int const rotary_dim, float const eps,
320
+ void const* q_weight, void const* k_weight,
321
+ void const* cos_sin_cache, bool const interleave,
322
+ int64_t const* position_ids, cudaStream_t stream) {
323
constexpr int blockSize = 256;
324
325
int const warpsPerBlock = blockSize / 32;
@@ -332,7 +336,7 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens,
336
fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 64, INTERLEAVE>
337
<<<gridDim, blockDim, 0, stream>>>(
338
qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight,
335
- k_weight, cos_sin_cache, position_ids, num_tokens);
339
+ k_weight, cos_sin_cache, position_ids, num_tokens, rotary_dim);
340
});
341
break;
342
case 128:
@@ -340,7 +344,7 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens,
344
fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 128, INTERLEAVE>
345
<<<gridDim, blockDim, 0, stream>>>(
346
qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight,
343
- k_weight, cos_sin_cache, position_ids, num_tokens);
347
+ k_weight, cos_sin_cache, position_ids, num_tokens, rotary_dim);
348
});
349
break;
350
case 256:
@@ -348,7 +352,7 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens,
352
fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 256, INTERLEAVE>
353
<<<gridDim, blockDim, 0, stream>>>(
354
qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight,
351
- k_weight, cos_sin_cache, position_ids, num_tokens);
355
+ k_weight, cos_sin_cache, position_ids, num_tokens, rotary_dim);
356
});
357
break;
358
default:
@@ -392,12 +396,16 @@ void fused_qk_norm_rope(
396
"Query weights size must match head dimension");
397
TORCH_CHECK(k_weight.size(0) == head_dim,
398
"Key weights size must match head dimension");
395
- TORCH_CHECK(cos_sin_cache.size(1) == head_dim,
396
- "Cos/sin cache dimension must match head_dim");
399
+
400
+ TORCH_CHECK(cos_sin_cache.size(1) % 2 == 0, "rotary_dim must be even");
401
+ TORCH_CHECK(cos_sin_cache.size(1) <= head_dim,
402
+ "rotary_dim must be less than or equal to head_dim");
403
+
404
TORCH_CHECK(qkv.scalar_type() == q_weight.scalar_type() &&
405
qkv.scalar_type() == k_weight.scalar_type(),
406
"qkv, q_weight and k_weight must have the same dtype");
407
408
+ int64_t rotary_dim = cos_sin_cache.size(1);
409
int64_t num_tokens = qkv.size(0);
410
TORCH_CHECK(position_ids.size(0) == num_tokens,
411
"Number of tokens in position_ids must match QKV");
@@ -419,7 +427,8 @@ void fused_qk_norm_rope(
427
qkv.data_ptr(), static_cast<int>(num_tokens),
428
static_cast<int>(num_heads_q), static_cast<int>(num_heads_k),
429
static_cast<int>(num_heads_v), static_cast<int>(head_dim),
422
- static_cast<float>(eps), q_weight.data_ptr(), k_weight.data_ptr(),
430
+ static_cast<int>(cos_sin_cache.size(1)), static_cast<float>(eps),
431
+ q_weight.data_ptr(), k_weight.data_ptr(),
432
cos_sin_cache.data_ptr(), !is_neox,
433
reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
434
stream);
tests/kernels/core/test_fused_qk_norm_rope.py
+5
-2
@@ -13,6 +13,7 @@ DTYPES = [torch.bfloat16, torch.float16]
13
IS_NEOX = [True, False]
14
EPS_VALUES = [1e-5, 1e-6]
15
SEEDS = [13]
16
+PARTIAL_ROPE = [True, False]
17
CUDA_DEVICES = ["cuda:0"]
18
19
@@ -52,6 +53,7 @@ def _apply_qk_norm_rope(
53
@pytest.mark.parametrize("is_neox", IS_NEOX)
54
@pytest.mark.parametrize("eps", EPS_VALUES)
55
@pytest.mark.parametrize("seed", SEEDS)
56
+@pytest.mark.parametrize("rotary_ratio", [1.0, 0.5, 0.25])
57
@torch.inference_mode()
58
def test_fused_qk_norm_rope_matches_reference(
59
device: str,
@@ -59,6 +61,7 @@ def test_fused_qk_norm_rope_matches_reference(
61
is_neox: bool,
62
eps: float,
63
seed: int,
64
+ rotary_ratio: float,
65
):
66
torch.set_default_device(device)
67
current_platform.seed_everything(seed)
@@ -76,10 +79,10 @@ def test_fused_qk_norm_rope_matches_reference(
79
k_norm.weight.data.normal_(mean=1.0, std=0.1)
80
q_weight = q_norm.weight.data
81
k_weight = k_norm.weight.data
79
-
82
+ rotary_dim = int(head_dim * rotary_ratio)
83
rope = RotaryEmbedding(
84
head_size=head_dim,
82
- rotary_dim=head_dim,
85
+ rotary_dim=rotary_dim,
86
max_position_embeddings=4096,
87
base=10000.0,
88
is_neox_style=is_neox,