已合并
[flash_attn_metadata] 添加防御检查 #9078
FiguraDoge创建于 7月23日
[flash_attn_metadata] 添加防御检查 #9078
已合并
Pull Request已成功合入, 合并人@CANN-robot
(感谢 FiguraDoge 的贡献)7月23日 创建了 pull request,commit cbd2b50e
atomgit-bot
7月23日 评论:
7月23日 评论:
变更摘要
本 PR 主要对 flash_attn_metadata 算子进行了防御性加固:一是在 Host 侧 FlashAttnMetadataCheck 中加强了参数范围校验与 PA 布局支持;二是在 AICPU 内核侧新增 ParamsCheck 独立校验流程,将 seq 数据的验证与实际序列载入分离;三是修复了全空输入时分核结果 bN2End 错误调用 NumToIndex 的 bug。
主要改动
- 修复全空输入分核结果错误:
SectionStreamKImpl::Compute中空算例分支的bN2End[0]赋值从NumToIndex(...)改为直接使用原始 batch 数乘 query 头数,避免对空序列进行索引转换导致的错误结果。 - 加强 BaseAttr 参数范围校验:
CheckBaseAttr中引入NONE_VALUE = -1常量,将batchSize校验从== -1 || > 0改为区间(0, 65536)的严格范围检查,并将maxSeqlenQ/maxSeqlenKv从> 0放宽为>= 0;同时修复了numHeadsKv错误消息中误用maxSeqlenKv变量的 bug。 - 新增 Mask 模式的差异化校验:
CheckMask中针对NO_MASK/CAUSAL_MASK强制要求winLeft/winRight必须为NONE_VALUE,对WINDOW_MASK则要求两者>= NONE_VALUE,避免无效窗口参数传入。 - 增加 PA 布局的 sequsedKv 必传校验与 Existency 零值边界修正:
CheckExistency中新增IsPA()判断,当layoutKv为 PA 格式时必须传入sequsedKvOptional;并将非 TND 布局下maxSeqlenQ/maxSeqlenKv的判据从>= 0改为> 0,避免零值绕过。 - AICPU 内核新增独立参数检查流程: 在
FlashAttnMetadataCpuKernel中新增ParamsCheck→CheckActualQuerySeq/CheckActualKvSeq调用链,对cuSeqlensQ/Kv的单调性与首元素为 0、sequsedQ/Kv的非负性进行校验;将原LoadActualQuerySeq/LoadActualKvSeq的序列处理逻辑分离并简化为成员变量直接赋值,使校验与初始化解耦。


7月23日 添加了label:stat/needs-squash
7月23日 添加了label:cann-cla/yes
CANN-robot
7月23日 评论:
7月23日 评论:
Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here
PR Approval Progress
✅ Congratulations! All modules have met the lgtm and approve requirements.
Module Approval Details
| module | lgtm status | approve status |
|---|---|---|
| attention | ✅ shen_weiling, haijie_699874, tang-hao-hw-gitcode, zzy__ (4/2) | ✅ haijie_699874, tang-hao-hw-gitcode, zzy__ (3/1) |
| attention/common/op_kernel/load_balance | ✅ zzy__, haijie_699874 (2/2) | ✅ zzy__, haijie_699874 (2/1) |
💡 Tip:
- Committer can comment
/approveor/lgtm- Commenting
/approveimplies both code review (lgtm) and intent to merge (approve)
CLA Signature Pass
ChengjiaWu, thanks for your pull request. All authors of the commits have signed the CLA. 👍


此处折叠了47条消息 查看更多
shen_weiling
7月24日 评论:
7月24日 评论:
/lgtm


tang-hao-hw-gitcode
7月24日 评论:
7月24日 评论:
/approve


7月24日 添加了label:lgtm
7月24日 合入了pull request
描述
关联的Issue
测试
文档更新
类型标签