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 10, 09:00 UTC
[Bugfix][Hardware][AMD] Use dynamic WARP_SIZE in sampler vectorized_process (#31295)
Signed-off-by: c0de128 <kevin.mckay@outlook.com>
Kevin McKay committed
Jan 9, 2026 at 21:57 UTC
c60578de0a389d261fcb4d3323c75387b98b1fbb
1 file changed
+7
-4
csrc/sampler.cu
+7
-4
@@ -1,3 +1,4 @@
1
+#include "cuda_compat.h"
2
#include "dispatch_utils.h"
3
4
#include <torch/cuda.h>
@@ -97,7 +98,9 @@ static inline __device__ bool isPartialMatch(float x, uint32_t pattern) {
98
template <typename T, typename idxT, typename Func>
99
__device__ void vectorized_process(size_t thread_rank, size_t num_threads,
100
const T* in, idxT len, Func f) {
100
- constexpr int WARP_SIZE = 32;
101
+ // Use dynamic WARP_SIZE from cuda_compat.h to support both
102
+ // Wave64 (MI300X/gfx942) and Wave32 (Strix Halo/gfx1151) architectures
103
+ constexpr int kWarpSize = WARP_SIZE;
104
using WideT = float4;
105
if constexpr (sizeof(T) >= sizeof(WideT)) {
106
for (idxT i = thread_rank; i < len; i += num_threads) {
@@ -132,8 +135,8 @@ __device__ void vectorized_process(size_t thread_rank, size_t num_threads,
135
}
136
}
137
135
- static_assert(WARP_SIZE >= items_per_scalar);
136
- // and because items_per_scalar > skip_cnt, WARP_SIZE > skip_cnt
138
+ static_assert(kWarpSize >= items_per_scalar);
139
+ // and because items_per_scalar > skip_cnt, kWarpSize > skip_cnt
140
// no need to use loop
141
if (thread_rank < skip_cnt) {
142
f(in[thread_rank], thread_rank);
@@ -142,7 +145,7 @@ __device__ void vectorized_process(size_t thread_rank, size_t num_threads,
145
// len_cast * items_per_scalar + items_per_scalar > len - skip_cnt;
146
// and so
147
// len - (skip_cnt + len_cast * items_per_scalar) < items_per_scalar <=
145
- // WARP_SIZE no need to use loop
148
+ // kWarpSize no need to use loop
149
const idxT remain_i = skip_cnt + len_cast * items_per_scalar + thread_rank;
150
if (remain_i < len) {
151
f(in[remain_i], remain_i);