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 15, 12:15 UTC
[ROCM] Add gfx950 to the custom attention archs (#16034)
Signed-off-by: jpvillam <Juan.Villamizar@amd.com> Signed-off-by: seungrokjung <seungrok.jung@amd.com> Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com> Co-authored-by: seungrokjung <seungrok.jung@amd.com> Co-authored-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Juan Villamizar committed
May 1, 2025 at 13:18 UTC
811a6c0972daa4c5cfa7e6280ba22a24324a1141
2 files changed
+12
-8
csrc/rocm/attention.cu
+6
-5
@@ -25,8 +25,9 @@
25
#include "../attention/dtype_fp8.cuh"
26
#include "../quantization/fp8/amd/quant_utils.cuh"
27
28
-#if defined(__HIPCC__) && (defined(__gfx90a__) || defined(__gfx942__))
29
- #define __HIP__MI300_MI250__
28
+#if defined(__HIPCC__) && \
29
+ (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__))
30
+ #define __HIP__GFX9__
31
#endif
32
33
#if defined(NDEBUG)
@@ -42,7 +43,7 @@
43
#define MIN(a, b) ((a) < (b) ? (a) : (b))
44
#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
45
45
-#if defined(__HIP__MI300_MI250__) // TODO: Add NAVI support
46
+#if defined(__HIP__GFX9__) // TODO: Add NAVI support
47
48
#define GCN_MFMA_INSTR1 __builtin_amdgcn_mfma_f32_16x16x4f32
49
#define GCN_MFMA_INSTR __builtin_amdgcn_mfma_f32_4x4x4f16
@@ -1479,7 +1480,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1480
}
1481
}
1482
1482
-#else // !defined(__HIP__MI300_MI250__) TODO: Add NAVI support
1483
+#else // !defined(__HIP__GFX9__) TODO: Add NAVI support
1484
1485
// clang-format off
1486
template <typename scalar_t, typename cache_t,
@@ -1552,7 +1553,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1553
}
1554
// clang-format on
1555
1555
-#endif // defined(__HIP__MI300_MI250__) TODO: Add NAVI support
1556
+#endif // defined(__HIP__GFX9__) TODO: Add NAVI support
1557
1558
#define LAUNCH_CUSTOM_ATTENTION_MFMA16(GQA_RATIO) \
1559
paged_attention_ll4mi_QKV_mfma16_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, \
vllm/platforms/rocm.py
+6
-3
@@ -106,11 +106,14 @@ def use_rocm_custom_paged_attention(qtype: torch.dtype, head_size: int,
106
max_seq_len: int,
107
sliding_window: int) -> bool:
108
109
+ GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
110
+ ON_GFX9 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
111
+
112
# rocm custom page attention not support on gfx1*
113
# custom paged attn always supported on V0. On V1, requires sliding window
114
# disabled due to observed numerical discrepancy.
112
- return (on_mi250_mi300() and (not envs.VLLM_USE_V1 or sliding_window == 0
113
- or sliding_window == (-1, -1))
115
+ return (ON_GFX9 and (not envs.VLLM_USE_V1 or sliding_window == 0
116
+ or sliding_window == (-1, -1))
117
and (qtype == torch.half or qtype == torch.bfloat16)
118
and (head_size == 64 or head_size == 128)
119
and (block_size == 16 or block_size == 32)
@@ -337,7 +340,7 @@ class RocmPlatform(Platform):
340
def use_custom_allreduce(cls) -> bool:
341
# We only enable custom allreduce for MI300 series
342
gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
340
- supported_archs = ['gfx94']
343
+ supported_archs = ['gfx94', 'gfx95']
344
return any(gfx in gcn_arch for gfx in supported_archs)
345
346
@classmethod