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 14, 05:30 UTC
Apply fixes for CUDA 13 (#24599)
Signed-off-by: Aidyn-A <aidyn.b.aitzhan@gmail.com>
Aidyn-A committed
Sep 17, 2025 at 17:15 UTC
bfe93801614b73ee5b4ac8ff65f977686a674bf2
8 files changed
+47
-56
CMakeLists.txt
+10
@@ -175,6 +175,16 @@ if(NVCC_THREADS AND VLLM_GPU_LANG STREQUAL "CUDA")
175
list(APPEND VLLM_GPU_FLAGS "--threads=${NVCC_THREADS}")
176
endif()
177
178
+#
179
+# Set CUDA include flags for CXX compiler.
180
+#
181
+if(VLLM_GPU_LANG STREQUAL "CUDA")
182
+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -I${CUDA_TOOLKIT_ROOT_DIR}/include")
183
+ if(CUDA_VERSION VERSION_GREATER_EQUAL 13.0)
184
+ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -I${CUDA_TOOLKIT_ROOT_DIR}/include/cccl")
185
+ endif()
186
+endif()
187
+
188
#
189
# Use FetchContent for C++ dependencies that are compiled as part of vLLM's build process.
190
# setup.py will override FETCHCONTENT_BASE_DIR to play nicely with sccache.
csrc/cub_helpers.h
new
+17
@@ -0,0 +1,17 @@
1
+#pragma once
2
+
3
+#ifndef USE_ROCM
4
+ #include <cub/cub.cuh>
5
+ #if CUB_VERSION >= 200800
6
+ #include <cuda/std/functional>
7
+using CubAddOp = cuda::std::plus<>;
8
+using CubMaxOp = cuda::maximum<>;
9
+ #else // if CUB_VERSION < 200800
10
+using CubAddOp = cub::Sum;
11
+using CubMaxOp = cub::Max;
12
+ #endif // CUB_VERSION
13
+#else
14
+ #include <hipcub/hipcub.hpp>
15
+using CubAddOp = cub::Sum;
16
+using CubMaxOp = cub::Max;
17
+#endif // USE_ROCM
csrc/layernorm_kernels.cu
+4
-9
@@ -1,15 +1,10 @@
1
#include "type_convert.cuh"
2
#include "dispatch_utils.h"
3
+#include "cub_helpers.h"
4
5
#include <torch/cuda.h>
6
#include <c10/cuda/CUDAGuard.h>
7
7
-#ifndef USE_ROCM
8
- #include <cub/cub.cuh>
9
-#else
10
- #include <hipcub/hipcub.hpp>
11
-#endif
12
-
8
namespace vllm {
9
10
// TODO(woosuk): Further optimize this kernel.
@@ -30,7 +25,7 @@ __global__ void rms_norm_kernel(
25
26
using BlockReduce = cub::BlockReduce<float, 1024>;
27
__shared__ typename BlockReduce::TempStorage reduceStore;
33
- variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
28
+ variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
29
30
if (threadIdx.x == 0) {
31
s_variance = rsqrtf(variance / hidden_size + epsilon);
@@ -85,7 +80,7 @@ fused_add_rms_norm_kernel(
80
81
using BlockReduce = cub::BlockReduce<float, 1024>;
82
__shared__ typename BlockReduce::TempStorage reduceStore;
88
- variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
83
+ variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
84
85
if (threadIdx.x == 0) {
86
s_variance = rsqrtf(variance / hidden_size + epsilon);
@@ -126,7 +121,7 @@ fused_add_rms_norm_kernel(
121
122
using BlockReduce = cub::BlockReduce<float, 1024>;
123
__shared__ typename BlockReduce::TempStorage reduceStore;
129
- variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
124
+ variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
125
126
if (threadIdx.x == 0) {
127
s_variance = rsqrtf(variance / hidden_size + epsilon);
csrc/layernorm_quant_kernels.cu
+4
-9
@@ -8,16 +8,11 @@
8
#include "type_convert.cuh"
9
#include "quantization/fp8/common.cuh"
10
#include "dispatch_utils.h"
11
+#include "cub_helpers.h"
12
13
#include <torch/cuda.h>
14
#include <c10/cuda/CUDAGuard.h>
15
15
-#ifndef USE_ROCM
16
- #include <cub/cub.cuh>
17
-#else
18
- #include <hipcub/hipcub.hpp>
19
-#endif
20
-
16
namespace vllm {
17
18
// TODO(woosuk): Further optimize this kernel.
@@ -39,7 +34,7 @@ __global__ void rms_norm_static_fp8_quant_kernel(
34
35
using BlockReduce = cub::BlockReduce<float, 1024>;
36
__shared__ typename BlockReduce::TempStorage reduceStore;
42
- variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
37
+ variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
38
39
if (threadIdx.x == 0) {
40
s_variance = rsqrtf(variance / hidden_size + epsilon);
@@ -100,7 +95,7 @@ fused_add_rms_norm_static_fp8_quant_kernel(
95
96
using BlockReduce = cub::BlockReduce<float, 1024>;
97
__shared__ typename BlockReduce::TempStorage reduceStore;
103
- variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
98
+ variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
99
100
if (threadIdx.x == 0) {
101
s_variance = rsqrtf(variance / hidden_size + epsilon);
@@ -149,7 +144,7 @@ fused_add_rms_norm_static_fp8_quant_kernel(
144
145
using BlockReduce = cub::BlockReduce<float, 1024>;
146
__shared__ typename BlockReduce::TempStorage reduceStore;
152
- variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
147
+ variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
148
149
if (threadIdx.x == 0) {
150
s_variance = rsqrtf(variance / hidden_size + epsilon);
csrc/moe/topk_softmax_kernels.cu
+3
-13
@@ -20,17 +20,7 @@
20
#include <ATen/cuda/CUDAContext.h>
21
#include <c10/cuda/CUDAGuard.h>
22
#include "../cuda_compat.h"
23
-
24
-#ifndef USE_ROCM
25
- #include <cub/util_type.cuh>
26
- #include <cub/cub.cuh>
27
- #include <cuda/std/functional>
28
- using AddOp = cuda::std::plus<float>;
29
-#else
30
- #include <hipcub/util_type.hpp>
31
- #include <hipcub/hipcub.hpp>
32
- using AddOp = cub::Sum;
33
-#endif
23
+#include "../cub_helpers.h"
24
25
#define MAX(a, b) ((a) > (b) ? (a) : (b))
26
#define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -79,7 +69,7 @@ __launch_bounds__(TPB) __global__
69
threadData = max(static_cast<float>(input[idx]), threadData);
70
}
71
82
- const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, cub::Max());
72
+ const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, CubMaxOp());
73
if (threadIdx.x == 0)
74
{
75
float_max = maxElem;
@@ -94,7 +84,7 @@ __launch_bounds__(TPB) __global__
84
threadData += exp((static_cast<float>(input[idx]) - float_max));
85
}
86
97
- const auto Z = BlockReduce(tmpStorage).Reduce(threadData, AddOp());
87
+ const auto Z = BlockReduce(tmpStorage).Reduce(threadData, CubAddOp());
88
89
if (threadIdx.x == 0)
90
{
csrc/quantization/compressed_tensors/int8_quant_kernels.cu
+2
-9
@@ -7,17 +7,10 @@
7
8
#include <cmath>
9
10
+#include "../../cub_helpers.h"
11
#include "../../dispatch_utils.h"
12
#include "../vectorization_utils.cuh"
13
13
-#ifndef USE_ROCM
14
- #include <cub/cub.cuh>
15
- #include <cub/util_type.cuh>
16
-#else
17
- #include <hipcub/hipcub.hpp>
18
- #include <hipcub/util_type.hpp>
19
-#endif
20
-
14
static inline __device__ int8_t float_to_int8_rn(float x) {
15
#ifdef USE_ROCM
16
static constexpr auto i8_min =
@@ -173,7 +166,7 @@ __global__ void dynamic_scaled_int8_quant_kernel(
166
});
167
using BlockReduce = cub::BlockReduce<float, 256>;
168
__shared__ typename BlockReduce::TempStorage tmp;
176
- float block_max = BlockReduce(tmp).Reduce(thread_max, cub::Max{}, blockDim.x);
169
+ float block_max = BlockReduce(tmp).Reduce(thread_max, CubMaxOp{}, blockDim.x);
170
__shared__ float absmax;
171
if (tid == 0) {
172
absmax = block_max;
csrc/quantization/fp8/common.cu
+2
-7
@@ -1,15 +1,10 @@
1
#include "common.cuh"
2
#include "dispatch_utils.h"
3
+#include "../../cub_helpers.h"
4
#include "../vectorization_utils.cuh"
5
#include <c10/cuda/CUDAGuard.h>
6
#include <ATen/cuda/Exceptions.h>
7
7
-#ifndef USE_ROCM
8
- #include <cub/cub.cuh>
9
-#else
10
- #include <hipcub/hipcub.hpp>
11
-#endif
12
-
8
namespace vllm {
9
10
template <typename scalar_t, typename fp8_type>
@@ -116,7 +111,7 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel_strided(
111
using BlockReduce = cub::BlockReduce<float, 256>;
112
__shared__ typename BlockReduce::TempStorage tmp;
113
const float block_max =
119
- BlockReduce(tmp).Reduce(absmax_val, cub::Max{}, blockDim.x);
114
+ BlockReduce(tmp).Reduce(absmax_val, CubMaxOp{}, blockDim.x);
115
116
__shared__ float token_scale;
117
if (tid == 0) {
csrc/quantization/fused_kernels/layernorm_utils.cuh
+5
-9
@@ -8,11 +8,7 @@
8
#include "quantization/utils.cuh"
9
#include "quant_conversions.cuh"
10
11
-#ifndef USE_ROCM
12
- #include <cub/cub.cuh>
13
-#else
14
- #include <hipcub/hipcub.hpp>
15
-#endif
11
+#include "../../cub_helpers.h"
12
13
namespace vllm {
14
@@ -36,7 +32,7 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
32
33
using BlockReduce = cub::BlockReduce<float, 1024>;
34
__shared__ typename BlockReduce::TempStorage reduceStore;
39
- ss = BlockReduce(reduceStore).Reduce(ss, cub::Sum{}, blockDim.x);
35
+ ss = BlockReduce(reduceStore).Reduce(ss, CubAddOp{}, blockDim.x);
36
37
__shared__ float s_rms;
38
if (threadIdx.x == 0) {
@@ -73,7 +69,7 @@ __device__ void compute_dynamic_per_token_scales(
69
__shared__ typename BlockReduce::TempStorage reduceStore;
70
block_absmax_val_maybe =
71
BlockReduce(reduceStore)
76
- .Reduce(block_absmax_val_maybe, cub::Max{}, blockDim.x);
72
+ .Reduce(block_absmax_val_maybe, CubMaxOp{}, blockDim.x);
73
74
__shared__ float s_token_scale;
75
if (threadIdx.x == 0) {
@@ -169,7 +165,7 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
165
166
using BlockReduce = cub::BlockReduce<float, 1024>;
167
__shared__ typename BlockReduce::TempStorage reduceStore;
172
- ss = BlockReduce(reduceStore).Reduce(ss, cub::Sum{}, blockDim.x);
168
+ ss = BlockReduce(reduceStore).Reduce(ss, CubAddOp{}, blockDim.x);
169
170
__shared__ float s_rms;
171
if (threadIdx.x == 0) {
@@ -240,7 +236,7 @@ __device__ void compute_dynamic_per_token_scales(
236
__shared__ typename BlockReduce::TempStorage reduceStore;
237
block_absmax_val_maybe =
238
BlockReduce(reduceStore)
243
- .Reduce(block_absmax_val_maybe, cub::Max{}, blockDim.x);
239
+ .Reduce(block_absmax_val_maybe, CubMaxOp{}, blockDim.x);
240
241
__shared__ float s_token_scale;
242
if (threadIdx.x == 0) {