Add vectorized rms_norm support for Navi31

- supports vectorized rms_norm_kernel
ROCm · Nov 12, 2024 · 867c12d · 867c12d
1 parent 8f3bf8b
commit 867c12d
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/csrc/layernorm_kernels.cu b/csrc/layernorm_kernels.cu
@@ -27,9 +27,10 @@ using __nv_bfloat162 = __hip_bfloat162;
   #include "quantization/fp8/nvidia/quant_utils.cuh"
 #endif
 
-#if defined(__HIPCC__) && (defined(__gfx90a__) || defined(__gfx940__) || \
-                           defined(__gfx941__) || defined(__gfx942__))
-  #define __HIP__MI300_MI250__
+#if defined(__HIPCC__) &&                                                 \
+    (defined(__gfx90a__) || defined(__gfx940__) || defined(__gfx941__) || \
+     defined(__gfx942__) || defined(__gfx1100__))
+  #define __HIP__MI300_MI250_Navi31__
 #endif
 
 namespace vllm {
@@ -72,7 +73,7 @@ struct __align__(16) vec8_t {
   __device__ scalar_t sum() const { return x + y + z + w + u + v + s + t; }
 };
 
-#ifdef __HIP__MI300_MI250__
+#ifdef __HIP__MI300_MI250_Navi31__
 
 // TODO(woosuk): Further optimize this kernel.
 template <typename scalar_t>