CANN/cannbot-skills Flash Attention内核实现

Deep Note:agent/example/kernels/a2/flash_attn_full_pj_hif8.py

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体，本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills

Open this file only after the short catalog entry confirmed the kernel is relevant.

What this kernel is really for

• the scaled-hif8 probability variant of normalized a2 online softmax
• a contract that intentionally changes the delayed value path while keepingrow_sumin float
• a kernel that exports finalrowmax/rowsumas part of the visible contract

Decisions worth copying

• updaterow_sumfrom the floatp_jtile before anyhalf/ hif8 cast
• keep stage-1 score scratch and stage-2pvscratch separate in the readable baseline
• implement the non-negative hif8 simulation without relying on unsupporteduint8 -> floatshortcuts
• copy[M,64]score slices into contiguous scratch beforereinterpret(...)when the quantized helper needs contiguous lanes
• handle non-alignedS2in score space with suffix invalidation and a sufficiently negative finite sentinel
• handle non-alignedS1separately fromS2; invalid rows should become zero contribution to delayedp @ vwhile GM still writes only valid rows

Prefer another kernel when

• you still want the plainp.half().float()value path
• you are debugging the normalized float/half baseline before introducing hif8 behavior

【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体，本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills