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 13, 09:30 UTC
[ROCm] [Bugfix] Fix `fused_qknorm_rope_kernel` rocm compatibility (#28500)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
TJian committed
Nov 12, 2025 at 05:01 UTC
edb59a9470f5c67ef11d52e7bb25fb8ea17f120f
6 files changed
+37
-38
csrc/fused_qknorm_rope_kernel.cu
+27
-27
@@ -35,10 +35,12 @@
35
CHECK_TH_CUDA(x); \
36
CHECK_CONTIGUOUS(x)
37
38
-#define FINAL_MASK 0xffffffff
38
+#ifdef USE_ROCM
39
+ #define FINAL_MASK 0xffffffffffffffffULL
40
+#else
41
+ #define FINAL_MASK 0xffffffff
42
+#endif
43
40
-// TODO: suport for AMD ROCM platform
41
-#ifndef USE_ROCM
44
namespace tensorrt_llm::common {
45
template <typename T, int num>
46
struct packed_as;
@@ -60,7 +62,7 @@ struct packed_as<uint, 4> {
62
63
template <typename T>
64
__inline__ __device__ T warpReduceSum(T val) {
63
- #pragma unroll
65
+#pragma unroll
66
for (int mask = 16; mask > 0; mask >>= 1)
67
val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
68
return val;
@@ -97,12 +99,12 @@ __global__ void fusedQKNormRopeKernel(
99
int64_t const* position_ids, // Position IDs for RoPE
100
int const num_tokens // Number of tokens
101
) {
100
- #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800
102
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
103
if constexpr ((std::is_same_v<scalar_t_in, c10::BFloat16>) ||
104
std::is_same_v<scalar_t_cache, c10::BFloat16>) {
105
return;
106
} else {
105
- #endif
107
+#endif
108
109
using Converter = vllm::_typeConvert<scalar_t_in>;
110
static_assert(Converter::exists,
@@ -179,7 +181,7 @@ __global__ void fusedQKNormRopeKernel(
181
{
182
vec_T vec = *reinterpret_cast<vec_T const*>(&qkv[offsetThread]);
183
constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
182
- #pragma unroll
184
+#pragma unroll
185
for (int i = 0; i < num_packed_elems; i++) {
186
// Interpret the generic vector chunk as the specific packed type
187
T2_in packed_val = *(reinterpret_cast<T2_in*>(&vec) + i);
@@ -200,7 +202,7 @@ __global__ void fusedQKNormRopeKernel(
202
float rms_rcp = rsqrtf(sumOfSquares / static_cast<float>(head_dim) + eps);
203
204
// Normalize elements
203
- #pragma unroll
205
+#pragma unroll
206
for (int i = 0; i < numElemsPerThread; i++) {
207
int dim = laneId * numElemsPerThread + i;
208
float weight = isQ ? Converter::convert(q_weight[dim])
@@ -222,7 +224,7 @@ __global__ void fusedQKNormRopeKernel(
224
225
if constexpr (interleave) {
226
// Perform interleaving. Use pre-computed cos/sin values.
225
- #pragma unroll
227
+#pragma unroll
228
for (int i = 0; i < numElemsPerThread / 2; ++i) {
229
int const idx0 = 2 * i;
230
int const idx1 = 2 * i + 1;
@@ -245,9 +247,9 @@ __global__ void fusedQKNormRopeKernel(
247
__syncwarp();
248
// Get the data from the other half of the warp. Use pre-computed cos/sin
249
// values.
248
- #pragma unroll
250
+#pragma unroll
251
for (int i = 0; i < numElemsPerThread; i++) {
250
- elements2[i] = __shfl_xor_sync(0xffffffff, elements[i], 16);
252
+ elements2[i] = __shfl_xor_sync(FINAL_MASK, elements[i], 16);
253
if (laneId < 16) {
254
elements2[i] = -elements2[i];
255
}
@@ -269,7 +271,7 @@ __global__ void fusedQKNormRopeKernel(
271
{
272
vec_T vec;
273
constexpr int num_packed_elems = elemSizeBytes / sizeof(T2_in);
272
- #pragma unroll
274
+#pragma unroll
275
for (int i = 0; i < num_packed_elems; i++) {
276
// Convert from float2 back to the specific packed type
277
T2_in packed_val = Converter::convert(
@@ -280,21 +282,21 @@ __global__ void fusedQKNormRopeKernel(
282
*reinterpret_cast<vec_T*>(&qkv[offsetThread]) = vec;
283
}
284
283
- #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800
285
+#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
286
}
285
- #endif
287
+#endif
288
}
289
288
- // Borrowed from
289
- // https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568
290
- #define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \
291
- if (interleave) { \
292
- const bool INTERLEAVE = true; \
293
- __VA_ARGS__ \
294
- } else { \
295
- const bool INTERLEAVE = false; \
296
- __VA_ARGS__ \
297
- }
290
+// Borrowed from
291
+// https://github.com/flashinfer-ai/flashinfer/blob/8125d079a43e9a0ba463a4ed1b639cefd084cec9/include/flashinfer/pos_enc.cuh#L568
292
+#define DISPATCH_INTERLEAVE(interleave, INTERLEAVE, ...) \
293
+ if (interleave) { \
294
+ const bool INTERLEAVE = true; \
295
+ __VA_ARGS__ \
296
+ } else { \
297
+ const bool INTERLEAVE = false; \
298
+ __VA_ARGS__ \
299
+ }
300
301
template <typename scalar_t_in, typename scalar_t_cache>
302
void launchFusedQKNormRope(void* qkv, int const num_tokens,
@@ -413,6 +415,4 @@ void fused_qk_norm_rope(
415
stream);
416
});
417
});
416
-}
417
-
418
-#endif // not USE_ROCM
\ No newline at end of file
418
+}
\ No newline at end of file
csrc/torch_bindings.cpp
-2
@@ -175,7 +175,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
175
"float epsilon) -> ()");
176
ops.impl("fused_add_rms_norm", torch::kCUDA, &fused_add_rms_norm);
177
178
-#ifndef USE_ROCM
178
// Function for fused QK Norm and RoPE
179
ops.def(
180
"fused_qk_norm_rope(Tensor! qkv, int num_heads_q, "
@@ -183,7 +182,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
182
"Tensor q_weight, Tensor k_weight, Tensor cos_sin_cache, "
183
"bool is_neox, Tensor position_ids) -> ()");
184
ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
186
-#endif
185
186
// Apply repetition penalties to logits in-place
187
ops.def(
csrc/type_convert.cuh
+4
-3
@@ -67,9 +67,9 @@ struct _typeConvert<c10::Half> {
67
}
68
};
69
70
- #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
70
+ #if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800) || defined(USE_ROCM)
71
// CUDA_ARCH < 800 does not have BF16 support
72
-// TODO: Add in ROCm support once public headers handle bf16 maturely
72
+// ROCm 7.0+ supports bfloat16
73
template <>
74
struct _typeConvert<c10::BFloat16> {
75
static constexpr bool exists = true;
@@ -89,7 +89,8 @@ struct _typeConvert<c10::BFloat16> {
89
return __float22bfloat162_rn(x);
90
}
91
};
92
- #endif // defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
92
+ #endif // (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800) ||
93
+ // defined(USE_ROCM)
94
#endif // defined(USE_ROCM) || (defined(CUDA_VERSION) && (CUDA_VERSION >=
95
// 12000))
96
tests/compile/test_qk_norm_rope_fusion.py
+2
-2
@@ -113,8 +113,8 @@ class QKNormRoPETestModel(torch.nn.Module):
113
@pytest.mark.parametrize("enable_rope_custom_op", [True])
114
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
115
@pytest.mark.skipif(
116
- not current_platform.is_cuda(),
117
- reason="Only test on cuda platform",
116
+ not current_platform.is_cuda_alike(),
117
+ reason="Only test on cuda and rocm platform",
118
)
119
def test_qk_norm_rope_fusion(
120
eps, is_neox, enable_rms_norm_custom_op, enable_rope_custom_op, dtype
tests/kernels/core/test_fused_qk_norm_rope.py
+2
-2
@@ -44,8 +44,8 @@ def _apply_qk_norm_rope(
44
45
46
@pytest.mark.skipif(
47
- not current_platform.is_cuda(),
48
- reason="fused_qk_norm_rope custom op requires cuda platform",
47
+ not current_platform.is_cuda_alike(),
48
+ reason="fused_qk_norm_rope custom op requires cuda and rocm platform",
49
)
50
@pytest.mark.parametrize("device", CUDA_DEVICES)
51
@pytest.mark.parametrize("dtype", DTYPES)
vllm/config/compilation.py
+2
-2
@@ -184,10 +184,10 @@ class PassConfig:
184
"Fusion enabled but reshape elimination disabled. "
185
"Allreduce + rms norm + quant (fp8) fusion might not work"
186
)
187
- if self.enable_qk_norm_rope_fusion and not current_platform.is_cuda():
187
+ if self.enable_qk_norm_rope_fusion and not current_platform.is_cuda_alike():
188
logger.warning_once(
189
"QK Norm + RoPE fusion enabled but the current platform is not "
190
- "CUDA. The fusion will be disabled."
190
+ "CUDA or ROCm. The fusion will be disabled."
191
)
192
self.enable_qk_norm_rope_fusion = False
193