ops-transformer/attention/flash_attn/op_kernel/utils · gitcode_lijd/ops-transformer

cann-robotflash_attn DN不合轴模板 s2BaseSize 切换为 256

Files	Last commit	Last update
attenmask_m_gs1.h	flash_attn mask and tiling refactor Co-authored-by: comic97<gaokeming@hisilicon.com> # message auto-generated for no-merge-commit merge: !14865 merge flash_attn_refactor into master flash_attn mask and tiling refactor Created-by: comic97 Commit-by: comic97 Merged-by: cann-robot Description: ## 描述 ### 背景 1. flash_attn host tiling 在 arch35/arch92 存在大量重复实现：`InitTilingInfo`/`IsCapable`/`CalcScheduleMode`/`DoOpTiling`/`SetPlatMemoryInfo`/`InitImplParam`/`UpdateTilingKeyLayout`/`UpdateTilingKeyKvLayout`/`GenTilingKey` 等公共流程两处各写一份，新增模板/配置需双改，易遗漏。 2. mask 处理存在多套模型：tiling data 与 kernel 结构体同时保留 `sparseMode`/`attenMaskBatch`/`attenMaskS1Size`/`isExistRowInvalid` 与 `maskMode`；`attenmask_gs1.h` 同时承载合轴（GS/SG）与不合轴（S1_ONLY）逻辑；`MaskMode`/`FAMaskMode`/`fa_base_vector::SparseMode` 枚举语义重叠，且 flash_attn kernel 直接引用 `fa_base_vector` 的 mask 枚举。 3. AICPU metadata 关键路径缺少维测日志，问题定位成本高。 ### 改动 1. Host tiling 归一（主要改动） - 新增 `op_host/flash_attn_tiling_base.h`：`FlashAttnTilingBase` 承接公共 tiling 流程（`DoOpTiling`/`SetPlatMemoryInfo`/`InitImplParam`/`GenTilingKey`/`UpdateTilingKeyLayout`/`UpdateTilingKeyKvLayout`/`FillTiling`/`CalcWorkspaceSize`/`CalcNumBlocks` 等）。 - `FlashAttnTilingImplArch35`/`FlashAttnTilingImplArch92` 继承基类，仅覆写差异部分： - arch35：`SplitPolicy`、`UpdateTilingKeyConfig`、`UpdateTilingKeyInfo`、`UpdateTilingKeyTemplateId`； - arch92：`SplitPolicy`、`UpdateTilingKeyConfig`、`UpdateTilingKeyInfo`。 - `fa_adjust_sinner_souter.h` 拆分为 `arch35/fa_adjust_sinner_souter_arch35.h` 与 `arch92/fa_adjust_sinner_souter_arch92.h`；去掉 `AdjustSinnerAndSouter` 未使用的 `qLayout` 形参及 `LAYOUT_`/`SOUTER_`/`DSIZE_*` 常量（改为数值直接判断）。 - `fa_tiling_info.h` 的 `MaskMode` 改为 `using MaskMode = FAMaskMode`，host 与 kernel 共用同一枚举定义。 2. mask 数据模型统一 - `FlashAttnAttenMaskParams`/`CommonConstInfo`：`sparseMode` → `maskMode`，删除 `attenMaskBatch`/`attenMaskS1Size`/`isExistRowInvalid`；host 侧不再推导 maskBatch/maskS1Size，仅保留 `attenMaskS2Size`。 - `attenmask_gs1.h` 拆分为： - `attenmask_m_gs1.h`：`MaskInfoMerged`（合轴 GS/SG），删除 `MaskInfo<MaskFormat>` 模板特化与 `S1_ONLY` 分支； - `attenmask_m_s1.h`：`MaskInfoUnmerged`（不合轴 S1_ONLY），承载 `IsSkipAttentionmask`/`ComputeAttenMaskOffsetCompressDn`/`AttentionmaskCopyInDn`。 - 与接口约束（`mask_mode` 仅 0/3/4，`attn_mask` 固定 2D `[2048,2048]`）对齐，删除不可达分支：`DEFAULT_MASK`/`ALL_MASK`/`LEFT_UP_CAUSAL`、`ComputeAttenMaskOffsetNoCompress` 及 batch mask 偏移、`MASK_BOOL`/`attenMaskType`/`maskValue`，以及 flashdecode 中 `DealInvalidMaskRows`/`negativeIntScalar_` 等无效行 mask 清理逻辑。 - DN 不合轴保持原有条件路由与 S1_ONLY mask 语义（`MASK_M_BASE`、`ProcessVec1VfDn` 路径不变），仅适配新 mask 模型，不改变业务功能。 - kernel 侧（kernel_dn/kernel_dn_unmerged/kernel_nd/flashdecode）统一使用 `constInfo_.maskMode` 判断，`batchNeedRowInvalid` 简化为仅依赖 `hasValidRow`。 3. mask 枚举归一与依赖整理 - `flash_attn_common_def.h` 定义唯一 `enum class FAMaskMode { NO_MASK = 0, CAUSAL = 3, BAND = 4 }`；host `MaskMode` 使用该定义。 - flash_attn kernel 内对 `fa_base_vector::RIGHT_DOWN_CAUSAL`/`BAND` 等 mask 枚举的引用替换为 `FAMaskMode::CAUSAL`/`BAND`；`GetSafeActToken` 等公共函数调用保留，但不再依赖 `fa_base_vector` 的 mask 枚举定义。 - `vector_common.h` 由真实使用方直接包含（flashdecode、kernel_nd、kernel_dn、kernel_dn_unmerged），移除 block_vec 文件的冗余传递依赖。 4. AICPU metadata 维测 - Compute 进入/成功、InitBaseInfo、InitLoadBalanceParams、BalanceSchedule、FlashDecode 开启处增加 info/debug 日志；同时移除未使用的 `qLayout` 计算。 ## 关联的Issue 关联 Issue #6956 ## 测试 - 代码走查：host tiling 路由/配置映射、DN 不合轴条件路由、S1_ONLY DN mask 与重构前逐项对齐，业务行为不变。 - 编译：ophost、ascend950 opkernel、flash_attn_metadata AICPU 已初步编译通过；mask 枚举归一改动后建议 CI 复核。 - 静态检查：行宽（<=120）/codespell/pre-commit 待补充。 - 算子验证：待补充（建议覆盖 arch35/arch92 的 NO_MASK/CAUSAL/BAND、DN 不合轴、PA/TND 等场景）。 - 二级冒烟：待补充。 ## 文档更新更新 `attention/flash_attn/README.md`：tiling 基类与 Dn/Nd 命名统一、`attenmask_m_gs1.h`/`attenmask_m_s1.h`、`DealInvalidMaskRows` 移除等描述同步。 ## 类型标签 - [ ] 🐛 Bug修复 - [ ] ✨ 新特性 - [ ] ⚡ 性能优化 - [x] ♻️ 重构 - [ ] 🧪 测试 - [ ] 📦 构建/CI - [ ] 🔧 配置变更 - [x] 📝 文档更新 - [ ] ⬆️ 依赖升级 - [ ] 🔒 安全修复 - [x] 🧹 代码清理 - [ ] ❓ 其他，请描述： See merge request: cann/ops-transformer!14865	10 days ago
attenmask_m_s1.h	flash_attn mask and tiling refactor Co-authored-by: comic97<gaokeming@hisilicon.com> # message auto-generated for no-merge-commit merge: !14865 merge flash_attn_refactor into master flash_attn mask and tiling refactor Created-by: comic97 Commit-by: comic97 Merged-by: cann-robot Description: ## 描述 ### 背景 1. flash_attn host tiling 在 arch35/arch92 存在大量重复实现：`InitTilingInfo`/`IsCapable`/`CalcScheduleMode`/`DoOpTiling`/`SetPlatMemoryInfo`/`InitImplParam`/`UpdateTilingKeyLayout`/`UpdateTilingKeyKvLayout`/`GenTilingKey` 等公共流程两处各写一份，新增模板/配置需双改，易遗漏。 2. mask 处理存在多套模型：tiling data 与 kernel 结构体同时保留 `sparseMode`/`attenMaskBatch`/`attenMaskS1Size`/`isExistRowInvalid` 与 `maskMode`；`attenmask_gs1.h` 同时承载合轴（GS/SG）与不合轴（S1_ONLY）逻辑；`MaskMode`/`FAMaskMode`/`fa_base_vector::SparseMode` 枚举语义重叠，且 flash_attn kernel 直接引用 `fa_base_vector` 的 mask 枚举。 3. AICPU metadata 关键路径缺少维测日志，问题定位成本高。 ### 改动 1. Host tiling 归一（主要改动） - 新增 `op_host/flash_attn_tiling_base.h`：`FlashAttnTilingBase` 承接公共 tiling 流程（`DoOpTiling`/`SetPlatMemoryInfo`/`InitImplParam`/`GenTilingKey`/`UpdateTilingKeyLayout`/`UpdateTilingKeyKvLayout`/`FillTiling`/`CalcWorkspaceSize`/`CalcNumBlocks` 等）。 - `FlashAttnTilingImplArch35`/`FlashAttnTilingImplArch92` 继承基类，仅覆写差异部分： - arch35：`SplitPolicy`、`UpdateTilingKeyConfig`、`UpdateTilingKeyInfo`、`UpdateTilingKeyTemplateId`； - arch92：`SplitPolicy`、`UpdateTilingKeyConfig`、`UpdateTilingKeyInfo`。 - `fa_adjust_sinner_souter.h` 拆分为 `arch35/fa_adjust_sinner_souter_arch35.h` 与 `arch92/fa_adjust_sinner_souter_arch92.h`；去掉 `AdjustSinnerAndSouter` 未使用的 `qLayout` 形参及 `LAYOUT_`/`SOUTER_`/`DSIZE_*` 常量（改为数值直接判断）。 - `fa_tiling_info.h` 的 `MaskMode` 改为 `using MaskMode = FAMaskMode`，host 与 kernel 共用同一枚举定义。 2. mask 数据模型统一 - `FlashAttnAttenMaskParams`/`CommonConstInfo`：`sparseMode` → `maskMode`，删除 `attenMaskBatch`/`attenMaskS1Size`/`isExistRowInvalid`；host 侧不再推导 maskBatch/maskS1Size，仅保留 `attenMaskS2Size`。 - `attenmask_gs1.h` 拆分为： - `attenmask_m_gs1.h`：`MaskInfoMerged`（合轴 GS/SG），删除 `MaskInfo<MaskFormat>` 模板特化与 `S1_ONLY` 分支； - `attenmask_m_s1.h`：`MaskInfoUnmerged`（不合轴 S1_ONLY），承载 `IsSkipAttentionmask`/`ComputeAttenMaskOffsetCompressDn`/`AttentionmaskCopyInDn`。 - 与接口约束（`mask_mode` 仅 0/3/4，`attn_mask` 固定 2D `[2048,2048]`）对齐，删除不可达分支：`DEFAULT_MASK`/`ALL_MASK`/`LEFT_UP_CAUSAL`、`ComputeAttenMaskOffsetNoCompress` 及 batch mask 偏移、`MASK_BOOL`/`attenMaskType`/`maskValue`，以及 flashdecode 中 `DealInvalidMaskRows`/`negativeIntScalar_` 等无效行 mask 清理逻辑。 - DN 不合轴保持原有条件路由与 S1_ONLY mask 语义（`MASK_M_BASE`、`ProcessVec1VfDn` 路径不变），仅适配新 mask 模型，不改变业务功能。 - kernel 侧（kernel_dn/kernel_dn_unmerged/kernel_nd/flashdecode）统一使用 `constInfo_.maskMode` 判断，`batchNeedRowInvalid` 简化为仅依赖 `hasValidRow`。 3. mask 枚举归一与依赖整理 - `flash_attn_common_def.h` 定义唯一 `enum class FAMaskMode { NO_MASK = 0, CAUSAL = 3, BAND = 4 }`；host `MaskMode` 使用该定义。 - flash_attn kernel 内对 `fa_base_vector::RIGHT_DOWN_CAUSAL`/`BAND` 等 mask 枚举的引用替换为 `FAMaskMode::CAUSAL`/`BAND`；`GetSafeActToken` 等公共函数调用保留，但不再依赖 `fa_base_vector` 的 mask 枚举定义。 - `vector_common.h` 由真实使用方直接包含（flashdecode、kernel_nd、kernel_dn、kernel_dn_unmerged），移除 block_vec 文件的冗余传递依赖。 4. AICPU metadata 维测 - Compute 进入/成功、InitBaseInfo、InitLoadBalanceParams、BalanceSchedule、FlashDecode 开启处增加 info/debug 日志；同时移除未使用的 `qLayout` 计算。 ## 关联的Issue 关联 Issue #6956 ## 测试 - 代码走查：host tiling 路由/配置映射、DN 不合轴条件路由、S1_ONLY DN mask 与重构前逐项对齐，业务行为不变。 - 编译：ophost、ascend950 opkernel、flash_attn_metadata AICPU 已初步编译通过；mask 枚举归一改动后建议 CI 复核。 - 静态检查：行宽（<=120）/codespell/pre-commit 待补充。 - 算子验证：待补充（建议覆盖 arch35/arch92 的 NO_MASK/CAUSAL/BAND、DN 不合轴、PA/TND 等场景）。 - 二级冒烟：待补充。 ## 文档更新更新 `attention/flash_attn/README.md`：tiling 基类与 Dn/Nd 命名统一、`attenmask_m_gs1.h`/`attenmask_m_s1.h`、`DealInvalidMaskRows` 移除等描述同步。 ## 类型标签 - [ ] 🐛 Bug修复 - [ ] ✨ 新特性 - [ ] ⚡ 性能优化 - [x] ♻️ 重构 - [ ] 🧪 测试 - [ ] 📦 构建/CI - [ ] 🔧 配置变更 - [x] 📝 文档更新 - [ ] ⬆️ 依赖升级 - [ ] 🔒 安全修复 - [x] 🧹 代码清理 - [ ] ❓ 其他，请描述： See merge request: cann/ops-transformer!14865	10 days ago
attn_softmax_dn_bn1s1s2_vf.h	feat(fa): softmax VF 的 P 布局重构, 尾块按 actS2 裁剪列数 Co-authored-by: jiang-lirui<jianglirui1@huawei.com> # message auto-generated for no-merge-commit merge: !15117 merge flash_attn_softmax_vf_tail into master feat(fa): softmax VF 的 P 布局重构, 尾块按 actS2 裁剪列数 Created-by: jiang-lirui Commit-by: jiang-lirui Merged-by: cann-robot Description: ## 描述 <!--在这里详细描述你的改动，包括改动的原因和所采取的方法。--> 解决尾块（actS2 < tileCols）无论实际几列都要按满宽算一遍的问题 ## 关联的Issue <!-- 如果这个PR是为了解决特定的Issue，请在这里提供Issue链接。例如：关联Issue #000--> <!-- 如果这个PR是为了解决特定的问题单，请在这里描述问题单单号。--> Issue [#7031](https://gitcode.com/cann/ops-transformer/issues/7031) ## 测试 <!--描述进行了哪些测试来验证你的改动。包括但不限于二级冒烟、算子泛化等。--> 二级冒烟精度OK ## 文档更新 <!--如果这个PR包含文档的更新，请在这里指出。例如：更新了README.md文件。--> ## 类型标签 <!-- [x] 表示选中 --> - [ ] 🐛 Bug修复 - [ ] ✨ 新特性 - [x] ⚡ 性能优化 - [ ] ♻️ 重构 - [ ] 🧪 测试 - [ ] 📦 构建/CI - [ ] 🔧 配置变更 - [ ] 📝 文档更新 - [ ] ⬆️ 依赖升级 - [ ] 🔒 安全修复 - [ ] 🧹 代码清理 - [ ] ❓ 其他，请描述： See merge request: cann/ops-transformer!15117	12 days ago
flash_attn_common_def.h	flash_attn DN不合轴模板 s2BaseSize 切换为 256 Co-authored-by: jiang-lirui<jianglirui1@huawei.com> # message auto-generated for no-merge-commit merge: !15177 merge flash_attn_s2base256 into master flash_attn DN不合轴模板 s2BaseSize 切换为 256 Created-by: jiang-lirui Commit-by: jiang-lirui Merged-by: cann-robot Description: ## 描述 <!--在这里详细描述你的改动，包括改动的原因和所采取的方法。--> flash_attn DN不合轴模板 s2BaseSize 切换为 256 ## 关联的Issue <!-- 如果这个PR是为了解决特定的Issue，请在这里提供Issue链接。例如：关联Issue #000--> <!-- 如果这个PR是为了解决特定的问题单，请在这里描述问题单单号。--> Issue [#7137](https://gitcode.com/cann/ops-transformer/issues/7137) ## 测试 <!--描述进行了哪些测试来验证你的改动。包括但不限于二级冒烟、算子泛化等。--> ## 文档更新 <!--如果这个PR包含文档的更新，请在这里指出。例如：更新了README.md文件。--> ## 类型标签 <!-- [x] 表示选中 --> - [ ] 🐛 Bug修复 - [ ] ✨ 新特性 - [x] ⚡ 性能优化 - [ ] ♻️ 重构 - [ ] 🧪 测试 - [ ] 📦 构建/CI - [ ] 🔧 配置变更 - [ ] 📝 文档更新 - [ ] ⬆️ 依赖升级 - [ ] 🔒 安全修复 - [ ] 🧹 代码清理 - [ ] ❓ 其他，请描述： See merge request: cann/ops-transformer!15177	9 days ago
flash_attn_type.h	fix(flash_attn): 输出拷出显式指定ROW_PAD模板参数，d=72时为true其余为false Co-authored-by: linengyao<linengyao@huawei.com> Co-authored-by: XiyueL<lixiyue5@huawei.com> # message auto-generated for no-merge-commit merge: !13657 merge fixbug/rowpad into master fix(flash_attn): 输出拷出显式指定ROW_PAD模板参数，d=72时为true其余为false Created-by: linengyao Commit-by: XiyueL;linengyao Merged-by: cann-robot Description: ## 描述 flash_attn 输出拷出 `CopyAttenOutUbToGm::operator()` 的 `ROW_PAD` 模板参数此前依赖 `FaUbTensor` 默认模板参数推导，恒为 true，导致 S1G 场景 BSNGD/TNGD 布局下的多行合并拷贝优化路径不可达。本 PR 按业务规则显式指定 `ROW_PAD`：d=72 时为 true，其余为 false。 ## 改动 - `ConfigParams` 新增 `originalD` 字段（与 s1/s2/d/dv 同构），全档位填充原始 D 值；配置宏统一命名为 `Config_..._OriginalD<n>`（OriginalD64/128/192/256/72） - 新增 OriginalD72 克隆档（config=8/9）：档位内容与 D=128 档（config=2/3）完全一致，仅区分原始 D；`ConfigValueArch92` 补占位克隆（arch92 不支持 Dk=72，永不下发） - host tiling：Dk=72 仍按 sOuter 分档（切分逻辑不变），路由到 OriginalD72 档；dnFlag 同步加 config=8，保持 D=72 无 mask 场景仍走 DN 模板 - `flash_attn.cpp` 按 S1/S2/D/DV 同构方式解析 `cfg.originalD`（无特判函数），FAType 新增 `ORIGINAL_D` 模板参数透传，kernel 侧 `rowPad = (FA_T::originalD == regbaseutil::DTemplateType::Aligned72)` - 9 处调用点（arch35/arch92 共 6 个文件）显式指定模板实参：`copyAttenOutUbToGm_.template operator()<FaGmTensorOut, rowPad>(...)`；对应 `FaUbTensor` 构造同步使用 `<OUTPUT_T, rowPad>`，类型不一致时编译期报错 - d=128 等对齐场景行为不变；d=72 场景 ROW_PAD=true 走 padding 拷出逻辑；其余 d 启用 S1G+BSNGD/TNGD 多行合并拷贝优化 ## 关联的Issue #5978 ## 测试 - ascend950 全量打包编译验证通过（`bash build.sh --pkg --ops=flash_attn --soc=ascend950`） - ascend960dt 编译受 master 上 flash_attn arch92 tiling 两个既有问题阻塞（`flash_attn_template_tiling_key_arch92.h` 失效 include、双架构 tiling 常量重定义），与本次改动无关，已通过 stash 对比验证 ## 类型标签 - [x] 🐛 Bug 修复 - [ ] ✨ 新特性 - [ ] ⚡ 性能优化 - [ ] ♻️ 重构 - [ ] 🧪 测试 - [ ] 📦 构建/CI - [ ] 🔧 配置变更 - [ ] 📝 文档更新 - [ ] ⬆️ 依赖升级 - [ ] 🔒 安全修复 - [ ] 🧹 代码清理 - [ ] ❓ 其他，请描述： See merge request: cann/ops-transformer!13657	10 days ago