[CI][AMD] Gate Aiter MHA FA on dtype (#60679)

Signed-off-by: Ryan Rock <ryan.rock@amd.com> Co-authored-by: Micah Williamson <micah.williamson@amd.com>
2026-10-13 18:55:19 -05:00
parent 7a9d3db896
commit f43d7e1a0f
1 changed files with 5 additions and 1 deletions
@@ -887,7 +887,11 @@ class RocmPlatform(Platform):

        from vllm._aiter_ops import rocm_aiter_ops

-        if rocm_aiter_ops.is_mha_enabled() and on_cdna():
+        if (
+            rocm_aiter_ops.is_mha_enabled()
+            and on_cdna()
+            and (dtype == torch.float16 or dtype == torch.bfloat16)
+        ):
            logger.info_once("Using AITER Flash Attention backend for ViT model.")
            return AttentionBackendEnum.ROCM_AITER_FA