已合并
[flash_attn_metadata] 添加防御检查 #9078
[flash_attn_metadata] 添加防御检查 #9078
已合并
FiguraDoge创建于 7月23日
FiguraDoge成员
7月23日

描述

  1. 添加更多的防御,与下游算子
  2. 修复全空输入时,返回错误分核结果的问题

关联的Issue

测试

文档更新

类型标签

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 FiguraDoge 的贡献)
FFiguraDoge成员
7月23日 创建了 pull request,commit cbd2b50e
atomgit-bot
atomgit-bot
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 中新增 ParamsCheckCheckActualQuerySeq/CheckActualKvSeq 调用链,对 cuSeqlensQ/Kv 的单调性与首元素为 0、sequsedQ/Kv 的非负性进行校验;将原 LoadActualQuerySeq/LoadActualKvSeq 的序列处理逻辑分离并简化为成员变量直接赋值,使校验与初始化解耦。
likedislike
CANN-robotCANN-robot成员
7月23日 添加了label:stat/needs-squash
CANN-robotCANN-robot成员
7月23日 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
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 /approve or /lgtm
  • Commenting /approve implies 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. 👍

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

/lgtm

likedislike
tang-hao-hw-gitcode成员
7月24日 评论:

/approve

likedislike
赵志勇成员
7月24日 评论:

/approve

likedislike
CANN-robotCANN-robot成员
7月24日 添加了label:lgtm
CANN-robotCANN-robot成员
7月24日 合入了pull request