| 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_utils.h" |
| 22 | |
| 23 | #include "async_util.cuh" |
| 24 | #include "../cuda_compat.h" |
| 25 | #include "type_convert.cuh" |
| 26 | #include "dispatch_utils.h" |
| 27 | |
| 28 | #define CHECK_TYPE(x, st) \ |
| 29 | STD_TORCH_CHECK(x.scalar_type() == st, #x " dtype is ", x.scalar_type(), \ |
| 30 | ", while ", st, " is expected") |
| 31 | #define CHECK_TH_CUDA(x) \ |
| 32 | STD_TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor") |
| 33 | #define CHECK_CONTIGUOUS(x) \ |
| 34 | STD_TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") |
| 35 | #define CHECK_INPUT(x) \ |
| 36 | CHECK_TH_CUDA(x); \ |
| 37 | CHECK_CONTIGUOUS(x) |
| 38 | |
| 39 | #ifdef USE_ROCM |
| 40 | #define FINAL_MASK 0xffffffffffffffffULL |
| 41 | |
| 42 | #if defined(HIP_VERSION) && HIP_VERSION < 70000000 |
| 43 | // On ROCm versions before 7.0, __syncwarp isn't defined. The below |
| 44 | // implementation is copy/pasted from the implementation in ROCm 7.0 |
| 45 | __device__ inline void __syncwarp() { |
| 46 | __builtin_amdgcn_fence(__ATOMIC_RELEASE, "wavefront"); |
| 47 | __builtin_amdgcn_wave_barrier(); |
| 48 | __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "wavefront"); |
| 49 | } |
| 50 | #endif |
| 51 | #else |
| 52 | #define FINAL_MASK 0xffffffff |
| 53 | #endif |
| 54 | |
| 55 | namespace tensorrt_llm::common { |
| 56 | template <typename T, int num> |
| 57 | struct packed_as; |
| 58 | // Specialization for packed_as used in this kernel. |
| 59 | template <> |
| 60 | struct packed_as<uint, 1> { |
| 61 | using type = uint; |
| 62 | }; |
| 63 | |
| 64 | template <> |
| 65 | struct packed_as<uint, 2> { |
| 66 | using type = uint2; |
| 67 | }; |
| 68 | |
| 69 | template <> |
| 70 | struct packed_as<uint, 4> { |
| 71 | using type = uint4; |
| 72 | }; |
| 73 | |
| 74 | template <typename T> |
| 75 | __inline__ __device__ T warpReduceSum(T val) { |
| 76 | #pragma unroll |
| 77 | for (int mask = 16; mask > 0; mask >>= 1) |
| 78 | val += __shfl_xor_sync(FINAL_MASK, val, mask, 32); |
| 79 | return val; |
| 80 | } |
| 81 | |
| 82 | template <typename T> |
| 83 | inline __device__ __host__ T divUp(T m, T n) { |
| 84 | return (m + n - 1) / n; |
| 85 | } |
| 86 | |
| 87 | } // namespace tensorrt_llm::common |
| 88 | |
| 89 | namespace tensorrt_llm::kernels { |
| 90 | |
| 91 | using namespace vllm::cuda_async; |
| 92 | |
| 93 | // NOTE(zhuhaoran): This kernel is adapted from TensorRT-LLM implementation, |
| 94 | // with added support for passing the cos_sin_cache as an input. |
| 95 | // https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/fusedQKNormRopeKernel.cu |
| 96 | |
| 97 | // Perform per-head QK Norm and RoPE in a single kernel. |
| 98 | // scalar_t_in: data type of QKV and RMSNorm weights |
| 99 | // scalar_t_cache: data type of cos/sin cache |
| 100 | // head_dim: the dimension of each head |
| 101 | // interleave: interleave=!is_neox. |
| 102 | template <typename scalar_t_in, typename scalar_t_cache, int head_dim, |
| 103 | bool interleave> |
| 104 | __global__ void fusedQKNormRopeKernel( |
| 105 | void* qkv_void, // Combined QKV tensor |
| 106 | int const num_heads_q, // Number of query heads |
| 107 | int const num_heads_k, // Number of key heads |
| 108 | int const num_heads_v, // Number of value heads |
| 109 | float const eps, // Epsilon for RMS normalization |
| 110 | void const* q_weight_void, // RMSNorm weights for query |
| 111 | void const* k_weight_void, // RMSNorm weights for key |
| 112 | void const* cos_sin_cache_void, // Pre-computed cos/sin cache |
| 113 | int64_t const* position_ids, // Position IDs for RoPE |
| 114 | int const num_tokens, // Number of tokens |
| 115 | int const rotary_dim // Dimension for RoPE |
| 116 | ) { |
| 117 | #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) |
| 118 | if constexpr ((std::is_same_v<scalar_t_in, c10::BFloat16>) || |
| 119 | std::is_same_v<scalar_t_cache, c10::BFloat16>) { |
| 120 | return; |
| 121 | } else { |
| 122 | #endif |
| 123 | |
| 124 | using Converter = vllm::_typeConvert<scalar_t_in>; |
| 125 | static_assert(Converter::exists, |
| 126 | "Input QKV data type is not supported for this CUDA " |
| 127 | "architecture or toolkit version."); |
| 128 | using T_in = typename Converter::hip_type; |
| 129 | using T2_in = typename Converter::packed_hip_type; |
| 130 | |
| 131 | using CacheConverter = vllm::_typeConvert<scalar_t_cache>; |
| 132 | static_assert(CacheConverter::exists, |
| 133 | "Cache data type is not supported for this CUDA architecture " |
| 134 | "or toolkit version."); |
| 135 | using T_cache = typename CacheConverter::hip_type; |
| 136 | |
| 137 | T_in* qkv = reinterpret_cast<T_in*>(qkv_void); |
| 138 | T_in const* q_weight = reinterpret_cast<T_in const*>(q_weight_void); |
| 139 | T_in const* k_weight = reinterpret_cast<T_in const*>(k_weight_void); |
| 140 | T_cache const* cos_sin_cache = |
| 141 | reinterpret_cast<T_cache const*>(cos_sin_cache_void); |
| 142 | |
| 143 | int const warpsPerBlock = blockDim.x / 32; |
| 144 | int const warpId = threadIdx.x / 32; |
| 145 | int const laneId = threadIdx.x % 32; |
| 146 | |
| 147 | // Calculate global warp index to determine which head/token this warp |
| 148 | // processes |
| 149 | int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId; |
| 150 | |
| 151 | // Total number of attention heads (Q and K) |
| 152 | int const total_qk_heads = num_heads_q + num_heads_k; |
| 153 | |
| 154 | // Determine which token and head type (Q or K) this warp processes |
| 155 | int const tokenIdx = globalWarpIdx / total_qk_heads; |
| 156 | int const localHeadIdx = globalWarpIdx % total_qk_heads; |
| 157 | |
| 158 | // Skip if this warp is assigned beyond the number of tokens |
| 159 | if (tokenIdx >= num_tokens) return; |
| 160 | |
| 161 | bool const isQ = localHeadIdx < num_heads_q; |
| 162 | int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q; |
| 163 | |
| 164 | int const num_heads = num_heads_q + num_heads_k + num_heads_v; |
| 165 | |
| 166 | static_assert(head_dim % (32 * 2) == 0, |
| 167 | "head_dim must be divisible by 64 (each warp processes one " |
| 168 | "head, and each thread gets even number of " |
| 169 | "elements)"); |
| 170 | constexpr int numElemsPerThread = head_dim / 32; |
| 171 | float elements[numElemsPerThread]; |
| 172 | constexpr int elemSizeBytes = numElemsPerThread * sizeof(__nv_bfloat16); |
| 173 | static_assert(elemSizeBytes % 4 == 0, |
| 174 | "numSizeBytes must be a multiple of 4"); |
| 175 | constexpr int vecSize = |
| 176 | elemSizeBytes / |
| 177 | 4; // Use packed_as<uint, vecSize> to perform loading/saving. |
| 178 | using vec_T = typename tensorrt_llm::common::packed_as<uint, vecSize>::type; |
| 179 | |
| 180 | int offsetWarp; // Offset for the warp |
| 181 | if (isQ) { |
| 182 | // Q segment: token offset + head offset within Q segment |
| 183 | offsetWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim; |
| 184 | } else { |
| 185 | // K segment: token offset + entire Q segment + head offset within K |
| 186 | // segment |
| 187 | offsetWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim + |
| 188 | headIdx * head_dim; |
| 189 | } |
| 190 | int offsetThread = offsetWarp + laneId * numElemsPerThread; |
| 191 | |
| 192 | // PDL: wait for the predecessor kernel (the QKV projection that produces |
| 193 | // `qkv`) to complete before touching any global memory. |
| 194 | #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) |
| 195 | cudaGridDependencySynchronize(); |
| 196 | #endif |
| 197 | |
| 198 | // Sum of squares for RMSNorm |
| 199 | float sumOfSquares = 0.0f; |
| 200 | |
| 201 | // Load. |
| 202 | { |
| 203 | vec_T vec = *reinterpret_cast<vec_T const*>(&qkv[offsetThread]); |
| 204 | constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in); |
| 205 | #pragma unroll |
| 206 | for (int i = 0; i < num_packed_elems; i++) { |
| 207 | // Interpret the generic vector chunk as the specific packed type |
| 208 | T2_in packed_val = *(reinterpret_cast<T2_in*>(&vec) + i); |
| 209 | // Convert to float2 for computation |
| 210 | float2 vals = Converter::convert(packed_val); |
| 211 | sumOfSquares += vals.x * vals.x; |
| 212 | sumOfSquares += vals.y * vals.y; |
| 213 | |
| 214 | elements[2 * i] = vals.x; |
| 215 | elements[2 * i + 1] = vals.y; |
| 216 | } |
| 217 | } |
| 218 | |
| 219 | // Reduce sum across warp using the utility function |
| 220 | sumOfSquares = tensorrt_llm::common::warpReduceSum(sumOfSquares); |
| 221 | |
| 222 | // Compute RMS normalization factor |
| 223 | float rms_rcp = rsqrtf(sumOfSquares / static_cast<float>(head_dim) + eps); |
| 224 | |
| 225 | // Normalize elements |
| 226 | #pragma unroll |
| 227 | for (int i = 0; i < numElemsPerThread; i++) { |
| 228 | int dim = laneId * numElemsPerThread + i; |
| 229 | float weight = isQ ? Converter::convert(q_weight[dim]) |
| 230 | : Converter::convert(k_weight[dim]); |
| 231 | elements[i] *= rms_rcp * weight; |
| 232 | } |
| 233 | |
| 234 | // Apply RoPE to normalized elements |
| 235 | float elements2[numElemsPerThread]; // Additional buffer required for RoPE. |
| 236 | |
| 237 | int64_t pos_id = position_ids[tokenIdx]; |
| 238 | |
| 239 | // Calculate cache pointer for this position - similar to |
| 240 | // pos_encoding_kernels.cu |
| 241 | T_cache const* cache_ptr = cos_sin_cache + pos_id * rotary_dim; |
| 242 | int const embed_dim = rotary_dim / 2; |
| 243 | T_cache const* cos_ptr = cache_ptr; |
| 244 | T_cache const* sin_ptr = cache_ptr + embed_dim; |
| 245 | int const rotary_lanes = rotary_dim / numElemsPerThread; // rotary range |
| 246 | if (laneId < rotary_lanes) { |
| 247 | if constexpr (interleave) { |
| 248 | // Perform interleaving. Use pre-computed cos/sin values. |
| 249 | #pragma unroll |
| 250 | for (int i = 0; i < numElemsPerThread / 2; ++i) { |
| 251 | int const idx0 = 2 * i; |
| 252 | int const idx1 = 2 * i + 1; |
| 253 | // Global dimension index in the head |
| 254 | int const dim_idx = laneId * numElemsPerThread + idx0; |
| 255 | |
| 256 | float const val0 = elements[idx0]; |
| 257 | float const val1 = elements[idx1]; |
| 258 | |
| 259 | int const half_dim = dim_idx / 2; |
| 260 | float const cos_val = |
| 261 | CacheConverter::convert(VLLM_LDG(cos_ptr + half_dim)); |
| 262 | float const sin_val = |
| 263 | CacheConverter::convert(VLLM_LDG(sin_ptr + half_dim)); |
| 264 | |
| 265 | elements[idx0] = val0 * cos_val - val1 * sin_val; |
| 266 | elements[idx1] = val0 * sin_val + val1 * cos_val; |
| 267 | } |
| 268 | } else { |
| 269 | // Before data exchange with in warp, we need to sync. |
| 270 | __syncwarp(); |
| 271 | int pairOffset = (rotary_dim / 2) / numElemsPerThread; |
| 272 | // Get the data from the other half of the warp. Use pre-computed |
| 273 | // cos/sin values. |
| 274 | #pragma unroll |
| 275 | for (int i = 0; i < numElemsPerThread; i++) { |
| 276 | elements2[i] = __shfl_xor_sync(FINAL_MASK, elements[i], pairOffset); |
| 277 | |
| 278 | if (laneId < pairOffset) { |
| 279 | elements2[i] = -elements2[i]; |
| 280 | } |
| 281 | int dim_idx = laneId * numElemsPerThread + i; |
| 282 | |
| 283 | dim_idx = (dim_idx * 2) % rotary_dim; |
| 284 | int half_dim = dim_idx / 2; |
| 285 | float cos_val = CacheConverter::convert(VLLM_LDG(cos_ptr + half_dim)); |
| 286 | float sin_val = CacheConverter::convert(VLLM_LDG(sin_ptr + half_dim)); |
| 287 | |
| 288 | elements[i] = elements[i] * cos_val + elements2[i] * sin_val; |
| 289 | } |
| 290 | // __shfl_xor_sync does not provide memfence. Need to sync again. |
| 291 | __syncwarp(); |
| 292 | } |
| 293 | } |
| 294 | // Store. |
| 295 | { |
| 296 | vec_T vec; |
| 297 | constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in); |
| 298 | #pragma unroll |
| 299 | for (int i = 0; i < num_packed_elems; i++) { |
| 300 | // Convert from float2 back to the specific packed type |
| 301 | T2_in packed_val = Converter::convert( |
| 302 | make_float2(elements[2 * i], elements[2 * i + 1])); |
| 303 | // Place it into the generic vector |
| 304 | *(reinterpret_cast<T2_in*>(&vec) + i) = packed_val; |
| 305 | } |
| 306 | *reinterpret_cast<vec_T*>(&qkv[offsetThread]) = vec; |
| 307 | } |
| 308 | |
| 309 | // PDL: signal completion so a dependent successor may launch early. |
| 310 | #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) |
| 311 | cudaTriggerProgrammaticLaunchCompletion(); |
| 312 | #endif |
| 313 | |
| 314 | #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) |
| 315 | } |
| 316 | #endif |
| 317 | } |
| 318 | |
| 319 | // Multi-token-head kernel: one warp processes HEADS_PER_WARP token-heads for |
| 320 | // the same token, sharing cos/sin from shared memory via cp.async. |
| 321 | // When HEADS_PER_WARP > 1 the warp reuses the loaded cos/sin across all heads, |
| 322 | // hiding global-memory latency and improving occupancy for large batches. |
| 323 | template <typename scalar_t_in, typename scalar_t_cache, int head_dim, |
| 324 | bool interleave, int HEADS_PER_WARP> |
| 325 | __global__ void fusedQKNormRopeKernelNTokenHeads( |
| 326 | void* qkv_void, int const num_heads_q, int const num_heads_k, |
| 327 | int const num_heads_v, float const eps, void const* q_weight_void, |
| 328 | void const* k_weight_void, void const* cos_sin_cache_void, |
| 329 | int64_t const* position_ids, int const num_tokens, int const rotary_dim) { |
| 330 | #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) |
| 331 | if constexpr ((std::is_same_v<scalar_t_in, c10::BFloat16>) || |
| 332 | std::is_same_v<scalar_t_cache, c10::BFloat16>) { |
| 333 | return; |
| 334 | } else { |
| 335 | #endif |
| 336 | |
| 337 | using Converter = vllm::_typeConvert<scalar_t_in>; |
| 338 | static_assert(Converter::exists, |
| 339 | "Input QKV data type is not supported for this CUDA " |
| 340 | "architecture or toolkit version."); |
| 341 | using T_in = typename Converter::hip_type; |
| 342 | using T2_in = typename Converter::packed_hip_type; |
| 343 | |
| 344 | using CacheConverter = vllm::_typeConvert<scalar_t_cache>; |
| 345 | static_assert(CacheConverter::exists, |
| 346 | "Cache data type is not supported for this CUDA architecture " |
| 347 | "or toolkit version."); |
| 348 | using T_cache = typename CacheConverter::hip_type; |
| 349 | |
| 350 | extern __shared__ char smem_storage[]; |
| 351 | // Shared memory layout: |
| 352 | // [0, cos_sin_bytes) : cos/sin for each warp (warpsPerBlock * |
| 353 | // rotary_dim * sizeof(T_cache)) |
| 354 | // [cos_sin_bytes, ...) : QKV tiles |
| 355 | // per warp (warpsPerBlock * HEADS_PER_WARP * 32 * elemSizeBytes) |
| 356 | T_cache* const smem = reinterpret_cast<T_cache*>(smem_storage); |
| 357 | |
| 358 | T_in* qkv = reinterpret_cast<T_in*>(qkv_void); |
| 359 | T_in const* q_weight = reinterpret_cast<T_in const*>(q_weight_void); |
| 360 | T_in const* k_weight = reinterpret_cast<T_in const*>(k_weight_void); |
| 361 | T_cache const* cos_sin_cache = |
| 362 | reinterpret_cast<T_cache const*>(cos_sin_cache_void); |
| 363 | |
| 364 | int const warpsPerBlock = blockDim.x / 32; |
| 365 | int const warpId = threadIdx.x / 32; |
| 366 | int const laneId = threadIdx.x % 32; |
| 367 | |
| 368 | int const total_qk_heads = num_heads_q + num_heads_k; |
| 369 | int const num_heads = num_heads_q + num_heads_k + num_heads_v; |
| 370 | int const head_chunks_per_token = |
| 371 | (total_qk_heads + HEADS_PER_WARP - 1) / HEADS_PER_WARP; |
| 372 | |
| 373 | int const warp_global = blockIdx.x * warpsPerBlock + warpId; |
| 374 | int const tokenIdx = warp_global / head_chunks_per_token; |
| 375 | int const headChunk = warp_global % head_chunks_per_token; |
| 376 | int const first_head = headChunk * HEADS_PER_WARP; |
| 377 | int const num_heads_this_warp = |
| 378 | (first_head + HEADS_PER_WARP <= total_qk_heads) |
| 379 | ? HEADS_PER_WARP |
| 380 | : (total_qk_heads - first_head); |
| 381 | |
| 382 | if (tokenIdx >= num_tokens) return; |
| 383 | |
| 384 | static_assert(head_dim % (32 * 2) == 0, "head_dim must be divisible by 64"); |
| 385 | constexpr int numElemsPerThread = head_dim / 32; |
| 386 | constexpr int elemSizeBytes = numElemsPerThread * sizeof(__nv_bfloat16); |
| 387 | static_assert(elemSizeBytes % 4 == 0, |
| 388 | "elemSizeBytes must be a multiple of 4"); |
| 389 | constexpr int vecSize = elemSizeBytes / 4; |
| 390 | using vec_T = typename tensorrt_llm::common::packed_as<uint, vecSize>::type; |
| 391 | |
| 392 | int const cos_sin_bytes = |
| 393 | warpsPerBlock * rotary_dim * static_cast<int>(sizeof(T_cache)); |
| 394 | int const qkv_tile_bytes = 32 * elemSizeBytes; |
| 395 | char* const this_warp_head_smem = |
| 396 | smem_storage + cos_sin_bytes + |
| 397 | warpId * (HEADS_PER_WARP * qkv_tile_bytes); |
| 398 | |
| 399 | // PDL: wait for the predecessor kernel (the QKV projection that produces |
| 400 | // `qkv`) to complete before touching any global memory -- cp.async reads |
| 401 | // global memory, so this must precede the first commit group. |
| 402 | #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) |
| 403 | cudaGridDependencySynchronize(); |
| 404 | #endif |
| 405 | |
| 406 | // === Group 0: async load all heads' QKV into smem (issued first). === |
| 407 | for (int k = 0; k < num_heads_this_warp; ++k) { |
| 408 | int const localHeadIdx = first_head + k; |
| 409 | bool const isQ = localHeadIdx < num_heads_q; |
| 410 | int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q; |
| 411 | int offWarp; |
| 412 | if (isQ) { |
| 413 | offWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim; |
| 414 | } else { |
| 415 | offWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim + |
| 416 | headIdx * head_dim; |
| 417 | } |
| 418 | int const offThread = offWarp + laneId * numElemsPerThread; |
| 419 | char* smem_dst = |
| 420 | this_warp_head_smem + k * qkv_tile_bytes + laneId * elemSizeBytes; |
| 421 | cp_async_shared_global_ca(smem_dst, |
| 422 | reinterpret_cast<const char*>(&qkv[offThread]), |
| 423 | elemSizeBytes); |
| 424 | } |
| 425 | cp_async_commit_group(); // commit group 0 (QKV) |
| 426 | |
| 427 | // === Group 1: async load cos/sin into smem (issued second). === |
| 428 | int64_t const pos_id = position_ids[tokenIdx]; |
| 429 | T_cache const* const cache_ptr = cos_sin_cache + pos_id * rotary_dim; |
| 430 | int const copy_bytes = rotary_dim * static_cast<int>(sizeof(T_cache)); |
| 431 | int const num_copies = (copy_bytes + 15) / 16; |
| 432 | for (int copyId = laneId; copyId < num_copies; copyId += 32) { |
| 433 | char* smem_ptr = |
| 434 | reinterpret_cast<char*>(&smem[warpId * rotary_dim]) + copyId * 16; |
| 435 | const char* glob_ptr = |
| 436 | reinterpret_cast<const char*>(cache_ptr) + copyId * 16; |
| 437 | cp_async_shared_global_16_cg(smem_ptr, glob_ptr); |
| 438 | } |
| 439 | cp_async_commit_group(); // commit group 1 (cos/sin) |
| 440 | |
| 441 | // wait<1>: allow at most 1 pending group (group 1) → group 0 (QKV) is done. |
| 442 | cp_async_wait_group<1>(); |
| 443 | |
| 444 | float elements[numElemsPerThread]; |
| 445 | float elements2[numElemsPerThread]; |
| 446 | int const rotary_lanes = rotary_dim / numElemsPerThread; |
| 447 | int const embed_dim = rotary_dim / 2; |
| 448 | T_cache const* const cos_smem = &smem[warpId * rotary_dim]; |
| 449 | T_cache const* const sin_smem = &smem[warpId * rotary_dim + embed_dim]; |
| 450 | |
| 451 | // Preload weights into registers once, reused across all heads. |
| 452 | float q_w[numElemsPerThread]; |
| 453 | float k_w[numElemsPerThread]; |
| 454 | #pragma unroll |
| 455 | for (int i = 0; i < numElemsPerThread; i++) { |
| 456 | int const dim = laneId * numElemsPerThread + i; |
| 457 | q_w[i] = Converter::convert(q_weight[dim]); |
| 458 | k_w[i] = Converter::convert(k_weight[dim]); |
| 459 | } |
| 460 | |
| 461 | for (int k = 0; k < num_heads_this_warp; ++k) { |
| 462 | int const localHeadIdx = first_head + k; |
| 463 | bool const isQ = localHeadIdx < num_heads_q; |
| 464 | int const headIdx = isQ ? localHeadIdx : localHeadIdx - num_heads_q; |
| 465 | |
| 466 | int offsetWarp; |
| 467 | if (isQ) { |
| 468 | offsetWarp = tokenIdx * num_heads * head_dim + headIdx * head_dim; |
| 469 | } else { |
| 470 | offsetWarp = tokenIdx * num_heads * head_dim + num_heads_q * head_dim + |
| 471 | headIdx * head_dim; |
| 472 | } |
| 473 | int const offsetThread = offsetWarp + laneId * numElemsPerThread; |
| 474 | |
| 475 | // === Part 1: QK Norm (read from smem; group 0 already done). === |
| 476 | float sumOfSquares = 0.0f; |
| 477 | { |
| 478 | char const* smem_src = |
| 479 | this_warp_head_smem + k * qkv_tile_bytes + laneId * elemSizeBytes; |
| 480 | vec_T vec = *reinterpret_cast<vec_T const*>(smem_src); |
| 481 | constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in); |
| 482 | #pragma unroll |
| 483 | for (int i = 0; i < num_packed_elems; i++) { |
| 484 | T2_in packed_val = *(reinterpret_cast<T2_in*>(&vec) + i); |
| 485 | float2 vals = Converter::convert(packed_val); |
| 486 | sumOfSquares += vals.x * vals.x; |
| 487 | sumOfSquares += vals.y * vals.y; |
| 488 | elements[2 * i] = vals.x; |
| 489 | elements[2 * i + 1] = vals.y; |
| 490 | } |
| 491 | } |
| 492 | |
| 493 | sumOfSquares = tensorrt_llm::common::warpReduceSum(sumOfSquares); |
| 494 | float rms_rcp = rsqrtf(sumOfSquares / static_cast<float>(head_dim) + eps); |
| 495 | |
| 496 | #pragma unroll |
| 497 | for (int i = 0; i < numElemsPerThread; i++) { |
| 498 | elements[i] *= rms_rcp * (isQ ? q_w[i] : k_w[i]); |
| 499 | } |
| 500 | |
| 501 | // On first head: wait for group 1 (cos/sin) before RoPE. |
| 502 | if (k == 0) cp_async_wait_group<0>(); |
| 503 | |
| 504 | // === Part 2: RoPE using cos/sin from shared memory. === |
| 505 | if (laneId < rotary_lanes) { |
| 506 | if constexpr (interleave) { |
| 507 | #pragma unroll |
| 508 | for (int i = 0; i < numElemsPerThread / 2; ++i) { |
| 509 | int const idx0 = 2 * i; |
| 510 | int const idx1 = 2 * i + 1; |
| 511 | int const dim_idx = laneId * numElemsPerThread + idx0; |
| 512 | float const val0 = elements[idx0]; |
| 513 | float const val1 = elements[idx1]; |
| 514 | int const half_dim = dim_idx / 2; |
| 515 | float const cos_val = CacheConverter::convert(cos_smem[half_dim]); |
| 516 | float const sin_val = CacheConverter::convert(sin_smem[half_dim]); |
| 517 | elements[idx0] = val0 * cos_val - val1 * sin_val; |
| 518 | elements[idx1] = val0 * sin_val + val1 * cos_val; |
| 519 | } |
| 520 | } else { |
| 521 | __syncwarp(); |
| 522 | int const pairOffset = (rotary_dim / 2) / numElemsPerThread; |
| 523 | #pragma unroll |
| 524 | for (int i = 0; i < numElemsPerThread; i++) { |
| 525 | elements2[i] = __shfl_xor_sync(FINAL_MASK, elements[i], pairOffset); |
| 526 | if (laneId < pairOffset) elements2[i] = -elements2[i]; |
| 527 | int dim_idx = laneId * numElemsPerThread + i; |
| 528 | dim_idx = (dim_idx * 2) % rotary_dim; |
| 529 | int const half_dim = dim_idx / 2; |
| 530 | float const cos_val = CacheConverter::convert(cos_smem[half_dim]); |
| 531 | float const sin_val = CacheConverter::convert(sin_smem[half_dim]); |
| 532 | elements[i] = elements[i] * cos_val + elements2[i] * sin_val; |
| 533 | } |
| 534 | __syncwarp(); |
| 535 | } |
| 536 | } |
| 537 | |
| 538 | // Store. |
| 539 | { |
| 540 | vec_T vec; |
| 541 | constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in); |
| 542 | #pragma unroll |
| 543 | for (int i = 0; i < num_packed_elems; i++) { |
| 544 | T2_in packed_val = Converter::convert( |
| 545 | make_float2(elements[2 * i], elements[2 * i + 1])); |
| 546 | *(reinterpret_cast<T2_in*>(&vec) + i) = packed_val; |
| 547 | } |
| 548 | *reinterpret_cast<vec_T*>(&qkv[offsetThread]) = vec; |
| 549 | } |
| 550 | } |
| 551 | |
| 552 | // PDL: signal completion so a dependent successor may launch early. |
| 553 | #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) |
| 554 | cudaTriggerProgrammaticLaunchCompletion(); |
| 555 | #endif |
| 556 | |
| 557 | #if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM) |
| 558 | } |
| 559 | #endif |
| 560 | } |
| 561 | |
| 562 | // Borrowed from |
| 563 | // https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568 |
| 564 | #define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \ |
| 565 | if (interleave) { \ |
| 566 | const bool INTERLEAVE = true; \ |
| 567 | __VA_ARGS__ \ |
| 568 | } else { \ |
| 569 | const bool INTERLEAVE = false; \ |
| 570 | __VA_ARGS__ \ |
| 571 | } |
| 572 | |
| 573 | namespace { |
| 574 | // The empirical value for small batch: overlapping with the predecessor QKV |
| 575 | // projection only pays off while this kernel is launch-bound. |
| 576 | constexpr int kPDLEnableTokens = 16; |
| 577 | |
| 578 | // Launch helper for PDL (programmatic dependent launch), which lets this |
| 579 | // kernel start while the predecessor QKV projection is still draining. The |
| 580 | // launch attribute is the only switch: with it off, the kernels' |
| 581 | // cudaGridDependencySynchronize() / cudaTriggerProgrammaticLaunchCompletion() |
| 582 | // are no-ops. |
| 583 | template <typename KernelFn, typename... Args> |
| 584 | void launchWithPDL(KernelFn kernel, dim3 gridDim, dim3 blockDim, |
| 585 | size_t smem_bytes, cudaStream_t stream, bool enable_pdl, |
| 586 | Args... args) { |
| 587 | #ifndef USE_ROCM |
| 588 | cudaLaunchConfig_t config; |
| 589 | config.gridDim = gridDim; |
| 590 | config.blockDim = blockDim; |
| 591 | config.dynamicSmemBytes = smem_bytes; |
| 592 | config.stream = stream; |
| 593 | cudaLaunchAttribute attrs[1]; |
| 594 | attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; |
| 595 | attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; |
| 596 | config.numAttrs = 1; |
| 597 | config.attrs = attrs; |
| 598 | cudaLaunchKernelEx(&config, kernel, args...); |
| 599 | #else |
| 600 | // ROCm: standard kernel launch syntax (no PDL/stream serialization). |
| 601 | (void)enable_pdl; |
| 602 | // clang-format off |
| 603 | kernel<<<gridDim, blockDim, smem_bytes, stream>>>(args...); |
| 604 | // clang-format on |
| 605 | #endif |
| 606 | } |
| 607 | } // namespace |
| 608 | |
| 609 | template <typename scalar_t_in, typename scalar_t_cache> |
| 610 | void launchFusedQKNormRope(void* qkv, int const num_tokens, |
| 611 | int const num_heads_q, int const num_heads_k, |
| 612 | int const num_heads_v, int const head_dim, |
| 613 | int const rotary_dim, float const eps, |
| 614 | void const* q_weight, void const* k_weight, |
| 615 | void const* cos_sin_cache, bool const interleave, |
| 616 | int64_t const* position_ids, cudaStream_t stream) { |
| 617 | constexpr int blockSize = 256; |
| 618 | int const warpsPerBlock = blockSize / 32; |
| 619 | int const totalQKHeads = num_heads_q + num_heads_k; |
| 620 | int const totalWarps = num_tokens * totalQKHeads; |
| 621 | int const gridSize = common::divUp(totalWarps, warpsPerBlock); |
| 622 | dim3 gridDim(gridSize); |
| 623 | dim3 blockDim(blockSize); |
| 624 | bool const enable_pdl = num_tokens <= kPDLEnableTokens; |
| 625 | switch (head_dim) { |
| 626 | case 64: |
| 627 | DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { |
| 628 | launchWithPDL( |
| 629 | fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 64, INTERLEAVE>, |
| 630 | gridDim, blockDim, 0, stream, enable_pdl, qkv, num_heads_q, |
| 631 | num_heads_k, num_heads_v, eps, q_weight, k_weight, cos_sin_cache, |
| 632 | position_ids, num_tokens, rotary_dim); |
| 633 | }); |
| 634 | break; |
| 635 | case 128: |
| 636 | DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { |
| 637 | launchWithPDL( |
| 638 | fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 128, INTERLEAVE>, |
| 639 | gridDim, blockDim, 0, stream, enable_pdl, qkv, num_heads_q, |
| 640 | num_heads_k, num_heads_v, eps, q_weight, k_weight, cos_sin_cache, |
| 641 | position_ids, num_tokens, rotary_dim); |
| 642 | }); |
| 643 | break; |
| 644 | case 256: |
| 645 | DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { |
| 646 | launchWithPDL( |
| 647 | fusedQKNormRopeKernel<scalar_t_in, scalar_t_cache, 256, INTERLEAVE>, |
| 648 | gridDim, blockDim, 0, stream, enable_pdl, qkv, num_heads_q, |
| 649 | num_heads_k, num_heads_v, eps, q_weight, k_weight, cos_sin_cache, |
| 650 | position_ids, num_tokens, rotary_dim); |
| 651 | }); |
| 652 | break; |
| 653 | default: |
| 654 | STD_TORCH_CHECK( |
| 655 | false, "Unsupported head dimension for fusedQKNormRope: ", head_dim); |
| 656 | } |
| 657 | } |
| 658 | |
| 659 | // Launch: one warp processes token_heads_per_warp token-heads (1, 2, 4, or 8). |
| 660 | // When token_heads_per_warp == 1, delegates to the 1-head baseline above. |
| 661 | template <typename scalar_t_in, typename scalar_t_cache> |
| 662 | void launchFusedQKNormRopeNTokenHeads( |
| 663 | void* qkv, int const num_tokens, int const num_heads_q, |
| 664 | int const num_heads_k, int const num_heads_v, int const head_dim, |
| 665 | int const rotary_dim, float const eps, void const* q_weight, |
| 666 | void const* k_weight, void const* cos_sin_cache, bool const interleave, |
| 667 | int64_t const* position_ids, int const token_heads_per_warp, |
| 668 | cudaStream_t stream) { |
| 669 | STD_TORCH_CHECK(token_heads_per_warp == 1 || token_heads_per_warp == 2 || |
| 670 | token_heads_per_warp == 4 || token_heads_per_warp == 8, |
| 671 | "token_heads_per_warp must be 1, 2, 4, or 8, got ", |
| 672 | token_heads_per_warp); |
| 673 | |
| 674 | // token_heads_per_warp == 1: delegate to the 1-head baseline kernel. |
| 675 | if (token_heads_per_warp == 1) { |
| 676 | launchFusedQKNormRope<scalar_t_in, scalar_t_cache>( |
| 677 | qkv, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim, |
| 678 | rotary_dim, eps, q_weight, k_weight, cos_sin_cache, interleave, |
| 679 | position_ids, stream); |
| 680 | return; |
| 681 | } |
| 682 | |
| 683 | // NTokenHeads kernel uses cp.async to load cos/sin in 16-byte chunks. |
| 684 | // If rotary_dim * sizeof(cache_dtype) is not a multiple of 16, the last |
| 685 | // cp.async would write past the shared memory allocation. |
| 686 | // Fall back to the base kernel instead of failing. |
| 687 | { |
| 688 | size_t const rotary_bytes = |
| 689 | static_cast<size_t>(rotary_dim) * |
| 690 | (std::is_same_v<scalar_t_cache, float> ? sizeof(float) : 2u); |
| 691 | if (rotary_bytes % 16 != 0) { |
| 692 | launchFusedQKNormRope<scalar_t_in, scalar_t_cache>( |
| 693 | qkv, num_tokens, num_heads_q, num_heads_k, num_heads_v, head_dim, |
| 694 | rotary_dim, eps, q_weight, k_weight, cos_sin_cache, interleave, |
| 695 | position_ids, stream); |
| 696 | return; |
| 697 | } |
| 698 | } |
| 699 | |
| 700 | constexpr int blockSize = 256; |
| 701 | int const warpsPerBlock = blockSize / 32; |
| 702 | int const totalQKHeads = num_heads_q + num_heads_k; |
| 703 | // Grid: one warp per (token, head_chunk); same token → reuse cos/sin in smem. |
| 704 | int const head_chunks_per_token = |
| 705 | (totalQKHeads + token_heads_per_warp - 1) / token_heads_per_warp; |
| 706 | int const total_warps = num_tokens * head_chunks_per_token; |
| 707 | int const gridSize = common::divUp(total_warps, warpsPerBlock); |
| 708 | dim3 gridDim(gridSize); |
| 709 | dim3 blockDim(blockSize); |
| 710 | bool const enable_pdl = num_tokens <= kPDLEnableTokens; |
| 711 | // Cache element size: float=4, bfloat16=2 (host-safe; kernel uses same |
| 712 | // layout). |
| 713 | size_t const cache_elem_size = |
| 714 | std::is_same_v<scalar_t_cache, float> ? sizeof(float) : 2u; |
| 715 | // QKV smem: token_heads_per_warp tiles per warp, each tile 32*(head_dim/32*2) |
| 716 | // = 2*head_dim bytes. |
| 717 | size_t const qkv_smem_per_warp = static_cast<size_t>(token_heads_per_warp) * |
| 718 | 2u * static_cast<size_t>(head_dim); |
| 719 | size_t const smem_bytes = |
| 720 | warpsPerBlock * static_cast<size_t>(rotary_dim) * cache_elem_size + |
| 721 | warpsPerBlock * qkv_smem_per_warp; |
| 722 | |
| 723 | #define LAUNCH_N_TOKEN_HEADS(N) \ |
| 724 | do { \ |
| 725 | switch (head_dim) { \ |
| 726 | case 64: \ |
| 727 | DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \ |
| 728 | launchWithPDL( \ |
| 729 | fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, \ |
| 730 | 64, INTERLEAVE, (N)>, \ |
| 731 | gridDim, blockDim, smem_bytes, stream, enable_pdl, qkv, \ |
| 732 | num_heads_q, num_heads_k, num_heads_v, eps, q_weight, k_weight, \ |
| 733 | cos_sin_cache, position_ids, num_tokens, rotary_dim); \ |
| 734 | }); \ |
| 735 | break; \ |
| 736 | case 128: \ |
| 737 | DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \ |
| 738 | launchWithPDL( \ |
| 739 | fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, \ |
| 740 | 128, INTERLEAVE, (N)>, \ |
| 741 | gridDim, blockDim, smem_bytes, stream, enable_pdl, qkv, \ |
| 742 | num_heads_q, num_heads_k, num_heads_v, eps, q_weight, k_weight, \ |
| 743 | cos_sin_cache, position_ids, num_tokens, rotary_dim); \ |
| 744 | }); \ |
| 745 | break; \ |
| 746 | case 256: \ |
| 747 | DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { \ |
| 748 | launchWithPDL( \ |
| 749 | fusedQKNormRopeKernelNTokenHeads<scalar_t_in, scalar_t_cache, \ |
| 750 | 256, INTERLEAVE, (N)>, \ |
| 751 | gridDim, blockDim, smem_bytes, stream, enable_pdl, qkv, \ |
| 752 | num_heads_q, num_heads_k, num_heads_v, eps, q_weight, k_weight, \ |
| 753 | cos_sin_cache, position_ids, num_tokens, rotary_dim); \ |
| 754 | }); \ |
| 755 | break; \ |
| 756 | default: \ |
| 757 | STD_TORCH_CHECK(false, "Unsupported head dimension: ", head_dim); \ |
| 758 | } \ |
| 759 | } while (0) |
| 760 | |
| 761 | if (token_heads_per_warp == 2) { |
| 762 | LAUNCH_N_TOKEN_HEADS(2); |
| 763 | } else if (token_heads_per_warp == 4) { |
| 764 | LAUNCH_N_TOKEN_HEADS(4); |
| 765 | } else if (token_heads_per_warp == 8) { |
| 766 | LAUNCH_N_TOKEN_HEADS(8); |
| 767 | } |
| 768 | #undef LAUNCH_N_TOKEN_HEADS |
| 769 | } |
| 770 | |
| 771 | } // namespace tensorrt_llm::kernels |
| 772 | |
| 773 | void fused_qk_norm_rope( |
| 774 | torch::stable::Tensor& |
| 775 | qkv, // Combined QKV tensor [num_tokens, |
| 776 | // (num_heads_q+num_heads_k+num_heads_v)*head_dim] |
| 777 | int64_t num_heads_q, // Number of query heads |
| 778 | int64_t num_heads_k, // Number of key heads |
| 779 | int64_t num_heads_v, // Number of value heads |
| 780 | int64_t head_dim, // Dimension per head |
| 781 | double eps, // Epsilon for RMS normalization |
| 782 | torch::stable::Tensor& q_weight, // RMSNorm weights for query [head_dim] |
| 783 | torch::stable::Tensor& k_weight, // RMSNorm weights for key [head_dim] |
| 784 | torch::stable::Tensor& cos_sin_cache, // Cos/sin cache [max_position, |
| 785 | // head_dim] |
| 786 | bool is_neox, // Whether RoPE is applied in Neox style |
| 787 | torch::stable::Tensor& position_ids, // Position IDs for RoPE [num_tokens] |
| 788 | int64_t forced_token_heads_per_warp // -1 = auto-select, >0 = forced value |
| 789 | ) { |
| 790 | // Input validation |
| 791 | CHECK_INPUT(qkv); |
| 792 | CHECK_INPUT(position_ids); |
| 793 | CHECK_INPUT(q_weight); |
| 794 | CHECK_INPUT(k_weight); |
| 795 | CHECK_INPUT(cos_sin_cache); |
| 796 | CHECK_TYPE(position_ids, torch::headeronly::ScalarType::Long); |
| 797 | |
| 798 | STD_TORCH_CHECK(qkv.dim() == 2, |
| 799 | "QKV tensor must be 2D: [num_tokens, " |
| 800 | "(num_heads_q+num_heads_k+num_heads_v)*head_dim]"); |
| 801 | STD_TORCH_CHECK(position_ids.dim() == 1, |
| 802 | "Position IDs must be 1D: [num_tokens]"); |
| 803 | STD_TORCH_CHECK(q_weight.dim() == 1, "Query weights must be 1D: [head_dim]"); |
| 804 | STD_TORCH_CHECK(k_weight.dim() == 1, "Key weights must be 1D: [head_dim]"); |
| 805 | STD_TORCH_CHECK(cos_sin_cache.dim() == 2, |
| 806 | "Cos/sin cache must be 2D: [max_position, head_dim]"); |
| 807 | STD_TORCH_CHECK(q_weight.size(0) == head_dim, |
| 808 | "Query weights size must match head dimension"); |
| 809 | STD_TORCH_CHECK(k_weight.size(0) == head_dim, |
| 810 | "Key weights size must match head dimension"); |
| 811 | |
| 812 | STD_TORCH_CHECK(cos_sin_cache.size(1) % 2 == 0, "rotary_dim must be even"); |
| 813 | STD_TORCH_CHECK(cos_sin_cache.size(1) <= head_dim, |
| 814 | "rotary_dim must be less than or equal to head_dim"); |
| 815 | |
| 816 | STD_TORCH_CHECK(qkv.scalar_type() == q_weight.scalar_type() && |
| 817 | qkv.scalar_type() == k_weight.scalar_type(), |
| 818 | "qkv, q_weight and k_weight must have the same dtype"); |
| 819 | |
| 820 | int64_t num_tokens = qkv.size(0); |
| 821 | STD_TORCH_CHECK(position_ids.size(0) == num_tokens, |
| 822 | "Number of tokens in position_ids must match QKV"); |
| 823 | |
| 824 | int64_t total_heads = num_heads_q + num_heads_k + num_heads_v; |
| 825 | STD_TORCH_CHECK( |
| 826 | qkv.size(1) == total_heads * head_dim, |
| 827 | "QKV tensor size must match total number of heads and head dimension"); |
| 828 | |
| 829 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 830 | qkv.get_device_index()); |
| 831 | auto stream = get_current_cuda_stream(qkv.get_device_index()); |
| 832 | |
| 833 | // Select token_heads_per_warp: forced value if >0, else auto-select. |
| 834 | // Auto thresholds are calibrated on SM 9.0 (H100). On other architectures, |
| 835 | // fall back to token_heads_per_warp=1 (base kernel) until profiled. |
| 836 | int token_heads_per_warp; |
| 837 | if (forced_token_heads_per_warp > 0) { // only support SM80+ |
| 838 | token_heads_per_warp = static_cast<int>(forced_token_heads_per_warp); |
| 839 | } else { |
| 840 | token_heads_per_warp = 1; |
| 841 | int sm_version = get_device_prop()->major * 10 + get_device_prop()->minor; |
| 842 | int64_t total_qk_units = num_tokens * (num_heads_q + num_heads_k); |
| 843 | if (sm_version == 90) { |
| 844 | if (head_dim >= 256) { |
| 845 | if (total_qk_units < 4096LL) { |
| 846 | token_heads_per_warp = 1; |
| 847 | } else if (total_qk_units < 8192LL) { |
| 848 | token_heads_per_warp = 2; |
| 849 | } else { |
| 850 | token_heads_per_warp = 4; |
| 851 | } |
| 852 | } else { |
| 853 | if (total_qk_units < 10240LL) { |
| 854 | token_heads_per_warp = 1; |
| 855 | } else if (total_qk_units < 40960LL) { |
| 856 | token_heads_per_warp = 4; |
| 857 | } else { |
| 858 | token_heads_per_warp = 8; |
| 859 | } |
| 860 | } |
| 861 | } |
| 862 | } |
| 863 | |
| 864 | VLLM_STABLE_DISPATCH_HALF_TYPES( |
| 865 | qkv.scalar_type(), "fused_qk_norm_rope_kernel", [&] { |
| 866 | using qkv_scalar_t = scalar_t; |
| 867 | VLLM_STABLE_DISPATCH_FLOATING_TYPES( |
| 868 | cos_sin_cache.scalar_type(), "fused_qk_norm_rope_kernel", [&] { |
| 869 | using cache_scalar_t = scalar_t; |
| 870 | tensorrt_llm::kernels::launchFusedQKNormRopeNTokenHeads< |
| 871 | qkv_scalar_t, cache_scalar_t>( |
| 872 | qkv.data_ptr(), static_cast<int>(num_tokens), |
| 873 | static_cast<int>(num_heads_q), static_cast<int>(num_heads_k), |
| 874 | static_cast<int>(num_heads_v), static_cast<int>(head_dim), |
| 875 | static_cast<int>(cos_sin_cache.size(1)), |
| 876 | static_cast<float>(eps), q_weight.data_ptr(), |
| 877 | k_weight.data_ptr(), cos_sin_cache.data_ptr(), !is_neox, |
| 878 | reinterpret_cast<int64_t const*>(position_ids.data_ptr()), |
| 879 | token_heads_per_warp, stream); |
| 880 | }); |
| 881 | }); |
| 882 | } |