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, 23:45 UTC
[Bugfix][ROCm] Fix for warp_size uses on host (#21205)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
Jul 24, 2025 at 03:37 UTC
90eeea8f8501f22151d08e57d6f18f639c9faf2f
9 files changed
+67
-31
csrc/attention/attention_kernels.cuh
+1
-1
@@ -24,7 +24,7 @@
24
25
#include "attention_dtypes.h"
26
#include "attention_utils.cuh"
27
-#include "cuda_compat.h"
27
+#include "../cuda_compat.h"
28
29
#ifdef USE_ROCM
30
#include <hip/hip_bf16.h>
csrc/attention/paged_attention_v1.cu
+2
-3
@@ -16,9 +16,8 @@
16
* See the License for the specific language governing permissions and
17
* limitations under the License.
18
*/
19
-
19
#include "attention_kernels.cuh"
21
-#include "cuda_compat.h"
20
+#include "../cuda_compat.h"
21
22
#define MAX(a, b) ((a) > (b) ? (a) : (b))
23
#define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -75,7 +74,7 @@ void paged_attention_v1_launcher(
74
const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
75
const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
76
78
- constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
77
+ const int NUM_WARPS = NUM_THREADS / WARP_SIZE;
78
int padded_max_seq_len =
79
DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE;
80
int logits_size = padded_max_seq_len * sizeof(float);
csrc/attention/paged_attention_v2.cu
+2
-3
@@ -16,9 +16,8 @@
16
* See the License for the specific language governing permissions and
17
* limitations under the License.
18
*/
19
-
19
#include "attention_kernels.cuh"
21
-#include "cuda_compat.h"
20
+#include "../cuda_compat.h"
21
22
#define MAX(a, b) ((a) > (b) ? (a) : (b))
23
#define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -79,7 +78,7 @@ void paged_attention_v2_launcher(
78
const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
79
const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
80
82
- constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
81
+ const int NUM_WARPS = NUM_THREADS / WARP_SIZE;
82
int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);
83
int logits_size = PARTITION_SIZE * sizeof(float);
84
int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
csrc/cuda_compat.h
+29
-2
@@ -4,8 +4,35 @@
4
#include <hip/hip_runtime.h>
5
#endif
6
7
-#if defined(USE_ROCM) && defined(__GFX9__)
8
- #define WARP_SIZE 64
7
+#ifdef USE_ROCM
8
+struct Utils {
9
+ static __host__ int get_warp_size() {
10
+ static bool is_cached = false;
11
+ static int result;
12
+
13
+ if (!is_cached) {
14
+ int device_id;
15
+ cudaDeviceProp deviceProp;
16
+ cudaGetDevice(&device_id);
17
+ cudaGetDeviceProperties(&deviceProp, device_id);
18
+
19
+ result = deviceProp.warpSize;
20
+ is_cached = true;
21
+ }
22
+
23
+ return result;
24
+ }
25
+
26
+ static __device__ constexpr int get_warp_size() {
27
+ #ifdef __GFX9__
28
+ return 64;
29
+ #else
30
+ return 32;
31
+ #endif
32
+ }
33
+};
34
+
35
+ #define WARP_SIZE Utils::get_warp_size()
36
#else
37
#define WARP_SIZE 32
38
#endif
csrc/moe/topk_softmax_kernels.cu
+29
-18
@@ -190,8 +190,8 @@ __launch_bounds__(TPB) __global__ void moeTopK(
190
2) This implementation assumes k is small, but will work for any k.
191
*/
192
193
-template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, typename IndType>
194
-__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
193
+template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WARP_SIZE_PARAM, typename IndType>
194
+__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
195
void topkGatingSoftmax(const float* input, const bool* finished, float* output, const int num_rows, IndType* indices,
196
int* source_rows, const int k, const int start_expert, const int end_expert)
197
{
@@ -209,12 +209,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
209
210
// Restrictions based on previous section.
211
static_assert(VPT % ELTS_PER_LDG == 0, "The elements per thread must be a multiple of the elements per ldg");
212
- static_assert(WARP_SIZE % THREADS_PER_ROW == 0, "The threads per row must cleanly divide the threads per warp");
212
+ static_assert(WARP_SIZE_PARAM % THREADS_PER_ROW == 0, "The threads per row must cleanly divide the threads per warp");
213
static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), "THREADS_PER_ROW must be power of 2");
214
- static_assert(THREADS_PER_ROW <= WARP_SIZE, "THREADS_PER_ROW can be at most warp size");
214
+ static_assert(THREADS_PER_ROW <= WARP_SIZE_PARAM, "THREADS_PER_ROW can be at most warp size");
215
216
// We have NUM_EXPERTS elements per row. We specialize for small #experts
217
- static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT;
217
+ static constexpr int ELTS_PER_WARP = WARP_SIZE_PARAM * VPT;
218
static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
219
static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
220
@@ -393,41 +393,51 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
393
namespace detail
394
{
395
// Constructs some constants needed to partition the work across threads at compile time.
396
-template <int EXPERTS, int BYTES_PER_LDG>
396
+template <int EXPERTS, int BYTES_PER_LDG, int WARP_SIZE_PARAM>
397
struct TopkConstants
398
{
399
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(float);
400
- static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 || EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0, "");
401
- static constexpr int VECs_PER_THREAD = MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE));
400
+ static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0 || EXPERTS % (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0, "");
401
+ static constexpr int VECs_PER_THREAD = MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM));
402
static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
403
static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
404
- static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW;
404
+ static const int ROWS_PER_WARP = WARP_SIZE_PARAM / THREADS_PER_ROW;
405
};
406
} // namespace detail
407
408
-template <int EXPERTS, int WARPS_PER_TB, typename IndType>
408
+template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, typename IndType>
409
void topkGatingSoftmaxLauncherHelper(const float* input, const bool* finished, float* output, IndType* indices,
410
int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, cudaStream_t stream)
411
{
412
static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
413
414
static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(float) * EXPERTS);
415
- using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG>;
415
+ using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM>;
416
static constexpr int VPT = Constants::VPT;
417
static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
418
const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
419
const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
420
421
- dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
422
- topkGatingSoftmax<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG><<<num_blocks, block_dim, 0, stream>>>(
421
+ dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
422
+ topkGatingSoftmax<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM><<<num_blocks, block_dim, 0, stream>>>(
423
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert);
424
}
425
426
-#define LAUNCH_SOFTMAX(NUM_EXPERTS, WARPS_PER_TB) \
427
- topkGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB>( \
428
- gating_output, nullptr, topk_weights, topk_indices, \
429
- token_expert_indices, num_tokens, topk, 0, num_experts, \
430
- stream);
426
+#define LAUNCH_SOFTMAX(NUM_EXPERTS, WARPS_PER_TB) \
427
+ switch (warpSize) { \
428
+ case 32: \
429
+ topkGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32>( \
430
+ gating_output, nullptr, topk_weights, topk_indices, \
431
+ token_expert_indices, num_tokens, topk, 0, num_experts, stream); \
432
+ break; \
433
+ case 64: \
434
+ topkGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64>( \
435
+ gating_output, nullptr, topk_weights, topk_indices, \
436
+ token_expert_indices, num_tokens, topk, 0, num_experts, stream); \
437
+ break; \
438
+ default: \
439
+ TORCH_CHECK(false, "Unsupported warp size: ", warpSize); \
440
+ }
441
442
template <typename IndType>
443
void topkGatingSoftmaxKernelLauncher(
@@ -441,6 +451,7 @@ void topkGatingSoftmaxKernelLauncher(
451
const int topk,
452
cudaStream_t stream) {
453
static constexpr int WARPS_PER_TB = 4;
454
+ auto warpSize = WARP_SIZE;
455
switch (num_experts) {
456
case 1:
457
LAUNCH_SOFTMAX(1, WARPS_PER_TB);
csrc/quantization/activation_kernels.cu
+1
-1
@@ -4,7 +4,7 @@
4
5
#include <cmath>
6
#include "core/math.hpp"
7
-#include "cuda_compat.h"
7
+#include "../cuda_compat.h"
8
#include "dispatch_utils.h"
9
10
#include "quantization/fp8/common.cuh"
csrc/quantization/gguf/gguf_kernel.cu
+1
-1
@@ -4,7 +4,7 @@
4
#include <torch/all.h>
5
#include <c10/cuda/CUDAGuard.h>
6
7
-#include "cuda_compat.h"
7
+#include "../../cuda_compat.h"
8
#include "dispatch_utils.h"
9
10
#include "ggml-common.h"
csrc/rocm/attention.cu
+1
-1
@@ -19,7 +19,7 @@
19
#include <c10/cuda/CUDAGuard.h>
20
#include <hip/hip_fp8.h>
21
#include <hip/hip_bf16.h>
22
-#include "cuda_compat.h"
22
+#include "../cuda_compat.h"
23
24
#include <algorithm>
25
#include "../attention/dtype_fp8.cuh"
csrc/rocm/skinny_gemms.cu
+1
-1
@@ -9,7 +9,7 @@
9
#include <stdexcept>
10
#include <algorithm>
11
12
-#include "cuda_compat.h"
12
+#include "../cuda_compat.h"
13
#include "dispatch_utils.h"
14
#include "quantization/fp8/common.cuh"
15