diff --git a/.gitmodules b/.gitmodules
index d33d0a64..f0cfb7ac 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -10,3 +10,6 @@
 [submodule "third_party/googletest"]
 	path = third_party/googletest
 	url = git@github.com:google/googletest.git
+[submodule "third_party/flash-attention"]
+	path = third_party/flash-attention
+	url = git@github.com:Dao-AILab/flash-attention.git
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6bd8069d..99d88e2c 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -5,6 +5,7 @@ option(PROFILE_MODE "ENABLE PROFILE MODE" OFF)
 option(USE_OMP "Use OpenMP as backend for Eigen" ON)
 option(USE_NCCL "Build project for distributed running" ON)
 option(BUILD_TEST "Build InfiniTrain tests" OFF)
+option(USE_FLASH_ATTENTION "Enable FlashAttention 2 CUDA backend" ON)
 
 project(infini_train VERSION 0.6.0 LANGUAGES CXX)
 
@@ -97,11 +98,70 @@ if(USE_CUDA)
   find_package(CUDAToolkit REQUIRED)
   include_directories(${CUDAToolkit_INCLUDE_DIRS})
 
+  if(USE_FLASH_ATTENTION)
+    set(FLASH_ATTN_SOURCE_DIR "${PROJECT_SOURCE_DIR}/third_party/flash-attention" CACHE PATH
+      "Path to vendored FlashAttention source tree")
+    if(NOT EXISTS "${FLASH_ATTN_SOURCE_DIR}/csrc/flash_attn/src/flash.h")
+      message(FATAL_ERROR "FlashAttention submodule not found at ${FLASH_ATTN_SOURCE_DIR}. "
+        "Run: git submodule update --init --recursive third_party/flash-attention")
+    endif()
+    if(NOT EXISTS "${FLASH_ATTN_SOURCE_DIR}/csrc/cutlass/include/cutlass/cutlass.h")
+      message(FATAL_ERROR "FlashAttention CUTLASS dependency not found. "
+        "Run: git -C third_party/flash-attention submodule update --init csrc/cutlass")
+    endif()
+
+    # Minimal training subset: causal FP16/BF16 kernels for GPT-2/LLaMA head dims.
+    # TODO: add other head dims and architectures when InfiniTrain models require them.
+    set(FLASH_ATTN_CUDA_SOURCES
+      "${FLASH_ATTN_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_hdim64_fp16_causal_sm80.cu"
+      "${FLASH_ATTN_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_hdim64_bf16_causal_sm80.cu"
+      "${FLASH_ATTN_SOURCE_DIR}/csrc/flash_attn/src/flash_bwd_hdim64_fp16_causal_sm80.cu"
+      "${FLASH_ATTN_SOURCE_DIR}/csrc/flash_attn/src/flash_bwd_hdim64_bf16_causal_sm80.cu"
+      "${FLASH_ATTN_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_hdim128_fp16_causal_sm80.cu"
+      "${FLASH_ATTN_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_hdim128_bf16_causal_sm80.cu"
+      "${FLASH_ATTN_SOURCE_DIR}/csrc/flash_attn/src/flash_bwd_hdim128_fp16_causal_sm80.cu"
+      "${FLASH_ATTN_SOURCE_DIR}/csrc/flash_attn/src/flash_bwd_hdim128_bf16_causal_sm80.cu"
+    )
+    add_library(flash_attn_native STATIC ${FLASH_ATTN_CUDA_SOURCES})
+    set_target_properties(flash_attn_native PROPERTIES
+      CUDA_STANDARD 17
+      CUDA_STANDARD_REQUIRED ON
+      CUDA_ARCHITECTURES "80"
+      POSITION_INDEPENDENT_CODE ON
+    )
+    target_compile_definitions(flash_attn_native PRIVATE
+      FLASHATTENTION_DISABLE_DROPOUT
+      FLASHATTENTION_DISABLE_ALIBI
+      FLASHATTENTION_DISABLE_SOFTCAP
+      FLASHATTENTION_DISABLE_LOCAL
+    )
+    target_include_directories(flash_attn_native PRIVATE
+      "${PROJECT_SOURCE_DIR}/infini_train/src/kernels/cuda/flash_attention_compat"
+      "${FLASH_ATTN_SOURCE_DIR}/csrc/flash_attn"
+      "${FLASH_ATTN_SOURCE_DIR}/csrc/flash_attn/src"
+      "${FLASH_ATTN_SOURCE_DIR}/csrc/cutlass/include"
+    )
+    target_compile_options(flash_attn_native PRIVATE
+      $<$<COMPILE_LANGUAGE:CUDA>:-O3>
+      $<$<COMPILE_LANGUAGE:CUDA>:-U__CUDA_NO_HALF_OPERATORS__>
+      $<$<COMPILE_LANGUAGE:CUDA>:-U__CUDA_NO_HALF_CONVERSIONS__>
+      $<$<COMPILE_LANGUAGE:CUDA>:-U__CUDA_NO_HALF2_OPERATORS__>
+      $<$<COMPILE_LANGUAGE:CUDA>:-U__CUDA_NO_BFLOAT16_CONVERSIONS__>
+      $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>
+      $<$<COMPILE_LANGUAGE:CUDA>:--expt-extended-lambda>
+      $<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>
+    )
+    target_link_libraries(flash_attn_native PUBLIC CUDA::cudart)
+  endif()
+
   # CUDA compilation options
   set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-extended-lambda --expt-relaxed-constexpr")
 
   # Only compile CUDA kernels / cuda sources here (your original used src/*.cu)
   file(GLOB_RECURSE CUDA_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/*.cu)
+  if(NOT USE_FLASH_ATTENTION)
+    list(FILTER CUDA_KERNELS EXCLUDE REGEX ".*/kernels/cuda/flash_attention\\.cu$")
+  endif()
 
   add_library(infini_train_cuda_kernels STATIC ${CUDA_KERNELS})
   set_target_properties(infini_train_cuda_kernels PROPERTIES CUDA_ARCHITECTURES "75;80;90")
@@ -114,6 +174,21 @@ if(USE_CUDA)
       CUDA::cuda_driver
   )
 
+  if(USE_FLASH_ATTENTION)
+    target_compile_definitions(infini_train_cuda_kernels PRIVATE
+      FLASHATTENTION_DISABLE_DROPOUT
+      FLASHATTENTION_DISABLE_ALIBI
+      FLASHATTENTION_DISABLE_SOFTCAP
+      FLASHATTENTION_DISABLE_LOCAL
+    )
+    target_include_directories(infini_train_cuda_kernels PRIVATE
+      "${PROJECT_SOURCE_DIR}/infini_train/src/kernels/cuda/flash_attention_compat"
+      "${FLASH_ATTN_SOURCE_DIR}/csrc/flash_attn/src"
+      "${FLASH_ATTN_SOURCE_DIR}/csrc/cutlass/include"
+    )
+    target_link_libraries(infini_train_cuda_kernels PUBLIC flash_attn_native)
+  endif()
+
   if(USE_NCCL)
     message(STATUS "Add USE_NCCL, use NCCL with CUDA")
     list(APPEND CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake)
@@ -151,6 +226,7 @@ if(USE_CUDA)
     # keep this. Otherwise it's harmless.
     target_link_libraries(infini_train PUBLIC nccl)
   endif()
+
 endif()
 
 # ------------------------------------------------------------------------------
diff --git a/docs/flash_attn_integration_design.md b/docs/flash_attn_integration_design.md
new file mode 100644
index 00000000..7515d956
--- /dev/null
+++ b/docs/flash_attn_integration_design.md
@@ -0,0 +1,347 @@
+# FlashAttention 2 后端设计
+
+## 目标与边界
+
+InfiniTrain 通过 attention backend 选项在原有 unfused causal attention 和
+FlashAttention 2 之间切换：
+
+```bash
+--attention_backend=unfused
+--attention_backend=flash
+```
+
+`unfused` 使用框架算子显式计算 attention score、causal mask、softmax 和 value
+聚合；`flash` 调用 FlashAttention 2 的 fused CUDA kernel，不生成完整的 `T x T`
+attention 中间张量。
+
+本接入面向训练期、固定长度、causal self-attention。它不是 PyTorch extension，也不以
+兼容 `torch.nn.functional.scaled_dot_product_attention` 的完整参数语义为目标。当前
+functional API 只表达 Q/K/V 和 scale，mask、dropout、local window、KV cache 等能力尚未
+进入接口。
+
+## 分层与调用路径
+
+当前调用链为：
+
+```text
+GPT-2 / LLaMA3 CLI --attention_backend
+  -> TransformerConfig::flash
+  -> CausalSelfAttention::Forward
+       -> QKV projection / optional RoPE / MHA or GQA layout handling
+       -> nn::function::ScaledDotProductAttention
+  -> autograd::ScaledDotProductAttention
+       -> Dispatcher availability check
+       -> ScaledDotProductAttentionForward / ScaledDotProductAttentionBackward
+  -> Flash_fwd_params / Flash_bwd_params
+  -> run_mha_fwd_ / run_mha_bwd_
+  -> FlashAttention 2 CUDA kernel
+```
+
+各层职责如下：
+
+- example main：解析 backend，检查 CLI device/dtype，并把选择写入模型配置；
+- transformer module：处理 projection、位置编码、local heads、GQA 和输出 projection；
+- functional/autograd：定义可求导接口、选择 compute dtype，并提供统一 availability guard；
+- CUDA adapter：校验 kernel contract、管理 workspace/context、填充上游参数并调度 kernel；
+- CMake：只实例化当前模型需要的上游 CUDA 模板组合。
+
+## 源码与依赖
+
+FlashAttention 作为 submodule 固定在：
+
+- `third_party/flash-attention`
+- tag：`v2.7.4.post1`
+- commit：`5231d95fe13733fb534c01895f7ea88c6a6c7793`
+- CUTLASS：`third_party/flash-attention/csrc/cutlass`
+- CUTLASS commit：`c506e16788cb08416a4a57e11a9067beeee29420`
+
+初始化命令：
+
+```bash
+git submodule update --init --recursive third_party/flash-attention
+```
+
+本接入不编译上游 `flash_api.cpp`，不创建 `at::Tensor`，也不链接 Torch、ATen、c10、
+Python 或 `flash-attn` wheel。运行时只依赖 CUDA、InfiniTrain Tensor/allocator、
+DeviceGuard、当前 CUDA stream，以及 submodule 中的 FlashAttention/CUTLASS 源码。
+
+### ATen/c10 兼容头
+
+上游 kernel/launch 头仍残留少量 PyTorch 类型和宏：
+
+- `flash.h` 的参数 ABI 包含 `at::PhiloxCudaState`；
+- `flash_fwd_kernel.h` 调用 `at::cuda::philox::unpack()`；
+- forward/backward launch template 使用 `C10_CUDA_CHECK` 和
+  `C10_CUDA_KERNEL_LAUNCH_CHECK`。
+
+`infini_train/src/kernels/cuda/flash_attention_compat` 提供同路径的最小替代头，让这些
+引用在不安装 Torch 的情况下编译：
+
+- Philox state 只保留当前上游头所需的两个 64-bit 字段；
+- `unpack()` 只返回 seed/offset；
+- CUDA 检查宏基于 CUDA runtime error API 执行 fail-fast。
+
+当前构建禁用 dropout，不会实际消费随机状态。adapter 仍创建并清零一个两元素
+`uint64` rng-state buffer，以满足上游参数 ABI。
+
+该兼容层与固定版本的上游内部 header/ABI 耦合，不是稳定接口。升级 FlashAttention 或
+未来引入真实 ATen headers 时，必须重新检查类型布局、include 顺序和宏定义冲突，不能把
+“无需链接 Torch”等同于“完全没有 ATen/c10 header 形状依赖”。
+
+## 构建接入
+
+### CMake 开关
+
+当前 `USE_FLASH_ATTENTION` 默认值为 `ON`，但只有同时启用 `USE_CUDA` 才会创建 CUDA
+backend：
+
+```bash
+cmake -S . -B build \
+  -DUSE_CUDA=ON \
+  -DUSE_NCCL=ON \
+  -DUSE_FLASH_ATTENTION=ON \
+  -DBUILD_TEST=ON
+cmake --build build -j
+```
+
+`FLASH_ATTN_SOURCE_DIR` 是一个 CMake cache path，默认指向仓库内 submodule。配置阶段会
+检查 `flash.h` 和 CUTLASS 主头是否存在，缺失时直接报错。
+
+开启 backend 时：
+
+1. 创建静态库 `flash_attn_native`；
+2. 为 `infini_train_cuda_kernels` 编译并链接 native adapter；
+3. 注册 `ScaledDotProductAttentionForward/Backward` CUDA dispatcher kernels。
+
+关闭 backend 时，CMake 从 CUDA source 列表排除 `flash_attention.cu`，不会创建或链接
+`flash_attn_native`。functional/autograd 接口仍可编译，但第一次执行 FlashAttention
+operator 时会通过 `Dispatcher::HasKernel` 明确提示使用
+`-DUSE_FLASH_ATTENTION=ON` 重新构建。
+
+### AOT kernel 实例
+
+`flash_attn_native` 当前固定 `CUDA_ARCHITECTURES=80`，只编译 8 个训练实例：
+
+```text
+dtype:      fp16, bf16
+head_dim:   64, 128
+mask:       causal
+direction:  forward, backward
+```
+
+同时定义：
+
+```text
+FLASHATTENTION_DISABLE_DROPOUT
+FLASHATTENTION_DISABLE_ALIBI
+FLASHATTENTION_DISABLE_SOFTCAP
+FLASHATTENTION_DISABLE_LOCAL
+```
+
+这限制了模板数量和编译范围。增加 dtype、head dimension、GPU architecture 或 attention
+模式时，必须同步修改：
+
+- `FLASH_ATTN_CUDA_SOURCES` 中的 AOT 实例；
+- adapter 的 runtime 校验；
+- `RunForward` / `RunBackward` 的 dispatch 分支；
+- 相应自动化测试。
+
+InfiniTrain 其他 CUDA kernels 的 architecture 列表更宽，不代表 FlashAttention native
+kernel 已支持这些架构；当前 Flash backend 的有效架构边界仍是 sm80。
+
+## API 与 availability
+
+framework 入口为：
+
+```cpp
+nn::function::ScaledDotProductAttention(q, k, v, scale)
+```
+
+输入 contract：
+
+- Q 使用 `(B, Hq, T, D)`；
+- K/V 使用 `(B, Hkv, T, D)`，且 shape 相同；
+- Q/K/V 位于同一 CUDA device；
+- `Hq % Hkv == 0`；
+- `D` 为 64 或 128；
+- Q 和 K/V 当前使用相同 batch、sequence length 和 head dimension。
+
+API 没有 `enable_gqa`。`Hq == Hkv` 自动执行 MHA；`Hq > Hkv` 且可整除时自动执行
+GQA/MQA；不能整除时在 autograd/operator 边界报错。
+
+availability 不通过公共宏或 `nn` namespace 下的全局布尔值暴露。autograd operator 根据
+Q 的 device 构造 dispatcher key，并在 forward 调用前执行 `HasKernel`。这把构建能力
+检查放在 framework/operator 边界：example main 不需要包含构建宏，其他 framework
+调用者也能得到相同错误。
+
+## dtype 与 autocast
+
+native kernels 支持 fp16 和 bf16。compute dtype 的选择顺序为：
+
+1. Q 已是 fp16/bf16 时直接使用 Q dtype；
+2. 否则要求当前 autocast context 启用，且 autocast dtype 为 fp16/bf16；
+3. 其他情况 fail-fast。
+
+adapter 使用 `Tensor::To` 把 Q/K/V 和 backward 的 `grad_output` 转成 compute dtype，输出
+保持 compute dtype。GPT-2/LLaMA3 当前 CLI 只开放 `--dtype=bfloat16` 的 Flash 路径；
+fp16 是 framework/kernel 能力，但不是这两个 example 的已开放训练选项。
+
+BF16 backward 末尾当前会把 dQ/dK/dV 显式提升为 FP32。这是框架 autocast/autograd dtype
+语义尚未集中处理前的临时兼容逻辑：forward 的 raw `Tensor::To` 没有建立通用的
+cast-backward edge，混合 dtype 梯度也没有在 autograd/accumulation 层统一归一化。未来若
+框架层补齐该语义，应删除 adapter 内的特殊 upcast，让 kernel 返回其自然梯度 dtype。
+
+## Transformer 接入
+
+Flash 和 unfused backend 共用 `CausalSelfAttention::Forward`，不再区分
+`ForwardStandard` 与 `ForwardWithRoPE`。位置编码由 `PositionEmbeddingType` 决定：
+
+- `kLearnedAbsolute`：模型前段添加 WPE；attention module 创建内部 causal-mask buffer，
+  供 unfused fallback 使用；
+- `kRoPE`：`ApplyRotaryEmbedding` 在 attention backend 之前处理 Q/K；当前 transformer
+  调用者提供 runtime mask。
+
+统一路径中的 QKV 处理为：
+
+1. ColumnParallelLinear 产生 packed QKV；
+2. MHA 的 Q/K/V 宽度相等，使用单个 `Split` autograd node；
+3. GQA 的 Q 和 K/V 宽度不同，使用三个 `Slice`；
+4. RoPE 模型在此后旋转 Q/K；
+5. unfused backend 对 K/V 执行 `RepeatKV`；
+6. Flash backend 保留原始 KV heads，由 native kernel 根据 `Hq/Hkv` 处理 GQA；
+7. Q/K/V 转换到 `(B, H, T, D)` 后进入相应 backend。
+
+`Split` fast path 不改变统一 Forward 的语义。它避免在普通 MHA 中创建三个独立 Slice
+autograd nodes；GQA 因为分段宽度不同，仍必须使用 Slice 路径。
+
+### mask 与 start_pos
+
+当前 Flash functional API 没有 mask/start_pos 参数。Transformer Forward 虽然已经解析
+这两个输入，但选择 Flash backend 后会忽略它们，并固定设置：
+
+```text
+is_causal        = true
+window_size_left = -1
+window_size_right = 0
+seqlen_q         = seqlen_k
+```
+
+因此当前 Flash 语义仅适用于从位置 0 开始的标准 causal self-attention。外部自定义 mask、
+padding mask、非零 start position、incremental decoding 和 cross-attention 均不受支持。
+在这些能力进入 functional API 之前，调用者不能假设传入 mask/start_pos 会影响 Flash
+结果。
+
+## Native CUDA adapter
+
+入口文件为 `infini_train/src/kernels/cuda/flash_attention.cu`。它直接包含上游 `flash.h`，
+其中的 ATen/c10 引用由前述兼容 include path 满足；adapter 本身不构造任何 PyTorch
+对象。
+
+### 布局与 stride
+
+InfiniTrain 进入 adapter 的 Q/K/V 物理布局为连续 `(B, H, T, D)`。FlashAttention kernel
+按逻辑 `(B, T, H, D)` 解释，因此 adapter 设置元素 stride：
+
+```text
+batch_stride = H * T * D
+row_stride   = D
+head_stride  = T * D
+```
+
+Q 与 K/V 分别使用各自的 head count 计算 batch stride。输出沿用 Q 的 shape/stride，物理
+布局仍为 `(B, Hq, T, D)`，无需额外 Tensor wrapper 或布局复制。
+
+### Forward
+
+Forward 执行：
+
+1. 校验 device、shape、head ratio、dtype 和 head dimension；
+2. 将 Q/K/V 转为选定 compute dtype；
+3. 分配 output、FP32 `softmax_lse` 和零初始化 rng-state；
+4. 填充 `Flash_fwd_params`，固定 causal/dropout=0/num_splits=1；
+5. 在 InfiniTrain 当前 CUDA stream 上调度 AOT kernel；
+6. 把 kernel backward 所需状态保存到 opaque `FlashAttentionContext`。
+
+context 保存转换后的 Q/K/V、detach 后的 output、`softmax_lse`、rng-state、原始 dtype 和
+compute dtype。保存 detach output 是为了避免形成
+`Function -> flash_ctx -> output -> Function` 引用环。
+
+### Backward
+
+Backward 执行：
+
+1. 将 `grad_output` 转为 compute dtype；
+2. 分配 dQ/dK/dV；
+3. 分配 FP32 `dsoftmax_sum` 和 `dq_accum` workspace，其中 sequence length 向上取整到 128；
+4. 设置 `deterministic=false` 并调度 causal backward；
+5. 根据当前 autocast 兼容策略恢复返回梯度 dtype。
+
+MHA 时 kernel 直接写入最终 dK/dV。GQA/MQA 时，上游 backward launch 需要 Q-head shape 的
+临时 dK/dV；adapter 随后用自定义 CUDA reduction 按 group 求和回原始 KV-head shape。
+reduction 在寄存器中使用 FP32 累加，再写回 fp16/bf16 tensor。该实现会额外占用两个
+Q-head shape 临时 buffer，MQA 或长序列下需要关注其峰值内存。
+
+### Stream、device 与生命周期
+
+forward/backward 都使用 `DeviceGuard` 绑定 Q 所在 device，并通过 InfiniTrain
+`CudaStream::cuda_stream()` 获取当前 stream。adapter 不创建额外 stream，也不进行
+PyTorch stream guard 转换。
+
+所有 tensor 和 workspace 都由 InfiniTrain allocator 管理。forward 后仍需存活的数据由
+opaque context 持有，其余 backward 临时量由当前调用栈持有。kernel launch 保持同一
+stream 顺序，不依赖额外 host synchronization。
+
+## 当前支持矩阵
+
+| 维度 | 当前状态 |
+| --- | --- |
+| Device | CUDA |
+| GPU architecture | sm80 AOT kernels |
+| Kernel dtype | fp16、bf16 |
+| GPT-2/LLaMA3 CLI dtype | bf16 |
+| Attention type | fixed-length causal self-attention |
+| Head dimension | 64、128 |
+| Head mapping | MHA、GQA、MQA，要求 `Hq % Hkv == 0` |
+| Sequence relation | Q/K/V batch 和 sequence length 相同 |
+| Dropout | 0，编译期禁用 |
+| Backward | non-deterministic |
+| Mask | 仅 kernel 内建 causal mask |
+| Position | `start_pos=0` 语义 |
+| Unsupported | varlen、padding/custom mask、cross-attention、KV cache、generation、local attention、ALiBi、softcap、split-KV |
+
+## 自动化测试与开发入口
+
+`tests/autograd/test_autograd_scaled_dot_product_attention.cc` 构建为独立 CUDA target
+`test_flash_attention_cuda`，仅在 `USE_FLASH_ATTENTION=ON` 时注册。测试覆盖：
+
+- native GQA 与显式展开 KV 的 forward/backward contract；
+- packed QKV 经 Slice autograd 回传的梯度；
+- native GQA 与 unfused causal reference；
+- BF16 backward 的当前 FP32 gradient contract；
+- fused BF16 路径相对 BF16/FP32 reference 的误差边界。
+
+运行聚焦测试：
+
+```bash
+ctest --test-dir build --output-on-failure -R '^test_flash_attention_cuda$'
+ctest --test-dir build --output-on-failure -R '^test_transformer_cuda$'
+```
+
+`scripts/test_config.json` 还提供 `flash` tag 的 GPT-2/LLaMA3 端到端训练 cases，用于验证
+不同 batch/sequence shape 能完整执行 forward、backward 和 optimizer step。具体运行
+结果保存在独立测试日志或报告中。
+
+## 已知技术债与扩展顺序
+
+建议按以下依赖关系扩展：
+
+1. 在 functional API 中明确 mask/start_pos contract，并对不支持的输入 fail-fast；
+2. 把 autocast cast-backward 和 mixed-dtype gradient 语义下沉到通用 autograd 基础设施，
+   删除 adapter 的 BF16 特殊 upcast；
+3. 为 deterministic backward 增加接口、workspace 和 AOT 实例；
+4. 优化 GQA backward，避免长期保留 Q-head shape 的 dK/dV 临时 buffer；
+5. 按模型需求扩展 head dimension 和 GPU architecture，并保持 CMake/runtime dispatch 同步；
+6. 若支持 dropout，引入框架 generator/Philox contract，而不是继续使用零 rng-state；
+7. 评估 varlen、padding mask、KV cache、local attention、ALiBi 和 softcap 所需的新参数与
+   kernel 实例；
+8. 升级 FlashAttention 时优先审计 compat headers，并重新确认最终链接不引入 Torch。
diff --git a/example/gpt2/checkpoint_loader.cc b/example/gpt2/checkpoint_loader.cc
index 95e54730..3940710e 100644
--- a/example/gpt2/checkpoint_loader.cc
+++ b/example/gpt2/checkpoint_loader.cc
@@ -57,7 +57,7 @@ std::tuple<int32_t, infini_train::DataType> DetermineAndCheckVersion(const std::
 
 namespace gpt2 {
 
-std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath) {
+std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath, bool use_flash_attention) {
     if (!std::filesystem::exists(filepath)) {
         LOG(FATAL) << "File not found: " << filepath;
     }
@@ -87,7 +87,9 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
     gpt2_config.original_vocab_size = vocab_size;
     gpt2_config.n_layer = n_layer;
     gpt2_config.n_head = n_head;
+    gpt2_config.n_kv_head = n_head;
     gpt2_config.n_embd = n_embd;
+    gpt2_config.flash = use_flash_attention;
     gpt2::SanitizeGPT2Config(gpt2_config);
     auto local_gpt2 = std::make_shared<nn::TransformerModel>(gpt2_config);
 
diff --git a/example/gpt2/checkpoint_loader.h b/example/gpt2/checkpoint_loader.h
index e80c356e..e06ceece 100644
--- a/example/gpt2/checkpoint_loader.h
+++ b/example/gpt2/checkpoint_loader.h
@@ -8,5 +8,6 @@ class TransformerModel;
 } // namespace infini_train::nn
 
 namespace gpt2 {
-std::shared_ptr<infini_train::nn::TransformerModel> LoadFromLLMC(const std::string &filepath);
+std::shared_ptr<infini_train::nn::TransformerModel> LoadFromLLMC(const std::string &filepath,
+                                                                 bool use_flash_attention = false);
 } // namespace gpt2
diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc
index 60c0c908..d4b3af1e 100644
--- a/example/gpt2/main.cc
+++ b/example/gpt2/main.cc
@@ -76,6 +76,7 @@ DEFINE_uint32(sample_every, 0, "how often to sample from the model?");
 DEFINE_bool(overfit_single_batch, true, "overfit just one batch of data");
 // memory management
 DEFINE_string(device, "cuda", "device type (cpu/cuda), useless if using parallel training mode");
+DEFINE_string(attention_backend, "unfused", "attention backend: unfused|flash");
 // parallel
 DEFINE_int32(
     nthread_per_process, 1,
@@ -118,6 +119,8 @@ constexpr char kDtypeFP32[] = "float32";
 constexpr char kDtypeBF16[] = "bfloat16";
 const std::unordered_set<std::string> kSupportedLRDecayStyles
     = {"none", "constant", "linear", "cosine", "inverse-square-root"};
+constexpr char kAttentionBackendUnfused[] = "unfused";
+constexpr char kAttentionBackendFlash[] = "flash";
 
 //
 const std::unordered_map<std::string, nn::TransformerConfig> kModelToConfigs = {
@@ -132,6 +135,9 @@ const std::unordered_map<std::string, nn::TransformerConfig> kModelToConfigs = {
 DEFINE_validator(model, [](const char *, const std::string &value) { return kSupportedModels.contains(value); });
 DEFINE_validator(device,
                  [](const char *, const std::string &value) { return value == kDeviceCPU || value == kDeviceCUDA; });
+DEFINE_validator(attention_backend, [](const char *, const std::string &value) {
+    return value == kAttentionBackendUnfused || value == kAttentionBackendFlash;
+});
 DEFINE_validator(zero_stage, [](const char *, int32_t value) { return value >= 0 && value <= 3; });
 DEFINE_validator(lr_decay_style,
                  [](const char *, const std::string &value) { return kSupportedLRDecayStyles.contains(value); });
@@ -209,6 +215,11 @@ void Train(const nn::parallel::Rank &rank) {
         device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0);
     }
 
+    const bool use_flash_attention = FLAGS_attention_backend == kAttentionBackendFlash;
+    if (use_flash_attention && device.type() != Device::DeviceType::kCUDA) {
+        LOG(FATAL) << "--attention_backend=flash requires --device=cuda";
+    }
+
     // calculate gradient accumulation from the desired total batch size and the current run configuration
     const auto tokens_per_fwdbwd = FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size;
     CHECK_EQ(FLAGS_total_batch_size % tokens_per_fwdbwd, 0);
@@ -224,9 +235,15 @@ void Train(const nn::parallel::Rank &rank) {
     std::shared_ptr<nn::Module> model = nullptr;
 
     if (!FLAGS_llmc_filepath.empty()) {
-        model = gpt2::LoadFromLLMC(FLAGS_llmc_filepath);
+        model = gpt2::LoadFromLLMC(FLAGS_llmc_filepath, use_flash_attention);
     } else if (kModelToConfigs.count(FLAGS_model)) {
         model_config = kModelToConfigs.at(FLAGS_model);
+        model_config.n_kv_head = model_config.n_head;
+        model_config.flash = use_flash_attention;
+        gpt2::SanitizeGPT2Config(model_config);
+        model = std::make_shared<nn::TransformerModel>(model_config);
+    } else {
+        model_config.flash = use_flash_attention;
         gpt2::SanitizeGPT2Config(model_config);
         model = std::make_shared<nn::TransformerModel>(model_config);
     }
@@ -268,6 +285,10 @@ void Train(const nn::parallel::Rank &rank) {
     } else {
         LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << FLAGS_dtype << " not supported.";
     }
+    if (use_flash_attention && dtype != DataType::kBFLOAT16) {
+        LOG(FATAL) << "--attention_backend=flash currently requires --dtype=bfloat16 because FlashAttention 2 "
+                      "supports fp16/bf16 kernels only";
+    }
 
     auto num_micro_batches = FLAGS_total_batch_size / (FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size);
 
diff --git a/example/llama3/checkpoint_loader.cc b/example/llama3/checkpoint_loader.cc
index f3590af6..3c666ff1 100644
--- a/example/llama3/checkpoint_loader.cc
+++ b/example/llama3/checkpoint_loader.cc
@@ -40,7 +40,7 @@ constexpr int32_t kLLaMA3FP32Version = 3;
 
 namespace llama3 {
 
-std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath) {
+std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath, bool use_flash_attention) {
     if (!std::filesystem::exists(filepath)) {
         LOG(FATAL) << "File not found: " << filepath;
     }
@@ -81,6 +81,7 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
     llama3_config.use_scaled_rope = static_cast<bool>(use_scaled_rope);
     llama3_config.norm_eps = norm_eps;
     llama3_config.max_gen_batch_size = max_gen_bs;
+    llama3_config.flash = use_flash_attention;
     llama3::SanitizeLLaMA3Config(llama3_config);
     auto llama3 = std::make_shared<nn::TransformerModel>(llama3_config);
 
diff --git a/example/llama3/checkpoint_loader.h b/example/llama3/checkpoint_loader.h
index d4aea3d0..c9eced99 100644
--- a/example/llama3/checkpoint_loader.h
+++ b/example/llama3/checkpoint_loader.h
@@ -8,5 +8,6 @@ class TransformerModel;
 } // namespace infini_train::nn
 
 namespace llama3 {
-std::shared_ptr<infini_train::nn::TransformerModel> LoadFromLLMC(const std::string &filepath);
+std::shared_ptr<infini_train::nn::TransformerModel> LoadFromLLMC(const std::string &filepath,
+                                                                 bool use_flash_attention = false);
 } // namespace llama3
diff --git a/example/llama3/main.cc b/example/llama3/main.cc
index 302e0808..4a33e912 100644
--- a/example/llama3/main.cc
+++ b/example/llama3/main.cc
@@ -75,6 +75,7 @@ DEFINE_uint32(sample_every, 0, "how often to sample from the model?");
 DEFINE_bool(overfit_single_batch, true, "overfit just one batch of data");
 // memory management
 DEFINE_string(device, "cuda", "device type (cpu/cuda), useless if using parallel training mode");
+DEFINE_string(attention_backend, "unfused", "attention backend: unfused|flash");
 // parallel
 DEFINE_int32(
     nthread_per_process, 1,
@@ -114,11 +115,16 @@ constexpr char kDtypeFP32[] = "float32";
 constexpr char kDtypeBF16[] = "bfloat16";
 const std::unordered_set<std::string> kSupportedLRDecayStyles
     = {"none", "constant", "linear", "cosine", "inverse-square-root"};
+constexpr char kAttentionBackendUnfused[] = "unfused";
+constexpr char kAttentionBackendFlash[] = "flash";
 } // namespace
 
 DEFINE_validator(model, [](const char *, const std::string &value) { return kSupportedModels.contains(value); });
 DEFINE_validator(device,
                  [](const char *, const std::string &value) { return value == kDeviceCPU || value == kDeviceCUDA; });
+DEFINE_validator(attention_backend, [](const char *, const std::string &value) {
+    return value == kAttentionBackendUnfused || value == kAttentionBackendFlash;
+});
 DEFINE_validator(zero_stage, [](const char *, int32_t value) { return value >= 0 && value <= 3; });
 DEFINE_validator(lr_decay_style,
                  [](const char *, const std::string &value) { return kSupportedLRDecayStyles.contains(value); });
@@ -195,6 +201,11 @@ void Train(const nn::parallel::Rank &rank) {
         device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0);
     }
 
+    const bool use_flash_attention = FLAGS_attention_backend == kAttentionBackendFlash;
+    if (use_flash_attention && device.type() != Device::DeviceType::kCUDA) {
+        LOG(FATAL) << "--attention_backend=flash requires --device=cuda";
+    }
+
     // calculate gradient accumulation from the desired total batch size and the current run configuration
     const auto tokens_per_fwdbwd = FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size;
     CHECK_EQ(FLAGS_total_batch_size % tokens_per_fwdbwd, 0);
@@ -210,8 +221,9 @@ void Train(const nn::parallel::Rank &rank) {
     nn::TransformerConfig model_config = llama3::LLaMA3Config();
     std::shared_ptr<nn::Module> model = nullptr;
     if (!FLAGS_llmc_filepath.empty()) {
-        model = llama3::LoadFromLLMC(FLAGS_llmc_filepath);
+        model = llama3::LoadFromLLMC(FLAGS_llmc_filepath, use_flash_attention);
     } else {
+        model_config.flash = use_flash_attention;
         llama3::SanitizeLLaMA3Config(model_config);
         model = std::make_shared<nn::TransformerModel>(model_config);
     }
@@ -249,6 +261,10 @@ void Train(const nn::parallel::Rank &rank) {
     } else {
         LOG(FATAL) << "Rank " << rank.GlobalRank() << ": Datatype " << FLAGS_dtype << " not supported.";
     }
+    if (use_flash_attention && dtype != DataType::kBFLOAT16) {
+        LOG(FATAL) << "--attention_backend=flash currently requires --dtype=bfloat16 because FlashAttention 2 "
+                      "supports fp16/bf16 kernels only";
+    }
 
     auto num_micro_batches = FLAGS_total_batch_size / (FLAGS_batch_size * FLAGS_sequence_length * ddp_world_size);
 
diff --git a/infini_train/include/autograd/scaled_dot_product_attention.h b/infini_train/include/autograd/scaled_dot_product_attention.h
new file mode 100644
index 00000000..ca82865b
--- /dev/null
+++ b/infini_train/include/autograd/scaled_dot_product_attention.h
@@ -0,0 +1,30 @@
+#pragma once
+
+#include <memory>
+#include <vector>
+
+#include "infini_train/include/autograd/function.h"
+
+namespace infini_train {
+class Tensor;
+}
+
+namespace infini_train::autograd {
+
+class ScaledDotProductAttention : public Function {
+public:
+    static constexpr char kType[] = "ScaledDotProductAttentionFunction";
+
+    explicit ScaledDotProductAttention(float scale) : Function(kType), scale_(scale) {}
+
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
+    void SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                      const std::vector<std::shared_ptr<Tensor>> &output_tensors) override;
+    std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
+
+private:
+    float scale_ = 1.0f;
+    std::shared_ptr<void> flash_ctx_;
+};
+
+} // namespace infini_train::autograd
diff --git a/infini_train/include/common/common.h b/infini_train/include/common/common.h
index 80cba728..86b10d81 100644
--- a/infini_train/include/common/common.h
+++ b/infini_train/include/common/common.h
@@ -21,6 +21,7 @@
 #define CAT_(a, b) a##b
 
 #define CEIL_DIV(x, y) (((x) + (y)-1) / (y))
+#define ROUND_UP(x, y) (CEIL_DIV((x), (y)) * (y))
 #define LOG_LOC(LEVEL, MSG) LOG(LEVEL) << MSG << " at " << __FILE__ << ":" << __LINE__
 
 inline std::vector<int64_t> ComputeStrides(const std::vector<int64_t> &dims) {
diff --git a/infini_train/include/dispatcher.h b/infini_train/include/dispatcher.h
index 2c390ec8..3b16ee88 100644
--- a/infini_train/include/dispatcher.h
+++ b/infini_train/include/dispatcher.h
@@ -55,9 +55,10 @@ class Dispatcher {
         return instance;
     }
 
+    bool HasKernel(const KeyT &key) const { return key_to_kernel_map_.contains(key); }
+
     const KernelFunction &GetKernel(KeyT key) const {
-        CHECK(key_to_kernel_map_.contains(key))
-            << "Kernel not found: " << key.second << " on device: " << static_cast<int>(key.first);
+        CHECK(HasKernel(key)) << "Kernel not found: " << key.second << " on device: " << static_cast<int>(key.first);
 #ifdef PROFILE_MODE
         SetProfileContext(key.second, key.first);
 #endif
diff --git a/infini_train/include/nn/functional.h b/infini_train/include/nn/functional.h
index e4354fd1..bb1f33f9 100644
--- a/infini_train/include/nn/functional.h
+++ b/infini_train/include/nn/functional.h
@@ -149,6 +149,12 @@ std::shared_ptr<Tensor> Sigmoid(const std::shared_ptr<Tensor> &input);
 //   A tensor with softmax applied along the specified dimension.
 std::shared_ptr<Tensor> Softmax(const std::shared_ptr<Tensor> &input, int64_t dim = -1);
 
+// Computes causal self-attention for q/k/v tensors in (B, H, T, D) layout.
+//
+// The current FlashAttention backend applies a causal mask (j <= i).
+std::shared_ptr<Tensor> ScaledDotProductAttention(const std::shared_ptr<Tensor> &q, const std::shared_ptr<Tensor> &k,
+                                                  const std::shared_ptr<Tensor> &v, float scale);
+
 // Returns a slice of the input tensor defined by start, end, and step per dimension.
 //
 // Args:
diff --git a/infini_train/src/autograd/scaled_dot_product_attention.cc b/infini_train/src/autograd/scaled_dot_product_attention.cc
new file mode 100644
index 00000000..a6c49d1f
--- /dev/null
+++ b/infini_train/src/autograd/scaled_dot_product_attention.cc
@@ -0,0 +1,97 @@
+#include "infini_train/include/autograd/scaled_dot_product_attention.h"
+
+#include <memory>
+#include <vector>
+
+#include "glog/logging.h"
+
+#include "infini_train/include/autocast.h"
+#include "infini_train/include/datatype.h"
+#include "infini_train/include/dispatcher.h"
+#include "infini_train/include/tensor.h"
+
+namespace infini_train::autograd {
+namespace {
+
+constexpr char kForwardKernel[] = "ScaledDotProductAttentionForward";
+constexpr char kBackwardKernel[] = "ScaledDotProductAttentionBackward";
+
+DataType SelectFlashAttentionDtype(const std::shared_ptr<Tensor> &q) {
+    if (q->Dtype() == DataType::kBFLOAT16 || q->Dtype() == DataType::kFLOAT16) {
+        return q->Dtype();
+    }
+    if (tls_autocast_context.enabled
+        && (tls_autocast_context.autocast_dtype == DataType::kBFLOAT16
+            || tls_autocast_context.autocast_dtype == DataType::kFLOAT16)) {
+        return tls_autocast_context.autocast_dtype;
+    }
+    LOG(FATAL) << "FlashAttention 2 supports fp16/bf16 only. Use --dtype=bfloat16 for --attention_backend=flash.";
+    return DataType::kBFLOAT16;
+}
+
+void CheckQKVHeads(const std::shared_ptr<Tensor> &q, const std::shared_ptr<Tensor> &k,
+                   const std::shared_ptr<Tensor> &v) {
+    CHECK_EQ(q->Dims().size(), 4) << "Q must use (B, H, T, D) layout";
+    CHECK_EQ(k->Dims().size(), 4) << "K must use (B, H, T, D) layout";
+    CHECK_EQ(v->Dims().size(), 4) << "V must use (B, H, T, D) layout";
+    CHECK(k->Dims() == v->Dims()) << "K and V must have the same shape";
+
+    const auto query_heads = q->Dims()[1];
+    const auto kv_heads = k->Dims()[1];
+    CHECK_GT(query_heads, 0) << "Q must have at least one head";
+    CHECK_GT(kv_heads, 0) << "K/V must have at least one head";
+    CHECK_EQ(query_heads % kv_heads, 0) << "Q heads must be divisible by KV heads for GQA/MQA";
+}
+
+} // namespace
+
+std::vector<std::shared_ptr<Tensor>>
+ScaledDotProductAttention::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
+    CHECK_EQ(input_tensors.size(), 3);
+    const auto &q = input_tensors[0];
+    const auto &k = input_tensors[1];
+    const auto &v = input_tensors[2];
+
+    CheckQKVHeads(q, k, v);
+    const auto device = q->GetDevice();
+    CHECK(device.IsCUDA()) << "FlashAttention backend requires CUDA tensors";
+
+    const Dispatcher::KeyT forward_key{device.type(), kForwardKernel};
+    CHECK(Dispatcher::Instance().HasKernel(forward_key))
+        << "FlashAttention backend is not available in this build; configure with -DUSE_FLASH_ATTENTION=ON";
+
+    const auto flash_dtype = SelectFlashAttentionDtype(q);
+    flash_ctx_.reset();
+    return {
+        Dispatcher::Instance().Call<std::shared_ptr<Tensor>>(forward_key, q, k, v, scale_, flash_dtype, &flash_ctx_)};
+}
+
+void ScaledDotProductAttention::SetupContext(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
+                                             const std::vector<std::shared_ptr<Tensor>> &output_tensors) {
+    CHECK_EQ(input_tensors.size(), 3);
+    CHECK_EQ(output_tensors.size(), 1);
+    ctx_.SaveForBackward({input_tensors[0], input_tensors[1], input_tensors[2], output_tensors[0]});
+}
+
+std::vector<std::shared_ptr<Tensor>>
+ScaledDotProductAttention::Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) {
+    CHECK_EQ(grad_outputs.size(), 1);
+    const auto &grad_output = grad_outputs[0];
+
+    auto saved_tensors = ctx_.GetSavedTensors();
+    CHECK_EQ(saved_tensors.size(), 4);
+    const auto &q = saved_tensors[0];
+    const auto &k = saved_tensors[1];
+    const auto &v = saved_tensors[2];
+    const auto &out = saved_tensors[3];
+
+    const Dispatcher::KeyT backward_key{q->GetDevice().type(), kBackwardKernel};
+    CHECK(Dispatcher::Instance().HasKernel(backward_key))
+        << "FlashAttention backward kernel is not available in this build; configure with "
+           "-DUSE_FLASH_ATTENTION=ON";
+
+    return Dispatcher::Instance().Call<std::vector<std::shared_ptr<Tensor>>>(backward_key, grad_output, q, k, v, out,
+                                                                             scale_, flash_ctx_);
+}
+
+} // namespace infini_train::autograd
diff --git a/infini_train/src/kernels/cuda/flash_attention.cu b/infini_train/src/kernels/cuda/flash_attention.cu
new file mode 100644
index 00000000..92a123f2
--- /dev/null
+++ b/infini_train/src/kernels/cuda/flash_attention.cu
@@ -0,0 +1,357 @@
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "flash.h"
+#include "glog/logging.h"
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cutlass/numeric_types.h>
+
+#include "infini_train/include/common/common.h"
+#include "infini_train/include/core/runtime/device_guard.h"
+#include "infini_train/include/datatype.h"
+#include "infini_train/include/dispatcher.h"
+#include "infini_train/include/tensor.h"
+#include "infini_train/src/core/runtime/cuda/cuda_runtime_common.h"
+
+namespace infini_train::kernels::cuda {
+namespace {
+
+constexpr int kSequenceAlignment = 128;
+
+struct FlashAttentionContext {
+    std::shared_ptr<Tensor> q;
+    std::shared_ptr<Tensor> k;
+    std::shared_ptr<Tensor> v;
+    std::shared_ptr<Tensor> out;
+    std::shared_ptr<Tensor> softmax_lse;
+    std::shared_ptr<Tensor> rng_state;
+    DataType q_original_dtype = DataType::kFLOAT32;
+    DataType k_original_dtype = DataType::kFLOAT32;
+    DataType v_original_dtype = DataType::kFLOAT32;
+    DataType flash_dtype = DataType::kBFLOAT16;
+};
+
+std::shared_ptr<Tensor> CastIfNeeded(const std::shared_ptr<Tensor> &tensor, DataType dtype) {
+    return tensor->Dtype() == dtype ? tensor : std::make_shared<Tensor>(tensor->To(dtype));
+}
+
+std::shared_ptr<Tensor> CastGrad(const std::shared_ptr<Tensor> &grad, DataType original_dtype, DataType flash_dtype) {
+    const auto grad_dtype = flash_dtype == DataType::kBFLOAT16 ? DataType::kFLOAT32 : original_dtype;
+    return CastIfNeeded(grad, grad_dtype);
+}
+
+cudaStream_t GetCudaStream(Device device) {
+    auto *stream = dynamic_cast<infini_train::core::cuda::CudaStream *>(
+        infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device));
+    CHECK(stream != nullptr);
+    return stream->cuda_stream();
+}
+
+void CheckQKV(const std::shared_ptr<Tensor> &q, const std::shared_ptr<Tensor> &k, const std::shared_ptr<Tensor> &v,
+              DataType flash_dtype) {
+    CHECK(q->GetDevice() == k->GetDevice());
+    CHECK(q->GetDevice() == v->GetDevice());
+    CHECK(q->GetDevice().IsCUDA()) << "FlashAttention backend requires CUDA tensors";
+    CHECK(k->Dims() == v->Dims());
+    CHECK_EQ(q->Dims().size(), 4);
+    CHECK_EQ(k->Dims().size(), 4);
+    CHECK(flash_dtype == DataType::kBFLOAT16 || flash_dtype == DataType::kFLOAT16)
+        << "FlashAttention supports fp16/bf16 only";
+
+    const auto &q_dims = q->Dims();
+    const auto &kv_dims = k->Dims();
+    CHECK_GT(q_dims[0], 0);
+    CHECK_GT(q_dims[1], 0);
+    CHECK_GT(q_dims[2], 0);
+    CHECK_EQ(q_dims[0], kv_dims[0]);
+    CHECK_EQ(q_dims[2], kv_dims[2]);
+    CHECK_EQ(q_dims[3], kv_dims[3]);
+    // FIXME(zbl): Extend supported head dimensions together with the AOT kernel instances in CMakeLists.txt;
+    //             the runtime checks/dispatch and FLASH_ATTN_CUDA_SOURCES must remain in sync.
+    CHECK(q_dims[3] == 64 || q_dims[3] == 128) << "Native FlashAttention currently supports head_dim 64/128 only";
+    CHECK_GT(kv_dims[1], 0);
+    CHECK_EQ(q_dims[1] % kv_dims[1], 0) << "Q heads must be divisible by KV heads for GQA/MQA";
+}
+
+void SetForwardParams(flash::Flash_fwd_params *params, const Tensor &q, const Tensor &k, const Tensor &v, Tensor *out,
+                      Tensor *softmax_lse, float scale) {
+    CHECK(params != nullptr);
+    const auto &q_dims = q.Dims();
+    const auto &kv_dims = k.Dims();
+    const int64_t batch = q_dims[0];
+    const int64_t q_heads = q_dims[1];
+    const int64_t kv_heads = kv_dims[1];
+    const int64_t seqlen = q_dims[2];
+    const int64_t head_dim = q_dims[3];
+    const int64_t q_batch_stride = q_heads * seqlen * head_dim;
+    const int64_t kv_batch_stride = kv_heads * seqlen * head_dim;
+    const int64_t head_stride = seqlen * head_dim;
+
+    *params = {};
+    params->q_ptr = const_cast<void *>(q.DataPtr());
+    params->k_ptr = const_cast<void *>(k.DataPtr());
+    params->v_ptr = const_cast<void *>(v.DataPtr());
+    params->o_ptr = out->DataPtr();
+    params->softmax_lse_ptr = softmax_lse->DataPtr();
+
+    params->q_batch_stride = q_batch_stride;
+    params->k_batch_stride = params->v_batch_stride = kv_batch_stride;
+    params->o_batch_stride = q_batch_stride;
+    params->q_row_stride = params->k_row_stride = params->v_row_stride = head_dim;
+    params->o_row_stride = head_dim;
+    params->q_head_stride = params->k_head_stride = params->v_head_stride = head_stride;
+    params->o_head_stride = head_stride;
+
+    params->b = batch;
+    params->h = q_heads;
+    params->h_k = kv_heads;
+    params->h_h_k_ratio = q_heads / kv_heads;
+    params->seqlen_q = seqlen;
+    params->seqlen_k = seqlen;
+    params->seqlen_q_rounded = ROUND_UP(seqlen, kSequenceAlignment);
+    params->seqlen_k_rounded = ROUND_UP(seqlen, kSequenceAlignment);
+    params->d = head_dim;
+    params->d_rounded = head_dim;
+    params->total_q = batch * seqlen;
+
+    params->scale_softmax = scale;
+    params->scale_softmax_log2 = scale * 1.4426950408889634F;
+    params->p_dropout = 1.0F;
+    params->p_dropout_in_uint8_t = 255;
+    params->rp_dropout = 1.0F;
+    params->scale_softmax_rp_dropout = scale;
+    params->is_bf16 = q.Dtype() == DataType::kBFLOAT16;
+    // TODO(zbl): Plumb mask type and local-window configuration through ScaledDotProductAttention.
+    params->window_size_left = -1;
+    params->window_size_right = 0;
+    params->is_causal = true;
+    params->is_seqlens_k_cumulative = true;
+    params->num_splits = 1;
+}
+
+template <typename T> void RunForward(flash::Flash_fwd_params &params, cudaStream_t stream) {
+    switch (params.d) {
+    case 64:
+        flash::run_mha_fwd_<T, 64, true>(params, stream);
+        return;
+    case 128:
+        flash::run_mha_fwd_<T, 128, true>(params, stream);
+        return;
+    default:
+        LOG(FATAL) << "Unsupported FlashAttention head_dim=" << params.d;
+    }
+}
+
+template <typename T> void RunBackward(flash::Flash_bwd_params &params, cudaStream_t stream) {
+    switch (params.d) {
+    case 64:
+        flash::run_mha_bwd_<T, 64, true>(params, stream);
+        return;
+    case 128:
+        flash::run_mha_bwd_<T, 128, true>(params, stream);
+        return;
+    default:
+        LOG(FATAL) << "Unsupported FlashAttention head_dim=" << params.d;
+    }
+}
+
+void DispatchForward(flash::Flash_fwd_params &params, cudaStream_t stream) {
+    if (params.is_bf16) {
+        RunForward<cutlass::bfloat16_t>(params, stream);
+    } else {
+        RunForward<cutlass::half_t>(params, stream);
+    }
+}
+
+void DispatchBackward(flash::Flash_bwd_params &params, cudaStream_t stream) {
+    if (params.is_bf16) {
+        RunBackward<cutlass::bfloat16_t>(params, stream);
+    } else {
+        RunBackward<cutlass::half_t>(params, stream);
+    }
+}
+
+template <typename T> struct GqaGradConvert;
+
+template <> struct GqaGradConvert<__half> {
+    __device__ static float ToFloat(__half value) { return __half2float(value); }
+    __device__ static __half FromFloat(float value) { return __float2half_rn(value); }
+};
+
+template <> struct GqaGradConvert<__nv_bfloat16> {
+    __device__ static float ToFloat(__nv_bfloat16 value) { return __bfloat162float(value); }
+    __device__ static __nv_bfloat16 FromFloat(float value) { return __float2bfloat16_rn(value); }
+};
+
+template <typename T>
+__global__ void ReduceGqaGradKernel(T *output, const T *expanded, int64_t num_elements, int64_t kv_heads,
+                                    int64_t query_heads, int64_t seqlen, int64_t head_dim) {
+    const int64_t output_idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (output_idx >= num_elements) {
+        return;
+    }
+
+    int64_t index = output_idx;
+    const int64_t dim_idx = index % head_dim;
+    index /= head_dim;
+    const int64_t sequence_idx = index % seqlen;
+    index /= seqlen;
+    const int64_t kv_head_idx = index % kv_heads;
+    const int64_t batch_idx = index / kv_heads;
+    const int64_t repeats = query_heads / kv_heads;
+
+    float sum = 0.0F;
+    for (int64_t group = 0; group < repeats; ++group) {
+        const int64_t query_head_idx = kv_head_idx * repeats + group;
+        const int64_t expanded_idx
+            = ((batch_idx * query_heads + query_head_idx) * seqlen + sequence_idx) * head_dim + dim_idx;
+        sum += GqaGradConvert<T>::ToFloat(expanded[expanded_idx]);
+    }
+    output[output_idx] = GqaGradConvert<T>::FromFloat(sum);
+}
+
+void ReduceGqaGrad(Tensor *output, const Tensor &expanded, cudaStream_t stream) {
+    CHECK(output != nullptr);
+    CHECK(output->Dtype() == expanded.Dtype());
+    const auto &output_dims = output->Dims();
+    const auto &expanded_dims = expanded.Dims();
+    CHECK_EQ(output_dims.size(), 4);
+    CHECK_EQ(expanded_dims.size(), 4);
+    CHECK_EQ(output_dims[0], expanded_dims[0]);
+    CHECK_EQ(output_dims[2], expanded_dims[2]);
+    CHECK_EQ(output_dims[3], expanded_dims[3]);
+    CHECK_EQ(expanded_dims[1] % output_dims[1], 0);
+
+    constexpr int threads = 256;
+    const int64_t num_elements = output->NumElements();
+    const int blocks = static_cast<int>((num_elements + threads - 1) / threads);
+    if (output->Dtype() == DataType::kBFLOAT16) {
+        ReduceGqaGradKernel<<<blocks, threads, 0, stream>>>(
+            static_cast<__nv_bfloat16 *>(output->DataPtr()), static_cast<const __nv_bfloat16 *>(expanded.DataPtr()),
+            num_elements, output_dims[1], expanded_dims[1], output_dims[2], output_dims[3]);
+    } else {
+        ReduceGqaGradKernel<<<blocks, threads, 0, stream>>>(
+            static_cast<__half *>(output->DataPtr()), static_cast<const __half *>(expanded.DataPtr()), num_elements,
+            output_dims[1], expanded_dims[1], output_dims[2], output_dims[3]);
+    }
+    CHECK_EQ(cudaGetLastError(), cudaSuccess);
+}
+
+} // namespace
+
+std::shared_ptr<Tensor> ScaledDotProductAttentionForward(const std::shared_ptr<Tensor> &q,
+                                                         const std::shared_ptr<Tensor> &k,
+                                                         const std::shared_ptr<Tensor> &v, float scale,
+                                                         DataType flash_dtype, std::shared_ptr<void> *opaque_ctx) {
+    CHECK(opaque_ctx != nullptr);
+    CheckQKV(q, k, v, flash_dtype);
+
+    const auto device = q->GetDevice();
+    infini_train::core::DeviceGuard guard(device);
+    auto ctx = std::make_shared<FlashAttentionContext>();
+    ctx->q_original_dtype = q->Dtype();
+    ctx->k_original_dtype = k->Dtype();
+    ctx->v_original_dtype = v->Dtype();
+    ctx->flash_dtype = flash_dtype;
+    ctx->q = CastIfNeeded(q, flash_dtype);
+    ctx->k = CastIfNeeded(k, flash_dtype);
+    ctx->v = CastIfNeeded(v, flash_dtype);
+    auto output = std::make_shared<Tensor>(q->Dims(), flash_dtype, device);
+    // Keep the forward buffer for backward without retaining the autograd output and
+    // forming Function -> flash_ctx -> output -> Function.
+    ctx->out = output->Detach();
+
+    const auto &dims = q->Dims();
+    ctx->softmax_lse
+        = std::make_shared<Tensor>(std::vector<int64_t>{dims[0], dims[1], dims[2]}, DataType::kFLOAT32, device);
+    ctx->rng_state = std::make_shared<Tensor>(std::vector<int64_t>{2}, DataType::kUINT64, device);
+    const auto stream = GetCudaStream(device);
+    CHECK_EQ(cudaMemsetAsync(ctx->rng_state->DataPtr(), 0, ctx->rng_state->SizeInBytes(), stream), cudaSuccess);
+
+    flash::Flash_fwd_params params{};
+    SetForwardParams(&params, *ctx->q, *ctx->k, *ctx->v, ctx->out.get(), ctx->softmax_lse.get(), scale);
+    params.rng_state = static_cast<uint64_t *>(ctx->rng_state->DataPtr());
+    DispatchForward(params, stream);
+
+    *opaque_ctx = ctx;
+    return output;
+}
+
+std::vector<std::shared_ptr<Tensor>>
+ScaledDotProductAttentionBackward(const std::shared_ptr<Tensor> &grad_output, const std::shared_ptr<Tensor> &,
+                                  const std::shared_ptr<Tensor> &, const std::shared_ptr<Tensor> &,
+                                  const std::shared_ptr<Tensor> &, float scale, std::shared_ptr<void> opaque_ctx) {
+    auto ctx = std::static_pointer_cast<FlashAttentionContext>(opaque_ctx);
+    CHECK(ctx != nullptr) << "Missing FlashAttention forward context";
+
+    const auto device = ctx->q->GetDevice();
+    infini_train::core::DeviceGuard guard(device);
+    auto dout = CastIfNeeded(grad_output, ctx->flash_dtype);
+    CHECK(dout->Dims() == ctx->out->Dims());
+
+    auto dq = std::make_shared<Tensor>(ctx->q->Dims(), ctx->flash_dtype, device);
+    auto dk = std::make_shared<Tensor>(ctx->k->Dims(), ctx->flash_dtype, device);
+    auto dv = std::make_shared<Tensor>(ctx->v->Dims(), ctx->flash_dtype, device);
+    const auto &dims = ctx->q->Dims();
+    const int64_t batch = dims[0];
+    const int64_t heads = dims[1];
+    const int64_t seqlen = dims[2];
+    const int64_t head_dim = dims[3];
+    const int64_t seqlen_rounded = ROUND_UP(seqlen, kSequenceAlignment);
+
+    auto softmax_d
+        = std::make_shared<Tensor>(std::vector<int64_t>{batch, heads, seqlen_rounded}, DataType::kFLOAT32, device);
+    auto dq_accum = std::make_shared<Tensor>(std::vector<int64_t>{batch, seqlen_rounded, heads, head_dim},
+                                             DataType::kFLOAT32, device);
+    const bool use_gqa = ctx->q->Dims()[1] != ctx->k->Dims()[1];
+    auto dk_kernel = use_gqa ? std::make_shared<Tensor>(ctx->q->Dims(), ctx->flash_dtype, device) : dk;
+    auto dv_kernel = use_gqa ? std::make_shared<Tensor>(ctx->q->Dims(), ctx->flash_dtype, device) : dv;
+
+    flash::Flash_bwd_params params{};
+    SetForwardParams(&params, *ctx->q, *ctx->k, *ctx->v, ctx->out.get(), ctx->softmax_lse.get(), scale);
+    const int64_t batch_stride = heads * seqlen * head_dim;
+    const int64_t head_stride = seqlen * head_dim;
+    params.do_ptr = dout->DataPtr();
+    params.dq_ptr = dq->DataPtr();
+    params.dk_ptr = dk_kernel->DataPtr();
+    params.dv_ptr = dv_kernel->DataPtr();
+    params.do_batch_stride = params.dq_batch_stride = params.dk_batch_stride = params.dv_batch_stride = batch_stride;
+    params.do_row_stride = params.dq_row_stride = params.dk_row_stride = params.dv_row_stride = head_dim;
+    params.do_head_stride = params.dq_head_stride = params.dk_head_stride = params.dv_head_stride = head_stride;
+    params.dq_accum_ptr = dq_accum->DataPtr();
+    params.dsoftmax_sum = softmax_d->DataPtr();
+    params.rng_state = static_cast<uint64_t *>(ctx->rng_state->DataPtr());
+    params.deterministic = false;
+    params.dq_accum_split_stride = 0;
+
+    const auto stream = GetCudaStream(device);
+    DispatchBackward(params, stream);
+    if (use_gqa) {
+        ReduceGqaGrad(dk.get(), *dk_kernel, stream);
+        ReduceGqaGrad(dv.get(), *dv_kernel, stream);
+    }
+
+    // FIXME(zbl): Forward autocast currently uses raw Tensor::To conversions and wires the Function directly
+    //             to the original autograd graph. Without an autograd cast-backward edge or generic mixed-dtype
+    //             grad normalization, native BF16 FlashAttention grads can be accumulated together with the FP32
+    //             grads propagated by Matmul/Linear backward paths. FlashAttention backward kernel performs in
+    //             BF16, so we manually promote them at the end of kernel to keep backward consistent with the
+    //             current forward autocast behavior. The proper fix belongs in autograd: once autocast and grad
+    //             accumulation preserve type semantics centrally, return grads in its native input dtype.
+    return {CastGrad(dq, ctx->q_original_dtype, ctx->flash_dtype),
+            CastGrad(dk, ctx->k_original_dtype, ctx->flash_dtype),
+            CastGrad(dv, ctx->v_original_dtype, ctx->flash_dtype)};
+}
+
+} // namespace infini_train::kernels::cuda
+
+#define REGISTER_CUDA_FLASH_ATTENTION_KERNEL(kernel_name)                                                              \
+    REGISTER_KERNEL(infini_train::Device::DeviceType::kCUDA, kernel_name, infini_train::kernels::cuda::kernel_name)
+
+REGISTER_CUDA_FLASH_ATTENTION_KERNEL(ScaledDotProductAttentionForward)
+REGISTER_CUDA_FLASH_ATTENTION_KERNEL(ScaledDotProductAttentionBackward)
+
+#undef REGISTER_CUDA_FLASH_ATTENTION_KERNEL
diff --git a/infini_train/src/kernels/cuda/flash_attention_compat/ATen/cuda/CUDAGeneratorImpl.h b/infini_train/src/kernels/cuda/flash_attention_compat/ATen/cuda/CUDAGeneratorImpl.h
new file mode 100644
index 00000000..5f77f9fd
--- /dev/null
+++ b/infini_train/src/kernels/cuda/flash_attention_compat/ATen/cuda/CUDAGeneratorImpl.h
@@ -0,0 +1,18 @@
+#pragma once
+
+#include <cstdint>
+
+// FlashAttention keeps this type in its parameter ABI even when dropout is
+// compiled out. InfiniTrain only supports dropout=0 in the native adapter, so
+// no generator state is acquired or consumed.
+// FIXME(zbl): This minimal stand-in is coupled to the pinned FlashAttention
+// headers. Revalidate it when upgrading FlashAttention or including real ATen,
+// since layout/API changes may break compilation or conflict with ATen's type.
+namespace at {
+
+struct PhiloxCudaState {
+    uint64_t seed = 0;
+    uint64_t offset = 0;
+};
+
+} // namespace at
diff --git a/infini_train/src/kernels/cuda/flash_attention_compat/ATen/cuda/detail/UnpackRaw.cuh b/infini_train/src/kernels/cuda/flash_attention_compat/ATen/cuda/detail/UnpackRaw.cuh
new file mode 100644
index 00000000..8a052605
--- /dev/null
+++ b/infini_train/src/kernels/cuda/flash_attention_compat/ATen/cuda/detail/UnpackRaw.cuh
@@ -0,0 +1,14 @@
+#pragma once
+
+#include <cstdint>
+#include <tuple>
+
+#include "ATen/cuda/CUDAGeneratorImpl.h"
+
+namespace at::cuda::philox {
+
+__host__ __device__ inline std::tuple<uint64_t, uint64_t> unpack(const at::PhiloxCudaState &state) {
+    return {state.seed, state.offset};
+}
+
+} // namespace at::cuda::philox
diff --git a/infini_train/src/kernels/cuda/flash_attention_compat/c10/cuda/CUDAException.h b/infini_train/src/kernels/cuda/flash_attention_compat/c10/cuda/CUDAException.h
new file mode 100644
index 00000000..261e052a
--- /dev/null
+++ b/infini_train/src/kernels/cuda/flash_attention_compat/c10/cuda/CUDAException.h
@@ -0,0 +1,23 @@
+#pragma once
+
+#include <cstdio>
+#include <cstdlib>
+
+#include <cuda_runtime.h>
+
+namespace infini_train::kernels::cuda::flash_attention_compat {
+
+inline void CheckCuda(cudaError_t status, const char *expression, const char *file, int line) {
+    if (status == cudaSuccess) {
+        return;
+    }
+    std::fprintf(stderr, "FlashAttention CUDA error at %s:%d: %s failed: %s\n", file, line, expression,
+                 cudaGetErrorString(status));
+    std::abort();
+}
+
+} // namespace infini_train::kernels::cuda::flash_attention_compat
+
+#define C10_CUDA_CHECK(expression)                                                                                     \
+    ::infini_train::kernels::cuda::flash_attention_compat::CheckCuda((expression), #expression, __FILE__, __LINE__)
+#define C10_CUDA_KERNEL_LAUNCH_CHECK() C10_CUDA_CHECK(cudaGetLastError())
diff --git a/infini_train/src/nn/functional.cc b/infini_train/src/nn/functional.cc
index c33e2368..e7749a48 100644
--- a/infini_train/src/nn/functional.cc
+++ b/infini_train/src/nn/functional.cc
@@ -7,6 +7,7 @@
 #include "infini_train/include/autograd/activations.h"
 #include "infini_train/include/autograd/elementwise.h"
 #include "infini_train/include/autograd/reduction.h"
+#include "infini_train/include/autograd/scaled_dot_product_attention.h"
 #include "infini_train/include/autograd/softmax.h"
 #include "infini_train/include/autograd/transform.h"
 #include "infini_train/include/nn/init.h"
@@ -75,6 +76,11 @@ std::shared_ptr<Tensor> Softmax(const std::shared_ptr<Tensor> &input, int64_t di
     return std::make_shared<autograd::Softmax>(dim)->Apply({input})[0];
 }
 
+std::shared_ptr<Tensor> ScaledDotProductAttention(const std::shared_ptr<Tensor> &q, const std::shared_ptr<Tensor> &k,
+                                                  const std::shared_ptr<Tensor> &v, float scale) {
+    return std::make_shared<autograd::ScaledDotProductAttention>(scale)->Apply({q, k, v})[0];
+}
+
 std::shared_ptr<Tensor> Sigmoid(const std::shared_ptr<Tensor> &input) {
     return std::make_shared<autograd::Sigmoid>()->Apply({input})[0];
 }
diff --git a/infini_train/src/nn/modules/transformer/causal_self_attention.cc b/infini_train/src/nn/modules/transformer/causal_self_attention.cc
index bc4c6bd4..1a0ff5f5 100644
--- a/infini_train/src/nn/modules/transformer/causal_self_attention.cc
+++ b/infini_train/src/nn/modules/transformer/causal_self_attention.cc
@@ -115,13 +115,23 @@ CausalSelfAttention::Forward(const std::vector<std::shared_ptr<infini_train::Ten
     //            use Slice() to work around here
     const int64_t q_size_local = H_local * D;
     const int64_t kv_size_local = KV_local * D;
-    // -> Split into q, k, v
-    // q: (B, T, H_local, D)
-    auto q = qkv->Slice(2, 0, q_size_local)->View({B, T, H_local, D});
-    // k: (B, T, KV_local, D)
-    auto k = qkv->Slice(2, q_size_local, q_size_local + kv_size_local)->View({B, T, KV_local, D});
-    // v: (B, T, KV_local, D)
-    auto v = qkv->Slice(2, q_size_local + kv_size_local, q_size_local + 2 * kv_size_local)->View({B, T, KV_local, D});
+    std::shared_ptr<Tensor> q;
+    std::shared_ptr<Tensor> k;
+    std::shared_ptr<Tensor> v;
+    if (q_size_local == kv_size_local) {
+        // Split uses one autograd node for equal-sized MHA projections. Three independent Slice nodes are
+        // considerably more expensive, especially in backward where each one materializes a full-size input grad.
+        auto qkv_chunks = qkv->Split(q_size_local, 2);
+        CHECK_EQ(qkv_chunks.size(), 3);
+        q = qkv_chunks[0]->View({B, T, H_local, D});
+        k = qkv_chunks[1]->View({B, T, KV_local, D});
+        v = qkv_chunks[2]->View({B, T, KV_local, D});
+    } else {
+        // GQA has unequal Q and K/V widths, which the fixed-size Split API cannot represent.
+        q = qkv->Slice(2, 0, q_size_local)->View({B, T, H_local, D});
+        k = qkv->Slice(2, q_size_local, q_size_local + kv_size_local)->View({B, T, KV_local, D});
+        v = qkv->Slice(2, q_size_local + kv_size_local, q_size_local + 2 * kv_size_local)->View({B, T, KV_local, D});
+    }
 
     if (config_.position_embedding_type == PositionEmbeddingType::kRoPE) {
         // q: (B, T, H_local, D), k: (B, T, KV_local, D)
@@ -131,38 +141,43 @@ CausalSelfAttention::Forward(const std::vector<std::shared_ptr<infini_train::Ten
     // TODO(zbl): use kv cache during inference
     // if (use_kv_) { ... }
 
-    // align n_head in GQA
-    // (B, T, KV_local, D) -> (B, T, H_local, D) via RepeatKV
-    k = RepeatKV(k, n_rep_);
-    v = RepeatKV(v, n_rep_);
+    if (!config_.flash) {
+        // (B, T, KV_local, D) -> (B, T, H_local, D) via RepeatKV
+        k = RepeatKV(k, n_rep_);
+        v = RepeatKV(v, n_rep_);
+    }
 
     // (B, T, H_local, D) -> (B, H_local, T, D)
     q = q->Transpose(1, 2);
     k = k->Transpose(1, 2);
     v = v->Transpose(1, 2);
 
-    // TODO(zbl): support flash attention later
-    // if (flash_) { ... }
-
-    // manual implementation of attention
-    // this materializes the large (T,T) matrix for all the queries and keys
-
-    // q: (B, H_local, T, D)
-    // k: (B, H_local, T, D) -> (B, H_local, D, T)
-    // q @ k.T: (B, H_local, T, T) -> mul 1.0 / sqrt(D) -> (B, H_local, T, T)
-    auto att = q->Matmul(k->Transpose(-2, -1)) * (1.0 / std::sqrt(static_cast<float>(D)));
-    if (mask) {
-        // mask: (1, 1, T, T)
-        att = att->MaskedFill(mask, std::numeric_limits<float>::lowest());
+    std::shared_ptr<Tensor> y;
+    if (config_.flash) {
+        // FIXME(zbl): FlashAttention assumes start_pos=0 and uses its built-in causal mask;
+        // start_pos and mask are ignored until incremental decoding and custom masks are supported.
+        y = nn::function::ScaledDotProductAttention(q, k, v, 1.0f / std::sqrt(static_cast<float>(D)));
     } else {
-        // fallback causal mask: (1, 1, T, T)
-        auto causal_mask = buffers_[kParamBiasName]->Slice({0, 0, 0, 0}, {1, 1, T, T}, {1, 1, 1, 1});
-        att = att->MaskedFill(causal_mask == 0, -std::numeric_limits<float>::infinity());
+        // manual implementation of attention
+        // this materializes the large (T,T) matrix for all the queries and keys
+
+        // q: (B, H_local, T, D)
+        // k: (B, H_local, T, D) -> (B, H_local, D, T)
+        // q @ k.T: (B, H_local, T, T) -> mul 1.0 / sqrt(D) -> (B, H_local, T, T)
+        auto att = q->Matmul(k->Transpose(-2, -1)) * (1.0 / std::sqrt(static_cast<float>(D)));
+        if (mask) {
+            // mask: (1, 1, T, T)
+            att = att->MaskedFill(mask, std::numeric_limits<float>::lowest());
+        } else {
+            // fallback causal mask: (1, 1, T, T)
+            auto causal_mask = buffers_[kParamBiasName]->Slice({0, 0, 0, 0}, {1, 1, T, T}, {1, 1, 1, 1});
+            att = att->MaskedFill(causal_mask == 0, -std::numeric_limits<float>::infinity());
+        }
+        // (B, H_local, T, T)
+        att = nn::function::Softmax(att, -1);
+        // att: (B, H_local, T, T) @ v: (B, H_local, T, D) -> y: (B, H_local, T, D)
+        y = att->Matmul(v);
     }
-    // (B, H_local, T, T)
-    att = nn::function::Softmax(att, -1);
-    // att: (B, H_local, T, T) @ v: (B, H_local, T, D) -> y: (B, H_local, T, D)
-    auto y = att->Matmul(v);
     // (B, H_local, T, D) -> Transpose(1, 2) -> (B, T, H_local, D) -> (B, T, C_local)
     y = y->Transpose(1, 2)->Contiguous()->View({B, T, C_local});
     // output projection
diff --git a/scripts/test_config.json b/scripts/test_config.json
index 992cf5d2..1ba0330d 100644
--- a/scripts/test_config.json
+++ b/scripts/test_config.json
@@ -12,17 +12,17 @@
         "CKPT_ROOT_DIR": "/data1/ckpt",
         "COMPARE_LOG_DIR": "",
         "RUN_CTEST": "true",
-        "RUN_PROFILE_TEST": "true",
+        "RUN_PROFILE_TEST": "false",
         "MIXTRAL_INPUT_BIN": "/data1/shared/InfiniTrain-dev/data/llmc/llama3/tinyshakespeare/tiny_shakespeare_train.bin",
         "MIXTRAL_LLMC_FILEPATH": "/data1/shared/InfiniTrain-dev/data/llmc/mixtral/mixtral_megatron_export.bin",
-        "GPT2_TEST_GROUPS": "basic,zero,lora,checkpoint,lr_scheduler",
-        "LLAMA3_TEST_GROUPS": "basic,zero,lora,checkpoint,lr_scheduler",
+        "GPT2_TEST_GROUPS": "basic,flash,zero,lora,checkpoint,lr_scheduler",
+        "LLAMA3_TEST_GROUPS": "basic,flash,zero,lora,checkpoint,lr_scheduler",
         "MIXTRAL_TEST_GROUPS": "moe"
     },
     "basic_compile_commands": [
         {
             "id": "build_1",
-            "cmd": "cmake -DUSE_CUDA=ON -DUSE_NCCL=ON .. && make -j"
+            "cmd": "cmake -DUSE_CUDA=ON -DUSE_NCCL=ON -DUSE_FLASH_ATTENTION=ON .. && make -j"
         }
     ],
     "test_groups": [
@@ -201,6 +201,107 @@
                 }
             ]
         },
+        {
+            "tag": "flash",
+            "tests": [
+                {
+                    "id": "dp8_bs4_seq128_tb4096",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 4,
+                        "sequence_length": 128,
+                        "total_batch_size": 4096,
+                        "attention_backend": "flash"
+                    }
+                },
+                {
+                    "id": "dp8_bs8_seq128_tb8192",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 8,
+                        "sequence_length": 128,
+                        "total_batch_size": 8192,
+                        "attention_backend": "flash"
+                    }
+                },
+                {
+                    "id": "dp8_bs16_seq128_tb16384",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 16,
+                        "sequence_length": 128,
+                        "total_batch_size": 16384,
+                        "attention_backend": "flash"
+                    }
+                },
+                {
+                    "id": "dp8_bs4_seq256_tb8192",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 4,
+                        "sequence_length": 256,
+                        "total_batch_size": 8192,
+                        "attention_backend": "flash"
+                    }
+                },
+                {
+                    "id": "dp8_bs4_seq256_tb16384",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 4,
+                        "sequence_length": 256,
+                        "total_batch_size": 16384,
+                        "attention_backend": "flash"
+                    }
+                },
+                {
+                    "id": "dp8_bs8_seq256_tb16384",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 8,
+                        "sequence_length": 256,
+                        "total_batch_size": 16384,
+                        "attention_backend": "flash"
+                    }
+                },
+                {
+                    "id": "dp8_bs4_seq512_tb16384",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 4,
+                        "sequence_length": 512,
+                        "total_batch_size": 16384,
+                        "attention_backend": "flash"
+                    }
+                },
+                {
+                    "id": "dp8_bs2_seq1024_tb16384",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 8,
+                        "num_iteration": 10,
+                        "batch_size": 2,
+                        "sequence_length": 1024,
+                        "total_batch_size": 16384,
+                        "attention_backend": "flash"
+                    }
+                }
+            ]
+        },
         {
             "tag": "zero",
             "tests": [
diff --git a/tests/autograd/CMakeLists.txt b/tests/autograd/CMakeLists.txt
index 50bd6096..9d20deef 100644
--- a/tests/autograd/CMakeLists.txt
+++ b/tests/autograd/CMakeLists.txt
@@ -3,7 +3,16 @@
 # ============================================================================
 
 file(GLOB AUTOGRAD_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc)
+list(FILTER AUTOGRAD_SOURCES EXCLUDE REGEX ".*/test_autograd_scaled_dot_product_attention\\.cc$")
 
 infini_train_add_test_suite(test_autograd
   SOURCES ${AUTOGRAD_SOURCES}
 )
+
+if(USE_FLASH_ATTENTION)
+  infini_train_add_test(test_flash_attention_cuda
+    SOURCES test_autograd_scaled_dot_product_attention.cc
+    LABELS cuda
+    TEST_TIMEOUT 60
+  )
+endif()
diff --git a/tests/autograd/test_autograd_scaled_dot_product_attention.cc b/tests/autograd/test_autograd_scaled_dot_product_attention.cc
new file mode 100644
index 00000000..5b32940e
--- /dev/null
+++ b/tests/autograd/test_autograd_scaled_dot_product_attention.cc
@@ -0,0 +1,376 @@
+#include <algorithm>
+#include <cmath>
+#include <cstddef>
+#include <limits>
+#include <string>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+#include "infini_train/include/autograd/function_hook.h"
+#include "infini_train/include/nn/functional.h"
+#include "infini_train/include/tensor.h"
+
+#include "tests/common/test_utils.h"
+
+using namespace infini_train;
+
+namespace {
+
+struct AttentionResult {
+    std::vector<float> output;
+    std::vector<float> dq;
+    std::vector<float> dk;
+    std::vector<float> dv;
+    DataType dq_dtype;
+    DataType dk_dtype;
+    DataType dv_dtype;
+};
+
+struct PackedAttentionResult {
+    std::vector<float> output;
+    std::vector<float> dqkv;
+};
+
+std::vector<float> MakeValues(size_t count, float frequency, float scale) {
+    std::vector<float> values(count);
+    for (size_t i = 0; i < count; ++i) { values[i] = std::sin(static_cast<float>(i) * frequency) * scale; }
+    return values;
+}
+
+std::vector<float> ToFloatVector(const std::shared_ptr<Tensor> &tensor) {
+    auto float_tensor = tensor->Dtype() == DataType::kFLOAT32 ? *tensor : tensor->To(DataType::kFLOAT32);
+    auto cpu_tensor = float_tensor.GetDevice().IsCPU() ? float_tensor : float_tensor.To(Device());
+    const auto *data = static_cast<const float *>(cpu_tensor.DataPtr());
+    return {data, data + cpu_tensor.NumElements()};
+}
+
+class CaptureGradHook : public autograd::PreAccumulateGradHook {
+public:
+    void operator()(const std::shared_ptr<Tensor> &) override {}
+
+    bool TryBypassAccumulate(const std::shared_ptr<Tensor> &, const std::shared_ptr<Tensor> &grad_output, bool,
+                             float) override {
+        dtype_ = grad_output->Dtype();
+        values_ = ToFloatVector(grad_output);
+        return true;
+    }
+
+    const std::vector<float> &Values() const { return values_; }
+    DataType Dtype() const { return dtype_; }
+
+private:
+    std::vector<float> values_;
+    DataType dtype_ = DataType::kFLOAT32;
+};
+
+std::shared_ptr<Tensor> MakeTensor(const std::vector<float> &values, const std::vector<int64_t> &dims, Device device,
+                                   DataType dtype = DataType::kBFLOAT16) {
+    auto cpu_float = std::make_shared<Tensor>(values.data(), dims, DataType::kFLOAT32);
+    auto cpu_tensor = std::make_shared<Tensor>(cpu_float->To(dtype));
+    return std::make_shared<Tensor>(cpu_tensor->To(device));
+}
+
+AttentionResult RunFlashAttention(Device device, bool expand_kv, float input_scale = 0.2F, float grad_scale = 0.1F) {
+    constexpr int64_t batch = 1;
+    constexpr int64_t query_heads = 8;
+    constexpr int64_t kv_heads = 2;
+    constexpr int64_t seqlen = 64;
+    constexpr int64_t head_dim = 64;
+    constexpr int64_t groups = query_heads / kv_heads;
+
+    const std::vector<int64_t> q_dims{batch, query_heads, seqlen, head_dim};
+    const std::vector<int64_t> kv_dims{batch, kv_heads, seqlen, head_dim};
+    const auto q_values = MakeValues(batch * query_heads * seqlen * head_dim, 0.013F, input_scale);
+    const auto k_values = MakeValues(batch * kv_heads * seqlen * head_dim, 0.017F, input_scale);
+    const auto v_values = MakeValues(batch * kv_heads * seqlen * head_dim, 0.019F, input_scale);
+    const auto grad_values = MakeValues(batch * query_heads * seqlen * head_dim, 0.023F, grad_scale);
+
+    auto q = MakeTensor(q_values, q_dims, device);
+    auto k = MakeTensor(k_values, kv_dims, device);
+    auto v = MakeTensor(v_values, kv_dims, device);
+    q->RequiresGrad();
+    k->RequiresGrad();
+    v->RequiresGrad();
+    auto q_grad = std::make_shared<CaptureGradHook>();
+    auto k_grad = std::make_shared<CaptureGradHook>();
+    auto v_grad = std::make_shared<CaptureGradHook>();
+    q->RegisterPreAccumulateGradHook(q_grad);
+    k->RegisterPreAccumulateGradHook(k_grad);
+    v->RegisterPreAccumulateGradHook(v_grad);
+
+    auto k_input = expand_kv ? k->RepeatInterleave(groups, 1) : k;
+    auto v_input = expand_kv ? v->RepeatInterleave(groups, 1) : v;
+    auto output = nn::function::ScaledDotProductAttention(q, k_input, v_input, 1.0F / std::sqrt(head_dim));
+    auto grad = MakeTensor(grad_values, q_dims, device);
+    output->Backward(grad);
+
+    return {
+        .output = ToFloatVector(output),
+        .dq = q_grad->Values(),
+        .dk = k_grad->Values(),
+        .dv = v_grad->Values(),
+        .dq_dtype = q_grad->Dtype(),
+        .dk_dtype = k_grad->Dtype(),
+        .dv_dtype = v_grad->Dtype(),
+    };
+}
+
+AttentionResult RunUnfusedAttention(Device device, bool upcast = false, float input_scale = 0.2F,
+                                    float grad_scale = 0.1F) {
+    constexpr int64_t batch = 1;
+    constexpr int64_t query_heads = 8;
+    constexpr int64_t kv_heads = 2;
+    constexpr int64_t seqlen = 64;
+    constexpr int64_t head_dim = 64;
+    constexpr int64_t groups = query_heads / kv_heads;
+
+    const std::vector<int64_t> q_dims{batch, query_heads, seqlen, head_dim};
+    const std::vector<int64_t> kv_dims{batch, kv_heads, seqlen, head_dim};
+    const auto q_values = MakeValues(batch * query_heads * seqlen * head_dim, 0.013F, input_scale);
+    const auto k_values = MakeValues(batch * kv_heads * seqlen * head_dim, 0.017F, input_scale);
+    const auto v_values = MakeValues(batch * kv_heads * seqlen * head_dim, 0.019F, input_scale);
+    const auto grad_values = MakeValues(batch * query_heads * seqlen * head_dim, 0.023F, grad_scale);
+
+    auto q = MakeTensor(q_values, q_dims, device);
+    auto k = MakeTensor(k_values, kv_dims, device);
+    auto v = MakeTensor(v_values, kv_dims, device);
+    if (upcast) {
+        q = std::make_shared<Tensor>(q->To(DataType::kFLOAT32));
+        k = std::make_shared<Tensor>(k->To(DataType::kFLOAT32));
+        v = std::make_shared<Tensor>(v->To(DataType::kFLOAT32));
+    }
+    q->RequiresGrad();
+    k->RequiresGrad();
+    v->RequiresGrad();
+    auto q_grad = std::make_shared<CaptureGradHook>();
+    auto k_grad = std::make_shared<CaptureGradHook>();
+    auto v_grad = std::make_shared<CaptureGradHook>();
+    q->RegisterPreAccumulateGradHook(q_grad);
+    k->RegisterPreAccumulateGradHook(k_grad);
+    v->RegisterPreAccumulateGradHook(v_grad);
+
+    auto k_expanded = k->RepeatInterleave(groups, 1);
+    auto v_expanded = v->RepeatInterleave(groups, 1);
+    auto scores = q->Matmul(k_expanded->Transpose(-2, -1)) * (1.0F / std::sqrt(head_dim));
+    std::vector<float> mask_values(seqlen * seqlen, 0.0F);
+    for (int64_t row = 0; row < seqlen; ++row) {
+        for (int64_t column = row + 1; column < seqlen; ++column) { mask_values[row * seqlen + column] = 1.0F; }
+    }
+    auto mask = MakeTensor(mask_values, {1, 1, seqlen, seqlen}, device, DataType::kBOOL);
+    auto probabilities = nn::function::Softmax(scores->MaskedFill(mask, std::numeric_limits<float>::lowest()), -1);
+    auto output = probabilities->Matmul(v_expanded);
+    auto grad = MakeTensor(grad_values, q_dims, device);
+    if (upcast) {
+        grad = std::make_shared<Tensor>(grad->To(DataType::kFLOAT32));
+    }
+    output->Backward(grad);
+
+    return {
+        .output = ToFloatVector(output),
+        .dq = q_grad->Values(),
+        .dk = k_grad->Values(),
+        .dv = v_grad->Values(),
+    };
+}
+
+PackedAttentionResult RunPackedFlashAttention(Device device, bool expand_kv) {
+    constexpr int64_t batch = 1;
+    constexpr int64_t query_heads = 8;
+    constexpr int64_t kv_heads = 2;
+    constexpr int64_t seqlen = 64;
+    constexpr int64_t head_dim = 64;
+    constexpr int64_t groups = query_heads / kv_heads;
+    constexpr int64_t query_width = query_heads * head_dim;
+    constexpr int64_t kv_width = kv_heads * head_dim;
+    constexpr int64_t packed_width = query_width + 2 * kv_width;
+
+    const std::vector<int64_t> packed_dims{batch, seqlen, packed_width};
+    const std::vector<int64_t> output_dims{batch, seqlen, query_width};
+    const auto packed_values = MakeValues(batch * seqlen * packed_width, 0.013F, 0.2F);
+    const auto grad_values = MakeValues(batch * seqlen * query_width, 0.023F, 0.1F);
+
+    auto packed = MakeTensor(packed_values, packed_dims, device);
+    packed->RequiresGrad();
+    auto packed_grad = std::make_shared<CaptureGradHook>();
+    packed->RegisterPreAccumulateGradHook(packed_grad);
+    auto q = packed->Slice(2, 0, query_width)->View({batch, seqlen, query_heads, head_dim});
+    auto k = packed->Slice(2, query_width, query_width + kv_width)->View({batch, seqlen, kv_heads, head_dim});
+    auto v = packed->Slice(2, query_width + kv_width, packed_width)->View({batch, seqlen, kv_heads, head_dim});
+    if (expand_kv) {
+        k = k->RepeatInterleave(groups, 2);
+        v = v->RepeatInterleave(groups, 2);
+    }
+    q = q->Transpose(1, 2);
+    k = k->Transpose(1, 2);
+    v = v->Transpose(1, 2);
+
+    auto output = nn::function::ScaledDotProductAttention(q, k, v, 1.0F / std::sqrt(head_dim));
+    output = output->Transpose(1, 2)->Contiguous()->View(output_dims);
+    auto grad = MakeTensor(grad_values, output_dims, device);
+    output->Backward(grad);
+
+    return {
+        .output = ToFloatVector(output),
+        .dqkv = packed_grad->Values(),
+    };
+}
+
+PackedAttentionResult RunPackedUnfusedAttention(Device device) {
+    constexpr int64_t batch = 1;
+    constexpr int64_t query_heads = 8;
+    constexpr int64_t kv_heads = 2;
+    constexpr int64_t seqlen = 64;
+    constexpr int64_t head_dim = 64;
+    constexpr int64_t groups = query_heads / kv_heads;
+    constexpr int64_t query_width = query_heads * head_dim;
+    constexpr int64_t kv_width = kv_heads * head_dim;
+    constexpr int64_t packed_width = query_width + 2 * kv_width;
+
+    const std::vector<int64_t> packed_dims{batch, seqlen, packed_width};
+    const std::vector<int64_t> output_dims{batch, seqlen, query_width};
+    const auto packed_values = MakeValues(batch * seqlen * packed_width, 0.013F, 0.2F);
+    const auto grad_values = MakeValues(batch * seqlen * query_width, 0.023F, 0.1F);
+
+    auto packed = MakeTensor(packed_values, packed_dims, device);
+    packed->RequiresGrad();
+    auto packed_grad = std::make_shared<CaptureGradHook>();
+    packed->RegisterPreAccumulateGradHook(packed_grad);
+    auto q = packed->Slice(2, 0, query_width)->View({batch, seqlen, query_heads, head_dim});
+    auto k = packed->Slice(2, query_width, query_width + kv_width)->View({batch, seqlen, kv_heads, head_dim});
+    auto v = packed->Slice(2, query_width + kv_width, packed_width)->View({batch, seqlen, kv_heads, head_dim});
+    k = k->RepeatInterleave(groups, 2);
+    v = v->RepeatInterleave(groups, 2);
+    q = q->Transpose(1, 2);
+    k = k->Transpose(1, 2);
+    v = v->Transpose(1, 2);
+
+    auto scores = q->Matmul(k->Transpose(-2, -1)) * (1.0F / std::sqrt(head_dim));
+    std::vector<float> mask_values(seqlen * seqlen, 0.0F);
+    for (int64_t row = 0; row < seqlen; ++row) {
+        for (int64_t column = row + 1; column < seqlen; ++column) { mask_values[row * seqlen + column] = 1.0F; }
+    }
+    auto mask = MakeTensor(mask_values, {1, 1, seqlen, seqlen}, device, DataType::kBOOL);
+    auto probabilities = nn::function::Softmax(scores->MaskedFill(mask, std::numeric_limits<float>::lowest()), -1);
+    auto output = probabilities->Matmul(v);
+    output = output->Transpose(1, 2)->Contiguous()->View(output_dims);
+    auto grad = MakeTensor(grad_values, output_dims, device);
+    output->Backward(grad);
+
+    return {
+        .output = ToFloatVector(output),
+        .dqkv = packed_grad->Values(),
+    };
+}
+
+void ExpectClose(const std::vector<float> &actual, const std::vector<float> &expected, float max_relative_l2,
+                 float min_cosine, const std::string &name) {
+    ASSERT_EQ(actual.size(), expected.size()) << name;
+    double diff_squared = 0.0;
+    double actual_squared = 0.0;
+    double expected_squared = 0.0;
+    double dot = 0.0;
+    float max_abs = 0.0F;
+    for (size_t i = 0; i < actual.size(); ++i) {
+        const double diff = static_cast<double>(actual[i]) - expected[i];
+        diff_squared += diff * diff;
+        actual_squared += static_cast<double>(actual[i]) * actual[i];
+        expected_squared += static_cast<double>(expected[i]) * expected[i];
+        dot += static_cast<double>(actual[i]) * expected[i];
+        max_abs = std::max(max_abs, std::abs(actual[i] - expected[i]));
+    }
+    const double relative_l2 = std::sqrt(diff_squared) / std::max(std::sqrt(expected_squared), 1e-30);
+    const double cosine = dot / std::max(std::sqrt(actual_squared * expected_squared), 1e-30);
+    EXPECT_LE(relative_l2, max_relative_l2) << name << " max_abs=" << max_abs << " cosine=" << cosine;
+    EXPECT_GE(cosine, min_cosine) << name << " max_abs=" << max_abs << " relative_l2=" << relative_l2;
+}
+
+float MaxAbsDiff(const std::vector<float> &actual, const std::vector<float> &expected) {
+    EXPECT_EQ(actual.size(), expected.size());
+    float max_abs = 0.0F;
+    for (size_t i = 0; i < std::min(actual.size(), expected.size()); ++i) {
+        max_abs = std::max(max_abs, std::abs(actual[i] - expected[i]));
+    }
+    return max_abs;
+}
+
+void ExpectErrorBoundedByBfloat16Reference(const std::vector<float> &flash, const std::vector<float> &bfloat16,
+                                           const std::vector<float> &float32, const std::string &name) {
+    const float flash_error = MaxAbsDiff(flash, float32);
+    const float bfloat16_error = MaxAbsDiff(bfloat16, float32);
+    EXPECT_LE(flash_error, 3.0F * bfloat16_error + 1e-5F)
+        << name << " flash_max_error=" << flash_error << " bfloat16_max_error=" << bfloat16_error;
+}
+
+} // namespace
+
+class AutogradScaledDotProductAttentionTest : public infini_train::test::InfiniTrainTest {};
+
+TEST_P(AutogradScaledDotProductAttentionTest, NativeGqaMatchesExpandedKv) {
+    ONLY_CUDA();
+
+    const auto native_gqa = RunFlashAttention(GetDevice(), false);
+    const auto expanded_kv = RunFlashAttention(GetDevice(), true);
+
+    ExpectClose(native_gqa.output, expanded_kv.output, 0.001F, 0.9999F, "output");
+    ExpectClose(native_gqa.dq, expanded_kv.dq, 0.001F, 0.9999F, "dQ");
+    ExpectClose(native_gqa.dk, expanded_kv.dk, 0.01F, 0.9999F, "dK");
+    ExpectClose(native_gqa.dv, expanded_kv.dv, 0.01F, 0.9999F, "dV");
+}
+
+TEST_P(AutogradScaledDotProductAttentionTest, Bfloat16BackwardPropagatesFloat32Gradients) {
+    ONLY_CUDA();
+
+    const auto result = RunFlashAttention(GetDevice(), false);
+
+    EXPECT_EQ(result.dq_dtype, DataType::kFLOAT32);
+    EXPECT_EQ(result.dk_dtype, DataType::kFLOAT32);
+    EXPECT_EQ(result.dv_dtype, DataType::kFLOAT32);
+}
+
+TEST_P(AutogradScaledDotProductAttentionTest, NativeGqaMatchesExpandedKvWithPackedInput) {
+    ONLY_CUDA();
+
+    const auto native_gqa = RunPackedFlashAttention(GetDevice(), false);
+    const auto expanded_kv = RunPackedFlashAttention(GetDevice(), true);
+
+    ExpectClose(native_gqa.output, expanded_kv.output, 0.001F, 0.9999F, "output");
+    ExpectClose(native_gqa.dqkv, expanded_kv.dqkv, 0.02F, 0.9999F, "packed dQKV");
+}
+
+TEST_P(AutogradScaledDotProductAttentionTest, NativeGqaMatchesUnfusedReferenceWithPackedInput) {
+    ONLY_CUDA();
+
+    const auto native_gqa = RunPackedFlashAttention(GetDevice(), false);
+    const auto unfused = RunPackedUnfusedAttention(GetDevice());
+
+    ExpectClose(native_gqa.output, unfused.output, 0.01F, 0.9995F, "output");
+    ExpectClose(native_gqa.dqkv, unfused.dqkv, 0.03F, 0.9995F, "packed dQKV");
+}
+
+TEST_P(AutogradScaledDotProductAttentionTest, NativeGqaMatchesUnfusedReference) {
+    ONLY_CUDA();
+
+    const auto native_gqa = RunFlashAttention(GetDevice(), false);
+    const auto unfused = RunUnfusedAttention(GetDevice());
+
+    ExpectClose(native_gqa.output, unfused.output, 0.01F, 0.9995F, "output");
+    ExpectClose(native_gqa.dq, unfused.dq, 0.01F, 0.9995F, "dQ");
+    ExpectClose(native_gqa.dk, unfused.dk, 0.03F, 0.9995F, "dK");
+    ExpectClose(native_gqa.dv, unfused.dv, 0.01F, 0.9995F, "dV");
+}
+
+TEST_P(AutogradScaledDotProductAttentionTest, NativeGqaErrorIsBoundedByUnfusedBfloat16Error) {
+    ONLY_CUDA();
+
+    const auto native_gqa = RunFlashAttention(GetDevice(), false, 1.0F, 1.0F);
+    const auto bfloat16 = RunUnfusedAttention(GetDevice(), false, 1.0F, 1.0F);
+    const auto float32 = RunUnfusedAttention(GetDevice(), true, 1.0F, 1.0F);
+
+    ExpectErrorBoundedByBfloat16Reference(native_gqa.output, bfloat16.output, float32.output, "output");
+    ExpectErrorBoundedByBfloat16Reference(native_gqa.dq, bfloat16.dq, float32.dq, "dQ");
+    ExpectErrorBoundedByBfloat16Reference(native_gqa.dk, bfloat16.dk, float32.dk, "dK");
+    ExpectErrorBoundedByBfloat16Reference(native_gqa.dv, bfloat16.dv, float32.dv, "dV");
+}
+
+INFINI_TRAIN_REGISTER_TEST(AutogradScaledDotProductAttentionTest);
diff --git a/third_party/flash-attention b/third_party/flash-attention
new file mode 160000
index 00000000..5231d95f
--- /dev/null
+++ b/third_party/flash-attention
@@ -0,0 +1 @@
+Subproject commit 5231d95fe13733fb534c01895f7ea88c6a6c7793