From 199b8873d7685676f0ef81f0cc983cb48a394a59 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B5=AE=E6=A2=A6?= <914814442@qq.com>
Date: Thu, 13 Aug 2026 20:39:01 +0800
Subject: [PATCH] [train] add fa3 supported (#10742)

---
 src/llamafactory/model/model_utils/attention.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/llamafactory/model/model_utils/attention.py b/src/llamafactory/model/model_utils/attention.py
index 290df74a4..1edf3ebcd 100644
--- a/src/llamafactory/model/model_utils/attention.py
+++ b/src/llamafactory/model/model_utils/attention.py
@@ -82,6 +82,14 @@ def configure_attn_implementation(config: "PretrainedConfig", model_args: "Model
             return
 
         requested_attn_implementation = "flash_attention_2"
+    elif model_args.flash_attn == AttentionFunction.FA3:
+        from transformers.utils import is_flash_attn_3_available
+
+        if not is_flash_attn_3_available():
+            logger.warning_rank0("FlashAttention-3 is not installed.")
+            return
+
+        requested_attn_implementation = "flash_attention_3"
     else:
         raise NotImplementedError(f"Unknown attention type: {model_args.flash_attn}")
 
@@ -109,6 +117,8 @@ def print_attn_implementation(config: "PretrainedConfig") -> None:
 
     if attn_implementation == "flash_attention_2":
         logger.info_rank0("Using FlashAttention-2 for faster training and inference.")
+    elif attn_implementation == "flash_attention_3":
+        logger.info_rank0("Using FlashAttention-3 for faster training and inference.")
     elif attn_implementation == "sdpa":
         logger.info_rank0("Using torch SDPA for faster training and inference.")
     else: