PyPTO Online Flash Attention 设计骨架 SK-01KV 分块在线 softmax 注意力算子的完整实战指南【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym本指南系统讲解 CANN / pypto-gym 仓库中cannbot-skills/ops/pypto-op-design/patterns/skeletons/SK-01-online-flash-attention.md这份设计骨架的全部内容它面向KV 序列分块迭代 online softmax 累积的注意力算子FA MHA / FA Score / BSA / PageAttn / SparseAttn 等给出 C1(QK^T) → V1(Softmax) → C2(PV) 的循环骨架、关键编码约束、变体矩阵、开箱性能配置与反向传播结构决策。读完本文你将能够在 PyPTO 算子设计阶段直接套用该骨架并结合仓库内 DeepSeek-V4 系列真实实现sparse_compress_flash_attention_impl.py、win_attention_impl.py、compress_flash_attention_impl.py完成配置选值与性能基线设计。一、适用场景与 CV 排布KV tile 循环内的 CVC 模式适用场景需要对 KV 序列分块迭代、使用 online softmax 累积的注意力计算。典型语义为QK^T → online softmax → PV覆盖 MHA / GQA / cross-attention / causal 变体。CV 排布C1(QK^T) → V1(Softmax) → C2(PV)在 KV tile 循环内重复形成CVC...循环模式。即每个 KV 分块都要经历一次 cube 矩阵乘算 scores、一次 vec 向量链online softmax 归一化、再一次 cube 矩阵乘加权求和循环结束后通过累积器O/L/M得到完整结果。展开因子候选为8、4、2、1初始设计每次只选一个值并验证不能整除时的余数处理其余值分别作为调优候选详见第五节变体矩阵与第七节性能导向粒度。从仓库实现看这一排布在 compress_flash_attention_impl.py 中得到了完整印证Sa_C1set_semantic_labelset_cube_tile_shapesmatmul→Sa_V1sg_set_scope2 amax/exp/sum/div 链→Sa_C2set_cube_tile_shapesmatmul逐阶段展开每个阶段前都切换对应的 TileShape与骨架的分阶段切换要求完全一致。二、骨架结构完整可复用的伪代码骨架给出了一个可直接落地的结构范本以下为原文骨架注释为便于理解补充的说明pypto.frontend.jit( pass_options{ cube_l1_reuse_setting: {-1: N}, # C1/C2 的 L1 复用设置见第五节 vec_nbuffer_setting: {-1: M}, # vec 双缓冲数 cube_nbuffer_setting: {-1: K}, # cube 双缓冲数 }, runtime_options{ stitch_function_max_num: 128, # 子图拼接上限 device_sched_mode: 1, # 设备侧并行调度 }, ) def flash_attention_kernel(Q, K, V, output, accumulators, seq_lens, ...): pypto.experimental.set_operation_options(combine_axisTrue) # 函数体首行必配 pypto.set_cube_tile_shapes(c1_global, k_global, c2_global) pypto.set_vec_tile_shapes(v1_global, v2_global) for b_idx in pypto.loop(batch_size): # Loop: Batch seq_q, seq_kv dynamic_seq_lens(...) # varlen 动态长度 q_tiles ceildiv(seq_q, Q_TILE) # Q 分块数 kv_tiles ceildiv(seq_kv, KV_TILE) # KV 分块数 for h_idx in pypto.loop(num_heads): # Loop: Head for q_idx in pypto.loop(q_tiles): # Loop: Q Tile oi pypto.tensor(...) # accumulators: 输出累积 li pypto.tensor(...) # accumulators: 分母累积 mi pypto.tensor(...) # accumulators: running max q_tile view_or_gather(Q, ...) # 按 Q tile 取数 for kv_idx in pypto.loop(kv_tiles): # Loop: KV Tile (C1-V1-C2) k_tile view_or_gather(K, ...) v_tile view_or_gather(V, ...) # C1: QK^T MatMul pypto.set_cube_tile_shapes(c1_tiles) # 阶段前必须切换 scores pypto.matmul(q_tile, k_tile, ...) # V1: Online Softmax pypto.set_vec_tile_shapes(v1_tiles) # 阶段前必须切换 ... online softmax computation ... # amax - sub - exp - sum # V: Quant (optional) # FP8 变体在此插入 P 量化 p_quant pypto.cast(p, ...) # C2: PV MatMul pypto.set_cube_tile_shapes(c2_tiles) # 阶段前必须切换 oij pypto.matmul(p_quant, v_tile, ...) # V: Online Accumulator Update (three-way branch) ... online softmax accumulator update ... # is_loop_begin/is_loop_end 三路分支结构要点三级循环嵌套Batch → Head → Q Tile最内层是 KV Tile 循环C1-V1-C2 的重复单元累积器分配位置O/L/M 三个累积器在 Q tile loop 内、KV tile loop 外分配保证跨 KV 分块累积动态边界seq_q/seq_kv来自动态seq_lensq_tiles/kv_tiles用ceildiv推导view valid_shape处理尾块边界量化钩子V1→C2 之间预留了可选的 P 量化位点。骨架中循环用pypto.loop表达动态循环当 KV 分块数编译期可确定且较小时可换用 Pythonrange静态展开见第七节性能导向形态注意与全动态pypto.loop的适用边界区分。该骨架与设计工作流的衔接方式见 pypto-op-design/SKILL.md先读 SK 索引 与 AT 索引选候选卡片再把结论写入 DESIGN.md。三、关键编码特征八条硬性规则特征规则TileShape 切换C1/V1/C2 各阶段前必须set_cube/vec_tile_shapes累积器位置在 Q tile loop 内、KV tile loop 外分配三路分支is_loop_beginis_loop_end组合判断数据加载viewvalid_shape处理边界数据存储assemble到全局输出张量FP32 累积O/L/M 全程 FP32仅最终 cast 输出 dtypePaged KV替换 K/V 加载为 AT-17 Block Table GatherFP8 量化在 V1→C2 之间插入 AT-08 (P 量化)C2 后插入 AT-06 (反量化)几条特征在仓库中的实证TileShape 切换compress_flash_attention_impl.py中 C1 用[[128,128],[256,256],[128,128]]级 cube tileV1 用[16, combine_s2_tile]级 vec tileC2 再切回 cube tileL121-L154。tiling 约束C-TILE-05要求矩阵乘 m/k/n 各轴使用[L0, L1]配置且0 L0 L1、L1 % L0 0vec 配置不能替代 cube 配置详见 constraints/tiling.md。Paged KVAT-17分页 KV 的拼装采用view 零搬运拼装而非显式 gather 指令范本见 AT-17-block-gather.mdsparse_compress_flash_attention_impl.py中通过pypto.viewassemble将物理分块的 KV 拼成连续kj再进入 C1L102-L119。softmax 链合V1 阶段用sg_set_scope2包裹 amax/sub/exp/sum/div 链、段尾复位-1见 sparse_compress_flash_attention_impl.py#L131-L144避免 vec 子图过小产生调度缝隙范本详见 AT-21-attention-scope.md。四、变体矩阵五种注意力的骨架落点变体Loop 4 策略MaskKV 来源量化FA MHApypto.loop无/因果连续 view无FA Scoreloop_unroll(unroll_list[4])掩码张量连续 view无BSApypto.loop稀疏掩码紧凑 gather无PageAttn FP8loop_unroll(unroll_list[8])无AT-17 分页 gatherFP8Sparse Attnpypto.looptopk 索引index_select无解读与实证FA MHA / FA Score最基础的连续 KV 形态区别仅在 KV 循环是否静态展开、是否有掩码张量BSA稀疏掩码 紧凑 gather适合局部窗口类注意力PageAttn FP8KV 来自分页缓存block_table 索引循环展开因子候选取 8、4、2、1V1→C2 之间走 FP8 量化AT-08C2 后反量化AT-06Sparse Attn用 topk 索引 index_select挑选稀疏 KV 块仓库对应实现见 sparse_compress_flash_attention_impl.py其 host 包装中unroll_list [4]见 L102-L106。注意设计工作流要求初始实现中的unroll_list只选择一个展开因子其余值作为调优候选见 SKILL.md 的输入与输出一节与骨架展开因子候选为 8、4、2、1的约定一致。五、开箱性能优化提示12 个配置维度全解析实证来源原文引用仓库内实际路径sparse_compress_flash_attention_impl.py:141-149、win_attention_impl.py:54-59、compress_flash_attention_impl.py:109-117维度推荐配置取值经验作用pass_options.cube_l1_reuse_setting必配{-1: 2~4}或{-1: 2, 0: 8}分阶段C1(QK^T) 与 C2(PV) 的权重/激活 L1 复用FA 的核心 cube 优化pass_options.cube_nbuffer_setting推荐{-1: 2}起步head_dim 大时升{1: 2}Cube 双缓冲掩盖 K/V 加载延迟pass_options.vec_nbuffer_setting必配{-1: 4}分页/稀疏可升{-1: 6}Online softmax 阶段向量算子多nbuffer 提升吞吐runtime_options.stitch_function_max_num必配128子图拼接上限FA 经验值runtime_options.device_sched_mode推荐1启用设备侧并行调度runtime_options.ready_on_host_tensors推荐[block_table, kv_act_seqs]等控制流读值 tensor控制流 tensor host 提前发射消除调度等待气泡paged/IFA 控制流必配pypto.experimental.set_operation_options(combine_axisTrue)必配在 jit 函数体首行尾轴 broadcast 内联 brcb如[M,N]*[M,1]、[M,N]/[M,1]见 F-15KV tile 循环展开按变体选择单值PageAttn FP8 候选8、4、2、1FA Score 候选4、2、1每次采用一个因子核对实际 KV 长度和余数处理TileShape 分阶段切换强制C1 / V1 / C2 各自set_*_tile_shapes不切换会导致 18000 表达式上限突破Online softmax 累积器 dtype强制全程 FP32mi/li/oi 必须 FP32最终仅 cast 输出Paged KV 模式配套AT-17 Block Gather view valid_shape避免单 block gather kernel 调用强制合图按场景softmax vec 链sg_set_scope2/ 状态更新1/ Cube 前后-1范本见 AT-21避免 Vector 图过小调度缝隙变得多三个仓库实现的配置对照可作复现起点compress_flash_attentionCFA连续/压缩 KVjit 装饰器 L111-L119 使用runtime_options{stitch_function_max_num: 128, device_sched_mode: 1}pass_options{cube_l1_reuse_setting: {-1: 3}, cube_nbuffer_setting: {-1: 2}, vec_nbuffer_setting: {-1: 4}}——正是本表必配/推荐档的典型取值Window Attention窗口注意力jit 装饰器 L55-L60 使用cube_l1_reuse_setting {0: 4}按阶段指定、cube_nbuffer_setting {1: 2}、vec_nbuffer_setting {-1: 4}Sparse Compress FA稀疏压缩注意力jit 装饰器 L157-L167 使用cube_l1_reuse_setting {-1: 2, 0: 8}分阶段、cube_nbuffer_setting {-1: 2}、vec_nbuffer_setting {-2: 1, -1: 8}vec 缓冲更激进因为分页/稀疏场景 vec 链更长、stitch_function_max_num 1024——可见稀疏变体在 vec 侧与子图上限上的取值更高。该骨架特有的性能方向C1-V1-C2 三阶段 nbuffer 解耦online softmax 累积器零冗余 cast。瓶颈通常出在 V1 的 amax/exp/sum 流水优先提升vec_nbuffer_settingKV 是稀疏/分页时优先配gather_in_ub。注意配置取值属于求稳范畴性能收益留待pypto-op-perf-tune调优阶段实测验证SKILL.md 中结构选择原则明确代码结构追求性能最优配置项求稳取参考值。六、结构性能特征调度等待主导时的诊断框架调度等待占总耗时较大、任务数量较多时检查每个任务处理的数据量及 root 数评估增大 tile 或减少计算趟数。等待源于跨 loop 累积器串行依赖且 head/batch 存在无依赖轴时按 AT-23 多链交错并行化AT-23-multi-chain-interleave.md。原因仍需结合性能报告确认。结构性能特征表| 维度 | 估算方式 | 检查内容 | ||---|---|---| | root 数 |batch × heads × ceildiv(seq_len_q, Q_TILE)| 结合设备并行度和调度耗时判断 | | 任务粒度 | 单 root 计算量 Q_TILE × seq_len_k × head_dim| 权衡任务开销、资源占用与并行度 | | 常驻 UB 状态 | 显式 bufferQ_TILE × head_dim × 4 Boi | ≤ UB 容量 | | 调度等待 | 从性能报告读取实际耗时 | 不仅凭 root 数预测等待比例 |root 数直接决定设备并行度root 过少则设备并行度不足过多则任务调度开销占比上升。UB 常驻状态预算须满足 tiling 约束C-TILE-03按目标设备可用 UB 与实际驻留张量估算容量计入中间结果、缓冲副本及数据类型。性能导向形态适用条件不满足时回退安全基线形态安全基线cube[128,128]×3、vec 各轴 [16,64]见 constraints/tiling.md。性能导向形态适用条件全部满足才可用任一不满足回退安全基线形态k 侧段长编译期可确定为统一常量varlen k 非均匀不适用head_dim ∈ {64, 128}root 数按上表 ∈ [16, 256]性能目标与基线差距 5×。结构不变量违反任一性能导向形态失效Q_TILE seq_len_q / 2每 (batch, head) 至多 2 个 Q 分块禁用固定小 tile 基线128/320直接套用k 侧循环静态展开k_tile_count ≤ 8 时必须用 Pythonrange禁用is_loop_begin/is_loop_end谓词状态 SSA 重绑定首分块oi oij禁止显式pypto.full常驻 buffer [:]写回否则 UB 溢出回退安全基线softmax 链合必须用sg_set_scope(6)包裹mul→amax→sub→exp→sum→cast链尾复位sg_set_scope(-1)l/m 直写assemble列写[total_q, num_heads]禁止 head-major 行写 转置桥。UB 预算强制设 tile 前列出常驻 buffer 字节数 ≤ UB 容量SSA 形态无额外常驻显式 buffer 形态按Q_TILE × head_dim × dtype_size计入。七、性能导向 Tile 粒度选型Performance handoff实测性能结论归pypto-op-perf-tune模块本骨架不持有实测定值以下为设计期的选型规则。适用条件算子语义 QK^T → 在线 softmax → PVMHA/GQA/cross-attention/causal 变体且性能目标与基线差距 5×。不满足时回退安全基线。Q_TILE 按 root 粒度反推目标 root 数 ∈ [16, 256]每 (batch, head) 至多 2 个 Q 分块Q_TILE seq_len_q / 2。禁用固定小 tile 基线128/320直接套用kv 循环静态化判据k_tile_count 编译期可确定且 ≤ 8 时k 侧循环必须用 Pythonrange静态展开禁用is_loop_begin/is_loop_end谓词状态 SSA 重绑定前提性能导向粒度必须配合 SSA 状态管理首分块oi oij禁止显式pypto.full常驻 buffer [:]写回否则 UB 溢出回退安全基线k 侧段长约束适用域 k 侧段长编译期可确定为统一常量varlen k 非均匀各 k 段不等不适用回退安全基线 动态 loopUB 预算强制设 tile 前列出常驻 buffer 字节数 ≤ UB 容量SSA 形态无额外常驻显式 buffer 形态按Q_TILE × head_dim × dtype_size计入任务数量使用第六节公式估算结合调度耗时评估 tile 大小及计算趟数。反向计算采用单趟还是多趟需同时满足依赖和精度要求。其中C-TILE-02尾轴 32B 对齐FP32 为 8 元素、FP16/BF16 为 16 元素、C-TILE-06TileShape 参数须为编译期整数不能来自运行时 shape、kernel 参数或 SymbolicScalar、C-TILE-08形状变化时重新核对 vec TileShape是落实上述选型时必须同步核对的硬约束全文见 constraints/tiling.md。八、反向传播结构决策单 Pass vs 双 Pass适用场景attention 反向传播算子输出 dQ/dK/dV。dQ 的最优循环方向为 q-outer/kv-innerdK/dV 为 kv-outer/q-inner该方向冲突决定结构选型必须在骨架匹配阶段显式决策并将结论记录到 DESIGN.md。结构循环组织梯度写回softmax 统计量QK^T 次数单 Pass默认推荐batch → head → q_tile → kv_tile单重嵌套三路梯度在同一块对内完成pypto.atomic_add直写 GM 输出host 侧torch.zeros预清零输入含 l/m 时直接消费零重算无则块内在线自算每 (q, kv) 块对 1 次双 Pass安全回退Pass1 q-outer 计算 dQPass2 kv-outer 计算 dK/dVUB 累积器 三路分支 assemble尾块写回Pass1 自算后经 GM scratch 传递至 Pass2每块对 2~3 次统计量趟/累加趟/Pass2 各一次决策规则默认选单 Passatomic_add消解循环方向冲突消除双 Pass 的 QK^T 重算与 GM scratch 往返FP32 输出与 varlen 动态 offset 均支持。单 Pass 的具体实现骨架见 SK-16-attention-backward.mdSingle-Pass Attention Backward——全动态四级pypto.loop l/m 直用 三路 atomic_add sg_set_scope 分段回退双 Pass 须写明合法理由平台不支持atomic_add、精度语义要求确定性归约顺序、或输出 buffer 不可预清零统计量来源联动签名含 l/m 统计量输入时禁止 kernel 内重算重算会为每个 q_tile 增加一趟完整 QK^T任务数预估任务数 ≈ batch × heads × ⌈sq/Q_TILE⌉ × ⌈skv/KV_TILE⌉ × Pass 趟数结合实际调度开销评估增大 tile 或减少计算趟数。SK-16 与该决策的衔接SK-16 适用条件is_backward true且输入含前向 l/m 统计量、matmul_count 5时骨架匹配阶段先做本节决策选单 Pass 即落入 SK-16。SK-16 的runtime_options与 SK-01 有系统性差异stitch_function_max_num需升至512、device_sched_mode用3、varlen 控制流 tensor 必须配ready_on_host_tensors、大 workspace 需配max_workspace_kbA3 平台实测25170624并禁止手写where掩码与 Pythonfor展开 batch/head见 SK-16 开箱性能优化提示。九、骨架落地流程与相关资源导航在设计工作流中见 pypto-op-design/SKILL.mdSK-01 的使用路径是对照 SPEC 与 golden 分析计算依赖确认算子语义属于QK^T → online softmax → PV家族从 SK 索引 与 AT 索引选出候选骨架SK-01与原子模式AT-17 / AT-21 / AT-08 / AT-06 / AT-23运行 estimate_decomposition.py 确定模块边界按 Tiling 约束 核对对齐与容量将 tiling、配置、布局量化、循环状态与数值稳定性决策写入 DESIGN.md模板见 DESIGN.md.tmpl用 validate_artifacts.py 做结构检查后交接给 develop 阶段实现。相关的仓库资源骨架与原子模式索引patterns/skeletons/index.md、patterns/atoms/index.md分页 KV 拼装与合图范本AT-17-block-gather.md、AT-21-attention-scope.md反向梯度骨架SK-16-attention-backward.md参考实现DeepSeek-V4 家族配置取值可作复现起点sparse_compress_flash_attention_impl.py、win_attention_impl.py、compress_flash_attention_impl.py相关测试tests/ops/deepseek_v4/下的test_sparse_compress_flash_attention.py、test_win_attention.py、test_compress_flash_attention.py可用于验证骨架语义与边界行为。总结SK-01 是 PyPTO 生态中 attention 家族算子设计的第一落点——前向用 C1-V1-C2 循环骨架配合 online softmax 三累积器与分阶段 TileShape 切换反向则依据第四节决策规则落到 SK-16 单 Pass 骨架配置上牢记cube_l1_reuse_setting必配、vec_nbuffer_setting优先、sg_set_scope分段合图、累积器全程 FP32四条主线即可在仓库参考实现的基础上快速产出可编译、可验证、可调优的 FA 系列算子设计。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
PyPTO Online Flash Attention 设计骨架 SK-01：KV 分块在线 softmax 注意力算子的完整实战指南
PyPTO Online Flash Attention 设计骨架 SK-01KV 分块在线 softmax 注意力算子的完整实战指南【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym本指南系统讲解 CANN / pypto-gym 仓库中cannbot-skills/ops/pypto-op-design/patterns/skeletons/SK-01-online-flash-attention.md这份设计骨架的全部内容它面向KV 序列分块迭代 online softmax 累积的注意力算子FA MHA / FA Score / BSA / PageAttn / SparseAttn 等给出 C1(QK^T) → V1(Softmax) → C2(PV) 的循环骨架、关键编码约束、变体矩阵、开箱性能配置与反向传播结构决策。读完本文你将能够在 PyPTO 算子设计阶段直接套用该骨架并结合仓库内 DeepSeek-V4 系列真实实现sparse_compress_flash_attention_impl.py、win_attention_impl.py、compress_flash_attention_impl.py完成配置选值与性能基线设计。一、适用场景与 CV 排布KV tile 循环内的 CVC 模式适用场景需要对 KV 序列分块迭代、使用 online softmax 累积的注意力计算。典型语义为QK^T → online softmax → PV覆盖 MHA / GQA / cross-attention / causal 变体。CV 排布C1(QK^T) → V1(Softmax) → C2(PV)在 KV tile 循环内重复形成CVC...循环模式。即每个 KV 分块都要经历一次 cube 矩阵乘算 scores、一次 vec 向量链online softmax 归一化、再一次 cube 矩阵乘加权求和循环结束后通过累积器O/L/M得到完整结果。展开因子候选为8、4、2、1初始设计每次只选一个值并验证不能整除时的余数处理其余值分别作为调优候选详见第五节变体矩阵与第七节性能导向粒度。从仓库实现看这一排布在 compress_flash_attention_impl.py 中得到了完整印证Sa_C1set_semantic_labelset_cube_tile_shapesmatmul→Sa_V1sg_set_scope2 amax/exp/sum/div 链→Sa_C2set_cube_tile_shapesmatmul逐阶段展开每个阶段前都切换对应的 TileShape与骨架的分阶段切换要求完全一致。二、骨架结构完整可复用的伪代码骨架给出了一个可直接落地的结构范本以下为原文骨架注释为便于理解补充的说明pypto.frontend.jit( pass_options{ cube_l1_reuse_setting: {-1: N}, # C1/C2 的 L1 复用设置见第五节 vec_nbuffer_setting: {-1: M}, # vec 双缓冲数 cube_nbuffer_setting: {-1: K}, # cube 双缓冲数 }, runtime_options{ stitch_function_max_num: 128, # 子图拼接上限 device_sched_mode: 1, # 设备侧并行调度 }, ) def flash_attention_kernel(Q, K, V, output, accumulators, seq_lens, ...): pypto.experimental.set_operation_options(combine_axisTrue) # 函数体首行必配 pypto.set_cube_tile_shapes(c1_global, k_global, c2_global) pypto.set_vec_tile_shapes(v1_global, v2_global) for b_idx in pypto.loop(batch_size): # Loop: Batch seq_q, seq_kv dynamic_seq_lens(...) # varlen 动态长度 q_tiles ceildiv(seq_q, Q_TILE) # Q 分块数 kv_tiles ceildiv(seq_kv, KV_TILE) # KV 分块数 for h_idx in pypto.loop(num_heads): # Loop: Head for q_idx in pypto.loop(q_tiles): # Loop: Q Tile oi pypto.tensor(...) # accumulators: 输出累积 li pypto.tensor(...) # accumulators: 分母累积 mi pypto.tensor(...) # accumulators: running max q_tile view_or_gather(Q, ...) # 按 Q tile 取数 for kv_idx in pypto.loop(kv_tiles): # Loop: KV Tile (C1-V1-C2) k_tile view_or_gather(K, ...) v_tile view_or_gather(V, ...) # C1: QK^T MatMul pypto.set_cube_tile_shapes(c1_tiles) # 阶段前必须切换 scores pypto.matmul(q_tile, k_tile, ...) # V1: Online Softmax pypto.set_vec_tile_shapes(v1_tiles) # 阶段前必须切换 ... online softmax computation ... # amax - sub - exp - sum # V: Quant (optional) # FP8 变体在此插入 P 量化 p_quant pypto.cast(p, ...) # C2: PV MatMul pypto.set_cube_tile_shapes(c2_tiles) # 阶段前必须切换 oij pypto.matmul(p_quant, v_tile, ...) # V: Online Accumulator Update (three-way branch) ... online softmax accumulator update ... # is_loop_begin/is_loop_end 三路分支结构要点三级循环嵌套Batch → Head → Q Tile最内层是 KV Tile 循环C1-V1-C2 的重复单元累积器分配位置O/L/M 三个累积器在 Q tile loop 内、KV tile loop 外分配保证跨 KV 分块累积动态边界seq_q/seq_kv来自动态seq_lensq_tiles/kv_tiles用ceildiv推导view valid_shape处理尾块边界量化钩子V1→C2 之间预留了可选的 P 量化位点。骨架中循环用pypto.loop表达动态循环当 KV 分块数编译期可确定且较小时可换用 Pythonrange静态展开见第七节性能导向形态注意与全动态pypto.loop的适用边界区分。该骨架与设计工作流的衔接方式见 pypto-op-design/SKILL.md先读 SK 索引 与 AT 索引选候选卡片再把结论写入 DESIGN.md。三、关键编码特征八条硬性规则特征规则TileShape 切换C1/V1/C2 各阶段前必须set_cube/vec_tile_shapes累积器位置在 Q tile loop 内、KV tile loop 外分配三路分支is_loop_beginis_loop_end组合判断数据加载viewvalid_shape处理边界数据存储assemble到全局输出张量FP32 累积O/L/M 全程 FP32仅最终 cast 输出 dtypePaged KV替换 K/V 加载为 AT-17 Block Table GatherFP8 量化在 V1→C2 之间插入 AT-08 (P 量化)C2 后插入 AT-06 (反量化)几条特征在仓库中的实证TileShape 切换compress_flash_attention_impl.py中 C1 用[[128,128],[256,256],[128,128]]级 cube tileV1 用[16, combine_s2_tile]级 vec tileC2 再切回 cube tileL121-L154。tiling 约束C-TILE-05要求矩阵乘 m/k/n 各轴使用[L0, L1]配置且0 L0 L1、L1 % L0 0vec 配置不能替代 cube 配置详见 constraints/tiling.md。Paged KVAT-17分页 KV 的拼装采用view 零搬运拼装而非显式 gather 指令范本见 AT-17-block-gather.mdsparse_compress_flash_attention_impl.py中通过pypto.viewassemble将物理分块的 KV 拼成连续kj再进入 C1L102-L119。softmax 链合V1 阶段用sg_set_scope2包裹 amax/sub/exp/sum/div 链、段尾复位-1见 sparse_compress_flash_attention_impl.py#L131-L144避免 vec 子图过小产生调度缝隙范本详见 AT-21-attention-scope.md。四、变体矩阵五种注意力的骨架落点变体Loop 4 策略MaskKV 来源量化FA MHApypto.loop无/因果连续 view无FA Scoreloop_unroll(unroll_list[4])掩码张量连续 view无BSApypto.loop稀疏掩码紧凑 gather无PageAttn FP8loop_unroll(unroll_list[8])无AT-17 分页 gatherFP8Sparse Attnpypto.looptopk 索引index_select无解读与实证FA MHA / FA Score最基础的连续 KV 形态区别仅在 KV 循环是否静态展开、是否有掩码张量BSA稀疏掩码 紧凑 gather适合局部窗口类注意力PageAttn FP8KV 来自分页缓存block_table 索引循环展开因子候选取 8、4、2、1V1→C2 之间走 FP8 量化AT-08C2 后反量化AT-06Sparse Attn用 topk 索引 index_select挑选稀疏 KV 块仓库对应实现见 sparse_compress_flash_attention_impl.py其 host 包装中unroll_list [4]见 L102-L106。注意设计工作流要求初始实现中的unroll_list只选择一个展开因子其余值作为调优候选见 SKILL.md 的输入与输出一节与骨架展开因子候选为 8、4、2、1的约定一致。五、开箱性能优化提示12 个配置维度全解析实证来源原文引用仓库内实际路径sparse_compress_flash_attention_impl.py:141-149、win_attention_impl.py:54-59、compress_flash_attention_impl.py:109-117维度推荐配置取值经验作用pass_options.cube_l1_reuse_setting必配{-1: 2~4}或{-1: 2, 0: 8}分阶段C1(QK^T) 与 C2(PV) 的权重/激活 L1 复用FA 的核心 cube 优化pass_options.cube_nbuffer_setting推荐{-1: 2}起步head_dim 大时升{1: 2}Cube 双缓冲掩盖 K/V 加载延迟pass_options.vec_nbuffer_setting必配{-1: 4}分页/稀疏可升{-1: 6}Online softmax 阶段向量算子多nbuffer 提升吞吐runtime_options.stitch_function_max_num必配128子图拼接上限FA 经验值runtime_options.device_sched_mode推荐1启用设备侧并行调度runtime_options.ready_on_host_tensors推荐[block_table, kv_act_seqs]等控制流读值 tensor控制流 tensor host 提前发射消除调度等待气泡paged/IFA 控制流必配pypto.experimental.set_operation_options(combine_axisTrue)必配在 jit 函数体首行尾轴 broadcast 内联 brcb如[M,N]*[M,1]、[M,N]/[M,1]见 F-15KV tile 循环展开按变体选择单值PageAttn FP8 候选8、4、2、1FA Score 候选4、2、1每次采用一个因子核对实际 KV 长度和余数处理TileShape 分阶段切换强制C1 / V1 / C2 各自set_*_tile_shapes不切换会导致 18000 表达式上限突破Online softmax 累积器 dtype强制全程 FP32mi/li/oi 必须 FP32最终仅 cast 输出Paged KV 模式配套AT-17 Block Gather view valid_shape避免单 block gather kernel 调用强制合图按场景softmax vec 链sg_set_scope2/ 状态更新1/ Cube 前后-1范本见 AT-21避免 Vector 图过小调度缝隙变得多三个仓库实现的配置对照可作复现起点compress_flash_attentionCFA连续/压缩 KVjit 装饰器 L111-L119 使用runtime_options{stitch_function_max_num: 128, device_sched_mode: 1}pass_options{cube_l1_reuse_setting: {-1: 3}, cube_nbuffer_setting: {-1: 2}, vec_nbuffer_setting: {-1: 4}}——正是本表必配/推荐档的典型取值Window Attention窗口注意力jit 装饰器 L55-L60 使用cube_l1_reuse_setting {0: 4}按阶段指定、cube_nbuffer_setting {1: 2}、vec_nbuffer_setting {-1: 4}Sparse Compress FA稀疏压缩注意力jit 装饰器 L157-L167 使用cube_l1_reuse_setting {-1: 2, 0: 8}分阶段、cube_nbuffer_setting {-1: 2}、vec_nbuffer_setting {-2: 1, -1: 8}vec 缓冲更激进因为分页/稀疏场景 vec 链更长、stitch_function_max_num 1024——可见稀疏变体在 vec 侧与子图上限上的取值更高。该骨架特有的性能方向C1-V1-C2 三阶段 nbuffer 解耦online softmax 累积器零冗余 cast。瓶颈通常出在 V1 的 amax/exp/sum 流水优先提升vec_nbuffer_settingKV 是稀疏/分页时优先配gather_in_ub。注意配置取值属于求稳范畴性能收益留待pypto-op-perf-tune调优阶段实测验证SKILL.md 中结构选择原则明确代码结构追求性能最优配置项求稳取参考值。六、结构性能特征调度等待主导时的诊断框架调度等待占总耗时较大、任务数量较多时检查每个任务处理的数据量及 root 数评估增大 tile 或减少计算趟数。等待源于跨 loop 累积器串行依赖且 head/batch 存在无依赖轴时按 AT-23 多链交错并行化AT-23-multi-chain-interleave.md。原因仍需结合性能报告确认。结构性能特征表| 维度 | 估算方式 | 检查内容 | ||---|---|---| | root 数 |batch × heads × ceildiv(seq_len_q, Q_TILE)| 结合设备并行度和调度耗时判断 | | 任务粒度 | 单 root 计算量 Q_TILE × seq_len_k × head_dim| 权衡任务开销、资源占用与并行度 | | 常驻 UB 状态 | 显式 bufferQ_TILE × head_dim × 4 Boi | ≤ UB 容量 | | 调度等待 | 从性能报告读取实际耗时 | 不仅凭 root 数预测等待比例 |root 数直接决定设备并行度root 过少则设备并行度不足过多则任务调度开销占比上升。UB 常驻状态预算须满足 tiling 约束C-TILE-03按目标设备可用 UB 与实际驻留张量估算容量计入中间结果、缓冲副本及数据类型。性能导向形态适用条件不满足时回退安全基线形态安全基线cube[128,128]×3、vec 各轴 [16,64]见 constraints/tiling.md。性能导向形态适用条件全部满足才可用任一不满足回退安全基线形态k 侧段长编译期可确定为统一常量varlen k 非均匀不适用head_dim ∈ {64, 128}root 数按上表 ∈ [16, 256]性能目标与基线差距 5×。结构不变量违反任一性能导向形态失效Q_TILE seq_len_q / 2每 (batch, head) 至多 2 个 Q 分块禁用固定小 tile 基线128/320直接套用k 侧循环静态展开k_tile_count ≤ 8 时必须用 Pythonrange禁用is_loop_begin/is_loop_end谓词状态 SSA 重绑定首分块oi oij禁止显式pypto.full常驻 buffer [:]写回否则 UB 溢出回退安全基线softmax 链合必须用sg_set_scope(6)包裹mul→amax→sub→exp→sum→cast链尾复位sg_set_scope(-1)l/m 直写assemble列写[total_q, num_heads]禁止 head-major 行写 转置桥。UB 预算强制设 tile 前列出常驻 buffer 字节数 ≤ UB 容量SSA 形态无额外常驻显式 buffer 形态按Q_TILE × head_dim × dtype_size计入。七、性能导向 Tile 粒度选型Performance handoff实测性能结论归pypto-op-perf-tune模块本骨架不持有实测定值以下为设计期的选型规则。适用条件算子语义 QK^T → 在线 softmax → PVMHA/GQA/cross-attention/causal 变体且性能目标与基线差距 5×。不满足时回退安全基线。Q_TILE 按 root 粒度反推目标 root 数 ∈ [16, 256]每 (batch, head) 至多 2 个 Q 分块Q_TILE seq_len_q / 2。禁用固定小 tile 基线128/320直接套用kv 循环静态化判据k_tile_count 编译期可确定且 ≤ 8 时k 侧循环必须用 Pythonrange静态展开禁用is_loop_begin/is_loop_end谓词状态 SSA 重绑定前提性能导向粒度必须配合 SSA 状态管理首分块oi oij禁止显式pypto.full常驻 buffer [:]写回否则 UB 溢出回退安全基线k 侧段长约束适用域 k 侧段长编译期可确定为统一常量varlen k 非均匀各 k 段不等不适用回退安全基线 动态 loopUB 预算强制设 tile 前列出常驻 buffer 字节数 ≤ UB 容量SSA 形态无额外常驻显式 buffer 形态按Q_TILE × head_dim × dtype_size计入任务数量使用第六节公式估算结合调度耗时评估 tile 大小及计算趟数。反向计算采用单趟还是多趟需同时满足依赖和精度要求。其中C-TILE-02尾轴 32B 对齐FP32 为 8 元素、FP16/BF16 为 16 元素、C-TILE-06TileShape 参数须为编译期整数不能来自运行时 shape、kernel 参数或 SymbolicScalar、C-TILE-08形状变化时重新核对 vec TileShape是落实上述选型时必须同步核对的硬约束全文见 constraints/tiling.md。八、反向传播结构决策单 Pass vs 双 Pass适用场景attention 反向传播算子输出 dQ/dK/dV。dQ 的最优循环方向为 q-outer/kv-innerdK/dV 为 kv-outer/q-inner该方向冲突决定结构选型必须在骨架匹配阶段显式决策并将结论记录到 DESIGN.md。结构循环组织梯度写回softmax 统计量QK^T 次数单 Pass默认推荐batch → head → q_tile → kv_tile单重嵌套三路梯度在同一块对内完成pypto.atomic_add直写 GM 输出host 侧torch.zeros预清零输入含 l/m 时直接消费零重算无则块内在线自算每 (q, kv) 块对 1 次双 Pass安全回退Pass1 q-outer 计算 dQPass2 kv-outer 计算 dK/dVUB 累积器 三路分支 assemble尾块写回Pass1 自算后经 GM scratch 传递至 Pass2每块对 2~3 次统计量趟/累加趟/Pass2 各一次决策规则默认选单 Passatomic_add消解循环方向冲突消除双 Pass 的 QK^T 重算与 GM scratch 往返FP32 输出与 varlen 动态 offset 均支持。单 Pass 的具体实现骨架见 SK-16-attention-backward.mdSingle-Pass Attention Backward——全动态四级pypto.loop l/m 直用 三路 atomic_add sg_set_scope 分段回退双 Pass 须写明合法理由平台不支持atomic_add、精度语义要求确定性归约顺序、或输出 buffer 不可预清零统计量来源联动签名含 l/m 统计量输入时禁止 kernel 内重算重算会为每个 q_tile 增加一趟完整 QK^T任务数预估任务数 ≈ batch × heads × ⌈sq/Q_TILE⌉ × ⌈skv/KV_TILE⌉ × Pass 趟数结合实际调度开销评估增大 tile 或减少计算趟数。SK-16 与该决策的衔接SK-16 适用条件is_backward true且输入含前向 l/m 统计量、matmul_count 5时骨架匹配阶段先做本节决策选单 Pass 即落入 SK-16。SK-16 的runtime_options与 SK-01 有系统性差异stitch_function_max_num需升至512、device_sched_mode用3、varlen 控制流 tensor 必须配ready_on_host_tensors、大 workspace 需配max_workspace_kbA3 平台实测25170624并禁止手写where掩码与 Pythonfor展开 batch/head见 SK-16 开箱性能优化提示。九、骨架落地流程与相关资源导航在设计工作流中见 pypto-op-design/SKILL.mdSK-01 的使用路径是对照 SPEC 与 golden 分析计算依赖确认算子语义属于QK^T → online softmax → PV家族从 SK 索引 与 AT 索引选出候选骨架SK-01与原子模式AT-17 / AT-21 / AT-08 / AT-06 / AT-23运行 estimate_decomposition.py 确定模块边界按 Tiling 约束 核对对齐与容量将 tiling、配置、布局量化、循环状态与数值稳定性决策写入 DESIGN.md模板见 DESIGN.md.tmpl用 validate_artifacts.py 做结构检查后交接给 develop 阶段实现。相关的仓库资源骨架与原子模式索引patterns/skeletons/index.md、patterns/atoms/index.md分页 KV 拼装与合图范本AT-17-block-gather.md、AT-21-attention-scope.md反向梯度骨架SK-16-attention-backward.md参考实现DeepSeek-V4 家族配置取值可作复现起点sparse_compress_flash_attention_impl.py、win_attention_impl.py、compress_flash_attention_impl.py相关测试tests/ops/deepseek_v4/下的test_sparse_compress_flash_attention.py、test_win_attention.py、test_compress_flash_attention.py可用于验证骨架语义与边界行为。总结SK-01 是 PyPTO 生态中 attention 家族算子设计的第一落点——前向用 C1-V1-C2 循环骨架配合 online softmax 三累积器与分阶段 TileShape 切换反向则依据第四节决策规则落到 SK-16 单 Pass 骨架配置上牢记cube_l1_reuse_setting必配、vec_nbuffer_setting优先、sg_set_scope分段合图、累积器全程 FP32四条主线即可在仓库参考实现的基础上快速产出可编译、可验证、可调优的 FA 系列算子设计。【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考