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 9, 17:06 UTC
[ROCm][ViT] Enable Flash Attention Triton backend on RDNA3/RDNA4 (#32944)
Signed-off-by: mohammad najafi <mohammad.najafi@amd.com>
monajafi-amd committed
Jan 23, 2026 at 19:03 UTC
97ef11dd3406ab7a9080921c62f001d60399d928
1 file changed
+34
-1
vllm/platforms/rocm.py
+34
-1
@@ -163,6 +163,28 @@ def use_rocm_custom_paged_attention(
163
)
164
165
166
+@cache
167
+def flash_attn_triton_available() -> bool:
168
+ if not on_gfx1x():
169
+ return False
170
+ try:
171
+ from importlib.util import find_spec
172
+
173
+ if find_spec("flash_attn") is None:
174
+ return False
175
+ if find_spec("flash_attn.flash_attn_triton_amd") is None:
176
+ return False
177
+ if os.environ.get("FLASH_ATTENTION_TRITON_AMD_ENABLE") != "TRUE":
178
+ logger.info_once(
179
+ "Set FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE to enable "
180
+ "Flash Attention Triton backend on RDNA."
181
+ )
182
+ return False
183
+ return True
184
+ except ImportError:
185
+ return False
186
+
187
+
188
class RocmPlatform(Platform):
189
_enum = PlatformEnum.ROCM
190
device_name: str = "rocm"
@@ -348,7 +370,7 @@ class RocmPlatform(Platform):
370
371
from vllm._aiter_ops import rocm_aiter_ops
372
351
- if rocm_aiter_ops.is_enabled():
373
+ if rocm_aiter_ops.is_enabled() and on_gfx9():
374
logger.info_once("Using AITER Flash Attention backend for ViT model.")
375
return AttentionBackendEnum.ROCM_AITER_FA
376
@@ -360,6 +382,17 @@ class RocmPlatform(Platform):
382
logger.info_once("Using Flash Attention backend for ViT model.")
383
return AttentionBackendEnum.FLASH_ATTN
384
385
+ # RDNA3/RDNA4 (gfx11xx/gfx12xx): Use Flash Attention Triton backend
386
+ if (
387
+ on_gfx1x()
388
+ and flash_attn_triton_available()
389
+ and (dtype == torch.float16 or dtype == torch.bfloat16)
390
+ ):
391
+ logger.info_once(
392
+ "Using Flash Attention (Triton backend) for ViT model on RDNA."
393
+ )
394
+ return AttentionBackendEnum.FLASH_ATTN
395
+
396
logger.info_once("Using Torch SDPA backend for ViT model.")
397
return AttentionBackendEnum.TORCH_SDPA
398