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 5, 06:00 UTC
fused qknorm+rope kernel optimization for SM9.0 (#37376)
fused qknorm+rope kernel optimization for SM9.0 (#37376) Signed-off-by: EricccYang <yangyang4991@gmail.com> Signed-off-by: Kaicheng Yang <53411596+EricccYang@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Kaicheng Yang committed
Apr 13, 2026 at 10:58 UTC
4beeb0689cb6d86c5dbc6a8688a468892196b540
7 files changed
+497
-12
csrc/async_util.cuh
new
+100
@@ -0,0 +1,100 @@
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
+#pragma once
18
+
19
+namespace vllm {
20
+namespace cuda_async {
21
+
22
+__device__ __forceinline__ void cp_async_shared_global_16_cg(
23
+ void* smem_ptr, const void* glob_ptr) {
24
+#if defined(USE_ROCM)
25
+ *reinterpret_cast<int4*>(smem_ptr) = *reinterpret_cast<const int4*>(glob_ptr);
26
+#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
27
+ uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
28
+ asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n"
29
+ :
30
+ : "r"(smem), "l"(glob_ptr));
31
+#elif defined(__CUDA_ARCH__)
32
+ *reinterpret_cast<int4*>(smem_ptr) = *reinterpret_cast<const int4*>(glob_ptr);
33
+#else
34
+ (void)smem_ptr;
35
+ (void)glob_ptr;
36
+#endif
37
+}
38
+
39
+__device__ __forceinline__ void cp_async_shared_global_ca(void* smem_ptr,
40
+ const void* glob_ptr,
41
+ int size_bytes) {
42
+#if defined(USE_ROCM)
43
+ if (size_bytes == 4) {
44
+ *reinterpret_cast<uint32_t*>(smem_ptr) =
45
+ *reinterpret_cast<const uint32_t*>(glob_ptr);
46
+ } else if (size_bytes == 8) {
47
+ *reinterpret_cast<uint64_t*>(smem_ptr) =
48
+ *reinterpret_cast<const uint64_t*>(glob_ptr);
49
+ } else {
50
+ *reinterpret_cast<int4*>(smem_ptr) =
51
+ *reinterpret_cast<const int4*>(glob_ptr);
52
+ }
53
+#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
54
+ uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
55
+ if (size_bytes == 4) {
56
+ asm volatile("cp.async.ca.shared.global [%0], [%1], 4;\n"
57
+ :
58
+ : "r"(smem), "l"(glob_ptr));
59
+ } else if (size_bytes == 8) {
60
+ asm volatile("cp.async.ca.shared.global [%0], [%1], 8;\n"
61
+ :
62
+ : "r"(smem), "l"(glob_ptr));
63
+ } else {
64
+ asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n"
65
+ :
66
+ : "r"(smem), "l"(glob_ptr));
67
+ }
68
+#elif defined(__CUDA_ARCH__)
69
+ if (size_bytes == 4) {
70
+ *reinterpret_cast<uint32_t*>(smem_ptr) =
71
+ *reinterpret_cast<const uint32_t*>(glob_ptr);
72
+ } else if (size_bytes == 8) {
73
+ *reinterpret_cast<uint64_t*>(smem_ptr) =
74
+ *reinterpret_cast<const uint64_t*>(glob_ptr);
75
+ } else {
76
+ *reinterpret_cast<int4*>(smem_ptr) =
77
+ *reinterpret_cast<const int4*>(glob_ptr);
78
+ }
79
+#else
80
+ (void)smem_ptr;
81
+ (void)glob_ptr;
82
+ (void)size_bytes;
83
+#endif
84
+}
85
+
86
+__device__ __forceinline__ void cp_async_commit_group() {
87
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 && !defined(USE_ROCM)
88
+ asm volatile("cp.async.commit_group;\n" ::);
89
+#endif
90
+}
91
+
92
+template <int n>
93
+__device__ __forceinline__ void cp_async_wait_group() {
94
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 && !defined(USE_ROCM)
95
+ asm volatile("cp.async.wait_group %0;\n" : : "n"(n));
96
+#endif
97
+}
98
+
99
+} // namespace cuda_async
100
+} // namespace vllm
csrc/fused_qknorm_rope_kernel.cu
+388
-9
@@ -19,8 +19,10 @@
19
#include <type_traits>
20
21
#include <torch/cuda.h>
22
+#include <ATen/cuda/CUDAContext.h>
23
#include <c10/cuda/CUDAGuard.h>
24
25
+#include "async_util.cuh"
26
#include "cuda_compat.h"
27
#include "dispatch_utils.h"
28
#include "type_convert.cuh"
@@ -86,6 +88,9 @@ inline __device__ __host__ T divUp(T m, T n) {
88
} // namespace tensorrt_llm::common
89
90
namespace tensorrt_llm::kernels {
91
+
92
+using namespace vllm::cuda_async;
93
+
94
// NOTE(zhuhaoran): This kernel is adapted from TensorRT-LLM implementation,
95
// with added support for passing the cos_sin_cache as an input.
96
// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu
@@ -301,6 +306,237 @@ __global__ void fusedQKNormRopeKernel(
306
#endif
307
}
308
309
+// Multi-token-head kernel: one warp processes HEADS_PER_WARP token-heads for
310
+// the same token, sharing cos/sin from shared memory via cp.async.
311
+// When HEADS_PER_WARP > 1 the warp reuses the loaded cos/sin across all heads,
312
+// hiding global-memory latency and improving occupancy for large batches.
313
+template <typename scalar_t_in, typename scalar_t_cache, int head_dim,
314
+ bool interleave, int HEADS_PER_WARP>
315
+__global__ void fusedQKNormRopeKernelNTokenHeads(
316
+ void* qkv_void, int const num_heads_q, int const num_heads_k,
317
+ int const num_heads_v, float const eps, void const* q_weight_void,
318
+ void const* k_weight_void, void const* cos_sin_cache_void,
319
+ int64_t const* position_ids, int const num_tokens, int const rotary_dim) {
320
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
321
+ if constexpr ((std::is_same_v<scalar_t_in, c10::BFloat16>) ||
322
+ std::is_same_v<scalar_t_cache, c10::BFloat16>) {
323
+ return;
324
+ } else {
325
+#endif
326
+
327
+ using Converter = vllm::_typeConvert<scalar_t_in>;
328
+ static_assert(Converter::exists,
329
+ "Input QKV data type is not supported for this CUDA "
330
+ "architecture or toolkit version.");
331
+ using T_in = typename Converter::hip_type;
332
+ using T2_in = typename Converter::packed_hip_type;
333
+
334
+ using CacheConverter = vllm::_typeConvert<scalar_t_cache>;
335
+ static_assert(CacheConverter::exists,
336
+ "Cache data type is not supported for this CUDA architecture "
337
+ "or toolkit version.");
338
+ using T_cache = typename CacheConverter::hip_type;
339
+
340
+ extern __shared__ char smem_storage[];
341
+ // Shared memory layout:
342
+ // [0, cos_sin_bytes) : cos/sin for each warp (warpsPerBlock *
343
+ // rotary_dim * sizeof(T_cache))
344
+ // [cos_sin_bytes, ...) : QKV tiles
345
+ // per warp (warpsPerBlock * HEADS_PER_WARP * 32 * elemSizeBytes)
346
+ T_cache* const smem = reinterpret_cast<T_cache*>(smem_storage);
347
+
348
+ T_in* qkv = reinterpret_cast<T_in*>(qkv_void);
349
+ T_in const* q_weight = reinterpret_cast<T_in const*>(q_weight_void);
350
+ T_in const* k_weight = reinterpret_cast<T_in const*>(k_weight_void);
351
+ T_cache const* cos_sin_cache =
352
+ reinterpret_cast<T_cache const*>(cos_sin_cache_void);
353
+
354
+ int const warpsPerBlock = blockDim.x / 32;
355
+ int const warpId = threadIdx.x / 32;
356
+ int const laneId = threadIdx.x % 32;
357
+
358
+ int const total_qk_heads = num_heads_q + num_heads_k;
359
+ int const num_heads = num_heads_q + num_heads_k + num_heads_v;
360
+ int const head_chunks_per_token =
361
+ (total_qk_heads + HEADS_PER_WARP - 1) / HEADS_PER_WARP;
362
+
363
+ int const warp_global = blockIdx.x * warpsPerBlock + warpId;
364
+ int const tokenIdx = warp_global / head_chunks_per_token;
365
+ int const headChunk = warp_global % head_chunks_per_token;
366
+ int const first_head = headChunk * HEADS_PER_WARP;
367
+ int const num_heads_this_warp =
368
+ (first_head + HEADS_PER_WARP <= total_qk_heads)
369
+ ? HEADS_PER_WARP
370
+ : (total_qk_heads - first_head);
371
+
372
+ if (tokenIdx >= num_tokens) return;
373
+
374
+ static_assert(head_dim % (32 * 2) == 0, "head_dim must be divisible by 64");
375
+ constexpr int numElemsPerThread = head_dim / 32;
376
+ constexpr int elemSizeBytes = numElemsPerThread * sizeof(__nv_bfloat16);
377
+ static_assert(elemSizeBytes % 4 == 0,
378
+ "elemSizeBytes must be a multiple of 4");
379
+ constexpr int vecSize = elemSizeBytes / 4;
380
+ using vec_T = typename tensorrt_llm::common::packed_as<uint, vecSize>::type;
381
+
382
+ int const cos_sin_bytes =
383
+ warpsPerBlock * rotary_dim * static_cast<int>(sizeof(T_cache));
384
+ int const qkv_tile_bytes = 32 * elemSizeBytes;
385
+ char* const this_warp_head_smem =
386
+ smem_storage + cos_sin_bytes +
387
+ warpId * (HEADS_PER_WARP * qkv_tile_bytes);
388
+
389
+ // === Group 0: async load all heads' QKV into smem (issued first). ===
390
+ for (int k = 0; k < num_heads_this_warp; ++k) {
391
+ int const localHeadIdx = first_head + k;
392
+ bool const isQ = localHeadIdx < num_heads_q;
393
+ int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q;
394
+ int offWarp;
395
+ if (isQ) {
396
+ offWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim;
397
+ } else {
398
+ offWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim +
399
+ headIdx * head_dim;
400
+ }
401
+ int const offThread = offWarp + laneId * numElemsPerThread;
402
+ char* smem_dst =
403
+ this_warp_head_smem + k * qkv_tile_bytes + laneId * elemSizeBytes;
404
+ cp_async_shared_global_ca(smem_dst,
405
+ reinterpret_cast<const char*>(&qkv[offThread]),
406
+ elemSizeBytes);
407
+ }
408
+ cp_async_commit_group(); // commit group 0 (QKV)
409
+
410
+ // === Group 1: async load cos/sin into smem (issued second). ===
411
+ int64_t const pos_id = position_ids[tokenIdx];
412
+ T_cache const* const cache_ptr = cos_sin_cache + pos_id * rotary_dim;
413
+ int const copy_bytes = rotary_dim * static_cast<int>(sizeof(T_cache));
414
+ int const num_copies = (copy_bytes + 15) / 16;
415
+ for (int copyId = laneId; copyId < num_copies; copyId += 32) {
416
+ char* smem_ptr =
417
+ reinterpret_cast<char*>(&smem[warpId * rotary_dim]) + copyId * 16;
418
+ const char* glob_ptr =
419
+ reinterpret_cast<const char*>(cache_ptr) + copyId * 16;
420
+ cp_async_shared_global_16_cg(smem_ptr, glob_ptr);
421
+ }
422
+ cp_async_commit_group(); // commit group 1 (cos/sin)
423
+
424
+ // wait<1>: allow at most 1 pending group (group 1) → group 0 (QKV) is done.
425
+ cp_async_wait_group<1>();
426
+
427
+ float elements[numElemsPerThread];
428
+ float elements2[numElemsPerThread];
429
+ int const rotary_lanes = rotary_dim / numElemsPerThread;
430
+ int const embed_dim = rotary_dim / 2;
431
+ T_cache const* const cos_smem = &smem[warpId * rotary_dim];
432
+ T_cache const* const sin_smem = &smem[warpId * rotary_dim + embed_dim];
433
+
434
+ // Preload weights into registers once, reused across all heads.
435
+ float q_w[numElemsPerThread];
436
+ float k_w[numElemsPerThread];
437
+#pragma unroll
438
+ for (int i = 0; i < numElemsPerThread; i++) {
439
+ int const dim = laneId * numElemsPerThread + i;
440
+ q_w[i] = Converter::convert(q_weight[dim]);
441
+ k_w[i] = Converter::convert(k_weight[dim]);
442
+ }
443
+
444
+ for (int k = 0; k < num_heads_this_warp; ++k) {
445
+ int const localHeadIdx = first_head + k;
446
+ bool const isQ = localHeadIdx < num_heads_q;
447
+ int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q;
448
+
449
+ int offsetWarp;
450
+ if (isQ) {
451
+ offsetWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim;
452
+ } else {
453
+ offsetWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim +
454
+ headIdx * head_dim;
455
+ }
456
+ int const offsetThread = offsetWarp + laneId * numElemsPerThread;
457
+
458
+ // === Part 1: QK Norm (read from smem; group 0 already done). ===
459
+ float sumOfSquares = 0.0f;
460
+ {
461
+ char const* smem_src =
462
+ this_warp_head_smem + k * qkv_tile_bytes + laneId * elemSizeBytes;
463
+ vec_T vec = *reinterpret_cast<vec_T const*>(smem_src);
464
+ constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
465
+#pragma unroll
466
+ for (int i = 0; i < num_packed_elems; i++) {
467
+ T2_in packed_val = *(reinterpret_cast<T2_in*>(&vec) + i);
468
+ float2 vals = Converter::convert(packed_val);
469
+ sumOfSquares += vals.x * vals.x;
470
+ sumOfSquares += vals.y * vals.y;
471
+ elements[2 * i] = vals.x;
472
+ elements[2 * i + 1] = vals.y;
473
+ }
474
+ }
475
+
476
+ sumOfSquares = tensorrt_llm::common::warpReduceSum(sumOfSquares);
477
+ float rms_rcp = rsqrtf(sumOfSquares / static_cast<float>(head_dim) + eps);
478
+
479
+#pragma unroll
480
+ for (int i = 0; i < numElemsPerThread; i++) {
481
+ elements[i] *= rms_rcp * (isQ ? q_w[i] : k_w[i]);
482
+ }
483
+
484
+ // On first head: wait for group 1 (cos/sin) before RoPE.
485
+ if (k == 0) cp_async_wait_group<0>();
486
+
487
+ // === Part 2: RoPE using cos/sin from shared memory. ===
488
+ if (laneId < rotary_lanes) {
489
+ if constexpr (interleave) {
490
+#pragma unroll
491
+ for (int i = 0; i < numElemsPerThread / 2; ++i) {
492
+ int const idx0 = 2 * i;
493
+ int const idx1 = 2 * i + 1;
494
+ int const dim_idx = laneId * numElemsPerThread + idx0;
495
+ float const val0 = elements[idx0];
496
+ float const val1 = elements[idx1];
497
+ int const half_dim = dim_idx / 2;
498
+ float const cos_val = CacheConverter::convert(cos_smem[half_dim]);
499
+ float const sin_val = CacheConverter::convert(sin_smem[half_dim]);
500
+ elements[idx0] = val0 * cos_val - val1 * sin_val;
501
+ elements[idx1] = val0 * sin_val + val1 * cos_val;
502
+ }
503
+ } else {
504
+ __syncwarp();
505
+ int const pairOffset = (rotary_dim / 2) / numElemsPerThread;
506
+#pragma unroll
507
+ for (int i = 0; i < numElemsPerThread; i++) {
508
+ elements2[i] = __shfl_xor_sync(FINAL_MASK, elements[i], pairOffset);
509
+ if (laneId < pairOffset) elements2[i] = -elements2[i];
510
+ int dim_idx = laneId * numElemsPerThread + i;
511
+ dim_idx = (dim_idx * 2) % rotary_dim;
512
+ int const half_dim = dim_idx / 2;
513
+ float const cos_val = CacheConverter::convert(cos_smem[half_dim]);
514
+ float const sin_val = CacheConverter::convert(sin_smem[half_dim]);
515
+ elements[i] = elements[i] * cos_val + elements2[i] * sin_val;
516
+ }
517
+ __syncwarp();
518
+ }
519
+ }
520
+
521
+ // Store.
522
+ {
523
+ vec_T vec;
524
+ constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
525
+#pragma unroll
526
+ for (int i = 0; i < num_packed_elems; i++) {
527
+ T2_in packed_val = Converter::convert(
528
+ make_float2(elements[2 * i], elements[2 * i + 1]));
529
+ *(reinterpret_cast<T2_in*>(&vec) + i) = packed_val;
530
+ }
531
+ *reinterpret_cast<vec_T*>(&qkv[offsetThread]) = vec;
532
+ }
533
+ }
534
+
535
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
536
+ }
537
+#endif
538
+}
539
+
540
// Borrowed from
541
// https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568
542
#define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \
@@ -321,15 +557,12 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens,
557
void const* cos_sin_cache, bool const interleave,
558
int64_t const* position_ids, cudaStream_t stream) {
559
constexpr int blockSize = 256;
324
-
560
int const warpsPerBlock = blockSize / 32;
561
int const totalQKHeads = num_heads_q + num_heads_k;
562
int const totalWarps = num_tokens * totalQKHeads;
328
-
563
int const gridSize = common::divUp(totalWarps, warpsPerBlock);
564
dim3 gridDim(gridSize);
565
dim3 blockDim(blockSize);
332
-
566
switch (head_dim) {
567
case 64:
568
DISPATCH_INTERLEAVE(interleave, INTERLEAVE, {
@@ -360,6 +593,118 @@ void launchFusedQKNormRope(void* qkv, int const num_tokens,
593
"Unsupported head dimension for fusedQKNormRope: ", head_dim);
594
}
595
}
596
+
597
+// Launch: one warp processes token_heads_per_warp token-heads (1, 2, 4, or 8).
598
+// When token_heads_per_warp == 1, delegates to the 1-head baseline above.
599
+template <typename scalar_t_in, typename scalar_t_cache>
600
+void launchFusedQKNormRopeNTokenHeads(
601
+ void* qkv, int const num_tokens, int const num_heads_q,
602
+ int const num_heads_k, int const num_heads_v, int const head_dim,
603
+ int const rotary_dim, float const eps, void const* q_weight,
604
+ void const* k_weight, void const* cos_sin_cache, bool const interleave,
605
+ int64_t const* position_ids, int const token_heads_per_warp,
606
+ cudaStream_t stream) {
607
+ TORCH_CHECK(token_heads_per_warp == 1 || token_heads_per_warp == 2 ||
608
+ token_heads_per_warp == 4 || token_heads_per_warp == 8,
609
+ "token_heads_per_warp must be 1, 2, 4, or 8, got ",
610
+ token_heads_per_warp);
611
+
612
+ // token_heads_per_warp == 1: delegate to the 1-head baseline kernel.
613
+ if (token_heads_per_warp == 1) {
614
+ launchFusedQKNormRope<scalar_t_in, scalar_t_cache>(
615
+ qkv, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim,
616
+ rotary_dim, eps, q_weight, k_weight, cos_sin_cache, interleave,
617
+ position_ids, stream);
618
+ return;
619
+ }
620
+
621
+ // NTokenHeads kernel uses cp.async to load cos/sin in 16-byte chunks.
622
+ // If rotary_dim * sizeof(cache_dtype) is not a multiple of 16, the last
623
+ // cp.async would write past the shared memory allocation.
624
+ // Fall back to the base kernel instead of failing.
625
+ {
626
+ size_t const rotary_bytes =
627
+ static_cast<size_t>(rotary_dim) *
628
+ (std::is_same_v<scalar_t_cache, float> ? sizeof(float) : 2u);
629
+ if (rotary_bytes % 16 != 0) {
630
+ launchFusedQKNormRope<scalar_t_in, scalar_t_cache>(
631
+ qkv, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim,
632
+ rotary_dim, eps, q_weight, k_weight, cos_sin_cache, interleave,
633
+ position_ids, stream);
634
+ return;
635
+ }
636
+ }
637
+
638
+ constexpr int blockSize = 256;
639
+ int const warpsPerBlock = blockSize / 32;
640
+ int const totalQKHeads = num_heads_q + num_heads_k;
641
+ // Grid: one warp per (token, head_chunk); same token → reuse cos/sin in smem.
642
+ int const head_chunks_per_token =
643
+ (totalQKHeads + token_heads_per_warp - 1) / token_heads_per_warp;
644
+ int const total_warps = num_tokens * head_chunks_per_token;
645
+ int const gridSize = common::divUp(total_warps, warpsPerBlock);
646
+ dim3 gridDim(gridSize);
647
+ dim3 blockDim(blockSize);
648
+ // Cache element size: float=4, bfloat16=2 (host-safe; kernel uses same
649
+ // layout).
650
+ size_t const cache_elem_size =
651
+ std::is_same_v<scalar_t_cache, float> ? sizeof(float) : 2u;
652
+ // QKV smem: token_heads_per_warp tiles per warp, each tile 32*(head_dim/32*2)
653
+ // = 2*head_dim bytes.
654
+ size_t const qkv_smem_per_warp = static_cast<size_t>(token_heads_per_warp) *
655
+ 2u * static_cast<size_t>(head_dim);
656
+ size_t const smem_bytes =
657
+ warpsPerBlock * static_cast<size_t>(rotary_dim) * cache_elem_size +
658
+ warpsPerBlock * qkv_smem_per_warp;
659
+
660
+#define LAUNCH_N_TOKEN_HEADS(N) \
661
+ do { \
662
+ switch (head_dim) { \
663
+ case 64: \
664
+ DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \
665
+ fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, 64, \
666
+ INTERLEAVE, (N)> \
667
+ <<<gridDim, blockDim, smem_bytes, stream>>>( \
668
+ qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \
669
+ k_weight, cos_sin_cache, position_ids, num_tokens, \
670
+ rotary_dim); \
671
+ }); \
672
+ break; \
673
+ case 128: \
674
+ DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \
675
+ fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, 128, \
676
+ INTERLEAVE, (N)> \
677
+ <<<gridDim, blockDim, smem_bytes, stream>>>( \
678
+ qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \
679
+ k_weight, cos_sin_cache, position_ids, num_tokens, \
680
+ rotary_dim); \
681
+ }); \
682
+ break; \
683
+ case 256: \
684
+ DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \
685
+ fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, 256, \
686
+ INTERLEAVE, (N)> \
687
+ <<<gridDim, blockDim, smem_bytes, stream>>>( \
688
+ qkv, num_heads_q, num_heads_k, num_heads_v, eps, q_weight, \
689
+ k_weight, cos_sin_cache, position_ids, num_tokens, \
690
+ rotary_dim); \
691
+ }); \
692
+ break; \
693
+ default: \
694
+ TORCH_CHECK(false, "Unsupported head dimension: ", head_dim); \
695
+ } \
696
+ } while (0)
697
+
698
+ if (token_heads_per_warp == 2) {
699
+ LAUNCH_N_TOKEN_HEADS(2);
700
+ } else if (token_heads_per_warp == 4) {
701
+ LAUNCH_N_TOKEN_HEADS(4);
702
+ } else if (token_heads_per_warp == 8) {
703
+ LAUNCH_N_TOKEN_HEADS(8);
704
+ }
705
+#undef LAUNCH_N_TOKEN_HEADS
706
+}
707
+
708
} // namespace tensorrt_llm::kernels
709
710
void fused_qk_norm_rope(
@@ -374,7 +719,8 @@ void fused_qk_norm_rope(
719
torch::Tensor& k_weight, // RMSNorm weights for key [head_dim]
720
torch::Tensor& cos_sin_cache, // Cos/sin cache [max_position, head_dim]
721
bool is_neox, // Whether RoPE is applied in Neox style
377
- torch::Tensor& position_ids // Position IDs for RoPE [num_tokens]
722
+ torch::Tensor& position_ids, // Position IDs for RoPE [num_tokens]
723
+ int64_t forced_token_heads_per_warp // -1 = auto-select, >0 = forced value
724
) {
725
// Input validation
726
CHECK_INPUT(qkv);
@@ -414,15 +760,48 @@ void fused_qk_norm_rope(
760
qkv.size(1) == total_heads * head_dim,
761
"QKV tensor size must match total number of heads and head dimension");
762
417
- auto stream = at::cuda::getCurrentCUDAStream(qkv.get_device());
763
+ auto device_id = qkv.get_device();
764
+ auto stream = at::cuda::getCurrentCUDAStream(device_id);
765
+
766
+ // Select token_heads_per_warp: forced value if >0, else auto-select.
767
+ // Auto thresholds are calibrated on SM 9.0 (H100). On other architectures,
768
+ // fall back to token_heads_per_warp=1 (base kernel) until profiled.
769
+ int token_heads_per_warp;
770
+ if (forced_token_heads_per_warp > 0) { // only support SM80+
771
+ token_heads_per_warp = static_cast<int>(forced_token_heads_per_warp);
772
+ } else {
773
+ token_heads_per_warp = 1;
774
+ auto* dev_prop = at::cuda::getDeviceProperties(device_id);
775
+ int sm_version = dev_prop->major * 10 + dev_prop->minor;
776
+ int64_t total_qk_units = num_tokens * (num_heads_q + num_heads_k);
777
+ if (sm_version == 90) {
778
+ if (head_dim >= 256) {
779
+ if (total_qk_units < 4096LL) {
780
+ token_heads_per_warp = 1;
781
+ } else if (total_qk_units < 8192LL) {
782
+ token_heads_per_warp = 2;
783
+ } else {
784
+ token_heads_per_warp = 4;
785
+ }
786
+ } else {
787
+ if (total_qk_units < 10240LL) {
788
+ token_heads_per_warp = 1;
789
+ } else if (total_qk_units < 40960LL) {
790
+ token_heads_per_warp = 4;
791
+ } else {
792
+ token_heads_per_warp = 8;
793
+ }
794
+ }
795
+ }
796
+ }
797
798
VLLM_DISPATCH_HALF_TYPES(qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
799
using qkv_scalar_t = scalar_t;
800
VLLM_DISPATCH_FLOATING_TYPES(
801
cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] {
802
using cache_scalar_t = scalar_t;
424
- tensorrt_llm::kernels::launchFusedQKNormRope<qkv_scalar_t,
425
- cache_scalar_t>(
803
+ tensorrt_llm::kernels::launchFusedQKNormRopeNTokenHeads<
804
+ qkv_scalar_t, cache_scalar_t>(
805
qkv.data_ptr(), static_cast<int>(num_tokens),
806
static_cast<int>(num_heads_q), static_cast<int>(num_heads_k),
807
static_cast<int>(num_heads_v), static_cast<int>(head_dim),
@@ -430,7 +809,7 @@ void fused_qk_norm_rope(
809
q_weight.data_ptr(), k_weight.data_ptr(),
810
cos_sin_cache.data_ptr(), !is_neox,
811
reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
433
- stream);
812
+ token_heads_per_warp, stream);
813
});
814
});
436
-}
\ No newline at end of file
815
+}
csrc/ops.h
+3
-2
@@ -96,7 +96,8 @@ 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);
99
+ bool is_neox, torch::Tensor& position_ids,
100
+ int64_t forced_token_heads_per_warp);
101
102
void apply_repetition_penalties_(torch::Tensor& logits,
103
const torch::Tensor& prompt_mask,
@@ -320,4 +321,4 @@ std::tuple<torch::Tensor, torch::Tensor> minimax_allreduce_rms_qk(
321
torch::Tensor const& norm_weight_k, torch::Tensor workspace,
322
int64_t const q_size, int64_t const kv_size, int64_t const rank,
323
int64_t const nranks, double const eps);
323
-#endif
\ No newline at end of file
324
+#endif
csrc/torch_bindings.cpp
+2
-1
@@ -173,7 +173,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
173
"fused_qk_norm_rope(Tensor! qkv, int num_heads_q, "
174
"int num_heads_k, int num_heads_v, int head_dim, float eps, "
175
"Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, "
176
- "bool is_neox, Tensor position_ids) -> ()");
176
+ "bool is_neox, Tensor position_ids, "
177
+ "int forced_token_heads_per_warp=-1) -> ()");
178
ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
179
180
// Apply repetition penalties to logits in-place
vllm/_custom_ops.py
+2
@@ -435,6 +435,7 @@ def fused_qk_norm_rope(
435
cos_sin_cache: torch.Tensor,
436
is_neox: bool,
437
position_ids: torch.Tensor,
438
+ forced_token_heads_per_warp: int = -1,
439
) -> None:
440
torch.ops._C.fused_qk_norm_rope(
441
qkv,
@@ -448,6 +449,7 @@ def fused_qk_norm_rope(
449
cos_sin_cache,
450
is_neox,
451
position_ids,
452
+ forced_token_heads_per_warp,
453
)
454
455
vllm/compilation/passes/fusion/qk_norm_rope_fusion.py
+1
@@ -164,6 +164,7 @@ class QkNormRopePattern:
164
cos_sin_cache=cos_sin_cache,
165
is_neox=self.is_neox,
166
position_ids=positions.view(-1),
167
+ forced_token_heads_per_warp=-1,
168
)
169
result_qkv = result[1]
170
vllm/compilation/passes/utility/fix_functionalization.py
+1
@@ -168,6 +168,7 @@ class FixFunctionalizationPass(VllmInductorPass):
168
"cos_sin_cache",
169
"is_neox",
170
"position_ids",
171
+ "forced_token_heads_per_warp",
172
)
173
self.defunctionalize(graph, node, mutated_args=mutated_args, args=args)
174
elif (