| 1 | #include <numeric> |
| 2 | |
| 3 | #include "torch_utils.h" |
| 4 | |
| 5 | #include "cub_helpers.h" |
| 6 | #include "../core/batch_invariant.hpp" |
| 7 | #include "type_convert.cuh" |
| 8 | #include "dispatch_utils.h" |
| 9 | #include "quantization/vectorization_utils.cuh" |
| 10 | |
| 11 | namespace vllm { |
| 12 | |
| 13 | // TODO(woosuk): Further optimize this kernel. |
| 14 | template <typename scalar_t, int VEC_SIZE, int NUM_DIMS, bool HasWeight> |
| 15 | __global__ void rms_norm_kernel( |
| 16 | scalar_t* __restrict__ out, // [..., hidden_size] |
| 17 | const scalar_t* __restrict__ input, // [..., hidden_size] |
| 18 | const int64_t input_stride_d2, // input.stride(-2) |
| 19 | const int64_t input_stride_d3, // input.stride(-3) |
| 20 | const int64_t input_stride_d4, // input.stride(-4) |
| 21 | const int64_t input_shape_d2, // input.size(-2) |
| 22 | const int64_t input_shape_d3, // input.size(-3) |
| 23 | const scalar_t* __restrict__ weight, // [hidden_size] or |
| 24 | // [num_groups, hidden_size]; |
| 25 | // null if !HasWeight |
| 26 | const int64_t weight_stride, // 0 or weight.stride(0) |
| 27 | const float epsilon, const int num_tokens, const int hidden_size) { |
| 28 | __shared__ float s_variance; |
| 29 | float variance = 0.0f; |
| 30 | const scalar_t* input_row; |
| 31 | const scalar_t* weight_row; |
| 32 | if constexpr (NUM_DIMS == 2) { |
| 33 | // 2D for layernorm normal case [batch_size, hidden] |
| 34 | input_row = input + blockIdx.x * input_stride_d2; |
| 35 | weight_row = weight + blockIdx.x * weight_stride; |
| 36 | } else if constexpr (NUM_DIMS == 3) { |
| 37 | // 3D for q/k norm [batch_size, num_heads, head_size] |
| 38 | int batch_idx = blockIdx.x / input_shape_d2; |
| 39 | int head_idx = blockIdx.x % input_shape_d2; |
| 40 | input_row = |
| 41 | input + batch_idx * input_stride_d3 + head_idx * input_stride_d2; |
| 42 | weight_row = weight + batch_idx * weight_stride; |
| 43 | } else if constexpr (NUM_DIMS == 4) { |
| 44 | // 4D for transformers model_impl qk norm [batch, seq, head, head_dim] |
| 45 | int batch_idx = blockIdx.x / (input_shape_d3 * input_shape_d2); |
| 46 | int remaining = blockIdx.x % (input_shape_d3 * input_shape_d2); |
| 47 | int seq_idx = remaining / input_shape_d2; |
| 48 | int head_idx = remaining % input_shape_d2; |
| 49 | input_row = input + batch_idx * input_stride_d4 + |
| 50 | seq_idx * input_stride_d3 + head_idx * input_stride_d2; |
| 51 | weight_row = weight + batch_idx * weight_stride; |
| 52 | } |
| 53 | |
| 54 | auto vec_op = [&variance](const vec_n_t<scalar_t, VEC_SIZE>& vec) { |
| 55 | #pragma unroll |
| 56 | for (int i = 0; i < VEC_SIZE; ++i) { |
| 57 | float x = static_cast<float>(vec.val[i]); |
| 58 | variance += x * x; |
| 59 | } |
| 60 | }; |
| 61 | auto scalar_op = [&variance](const scalar_t& val) { |
| 62 | float x = static_cast<float>(val); |
| 63 | variance += x * x; |
| 64 | }; |
| 65 | vllm::vectorize_read_with_alignment<VEC_SIZE>( |
| 66 | input_row, hidden_size, threadIdx.x, blockDim.x, vec_op, scalar_op); |
| 67 | |
| 68 | using BlockReduce = cub::BlockReduce<float, 1024>; |
| 69 | __shared__ typename BlockReduce::TempStorage reduceStore; |
| 70 | variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); |
| 71 | |
| 72 | if (threadIdx.x == 0) { |
| 73 | s_variance = rsqrtf(variance / hidden_size + epsilon); |
| 74 | } |
| 75 | __syncthreads(); |
| 76 | |
| 77 | scalar_t* out_row = out + blockIdx.x * hidden_size; |
| 78 | auto* v_in = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(input_row); |
| 79 | auto* v_w = reinterpret_cast<const vec_n_t<scalar_t, VEC_SIZE>*>(weight_row); |
| 80 | auto* v_out = reinterpret_cast<vec_n_t<scalar_t, VEC_SIZE>*>(out_row); |
| 81 | for (int i = threadIdx.x; i < hidden_size / VEC_SIZE; i += blockDim.x) { |
| 82 | vec_n_t<scalar_t, VEC_SIZE> dst; |
| 83 | vec_n_t<scalar_t, VEC_SIZE> src1 = v_in[i]; |
| 84 | vec_n_t<scalar_t, VEC_SIZE> src2; |
| 85 | if constexpr (HasWeight) { |
| 86 | src2 = v_w[i]; |
| 87 | } |
| 88 | #pragma unroll |
| 89 | for (int j = 0; j < VEC_SIZE; j++) { |
| 90 | float x = static_cast<float>(src1.val[j]); |
| 91 | if constexpr (HasWeight) { |
| 92 | float w = static_cast<float>(src2.val[j]); |
| 93 | dst.val[j] = static_cast<scalar_t>(x * s_variance * w); |
| 94 | } else { |
| 95 | dst.val[j] = static_cast<scalar_t>(x * s_variance); |
| 96 | } |
| 97 | } |
| 98 | v_out[i] = dst; |
| 99 | } |
| 100 | } |
| 101 | |
| 102 | /* Function specialization in the case of FP16/BF16 tensors. |
| 103 | Additional optimizations we can make in this case are |
| 104 | packed and vectorized operations, which help with the |
| 105 | memory latency bottleneck. */ |
| 106 | template <typename scalar_t, int width, bool HasWeight> |
| 107 | __global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::exists> |
| 108 | fused_add_rms_norm_kernel( |
| 109 | scalar_t* __restrict__ input, // [..., hidden_size] |
| 110 | const int64_t input_stride, |
| 111 | scalar_t* __restrict__ residual, // [..., hidden_size] |
| 112 | const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight |
| 113 | const float epsilon, const int num_tokens, const int hidden_size, |
| 114 | const int64_t residual_stride) { |
| 115 | // Sanity checks on our vector struct and type-punned pointer arithmetic |
| 116 | static_assert(std::is_pod_v<_f16Vec<scalar_t, width>>); |
| 117 | static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width); |
| 118 | |
| 119 | const int vec_hidden_size = hidden_size / width; |
| 120 | const int64_t vec_input_stride = input_stride / width; |
| 121 | __shared__ float s_variance; |
| 122 | float variance = 0.0f; |
| 123 | /* These and the argument pointers are all declared `restrict` as they are |
| 124 | not aliased in practice. Argument pointers should not be dereferenced |
| 125 | in this kernel as that would be undefined behavior */ |
| 126 | auto* __restrict__ input_v = |
| 127 | reinterpret_cast<_f16Vec<scalar_t, width>*>(input); |
| 128 | auto* __restrict__ residual_v = |
| 129 | reinterpret_cast<_f16Vec<scalar_t, width>*>(residual); |
| 130 | auto* __restrict__ weight_v = |
| 131 | reinterpret_cast<const _f16Vec<scalar_t, width>*>(weight); |
| 132 | |
| 133 | for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { |
| 134 | int64_t id = blockIdx.x * residual_stride / width + idx; |
| 135 | int64_t strided_id = blockIdx.x * vec_input_stride + idx; |
| 136 | _f16Vec<scalar_t, width> temp = input_v[strided_id]; |
| 137 | temp += residual_v[id]; |
| 138 | variance += temp.sum_squares(); |
| 139 | residual_v[id] = temp; |
| 140 | } |
| 141 | |
| 142 | using BlockReduce = cub::BlockReduce<float, 1024>; |
| 143 | __shared__ typename BlockReduce::TempStorage reduceStore; |
| 144 | variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); |
| 145 | |
| 146 | if (threadIdx.x == 0) { |
| 147 | s_variance = rsqrtf(variance / hidden_size + epsilon); |
| 148 | } |
| 149 | __syncthreads(); |
| 150 | |
| 151 | for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) { |
| 152 | int64_t id = blockIdx.x * residual_stride / width + idx; |
| 153 | int64_t strided_id = blockIdx.x * vec_input_stride + idx; |
| 154 | _f16Vec<scalar_t, width> res = residual_v[id]; |
| 155 | _f16Vec<scalar_t, width> out; |
| 156 | using Converter = _typeConvert<scalar_t>; |
| 157 | if constexpr (HasWeight) { |
| 158 | _f16Vec<scalar_t, width> w = weight_v[idx]; |
| 159 | #pragma unroll |
| 160 | for (int j = 0; j < width; ++j) { |
| 161 | float x = Converter::convert(res.data[j]); |
| 162 | float wf = Converter::convert(w.data[j]); |
| 163 | out.data[j] = Converter::convert(x * s_variance * wf); |
| 164 | } |
| 165 | } else { |
| 166 | #pragma unroll |
| 167 | for (int j = 0; j < width; ++j) { |
| 168 | float x = Converter::convert(res.data[j]); |
| 169 | out.data[j] = Converter::convert(x * s_variance); |
| 170 | } |
| 171 | } |
| 172 | input_v[strided_id] = out; |
| 173 | } |
| 174 | } |
| 175 | |
| 176 | /* Generic fused_add_rms_norm_kernel |
| 177 | The width field is not used here but necessary for other specializations. |
| 178 | */ |
| 179 | template <typename scalar_t, int width, bool HasWeight> |
| 180 | __global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::exists> |
| 181 | fused_add_rms_norm_kernel( |
| 182 | scalar_t* __restrict__ input, // [..., hidden_size] |
| 183 | const int64_t input_stride, |
| 184 | scalar_t* __restrict__ residual, // [..., hidden_size] |
| 185 | const scalar_t* __restrict__ weight, // [hidden_size], null if !HasWeight |
| 186 | const float epsilon, const int num_tokens, const int hidden_size, |
| 187 | const int64_t residual_stride) { |
| 188 | __shared__ float s_variance; |
| 189 | float variance = 0.0f; |
| 190 | |
| 191 | for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { |
| 192 | scalar_t z = input[blockIdx.x * input_stride + idx]; |
| 193 | z += residual[blockIdx.x * residual_stride + idx]; |
| 194 | float x = (float)z; |
| 195 | variance += x * x; |
| 196 | residual[blockIdx.x * residual_stride + idx] = z; |
| 197 | } |
| 198 | |
| 199 | using BlockReduce = cub::BlockReduce<float, 1024>; |
| 200 | __shared__ typename BlockReduce::TempStorage reduceStore; |
| 201 | variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x); |
| 202 | |
| 203 | if (threadIdx.x == 0) { |
| 204 | s_variance = rsqrtf(variance / hidden_size + epsilon); |
| 205 | } |
| 206 | __syncthreads(); |
| 207 | |
| 208 | for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { |
| 209 | float x = (float)residual[blockIdx.x * residual_stride + idx]; |
| 210 | if constexpr (HasWeight) { |
| 211 | float w = (float)weight[idx]; |
| 212 | input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w); |
| 213 | } else { |
| 214 | input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance); |
| 215 | } |
| 216 | } |
| 217 | } |
| 218 | |
| 219 | } // namespace vllm |
| 220 | |
| 221 | void rms_norm(torch::stable::Tensor& out, // [..., hidden_size] |
| 222 | torch::stable::Tensor& input, // [..., hidden_size] |
| 223 | std::optional<torch::stable::Tensor> weight, double epsilon) { |
| 224 | STD_TORCH_CHECK(out.is_contiguous()); |
| 225 | if (input.stride(-1) != 1) { |
| 226 | input = torch::stable::contiguous(input); |
| 227 | } |
| 228 | STD_TORCH_CHECK(input.stride(-1) == 1); |
| 229 | int64_t weight_stride = 0; |
| 230 | if (weight.has_value()) { |
| 231 | STD_TORCH_CHECK(weight->is_contiguous()); |
| 232 | if (weight->dim() == 1) { |
| 233 | STD_TORCH_CHECK(weight->size(0) == input.size(-1)); |
| 234 | } else if (weight->dim() == 2) { |
| 235 | STD_TORCH_CHECK(weight->size(0) == input.size(0)); |
| 236 | STD_TORCH_CHECK(weight->size(-1) == input.size(-1)); |
| 237 | weight_stride = weight->stride(0); |
| 238 | } else { |
| 239 | STD_TORCH_CHECK(false, "rms_norm weight must be 1D or 2D"); |
| 240 | } |
| 241 | } |
| 242 | |
| 243 | int hidden_size = input.size(-1); |
| 244 | |
| 245 | int num_tokens = input.numel() / hidden_size; |
| 246 | int num_dims = input.dim(); |
| 247 | int64_t input_stride_d2 = input.stride(-2); |
| 248 | int64_t input_stride_d3 = (num_dims >= 3) ? input.stride(-3) : 0; |
| 249 | int64_t input_stride_d4 = (num_dims >= 4) ? input.stride(-4) : 0; |
| 250 | int64_t input_shape_d2 = (num_dims >= 3) ? input.size(-2) : 0; |
| 251 | int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0; |
| 252 | |
| 253 | // For large num_tokens, use smaller blocks to increase SM concurrency. |
| 254 | const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); |
| 255 | const int max_block_size = |
| 256 | batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256); |
| 257 | dim3 grid(num_tokens); |
| 258 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 259 | input.get_device_index()); |
| 260 | const cudaStream_t stream = get_current_cuda_stream(); |
| 261 | const bool has_weight = weight.has_value(); |
| 262 | VLLM_STABLE_DISPATCH_RANK234(num_dims, [&] { |
| 263 | VLLM_STABLE_DISPATCH_FLOATING_TYPES( |
| 264 | input.scalar_type(), "rms_norm_kernel", [&] { |
| 265 | const scalar_t* weight_ptr = |
| 266 | has_weight ? weight->const_data_ptr<scalar_t>() : nullptr; |
| 267 | const int calculated_vec_size = |
| 268 | std::gcd(16 / sizeof(scalar_t), hidden_size); |
| 269 | const int block_size = |
| 270 | std::min(hidden_size / calculated_vec_size, max_block_size); |
| 271 | dim3 block(block_size); |
| 272 | VLLM_STABLE_DISPATCH_VEC_SIZE(calculated_vec_size, [&] { |
| 273 | if (has_weight) { |
| 274 | vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank, true> |
| 275 | <<<grid, block, 0, stream>>>( |
| 276 | out.mutable_data_ptr<scalar_t>(), |
| 277 | input.const_data_ptr<scalar_t>(), input_stride_d2, |
| 278 | input_stride_d3, input_stride_d4, input_shape_d2, |
| 279 | input_shape_d3, weight_ptr, weight_stride, epsilon, |
| 280 | num_tokens, hidden_size); |
| 281 | } else { |
| 282 | vllm::rms_norm_kernel<scalar_t, vec_size, tensor_rank, false> |
| 283 | <<<grid, block, 0, stream>>>( |
| 284 | out.mutable_data_ptr<scalar_t>(), |
| 285 | input.const_data_ptr<scalar_t>(), input_stride_d2, |
| 286 | input_stride_d3, input_stride_d4, input_shape_d2, |
| 287 | input_shape_d3, weight_ptr, /*weight_stride=*/0, epsilon, |
| 288 | num_tokens, hidden_size); |
| 289 | } |
| 290 | }); |
| 291 | }); |
| 292 | }); |
| 293 | } |
| 294 | |
| 295 | #define LAUNCH_FUSED_ADD_RMS_NORM(width, has_weight) \ |
| 296 | VLLM_STABLE_DISPATCH_FLOATING_TYPES( \ |
| 297 | input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \ |
| 298 | if (has_weight) { \ |
| 299 | vllm::fused_add_rms_norm_kernel<scalar_t, width, true> \ |
| 300 | <<<grid, block, 0, stream>>>( \ |
| 301 | input.mutable_data_ptr<scalar_t>(), input_stride, \ |
| 302 | residual.mutable_data_ptr<scalar_t>(), \ |
| 303 | weight->const_data_ptr<scalar_t>(), epsilon, num_tokens, \ |
| 304 | hidden_size, residual_stride); \ |
| 305 | } else { \ |
| 306 | vllm::fused_add_rms_norm_kernel<scalar_t, width, false> \ |
| 307 | <<<grid, block, 0, stream>>>( \ |
| 308 | input.mutable_data_ptr<scalar_t>(), input_stride, \ |
| 309 | residual.mutable_data_ptr<scalar_t>(), nullptr, epsilon, \ |
| 310 | num_tokens, hidden_size, residual_stride); \ |
| 311 | } \ |
| 312 | }); |
| 313 | |
| 314 | void fused_add_rms_norm(torch::stable::Tensor& input, // [..., hidden_size] |
| 315 | torch::stable::Tensor& residual, // [..., hidden_size] |
| 316 | std::optional<torch::stable::Tensor> weight, |
| 317 | double epsilon) { |
| 318 | STD_TORCH_CHECK(input.scalar_type() == residual.scalar_type()); |
| 319 | STD_TORCH_CHECK(residual.stride(-1) == 1); |
| 320 | if (weight.has_value()) { |
| 321 | STD_TORCH_CHECK(weight->scalar_type() == input.scalar_type()); |
| 322 | STD_TORCH_CHECK(weight->is_contiguous()); |
| 323 | } |
| 324 | int hidden_size = input.size(-1); |
| 325 | int64_t input_stride = input.stride(-2); |
| 326 | int64_t residual_stride = residual.stride(-2); |
| 327 | int num_tokens = input.numel() / hidden_size; |
| 328 | |
| 329 | dim3 grid(num_tokens); |
| 330 | /* This kernel is memory-latency bound in many scenarios. |
| 331 | When num_tokens is large, a smaller block size allows |
| 332 | for increased block occupancy on CUs and better latency |
| 333 | hiding on global mem ops. In batch-invariant mode the block size must |
| 334 | not depend on num_tokens, otherwise the same token would use a different |
| 335 | reduction width (and thus a different floating-point summation order) |
| 336 | across batches; lock it to 1024 to keep results bit-exact. */ |
| 337 | const bool batch_invariant_launch = vllm::vllm_is_batch_invariant(); |
| 338 | const int max_block_size = |
| 339 | batch_invariant_launch ? 1024 : ((num_tokens < 256) ? 1024 : 256); |
| 340 | dim3 block(std::min(hidden_size, max_block_size)); |
| 341 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 342 | input.get_device_index()); |
| 343 | const cudaStream_t stream = get_current_cuda_stream(); |
| 344 | constexpr int vector_width = 8; |
| 345 | constexpr int req_alignment_bytes = vector_width * 2; |
| 346 | auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr()); |
| 347 | auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr()); |
| 348 | bool offsets_are_multiple_of_vector_width = |
| 349 | hidden_size % vector_width == 0 && input_stride % vector_width == 0 && |
| 350 | residual_stride % vector_width == 0; |
| 351 | const bool has_weight = weight.has_value(); |
| 352 | if (has_weight) { |
| 353 | auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight->data_ptr()); |
| 354 | bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 && |
| 355 | res_ptr % req_alignment_bytes == 0 && |
| 356 | wt_ptr % req_alignment_bytes == 0; |
| 357 | if (ptrs_are_aligned && offsets_are_multiple_of_vector_width && |
| 358 | !batch_invariant_launch) { |
| 359 | LAUNCH_FUSED_ADD_RMS_NORM(8, true); |
| 360 | } else { |
| 361 | LAUNCH_FUSED_ADD_RMS_NORM(0, true); |
| 362 | } |
| 363 | } else { |
| 364 | bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 && |
| 365 | res_ptr % req_alignment_bytes == 0; |
| 366 | if (ptrs_are_aligned && offsets_are_multiple_of_vector_width && |
| 367 | !batch_invariant_launch) { |
| 368 | LAUNCH_FUSED_ADD_RMS_NORM(8, false); |
| 369 | } else { |
| 370 | LAUNCH_FUSED_ADD_RMS_NORM(0, false); |
| 371 | } |
| 372 | } |
| 373 | } |