feat: Flash Attention

[Chamberlain0w0]

背景

当前 unfused causal attention 会显式生成完整的 T x T attention score。该 PR 为 InfiniTrain 接入 FlashAttention 2 native CUDA backend，并适配最新的统一 MHA/GQA Forward 路径。实现直接调用 FlashAttention CUDA kernels，在 CMakeLists 中不依赖 PyTorch extension，不链接 Torch、ATen、c10 或 Python runtime。

主要改动

FlashAttention native backend

• 以 submodule 引入 FlashAttention 2 v2.7.4.post1。
• 新增 flash_attn_native 静态库和 USE_FLASH_ATTENTION CMake 开关。
• 当前 AOT 编译 sm80 causal forward/backward kernels：
  • dtype：FP16、BF16
  • head dimension：64、128
• 增加最小 ATen/c10 兼容头，位于 src/kernels/cuda/flash_attention_compat，满足上游 kernel headers 中残留的 Philox 和 CUDA 检查接口，无需链接 Torch。

Framework 与 autograd

• 新增 nn::function::ScaledDotProductAttention(q, k, v, scale)，以及对应 autograd Function 和 CUDA dispatcher kernels。
• 在 autograd operator 边界通过 Dispatcher::HasKernel 检查 flash attn 的 backend availability。
• 根据 Q/KV head 数自动选择 MHA、GQA 或 MQA：
  • Hq == Hkv：MHA
  • Hq % Hkv == 0：GQA/MQA
  • 其他情况直接报错
• 暂时显式提升 BF16 backward gradients 到 FP32，以兼容当前 autocast/autograd dtype 语义。

Transformer 接入

• GPT-2 和 LLaMA3 增加 --attention_backend=unfused|flash。
• Flash backend 与 unfused backend 共用统一的 CausalSelfAttention::Forward；unfused GQA 继续通过 RepeatKV 展开 K/V；Flash 保留原始 KV heads，交由 kernel 来处理。
• 普通 MHA 的 packed QKV 使用单个 Split autograd node；不等宽 GQA 使用 Slice。
• checkpoint loader 会将 backend 选择写入 TransformerConfig::flash。

测试与文档

• 新增独立 CUDA 测试目标 test_flash_attention_cuda。
• 覆盖 native GQA、显式 KV 展开、packed QKV、unfused reference、forward/backward 和当前 gradient dtype contract。
• 在 scripts/test_config.json 中增加 flash 测试组，覆盖 GPT-2/LLaMA3 的多组 batch/sequence shapes。
• 新增 docs/flash_attn_integration_design.md，记录架构、构建方式、支持矩阵和已知技术债。

TODO

暂不支持：

• 外部 custom/padding mask
• 非零 start_pos
• cross-attention
• varlen
• KV cache 和 generation
• local attention、ALiBi、softcap
• deterministic backward

当前 Transformer Flash 路径会忽略已解析的 mask/start_pos，并固定使用 kernel 内建 causal mask。

[chen2021673]

那这里不支持的话是不是应该加个报错或者回退

[chen2021673]

这里默认把 BF16 梯度提升为 FP32，上次好像说应该再看看这里的合理性。

[chen2021673]

GenerateText 里会调用model.Forward()，但这个forward不在AutocastGuard里面，所以没有BF16数据类型参与计算，在flashatten入口会报错

[chen2021673]

这里默认ON合适吗？我们要默认编译FlashAttention吗

[chen2021673]

之前讨论过还是保留true配置，这里改成false是有需要吗