fix(attention): fix unflatten crash with flash-attn >= 2.5.9 due to API return type change

[zhj2300380]

🐛 Bug Description / 问题描述

When using flash-attn >= 2.5.9, the application crashes with:

RuntimeError: unflatten: Provided sizes [1, 32760] don't multiply up to the size of dim 0 (40)

使用 flash-attn >= 2.5.9 时，应用程序崩溃并报错：

RuntimeError: unflatten: Provided sizes [1, 32760] don't multiply up to the size of dim 0 (40)

🔍 Root Cause / 根本原因

flash-attn >= 2.5.9 changed the default return type of flash_attn_varlen_func:

• Old API (< 2.5.9): Returns Tuple[Tensor, ...] → requires [0] unpacking
• New API (>= 2.5.9): Returns Tensor directly → [0] incorrectly indexes the tensor

flash-attn >= 2.5.9 改变了 flash_attn_varlen_func 的默认返回类型：

• 旧版 API (< 2.5.9)：返回 Tuple[Tensor, ...] → 需要 [0] 解包
• 新版 API (>= 2.5.9)：直接返回 Tensor → [0] 会错误索引 tensor

✅ Fix Summary / 修复摘要

Change / 修改	Description / 说明
_extract_output() helper	兼容函数处理新旧 API 返回类型差异
FA3 import fix	from flash_attn import flash_attn_interface
Deprecated params	移除 seqused_q/seqused_k 参数
Fallback chain	FA3 → FA2 → PyTorch 原生注意力三级降级
Remove hard assert	允许 PyTorch fallback，不再强制要求 FA2

🧪 Testing / 测试验证

• ✅ Verified with flash-attn==2.7.4.post1
• ✅ Backward compatible with flash-attn < 2.5.9
• ✅ Fallback to FA2/PyTorch when FA3 unavailable
• ✅ 使用 RTX 5090 D (31.8GB) 测试通过

📦 Environment / 测试环境

Component / 组件	Version / 版本
Python	3.12.3
PyTorch	2.8.0+cu128
CUDA	12.8
flash-attn	2.7.4.post1
GPU	NVIDIA GeForce RTX 5090 D

📝 Checklist / 检查清单

• Code follows project style / 代码遵循项目风格
• Backward compatible / 向后兼容
• Tested on actual hardware / 已在实际硬件测试
• Commit message follows Conventional Commits / 提交信息遵循规范