feat: Add template definitions for JD.com's core retail operations.-catlass-AtomGit

Adocs/zh/2_Design/04_contrib/00_XAttention/00_summary.md+58-0

@@ -0,0 +1,58 @@
	1	+# Flash Attention 新增模板总结
	2	+
	3	+本目录（04_contrib/00_XAttention）收录个人仓在 CATLASS 主干仓基础上新增的 23 个 Flash Attention 相关模板（10 个 GEMM + 13 个 Epilogue）的设计文档。23 个模板按算子形态划分为五个系列，每个系列配备一篇设计文档。
	4	+
	5	+## 系列文档索引
	6	+
	7	+\| 文档 \| 系列 \| 场景 \| 新增模板 \|
	8	+\| --- \| --- \| --- \| --- \|
	9	+\| [01_fa_unshared_kernel.md](./01_fa_unshared_kernel.md) \| FA Unshared \| 非 shared expert（unshared）场景的 Flash Attention，MHA/GQA 基础形态 \| 3 个 \|
	10	+\| [02_fai_split_row_kernel.md](./02_fai_split_row_kernel.md) \| FAI SplitRow \| 推理 Flash Attention，按行切分免 Combine 的 split-KV 变体 \| 4 个 \|
	11	+\| [03_xfai_kernel.md](./03_xfai_kernel.md) \| XFAI \| Atlas A2 推理 Flash Attention，AIC QK→AIV Softmax→AIC PV 跨核流水，支持 Paged KV \| 5 个 \|
	12	+\| [04_fd_kernel.md](./04_fd_kernel.md) \| FD \| XFAI 的深度性能演化版（FlashAttention-Decode 形态），含通用 CombineScale \| 3 个 \|
	13	+\| [05_xa_tla_kernel.md](./05_xa_tla_kernel.md) \| XA TLA \| Atlas A5（Ascend950）TLA 指令版 Flash Attention，shared/unshared 双形态 \| 8 个 \|
	14	+
	15	+## 新增模板清单
	16	+
	17	+### GEMM 侧（10 个，`include/catlass/gemm/block/`）
	18	+
	19	+\| 模板文件 \| 所属系列 \|
	20	+\| --- \| --- \|
	21	+\| `block_mmad_unshared_fa_qk.hpp` \| FA Unshared \|
	22	+\| `block_mmad_unshared_fa_pv.hpp` \| FA Unshared \|
	23	+\| `block_mmad_fai_qk_split_row.hpp` \| FAI SplitRow \|
	24	+\| `block_mmad_fai_pv_split_row.hpp` \| FAI SplitRow \|
	25	+\| `block_mmad_xfai_qk.hpp` \| XFAI \|
	26	+\| `block_mmad_xfai_pv.hpp` \| XFAI \|
	27	+\| `block_mmad_xa_shared_qk_tla.hpp` \| XA TLA（shared） \|
	28	+\| `block_mmad_xa_shared_pv_tla.hpp` \| XA TLA（shared） \|
	29	+\| `block_mmad_xa_unshared_qk_tla.hpp` \| XA TLA（unshared） \|
	30	+\| `block_mmad_xa_unshared_pv_tla.hpp` \| XA TLA（unshared） \|
	31	+
	32	+### Epilogue 侧（13 个，`include/catlass/epilogue/block/`）
	33	+
	34	+\| 模板文件 \| 所属系列 \|
	35	+\| --- \| --- \|
	36	+\| `block_epilogue_fa_unshared_softmax.hpp` \| FA Unshared \|
	37	+\| `block_epilogue_online_softmax_copy_glm.hpp` \| FAI SplitRow \|
	38	+\| `block_epilogue_rescale_o_no_split_row.hpp` \| FAI SplitRow \|
	39	+\| `block_epilogue_xfai_online_softmax.hpp` \| XFAI \|
	40	+\| `block_epilogue_xfai_rescale_o.hpp` \| XFAI \|
	41	+\| `block_epilogue_xfai_combine_scale.hpp` \| XFAI \|
	42	+\| `block_epilogue_online_softmax_FD.hpp` \| FD \|
	43	+\| `block_epilogue_rescale_o_FD.hpp` \| FD \|
	44	+\| `block_epilogue_combine_scale.hpp` \| FD \|
	45	+\| `block_epilogue_xa_shared_softmax_ascend950.hpp` \| XA TLA \|
	46	+\| `block_epilogue_xa_unshared_softmax_ascend950.hpp` \| XA TLA \|
	47	+\| `block_epilogue_xa_shared_rescale_ascend950.hpp` \| XA TLA \|
	48	+\| `block_epilogue_xa_combine_scale_ascend950.hpp` \| XA TLA \|
	49	+
	50	+## 系列速查
	51	+
	52	+- 01 FA Unshared：QK/PV 两个 `BlockMmad` + 一个 UnsharedSoftmax `BlockEpilogue`，结构最简的入门形态；
	53	+- 02 FAI SplitRow：按行切分的 split-KV 免合并变体（`CopySumMax` 记录行统计量替代 Combine 步骤）；
	54	+- 03 XFAI：AIC/AIV 跨核流水"QK→OnlineSoftmax→PV→RescaleO"，`l1BufAddrStart` 共享 L1，`PAGED_CACHE_FLAG` 原生分页 KV，`CombineScale` 收口 split-KV 合并；
	55	+- 04 FD：XFAI 的深度性能演化版，dm 按周期分区等优化，`CombineScale` 通用化为 `EpilogueAtlasA2CombineScale`；
	56	+- 05 XA TLA：面向 Atlas A5（Ascend950）的 TLA 指令实现，shared/unshared 两种 expert 形态各配 QK/PV/Softmax 模板。
	57	+
	58	+各系列的 DispatchPolicy 定义、算法设计（数据布局/流水组织/跨核同步）、系列间差异对比与基于 xllm-ops（https://gitcode.com/xLLM-AI/xllm_ops，`x_attention` / `x_flash_attention_infer`）的真实工程使用示例，详见对应设计文档。

Adocs/zh/2_Design/04_contrib/00_XAttention/01_fa_unshared_kernel.md+208-0

@@ -0,0 +1,208 @@
	1	+# Unshared FA 推理系列模板设计
	2	+
	3	+本系列模板面向 Unshared（非共享 KV）FlashAttention 推理场景：每个 attention head 的 KV 有效序列长度各不相同（如 Multi-LoRA / 独立上下文场景），因此 softmax 阶段需要按 head 构造独立的 mask，且 KV 序列较短、可一次性处理完，无需在线 rescale 累积输出 O。
	4	+
	5	+系列包含以下模板：
	6	+
	7	+\| 模板 \| 层级 \| Policy \| 源文件 \|
	8	+\| --- \| --- \| --- \| --- \|
	9	+\| BlockMmad（QK） \| GEMM \| `Gemm::MmadAtlasA2UnsharedFAQK` \| `include/catlass/gemm/block/block_mmad_unshared_fa_qk.hpp` \|
	10	+\| BlockMmad（PV） \| GEMM \| `Gemm::MmadAtlasA2UnsharedFAPV` \| `include/catlass/gemm/block/block_mmad_unshared_fa_pv.hpp` \|
	11	+\| BlockEpilogue（Softmax） \| Epilogue \| `Epilogue::EpilogueAtlasA2FAUnsharedSoftmax` / `EpilogueAscend950FAUnsharedSoftmax` \| `include/catlass/epilogue/block/block_epilogue_fa_unshared_softmax.hpp` \|
	12	+
	13	+三者的协作关系：
	14	+
	15	+```
	16	+Q @ K^T ──► (S) ──► FAUnsharedSoftmax ──► P(f16/bf16), gm(rowMax), gl(rowSum)
	17	+ UnsharedFAQK │
	18	+ └─(跨核通知 softmaxReady)─► P @ V ──► OTmp(f32) ──► 外部 combine: O = OTmp / gl
	19	+ UnsharedFAPV
	20	+```
	21	+
	22	+## 1. DispatchPolicy 定义
	23	+
	24	+`include/catlass/gemm/dispatch_policy.hpp`：
	25	+
	26	+```cpp
	27	+struct MmadAtlasA2UnsharedFAQK : public MmadAtlasA2 {
	28	+ static constexpr uint32_t STAGES = 2;
	29	+};
	30	+
	31	+struct MmadAtlasA2UnsharedFAPV : public MmadAtlasA2 {
	32	+ static constexpr uint32_t STAGES = 2;
	33	+};
	34	+```
	35	+
	36	+`include/catlass/epilogue/dispatch_policy.hpp`：
	37	+
	38	+```cpp
	39	+struct EpilogueAtlasA2FAUnsharedSoftmax {
	40	+ using ArchTag = Arch::AtlasA2;
	41	+};
	42	+
	43	+struct EpilogueAscend950FAUnsharedSoftmax {
	44	+ using ArchTag = Arch::Ascend950;
	45	+};
	46	+```
	47	+
	48	+## 2. BlockMmad MmadAtlasA2UnsharedFAQK 设计方案
	49	+
	50	+### 2.1 模板参数
	51	+
	52	+```cpp
	53	+template <class L1TileShape_, class L0TileShape_, class AType_, class BType_, class CType_,
	54	+ class BiasType_, class TileCopy_, class TileMmad_>
	55	+struct BlockMmad<MmadAtlasA2UnsharedFAQK, L1TileShape_, L0TileShape_, AType_, BType_, CType_,
	56	+ BiasType_, TileCopy_, TileMmad_>;
	57	+```
	58	+
	59	+- `L1TileShape_`：L1 级 tile 形状，QK 场景典型值 `GemmShape<128, 256, 128>`（N 为单次搬运的 KV tile 长度）。
	60	+- `L0TileShape_`：L0 级 tile 形状，一般与 L1 相同。
	61	+- `AType_`/`BType_`：Q（RowMajor）与 K（ColumnMajor）的 `GemmType<Element, Layout>`。
	62	+- `CType_`：输出 S 矩阵，仅支持 RowMajor（`static_assert` 约束）。
	63	+
	64	+### 2.2 内存布局
	65	+
	66	+- L1：`l1A` 与 `l1B` 连续分配（`l1B` 起始偏移为 `L1A_SIZE = MKsizeof(ElementA)`）。
	67	+- L0A/L0B/L0C：按 `STAGES = 2` 计算 pingpong buffer 尺寸（本模板单次调用完成一个 block，实际按单缓冲使用）。
	68	+- L0C 中 S 的布局为 `layout::zN`（MMAD 原生输出布局）。
	69	+
	70	+### 2.3 事件同步与执行流程
	71	+
	72	+构造函数完成缓冲初始化，并预置三组硬件事件（`EVENT_ID0`）：`MTE1_MTE2`、`M_MTE1`、`FIX_M`；析构函数对称 Wait，保证 Kernel 退出时流水排空。
	73	+
	74	+`operator()` 单次调用内完成一个 block 的完整搬运与计算：
	75	+
	76	+1. `WaitFlag(MTE1_MTE2)` → GM→L1 搬运 A（Q）→ `SetFlag(MTE2_MTE1)` → Wait；
	77	+2. `WaitFlag(M_MTE1)` → L1→L0A（`copyL1ToL0A`）；
	78	+3. GM→L1 搬运 B（K^T）→ L1→L0B；
	79	+4. `tileMmad(l0C, l0A, l0B, mRound, nRound, actualShape.k())` 执行 MMAD；
	80	+5. FIX 通路：`copyL0CToGm` 将 S 写回 GM（供 softmax epilogue 读取）。
	81	+
	82	+m/n/k 维度均按 `L1AlignHelper` 对齐规则 `RoundUp` 后参与 L0 布局，`tileMmad` 的 K 用真实值 `actualShape.k()`。
	83	+
	84	+## 3. BlockMmad MmadAtlasA2UnsharedFAPV 设计方案
	85	+
	86	+### 3.1 与 FAQK 的差异
	87	+
	88	+- 数据通路相同（GM→L1→L0→MMAD→L0C→GM），差异在于 A（即 P 矩阵）的搬运时机由跨核同步控制。
	89	+
	90	+### 3.2 跨核同步设计
	91	+
	92	+```cpp
	93	+void operator()(..., GemmCoord actualShape, Arch::CrossCoreFlag softmaxReady)
	94	+{
	95	+ // 1. 先搬运 B（V）：无需依赖 softmax 结果
	96	+ copyGmToL1B(l1BTensor, gB, layoutBInL1, layoutTileB);
	97	+ ...
	98	+ copyL1ToL0B(...);
	99	+
	100	+ // 2. 跨核等待 softmax 完成通知（P 已写出 GM）
	101	+ Arch::CrossCoreWaitFlag(softmaxReady);
	102	+ copyGmToL1A(l1ATensor, gA, layoutAInL1, layoutTileA); // gA 即 P
	103	+ ...
	104	+}
	105	+```
	106	+
	107	+PV GEMM 与 softmax epilogue 通常运行在不同核（或不同 subBlock）上：V 的搬运与 softmax 计算重叠执行，等 softmax 发出 `softmaxReady` 通知后才开始搬 P，从而隐藏 softmax 时延。构造/析构仅维护 `MTE1_MTE2 (EVENT_ID2)` 一组事件，QK 阶段使用的 `EVENT_ID0` 组保留给 MMAD/FIX 流水。
	108	+
	109	+PV 场景典型 L1TileShape 为 `GemmShape<128, 128, 256>`（K 维为 KV 序列长度方向，单 tile 内完成整个短序列）。
	110	+
	111	+## 4. BlockEpilogue EpilogueAtlasA2FAUnsharedSoftmax 设计方案
	112	+
	113	+### 4.1 模板参数与构造
	114	+
	115	+```cpp
	116	+BlockEpilogue(Arch::Resource<ArchTag>& resource, float tor_,
	117	+ uint32_t unsharedKvSeqLen, uint32_t maxDecodeStep,
	118	+ uint32_t headNum, uint32_t groupSize);
	119	+```
	120	+
	121	+- `OutputType_`：P 矩阵类型（fp16/bfloat16）。
	122	+- `InputType_`：S 矩阵类型（float）。
	123	+- `MaskType_`：mask 类型。
	124	+- 构造参数：softmax scale `tor`、每个 head 的 KV 有效长度 `unsharedKvSeqLen`、最大 decode 步数、head 数与 GQA group 大小。
	125	+
	126	+### 4.2 UB 内存布局
	127	+
	128	+\| Tensor \| 元素类型 \| 起始偏移（字节） \| 用途 \|
	129	+\| --- \| --- \| --- \| --- \|
	130	+\| lsUbTensor \| float \| 0 \| S 分数矩阵（加 scale 后） \|
	131	+\| lpUbTensor32 / tvUbTensor16 \| float / Output \| `2 * 32768` \| P（cast 后）复用区 \|
	132	+\| lmUbTensor \| float \| `3 * 32768` \| 行最大值 rowMax \|
	133	+\| llUbTensor \| float \| `3 * 32768 + 4 * 512` \| 行和 rowSum \|
	134	+\| tvUbTensor \| float \| `3 * 32768 + 8 * 512` \| 临时向量（Brcb 展开等） \|
	135	+\| unsharedMaskUbTensor \| float \| `3 * 32768 + 12 * 512` \| 按 head 构造的加法 mask \|
	136	+
	137	+### 4.3 Unshared mask 的设备侧构造（InitUnsharedMaskV2）
	138	+
	139	+mask 尺寸为 `[headNum * groupSize, kSeqTileRound]`（`kSeqTileRound = ceil(maxDecodeStepheadNum / 8) 8`）。构造逻辑：
	140	+
	141	+1. 先整体填充 `lowest()`（负无穷，经 Add 加到 S 上等效屏蔽）；
	142	+2. 对每个 head（`colOffset = (headOffset + round) * maxDecodeStep`），将 `[colOffset, colOffset + unsharedKvSeqLen)` 区间填充 0（等效不屏蔽）；
	143	+3. 考虑 8 元素 block 对齐（`FLOAT_BLOCK_SIZE`），未对齐尾部用 0/lowest 二次修正；
	144	+4. 按 subBlock（双核）将 head 数一分为二，各自构造本核负责的行段。
	145	+
	146	+由于每个 head 的有效区间起点随 `round * maxDecodeStep` 平移，不同 head 的 mask 不同——这正是 "Unshared" 的含义。
	147	+
	148	+### 4.4 算法流程（SubCoreCompute）
	149	+
	150	+```
	151	+S = DataCopy(GM) // MTE2, EVENT_ID3
	152	+S = S * tor // Muls scale
	153	+S = S + unsharedMask // 加法 mask（越界位置变 -inf）
	154	+lm = ReduceMax(S, row) // WholeReduceMax + Max 折叠（>128 列时）
	155	+S = S - broadcast(lm) // Brcb 展开行最大值后逐行相减
	156	+ls = Exp(S)
	157	+lp = Cast(ls, f32 -> f16/bf16) // P 矩阵
	158	+ll = ReduceSum(ls, row) // WholeReduceSum + Add 折叠
	159	+写出: P -> GM（DataCopy, T_BLOCK_SIZE=16 对齐）
	160	+ gm(lm), gl(ll) -> GM（按 head 偏移，非 8 对齐时用 DataCopyPad）
	161	+```
	162	+
	163	+由于 KV 序列一次处理完，本模板不做 O 的 rescale 与归一化，`gm/gl` 交由外部（上层 kernel 或 combine kernel）完成 `O = OTmp / gl`。双 subBlock 场景下按 head 切分行区间并行计算，事件全部使用 `EVENT_ID3` 避免与 GEMM 侧冲突。
	164	+
	165	+## 5. Ascend950 特化
	166	+
	167	+`EpilogueAscend950FAUnsharedSoftmax` 直接继承 AtlasA2 实现（Ascend C 向量 API 兼容）：
	168	+
	169	+```cpp
	170	+template <class OutputType_, class InputType_, class MaskType_>
	171	+class BlockEpilogue<EpilogueAscend950FAUnsharedSoftmax, OutputType_, InputType_, MaskType_>
	172	+ : public BlockEpilogue<EpilogueAtlasA2FAUnsharedSoftmax, OutputType_, InputType_, MaskType_> {
	173	+ using Base::Base;
	174	+};
	175	+```
	176	+
	177	+## 6. 使用示例
	178	+
	179	+摘自 xllm-ops（https://gitcode.com/xLLM-AI/xllm_ops）`x_attention/op_kernel/x_attention_catlass_helper.h` 的 `CallUnsharedInferKernel`（真实工程用法）：
	180	+
	181	+```cpp
	182	+using QKL1TileShape = GemmShape<128, 256, 128>;
	183	+using QKL0TileShape = QKL1TileShape;
	184	+using MmadDispatchPolicyQK = Gemm::MmadAtlasA2UnsharedFAQK;
	185	+using BlockMmadQK = Gemm::Block::BlockMmad<MmadDispatchPolicyQK, QKL1TileShape, QKL0TileShape,
	186	+ QType, KType, SType>;
	187	+
	188	+using DispatchPolicyFAUnsharedSoftmax = Epilogue::EpilogueAtlasA2FAUnsharedSoftmax;
	189	+using EpilogueFAUnsharedSoftmax = Epilogue::Block::BlockEpilogue<DispatchPolicyFAUnsharedSoftmax,
	190	+ PType, SType, maskType>;
	191	+
	192	+using PVL1TileShape = GemmShape<128, 128, 256>;
	193	+using MmadDispatchPolicyPV = Gemm::MmadAtlasA2UnsharedFAPV;
	194	+using BlockMmadPV = Gemm::Block::BlockMmad<MmadDispatchPolicyPV, PVL1TileShape, PVL0TileShape,
	195	+ PType, VType, OTmpType>;
	196	+
	197	+using UnsharedFAInferKernel = UnsharedFAInferKernel<BlockMmadQK, BlockMmadPV,
	198	+ EpilogueFAUnsharedSoftmax, isPAEnabled>;
	199	+UnsharedFAInferKernel unsharedInferKernel(tilingData);
	200	+unsharedInferKernel(params);
	201	+```
	202	+
	203	+类型约定：Q/K/V/P/O/Mask 为 `INPUT_T`（fp16/bf16），S/OTmp 为 float；上层 `UnsharedFAInferKernel` 负责 QK→softmax→PV 的调度与 `softmaxReady` 跨核通知。
	204	+
	205	+## 7. 与其他系列的差异
	206	+
	207	+- 相比 Shared FA 系列（`FAIQKSplitRow` + `OnlineSoftmaxCopySumMax` + `RescaleOWithoutDivSum`）：Unshared 场景 KV 短且每 head 独立，softmax 单趟完成，不维护跨 KV tile 的在线状态（rowMax/rowSum 的历史累积与 O rescale），因此无需 SplitRow 与 RescaleO 模板。
	208	+- 相比 FAI（PagedAttention 共享前缀）系列：不依赖 blockTables 分页 KV 索引，mask 在设备侧按 head 静态构造。

Adocs/zh/2_Design/04_contrib/00_XAttention/02_fai_split_row_kernel.md+301-0

@@ -0,0 +1,301 @@
	1	+# FAI SplitRow 系列模板设计文档
	2	+
	3	+本文档覆盖 Flash Attention 推理（FAI）SplitRow 系列的 4 个新增模板：
	4	+
	5	+\| 模板 \| 类型 \| 源文件 \|
	6	+\| ---- \| ---- \| ---- \|
	7	+\| `Gemm::MmadAtlasA2FAIQKSplitRow` \| BlockMmad 偏特化 \| `include/catlass/gemm/block/block_mmad_fai_qk_split_row.hpp` \|
	8	+\| `Gemm::MmadAtlasA2FAIPVSplitRow` \| BlockMmad 偏特化 \| `include/catlass/gemm/block/block_mmad_fai_pv_split_row.hpp` \|
	9	+\| `Epilogue::EpilogueAtlasA2OnlineSoftmaxCopySumMax` \| BlockEpilogue 偏特化 \| `include/catlass/epilogue/block/block_epilogue_online_softmax_copy_glm.hpp` \|
	10	+\| `Epilogue::EpilogueAtlasA2RescaleOWithoutDivSum` \| BlockEpilogue 偏特化 \| `include/catlass/epilogue/block/block_epilogue_rescale_o_no_div_rowsum.hpp` \|
	11	+
	12	+`EpilogueAscend950OnlineSoftmaxCopySumMax` 为 Ascend950 特化，直接继承 AtlasA2 实现（`using Base::Base`），不单独展开。
	13	+
	14	+## 1. 系列概述
	15	+
	16	+SplitRow 系列服务于共享式（Shared）FA 推理 kernel：QK GEMM 与 PV GEMM 分属两个 Cube 核，Softmax / RescaleO 作为 Vector epilogue 在 GM 中转数据上工作，通过 stackTile 外层循环遍历超长 KV 序列。相比 Unshared 系列（单核内完成 QK→softmax→PV），它把 softmax 卸载到 Vector 侧，两个 Cube 核可以分别满载 MMAD。
	17	+
	18	+```mermaid
	19	+graph TB
	20	+ Q[Q from GM] --> QK[BlockMmad FAIQKSplitRow<br/>核0: QK^T]
	21	+ QK -->\|S fp32 写 GM\| OS[BlockEpilogue OnlineSoftmaxCopySumMax<br/>Vector核: 在线softmax]
	22	+ OS -->\|P fp16 写 GM<br/>max/sum 写 GM\| PV[BlockMmad FAIPVSplitRow<br/>核1: PV]
	23	+ PV -->\|OTmp fp32 写 GM\| RO[BlockEpilogue RescaleOWithoutDivSum<br/>Vector核: 重标度累加]
	24	+ RO -->\|lastStackTile: O fp32\| CS[后续 CombineScale<br/>除 rowSum + cast]
	25	+```
	26	+
	27	+四模板协作要点：
	28	+
	29	+- QK SplitRow 沿 N 维（KV 序列长）切 stackTile，每个 stackTile 产出一块 S 矩阵；
	30	+- OnlineSoftmaxCopySumMax 对每块 S 做在线 softmax，输出 P（供 PV 使用）并维护跨 stackTile 的 rowmax/rowsum，仅在 lastStackTile 时将全局 max/sum 写回 GM；
	31	+- PV SplitRow 预载全部 V，等待 softmax 侧就绪后按 stackTile 消费 P，累加输出 OTmp；
	32	+- RescaleOWithoutDivSum 用 `dm = exp(oldMax - newMax)` 对历史累加结果重标度并累加新 OTmp，lastStackTile 时只输出 fp32 累加和，不除 rowSum、不做 cast，归一化交给后续 CombineScale epilogue。
	33	+
	34	+## 2. DispatchPolicy 定义
	35	+
	36	+```cpp
	37	+// gemm/dispatch_policy.hpp
	38	+template <bool PAGED_CACHE_FLAG_ = false, bool ENABLE_UNIT_FLAG_ = false>
	39	+struct MmadAtlasA2FAIQKSplitRow {
	40	+ using ArchTag = Arch::AtlasA2;
	41	+ static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_;
	42	+ static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
	43	+};
	44	+template <bool PAGED_CACHE_FLAG_ = false, bool ENABLE_UNIT_FLAG_ = false>
	45	+struct MmadAtlasA2FAIPVSplitRow { /* 同上 */ };
	46	+
	47	+// epilogue/dispatch_policy.hpp
	48	+struct EpilogueAtlasA2OnlineSoftmaxCopySumMax { using ArchTag = Arch::AtlasA2; };
	49	+struct EpilogueAtlasA2RescaleOWithoutDivSum { using ArchTag = Arch::AtlasA2; };
	50	+struct EpilogueAscend950OnlineSoftmaxCopySumMax { using ArchTag = Arch::Ascend950; };
	51	+```
	52	+
	53	+模板参数含义：
	54	+
	55	+- `PAGED_CACHE_FLAG`：KV 是否走 blockTable 分页缓存（PagedAttention）。`true` 时通过 `gBlockTable.GetValue(nowNIdx)` 查表计算 KV 偏移；`false` 时按 `nowNIdx * blockSize * strideKV` 连续寻址。
	56	+- `ENABLE_UNIT_FLAG`：预留开关（当前示例均传 `false`）。
	57	+
	58	+## 3. BlockMmad FAIQKSplitRow：Q @ K^T
	59	+
	60	+### 3.1 模板签名与约束
	61	+
	62	+```cpp
	63	+template <bool PAGED_CACHE_FLAG_, bool ENABLE_UNIT_FLAG_,
	64	+ class L1TileShape_, class L0TileShape_, class A_, class B_, class C_>
	65	+class BlockMmad<MmadAtlasA2FAIQKSplitRow<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>,
	66	+ L1TileShape_, L0TileShape_, A_, B_, C_>
	67	+```
	68	+
	69	+static_assert 约束：
	70	+
	71	+- `LayoutC` 仅支持 `RowMajor`（S 矩阵按行写 GM，供 softmax 按行读）；
	72	+- `N * K <= 32768`（L1B 双缓冲容量限制，K 即 embedding 维）。
	73	+
	74	+### 3.2 内存布局
	75	+
	76	+\| 缓冲 \| 大小 \| 组织 \|
	77	+\| ---- \| ---- \| ---- \|
	78	+\| L1A \| `M * K * sizeof(A)` \| 单缓冲，Q 一次性常驻 \|
	79	+\| L1B \| `32768 * sizeof(B)` \| 双缓冲（乒乓），KV 按 stackTile 分块流入 \|
	80	+
	81	+Q 通过 `loadQGM()` 一次搬运，采用扩展签名 `copyGmToL1A`，携带 `tokenNumPerGroup / qHeads * embed / BLOCK_SIZE` 参数实现 GQA 分组搬运（`layoutA.GetTileLayout(MakeCoord(singleGroupHeads, embed))`），即按"每组 Q 头"重排数据，避免逐头多次发起搬运。事件号 `EVENT_ID3`。
	82	+
	83	+### 3.3 主循环与流水
	84	+
	85	+`operator()` 三层循环结构：
	86	+
	87	+```
	88	+nL1Loop 按 L1TileShape::N 切 stackSeqTile（末轮 getBlockShape 取余量）
	89	+ └─ mL0Loop 按 L0TileShape::M 切行块
	90	+ └─ kL0Loop 按 L0TileShape::K 切 K 维
	91	+ └─ tileMmad(initMmad = (kL0Idx == 0))
	92	+```
	93	+
	94	+- 每个外层迭代调用 `getKVOffset(nowNIdx)`：分页模式查 blockTable 得 blockId 后计算 `blockId * blockSize * strideKV` 类偏移；连续模式直接线性寻址。
	95	+- 三组乒乓事件支撑全流水：
	96	+ - `l1KPPingPongFlag`：`MTE1_MTE2 / MTE2_MTE1`，控制 L1B 的 KV 装载与消费；
	97	+ - `l0ABPingPongFlag`：`M_MTE1 / MTE1_M`，控制 L0A/L0B 的装载与 MMAD 消费；
	98	+ - `l0CPingPongFlag`：`M_FIX / FIX_M`，控制 L0C 结果经 FixPipe 写 GM 与下一轮覆写。
	99	+- `tileMmad` 的 `initMmad` 标志保证 K 维首次累加时初始化 L0C，后续累加。
	100	+
	101	+## 4. BlockMmad FAIPVSplitRow：P @ V
	102	+
	103	+### 4.1 模板签名与约束
	104	+
	105	+```cpp
	106	+template <...同上...>
	107	+class BlockMmad<MmadAtlasA2FAIPVSplitRow<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>, ...>
	108	+```
	109	+
	110	+static_assert：`M * K <= 32768`（L1A 双缓冲容量限制）。
	111	+
	112	+### 4.2 内存布局
	113	+
	114	+\| 缓冲 \| 大小 \| 组织 \|
	115	+\| ---- \| ---- \| ---- \|
	116	+\| L1A \| `32768 * sizeof(A)` \| 双缓冲，偏移 `l1BufAddrStart + L1A_SIZE * i`，P 分块流入 \|
	117	+\| L1B \| `N * K * sizeof(B)` \| 单缓冲（全部 V 一次性常驻），偏移 `l1BufAddrStart + L1A_SIZE * 2` \|
	118	+
	119	+### 4.3 执行时序（核心设计）
	120	+
	121	+```
	122	+1. kLoop = CeilDiv(stackSeqTile, blockSize)
	123	+ └─ 循环将全部 V 按 block 分页搬入 L1B（分页时查 blockTable），事件 EVENT_ID2
	124	+2. Arch::CrossCoreWaitFlag(softmaxFlag)
	125	+ └─ 跨核等待：softmax 核已产出可用的 P（GM 中转）
	126	+3. mL1Loop × kL1Loop × kL0Loop
	127	+ └─ 搬 P 进 L1A（乒乓）→ MMAD 累加（initMmad = (kL1Idx==0 && kL0Idx==0)）
	128	+4. 输出 OTmp（fp32）写 GM
	129	+5. SetFlag(MTE1_MTE2, EVENT_ID2)
	130	+ └─ 结束后释放 V 搬运通道，供下一 stackTile 复用
	131	+```
	132	+
	133	+相比 QK 的三组乒乓，PV 侧的同步重点是先 V 后 P：V 预载可提前于 softmax 完成进行，`CrossCoreWaitFlag` 只阻塞 P 消费路径，V 装载时间被完全隐藏。
	134	+
	135	+## 5. BlockEpilogue OnlineSoftmaxCopySumMax：在线 Softmax
	136	+
	137	+### 5.1 模板签名
	138	+
	139	+```cpp
	140	+template <class OutputType_, class InputType_, class MaskType_>
	141	+class BlockEpilogue<EpilogueAtlasA2OnlineSoftmaxCopySumMax, OutputType_, InputType_, MaskType_>
	142	+```
	143	+
	144	+- `OutputType_`：P 矩阵类型（fp16/bf16，供 PV GEMM 消费）；
	145	+- `InputType_`：S 矩阵类型（fp32，来自 QK GEMM）；
	146	+- `MaskType_`：attention mask 类型，支持 `NO_MASK` 与常规 mask 两种路径。
	147	+
	148	+构造函数签名 `BlockEpilogue(resource, scaleValue_)`，`scaleValue_` 即 softmax 前的 `1/sqrt(d)` 缩放系数。
	149	+
	150	+### 5.2 UB 布局
	151	+
	152	+按 `UB_UINT8_BLOCK_SIZE`（16384 字节块）组织，关键偏移：
	153	+
	154	+\| 张量 \| UB 偏移 \| 说明 \|
	155	+\| ---- \| ---- \| ---- \|
	156	+\| ls（S 本轮值） \| `0` \| 当前 stackTile 的 scale·S，fp32，容量 `MAX_UB_S_ELEM_NUM = 8192` \|
	157	+\| lp（P 输出）/ mask32 \| `4 * block` \| 与 mask 的 32 位视图共享空间 \|
	158	+\| tmp（tv） \| `10 * block` \| 归约中间量 \|
	159	+\| lm / hm \| `10 * block + 8/9 * vec` \| 本轮/全局行最大值 \|
	160	+\| gm / ll / gl / dm \| `10 * block + 10~13 * vec` \| 全局 max、本轮行和、全局行和、重标度系数 \|
	161	+\| mask \| `11 * block` \| mask 原始数据 \|
	162	+
	163	+相邻行和/行最大采用乒乓布局（`ROW_SUM_PINGPONG_OFFSET = 64 * 8`），配合行分块循环隐藏 MTE2 装载。`MAX_ROW_NUM_SUB_CORE = 128` 限定单个 SubBlock 处理的最大行数。
	164	+
	165	+### 5.3 在线 Softmax 迭代式
	166	+
	167	+对每个 stackTile，按行维护全局 `gm`（max）与 `gl`（sum）：
	168	+
	169	+```
	170	+lm = rowmax(scale · S) // 本 tile 行最大
	171	+hm = isFirst ? lm : max(lm, gm) // 更新全局最大
	172	+dm = isFirst ? 1 : exp(gm - hm) // 历史缩放系数
	173	+ls = exp(ls - hm) // 数值稳定的指数
	174	+ll = rowsum(ls) // 本 tile 行和
	175	+gl = isFirst ? ll : dm * gl + ll // 重标度累加
	176	+gm = hm
	177	+```
	178	+
	179	+`isLastStackTile` 时：`gm`/`gl` 经 `Brcb` 展开 + `DataCopy(rowNum, 1, 0, headNum - 1)`（stride 间隔写）输出到 `gSharedMax`/`gSharedSum`，供 RescaleO 与后续 CombineScale 使用。非末 tile 期间 max/sum 只在 UB/寄存器中滚动，不产生 GM 流量。
	180	+
	181	+### 5.4 行归约三分支
	182	+
	183	+`Rowmax`/`Rowsum` 按 `columnNum`（stackTile 序列长）分三档实现：
	184	+
	185	+\| 分支 \| 条件 \| 手段 \|
	186	+\| ---- \| ---- \| ---- \|
	187	+\| `SPECTILE512` \| columnNum == 512 \| 3 次 `BlockReduceMax/Sum` 级联 \|
	188	+\| `SPECTILE256` \| columnNum == 256 \| `SetVecMask(32)` + `SetBlockReduceMask(4)` \|
	189	+\| `TAILTILE` \| 其他 \| 整 64 元素向量循环 + 尾部 `SetVecMask` 掩码处理 \|
	190	+
	191	+### 5.5 P 的降精度输出
	192	+
	193	+`CalcExp` 中 `hm` 经 `Brcb` 广播到整行后计算 `exp`，随后 `DownCastP` 将 fp32 的 P 转为 fp16（`CAST_NONE`）或 bf16（`CAST_RINT`，避免溢出），再 `CopyPUbToGm` 写 GM。bf16 场景选用 RINT 舍入是精度关键点。
	194	+
	195	+### 5.6 SubBlock 切分与预取流水
	196	+
	197	+`operator()` 将行维度对半切给两个 SubBlock（`qNBlockSize == 1` 时 `qSBlockSize / 2` 对半；否则按 qN 乘子扩大）。行方向再按 `maxRowNumPerLoop`（由 8192 元素容量折算）分块，采用 `preLoad = 1` 的乒乓预取：第 i 块计算时预搬第 i+1 块的 S，事件族 `V_MTE2 / MTE2_V / V_MTE3 / MTE3_V`（HardEvent）保证搬运、计算、写出三级流水。
	198	+
	199	+## 6. BlockEpilogue RescaleOWithoutDivSum：O 累加
	200	+
	201	+### 6.1 模板签名
	202	+
	203	+```cpp
	204	+template <class OutputType_, class InputType_, class UpdateType_>
	205	+class BlockEpilogue<EpilogueAtlasA2RescaleOWithoutDivSum, OutputType_, InputType_, UpdateType_>
	206	+```
	207	+
	208	+典型实例化：`OutputType_ = fp16`、`InputType_ = fp32`（OTmp）、`UpdateType_ = fp32`。构造仅接收 `resource`，无额外参数。
	209	+
	210	+### 6.2 UB 布局与共享设计
	211	+
	212	+\| 张量 \| UB 偏移 \|
	213	+\| ---- \| ---- \|
	214	+\| lo（本轮 OTmp） \| `6 * block` \|
	215	+\| go（累加 O） \| `8 * block`，`goUbTensor16/goUbTensor32` 同址双视图 \|
	216	+\| tmp \| `10 * block` \|
	217	+\| hm / gl / dm \| `10 * block + 9/12/13 * vec` \|
	218	+
	219	+注意 hm/gl/dm 偏移与 OnlineSoftmax epilogue 的布局完全一致——两者在同一 kernel 的不同阶段运行，复用同一份 UB 规划，避免跨 epilogue 的布局冲突。`MAX_UB_O_ELEM_NUM = 4096` 限制单块行数 × 列数。
	220	+
	221	+### 6.3 核心算法
	222	+
	223	+```
	224	+WaitFlag(V_MTE2, EVENT_ID3) // 等待 OTmp 可读（与 PV GEMM 握手）
	225	+if (isFirstStackTile):
	226	+ go = lo // 首块直接落位
	227	+else:
	228	+ lo = GM -> UB（上一 stackTile 的累加结果）
	229	+ dm[curStackTileMod * 128] 经 Brcb 广播为 dm_block
	230	+ go = go * dm_block // 按 FLOAT_VECTOR_SIZE=64 分段 Mul
	231	+ go = go + lo // 累加
	232	+SetFlag(V_MTE2, EVENT_ID3) // 释放握手
	233	+if (isLastStackTile):
	234	+ CopyFloatOToGm -> gSharedOut // float 直出，不除 rowSum、不 cast
	235	+```
	236	+
	237	+- `EVENT_ID3`（`V_MTE2` 事件对）是本 epilogue 与数据生产方的跨模块握手信号；`EVENT_ID0` 为内部 MTE2↔V 同步。
	238	+- `CopyFloatOToGm` 内部按 `qNBlockSize == 0` 走单块 `DataCopyPad`，否则逐 qN 块输出。
	239	+- 末 tile 不做 `O / rowSum`、不做 fp32→fp16 cast——除法与降精度统一推迟到下游 CombineScale epilogue，一次性完成，减少中间精度损失与 GM 读写次数。
	240	+
	241	+### 6.4 SubBlock 切分
	242	+
	243	+与 OnlineSoftmax 的行对半不同，本 epilogue 按 qN 维度自适应：`qNBlockSize == 1` 时两个 SubBlock 对分行；`qNBlockSize > 1` 时对分列（outCol 对半），inRow 则按 qN 乘子整体扩大。切分维度与数据的 GM 排布（qN 在行维展开时分行连续）对齐，保证每个 SubBlock 的搬运都是连续段。
	244	+
	245	+## 7. 使用示例
	246	+
	247	+### 7.1 真实工程组装（xllm_ops x_attention）
	248	+
	249	+摘自 xllm_ops（https://gitcode.com/xLLM-AI/xllm_ops）`x_attention/op_kernel/x_attention_catlass_helper.h` 的 `CallSharedInferKernelShort`：
	250	+
	251	+```cpp
	252	+using INPUT_T = ...; // Q/K/V/P/O 与 mask 的元素类型（如 fp16）
	253	+using L0TileShape = GemmShape<16, 16, 16>;
	254	+
	255	+using L1TileShapeQK = GemmShape<128, 128, 128>; // L1TileShape::K must be embedding
	256	+using DispatchPolicyQK = Gemm::MmadAtlasA2FAIQKSplitRow<isPAEnabled, false>;
	257	+using BlockMmadQK = Gemm::Block::BlockMmad<DispatchPolicyQK, L1TileShapeQK,
	258	+ L0TileShape, QType, KType, SType>;
	259	+using DispatchPolicyOnlineSoftmax = Epilogue::EpilogueAtlasA2OnlineSoftmaxCopySumMax;
	260	+using EpilogueOnlineSoftmax = Epilogue::Block::BlockEpilogue<
	261	+ DispatchPolicyOnlineSoftmax, PType, SType, maskType>;
	262	+// update rowsum rowmax and copyOut on lastStackTile
	263	+
	264	+using L1TileShapePV = GemmShape<128, 128, 128>;
	265	+using DispatchPolicyPV = Gemm::MmadAtlasA2FAIPVSplitRow<isPAEnabled, false>;
	266	+using BlockMmadPV = Gemm::Block::BlockMmad<DispatchPolicyPV, L1TileShapePV,
	267	+ L0TileShape, PType, VType, OTmpType>;
	268	+using DispatchPolicyRescaleO = Epilogue::EpilogueAtlasA2RescaleOWithoutDivSum;
	269	+using EpilogueRescaleO = Epilogue::Block::BlockEpilogue<
	270	+ DispatchPolicyRescaleO, OType, OTmpType, OUpdateType>;
	271	+// do not div rowSum or cast on lastStackTile
	272	+
	273	+using SharedFAInferKernel = SharedFAInferKernelShort<
	274	+ BlockMmadQK, BlockMmadPV, EpilogueOnlineSoftmax, EpilogueRescaleO, isPAEnabled>;
	275	+SharedFAInferKernel kernel(...);
	276	+kernel(); // 外层 stackTile 循环内由 kernel 依次驱动四个模板
	277	+```
	278	+
	279	+组装要点：
	280	+
	281	+1. `L1TileShape::K` 必须等于 embedding 维（128），QK/PV 的 K 都是"收缩维=embedding"；
	282	+2. QK 的 C 类型（SType）为 fp32（`LayoutC = RowMajor`），PV 的 C 类型为 fp32 的 OTmp；
	283	+3. `isPAEnabled` 直接透传给两个 BlockMmad 的 `PAGED_CACHE_FLAG`，blockTable 指针随 kernel 参数传入；
	284	+4. Mask 类型经 `maskType` 传给 OnlineSoftmax epilogue，`NO_MASK` 时走无掩码快路径。
	285	+
	286	+### 7.2 仓内测试
	287	+
	288	+`tests/optest/kernels/23_flash_attention_infer/flash_attention_infer.cpp:741-768` 提供了 SplitRow 路径的完整调用样例（构造、传参、启动），可与上文组装代码互相印证。
	289	+
	290	+## 8. 与 FA Unshared 系列的差异
	291	+
	292	+\| 维度 \| FAI SplitRow（本文） \| FA Unshared（见 01 文档） \|
	293	+\| ---- \| ---- \| ---- \|
	294	+\| 核拓扑 \| QK、PV 两个 Cube 核 + Vector epilogue \| 单 Cube 核内 QK→softmax→PV \|
	295	+\| S/P 数据通路 \| 经 GM 中转，各模块独立乒乓 \| 核内 L0C/UB 直连，无 GM 往返 \|
	296	+\| 长序列处理 \| stackTile 外层循环，天然支持超长 KV \| 单轮固定 tile，序列长受 L1 约束 \|
	297	+\| softmax 位置 \| 独立 epilogue 核，OnlineSoftmax 在线式 \| `EpilogueAtlasA2FAUnsharedSoftmax` 单次 softmax \|
	298	+\| 归一化时机 \| RescaleO 不除 sum，推迟到 CombineScale \| 核内直接完成 \|
	299	+\| 适用场景 \| 长序列 / PagedAttention 推理 \| 短序列、tile 可整装的推理 \|
	300	+
	301	+选型建议：KV 序列长超出单核 L1 容量（`N * K > 32768`）或需要分页 KV 时选 SplitRow；序列短且可整装时 Unshared 的核内直连更省 GM 带宽。

Adocs/zh/2_Design/04_contrib/00_XAttention/03_xfai_kernel.md+424-0

@@ -0,0 +1,424 @@
	1	+# XFAI Kernel 设计（Atlas A2 推理 Flash Attention）
	2	+
	3	+## 1. 系列概述
	4	+
	5	+XFAI（X Flash Attention Infer）系列面向 Atlas A2（DaVinci C220）推理场景的 Flash Attention 算子，与 [FAI SplitRow 系列](./02_fai_split_row_kernel.md) 同属"AIC 算 + AIV 归"的跨核流水形态，但在模板组织上做了更强的面向封装：
	6	+
	7	+- QK 与 PV 两段 GEMM 分别独立成 `BlockMmad` 模板，中间的 `S`/`P` 矩阵经 GM 中转；
	8	+- Online Softmax、RescaleO 作为 `BlockEpilogue` 偏特化模板，RescaleO 引入独立的 `UpdateType_`/`LseType_` 类型参数；
	9	+- 原生支持 Paged KV Cache（`PAGED_CACHE_FLAG_` 策略参数，QK/PV 侧均可开启）；
	10	+- 提供 `CombineScale` 独立模板类，用于 split-KV 多份局部结果的 LSE 合并归一。
	11	+
	12	+### 1.1 模板清单
	13	+
	14	+\| 模板 \| 策略（DispatchPolicy） \| 位置 \|
	15	+\| --- \| --- \| --- \|
	16	+\| BlockMmadXFAIQK \| `MmadAtlasA2XFAIQK<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>` \| `include/catlass/gemm/block/block_mmad_xfai_qk.hpp` \|
	17	+\| BlockMmadXFAIPV \| `MmadAtlasA2XFAIPV<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>` \| `include/catlass/gemm/block/block_mmad_xfai_pv.hpp` \|
	18	+\| BlockEpilogueXFAIOnlineSoftmax \| `EpilogueAtlasA2XFAIOnlineSoftmax<LSE_MODE_>` \| `include/catlass/epilogue/block/block_epilogue_xfai_online_softmax.hpp` \|
	19	+\| BlockEpilogueXFAIRescaleO \| `EpilogueAtlasA2XFAIRescaleO<LSE_MODE_>` \| `include/catlass/epilogue/block/block_epilogue_xfai_rescale_o.hpp` \|
	20	+\| CombineScale \| `CombineScale<OutputType_, LseType_>`（独立类，非 BlockEpilogue 特化） \| `include/catlass/epilogue/block/block_epilogue_xfai_combine_scale.hpp` \|
	21	+
	22	+### 1.2 跨核流水全景
	23	+
	24	+```mermaid
	25	+graph LR
	26	+ Q[Q/K/V GM] --> AIC1[AIC: BlockMmadXFAIQK<br/>QK GEMM]
	27	+ AIC1 -- "S (fp32) 写 GM<br/>qkReady" --> AIV1[AIV: OnlineSoftmax<br/>scale/mask/exp/sum]
	28	+ AIV1 -- "P 写 GM<br/>softmaxReady" --> AIC2[AIC: BlockMmadXFAIPV<br/>PV GEMM]
	29	+ AIC2 -- "OTmp 写 GM<br/>pvReady" --> AIV2[AIV: RescaleO<br/>rescale/归一/LSE]
	30	+ AIV2 -- split-KV 场景 --> CS[CombineScale<br/>LSE 合并归一]
	31	+```
	32	+
	33	+核心思想：CUBE 核负责 QK、PV 两段 GEMM，Vector 核负责 softmax 与 rescale；`preLoad=1` 的软件流水使 QK 领先 PV 一个 stackTile，两组乒乓掩盖 GM 中转开销。
	34	+
	35	+## 2. DispatchPolicy 定义
	36	+
	37	+GEMM 侧（`include/catlass/gemm/dispatch_policy.hpp:175-187`）：
	38	+
	39	+```cpp
	40	+// 推理 Flash Attention 的 QK GEMM 策略
	41	+template <bool PAGED_CACHE_FLAG_ = false, bool ENABLE_UNIT_FLAG_ = false>
	42	+struct MmadAtlasA2XFAIQK : public MmadAtlasA2 {
	43	+ using ArchTag = AtlasA2;
	44	+ static constexpr uint32_t STAGES = 2; // L1/L0 二级乒乓
	45	+ static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_;
	46	+ static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
	47	+};
	48	+
	49	+// 推理 Flash Attention 的 PV GEMM 策略（模板参数同上）
	50	+template <bool PAGED_CACHE_FLAG_ = false, bool ENABLE_UNIT_FLAG_ = false>
	51	+struct MmadAtlasA2XFAIPV : public MmadAtlasA2 { ... };
	52	+```
	53	+
	54	+Epilogue 侧（`include/catlass/epilogue/dispatch_policy.hpp:283-306`）：
	55	+
	56	+```cpp
	57	+// Online Softmax 策略，LSE_MODE_ 控制 logsumexp 输出行为
	58	+template <bool LSE_MODE_ = false>
	59	+struct EpilogueAtlasA2XFAIOnlineSoftmax {
	60	+ using ArchTag = AtlasA2;
	61	+ static constexpr bool LSE_MODE = LSE_MODE_;
	62	+};
	63	+
	64	+// RescaleO 策略，模板参数含义同上
	65	+template <bool LSE_MODE_ = false>
	66	+struct EpilogueAtlasA2XFAIRescaleO { ... };
	67	+```
	68	+
	69	+## 3. BlockMmadXFAIQK —— QK GEMM 设计
	70	+
	71	+类形态（`block_mmad_xfai_qk.hpp:32`）：
	72	+
	73	+```cpp
	74	+template <typename DispatchPolicy, typename L1TileShape_, typename L0TileShape_,
	75	+ typename AType_, typename BType_, typename CType_>
	76	+class BlockMmad<MmadAtlasA2XFAIQK<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>,
	77	+ L1TileShape_, L0TileShape_, AType_, BType_, CType_> { ... };
	78	+```
	79	+
	80	+### 3.1 关键常量与存储布局
	81	+
	82	+\| 常量 \| 值 \| 含义 \|
	83	+\| --- \| --- \| --- \|
	84	+\| BLOCK_SIZE \| 16 \| KV 分页基础块（head 维分块粒度） \|
	85	+\| EMBED_SPLIT_SIZE \| 128 \| embed 维单块大小 \|
	86	+\| UNIT_BLOCK_STACK_NUM \| 4 \| 单次进栈的 unit block 数 \|
	87	+\| KV_BASE_BLOCK \| 512 \| KV cache 分页基准块 \|
	88	+\| KV_SPLIT_SIZE \| 128 \| KV split 切分粒度 \|
	89	+
	90	+- 构造函数：`BlockMmad(resource, nDyn, kDyn, l1BufAddrStart = 0)`。Q 矩阵在 L1A 单拷常驻（`l1ATensor` 起始于 `l1BufAddrStart`），B（K）以及 L0A/L0B/L0C 按 `STAGES=2` 双缓冲乒乓——这一布局是 PV 与 QK 共用同一个 CUBE 核 L1 空间的前提（PV 的 B 区紧跟 QK 占用之后）。
	91	+- 输出 C（即 S 矩阵）经 `copyL0CToGm` 从 L0C 直写 GM，交由 AIV 侧消费。
	92	+
	93	+### 3.2 Q 的分组加载：loadQGM
	94	+
	95	+```cpp
	96	+void loadQGM(const GlobalTensor<AType>& gA, const LayoutA& layoutA,
	97	+ uint32_t rowNum, uint32_t singleGroupHeads, uint32_t qHeads);
	98	+```
	99	+
	100	+Q 按 head 分组加载：`tokenNumPerGroup = rowNum / singleGroupHeads`，组间通过偏移跳转，一次 MTE2→MTE1 搬运完成整组入 L1A（`EVENT_ID3` 同步）。分组的目的是让同一份 Q 服务多个 KV split 的循环复用。
	101	+
	102	+### 3.3 分页寻址：getKVOffset
	103	+
	104	+KV 偏移按 `PAGED_CACHE_FLAG_` 走双路径：
	105	+
	106	+- Paged 路径：经 `gBlockTable` 查表，`kOffset = blockTableId * KV_BASE_BLOCK * strideKV + ...`，支持 KV cache 稀疏分页；
	107	+- 非 Paged 路径：直接 `nowNIdx * BLOCK_SIZE * strideKV` 连续寻址。
	108	+
	109	+### 3.4 主循环：三层乒乓
	110	+
	111	+`operator()` 内部为 nL1×mL0×kL0 三层循环：nL1 层按 `stackSeqTile` 以 `l1NDynamic` 切分 KV split 块；`l1KvPingPongFlag`（L1 B 区）、`l0ABPingPongFlag`（L0A/L0B）、`l0CPingPongFlag`（L0C）三组乒乓独立翻转，使"加载 K / 计算 QK / 回写 S"三段充分重叠。S 写 GM 后由调用方置 `qkReady` 跨核 flag 通知 AIV。
	112	+
	113	+## 4. BlockMmadXFAIPV —— PV GEMM 设计
	114	+
	115	+类形态（`block_mmad_xfai_pv.hpp`）：偏特化签名与 QK 相同，仅 Policy 换为 `MmadAtlasA2XFAIPV`。
	116	+
	117	+### 4.1 L1 布局与构造
	118	+
	119	+```cpp
	120	+BlockMmad(resource, nDyn, kDyn, l1BufAddrStart);
	121	+// 内部：l1B 置于 l1BufAddrStart + M * kDyn * sizeof(A) * 2
	122	+```
	123	+
	124	+PV 与 QK 共用 CUBE 核的 L1：QK 的 Q 常驻区之后紧跟 PV 的 B（V）区。V 区单拷加载、P 区（`l1ATensor[STAGES]`）双缓冲。
	125	+
	126	+### 4.2 operator() 与跨核门控
	127	+
	128	+```cpp
	129	+// 节选调用签名
	130	+blockMmadPV(gP, gV, gOTmp, ..., blockStackNum, /* Arch::CrossCoreFlag */ softmaxFlag,
	131	+ nIdx, nLoop /* 引用出参：当前块序号与总块数 */);
	132	+```
	133	+
	134	+执行序：
	135	+
	136	+1. 批量加载 V：一次将 `blockStackNum` 个 KV 块的 V 预取入 L1（`EVENT_ID4` + `EVENT_ID0` 同步），为后续多个 stackTile 的 PV 复用；
	137	+2. 跨核等待：`CrossCoreWaitFlag(softmaxFlag)` —— 阻塞直到 AIV 侧 Online Softmax 把当前块的 P 写入 GM。这是"QK → softmax → PV"流水衔接的关键门控；
	138	+3. P 流式乒乓：主循环 nL1×mL1×kL1×kL0 中，P 从 GM 按 `l1PPingPongFlag` 乒乓流入 L1A，`LOAB_BLOCK = 1`；
	139	+4. PV 直写 GM：结果 O 中间量（OTmp）从 L0C 直写 GM，尾部 `SetFlag`（MTE1_MTE2，`EVENT_ID4`）供 AIV 侧 RescaleO 对齐。
	140	+
	141	+## 5. BlockEpilogueXFAIOnlineSoftmax 设计
	142	+
	143	+类形态（`block_epilogue_xfai_online_softmax.hpp:25`）：
	144	+
	145	+```cpp
	146	+template <typename OutputType_, typename InputType_, typename MaskType_>
	147	+class BlockEpilogue<EpilogueAtlasA2XFAIOnlineSoftmax<LSE_MODE_>,
	148	+ OutputType_, InputType_, MaskType_> { ... };
	149	+// OutputType_ = P（下传 PV 的类型），InputType_ = S（fp32），MaskType_ 一般为 int8_t
	150	+```
	151	+
	152	+构造函数：`BlockEpilogue(resource, scaleValue)`，`scaleValue` 即 `1/sqrt(d)` 缩放因子。
	153	+
	154	+### 5.1 UB 布局
	155	+
	156	+\| 缓冲 \| 规格 \| 用途 \|
	157	+\| --- \| --- \| --- \|
	158	+\| lsUbTensor \| 8192 元素乒乓 \| S 矩阵（fp32），双缓冲配合 preLoad 流水 \|
	159	+\| lpUbTensor \| 同 ls \| P 矩阵（downcast 后） \|
	160	+\| maskUbTensor / mask16 / mask32 \| 三级 \| token mask 的原始/16bit/32bit 视图 \|
	161	+\| lm / hm / gm / dm \| 行向量标量 \| 局部 max / 展开后 max / 全局 max / dm=exp(gm-hm) \|
	162	+\| ll / gl \| 行向量标量 \| 局部 sum / 全局 sum \|
	163	+\| tv \| 行向量标量 \| 广播暂存（LSE 输出用） \|
	164	+
	165	+### 5.2 Rowsum / Rowmax 三分派
	166	+
	167	+按当前块的列宽分派三种归约特化：SPECTILE512 / SPECTILE256 / TAILTILE。前两种对齐列宽走整块 `BlockReduceSum/Max`（配合 `SetVecMask`/`SetBlockReduceMask`），TAILTILE 处理尾块非对齐列，保证跨块的行最大/行求和语义正确。
	168	+
	169	+### 5.3 Online Softmax 核心流程
	170	+
	171	+```
	172	+CopySGmToUb // S 从 GM 入 UB
	173	+ScaleS // S *= scaleValue
	174	+[ApplyMask] // 可选：mask32 * (-3e38) 加到 S 上实现屏蔽
	175	+CalcLocalRowMax // lm = rowmax(S)
	176	+UpdateGlobalRowMax // hm = max(lm, gm); dm = exp(gm - hm); gm ← hm
	177	+CalcExp // ls = exp(ls - hm) (Brcb 行广播 hm)
	178	+DownCastP // fp32 → OutputType_（bf16 走 CAST_RINT 取整策略）
	179	+CopyPUbToGm // P 直写 GM，供 AIC 侧 PV 消费
	180	+CalcLocalRowSum // ll = rowsum(P)
	181	+UpdateGlobalRowSum // gl = dm * gl + ll
	182	+```
	183	+
	184	+`dm = exp(gm_old - gm_new)` 即上一轮 P 需要整体乘的衰减因子，交给 RescaleO 使用——这就是经典的 FlashAttention online softmax 递推。
	185	+
	186	+### 5.4 双 operator()：无 mask 版与带 mask 版
	187	+
	188	+```cpp
	189	+// 无 mask 版（参数节选）
	190	+epilogueOnlineSoftmax(gS, gP, layoutS, layoutP, ..., stackSeqTile);
	191	+
	192	+// 带 mask 版（causal / token mask，13 参，参数节选）
	193	+epilogueOnlineSoftmax(gS, gP, gMask, layouts..., qkReady,
	194	+ triUp, triDown, kvSStartIdx, kvSEndIdx, ...);
	195	+```
	196	+
	197	+带 mask 版要点：
	198	+
	199	+- 入口先 `CrossCoreWaitFlag(qkReady)` 等 AIC 侧 S 就绪；
	200	+- mask 按行三段复制（proTokenNum / integralHeadNum / epiTokenNum），适配 token 级压缩布局；
	201	+- `triUp/triDown` 描述 causal 三角的上下边界，配合 `kvSStartIdx/kvSEndIdx` 判定当前 KV 块是否整块可跳过或部分遮蔽；
	202	+- `LSE_MODE == LSE_OUT` 时末块借 `tv` 缓冲把 LSE 行值写出。
	203	+
	204	+### 5.5 subBlock 切分与流水
	205	+
	206	+- `qNBlockSize == 1`（每 subBlock 一行 Q）时行数对半切给两个 subBlock 并行处理；
	207	+- `preLoad = 1`：S 的 GM 读取与上一块的 softmax 计算重叠，掩盖搬运延迟。
	208	+
	209	+## 6. BlockEpilogueXFAIRescaleO 设计
	210	+
	211	+类形态（`block_epilogue_xfai_rescale_o.hpp`）：五类型参数，是本系列对基类模板的扩展点：
	212	+
	213	+```cpp
	214	+template <typename OutputType_, typename InputType_, typename UpdateType_, typename LseType_>
	215	+class BlockEpilogue<EpilogueAtlasA2XFAIRescaleO<LSE_MODE_>, OutputType_, InputType_, UpdateType_, LseType_>;
	216	+// OutputType_ = O（最终输出），InputType_ = OTmp（PV 累加结果），
	217	+// UpdateType_ = 中间轮回写类型（fp32），LseType_ = LSE 输出类型
	218	+```
	219	+
	220	+构造函数：`BlockEpilogue(resource)`。
	221	+
	222	+### 6.1 核心算法：O 的新旧融合
	223	+
	224	+```
	225	+lo = gInput[当前块] // 读上一轮的 O（或首轮读 OTmp 直通）
	226	+go = gUpdate * dm_block + lo // dm 经 Brcb 行广播；即 O_new = exp(gm_old-gm_new)*O_old + PV_cur
	227	+if (!isLastStackTile):
	228	+ go → 回写 gUpdate（fp32 中间量，needRowLoop 场景）
	229	+else:
	230	+ go = go / gl_block // 末块按全局行和归一
	231	+ go → downcast → CopyOToGm // 按 pro/integral/epi token 三段写出最终 O
	232	+ if (LSE_MODE == LSE_OUT):
	233	+ lse = ln(gl) + gm // Brcb 行广播后 DataCopyPad 写 gLse
	234	+```
	235	+
	236	+要点：
	237	+
	238	+- 首块直通：第一个 stackTile 无旧 O 可乘，`go = lo` 直接进入累加；
	239	+- 末块归一：除以全局行和 `gl` 完成 softmax 分母归一，downcast 后按 token 三段（pro/integral/epi）写 GM，兼容变长 seq；
	240	+- LSE 输出：`lse = ln(gl) + gm`，配合 OnlineSoftmax 侧的 gm/gl 传递即可在末块一次性得到 logsumexp；
	241	+- 事件体系：`EVENT_ID0/1/3/5/6` 管理 GM 读入、广播、写出的多级同步。
	242	+
	243	+## 7. CombineScale 设计（split-KV 合并）
	244	+
	245	+类形态（`block_epilogue_xfai_combine_scale.hpp`）：独立模板类，不是 `BlockEpilogue` 的偏特化，因为它消费的是"多份局部结果"而非单个 tile：
	246	+
	247	+```cpp
	248	+template <typename OutputType_, typename LseType_>
	249	+class CombineScale { ... };
	250	+```
	251	+
	252	+### 7.1 operator() 签名
	253	+
	254	+```cpp
	255	+void operator()(uint32_t qHeads, uint32_t kvSplitCoreNum, uint32_t headSizeV,
	256	+ __gm__ SplitKvExtraInfo* extraInfo,
	257	+ GlobalTensor lGmTensor, // 各 split 的 LSE（局部）
	258	+ GlobalTensor oCoreTmpGmTensor, // 各 split 的 O（局部，未归一）
	259	+ GlobalTensor oGmTensor, // 合并后最终 O
	260	+ GlobalTensor gActualQseqlen, bool inputLayoutTND = true);
	261	+```
	262	+
	263	+任务遍历：`for (process = subBlockID; process < kvSplitCoreNum * 2; process += subBlockNum)`，从 `extraInfo->splitInfo` 解析 batchIdx/headStartIdx/headEndIdx/qStartIdx/qEndIdx/splitNum 及 LSE/O 任务偏移。
	264	+
	265	+### 7.2 LSE 合并算法
	266	+
	267	+对每个 q 行，把 splitNum 份局部 LSE（记 `ll[n]`）合并为全局值：
	268	+
	269	+```
	270	+lm = ReduceMax(ll) // 各 split 局部 max 的最大值
	271	+tl = exp(ll - lm) // 平移防溢出
	272	+rs = Ln(ReduceSum(tl)) // log-sum-exp 的核心一步
	273	+ts = rs + lm // 最终 LSE
	274	+gl = exp(ll - ts) // 每个 split 的最终缩放系数
	275	+```
	276	+
	277	+### 7.3 O 合并
	278	+
	279	+```
	280	+O_final = Σ_n O_n * gl[n] // BroadCast + Mul + Add 循环，loFloat 乒乓累加
	281	+```
	282	+
	283	+`toUbTensor` 与 `broadCastOTensor` 共用 UB 空间（时间上错开）以节省 UB。实现细节：
	284	+
	285	+- TND 布局：按 `prevQSeqlenSum` 前缀和计算各 batch 的 q 行偏移；
	286	+- q_len == 1 特化：decode 场景前后半分 subBlock 独立寻址；
	287	+- 对齐：`splitNumAlign` / LSE 块按 8 元素（32bit 语义）对齐，规避 DataCopy 非对齐限制。
	288	+
	289	+## 8. 跨核同步链与软件流水
	290	+
	291	+调用方（kernel 层）持有三个 `Arch::CrossCoreFlag`：
	292	+
	293	+\| Flag \| 生产者 \| 消费者 \| 含义 \|
	294	+\| --- \| --- \| --- \| --- \|
	295	+\| qkReady \| AIC（QK，PIPE_FIX） \| AIV（OnlineSoftmax 带 mask 版） \| 当前块 S 已写 GM \|
	296	+\| softmaxReady \| AIV（Softmax，PIPE_MTE3） \| AIC（PV 内部 WaitFlag） \| 当前块 P 已写 GM \|
	297	+\| pvReady \| AIC（PV，PIPE_FIX） \| AIV（RescaleO） \| 当前块 OTmp 已写 GM \|
	298	+
	299	+外层循环以 `kvSIdx ∈ [0, kvSLoopNumTotal + preKVNum)` 驱动软件流水：`kvSIdx < kvSLoopNumTotal` 时执行 QK+Softmax，`kvSIdx >= preKVNum` 时执行 PV+RescaleO（`nowkvSIdx = kvSIdx - preKVNum`），使 QK 领先 PV 一个 stackTile，实现"加载-计算-跨核"三线重叠。kernel 尾部 CUBE/VEC 两侧各自 `WaitFlag` 清理事件并 `PipeBarrier<PIPE_ALL>()` 收尾。
	300	+
	301	+## 9. 与 FAI / FD 系列的差异对比
	302	+
	303	+与 [FAI SplitRow](./02_fai_split_row_kernel.md)、[FD 系列](./04_fd_kernel.md) 同场竞技时的选型要点（更完整的七维对比见 FD 文档 §7）：
	304	+
	305	+\| 维度 \| XFAI（本系列） \| FAI SplitRow \| FD \|
	306	+\| --- \| --- \| --- \| --- \|
	307	+\| 模板组织 \| QK/PV 独立 BlockMmad + 双 Epilogue + CombineScale \| QK/PV/TailQK + Epilogue 一体化 \| XFAI 的深度演化版（dm 按周期分区等） \|
	308	+\| L1 布局 \| PV 显式接续 QK 的 L1 尾部（`l1BufAddrStart` 参数） \| 固定布局 \| 继承并细化 \|
	309	+\| RescaleO 类型参数 \| 5 参数（含 UpdateType_/LseType_） \| 4 参数 \| 5 参数 + isSplitkv 路由 \|
	310	+\| 分页 KV \| `PAGED_CACHE_FLAG_` 策略参数原生支持 \| 样例层支持 \| 同 XFAI \|
	311	+\| split-KV 合并 \| `CombineScale` 独立类（XFAI 专用） \| 无（SplitRow 免合并） \| `EpilogueAtlasA2CombineScale` 通用化 \|
	312	+
	313	+## 10. 使用示例（摘自 xllm-ops 真实工程）
	314	+
	315	+以下代码摘自 xllm_ops（https://gitcode.com/xLLM-AI/xllm_ops）`x_flash_attention_infer/op_kernel/x_flash_attention_infer.h`，演示五个模板的完整组装与调用。
	316	+
	317	+### 10.1 类型组装（FAInfer 入口）
	318	+
	319	+```cpp
	320	+using namespace AtlasA2T;
	321	+
	322	+// ---- QK GEMM：L1/L0 tile 与策略 ----
	323	+using L1TileShapeQK = Gemm::GemmShape<Q_TILE_CEIL, 128, 128>;
	324	+using L0TileShapeQK = Gemm::GemmShape<128, 128, 128>;
	325	+using DispatchPolicyQK = Gemm::MmadAtlasA2XFAIQK<PagedCacheFlag, false>;
	326	+using ElementS = float;
	327	+using BlockMmadQK = Gemm::BlockMmad<DispatchPolicyQK, L1TileShapeQK, L0TileShapeQK,
	328	+ QType, KType, ElementS>;
	329	+
	330	+// ---- PV GEMM ----
	331	+using L1TileShapePV = Gemm::GemmShape<128, 128, 256>;
	332	+using DispatchPolicyPV = Gemm::MmadAtlasA2XFAIPV<PagedCacheFlag, false>;
	333	+using ElementP = QType; // P 与 Q 同精度（bf16/fp16）
	334	+using BlockMmadPV = Gemm::BlockMmad<DispatchPolicyPV, L1TileShapePV, L1TileShapePV,
	335	+ ElementP, VType, ElementOTmp>;
	336	+
	337	+// ---- Epilogue：OnlineSoftmax 与 RescaleO ----
	338	+using EpilogueOnlineSoftmax = Epilogue::BlockEpilogue<
	339	+ Epilogue::EpilogueAtlasA2XFAIOnlineSoftmax<lseMode>, ElementP, ElementS, ElementMask>;
	340	+using EpilogueRescaleO = Epilogue::BlockEpilogue<
	341	+ Epilogue::EpilogueAtlasA2XFAIRescaleO<lseMode>, OType, ElementOTmp, ElementUpdate, LseType>;
	342	+
	343	+// ---- 顶层组装 ----
	344	+using FAInferKernel = FlashAttention::FAInferKernel<BlockMmadQK, BlockMmadPV,
	345	+ EpilogueOnlineSoftmax, EpilogueRescaleO, PagedCacheFlag, maskCategory, inLayout>;
	346	+```
	347	+
	348	+### 10.2 构造与跨核 flag
	349	+
	350	+```cpp
	351	+BlockMmadQK blockMmadQK(resource, nDynNum, kDynNum);
	352	+BlockMmadPV blockMmadPV(resource, nDynNum, kPVDynNum, L1_QK_SIZE); // L1 接续 QK
	353	+EpilogueOnlineSoftmax epilogueOnlineSoftmax(resource, scaleValue);
	354	+EpilogueRescaleO epilogueRescaleO(resource);
	355	+
	356	+Arch::CrossCoreFlag qkReady{READY_ID};
	357	+Arch::CrossCoreFlag softmaxReady{READY_ID};
	358	+Arch::CrossCoreFlag pvReady{READY_ID};
	359	+```
	360	+
	361	+### 10.3 CUBE 侧：QK → PV
	362	+
	363	+```cpp
	364	+// 软件流水：QK 领先 PV 一个 stackTile
	365	+for (uint32_t kvSIdx = 0; kvSIdx < kvSLoopNumTotal + preKVNum; kvSIdx += blockStackNum) {
	366	+ if (kvSIdx < kvSLoopNumTotal) {
	367	+ // Q 分组加载 + QK 主计算（paged 与非 paged 传参不同）
	368	+ blockMmadQK.loadQGM(gQ, layoutQ, rowNum, singleGroupHeads, qHeads);
	369	+ blockMmadQK(gQ, gK, gS, gBlockTable, layouts...,
	370	+ actualBlockShapeQK, kvSIdx, kvSLoopNumTotal, pagedBlockSize, strideK);
	371	+ CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady); // 通知 AIV：S 就绪
	372	+ }
	373	+ if (kvSIdx >= preKVNum) {
	374	+ uint32_t nowkvSIdx = kvSIdx - preKVNum;
	375	+ blockMmadPV(gP, gV, gOTmp, ..., blockStackNum, softmaxReady, nIdx, nLoop);
	376	+ CrossCoreSetFlag<0x2, PIPE_FIX>(pvReady); // 通知 AIV：OTmp 就绪
	377	+ }
	378	+}
	379	+```
	380	+
	381	+### 10.4 VECTOR 侧：OnlineSoftmax → RescaleO
	382	+
	383	+```cpp
	384	+if (causal) {
	385	+ // causal 分支：计算三角边界后走带 mask 的 13 参版本
	386	+ uint32_t triUp = noSkipKvS - qSBlockSize;
	387	+ uint32_t triDown = noSkipKvS;
	388	+ bool doTriUMask = triUp < kvSEndIdx - 1;
	389	+ epilogueOnlineSoftmax(gS, gP, gMask, ..., qkReady, triUp, triDown,
	390	+ kvSStartIdx, kvSEndIdx, ...);
	391	+} else {
	392	+ epilogueOnlineSoftmax(gS, gP, ...); // 无 mask 版
	393	+}
	394	+CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxReady); // 通知 AIC：P 就绪
	395	+
	396	+CrossCoreWaitFlag(pvReady); // 等 OTmp
	397	+epilogueRescaleO(gO, gOTmp, gOUpdate, gLse, layouts..., actualBlockShapePV,
	398	+ qSBlockSize, qNBlockSize,
	399	+ /isFirstStackTile=/(stackSeqCount - PRE_LAUNCH == 0),
	400	+ /isLastStackTile=/..., curStackTileMod);
	401	+```
	402	+
	403	+### 10.5 CombineScale 的两种调用形态
	404	+
	405	+形态一：FD kernel 核内调用（`x_flash_attention_infer_fd.h`）——`AscendC::SyncAll()` 后在 VEC 侧直接调用：
	406	+
	407	+```cpp
	408	+using CombineScale = Epilogue::Block::CombineScale<OType, LseType>; // FAInferKernelFD 第 5 模板参
	409	+...
	410	+combineScale(qHeads, extraInfo->totalSplitNodeNum, embedV, extraInfo,
	411	+ gmlse, gmlo, gO, gActualQseqlen, /inputLayoutTND=/true);
	412	+```
	413	+
	414	+形态二：独立 kernel（`x_attention/op_kernel/x_attention_catlass_helper.h` 的 `CallCombineScale`）——通用的 split-KV 后处理算子：
	415	+
	416	+```cpp
	417	+using Policy = Epilogue::EpilogueAtlasA2CombineScale;
	418	+using CombineScaleKernel = Epilogue::BlockEpilogue<Policy, OType(INPUT_T), float>;
	419	+// 在 vec 核 SyncAll 后调用，与主 kernel 解耦
	420	+```
	421	+
	422	+## 11. 小结
	423	+
	424	+XFAI 系列把推理 Flash Attention 的"AIC/AIV 跨核流水"沉淀为五个可组合模板：两个 GEMM 模板以 `l1BufAddrStart` 显式共享 CUBE 核 L1，两个 Epilogue 模板以 gm/gl/dm 标量链衔接 online softmax 递推，CombineScale 收口 split-KV 合并；三个 `CrossCoreFlag` 加 `preLoad=1` 软件流水构成完整的流水闭环。FD 系列在此基础上做了进一步性能演化（见 [04_fd_kernel.md](./04_fd_kernel.md)）。

Adocs/zh/2_Design/04_contrib/00_XAttention/04_fd_kernel.md+306-0

@@ -0,0 +1,306 @@
	1	+# FD 系列 Flash Attention 推理核设计文档
	2	+
	3	+## 1. 系列概述
	4	+
	5	+FD 系列是 XFAI 系列（见 [03_xfai_kernel.md](./03_xfai_kernel.md)）在 causal 三角 mask 场景下的增强实现。其核心思路是：
	6	+
	7	+- GEMM 层完全复用 XFAI：QK 与 PV 两个 BlockMmad 直接使用 `MmadAtlasA2XFAIQK` / `MmadAtlasA2XFAIPV`，不新增矩阵乘模板；
	8	+- Epilogue 层新增 FD 专属模板：`EpilogueAtlasA2OnlineSoftmax_FD`（在线 softmax，带参数化三角 mask）与 `EpilogueAtlasA2RescaleO_FD`（O 重缩放，支持 SplitKV），替换 XFAI 的对应 epilogue；
	9	+- 通用 CombineScale：`EpilogueAtlasA2CombineScale` 是一个不绑定 FD 的通用合并模板，负责 shared（causal 三角部分）与 unshared（非共享上下文部分）两路 softmax 中间结果的合并，在 FD 推理场景与 `EpilogueAtlasA2RescaleO_FD` 配合完成最终输出。
	10	+
	11	+FD 系列共 3 个新增模板：
	12	+
	13	+\| 模板 \| 类型 \| 文件 \|
	14	+\| --- \| --- \| --- \|
	15	+\| `EpilogueAtlasA2OnlineSoftmax_FD<LSE_MODE_>` \| BlockEpilogue \| `include/catlass/epilogue/block/block_epilogue_online_softmax_FD.hpp` \|
	16	+\| `EpilogueAtlasA2RescaleO_FD<LSE_MODE_>` \| BlockEpilogue \| `include/catlass/epilogue/block/block_epilogue_rescale_o_FD.hpp` \|
	17	+\| `EpilogueAtlasA2CombineScale` \| BlockEpilogue \| `include/catlass/epilogue/block/block_epilogue_combine_scale.hpp` \|
	18	+
	19	+其中 OnlineSoftmax_FD 与 CombineScale 在文件末尾均提供了 `EpilogueAscend950*` 同名特化（直接继承 AtlasA2 版本），用于 950 平台的 Policy 名注册。
	20	+
	21	+## 2. DispatchPolicy 定义
	22	+
	23	+```cpp
	24	+// 1. FD 在线 softmax epilogue（QK 后：行归约 + P 下搬运）
	25	+template <LseMode LSE_MODE_ = LseMode::OUT_AND_LSE>
	26	+struct EpilogueAtlasA2OnlineSoftmax_FD {
	27	+ static constexpr LseMode LSE_MODE = LSE_MODE_;
	28	+ static constexpr EpilogueType TYPE = EpilogueType::ATLAS_A2;
	29	+};
	30	+
	31	+// 2. FD 输出重缩放 epilogue（PV 后：O 缩放累加 + LSE 写出）
	32	+template <LseMode LSE_MODE_ = LseMode::OUT_AND_LSE>
	33	+struct EpilogueAtlasA2RescaleO_FD {
	34	+ static constexpr LseMode LSE_MODE = LSE_MODE_;
	35	+ static constexpr EpilogueType TYPE = EpilogueType::ATLAS_A2;
	36	+};
	37	+
	38	+// 3. 通用两路合并 epilogue（shared/unshared softmax 结果合并）
	39	+struct EpilogueAtlasA2CombineScale {
	40	+ static constexpr EpilogueType TYPE = EpilogueType::ATLAS_A2;
	41	+};
	42	+
	43	+// 950 平台特化（继承 AtlasA2 实现）
	44	+template <LseMode LSE_MODE_>
	45	+struct EpilogueAscend950OnlineSoftmax_FD : EpilogueAtlasA2OnlineSoftmax_FD<LSE_MODE_> {};
	46	+struct EpilogueAscend950CombineScale : EpilogueAtlasA2CombineScale {};
	47	+```
	48	+
	49	+## 3. EpilogueAtlasA2OnlineSoftmax_FD 设计
	50	+
	51	+### 3.1 类形态与构造
	52	+
	53	+```cpp
	54	+template <typename Policy_, typename OutputType_, typename InputType_, typename MaskType_>
	55	+class BlockEpilogue<EpilogueAtlasA2OnlineSoftmax_FD<LSE_MODE_>, OutputType_, InputType_, MaskType_> {
	56	+public:
	57	+ BlockEpilogue(Resource *resource, float scaleValue_); // scaleValue = softmax 缩放因子 1/sqrt(d)
	58	+};
	59	+```
	60	+
	61	+与 XFAI OnlineSoftmax 一致，构造时传入 softmax scale；区别在于 operator() 提供了带三角 mask 参数的版本。
	62	+
	63	+### 3.2 UB 空间布局
	64	+
	65	+\| 区域 \| 偏移（单位：fp32 元素） \| 大小 \| 用途 \|
	66	+\| --- \| --- \| --- \| --- \|
	67	+\| ls \| 0 \| 8192（乒乓） \| S 子块行 softmax 中间量（分核暂存） \|
	68	+\| lp / mask / mask32 \| 416384 \| 共 416384 \| P 中间量 / mask 比特 / mask 展开 \|
	69	+\| tv \| 10*16384 \| - \| 向量计算临时区 \|
	70	+\| lm / hm / gm / gl / dm \| 1016384 + {0..5}1024 \| 各 1KB \| 行内 max / 全局 max / 全局 sum / 重缩放系数 \|
	71	+\| mask16 \| 11*16384 \| - \| 16bit mask 暂存 \|
	72	+
	73	+关键点：dm（重缩放系数区）按 stackTile 周期分区，`dmUbOffsetCurCycle = curStackTileMod * MAX_ROW_NUM_SUB_CORE(256) + rowOffset`，即每个 stackTile 周期拥有独立的 256 行 dm 槽位，避免跨周期覆盖。
	74	+
	75	+### 3.3 行归约三分支
	76	+
	77	+按行块长度选择归约策略（与 XFAI 相同）：
	78	+
	79	+- SPECTILE512：三次级联 BlockReduceMax（512→256→…→1），适用于满行宽；
	80	+- SPECTILE256：`SetVecMask(32)` + `SetBlockReduceMask(4)` 的短行归约；
	81	+- TAILTILE：整段归约 + 尾段 `SetMask` 处理非对齐行。
	82	+
	83	+### 3.4 SubCoreCompute 六步流程（doTriUMask 模板参数）
	84	+
	85	+```cpp
	86	+template <bool doTriUMask>
	87	+void SubCoreCompute(...) {
	88	+ // ① CalcLocalRowMax: 行内局部 max（lm）
	89	+ // ② UpdateGlobalRowMax: hm = max(lm, gm); dm = exp(gm - hm) ← 产生本周期重缩放系数
	90	+ // ③ CalcExp: p = exp(s * scale - hm)
	91	+ // ④ WaitFlag(V_MTE2)（非三角 mask 时）；DownCastP: p 降精度到 P 类型 + SetFlag(V_MTE3)
	92	+ // ⑤ CalcLocalRowSum: gl += Σ exp(...)
	93	+ // ⑥ WaitFlag(V_MTE3) + CopyPUbToGm: P 写回 GM 供 PV GEMM 消费
	94	+}
	95	+```
	96	+
	97	+### 3.5 双 operator() 与三角 mask 机制
	98	+
	99	+提供两个调用入口：
	100	+
	101	+1. 无 mask 版：`SubCoreCompute<false>`，用于非 causal 或整块免 mask 的 stackTile；
	102	+2. 带 mask 版：`SubCoreCompute<true>`，额外参数 `triUp / triDown / kvSStartIdx / kvSEndIdx / qkReady`。
	103	+
	104	+三角 mask 偏移计算（causal 语义）：
	105	+
	106	+```
	107	+if (triUp >= kvSStartIdx) {
	108	+ maskStart = RoundDown(triUp - kvSStartIdx, BLOCK); // 三角起点折算到本 stackTile 内偏移
	109	+} else {
	110	+ 全列有效（无需 mask）
	111	+}
	112	+if (triDown < kvSEndIdx) { maskEnd = ...; } else { 全列有效; }
	113	+```
	114	+
	115	+即：triUp 为左上三角起点（query 相对位置），triDown 为右下止点，二者把当前 KV stackTile 划分为「全 mask / 部分三角 / 全有效」三段，仅部分三角段执行逐列 mask 计算。
	116	+
	117	+## 4. EpilogueAtlasA2RescaleO_FD 设计
	118	+
	119	+### 4.1 类形态与构造
	120	+
	121	+```cpp
	122	+template <typename Policy_, typename OutputType_, typename InputType_, typename UpdateType_, typename LseType_>
	123	+class BlockEpilogue<EpilogueAtlasA2RescaleO_FD<LSE_MODE_>, OutputType_, InputType_, UpdateType_, LseType_> {
	124	+public:
	125	+ BlockEpilogue(Resource *resource); // 注意：无 scaleValue 参数
	126	+};
	127	+```
	128	+
	129	+比 OnlineSoftmax 多两个类型参数：`UpdateType_`（中间结果 gOUpdate 的精度，通常 fp32）与 `LseType_`（LSE 输出精度）。
	130	+
	131	+### 4.2 UB 空间布局
	132	+
	133	+\| 区域 \| 偏移（fp32 元素） \| 用途 \|
	134	+\| --- \| --- \| --- \|
	135	+\| lo \| 6*16384 \| 上一周期 O 暂存（旧 O） \|
	136	+\| go \| 8*16384 \| 累计 O（16/32bit 双视图） \|
	137	+\| tv \| 10*16384 \| 向量临时区 \|
	138	+\| hm / gm \| 1016384 + {9,10}1024 \| 行 max \|
	139	+\| gl 与 lse32 共享 \| 1016384 + 121024 \| 行 sum / LSE 中间量 \|
	140	+\| dm \| 1016384 + 131024 \| 重缩放系数（与 OnlineSoftmax 的 dm 分区对齐） \|
	141	+
	142	+### 4.3 核心 algorithm：首 / 中 / 末 tile 三分支
	143	+
	144	+对每个 stackTile 周期：
	145	+
	146	+```
	147	+非首 tile:
	148	+ WaitFlag(V_MTE3, EVENT_ID3) → DataCopy lo(旧O) → SetFlag(MTE2_V, EVENT_ID0)
	149	+ Brcb dm 广播到 tv
	150	+ go = go * dm_block // 旧累计缩放
	151	+ WaitFlag(EVENT_ID0) 后 go = lo + go // 累加当前 PV 结果
	152	+ SetFlag(V_MTE3, EVENT_ID3)
	153	+
	154	+首 tile:
	155	+ go = lo（直接 DataCopy，无缩放）
	156	+
	157	+末 tile (isLastStackTile):
	158	+ Brcb gl 广播 → go = go / gl // 归一化
	159	+ Cast 到输出精度（bf16 用 CAST_RINT，否则 CAST_NONE；仅 !isSplitkv）
	160	+ !isSplitkv: CopyOToGm → gOutput（bf16）
	161	+ isSplitkv: CopyOToGmFp32 → gCombineo（fp32，不 cast，供后续 CombineScale 合并）
	162	+```
	163	+
	164	+### 4.4 CopyOToGm 三段式搬运
	165	+
	166	+`CopyOToGm` / `CopyOToGmFp32` 均按 prologue（前缀 token）/ integral（整 head）/ epilogue（尾 token）三段执行 `DataCopyPad`：每行搬运 `embed` 列有效数据 + `oHiddenSize - embed` 列 pad。行切分由 `rowNumTile = RoundDown(8192/embed, 8)` 决定，行循环内按 token-head 折算 `proTokenIdx / proTokenNum / integralHeadNum / epiTokenNum`。
	167	+
	168	+### 4.5 LSE 处理与中间结果回写
	169	+
	170	+- LSE_OUT 模式：`isLastRowLoop` 时 `lse = ln(gl) + gm`，Brcb 广播后写 `gLse`（DataCopyPad，带 `(qHeads-1)*4` 列 pad）；
	171	+- isSplitkv：仅写 `gCombineLse`（fp32 中间量，供 CombineScale 合并）；
	172	+- needRowLoop 且非末 tile：`goUbTensor32` DataCopy 回 `gOUpdate`（fp32 中间结果回写 GM，供下个周期读入 lo）。
	173	+
	174	+事件对：EVENT_ID0/1/3/4/5/6，其中 `MTE3_MTE2(EVENT_ID6)` 作为跨周期栅栏。
	175	+
	176	+## 5. EpilogueAtlasA2CombineScale 设计
	177	+
	178	+### 5.1 类形态
	179	+
	180	+```cpp
	181	+template <typename Policy_, typename OutputType_, typename InputType_>
	182	+class BlockEpilogue<EpilogueAtlasA2CombineScale, OutputType_, InputType_> {
	183	+public:
	184	+ BlockEpilogue(Resource *resource);
	185	+};
	186	+```
	187	+
	188	+仅两个类型参数（无 mask / lse 模板参数），是纯通用合并模板。
	189	+
	190	+### 5.2 UB 空间布局（按字节）
	191	+
	192	+\| 区域 \| 偏移 \| 说明 \|
	193	+\| --- \| --- \| --- \|
	194	+\| sharedOut \| 0 - 64k \| shared 路 O 中间量（32k 乒乓） \|
	195	+\| unsharedOut \| 64k - 128k \| unshared 路 O 中间量 \|
	196	+\| sharedGl / unsharedGl \| 128k / 132k \| 两路行 sum \|
	197	+\| sharedGm / unsharedGm \| 136k / 140k \| 两路行 max \|
	198	+\| realGm / realGl \| 144k / 148k \| 合并后全局 max / sum \|
	199	+\| out \| 152k - 184k \| 输出暂存 \|
	200	+
	201	+### 5.3 合并算法（五步）
	202	+
	203	+```
	204	+① BlockReduceMax(8→1) 压缩 sharedGm/Gl（shared 路每行按 SOFTMAX_BROAD_SIZE=8 重复存放）
	205	+② realGm = max(sharedGm, unsharedGm)
	206	+③ α = exp(sharedGm - realGm); β = exp(unsharedGm - realGm)
	207	+ gl = sharedGl * α + unsharedGl * β
	208	+④ Brcb 广播 α/β → Out = sharedOut * α + unsharedOut * β
	209	+⑤ Brcb realGl 广播 → Out = Out / realGl（归一化）
	210	+→ Cast 到 ElementOutput（bf16 用 CAST_RINT）→ DataCopy gFinalOutput
	211	+```
	212	+
	213	+### 5.4 流水组织
	214	+
	215	+主循环 `rowLoopNum + preLoad=1` 双缓冲（pingpongFlag 切换）：加载段 DataCopy 两路 gm/gl（shared 路每行 8 个，`sumMaxOffsetIoShared = row * 8`；unshared 路整行，非 8 对齐用 DataCopyPad）与两路 Out；计算段执行上述五步。事件对 `MTE3_MTE2` 乒乓 + `EVENT_ID4`（MTE2_V / V_MTE3）。
	216	+
	217	+## 6. 使用示例
	218	+
	219	+以下摘自 xllm-ops（https://gitcode.com/xLLM-AI/xllm_ops）真实工程 `x_flash_attention_infer/op_kernel/x_flash_attention_infer_fd.h`。
	220	+
	221	+### 6.1 类型组装（FDInfer 入口）
	222	+
	223	+```cpp
	224	+// GEMM 层：复用 XFAI Policy
	225	+using DispatchPolicyQK = MmadAtlasA2XFAIQK<PagedCacheFlag, false>;
	226	+using DispatchPolicyPV = MmadAtlasA2XFAIPV;
	227	+
	228	+// Epilogue 层：FD 专属 + 通用 CombineScale
	229	+using EpilogueOnlineSoftmax = BlockEpilogue<
	230	+ EpilogueAtlasA2OnlineSoftmax_FD<lseMode>, PType, SType, maskType>;
	231	+using EpilogueRescaleO = BlockEpilogue<
	232	+ EpilogueAtlasA2RescaleO_FD<lseMode>, OType, OTmpType, OUpdateType, LseType>;
	233	+using CombineScale = Epilogue::Block::CombineScale<OType, LseType>; // 预组装别名
	234	+
	235	+// 组装 FD 推理核
	236	+template <...> using FAInferKernelFD = ...;
	237	+```
	238	+
	239	+### 6.2 构造与 causal 三分支调用（vec 核）
	240	+
	241	+```cpp
	242	+// 构造：OnlineSoftmax 带 scale，其余仅传 resource
	243	+EpilogueOnlineSoftmax epilogueOnlineSoftmax(resource, scaleValue);
	244	+EpilogueRescaleO epilogueRescaleO(resource);
	245	+CombineScale combineScale(resource);
	246	+
	247	+// causal 分支参数
	248	+int32_t triUp = noSkipKvS - qSBlockSize; // 左上三角起点
	249	+int32_t triDown = noSkipKvS; // 右下止点
	250	+int32_t kvSStartIdx = kvSIdx * pagedBlockSize;
	251	+bool doTriUMask = triUp < kvSEndIdx - 1; // mask 长度仅 1 时相当于不加
	252	+
	253	+if (doTriUMask) {
	254	+ // 带 mask 版：13 参数，含 qkReady/triUp/triDown/kvSStartIdx/kvSEndIdx
	255	+ epilogueOnlineSoftmax(gS, gPUb, ..., qkReady, triUp, triDown, kvSStartIdx, kvSEndIdx, ...);
	256	+} else {
	257	+ // 无 mask 版
	258	+ epilogueOnlineSoftmax(gS, gPUb, ..., isLastStackTile, ...);
	259	+}
	260	+// softmaxReady 通过 CrossCoreSetFlag<0x2, PIPE_MTE3> 通知 PV GEMM
	261	+```
	262	+
	263	+### 6.3 RescaleO 调用（PV 后，vec 核）
	264	+
	265	+```cpp
	266	+Arch::CrossCoreWaitFlag(pvReady); // 等待 PV GEMM 完成
	267	+epilogueRescaleO(
	268	+ gO[gmOffsetO], gOTmp[gmOffsetOTmp], gOUpdate[gmOffsetUpdate],
	269	+ gLse[gmOffsetLse], gmlse[gmlse0ffset], gmlo[gmlooffset],
	270	+ layoutO, layoutOTmp, layoutUpdate, layoutLse,
	271	+ actualBlockShapePV, qSBlockSize, qNBlockSize,
	272	+ (stackSeqCount - PRE_LAUNCH == 0), // isFirstStackTile
	273	+ nowkvSIdx + blockStackNum >= kvEnd, // isLastStackTile
	274	+ curStackTileMod, isSplitKV,
	275	+ layoutgmLse, layoutgmLo); // splitkv 中间量布局
	276	+```
	277	+
	278	+### 6.4 CombineScale 调用（全核 SyncAll 后）
	279	+
	280	+```cpp
	281	+AscendC::SyncAll();
	282	+combineScale(
	283	+ qHeads, // 头数
	284	+ extraInfo->totalSplitNodeNum, // split 槽数
	285	+ embedV, // V 头维度
	286	+ extraInfo, // SplitInfo（TilingData 中定义）
	287	+ gmlse, // LSE 中间量
	288	+ gmlo, // O 中间量（workspace）
	289	+ gO, // 最终输出
	290	+ gActualQseqlen, // 实际 seq len
	291	+ true);
	292	+```
	293	+
	294	+另见 `x_attention/op_kernel/x_attention_catlass_helper.h` 中 `CallCombineScale`：以 `EpilogueAtlasA2CombineScale` + `BlockEpilogue<Policy, OutputType(INPUT_T), InputType(float)>` 组装独立 CombineScaleKernel，在 vec 核 SyncAll 后调用。
	295	+
	296	+## 7. 与 XFAI 系列差异对比
	297	+
	298	+\| 维度 \| XFAI（EpilogueAtlasA2OnlineSoftmax/RescaleO） \| FD（本系列） \|
	299	+\| --- \| --- \| --- \|
	300	+\| mask 能力 \| token mask 布尔式（整块生效） \| triUp/triDown 参数化三角 mask（逐列生效） \|
	301	+\| operator() 入口 \| 单一版本 \| 无 mask 版 + 带 mask 版（doTriUMask 模板分支） \|
	302	+\| dm 存储 \| 单周期覆盖 \| 按 stackTile 周期分区（256 行/周期） \|
	303	+\| RescaleO 类型参数 \| 4 个 \| 5 个（新增 UpdateType_/LseType_） \|
	304	+\| SplitKV 输出路由 \| gCombineo/gCombineLse \| isSplitkv 分支：fp32→gCombineo / bf16→gOutput \|
	305	+\| CombineScale \| SplitKV N-way 合并（RescaleO 内） \| 独立通用模板，shared/unshared 两路合并 \|
	306	+\| GEMM Policy \| XFAI 专属 \| 完全复用 XFAI（无新增 GEMM） \|

Adocs/zh/2_Design/04_contrib/00_XAttention/05_xa_tla_kernel.md+374-0

@@ -0,0 +1,374 @@
	1	+# XA TLA 系列 x_attention 推理核设计文档
	2	+
	3	+## 1. 系列概述
	4	+
	5	+XA TLA 系列是面向 Ascend 950（Arch::Ascend950）x_attention 推理场景的 TLA 指令实现。与 FD/XFAI 系列（Atlas A2 硬件路径）不同，本系列：
	6	+
	7	+- 直接特化 `BlockMmadTla`，使用 Cube 侧 TLA 硬件指令（`MmadTla`）完成矩阵乘；
	8	+- 面向 shared / unshared 双路 KV 架构：shared 路处理 batch 间共享的系统前缀 KV（流式 online softmax），unshared 路处理每 batch 独立的解码 KV（Paged KV Cache + 逐步 mask），最终由 Combine 路合并两路 partial 输出；
	9	+- AIC（Cube）与 AIV（Vector）通过跨核 flag 深度协同，形成多级软件流水。
	10	+
	11	+XA TLA 系列共 8 个新增模板（4 GEMM + 4 Epilogue）：
	12	+
	13	+\| 模板 \| 类型 \| 文件 \|
	14	+\| --- \| --- \| --- \|
	15	+\| `MmadXASharedQK` \| BlockMmadTla \| `include/catlass/gemm/block/block_mmad_tla_xa_shared_qk_ascend950.hpp` \|
	16	+\| `MmadXAUnsharedQK` \| BlockMmadTla \| `include/catlass/gemm/block/block_mmad_tla_xa_unshared_qk_ascend950.hpp` \|
	17	+\| `MmadXASharedPV` \| BlockMmadTla \| `include/catlass/gemm/block/block_mmad_tla_xa_shared_pv_ascend950.hpp` \|
	18	+\| `MmadXAUnsharedPV` \| BlockMmadTla \| `include/catlass/gemm/block/block_mmad_tla_xa_unshared_pv_ascend950.hpp` \|
	19	+\| `EpilogueXASharedSoftmax` \| BlockEpilogue \| `include/catlass/epilogue/block/block_epilogue_xa_shared_softmax_ascend950.hpp` \|
	20	+\| `EpilogueXAUnsharedSoftmax` \| BlockEpilogue \| `include/catlass/epilogue/block/block_epilogue_xa_unshared_softmax_ascend950.hpp` \|
	21	+\| `EpilogueXASharedRescaleO` \| BlockEpilogue \| `include/catlass/epilogue/block/block_epilogue_xa_shared_rescale_ascend950.hpp` \|
	22	+\| `EpilogueXACombineScale` \| BlockEpilogue \| `include/catlass/epilogue/block/block_epilogue_xa_combine_scale_ascend950.hpp` \|
	23	+
	24	+三路整体数据流：
	25	+
	26	+```mermaid
	27	+graph LR
	28	+ subgraph Shared路
	29	+ A[Q] --> QK1[MmadXASharedQK]
	30	+ K1[K shared] --> QK1
	31	+ QK1 --> SM1[SharedSoftmax<br/>online softmax]
	32	+ SM1 -->\|P 直写 L1\| PV1[MmadXASharedPV]
	33	+ V1[V shared] --> PV1
	34	+ PV1 --> RC[SharedRescaleO<br/>O 累积]
	35	+ end
	36	+ subgraph Unshared路
	37	+ A2[Q] --> QK2[MmadXAUnsharedQK]
	38	+ K2[K unshared<br/>Paged KV] --> QK2
	39	+ QK2 --> SM2[UnsharedSoftmax<br/>mask+softmax]
	40	+ SM2 -->\|P 直写 L1\| PV2[MmadXAUnsharedPV]
	41	+ V2[V unshared] --> PV2
	42	+ PV2 --> O2[partial O 直写 GM]
	43	+ end
	44	+ RC -->\|O/max/sum 写 GM\| CB[EpilogueXACombineScale]
	45	+ O2 -->\|O/max/sum 写 GM\| CB
	46	+ CB --> OUT[最终 Attention 输出]
	47	+```
	48	+
	49	+> Shared 路通过流式 online softmax（isFirstKv/isLastKv 控制）直接在 UB 中累积出 partial O，仅在最后一块 KV 时写出；Unshared 路每个任务独立完成一次完整 softmax（含变长 mask），partial O 直写 GM。两路各写一份 O/max/sum 到 GM，最终由 CombineScale 按 LSE 数学完成合并。
	50	+
	51	+## 2. 注册机制：xa_register.hpp
	52	+
	53	+XA TLA 系列在工程侧（xllm-ops 的 `common/catlass/include/catlass_patch/xa_register.hpp`）通过"空壳 Policy + 条件 include"的方式注册：
	54	+
	55	+```cpp
	56	+// Gemm 命名空间：4 个空壳 Policy（仅携带 ArchTag 与 STAGES 元信息）
	57	+namespace Catlass::Gemm {
	58	+template <typename ArchTag_>
	59	+struct MmadXASharedQK : MmadBase<ArchTag_, false> {
	60	+ static constexpr int32_t STAGES = 2;
	61	+};
	62	+// MmadXAUnsharedQK / MmadXASharedPV / MmadXAUnsharedPV 同理
	63	+} // namespace Catlass::Gemm
	64	+
	65	+// Epilogue 命名空间：4 个空壳结构体
	66	+namespace Catlass::Epilogue {
	67	+template <typename ArchTag_>
	68	+struct EpilogueXASharedSoftmax { using ArchTag = ArchTag_; };
	69	+// EpilogueXAUnsharedSoftmax / EpilogueXASharedRescaleO / EpilogueXACombineScale 同理
	70	+} // namespace Catlass::Epilogue
	71	+
	72	+// 950 架构下引入全部 8 个 block 特化实现
	73	+#if CATLASS_ARCH == 3510
	74	+#include "catlass/gemm/block/block_mmad_tla_xa_shared_qk_ascend950.hpp"
	75	+// ... 其余 7 个头文件
	76	+#endif
	77	+```
	78	+
	79	+使用要点：
	80	+
	81	+1. 必须 include 在 catlass 聚合头之后：空壳 Policy 结构体只做命名空间占位，真正的 `BlockMmadTla<Policy<...>>` / `BlockEpilogue<Policy<...>>` 偏特化由上述 8 个头文件补齐；
	82	+2. 通过 `-I` 搜索顺序覆盖：工程 CMake 中 `common/catlass/include` 路径优先于第三方 `catlass/include`，使第三方子模块保持 xa-free 状态，XA 扩展全部收敛在项目自有代码内；
	83	+3. 该机制是把"未提交到 CATLASS 主干的扩展"以工程 patch 形式落地的标准范式。
	84	+
	85	+## 3. GEMM 模板设计
	86	+
	87	+### 3.1 共性骨架
	88	+
	89	+4 个 GEMM 均为如下偏特化形式：
	90	+
	91	+```cpp
	92	+template <typename ArchTag_, typename L1TileShape_, typename L0TileShape_,
	93	+ typename ElementA_, typename ElementB_, typename ElementC_, typename ElementBias_,
	94	+ typename TileCopy_, typename TileMmad_>
	95	+class BlockMmadTla<MmadXASharedQK<ArchTag_>, L1TileShape_, L0TileShape_,
	96	+ ElementA_, ElementB_, ElementC_, ElementBias_, TileCopy_, TileMmad_> {
	97	+ // STAGES = 2；双缓冲 L1A / L1B / L0A / L0B / L0C
	98	+};
	99	+```
	100	+
	101	+- 三级搬运：GM → L1（`CopyGm2L1`）→ L0（`CopyL12L0`）→ TLA 矩阵乘（`MmadTla`）→ L0C →（由 epilogue 或上层消费 L0C/UB 结果）；
	102	+- 双缓冲事件 ID 规律（`i = 0/1` 为槽位）：
	103	+ - `l1BEvent = BLOCK_EVENT_ID + i + STAGES`、`l0AEvent = i`、`l0BEvent = i + STAGES`、`l0CEvent = BLOCK_EVENT_ID + i`；
	104	+ - 构造函数中对全部事件 `SetFlag` 预置，使首轮搬运无需等待；
	105	+- QK/PV 分组事件段：QK 类 GEMM `BLOCK_EVENT_ID = 0`，PV 类 `BLOCK_EVENT_ID = 4`，同一 kernel 内两组 GEMM 互不干扰；
	106	+- 跨核同步：`SYNC_MODE = 4`，通过 `CrossCoreWaitFlag<SYNC_MODE, PIPE_FIX>`（+16 偏移的 AIV1 事件）与 AIV 侧同步，等待 epilogue 释放 UB/L1 资源的 flag。
	107	+
	108	+### 3.2 差异矩阵
	109	+
	110	+\| 维度 \| MmadXASharedQK \| MmadXAUnsharedQK \| MmadXASharedPV \| MmadXAUnsharedPV \|
	111	+\| --- \| --- \| --- \| --- \| --- \|
	112	+\| BLOCK_EVENT_ID \| 0 \| 0 \| 4 \| 4 \|
	113	+\| A 矩阵供给 \| Q 仅 isFirstKv 时 GM→L1，跨 kv 块复用 \| 每次调用 GM 加载 \| P 常驻 L1（构造不分配 l1A） \| P 从传入 L1 tensor 取 \|
	114	+\| B 矩阵供给 \| GM→L1 \| GM→L1 \| GM→L1 \| GM→L1 \|
	115	+\| TileShape 约束 \| 三轴可不同 \| M/K 轴相同；L1_N 是 L0_N 整数倍 \| L1/L0 三轴必须相同 \| M/N 轴相同；L1_K 是 L0_K 整数倍 \|
	116	+\| 内循环切分 \| 单次 mmad（L1 粒度） \| nLoops：按 L0_TILE_N 切 N 轴 \| 单次 mmad \| kLoops：按 L0_TILE_K 切 K 轴（kIdx==0 时 init） \|
	117	+\| L0C 空间 \| 常规 \| 按 L1_TILE_N 计 L0C 大小 \| 常规 \| 复用 QK 的 128×256 L0C 前 128×128 区域 \|
	118	+\| 跨核 flag \| QK_UB_RELEASE_FLAG（uint64_t） \| QK_UB_RELEASE_FLAG（uint16_t） \| PV_UB_RELEASE_FLAG（uint16_t） \| 无 \|
	119	+\| operator() 附加参数 \| isFirstKv / isLastKv / releaseFlag \| releaseFlag / taskIdL0C \| releaseFlag \| 无（三 taskId） \|
	120	+
	121	+### 3.3 MmadXASharedQK：跨 KV 块 Q 复用
	122	+
	123	+- Q（beam 维 M 轴）对同一 (batch, qHead) 的所有 KV 块不变，因此仅在 `isFirstKv` 时执行 `CopyGm2L1` 加载 Q 到 L1A；后续 KV 块直接复用，省去重复搬运；
	124	+- `isLastKv` 时翻转 l1A 双缓冲槽位（`l1AEvent = 1 - l1AEvent`），保证下一任务的 Q 加载与当前任务的消费不冲突；
	125	+- K（N 轴 KV 块）每次从 GM 加载至 L1B；
	126	+- 结果 S 写入 UB 的 `qkTensorList[taskIdMod2]` 双缓冲，通过 `QK_UB_RELEASE_FLAG`（uint64_t）向 AIV 侧授权消费。
	127	+
	128	+### 3.4 MmadXAUnsharedQK：N 轴内循环
	129	+
	130	+- 每次调用 Q/K 均从 GM 加载（不同任务 Q 不同）；
	131	+- L1 一次装载 `L1_TILE_N` 列，L0 仅 `L0_TILE_N` 列，`nLoops = L1_TILE_N / L0_TILE_N` 次内循环：
	132	+ - 每轮 `CopyL12L0` 搬 B 的 L0 列块 → `MmadTla` 累加至同一 L0C 槽；
	133	+ - `nLoops` 结束后 L0C → UB；
	134	+- 适配 unshared 路径 `blockKvLen = groupCountPerLoop × maxDecodeStep` 大 N 块（如 128×256）场景；
	135	+- `operator()` 为 `(tensorA, tensorB, tensorC, actualShape, releaseFlag, taskIdL0A, taskIdL0B, taskIdL0C)` 三 taskId 形式。
	136	+
	137	+### 3.5 MmadXASharedPV：P 常驻 L1
	138	+
	139	+- A 矩阵 P（softmax 输出）由 epilogue 直写 L1（`CopyUb2L1Tla`），因此构造函数不分配 l1A 缓冲，仅分配 l1B/l0A/l0B/l0C；
	140	+- L1/L0 三轴 TileShape 必须一致（P 在 L1 中按 L1 粒度整块布局）；
	141	+- V（B 矩阵）每次 GM→L1→L0；
	142	+- 计算结果 O_tmp 写入 `pvTensorList[taskIdMod2]` UB 双缓冲，`PV_UB_RELEASE_FLAG` 由 SharedRescaleO epilogue 消费后置位释放。
	143	+
	144	+### 3.6 MmadXAUnsharedPV：K 轴内循环 + L0C 复用
	145	+
	146	+- P（A）从传入的 L1 tensor 获取（epilogue 直写），V（B）GM→L1→L0；
	147	+- `kLoops = L1_TILE_K / L0_TILE_K` 次 K 轴内循环切分，`kIdx == 0` 时 `init=false`（首矩阵），其后 `init=true` 累加；
	148	+- L0C 复用：`SHARED_L0C_STAGE_SIZE = L1_TILE_M × L1_TILE_K × sizeof(acc)`，直接复用 QK GEMM 已申请的 128×256 L0C 区域的前 128×128 子区，PV 构造时 L0C 指针指向该区域，节省 L0C 总量（L0C 有限，QK 与 PV 分时复用是 950 上常见手法）；
	149	+- 无跨核 release flag：O_tmp 结果由 AIC 直接经 FIX 通道写 GM（unshared 路 partial O 直写）。
	150	+
	151	+## 4. Epilogue 模板设计
	152	+
	153	+### 4.1 EpilogueXASharedSoftmax：流式在线 softmax
	154	+
	155	+```cpp
	156	+template <typename Policy_, typename L1TileShape_, typename PType_, typename SType_>
	157	+class BlockEpilogue<EpilogueXASharedSoftmax<...>, ...> {
	158	+ BlockEpilogue(Resource *resource, uint32_t &ubBufAddrStart,
	159	+ float scaleValue, int32_t qHeads);
	160	+ // operator() 共 18 参数
	161	+};
	162	+```
	163	+
	164	+- UB 资源：`pNzOutTensorList[2]`（P 的 NZ 格式乒乓）+ `maxBrcb/sumBrcb` 广播中间量；
	165	+- operator() 关键参数：五组 softmax 标量（`lastExpSum/nowExpSum/nowExpMax/lastMax/nowMax`，三缓冲随 taskIdMod3 轮转）、三组 flag（`SYNC_QK_READY / SYNC_SOFTMAX_READY / QK_UB_RELEASE`）、`isUpdate / isLastKv`；
	166	+- 执行流程：
	167	+ 1. `CrossCoreWaitFlag(SYNC_QK_READY)` 等 AIC 完成 QK；
	168	+ 2. `ComputeScaleandMax`：S×scale 后行 ReduceMax，`UpdateMax` 融合历史 max（`nowMax = max(lastMax, rowMax)`）；
	169	+ 3. `ComputeExpSubSum`：S 经 `exp(nowMax - S)` 得 P，以 NZ 格式（Half 非 16 对齐的奇偶分离布局）经 `DATA_BLOCK_COPY` 写 UB，再求行 exp 部分和；
	170	+ 4. `CrossCoreSetFlag(QK_UB_RELEASE)` 释放 QK 的 UB 双缓冲；
	171	+ 5. `CopyUb2L1Tla`：P 直接从 UB 写入 L1 的 PV A 矩阵区域；
	172	+ 6. `CrossCoreSetFlag(SYNC_SOFTMAX_READY)` 通知 AIC 启动 PV；
	173	+ 7. `UpdateExpSumAndExpMax`：`nowExpSum = lastExpSum×exp(lastMax-nowMax) + curSum`；
	174	+ 8. `isLastKv` 时 `CopyOutMaxAndSum`：Brcb + `DataCopyPad` 按 `qHeads` stride 把每行 max/sum 写 GM（供 CombineScale 消费）。
	175	+
	176	+### 4.2 EpilogueXAUnsharedSoftmax：mask + 独立 softmax
	177	+
	178	+```cpp
	179	+BlockEpilogue(Resource *resource, uint32_t &ubBufAddrStart, float scaleValue,
	180	+ int32_t unsharedKvLen, int32_t maxDecodeStep,
	181	+ int32_t groupCountPerLoop, int32_t groupSize);
	182	+```
	183	+
	184	+- mask 机制（核心差异）：unshared 路 KV 为变长解码步，构造时 `InitUnsharedMask` 按 `groupCountPerLoop / groupSize / maxDecodeStep` 用 `Duplicate` 预生成 uint8 mask（0/1），加载到 `maskUbTensor`（HALF_MASK_BLOCK_SIZE）；有效 KV 长度内为 1，超出为 0；支持双 SubBlock 分工（`UINT8_BLOCK_SIZE=256` 对齐 + floorSub 处理奇数行拆分）；
	185	+- ComputeMaskandScale：`Select(mask, S×scale, MIN_VALUE=-3e38)` 后行 ReduceMax，实现变长 KV 屏蔽；N 轴分档 `N128 / N0_64 / N65_127` 三种分支处理；mask 本身经 `LoadAlign<MASK>` 的 `pregCompare` 加载；
	186	+- 每步独立：无 `isUpdate/isLastKv` 流式逻辑，每个任务一次完整 softmax（exp/max/sum 均为本步独立值），`CopyOutMaxAndSum` 每次调用都执行（`DataCopyPad` 写 GM，V_MTE3 EVENT_ID7 同步）；
	187	+- `QK_UB_RELEASE_FLAG` 为 uint16_t（与 Shared 路的 uint64_t 区分）。
	188	+
	189	+### 4.3 EpilogueXASharedRescaleO：O 累积 + PV 释放闭环
	190	+
	191	+```cpp
	192	+BlockEpilogue(Resource *resource, uint32_t &ubBufAddrStart); // 仅 2 参数
	193	+// operator()(attenOutGm, expMaxUb, pvRes, isFirstKv, isLastKv, PV_RELEASE_FLAG)
	194	+```
	195	+
	196	+- UB 资源：仅 `attnTmpBuf`（`VEC2_UB_SIZE = HALF_S1 × D × sizeof`，O 的累积缓冲）；
	197	+- 执行流程：
	198	+ 1. `isFirstKv`：`DataCopy(attnTmp ← pvRes)` 初始化；
	199	+ 2. 否则 `FlashUpdateNew`：`O_new = expMax × O_old + PV_cur`，expMax 逐行 Brcb 广播加载，nLoops 按 vlSize 切 D 轴；
	200	+ 3. `isLastKv`：`CopyUbToGmO` 写出最终 partial O；
	201	+ 4. `CrossCoreSetFlag(PV_RELEASE_FLAG)`：消费完 `pvRes` 后释放 PV GEMM 的 UB 双缓冲——这是 `MmadXASharedPV` 中 `PV_UB_RELEASE_FLAG` 的消费者，形成完整生产者-消费者闭环。
	202	+
	203	+### 4.4 EpilogueXACombineScale：双路 LSE 合并
	204	+
	205	+```cpp
	206	+BlockEpilogue(Resource *resource, uint32_t &ubBufAddrStart,
	207	+ int32_t rowNumPerLoop, int32_t headDim);
	208	+// operator()(sharedMax/Sum/Gm, unsharedMax/Sum/Gm, gFinalOutput, m, taskId&) 共 8 参数
	209	+```
	210	+
	211	+- 模板仅 2 参数（`OutputType/InputType`），无 TileShape 依赖，是纯 Vector 合并算子；
	212	+- UB 资源：7 组乒乓缓冲（shared/unshared 的 attn/gm/gl + finalAttn，各 `[2]`）+ 3 个 tmp（finalGl/expMaxShared/expMaxUnShared）；
	213	+- 合并数学（标准 LSE 合并）：
	214	+
	215	+```
	216	+finalMax = max(sharedMax, unsharedMax)
	217	+finalGl = sharedSum × exp(sharedMax − finalMax) + unsharedSum × exp(unsharedMax − finalMax)
	218	+O_final = (O_shared × exp(sharedMax − finalMax) + O_unshared × exp(unsharedMax − finalMax)) / finalGl
	219	+```
	220	+
	221	+- 执行流程：DataCopy 搬入双路 O/max/sum → `ComputeExpSumAndExpMax` → `ComputeFinalAttn`（nLoops 按 vlSize 切 headDim）→ Cast 到 ElementOutput → DataCopy 写 `gFinalOutput` → `taskId = 1 - taskId` 翻转乒乓。
	222	+
	223	+## 5. 跨核同步 flag 总表
	224	+
	225	+\| Flag \| 生产者 \| 消费者 \| 语义 \|
	226	+\| --- \| --- \| --- \| --- \|
	227	+\| `SYNC_QK_READY_FLAG[i]`（+16 AIV1 镜像） \| AIC：QK GEMM 完成 \| AIV：softmax 开始读 QK 结果 \| QK S 就绪 \|
	228	+\| `QK_UB_RELEASE_FLAG[i]` \| AIV：softmax 消费完 S / kernel 尾部预置 \| AIC：QK GEMM 复用 qkTensorList 槽 \| QK UB 槽空闲 \|
	229	+\| `SYNC_SOFTMAX_READY_FLAG[i]`（+16） \| AIV：P 已写 L1 \| AIC：PV GEMM 开始读 L1A \| P 就绪 \|
	230	+\| `SYNC_PV_READY_FLAG[i]`（+16） \| AIC：PV GEMM 完成 \| AIV：Rescale 开始读 pvRes \| PV O 就绪 \|
	231	+\| `PV_UB_RELEASE_FLAG[i]` \| AIV：Rescale 消费完 pvRes / kernel 尾部预置 \| AIC：PV GEMM 复用 pvTensorList 槽 \| PV UB 槽空闲 \|
	232	+
	233	+## 6. Kernel 组装与软件流水
	234	+
	235	+三个组装 kernel 位于 xllm-ops 的 `x_attention/op_kernel/arch35/` 目录。AIC 与 AIV 以"双核启动"方式绑定（`CV_RATIO = 2` 表示 1 个 Cube 核带 2 个 Vector 核），通过 `sharedInfo.usedCoreNum` 在同一批核上错峰启动 shared 与 unshared kernel。
	236	+
	237	+### 6.1 SharedFaInferKernel：4 级软件流水（shared_infer_catlass_kernel.h，444 行）
	238	+
	239	+```cpp
	240	+SharedFaInferKernel<BlockMmadQK, BlockMmadPV, EpilogueOnlineSoftmax, EpilogueRescaleO, KVLEN_T>;
	241	+// BlockMmadQK = BlockMmadTla<MmadXASharedQK<...>, Shape<_128,_128,_128>, Shape<_128,_128,_128>, ...>
	242	+// BlockMmadPV = BlockMmadTla<MmadXASharedPV<...>, ...>
	243	+// EpilogueOnlineSoftmax = BlockEpilogue<EpilogueXASharedSoftmax<...>, ...>
	244	+// EpilogueRescaleO = BlockEpilogue<EpilogueXASharedRescaleO<...>, ...>
	245	+```
	246	+
	247	+采用 `taskArgList[4]` 环形队列的 4 级软件流水，主循环条件为 `qTaskId < taskEndId + 3`（多跑 3 轮排空流水）：
	248	+
	249	+\| 流水级 \| 执行体 \| 任务 \| 关键同步 \|
	250	+\| --- \| --- \| --- \| --- \|
	251	+\| 级 0 \| AIC \| 当前 qTaskId 的 QK GEMM \| 等 `QK_UB_RELEASE_FLAG` \|
	252	+\| 级 1 \| AIV \| 滞后 1 拍（taskId-1）的 softmax \| 等 `SYNC_QK_READY_FLAG` \|
	253	+\| 级 2 \| AIC \| 滞后 2 拍（taskId-2）的 PV GEMM \| 先等 `SYNC_SOFTMAX_READY` 再 Mmad，完成后置 `SYNC_PV_READY` \|
	254	+\| 级 3 \| AIV \| 滞后 3 拍（taskId-3）的 Rescale \| 等 `SYNC_PV_READY_FLAG` \|
	255	+
	256	+```mermaid
	257	+graph LR
	258	+ T0[AIC: QK task i] --> T1[AIV: softmax task i-1]
	259	+ T1 --> T2[AIC: PV task i-2]
	260	+ T2 --> T3[AIV: rescale task i-3]
	261	+```
	262	+
	263	+- UB 多缓冲：`qkTensorList[2]` + `pvTensorList[2]` 乒乓；softmax 标量 `expSumUb[3] / expMaxUb[3] / maxUb[3]` 三缓冲（与滞后拍数匹配）；
	264	+- L1 多缓冲：`pL1TensorList[3]`（P 的 L1 三缓冲）；
	265	+- AIV 双核分工：`coreIdx = coreIdx / CV_RATIO` 得逻辑核号，`subVecIdx = coreIdx % CV_RATIO` 决定处理 `halfBlockQLen` 的前/后半；
	266	+- 任务切分：`GetQTaskInfo` 按 batch / qHead / qBlock 切 Q 任务；`GetKvTaskInfo` 计算 `isFirstKv / isUpdate / isLastKv / taskIdMod2 / taskIdMod3`；
	267	+- 预热 SetFlag：kernel 入口 AIC 置 M_MTE1 EVENT_ID 0-3，AIV 置 4 个 UB_RELEASE flag，消除首轮等待；
	268	+- 所有 `SYNC_*_READY` flag 均双份 Set（含 +16 偏移的 AIV1 镜像事件）。
	269	+
	270	+### 6.2 UnSharedInferKernel：3 级软件流水 + 页表寻址（unshared_infer_catlass_kernel.h，375 行）
	271	+
	272	+```cpp
	273	+UnSharedInferKernel<BlockMmadQK, BlockMmadPV, EpilogueSoftmax, KVLEN_T, TABLE_T>;
	274	+// BlockMmadQK = BlockMmadTla<MmadXAUnsharedQK<...>, ...>
	275	+// BlockMmadPV = BlockMmadTla<MmadXAUnsharedPV<...>, ...>
	276	+// EpilogueSoftmax = BlockEpilogue<EpilogueXAUnsharedSoftmax<...>, ...>
	277	+```
	278	+
	279	+`taskArgList[3]` 环形队列的 3 级软件流水，主循环条件 `groupTaskId < taskEndId + 2`：
	280	+
	281	+\| 流水级 \| 执行体 \| 任务 \| 说明 \|
	282	+\| --- \| --- \| --- \| --- \|
	283	+\| 级 0 \| AIC \| 当前 groupTaskId 的 QK GEMM \| — \|
	284	+\| 级 1 \| AIV \| 滞后 1 拍的 mask softmax \| 每步写 GM max/sum \|
	285	+\| 级 2 \| AIC \| 滞后 2 拍的 PV GEMM \| O 直写 GM（tensorOTile），无 rescale 级、无 PV flag \|
	286	+
	287	+- Paged KV Cache 寻址：`blockTableGm.GetValue(batchId)` 得 `cacheBlockId`，再以 `gK[cacheBlockId * kvBatchStride]` 定位物理页基地址——K/V 的 GM tensor 在任务循环内动态计算；
	288	+- 任务切分：`groupTaskId` 切 batch，`groupCountBlockId` 切组；`blockQLen = groupCountPerLoop × groupSize`，`blockKvLen = groupCountPerLoop × maxDecodeStep`；
	289	+- AIV 双 SubBlock 分工：`halfVecGroupCount` 切组数，`halfVecRowCount / halfVecRowOffset` 切行；
	290	+- 核偏移：`coreIdx` 减去 `sharedInfo.usedCoreNum`——与 shared kernel 在同一物理核域内错峰启动；
	291	+- UB 仅 `qkTensorList[2]`，L1 为 `pL1TensorList[3]`。
	292	+
	293	+### 6.3 CombineScaleKernel：纯 AIV 合并（combine_kernel.h，151 行）
	294	+
	295	+```cpp
	296	+CombineScaleKernel<EpilogueCombineScale>;
	297	+// EpilogueCombineScale = BlockEpilogue<EpilogueXACombineScale<...>, InputType, OutputType>
	298	+```
	299	+
	300	+- 纯 AIV kernel：AIC 特化为空 `return`，全部工作在 Vector 核完成；
	301	+- former/tail 两级任务切分：`formerCoreNum / formerTaskNum / tailCoreNum` 做核间负载均衡，每核循环 `coreTaskNum` 次，每次按 `rowNumPerLoop` 行调用 epilogue（`realRowNum` 按剩余行数钳制）；
	302	+- 7 个 GM tensor（shared/unshared 各 max/sum/O + gFinalOut）按 `gmglOffsetPerCore / attnOffsetPerCore` 计算各核偏移。
	303	+
	304	+## 7. 使用示例
	305	+
	306	+以下摘自 xllm-ops（https://gitcode.com/xLLM-AI/xllm_ops）`x_attention/op_kernel/x_attention_catlass_helper.h`（124 行），展示三个 kernel 的组装入口。
	307	+
	308	+### 7.1 Shared 路入口：CallSharedInferKernel
	309	+
	310	+```cpp
	311	+template <typename INPUT_T, typename TILING_T>
	312	+__aicore__ inline void CallSharedInferKernel(const TILING_T &tiling, SharedInfo &sharedInfo,
	313	+ const INPUT_T query, const INPUT_T key,
	314	+ const INPUT_T value, float attnOut,
	315	+ float attnMax, float attnSum, ...) {
	316	+ using ShapeType = Shape<_128, _128, _128>; // L1TileShape
	317	+ using L0ShapeType = Shape<_128, _128, _128>; // L0TileShape
	318	+ using LayoutK = ACL_CALL_FUNC_NAMESPACE::Layout::ColumnMajor;
	319	+ using ElementS = float;
	320	+ using ElementOTmp = float;
	321	+
	322	+ using BlockMmadQK = Catlass::Gemm::BlockMmadTla<Catlass::Gemm::MmadXASharedQK<Arch::Ascend950>,
	323	+ ShapeType, L0ShapeType, INPUT_T, INPUT_T, ElementS, Catlass::Gemm::Tag, ...>;
	324	+ using BlockMmadPV = Catlass::Gemm::BlockMmadTla<Catlass::Gemm::MmadXASharedPV<Arch::Ascend950>,
	325	+ ShapeType, L0ShapeType, ElementS, INPUT_T, ElementOTmp, Catlass::Gemm::Tag, ...>;
	326	+ using EpilogueOnlineSoftmax = Catlass::Epilogue::BlockEpilogue<
	327	+ Catlass::Epilogue::EpilogueXASharedSoftmax<Arch::Ascend950>, ...>;
	328	+ using EpilogueRescaleO = Catlass::Epilogue::BlockEpilogue<
	329	+ Catlass::Epilogue::EpilogueXASharedRescaleO<Arch::Ascend950>, ...>;
	330	+
	331	+ using Kernel = SharedFaInferKernel<BlockMmadQK, BlockMmadPV,
	332	+ EpilogueOnlineSoftmax, EpilogueRescaleO, KVLEN_T>;
	333	+ Kernel::Invoke(tiling, sharedInfo, query, key, value, attnOut, attnMax, attnSum, ...);
	334	+}
	335	+```
	336	+
	337	+### 7.2 Unshared 路入口：CallUnsharedInferKernel
	338	+
	339	+```cpp
	340	+using BlockMmadQK = Catlass::Gemm::BlockMmadTla<Catlass::Gemm::MmadXAUnsharedQK<Arch::Ascend950>,
	341	+ ShapeType, L0ShapeType, INPUT_T, INPUT_T, ElementS, Catlass::Gemm::Tag, ...>;
	342	+using BlockMmadPV = Catlass::Gemm::BlockMmadTla<Catlass::Gemm::MmadXAUnsharedPV<Arch::Ascend950>,
	343	+ ShapeType, L0ShapeType, ElementS, INPUT_T, ElementOTmp, Catlass::Gemm::Tag, ...>;
	344	+using EpilogueSoftmax = Catlass::Epilogue::BlockEpilogue<
	345	+ Catlass::Epilogue::EpilogueXAUnsharedSoftmax<Arch::Ascend950>, ...>;
	346	+
	347	+using Kernel = UnSharedInferKernel<BlockMmadQK, BlockMmadPV, EpilogueSoftmax, KVLEN_T, TABLE_T>;
	348	+Kernel::Invoke(...); // blockTableGm 页表 + unshared K/V
	349	+```
	350	+
	351	+### 7.3 Combine 入口：CallCombineScale
	352	+
	353	+```cpp
	354	+using EpilogueCombineScale = Catlass::Epilogue::BlockEpilogue<
	355	+ Catlass::Epilogue::EpilogueXACombineScale<Arch::Ascend950>, OutputType, InputType>;
	356	+using Kernel = CombineScaleKernel<EpilogueCombineScale>;
	357	+Kernel::Invoke(sharedMax, sharedSum, sharedAttn, unsharedMax, unsharedSum,
	358	+ unsharedAttn, finalOut, totalRowNum, rowNumPerLoop, headDim);
	359	+```
	360	+
	361	+三个入口在算子层（`x_attention` 的 `OpKernel` 实现）中按顺序调用：先 shared 与 unshared（同核域错峰），最后 combine 合并输出。调用方只需保证 tiling 中给出 `sharedKvLen / unsharedKvLen / groupSize / maxDecodeStep / qHeads` 等字段与 GM 张量布局一致。
	362	+
	363	+## 8. 约束与注意事项
	364	+
	365	+1. 架构约束：仅支持 Ascend 950（`CATLASS_ARCH == 3510`），注册头需置于 catlass 聚合头之后；
	366	+2. TileShape 约束：
	367	+ - SharedQK：L1/L0 无耦合约束（典型均为 128×128×128）；
	368	+ - UnsharedQK：M/K 轴 L1 与 L0 相同，`L1_TILE_N` 必须是 `L0_TILE_N` 整数倍；
	369	+ - SharedPV：L1/L0 三轴必须相同；
	370	+ - UnsharedPV：M/N 轴相同，`L1_TILE_K` 必须是 `L0_TILE_K` 整数倍；
	371	+3. L0C 复用依赖：UnsharedPV 复用 QK 的 L0C 区域，两个 GEMM 必须在同一 kernel 内按 QK→PV 顺序构造，且 `L1_TILE_M × L1_TILE_K` 需落在 QK L0C 预算内；
	372	+4. 元素类型：典型组合为 Q/K/V/P = `INPUT_T`（half/bf16），S/OTmp = `float`，K 布局为 `ColumnMajor`；
	373	+5. flag 类型：Shared 路 `QK_UB_RELEASE_FLAG` 为 uint64_t，Unshared 路为 uint16_t，跨核事件号规划时不可混用；
	374	+6. 流水深度即缓冲深度：Shared 路 softmax 标量三缓冲、P 的 L1 三缓冲与 4 级流水（滞后最多 3 拍）严格对应，修改流水结构时需同步调整缓冲数量。

Minclude/catlass/epilogue/block/block_epilogue.hpp+12-1

	@@ -22,6 +22,10 @@ class BlockEpilogue {
22		22
23	} // namespace Catlass::Epilogue::Block	23	} // namespace Catlass::Epilogue::Block
24		24
		25	+#include "catlass/epilogue/block/block_epilogue_fa_unshared_softmax.hpp"
		26	+#include "catlass/epilogue/block/block_epilogue_rescale_o_no_div_rowsum.hpp"
		27	+#include "catlass/epilogue/block/block_epilogue_combine_scale.hpp"
		28	+#include "catlass/epilogue/block/block_epilogue_online_softmax_copy_glm.hpp"
25	#include "catlass/epilogue/block/block_epilogue_elemwise_no_source.hpp"	29	#include "catlass/epilogue/block/block_epilogue_elemwise_no_source.hpp"
26	#include "catlass/epilogue/block/block_epilogue_elemwise_one_source.hpp"	30	#include "catlass/epilogue/block/block_epilogue_elemwise_one_source.hpp"
27	#include "catlass/epilogue/block/block_epilogue_fa_softmax.hpp"	31	#include "catlass/epilogue/block/block_epilogue_fa_softmax.hpp"
	@@ -40,7 +44,10 @@ class BlockEpilogue {
40	#include "catlass/epilogue/block/block_epilogue_online_softmax_no_mask.hpp"	44	#include "catlass/epilogue/block/block_epilogue_online_softmax_no_mask.hpp"
41	#include "catlass/epilogue/block/block_epilogue_rescale_o_no_split_row.hpp"	45	#include "catlass/epilogue/block/block_epilogue_rescale_o_no_split_row.hpp"
42	#include "catlass/epilogue/block/block_epilogue_w4a4_per_token_per_channel_dequant.hpp"	46	#include "catlass/epilogue/block/block_epilogue_w4a4_per_token_per_channel_dequant.hpp"
43	-	47	+#include "catlass/epilogue/block/block_epilogue_online_softmax_FD.hpp"
		48	+#include "catlass/epilogue/block/block_epilogue_rescale_o_FD.hpp"
		49	+#include "catlass/epilogue/block/block_epilogue_xfai_online_softmax.hpp"
		50	+#include "catlass/epilogue/block/block_epilogue_xfai_rescale_o.hpp"
44	#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510)	51	#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510)
45	#include "catlass/epilogue/block/block_epilogue_fa_softmax_ascend950.hpp"	52	#include "catlass/epilogue/block/block_epilogue_fa_softmax_ascend950.hpp"
46	#include "catlass/epilogue/block/block_epilogue_fa_rescale_o_ascend950.hpp"	53	#include "catlass/epilogue/block/block_epilogue_fa_rescale_o_ascend950.hpp"
	@@ -59,6 +66,10 @@ class BlockEpilogue {
59	#include "catlass/epilogue/block/block_epilogue_flash_attention_online_softmax_low_prec.hpp"	66	#include "catlass/epilogue/block/block_epilogue_flash_attention_online_softmax_low_prec.hpp"
60	#include "catlass/epilogue/block/block_epilogue_flash_attention_rescale_o.hpp"	67	#include "catlass/epilogue/block/block_epilogue_flash_attention_rescale_o.hpp"
61	#include "catlass/epilogue/block/block_epilogue_elemwise_no_source_from_ub.hpp"	68	#include "catlass/epilogue/block/block_epilogue_elemwise_no_source_from_ub.hpp"
		69	+#include "catlass/epilogue/block/block_epilogue_xa_shared_softmax_ascend950.hpp"
		70	+#include "catlass/epilogue/block/block_epilogue_xa_shared_rescale_ascend950.hpp"
		71	+#include "catlass/epilogue/block/block_epilogue_xa_unshared_softmax_ascend950.hpp"
		72	+#include "catlass/epilogue/block/block_epilogue_xa_combine_scale_ascend950.hpp"
62	#endif	73	#endif
63		74
64	#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_HPP	75	#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_HPP

Ainclude/catlass/epilogue/block/block_epilogue_combine_scale.hpp+403-0

@@ -0,0 +1,403 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_COMBINE_SCALE_HPP
	13	+#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_COMBINE_SCALE_HPP
	14	+
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/arch/resource.hpp"
	17	+#include "catlass/epilogue/dispatch_policy.hpp"
	18	+#include "catlass/epilogue/tile/tile_copy.hpp"
	19	+#include "catlass/gemm_coord.hpp"
	20	+#include "catlass/matrix_coord.hpp"
	21	+
	22	+namespace Catlass::Epilogue::Block {
	23	+
	24	+template <class OutputType_, class InputType_>
	25	+class BlockEpilogue<EpilogueAtlasA2CombineScale, OutputType_, InputType_> {
	26	+public:
	27	+ // Type aliases
	28	+ using DispatchPolicy = EpilogueAtlasA2CombineScale;
	29	+ using ArchTag = typename DispatchPolicy::ArchTag;
	30	+
	31	+ using ElementOutput = typename OutputType_::Element;
	32	+ using ElementInput = typename InputType_::Element;
	33	+
	34	+ using LayoutOutput = typename OutputType_::Layout;
	35	+ using LayoutInput = typename InputType_::Layout;
	36	+
	37	+ // 常量定义
	38	+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
	39	+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
	40	+ static constexpr uint32_t VECTOR_SIZE = 128;
	41	+ static constexpr uint32_t MAX_UB_ELEM_NUM = 32 * 1024 / sizeof(float); // output使用32;
	42	+ static constexpr uint32_t SUM_MAX_UB_PINGPONG_OFFSET = 2 * 1024;
	43	+ static constexpr uint32_t ATTENTION_TEMP_PINGPONG_OFFSET = 32 * 1024;
	44	+ static constexpr uint32_t ATTENTION_TEMP_OFFSET = 2 * ATTENTION_TEMP_PINGPONG_OFFSET; // 64k
	45	+ static constexpr uint32_t ATTEN_OUT_OFFSET = 128 * 1024;
	46	+ static constexpr uint32_t SOFTMAX_BROAD_SIZE = 8;
	47	+ CATLASS_DEVICE
	48	+ BlockEpilogue(Arch::Resource<ArchTag>& resource)
	49	+ {
	50	+ // Allocate UB space
	51	+ /*
	52	+ shared_gl: 4kb [2kb ping, 2kb pong]
	53	+ unshared_gl: 4kb
	54	+ shared_gm: 4kb
	55	+ unshared_gm: 4kb
	56	+ shared_out: 64kb [32kb ping, 32kb pong]
	57	+ unshared_out: 64kb
	58	+ middle Tenosr : 192 - 16 - 128 = 48kb
	59	+ */
	60	+ sharedOutUbTensor = resource.ubBuf.template GetBufferByByte<float>(0); // 0 - 64k
	61	+ unsharedOutUbTensor = resource.ubBuf.template GetBufferByByte<float>(ATTENTION_TEMP_OFFSET); // 64k - 128k
	62	+ sharedGlUbTensor = resource.ubBuf.template GetBufferByByte<float>(ATTEN_OUT_OFFSET); // 128k - 132k
	63	+ unsharedGlUbTensor = resource.ubBuf.template GetBufferByByte<float>(
	64	+ ATTEN_OUT_OFFSET + 2 * SUM_MAX_UB_PINGPONG_OFFSET); // 132k - 136k
	65	+ sharedGmUbTensor = resource.ubBuf.template GetBufferByByte<float>(
	66	+ ATTEN_OUT_OFFSET + 4 * SUM_MAX_UB_PINGPONG_OFFSET); // 136k - 140k
	67	+ unsharedGmUbTensor = resource.ubBuf.template GetBufferByByte<float>(
	68	+ ATTEN_OUT_OFFSET + 6 * SUM_MAX_UB_PINGPONG_OFFSET); // 140k - 144k
	69	+ realGmUbTensor = resource.ubBuf.template GetBufferByByte<float>(
	70	+ ATTEN_OUT_OFFSET + 8 * SUM_MAX_UB_PINGPONG_OFFSET); // 144k - 148k
	71	+ realGlUbTensor = resource.ubBuf.template GetBufferByByte<float>(
	72	+ ATTEN_OUT_OFFSET + 10 * SUM_MAX_UB_PINGPONG_OFFSET); // 148k - 152k
	73	+ outUbTensor = resource.ubBuf.template GetBufferByByte<ElementOutput>(
	74	+ ATTEN_OUT_OFFSET + 12 * SUM_MAX_UB_PINGPONG_OFFSET); // 152k - 184k
	75	+ }
	76	+
	77	+ CATLASS_DEVICE
	78	+ ~BlockEpilogue()
	79	+ {}
	80	+
	81	+ CATLASS_DEVICE
	82	+ void SetMask(int32_t len)
	83	+ {
	84	+ uint64_t mask = 0;
	85	+ uint64_t one = 1;
	86	+ uint64_t temp = len % FLOAT_VECTOR_SIZE;
	87	+ for (int64_t i = 0; i < temp; i++) {
	88	+ mask \|= one << i;
	89	+ }
	90	+
	91	+ if (len == VECTOR_SIZE) {
	92	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	93	+ } else if (len >= FLOAT_VECTOR_SIZE) {
	94	+ AscendC::SetVectorMask<int8_t>(mask, (uint64_t)-1);
	95	+ } else {
	96	+ AscendC::SetVectorMask<int8_t>(0x0, mask);
	97	+ }
	98	+ }
	99	+
	100	+ CATLASS_DEVICE
	101	+ void operator()(
	102	+ AscendC::GlobalTensor<float> gSharedGm, AscendC::GlobalTensor<float> gUnsharedGm,
	103	+ AscendC::GlobalTensor<float> gSharedGl, AscendC::GlobalTensor<float> gUnsharedGl,
	104	+ AscendC::GlobalTensor<float> gSharedOut, AscendC::GlobalTensor<float> gUnsharedOut,
	105	+ AscendC::GlobalTensor<ElementOutput> gFinalOutput, const MatrixCoord& actualBlockShape)
	106	+ {
	107	+ // 伪代码实现：
	108	+ // 1. 计算分块参数
	109	+ uint32_t rowNum = actualBlockShape.row();
	110	+ uint32_t columnNum = actualBlockShape.column();
	111	+ uint32_t columnNumRound = RoundUp(columnNum, FLOAT_BLOCK_SIZE);
	112	+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
	113	+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
	114	+ uint32_t rowSplitSubBlock = rowNum / subBlockNum;
	115	+ uint32_t rowActualThisSubBlock = (subBlockIdx == 0) ? rowSplitSubBlock : (rowNum - rowSplitSubBlock);
	116	+ uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock;
	117	+ uint32_t rowNumTile = RoundDown(MAX_UB_ELEM_NUM / columnNum, FLOAT_BLOCK_SIZE);
	118	+ uint32_t rowLoopNum = CeilDiv(rowActualThisSubBlock, rowNumTile);
	119	+ uint32_t preLoad = 1;
	120	+
	121	+ if (rowActualThisSubBlock == 0) {
	122	+ return;
	123	+ }
	124	+
	125	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(0);
	126	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(1);
	127	+
	128	+ // 4. 主循环 - preload 和 pingpong 优化
	129	+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) {
	130	+ if (rowLoopIdx < rowLoopNum) {
	131	+ // 数据加载阶段
	132	+ uint32_t pingpongFlag = rowLoopIdx % 2;
	133	+ uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile;
	134	+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
	135	+ uint32_t rowNumCurLoop =
	136	+ (rowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
	137	+ uint32_t rowNumCurLoopRound = RoundUp(rowNumCurLoop, FLOAT_BLOCK_SIZE);
	138	+ uint32_t pingPongOffset = pingpongFlag * SUM_MAX_UB_PINGPONG_OFFSET / sizeof(float);
	139	+ uint32_t attnPingPongOffset = pingpongFlag * ATTENTION_TEMP_PINGPONG_OFFSET / sizeof(float);
	140	+ uint32_t sumMaxOffsetIoGm = rowOffsetIoGm;
	141	+ uint32_t sumMaxOffsetIoShared = rowOffsetIoGm * SOFTMAX_BROAD_SIZE;
	142	+ uint32_t attnOffsetIoGm = rowOffsetIoGm * columnNum;
	143	+
	144	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(pingpongFlag);
	145	+ AscendC::DataCopy(
	146	+ sharedGmUbTensor[pingPongOffset], gSharedGm[sumMaxOffsetIoShared],
	147	+ rowNumCurLoop * SOFTMAX_BROAD_SIZE);
	148	+ // Copy GL
	149	+ AscendC::DataCopy(
	150	+ sharedGlUbTensor[pingPongOffset], gSharedGl[sumMaxOffsetIoShared],
	151	+ rowNumCurLoop * SOFTMAX_BROAD_SIZE);
	152	+
	153	+ // Copy Shared and Unshared
	154	+ // [rowNum, column]
	155	+ if (rowNumCurLoop % FLOAT_BLOCK_SIZE == 0) {
	156	+ AscendC::DataCopy(unsharedGmUbTensor[pingPongOffset], gUnsharedGm[sumMaxOffsetIoGm], rowNumCurLoop);
	157	+ AscendC::DataCopy(unsharedGlUbTensor[pingPongOffset], gUnsharedGl[sumMaxOffsetIoGm], rowNumCurLoop);
	158	+ } else {
	159	+ AscendC::DataCopyExtParams copyInParams{
	160	+ 1, static_cast<uint32_t>(rowNumCurLoop * sizeof(float)), 0, 0, 0};
	161	+ AscendC::DataCopyPadExtParams<float> padParams{false, 0, 0, 0};
	162	+ AscendC::DataCopyPad(
	163	+ unsharedGmUbTensor[pingPongOffset], gUnsharedGm[sumMaxOffsetIoGm], copyInParams, padParams);
	164	+ AscendC::DataCopyPad(
	165	+ unsharedGlUbTensor[pingPongOffset], gUnsharedGl[sumMaxOffsetIoGm], copyInParams, padParams);
	166	+ }
	167	+ AscendC::DataCopy(
	168	+ sharedOutUbTensor[attnPingPongOffset], gSharedOut[attnOffsetIoGm], rowNumCurLoop * columnNum);
	169	+ AscendC::DataCopy(
	170	+ unsharedOutUbTensor[attnPingPongOffset], gUnsharedOut[attnOffsetIoGm], rowNumCurLoop * columnNum);
	171	+ }
	172	+
	173	+ if (rowLoopIdx >= preLoad) {
	174	+ // 计算阶段
	175	+ uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad;
	176	+ uint32_t pingpongFlag = delayedRowLoopIdx % 2;
	177	+ uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile;
	178	+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
	179	+ uint32_t rowNumCurLoop =
	180	+ (delayedRowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
	181	+ uint32_t rowNumCurLoopRound = RoundUp(rowNumCurLoop, FLOAT_BLOCK_SIZE);
	182	+ uint32_t pingPongOffset = pingpongFlag * SUM_MAX_UB_PINGPONG_OFFSET / sizeof(float);
	183	+ uint32_t attnPingPongOffset = pingpongFlag * ATTENTION_TEMP_PINGPONG_OFFSET / sizeof(float);
	184	+ // 调用 SubCoreCompute 进行核心计算
	185	+ uint32_t attnOffsetIoGm = rowOffsetIoGm * columnNum;
	186	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID4);
	187	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID4);
	188	+ SubCoreCompute(
	189	+ sharedGmUbTensor[pingPongOffset], unsharedGmUbTensor[pingPongOffset],
	190	+ realGmUbTensor[pingPongOffset], sharedGlUbTensor[pingPongOffset],
	191	+ unsharedGlUbTensor[pingPongOffset], realGlUbTensor[pingPongOffset],
	192	+ sharedOutUbTensor[attnPingPongOffset], unsharedOutUbTensor[attnPingPongOffset], pingpongFlag,
	193	+ rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound);
	194	+ // 6. Cast 到目标精度 (FP16/BF16)
	195	+ if constexpr (std::is_same_v<ElementOutput, bfloat16_t>) {
	196	+ AscendC::Cast<ElementOutput, float, false>(
	197	+ outUbTensor.template ReinterpretCast<bfloat16_t>()[attnPingPongOffset],
	198	+ sharedOutUbTensor[attnPingPongOffset], AscendC::RoundMode::CAST_RINT, (uint64_t)0,
	199	+ (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	200	+ AscendC::UnaryRepeatParams(1, 1, 4, 8));
	201	+ } else {
	202	+ AscendC::Cast<ElementOutput, float, false>(
	203	+ outUbTensor.template ReinterpretCast<half>()[attnPingPongOffset],
	204	+ sharedOutUbTensor[attnPingPongOffset], AscendC::RoundMode::CAST_NONE, (uint64_t)0,
	205	+ (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	206	+ AscendC::UnaryRepeatParams(1, 1, 4, 8));
	207	+ }
	208	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID4);
	209	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID4);
	210	+ if constexpr (std::is_same_v<ElementOutput, bfloat16_t>) {
	211	+ AscendC::DataCopy(
	212	+ gFinalOutput[attnOffsetIoGm],
	213	+ outUbTensor.template ReinterpretCast<bfloat16_t>()[attnPingPongOffset],
	214	+ rowNumCurLoop * columnNum);
	215	+ } else {
	216	+ AscendC::GlobalTensor<ElementOutput> gFinal = gFinalOutput[attnOffsetIoGm];
	217	+ AscendC::DataCopy(
	218	+ gFinal, outUbTensor.template ReinterpretCast<half>()[attnPingPongOffset],
	219	+ rowNumCurLoop * columnNum);
	220	+ }
	221	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(pingpongFlag);
	222	+ }
	223	+ }
	224	+
	225	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(0);
	226	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(1);
	227	+ }
	228	+
	229	+private:
	230	+ // SubCoreCompute 函数：执行核心计算逻辑
	231	+ CATLASS_DEVICE
	232	+ void SubCoreCompute(
	233	+ AscendC::LocalTensor<float> sharedGmUbLoopTensor, AscendC::LocalTensor<float> unsharedGmUbLoopTensor,
	234	+ AscendC::LocalTensor<float> realGmUbLoopTensor, AscendC::LocalTensor<float> sharedGlUbLoopTensor,
	235	+ AscendC::LocalTensor<float> unsharedGlUbLoopTensor, AscendC::LocalTensor<float> realGlUbLoopTensor,
	236	+ AscendC::LocalTensor<float> sharedOutUbLoopTensor, AscendC::LocalTensor<float> unsharedOutUbLoopTensor,
	237	+ uint32_t pingpongFlag, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum,
	238	+ uint32_t columnNumRound)
	239	+ {
	240	+ // 1. gm = max(shared_gm, unshared_gm)
	241	+ // sharedGmUbLoopTensor [row, 8]
	242	+ // unsharedGmUbLoopTensor [row, 1]
	243	+ AscendC::BlockReduceMax<float, false>(
	244	+ sharedGmUbLoopTensor, sharedGmUbLoopTensor, rowNumCurLoop, uint64_t(0), 1, 1, 8);
	245	+ AscendC::PipeBarrier<PIPE_V>();
	246	+
	247	+ AscendC::BlockReduceMax<float, false>(
	248	+ sharedGlUbLoopTensor, sharedGlUbLoopTensor, rowNumCurLoop, uint64_t(0), 1, 1, 8);
	249	+ AscendC::PipeBarrier<PIPE_V>();
	250	+
	251	+ AscendC::Max<float>(realGmUbLoopTensor, sharedGmUbLoopTensor, unsharedGmUbLoopTensor, rowNumCurLoop);
	252	+ AscendC::PipeBarrier<PIPE_V>();
	253	+
	254	+ // 2. update_shared_expgm = exp(shared_gm - gm)
	255	+ // update_unshared_expgm = exp(unshared_gm - gm)
	256	+ AscendC::Sub<float>(sharedGmUbLoopTensor, sharedGmUbLoopTensor, realGmUbLoopTensor, rowNumCurLoop);
	257	+ AscendC::Sub<float>(unsharedGmUbLoopTensor, unsharedGmUbLoopTensor, realGmUbLoopTensor, rowNumCurLoop);
	258	+ AscendC::PipeBarrier<PIPE_V>();
	259	+
	260	+ AscendC::Exp<float>(sharedGmUbLoopTensor, sharedGmUbLoopTensor, rowNumCurLoop);
	261	+ AscendC::Exp<float>(unsharedGmUbLoopTensor, unsharedGmUbLoopTensor, rowNumCurLoop);
	262	+ AscendC::PipeBarrier<PIPE_V>();
	263	+
	264	+ // 3. update_unshared_expgm = exp(unshared_gm - gm)
	265	+
	266	+ // 4. gl = shared_gl * update_shared_expgm + unshared_gl * update_unshared_expgm
	267	+ // 先计算 shared_gl * update_shared_expgm
	268	+ AscendC::Mul<float>(sharedGlUbLoopTensor, sharedGlUbLoopTensor, sharedGmUbLoopTensor, rowNumCurLoop);
	269	+
	270	+ // 计算 unshared_gl * update_unshared_expgm 并累加到 gl
	271	+ AscendC::Mul<float>(unsharedGlUbLoopTensor, unsharedGlUbLoopTensor, unsharedGmUbLoopTensor, rowNumCurLoop);
	272	+ AscendC::PipeBarrier<PIPE_V>();
	273	+
	274	+ AscendC::Add<float>(realGlUbLoopTensor, sharedGlUbLoopTensor, unsharedGlUbLoopTensor, rowNumCurLoop);
	275	+ AscendC::PipeBarrier<PIPE_V>();
	276	+
	277	+ // BroadCast update_shared_expgm and update_unshared_expgm
	278	+ AscendC::Brcb(
	279	+ sharedGmUbLoopTensor, sharedGmUbLoopTensor, rowNumCurLoopRound / FLOAT_BLOCK_SIZE,
	280	+ AscendC::BrcbRepeatParams(1, 8));
	281	+ AscendC::PipeBarrier<PIPE_V>();
	282	+ AscendC::Brcb(
	283	+ unsharedGmUbLoopTensor, unsharedGmUbLoopTensor, rowNumCurLoopRound / FLOAT_BLOCK_SIZE,
	284	+ AscendC::BrcbRepeatParams(1, 8));
	285	+ AscendC::PipeBarrier<PIPE_V>();
	286	+
	287	+ // 5. final_true_out = (shared_true_out * update_shared_expgm + unshared_true_out * update_unshared_expgm) / gl
	288	+ // 计算 shared_true_out * update_shared_expgm
	289	+ for (uint32_t mulIdx = 0; mulIdx < columnNum / FLOAT_VECTOR_SIZE; ++mulIdx) {
	290	+ AscendC::Mul<float, false>(
	291	+ sharedOutUbLoopTensor[mulIdx * FLOAT_VECTOR_SIZE], sharedOutUbLoopTensor[mulIdx * FLOAT_VECTOR_SIZE],
	292	+ sharedGmUbLoopTensor, AscendC::MASK_PLACEHOLDER, rowNumCurLoop,
	293	+ AscendC::BinaryRepeatParams(
	294	+ 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1));
	295	+ }
	296	+ AscendC::PipeBarrier<PIPE_V>();
	297	+ if (columnNum % FLOAT_VECTOR_SIZE > 0) {
	298	+ SetMask(columnNum % FLOAT_VECTOR_SIZE);
	299	+ AscendC::Mul<float, false>(
	300	+ sharedOutUbLoopTensor[columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
	301	+ sharedOutUbLoopTensor[columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], sharedGmUbLoopTensor,
	302	+ AscendC::MASK_PLACEHOLDER, rowNumCurLoop,
	303	+ AscendC::BinaryRepeatParams(
	304	+ 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1));
	305	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	306	+ }
	307	+ AscendC::PipeBarrier<PIPE_V>();
	308	+
	309	+ for (uint32_t mulIdx = 0; mulIdx < columnNum / FLOAT_VECTOR_SIZE; ++mulIdx) {
	310	+ AscendC::Mul<float, false>(
	311	+ unsharedOutUbLoopTensor[mulIdx * FLOAT_VECTOR_SIZE],
	312	+ unsharedOutUbLoopTensor[mulIdx * FLOAT_VECTOR_SIZE], unsharedGmUbLoopTensor, (uint64_t)0, rowNumCurLoop,
	313	+ AscendC::BinaryRepeatParams(
	314	+ 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1));
	315	+ }
	316	+ AscendC::PipeBarrier<PIPE_V>();
	317	+ if (columnNum % FLOAT_VECTOR_SIZE > 0) {
	318	+ SetMask(columnNum % FLOAT_VECTOR_SIZE);
	319	+ AscendC::Mul<float, false>(
	320	+ unsharedOutUbLoopTensor[columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
	321	+ unsharedOutUbLoopTensor[columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], unsharedGmUbLoopTensor,
	322	+ (uint64_t)0, rowNumCurLoop,
	323	+ AscendC::BinaryRepeatParams(
	324	+ 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1));
	325	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	326	+ }
	327	+ AscendC::PipeBarrier<PIPE_V>();
	328	+
	329	+ // 累加到 sharedOutUbTensor
	330	+ for (uint32_t addIdx = 0; addIdx < columnNum / FLOAT_VECTOR_SIZE; ++addIdx) {
	331	+ AscendC::Add<float, false>(
	332	+ sharedOutUbLoopTensor[addIdx * FLOAT_VECTOR_SIZE], sharedOutUbLoopTensor[addIdx * FLOAT_VECTOR_SIZE],
	333	+ unsharedOutUbLoopTensor[addIdx * FLOAT_VECTOR_SIZE], (uint64_t)0, rowNumCurLoop,
	334	+ AscendC::BinaryRepeatParams(
	335	+ 1, 1, 1, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE,
	336	+ columnNumRound / FLOAT_BLOCK_SIZE));
	337	+ }
	338	+ AscendC::PipeBarrier<PIPE_V>();
	339	+ if (columnNum % FLOAT_VECTOR_SIZE > 0) {
	340	+ SetMask(columnNum % FLOAT_VECTOR_SIZE);
	341	+ AscendC::Add<float, false>(
	342	+ sharedOutUbLoopTensor[columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
	343	+ sharedOutUbLoopTensor[columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
	344	+ unsharedOutUbLoopTensor[columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], (uint64_t)0, rowNumCurLoop,
	345	+ AscendC::BinaryRepeatParams(
	346	+ 1, 1, 1, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE,
	347	+ columnNumRound / FLOAT_BLOCK_SIZE));
	348	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	349	+ }
	350	+ AscendC::PipeBarrier<PIPE_V>();
	351	+
	352	+ // 除以 gl
	353	+ AscendC::Brcb(
	354	+ realGlUbLoopTensor, realGlUbLoopTensor, rowNumCurLoopRound / FLOAT_BLOCK_SIZE,
	355	+ AscendC::BrcbRepeatParams(1, 8));
	356	+ AscendC::PipeBarrier<PIPE_V>();
	357	+ for (uint32_t divIdx = 0; divIdx < columnNum / FLOAT_VECTOR_SIZE; ++divIdx) {
	358	+ AscendC::Div<float, false>(
	359	+ sharedOutUbLoopTensor[divIdx * FLOAT_VECTOR_SIZE], sharedOutUbLoopTensor[divIdx * FLOAT_VECTOR_SIZE],
	360	+ realGlUbLoopTensor, (uint64_t)0, rowNumCurLoop,
	361	+ AscendC::BinaryRepeatParams(
	362	+ 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1));
	363	+ }
	364	+ AscendC::PipeBarrier<PIPE_V>();
	365	+ if (columnNum % FLOAT_VECTOR_SIZE > 0) {
	366	+ SetMask(columnNum % FLOAT_VECTOR_SIZE);
	367	+ AscendC::Div<float, false>(
	368	+ sharedOutUbLoopTensor[columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
	369	+ sharedOutUbLoopTensor[columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], realGlUbLoopTensor,
	370	+ (uint64_t)0, rowNumCurLoop,
	371	+ AscendC::BinaryRepeatParams(
	372	+ 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1));
	373	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	374	+ }
	375	+ AscendC::PipeBarrier<PIPE_V>();
	376	+ }
	377	+
	378	+ // UB 张量定义
	379	+ AscendC::LocalTensor<float> sharedGmUbTensor;
	380	+ AscendC::LocalTensor<float> unsharedGmUbTensor;
	381	+ AscendC::LocalTensor<float> realGmUbTensor;
	382	+ AscendC::LocalTensor<float> realGlUbTensor;
	383	+ AscendC::LocalTensor<float> sharedGlUbTensor;
	384	+ AscendC::LocalTensor<float> unsharedGlUbTensor;
	385	+ AscendC::LocalTensor<float> sharedOutUbTensor;
	386	+ AscendC::LocalTensor<float> unsharedOutUbTensor;
	387	+ AscendC::LocalTensor<ElementOutput> outUbTensor;
	388	+};
	389	+
	390	+// Ascend950 specialization - same implementation as AtlasA2 since Ascend C APIs are compatible
	391	+template <class OutputType_, class InputType_>
	392	+class BlockEpilogue<EpilogueAscend950CombineScale, OutputType_, InputType_>
	393	+ : public BlockEpilogue<EpilogueAtlasA2CombineScale, OutputType_, InputType_> {
	394	+public:
	395	+ using Base = BlockEpilogue<EpilogueAtlasA2CombineScale, OutputType_, InputType_>;
	396	+ using DispatchPolicy = EpilogueAscend950CombineScale;
	397	+ using ArchTag = typename DispatchPolicy::ArchTag;
	398	+ using Base::Base;
	399	+};
	400	+
	401	+} // namespace Catlass::Epilogue::Block
	402	+
	403	+#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_COMBINE_SCALE_HPP

Ainclude/catlass/epilogue/block/block_epilogue_fa_unshared_softmax.hpp+405-0

@@ -0,0 +1,405 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_UNSHARED_SOFTMAX_HPP
	13	+#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_UNSHARED_SOFTMAX_HPP
	14	+
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/arch/cross_core_sync.hpp"
	17	+#include "catlass/arch/resource.hpp"
	18	+#include "catlass/epilogue/dispatch_policy.hpp"
	19	+#include "catlass/epilogue/tile/tile_copy.hpp"
	20	+#include "catlass/gemm_coord.hpp"
	21	+#include "catlass/matrix_coord.hpp"
	22	+#include <limits>
	23	+#include <climits>
	24	+
	25	+namespace Catlass::Epilogue::Block {
	26	+
	27	+template <class OutputType_, class InputType_, class MaskType_>
	28	+class BlockEpilogue<EpilogueAtlasA2FAUnsharedSoftmax, OutputType_, InputType_, MaskType_> {
	29	+public:
	30	+ // Type aliases
	31	+ using DispatchPolicy = EpilogueAtlasA2FAUnsharedSoftmax;
	32	+ using ArchTag = typename DispatchPolicy::ArchTag;
	33	+ using ElementOutput = typename OutputType_::Element;
	34	+ using ElementInput = typename InputType_::Element;
	35	+ using ElementMask = typename MaskType_::Element;
	36	+
	37	+ using LayoutOutput = typename OutputType_::Layout;
	38	+ using LayoutInput = typename InputType_::Layout;
	39	+ using LayoutMask = typename MaskType_::Layout;
	40	+
	41	+ using CopyGmToUbInput = Tile::CopyGm2Ub<ArchTag, InputType_>;
	42	+ using CopyGmToUbMask = Tile::CopyGm2Ub<ArchTag, MaskType_>;
	43	+ using CopyUbToGmOutput = Tile::CopyUb2Gm<ArchTag, OutputType_>;
	44	+
	45	+ static constexpr uint32_t HALF_ELENUM_PER_BLK = 16;
	46	+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
	47	+ static constexpr uint32_t BRCB_REPEAT_SIZE = 8;
	48	+ static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128;
	49	+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
	50	+ static constexpr uint32_t UB_TILE_SIZE = 16384; // 64 * 128 * 2B
	51	+ static constexpr uint32_t UB_LINE_SIZE = 512; // 128 * 2 * 2B
	52	+ static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; // 128 * 2
	53	+ static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; // 128
	54	+ static constexpr uint32_t MULTIPLIER = 2;
	55	+ static constexpr uint32_t HALF_VECTOR_SIZE = 128;
	56	+ static constexpr uint32_t BLOCK_SIZE = 16;
	57	+ static constexpr uint32_t T_BLOCK_SIZE = 32 / 2;
	58	+ static constexpr uint32_t UB_UINT8_LINE_SIZE = 512;
	59	+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE_UNSHARED = 32768; // 64 * 256 * 2B
	60	+ static constexpr uint32_t HALF_DM_UB_SIZE = 128;
	61	+ static constexpr uint32_t VECTOR_SIZE = 128;
	62	+ static constexpr uint32_t HALF_LL_UB_SIZE = 256;
	63	+ static constexpr uint32_t SEQ_TILE_SIZE = 8;
	64	+
	65	+ CATLASS_DEVICE
	66	+ BlockEpilogue(
	67	+ Arch::Resource<ArchTag>& resource, float tor_, uint32_t unsharedKvSeqLen, uint32_t maxDecodeStep,
	68	+ uint32_t headNum, uint32_t groupSize)
	69	+ {
	70	+ // Allocate UB space
	71	+ constexpr uint32_t LS_UB_TENSOR_OFFSET = 0;
	72	+ constexpr uint32_t LP_UB_TENSOR_OFFSET = 2 * UB_UINT8_BLOCK_SIZE_UNSHARED;
	73	+ constexpr uint32_t LM_UB_TENSOR_OFFSET = 3 * UB_UINT8_BLOCK_SIZE_UNSHARED;
	74	+ constexpr uint32_t LL_UB_TENSOR_OFFSET = 3 * UB_UINT8_BLOCK_SIZE_UNSHARED + 4 * UB_UINT8_LINE_SIZE;
	75	+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 3 * UB_UINT8_BLOCK_SIZE_UNSHARED + 8 * UB_UINT8_LINE_SIZE;
	76	+ ;
	77	+ constexpr uint32_t MASK_UB_TENSOR_OFFSET = 3 * UB_UINT8_BLOCK_SIZE_UNSHARED + 12 * UB_UINT8_LINE_SIZE;
	78	+
	79	+ tor = tor_;
	80	+ tvUbTensor16 = resource.ubBuf.template GetBufferByByte<ElementOutput>(LP_UB_TENSOR_OFFSET);
	81	+ lpUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(LP_UB_TENSOR_OFFSET);
	82	+ lsUbTensor = resource.ubBuf.template GetBufferByByte<float>(LS_UB_TENSOR_OFFSET);
	83	+ lmUbTensor = resource.ubBuf.template GetBufferByByte<float>(LM_UB_TENSOR_OFFSET);
	84	+ llUbTensor = resource.ubBuf.template GetBufferByByte<float>(LL_UB_TENSOR_OFFSET);
	85	+ tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
	86	+
	87	+ // init unsharedMask: check could generate mask on host
	88	+ unsharedMaskUbTensor = resource.ubBuf.template GetBufferByByte<float>(MASK_UB_TENSOR_OFFSET);
	89	+ InitUnsharedMaskV2(unsharedKvSeqLen, maxDecodeStep, headNum, groupSize);
	90	+ AscendC::PipeBarrier<PIPE_V>();
	91	+ }
	92	+
	93	+ CATLASS_DEVICE
	94	+ ~BlockEpilogue()
	95	+ {}
	96	+
	97	+ CATLASS_DEVICE
	98	+ void InitUnsharedMaskV2(uint32_t unsharedKvSeqLen, uint32_t maxDecodeStep, uint32_t headNum, uint32_t groupSize)
	99	+ {
	100	+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
	101	+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
	102	+ uint32_t curHeadSplitSubBlock = headNum / subBlockNum;
	103	+ uint32_t curHeadThisSubBlock = (subBlockIdx == 0) ? curHeadSplitSubBlock : (headNum - curHeadSplitSubBlock);
	104	+ uint64_t headOffset = (subBlockIdx == 0) ? 0 : curHeadSplitSubBlock;
	105	+
	106	+ uint32_t kSeqTileRound = (maxDecodeStep * headNum + SEQ_TILE_SIZE - 1) / SEQ_TILE_SIZE * SEQ_TILE_SIZE;
	107	+ AscendC::Duplicate(
	108	+ unsharedMaskUbTensor, std::numeric_limits<float>::lowest(),
	109	+ kSeqTileRound * curHeadThisSubBlock * groupSize);
	110	+ AscendC::PipeBarrier<PIPE_V>();
	111	+ for (uint32_t round = 0; round < curHeadThisSubBlock; ++round) {
	112	+ uint8_t repeatStride = kSeqTileRound * sizeof(ElementInput) / 32;
	113	+ uint32_t colOffset = (headOffset + round) * maxDecodeStep;
	114	+ uint32_t colOffsetFloor = colOffset / FLOAT_BLOCK_SIZE * FLOAT_BLOCK_SIZE;
	115	+ uint32_t floorSub = colOffset - colOffsetFloor;
	116	+ uint64_t rowOffset = round * groupSize * kSeqTileRound;
	117	+ uint64_t totalOffset = rowOffset + colOffsetFloor;
	118	+ auto totalDupLen = unsharedKvSeqLen + floorSub;
	119	+ if (totalDupLen > FLOAT_VECTOR_SIZE) {
	120	+ for (uint32_t loopIdx = 0; loopIdx < groupSize; ++loopIdx) {
	121	+ AscendC::Duplicate(
	122	+ unsharedMaskUbTensor[totalOffset + loopIdx * kSeqTileRound], float(0.f),
	123	+ static_cast<int32_t>(totalDupLen));
	124	+ }
	125	+ } else {
	126	+ AscendC::Duplicate(
	127	+ unsharedMaskUbTensor[totalOffset], float(0.f), static_cast<uint64_t>(totalDupLen), groupSize, 1,
	128	+ repeatStride);
	129	+ }
	130	+ AscendC::PipeBarrier<PIPE_V>();
	131	+
	132	+ if (floorSub > 0) {
	133	+ if (floorSub > FLOAT_VECTOR_SIZE) {
	134	+ for (uint32_t loopIdx = 0; loopIdx < groupSize; ++loopIdx) {
	135	+ AscendC::Duplicate(
	136	+ unsharedMaskUbTensor[totalOffset + loopIdx * kSeqTileRound],
	137	+ std::numeric_limits<float>::lowest(), static_cast<int32_t>(floorSub));
	138	+ }
	139	+ } else {
	140	+ AscendC::Duplicate(
	141	+ unsharedMaskUbTensor[totalOffset], std::numeric_limits<float>::lowest(),
	142	+ static_cast<uint64_t>(floorSub), groupSize, 1, repeatStride);
	143	+ }
	144	+ AscendC::PipeBarrier<PIPE_V>();
	145	+ }
	146	+ }
	147	+ AscendC::ResetMask();
	148	+ }
	149	+
	150	+ CATLASS_DEVICE
	151	+ void SetMask(int32_t len)
	152	+ {
	153	+ uint64_t mask = 0;
	154	+ uint64_t one = 1;
	155	+ uint64_t temp = len % FLOAT_VECTOR_SIZE;
	156	+ for (int64_t i = 0; i < temp; i++) {
	157	+ mask \|= one << i;
	158	+ }
	159	+ if (len == VECTOR_SIZE) {
	160	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	161	+ } else if (len >= FLOAT_VECTOR_SIZE) {
	162	+ AscendC::SetVectorMask<int8_t>(mask, (uint64_t)-1);
	163	+ } else {
	164	+ AscendC::SetVectorMask<int8_t>(0x0, mask);
	165	+ }
	166	+ }
	167	+
	168	+ CATLASS_DEVICE
	169	+ void ReduceSumRepeatM(
	170	+ const AscendC::LocalTensor<float>& dst, const AscendC::LocalTensor<float>& src, uint32_t curRowNum,
	171	+ uint32_t kSeqTile, uint32_t kSeqTileRound)
	172	+ {
	173	+ if (kSeqTile <= FLOAT_VECTOR_SIZE) {
	174	+ AscendC::WholeReduceSum(dst, src, kSeqTile, curRowNum, 1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE);
	175	+ } else {
	176	+ auto repeatTimes = kSeqTile / FLOAT_VECTOR_SIZE;
	177	+ auto tailLen = kSeqTile % FLOAT_VECTOR_SIZE;
	178	+ for (uint32_t rowsum_idx = 1; rowsum_idx < repeatTimes; ++rowsum_idx) {
	179	+ AscendC::Add(
	180	+ src, src, src[rowsum_idx * FLOAT_VECTOR_SIZE], FLOAT_VECTOR_SIZE, curRowNum,
	181	+ AscendC::BinaryRepeatParams(
	182	+ 1, 1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE,
	183	+ kSeqTileRound / FLOAT_BLOCK_SIZE));
	184	+ AscendC::PipeBarrier<PIPE_V>();
	185	+ }
	186	+ if (tailLen > 0) {
	187	+ AscendC::Add(
	188	+ src, src, src[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tailLen, curRowNum,
	189	+ AscendC::BinaryRepeatParams(
	190	+ 1, 1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE,
	191	+ kSeqTileRound / FLOAT_BLOCK_SIZE));
	192	+ }
	193	+ AscendC::PipeBarrier<PIPE_V>();
	194	+ uint32_t mask = repeatTimes > 0 ? FLOAT_VECTOR_SIZE : tailLen;
	195	+ AscendC::WholeReduceSum(dst, src, mask, curRowNum, 1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE);
	196	+ }
	197	+ AscendC::PipeBarrier<PIPE_V>();
	198	+ }
	199	+
	200	+ CATLASS_DEVICE
	201	+ void TensorSubValueRepeatM(
	202	+ const AscendC::LocalTensor<float>& dst, const AscendC::LocalTensor<float>& src,
	203	+ const AscendC::LocalTensor<float>& MaxTensor, const AscendC::LocalTensor<float>& tempMaxTensor,
	204	+ uint32_t curRowNum, uint32_t subMRound, uint32_t kSeqTile, uint32_t kSeqTileRound)
	205	+ {
	206	+ uint8_t brcbRepeatTimes = static_cast<uint8_t>((curRowNum + BRCB_REPEAT_SIZE - 1) / BRCB_REPEAT_SIZE);
	207	+ AscendC::Brcb(tempMaxTensor, MaxTensor, brcbRepeatTimes, AscendC::BrcbRepeatParams(1, 8));
	208	+ AscendC::PipeBarrier<PIPE_V>();
	209	+ for (uint32_t subIdx = 0; subIdx < kSeqTile / FLOAT_VECTOR_SIZE; ++subIdx) {
	210	+ AscendC::Sub(
	211	+ dst[subIdx * FLOAT_VECTOR_SIZE], src[subIdx * FLOAT_VECTOR_SIZE], tempMaxTensor, FLOAT_VECTOR_SIZE,
	212	+ curRowNum,
	213	+ AscendC::BinaryRepeatParams(
	214	+ 1, 1, 0, kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE, 1));
	215	+ }
	216	+
	217	+ if (kSeqTile % FLOAT_VECTOR_SIZE > 0) {
	218	+ AscendC::Sub(
	219	+ dst[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
	220	+ src[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tempMaxTensor, kSeqTile % FLOAT_VECTOR_SIZE,
	221	+ curRowNum,
	222	+ AscendC::BinaryRepeatParams(
	223	+ 1, 1, 0, kSeqTileRound / FLOAT_BLOCK_SIZE, kSeqTileRound / FLOAT_BLOCK_SIZE, 1));
	224	+ }
	225	+ AscendC::PipeBarrier<PIPE_V>();
	226	+ }
	227	+
	228	+ CATLASS_DEVICE
	229	+ void ReduceMaxRepeatM(
	230	+ const AscendC::LocalTensor<float>& dst, const AscendC::LocalTensor<float>& src,
	231	+ const AscendC::LocalTensor<float>& tempTensor, uint32_t curRowNum, uint32_t kSeqTile, uint32_t kSeqTileRound)
	232	+ {
	233	+ if (kSeqTile <= FLOAT_VECTOR_SIZE) {
	234	+ AscendC::WholeReduceMax(
	235	+ dst, src, kSeqTile, curRowNum, 1, 1, kSeqTileRound / FLOAT_BLOCK_SIZE,
	236	+ AscendC::ReduceOrder::ORDER_ONLY_VALUE);
	237	+ } else {
	238	+ AscendC::DataCopy(
	239	+ tempTensor, src,
	240	+ AscendC::DataCopyParams(
	241	+ curRowNum, FLOAT_VECTOR_SIZE / FLOAT_BLOCK_SIZE,
	242	+ (kSeqTileRound - FLOAT_VECTOR_SIZE) / FLOAT_BLOCK_SIZE, 0));
	243	+ AscendC::PipeBarrier<PIPE_V>();
	244	+ auto repeatTimes = kSeqTile / FLOAT_VECTOR_SIZE;
	245	+ auto tailLen = kSeqTile % FLOAT_VECTOR_SIZE;
	246	+ for (uint32_t rowmaxIdx = 1; rowmaxIdx < repeatTimes; ++rowmaxIdx) {
	247	+ AscendC::Max(
	248	+ tempTensor, tempTensor, src[rowmaxIdx * FLOAT_VECTOR_SIZE], FLOAT_VECTOR_SIZE, curRowNum,
	249	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, kSeqTileRound / FLOAT_BLOCK_SIZE));
	250	+ AscendC::PipeBarrier<PIPE_V>();
	251	+ }
	252	+ if (tailLen > 0) {
	253	+ AscendC::Max(
	254	+ tempTensor, tempTensor, src[kSeqTile / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tailLen, curRowNum,
	255	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, kSeqTileRound / FLOAT_BLOCK_SIZE));
	256	+ }
	257	+ AscendC::PipeBarrier<PIPE_V>();
	258	+ uint32_t mask = repeatTimes > 0 ? FLOAT_VECTOR_SIZE : tailLen;
	259	+ AscendC::WholeReduceMax(dst, tempTensor, mask, curRowNum, 1, 1, 8, AscendC::ReduceOrder::ORDER_ONLY_VALUE);
	260	+ }
	261	+ AscendC::PipeBarrier<PIPE_V>();
	262	+ }
	263	+
	264	+ CATLASS_DEVICE
	265	+ void SubCoreCompute(
	266	+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
	267	+ AscendC::GlobalTensor<ElementInput> gmOutput, AscendC::GlobalTensor<ElementInput> glOutput,
	268	+ const LayoutOutput& layoutOutput, const LayoutInput& layoutInput, uint32_t curHeadNum, uint32_t headOffset)
	269	+ {
	270	+ uint32_t curRowNum = layoutInput.shape(0);
	271	+ uint32_t kSeqTile = layoutInput.shape(1);
	272	+ uint32_t kSeqTileRound = layoutInput.stride(0);
	273	+ uint32_t subMRound = (curRowNum + 15) / 16 * 16;
	274	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
	275	+ AscendC::DataCopy(
	276	+ lsUbTensor, gInput, AscendC::DataCopyParams(1, curRowNum * kSeqTileRound / FLOAT_BLOCK_SIZE, 0, 0));
	277	+
	278	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID3);
	279	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID3);
	280	+
	281	+ // muls scale_value
	282	+ AscendC::Muls(lsUbTensor, lsUbTensor, tor, curRowNum * kSeqTileRound);
	283	+ AscendC::PipeBarrier<PIPE_V>();
	284	+ uint32_t groupSize = curRowNum / curHeadNum;
	285	+ AscendC::Add(lsUbTensor, lsUbTensor, unsharedMaskUbTensor, curRowNum * kSeqTileRound);
	286	+ AscendC::PipeBarrier<PIPE_V>();
	287	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID3);
	288	+ // *** lm = rowmax(ls)
	289	+ ReduceMaxRepeatM(lmUbTensor, lsUbTensor, lpUbTensor32, curRowNum, kSeqTile, kSeqTileRound);
	290	+
	291	+ // *** ls = ls - hm_block
	292	+ TensorSubValueRepeatM(
	293	+ lsUbTensor, lsUbTensor, lmUbTensor, tvUbTensor, curRowNum, subMRound, kSeqTile, kSeqTileRound);
	294	+ AscendC::Exp(lsUbTensor, lsUbTensor, curRowNum * kSeqTileRound);
	295	+ AscendC::PipeBarrier<PIPE_V>();
	296	+ // *** lp = castfp32to16(ls)
	297	+ if (std::is_same<ElementOutput, bfloat16_t>::value) {
	298	+ AscendC::Cast(tvUbTensor16, lsUbTensor, AscendC::RoundMode::CAST_RINT, curRowNum * kSeqTileRound);
	299	+ } else {
	300	+ AscendC::Cast(tvUbTensor16, lsUbTensor, AscendC::RoundMode::CAST_NONE, curRowNum * kSeqTileRound);
	301	+ }
	302	+ AscendC::PipeBarrier<PIPE_V>();
	303	+ // *** ll = rowsum(ls32)
	304	+ ReduceSumRepeatM(llUbTensor, lsUbTensor, curRowNum, kSeqTile, kSeqTileRound);
	305	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID3);
	306	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID3);
	307	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
	308	+
	309	+ uint16_t blockCount = 1;
	310	+ uint16_t blockLen = curRowNum * kSeqTileRound / T_BLOCK_SIZE;
	311	+ uint16_t srcStride = 0;
	312	+ uint16_t dstStride = 0;
	313	+ AscendC::DataCopy(
	314	+ gOutput, tvUbTensor16,
	315	+ AscendC::DataCopyParams(
	316	+ blockCount, // blockCount
	317	+ blockLen, // blockLen
	318	+ srcStride, // srcGap
	319	+ dstStride));
	320	+
	321	+ auto copyLen = curHeadNum * groupSize;
	322	+ if (copyLen % FLOAT_BLOCK_SIZE == 0) {
	323	+ AscendC::DataCopy(gmOutput[headOffset * groupSize], lmUbTensor, copyLen);
	324	+ AscendC::DataCopy(glOutput[headOffset * groupSize], llUbTensor, copyLen);
	325	+ } else {
	326	+ AscendC::DataCopyExtParams copyOutParams;
	327	+ copyOutParams.blockCount = 1;
	328	+ copyOutParams.blockLen = static_cast<uint32_t>(copyLen * sizeof(float));
	329	+ copyOutParams.srcStride = 0;
	330	+ copyOutParams.dstStride = 0;
	331	+ copyOutParams.rsv = 0;
	332	+ AscendC::DataCopyPad(gmOutput[headOffset * groupSize], lmUbTensor, copyOutParams);
	333	+ AscendC::DataCopyPad(glOutput[headOffset * groupSize], llUbTensor, copyOutParams);
	334	+ }
	335	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID3);
	336	+ }
	337	+
	338	+ CATLASS_DEVICE
	339	+ void operator()(
	340	+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
	341	+ AscendC::GlobalTensor<ElementInput> gmOutput, AscendC::GlobalTensor<ElementInput> glOutput,
	342	+ const LayoutOutput& layoutOutput, const LayoutInput& layoutInput, GemmCoord actualBlockShape,
	343	+ uint32_t curHeadNum)
	344	+ {
	345	+ uint32_t rowActual = actualBlockShape.m();
	346	+ uint32_t nActual = actualBlockShape.n();
	347	+ uint32_t tokenNumPerHead = rowActual / curHeadNum;
	348	+
	349	+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
	350	+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
	351	+
	352	+ uint32_t curHeadSplitSubBlock = curHeadNum / subBlockNum;
	353	+ uint32_t curHeadThisSubBlock = (subBlockIdx == 0) ? curHeadSplitSubBlock : (curHeadNum - curHeadSplitSubBlock);
	354	+
	355	+ uint32_t rowActualThisSubBlock = curHeadThisSubBlock * tokenNumPerHead;
	356	+ uint32_t rowOffsetSubBlock = subBlockIdx * curHeadSplitSubBlock * tokenNumPerHead;
	357	+
	358	+ if (rowActualThisSubBlock > 0) {
	359	+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetSubBlock, 0));
	360	+ auto gInputThisSubBlock = gInput[offsetInput];
	361	+ auto layoutInputThisSubBlock = layoutInput.GetTileLayout(MatrixCoord(rowActualThisSubBlock, nActual));
	362	+ int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetSubBlock, 0));
	363	+ auto gOutputThisSubBlock = gOutput[offsetOutput];
	364	+ auto layoutOutputThisSubBlock = layoutOutput.GetTileLayout(MatrixCoord(rowActualThisSubBlock, nActual));
	365	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
	366	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID3);
	367	+ SubCoreCompute(
	368	+ gOutputThisSubBlock, gInputThisSubBlock, gmOutput, glOutput, layoutOutputThisSubBlock,
	369	+ layoutInputThisSubBlock, curHeadThisSubBlock, (subBlockIdx == 0) ? 0 : curHeadSplitSubBlock);
	370	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
	371	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID3);
	372	+ }
	373	+ }
	374	+
	375	+private:
	376	+ float tor;
	377	+ uint32_t pingpongFlag = 0;
	378	+ uint32_t kvSplitCoreNum = 1;
	379	+ AscendC::LocalTensor<ElementOutput> tvUbTensor16;
	380	+ AscendC::LocalTensor<float> lpUbTensor32;
	381	+ AscendC::LocalTensor<float> lsUbTensor;
	382	+ AscendC::LocalTensor<float> lmUbTensor;
	383	+ AscendC::LocalTensor<float> llUbTensor;
	384	+ AscendC::LocalTensor<float> tvUbTensor;
	385	+ AscendC::LocalTensor<float> unsharedMaskUbTensor;
	386	+
	387	+ CopyGmToUbInput copyGmToUbInput;
	388	+ CopyGmToUbMask copyGmToUbMask;
	389	+ CopyUbToGmOutput copyUbToGmOutput;
	390	+};
	391	+
	392	+// Ascend950 specialization - same implementation as AtlasA2 since Ascend C APIs are compatible
	393	+template <class OutputType_, class InputType_, class MaskType_>
	394	+class BlockEpilogue<EpilogueAscend950FAUnsharedSoftmax, OutputType_, InputType_, MaskType_>
	395	+ : public BlockEpilogue<EpilogueAtlasA2FAUnsharedSoftmax, OutputType_, InputType_, MaskType_> {
	396	+public:
	397	+ using Base = BlockEpilogue<EpilogueAtlasA2FAUnsharedSoftmax, OutputType_, InputType_, MaskType_>;
	398	+ using DispatchPolicy = EpilogueAscend950FAUnsharedSoftmax;
	399	+ using ArchTag = typename DispatchPolicy::ArchTag;
	400	+ using Base::Base;
	401	+};
	402	+
	403	+} // namespace Catlass::Epilogue::Block
	404	+
	405	+#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_FA_UNSHARED_SOFTMAX_HPP

Ainclude/catlass/epilogue/block/block_epilogue_online_softmax_FD.hpp+856-0

@@ -0,0 +1,856 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_FD_HPP
	13	+#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_FD_HPP
	14	+
	15	+#include "catlass/arch/cross_core_sync.hpp"
	16	+#include "catlass/arch/resource.hpp"
	17	+#include "catlass/epilogue/dispatch_policy.hpp"
	18	+#include "catlass/epilogue/tile/tile_copy.hpp"
	19	+#include "catlass/gemm_coord.hpp"
	20	+#include "catlass/matrix_coord.hpp"
	21	+
	22	+namespace Catlass::Epilogue::Block {
	23	+
	24	+template <class OutputType_, class InputType_, class MaskType_, LseMode LSE_MODE_>
	25	+class BlockEpilogue<EpilogueAtlasA2OnlineSoftmax_FD<LSE_MODE_>, OutputType_, InputType_, MaskType_> {
	26	+public:
	27	+ using DispatchPolicy = EpilogueAtlasA2OnlineSoftmax_FD<LSE_MODE_>;
	28	+ using ArchTag = typename DispatchPolicy::ArchTag;
	29	+ using ElementOutput = typename OutputType_::Element;
	30	+ using ElementInput = typename InputType_::Element;
	31	+ using ElementMask = typename MaskType_::Element;
	32	+
	33	+ using LayoutOutput = typename OutputType_::Layout;
	34	+ using LayoutInput = typename InputType_::Layout;
	35	+ using LayoutMask = typename MaskType_::Layout;
	36	+
	37	+ static constexpr LseMode LSE_MODE = DispatchPolicy::LSE_MODE;
	38	+
	39	+ static constexpr uint32_t BLOCK_SIZE_IN_BYTE = 32;
	40	+ static constexpr uint32_t REPEAT_SIZE_IN_BYTE = 256;
	41	+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
	42	+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
	43	+ static constexpr uint32_t HALF_VECTOR_SIZE = 128;
	44	+ static constexpr uint32_t BLOCK_SIZE = 16;
	45	+ static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024;
	46	+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384;
	47	+ static constexpr uint32_t VECTOR_SIZE = 128;
	48	+ static constexpr uint32_t MAX_UB_S_ELEM_NUM = 8192;
	49	+
	50	+ static constexpr uint32_t REDUCE_UB_SIZE = 1024;
	51	+ static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32;
	52	+ static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4;
	53	+ static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 256;
	54	+ static constexpr int64_t UB_FLOAT_LINE_SIZE = 64;
	55	+
	56	+ __aicore__ inline BlockEpilogue(Arch::Resource<ArchTag>& resource, float scaleValue_)
	57	+ {
	58	+ // Allocate UB space
	59	+ constexpr uint32_t LS_UB_TENSOR_OFFSET = 0;
	60	+ constexpr uint32_t LP_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE;
	61	+ constexpr uint32_t MASK_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE;
	62	+ constexpr uint32_t MASK32_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE;
	63	+
	64	+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE;
	65	+ constexpr uint32_t LM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 8 * UB_UINT8_VECTOR_SIZE;
	66	+
	67	+ constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE;
	68	+ constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE;
	69	+ constexpr uint32_t LL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 11 * UB_UINT8_VECTOR_SIZE;
	70	+ constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE;
	71	+ constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE;
	72	+
	73	+ constexpr uint32_t MASK16_UB_TENSOR_OFFSET = 11 * UB_UINT8_BLOCK_SIZE;
	74	+
	75	+ scaleValue = scaleValue_;
	76	+ lsUbTensor = resource.ubBuf.template GetBufferByByte<float>(LS_UB_TENSOR_OFFSET);
	77	+ lpUbTensor = resource.ubBuf.template GetBufferByByte<ElementOutput>(LP_UB_TENSOR_OFFSET);
	78	+ maskUbTensor = resource.ubBuf.template GetBufferByByte<ElementMask>(MASK_UB_TENSOR_OFFSET);
	79	+ maskUbTensor16 = resource.ubBuf.template GetBufferByByte<half>(MASK16_UB_TENSOR_OFFSET);
	80	+ maskUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(MASK32_UB_TENSOR_OFFSET);
	81	+ lmUbTensor = resource.ubBuf.template GetBufferByByte<float>(LM_UB_TENSOR_OFFSET);
	82	+ hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);
	83	+ gmUbTensor = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
	84	+ dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
	85	+ llUbTensor = resource.ubBuf.template GetBufferByByte<float>(LL_UB_TENSOR_OFFSET);
	86	+ tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
	87	+ glUbTensor = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
	88	+ }
	89	+
	90	+ __aicore__ inline ~BlockEpilogue()
	91	+ {}
	92	+
	93	+ template <typename T>
	94	+ __aicore__ inline T Min(T a, T b)
	95	+ {
	96	+ return (a > b) ? b : a;
	97	+ }
	98	+
	99	+ __aicore__ inline void SetVecMask(int32_t len)
	100	+ {
	101	+ uint64_t mask = 0;
	102	+ uint64_t one = 1;
	103	+ uint64_t temp = len % FLOAT_VECTOR_SIZE;
	104	+ for (int64_t i = 0; i < temp; i++) {
	105	+ mask \|= one << i;
	106	+ }
	107	+
	108	+ if (len == VECTOR_SIZE \|\| len == 0) {
	109	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	110	+ } else if (len >= FLOAT_VECTOR_SIZE) {
	111	+ AscendC::SetVectorMask<int8_t>(mask, (uint64_t)-1);
	112	+ } else {
	113	+ AscendC::SetVectorMask<int8_t>(0x0, mask);
	114	+ }
	115	+ }
	116	+
	117	+ __aicore__ inline void SetBlockReduceMask(int32_t len)
	118	+ {
	119	+ if (len > 8 \|\| len < 1) {
	120	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	121	+ return;
	122	+ }
	123	+ uint64_t subMask = ((uint64_t)1 << len) - 1;
	124	+ uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask + (subMask << 56) +
	125	+ (subMask << 40) + (subMask << 24) + (subMask << 8);
	126	+ AscendC::SetVectorMask<int8_t>(maskValue, maskValue);
	127	+ }
	128	+
	129	+ __aicore__ inline void RowsumSPECTILE512(
	130	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowsumUb,
	131	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	132	+ uint32_t numElemsAligned)
	133	+ {
	134	+ AscendC::BlockReduceSum<float, false>(
	135	+ tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	136	+ AscendC::PipeBarrier<PIPE_V>();
	137	+
	138	+ AscendC::BlockReduceSum<float, false>(
	139	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
	140	+ numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	141	+ AscendC::PipeBarrier<PIPE_V>();
	142	+ AscendC::BlockReduceSum<float, false>(
	143	+ rowsumUb, tvUbTensor[REDUCE_UB_SIZE],
	144	+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	145	+ AscendC::PipeBarrier<PIPE_V>();
	146	+ }
	147	+
	148	+ __aicore__ inline void RowsumSPECTILE256(
	149	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowsumUb,
	150	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	151	+ uint32_t numElemsAligned)
	152	+ {
	153	+ AscendC::BlockReduceSum<float, false>(
	154	+ tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	155	+ AscendC::PipeBarrier<PIPE_V>();
	156	+ SetVecMask(ROW_OPS_SPEC_MASK_32);
	157	+ AscendC::BlockReduceSum<float, false>(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4);
	158	+ AscendC::PipeBarrier<PIPE_V>();
	159	+ SetBlockReduceMask(ROW_OPS_SPEC_MASK_4);
	160	+ AscendC::BlockReduceSum<float, false>(
	161	+ rowsumUb, tvUbTensor[REDUCE_UB_SIZE], CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), 0, 1, 1,
	162	+ 8);
	163	+ AscendC::PipeBarrier<PIPE_V>();
	164	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	165	+ }
	166	+
	167	+ __aicore__ inline void RowsumTAILTILE(
	168	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowsumUb,
	169	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	170	+ uint32_t numElemsAligned)
	171	+ {
	172	+ if (numElems >= FLOAT_VECTOR_SIZE) {
	173	+ AscendC::BlockReduceSum<float, false>(
	174	+ tvUbTensor, srcUb, numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
	175	+ AscendC::PipeBarrier<PIPE_V>();
	176	+ AscendC::BlockReduceSum<float, false>(
	177	+ rowsumUb, tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), 0, 1, 1, 8);
	178	+ AscendC::PipeBarrier<PIPE_V>();
	179	+ for (uint64_t rowSumIdx = 1; rowSumIdx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowSumIdx) {
	180	+ AscendC::BlockReduceSum<float, false>(
	181	+ tvUbTensor, srcUb[rowSumIdx * FLOAT_VECTOR_SIZE], numRowsRound, 0, 1, 1,
	182	+ numElemsAligned / FLOAT_BLOCK_SIZE);
	183	+ AscendC::PipeBarrier<PIPE_V>();
	184	+ AscendC::BlockReduceSum<float, false>(
	185	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
	186	+ 0, 1, 1, 8);
	187	+ AscendC::PipeBarrier<PIPE_V>();
	188	+ SetVecMask(numRowsRound);
	189	+ AscendC::Add<float, false>(
	190	+ rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
	191	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	192	+ AscendC::PipeBarrier<PIPE_V>();
	193	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	194	+ }
	195	+ }
	196	+ if (numElems % FLOAT_VECTOR_SIZE > 0) {
	197	+ SetVecMask(numElems % FLOAT_VECTOR_SIZE);
	198	+ AscendC::BlockReduceSum<float, false>(
	199	+ tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], numRowsRound, 0, 1, 1,
	200	+ numElemsAligned / FLOAT_BLOCK_SIZE);
	201	+ AscendC::PipeBarrier<PIPE_V>();
	202	+ SetBlockReduceMask(CeilDiv(numElems % FLOAT_VECTOR_SIZE, FLOAT_BLOCK_SIZE));
	203	+ if (numElems < FLOAT_VECTOR_SIZE) {
	204	+ AscendC::BlockReduceSum<float, false>(
	205	+ rowsumUb, tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), 0, 1, 1, 8);
	206	+ AscendC::PipeBarrier<PIPE_V>();
	207	+ } else {
	208	+ AscendC::BlockReduceSum<float, false>(
	209	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
	210	+ 0, 1, 1, 8);
	211	+ AscendC::PipeBarrier<PIPE_V>();
	212	+ SetVecMask(numRowsRound);
	213	+ AscendC::Add<float, false>(
	214	+ rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
	215	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	216	+ AscendC::PipeBarrier<PIPE_V>();
	217	+ }
	218	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	219	+ }
	220	+ }
	221	+
	222	+ __aicore__ inline void RowmaxSPECTILE512(
	223	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowmaxUb,
	224	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	225	+ uint32_t numElemsAligned)
	226	+ {
	227	+ AscendC::BlockReduceMax<float, false>(
	228	+ tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	229	+ AscendC::PipeBarrier<PIPE_V>();
	230	+ AscendC::BlockReduceMax<float, false>(
	231	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
	232	+ numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	233	+ AscendC::PipeBarrier<PIPE_V>();
	234	+ AscendC::BlockReduceMax<float, false>(
	235	+ rowmaxUb, tvUbTensor[REDUCE_UB_SIZE],
	236	+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	237	+ AscendC::PipeBarrier<PIPE_V>();
	238	+ }
	239	+
	240	+ __aicore__ inline void RowmaxSPECTILE256(
	241	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowmaxUb,
	242	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	243	+ uint32_t numElemsAligned)
	244	+ {
	245	+ AscendC::BlockReduceMax<float, false>(
	246	+ tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	247	+ AscendC::PipeBarrier<PIPE_V>();
	248	+ SetVecMask(ROW_OPS_SPEC_MASK_32);
	249	+ AscendC::BlockReduceMax<float, false>(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4);
	250	+ AscendC::PipeBarrier<PIPE_V>();
	251	+ SetBlockReduceMask(ROW_OPS_SPEC_MASK_4);
	252	+ AscendC::BlockReduceMax<float, false>(
	253	+ rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), 0, 1, 1,
	254	+ 8);
	255	+ AscendC::PipeBarrier<PIPE_V>();
	256	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	257	+ }
	258	+
	259	+ __aicore__ inline void RowmaxTAILTILE(
	260	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowmaxUb,
	261	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	262	+ uint32_t numElemsAligned)
	263	+ {
	264	+ if (numElems >= FLOAT_VECTOR_SIZE) {
	265	+ AscendC::BlockReduceMax<float, false>(
	266	+ tvUbTensor, srcUb, numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
	267	+ AscendC::PipeBarrier<PIPE_V>();
	268	+ AscendC::BlockReduceMax<float, false>(
	269	+ rowmaxUb, tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), 0, 1, 1, 8);
	270	+ AscendC::PipeBarrier<PIPE_V>();
	271	+ for (uint64_t rowmax_idx = 1; rowmax_idx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowmax_idx) {
	272	+ AscendC::BlockReduceMax<float, false>(
	273	+ tvUbTensor, srcUb[rowmax_idx * FLOAT_VECTOR_SIZE], numRowsRound, 0, 1, 1,
	274	+ numElemsAligned / FLOAT_BLOCK_SIZE);
	275	+ AscendC::PipeBarrier<PIPE_V>();
	276	+ AscendC::BlockReduceMax<float, false>(
	277	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
	278	+ 0, 1, 1, 8);
	279	+ AscendC::PipeBarrier<PIPE_V>();
	280	+ SetVecMask(numRowsRound);
	281	+ AscendC::Max<float, false>(
	282	+ rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
	283	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	284	+ AscendC::PipeBarrier<PIPE_V>();
	285	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	286	+ }
	287	+ }
	288	+ if (numElems % FLOAT_VECTOR_SIZE > 0) {
	289	+ SetVecMask(numElems % FLOAT_VECTOR_SIZE);
	290	+ AscendC::BlockReduceMax<float, false>(
	291	+ tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], numRowsRound, 0, 1, 1,
	292	+ numElemsAligned / FLOAT_BLOCK_SIZE);
	293	+ AscendC::PipeBarrier<PIPE_V>();
	294	+ SetBlockReduceMask(CeilDiv(numElems % FLOAT_VECTOR_SIZE, FLOAT_BLOCK_SIZE));
	295	+ if (numElems < FLOAT_VECTOR_SIZE) {
	296	+ AscendC::BlockReduceMax<float, false>(
	297	+ rowmaxUb, tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), 0, 1, 1, 8);
	298	+ AscendC::PipeBarrier<PIPE_V>();
	299	+ } else {
	300	+ AscendC::BlockReduceMax<float, false>(
	301	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
	302	+ 0, 1, 1, 8);
	303	+ AscendC::PipeBarrier<PIPE_V>();
	304	+ SetVecMask(numRowsRound);
	305	+ AscendC::Max<float, false>(
	306	+ rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
	307	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	308	+ AscendC::PipeBarrier<PIPE_V>();
	309	+ }
	310	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	311	+ }
	312	+ }
	313	+
	314	+ __aicore__ inline void CopySGmToUb(
	315	+ AscendC::GlobalTensor<ElementInput> gInput, uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound,
	316	+ uint32_t columnNumPad)
	317	+ {
	318	+ AscendC::DataCopy(
	319	+ lsUbTensor[sUbOffset], gInput,
	320	+ AscendC::DataCopyParams(
	321	+ rowNumCurLoop, columnNumRound / FLOAT_BLOCK_SIZE, (columnNumPad - columnNumRound) / FLOAT_BLOCK_SIZE,
	322	+ 0));
	323	+ }
	324	+
	325	+ __aicore__ inline void CopyMaskGmToUb(
	326	+ AscendC::GlobalTensor<ElementMask> gMask, uint32_t columnNum, uint32_t columnNumRound, uint32_t maskStride,
	327	+ uint32_t tokenNumPerHead, uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t integralHeadNum,
	328	+ uint32_t epiTokenNum)
	329	+ {
	330	+ uint32_t innerUbRowOffset = 0;
	331	+ if (proTokenNum != 0) {
	332	+ AscendC::DataCopyPad(
	333	+ maskUbTensor[innerUbRowOffset], gMask[proTokenIdx * maskStride],
	334	+ AscendC::DataCopyExtParams(
	335	+ proTokenNum, columnNum * sizeof(ElementMask), (maskStride - columnNum) * sizeof(ElementMask), 0, 0),
	336	+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
	337	+ innerUbRowOffset += proTokenNum * columnNumRound;
	338	+ }
	339	+ for (uint32_t headIdx = 0; headIdx < integralHeadNum; headIdx++) {
	340	+ AscendC::DataCopyPad(
	341	+ maskUbTensor[innerUbRowOffset], gMask,
	342	+ AscendC::DataCopyExtParams(
	343	+ tokenNumPerHead, columnNum * sizeof(ElementMask), (maskStride - columnNum) * sizeof(ElementMask), 0,
	344	+ 0),
	345	+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
	346	+ innerUbRowOffset += tokenNumPerHead * columnNumRound;
	347	+ }
	348	+ if (epiTokenNum != 0) {
	349	+ AscendC::DataCopyPad(
	350	+ maskUbTensor[innerUbRowOffset], gMask,
	351	+ AscendC::DataCopyExtParams(
	352	+ epiTokenNum, columnNum * sizeof(ElementMask), (maskStride - columnNum) * sizeof(ElementMask), 0, 0),
	353	+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
	354	+ }
	355	+ }
	356	+
	357	+ __aicore__ inline void ScaleS(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound)
	358	+ {
	359	+ AscendC::Muls<float, false>(
	360	+ lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], scaleValue, (uint64_t)0,
	361	+ CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE), AscendC::UnaryRepeatParams(1, 1, 8, 8));
	362	+
	363	+ AscendC::PipeBarrier<PIPE_V>();
	364	+ }
	365	+
	366	+ template <typename ElementMaskDst, typename ElementMaskSrc>
	367	+ __aicore__ inline void UpCastMask(
	368	+ const AscendC::LocalTensor<ElementMaskDst>& maskUbTensorDst,
	369	+ const AscendC::LocalTensor<ElementMaskSrc>& maskUbTensorSrc, uint32_t rowNumCurLoop, uint32_t columnNumRound)
	370	+ {
	371	+ AscendC::Cast<ElementMaskDst, ElementMaskSrc, false>(
	372	+ maskUbTensorDst, maskUbTensorSrc, AscendC::RoundMode::CAST_NONE, (uint64_t)0,
	373	+ CeilDiv(rowNumCurLoop * columnNumRound, (uint32_t)(REPEAT_SIZE_IN_BYTE / sizeof(ElementMaskDst))),
	374	+ AscendC::UnaryRepeatParams(1, 1, 8, 4));
	375	+ AscendC::PipeBarrier<PIPE_V>();
	376	+ }
	377	+
	378	+ __aicore__ inline void ApplyMask(
	379	+ uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound, uint32_t maskColumnRound,
	380	+ uint32_t addMaskUbOffset)
	381	+ {
	382	+ AscendC::Muls<float, false>(
	383	+ maskUbTensor32, maskUbTensor32, (float)-3e38, (uint64_t)0,
	384	+ CeilDiv(rowNumCurLoop * maskColumnRound, FLOAT_VECTOR_SIZE), AscendC::UnaryRepeatParams(1, 1, 8, 8));
	385	+ AscendC::PipeBarrier<PIPE_V>();
	386	+ if (maskColumnRound == columnNumRound) {
	387	+ AscendC::Add<float, false>(
	388	+ lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], maskUbTensor32, (uint64_t)0,
	389	+ CeilDiv(rowNumCurLoop * maskColumnRound, FLOAT_VECTOR_SIZE),
	390	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	391	+ } else {
	392	+ uint32_t loop = maskColumnRound / FLOAT_VECTOR_SIZE;
	393	+ for (uint32_t i = 0; i < loop; i++) {
	394	+ AscendC::Add<float, false>(
	395	+ lsUbTensor[sUbOffset][addMaskUbOffset + i * FLOAT_VECTOR_SIZE],
	396	+ lsUbTensor[sUbOffset][addMaskUbOffset + i * FLOAT_VECTOR_SIZE],
	397	+ maskUbTensor32[i * FLOAT_VECTOR_SIZE], (uint64_t)0, rowNumCurLoop,
	398	+ AscendC::BinaryRepeatParams(
	399	+ 1, 1, 1, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE,
	400	+ maskColumnRound / FLOAT_BLOCK_SIZE));
	401	+ }
	402	+ if (maskColumnRound % FLOAT_VECTOR_SIZE > 0) {
	403	+ SetVecMask(maskColumnRound % FLOAT_VECTOR_SIZE);
	404	+ AscendC::Add<float, false>(
	405	+ lsUbTensor[sUbOffset][addMaskUbOffset + loop * FLOAT_VECTOR_SIZE],
	406	+ lsUbTensor[sUbOffset][addMaskUbOffset + loop * FLOAT_VECTOR_SIZE],
	407	+ maskUbTensor32[loop * FLOAT_VECTOR_SIZE], (uint64_t)0, rowNumCurLoop,
	408	+ AscendC::BinaryRepeatParams(
	409	+ 1, 1, 1, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE,
	410	+ maskColumnRound / FLOAT_BLOCK_SIZE));
	411	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	412	+ }
	413	+ }
	414	+ AscendC::PipeBarrier<PIPE_V>();
	415	+ }
	416	+
	417	+ __aicore__ inline void CalcLocalRowMax(
	418	+ uint32_t sUbOffset, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound,
	419	+ uint32_t rowOffset)
	420	+ {
	421	+ if (columnNum == 512) {
	422	+ RowmaxSPECTILE512(
	423	+ lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	424	+ columnNumRound);
	425	+ } else if (columnNum == 256) {
	426	+ RowmaxSPECTILE256(
	427	+ lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	428	+ columnNumRound);
	429	+ } else {
	430	+ RowmaxTAILTILE(
	431	+ lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	432	+ columnNumRound);
	433	+ }
	434	+ }
	435	+
	436	+ __aicore__ inline void UpdateGlobalRowMax(
	437	+ uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound,
	438	+ uint32_t dmUbOffsetCurCycle, uint32_t rowOffset, uint32_t isFirstStackTile)
	439	+ {
	440	+ if (isFirstStackTile) {
	441	+ AscendC::DataCopy(
	442	+ hmUbTensor[rowOffset], lmUbTensor[rowOffset],
	443	+ AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0));
	444	+ AscendC::PipeBarrier<PIPE_V>();
	445	+ } else {
	446	+ SetVecMask(rowNumCurLoop);
	447	+ // *** hm = vmax(lm, gm)
	448	+ AscendC::Max<float, false>(
	449	+ hmUbTensor[rowOffset], lmUbTensor[rowOffset], gmUbTensor[rowOffset], (uint64_t)0, 1,
	450	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	451	+ AscendC::PipeBarrier<PIPE_V>();
	452	+ // *** dm = gm - hm
	453	+ AscendC::Sub<float, false>(
	454	+ dmUbTensor[dmUbOffsetCurCycle], gmUbTensor[rowOffset], hmUbTensor[rowOffset], (uint64_t)0, 1,
	455	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	456	+ AscendC::PipeBarrier<PIPE_V>();
	457	+ // *** dm = exp(dm)
	458	+ AscendC::Exp<float, false>(
	459	+ dmUbTensor[dmUbOffsetCurCycle], dmUbTensor[dmUbOffsetCurCycle], (uint64_t)0, 1,
	460	+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
	461	+ }
	462	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	463	+ AscendC::PipeBarrier<PIPE_V>();
	464	+ // *** gm = hm
	465	+ AscendC::DataCopy(
	466	+ gmUbTensor[rowOffset], hmUbTensor[rowOffset],
	467	+ AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0));
	468	+ AscendC::PipeBarrier<PIPE_V>();
	469	+ }
	470	+
	471	+ __aicore__ inline void CalcExp(
	472	+ uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum,
	473	+ uint32_t columnNumRound, uint32_t rowOffset)
	474	+ {
	475	+ // *** hm_block = expand_to_block(hm), 存放于 tv
	476	+ AscendC::Brcb(
	477	+ tvUbTensor.template ReinterpretCast<uint32_t>(), hmUbTensor[rowOffset].template ReinterpretCast<uint32_t>(),
	478	+ rowNumCurLoopRound / FLOAT_BLOCK_SIZE, AscendC::BrcbRepeatParams(1, 8));
	479	+ AscendC::PipeBarrier<PIPE_V>();
	480	+ // *** ls = ls - hm_block
	481	+ for (uint32_t subIdx = 0; subIdx < columnNum / FLOAT_VECTOR_SIZE; ++subIdx) {
	482	+ AscendC::Sub<float, false>(
	483	+ lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE],
	484	+ tvUbTensor, (uint64_t)0, rowNumCurLoop,
	485	+ AscendC::BinaryRepeatParams(
	486	+ 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1));
	487	+ }
	488	+ if (columnNum % FLOAT_VECTOR_SIZE > 0) {
	489	+ SetVecMask(columnNum % FLOAT_VECTOR_SIZE);
	490	+ AscendC::Sub<float, false>(
	491	+ lsUbTensor[sUbOffset][columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
	492	+ lsUbTensor[sUbOffset][columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0,
	493	+ rowNumCurLoop,
	494	+ AscendC::BinaryRepeatParams(
	495	+ 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1));
	496	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	497	+ }
	498	+ AscendC::PipeBarrier<PIPE_V>();
	499	+ // *** ls = exp(ls)
	500	+ AscendC::Exp<float, false>(
	501	+ lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], (uint64_t)0,
	502	+ CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE), AscendC::UnaryRepeatParams(1, 1, 8, 8));
	503	+ AscendC::PipeBarrier<PIPE_V>();
	504	+ }
	505	+
	506	+ __aicore__ inline void CalcLocalRowSum(
	507	+ uint32_t sUbOffset, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound,
	508	+ uint32_t rowOffset)
	509	+ {
	510	+ // *** ll = rowsum(ls32)
	511	+ if (columnNum == 512) {
	512	+ RowsumSPECTILE512(
	513	+ lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	514	+ columnNumRound);
	515	+ } else if (columnNum == 256) {
	516	+ RowsumSPECTILE256(
	517	+ lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	518	+ columnNumRound);
	519	+ } else {
	520	+ RowsumTAILTILE(
	521	+ lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	522	+ columnNumRound);
	523	+ }
	524	+ }
	525	+
	526	+ __aicore__ inline void UpdateGlobalRowSum(
	527	+ uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t dmUbOffsetCurCycle,
	528	+ uint32_t rowOffset, uint32_t isFirstStackTile)
	529	+ {
	530	+ if (isFirstStackTile) {
	531	+ // *** gl = ll
	532	+ AscendC::DataCopy(
	533	+ glUbTensor[rowOffset], llUbTensor[rowOffset],
	534	+ AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0));
	535	+ AscendC::PipeBarrier<PIPE_V>();
	536	+ } else {
	537	+ SetVecMask(rowNumCurLoop);
	538	+ // *** gl = dm * gl
	539	+ AscendC::Mul<float, false>(
	540	+ glUbTensor[rowOffset], dmUbTensor[dmUbOffsetCurCycle], glUbTensor[rowOffset], (uint64_t)0, 1,
	541	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	542	+ AscendC::PipeBarrier<PIPE_V>();
	543	+ // *** gl = ll + gl
	544	+ AscendC::Add<float, false>(
	545	+ glUbTensor[rowOffset], glUbTensor[rowOffset], llUbTensor[rowOffset], (uint64_t)0, 1,
	546	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	547	+ AscendC::PipeBarrier<PIPE_V>();
	548	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	549	+ }
	550	+ }
	551	+
	552	+ __aicore__ inline void DownCastP(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound)
	553	+ {
	554	+ // *** lp = castfp32to16(ls)
	555	+ if (std::is_same<ElementOutput, bfloat16_t>::value) {
	556	+ AscendC::Cast<ElementOutput, float, false>(
	557	+ lpUbTensor[sUbOffset], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_RINT, (uint64_t)0,
	558	+ CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE), AscendC::UnaryRepeatParams(1, 1, 4, 8));
	559	+ } else {
	560	+ AscendC::Cast<ElementOutput, float, false>(
	561	+ lpUbTensor[sUbOffset], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_NONE, (uint64_t)0,
	562	+ CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE), AscendC::UnaryRepeatParams(1, 1, 4, 8));
	563	+ }
	564	+ }
	565	+
	566	+ __aicore__ inline void CopyPUbToGm(
	567	+ AscendC::GlobalTensor<ElementOutput> gOutput, uint32_t sUbOffset, uint32_t rowNumCurLoop,
	568	+ uint32_t columnNumRound, uint32_t columnNumPad)
	569	+ {
	570	+ AscendC::DataCopy(
	571	+ gOutput, lpUbTensor[sUbOffset],
	572	+ AscendC::DataCopyParams(
	573	+ rowNumCurLoop, columnNumRound / BLOCK_SIZE, 0, (columnNumPad - columnNumRound) / BLOCK_SIZE));
	574	+ }
	575	+
	576	+ template <bool doTriUMask>
	577	+ __aicore__ inline void SubCoreCompute(
	578	+ AscendC::GlobalTensor<ElementOutput> gOutput, const LayoutOutput& layoutOutput, uint32_t rowOffset,
	579	+ uint32_t isFirstStackTile, uint32_t isLastNoMaskStackTile, uint32_t isFirstRowLoop, uint32_t isLastRowLoop,
	580	+ uint32_t columnNumRound, uint32_t pingpongFlag, uint32_t curStackTileMod)
	581	+ {
	582	+ uint32_t rowNumCurLoop = layoutOutput.shape(0);
	583	+ uint32_t rowNumCurLoopRound = RoundUp(rowNumCurLoop, FLOAT_BLOCK_SIZE);
	584	+ uint32_t columnNum = layoutOutput.shape(1);
	585	+ uint32_t columnNumPad = layoutOutput.stride(0);
	586	+ uint32_t sUbOffset = pingpongFlag * MAX_UB_S_ELEM_NUM;
	587	+ uint32_t dmUbOffsetCurCycle = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffset;
	588	+
	589	+ if constexpr (LSE_MODE_ == LseMode::LSE_OUT) {
	590	+ // In lse out-only mode, tv is used in the last stack tile to transport lse
	591	+ if (isFirstStackTile && isFirstRowLoop) {
	592	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
	593	+ }
	594	+ } else {
	595	+ if (isFirstStackTile && isFirstRowLoop) {
	596	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
	597	+ }
	598	+ }
	599	+ CalcLocalRowMax(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
	600	+ UpdateGlobalRowMax(
	601	+ rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, dmUbOffsetCurCycle, rowOffset,
	602	+ isFirstStackTile);
	603	+
	604	+ CalcExp(sUbOffset, rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
	605	+ if constexpr (!doTriUMask) {
	606	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(pingpongFlag);
	607	+ }
	608	+
	609	+ DownCastP(sUbOffset, rowNumCurLoop, columnNumRound);
	610	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
	611	+
	612	+ CalcLocalRowSum(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
	613	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(pingpongFlag);
	614	+
	615	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
	616	+ CopyPUbToGm(gOutput, sUbOffset, rowNumCurLoop, columnNumRound, columnNumPad);
	617	+ if constexpr (!doTriUMask) {
	618	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(pingpongFlag);
	619	+ if (isLastNoMaskStackTile && isLastRowLoop) {
	620	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
	621	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
	622	+ }
	623	+ } else {
	624	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
	625	+ }
	626	+ UpdateGlobalRowSum(
	627	+ sUbOffset, rowNumCurLoop, rowNumCurLoopRound, dmUbOffsetCurCycle, rowOffset, isFirstStackTile);
	628	+ }
	629	+
	630	+ __aicore__ inline void operator()(
	631	+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
	632	+ const LayoutOutput& layoutOutput, const LayoutInput& layoutInput, GemmCoord actualBlockShape,
	633	+ uint32_t isFirstStackTile, uint32_t isLastNoMaskStackTile, uint32_t qSBlockSize, uint32_t qNBlockSize,
	634	+ uint32_t curStackTileMod)
	635	+ {
	636	+ uint32_t rowNum = actualBlockShape.m();
	637	+ uint32_t columnNum = actualBlockShape.n();
	638	+ uint32_t columnNumRound = RoundUp(columnNum, BLOCK_SIZE);
	639	+ uint32_t columnNumPad = layoutInput.stride(0);
	640	+
	641	+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
	642	+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
	643	+
	644	+ uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum;
	645	+ uint32_t rowSplitSubBlock = (qNBlockSize == 1) ? (qSBlockSize / 2) : (qSBlockSize * qNSplitSubBlock);
	646	+ uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? (rowNum - rowSplitSubBlock) : rowSplitSubBlock;
	647	+ uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock;
	648	+ uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound;
	649	+ uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE);
	650	+ rowNumTile = AscendC::Std::min(rowNumTile, FLOAT_VECTOR_SIZE);
	651	+ uint32_t rowLoopNum = CeilDiv(rowActualThisSubBlock, rowNumTile);
	652	+ uint32_t preLoad = 1;
	653	+
	654	+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) {
	655	+ if (rowLoopIdx < rowLoopNum) {
	656	+ uint32_t pingpongFlag = rowLoopIdx % 2;
	657	+ uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile;
	658	+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
	659	+ uint32_t rowNumCurLoop =
	660	+ (rowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
	661	+
	662	+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
	663	+ auto gInputCurLoop = gInput[offsetInput];
	664	+
	665	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(pingpongFlag);
	666	+ CopySGmToUb(
	667	+ gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad);
	668	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
	669	+ }
	670	+ if (rowLoopIdx >= preLoad) {
	671	+ uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad;
	672	+ uint32_t pingpongFlag = delayedRowLoopIdx % 2;
	673	+ uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile;
	674	+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
	675	+ uint32_t rowNumCurLoop =
	676	+ (delayedRowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
	677	+
	678	+ int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
	679	+ auto gOutputCurLoop = gOutput[offsetOutput];
	680	+ auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum));
	681	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
	682	+ ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound);
	683	+ SubCoreCompute<false>(
	684	+ gOutputCurLoop, layoutOutputCurLoop, rowOffsetCurLoop, isFirstStackTile, isLastNoMaskStackTile,
	685	+ delayedRowLoopIdx == 0, delayedRowLoopIdx == rowLoopNum - 1, columnNumRound, pingpongFlag,
	686	+ curStackTileMod);
	687	+ }
	688	+ }
	689	+ }
	690	+
	691	+ __aicore__ inline void operator()(
	692	+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
	693	+ AscendC::GlobalTensor<ElementMask> gMask, const LayoutOutput& layoutOutput, const LayoutInput& layoutInput,
	694	+ const LayoutInput& layoutMask, GemmCoord actualBlockShape, uint32_t isFirstStackTile, uint32_t qSBlockSize,
	695	+ uint32_t qNBlockSize, uint32_t curStackTileMod, Arch::CrossCoreFlag qkReady, uint32_t triUp, uint32_t triDown,
	696	+ uint32_t kvSStartIdx, uint32_t kvSEndIdx)
	697	+ {
	698	+ uint32_t rowNum = actualBlockShape.m();
	699	+ uint32_t columnNum = actualBlockShape.n();
	700	+ uint32_t columnNumRound = RoundUp(columnNum, BLOCK_SIZE_IN_BYTE);
	701	+ uint32_t columnNumPad = layoutInput.stride(0);
	702	+ uint32_t maskStride = layoutMask.stride(0);
	703	+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
	704	+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
	705	+
	706	+ uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum;
	707	+ uint32_t rowSplitSubBlock = (qNBlockSize == 1) ? (qSBlockSize / 2) : (qSBlockSize * qNSplitSubBlock);
	708	+ uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? (rowNum - rowSplitSubBlock) : rowSplitSubBlock;
	709	+ uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock;
	710	+
	711	+ uint32_t tokenNumPerHeadThisSubBlock = Min(qSBlockSize, rowActualThisSubBlock);
	712	+ uint32_t maskOffsetThisSubBlock = (qNBlockSize == 1) ? rowOffsetThisSubBlock : 0;
	713	+
	714	+ // calc mask shift in gm
	715	+ uint32_t gmOffsetMaskRow;
	716	+ uint32_t gmOffsetMaskColumn;
	717	+ uint32_t maskColumn;
	718	+ uint32_t addMaskUbOffset;
	719	+ if (triUp >= kvSStartIdx) {
	720	+ uint32_t triUpRoundDown = RoundDown(triUp, BLOCK_SIZE_IN_BYTE);
	721	+ gmOffsetMaskRow = triUp - triUpRoundDown;
	722	+ gmOffsetMaskColumn = 0;
	723	+ maskColumn = kvSEndIdx - triUpRoundDown;
	724	+ addMaskUbOffset = triUpRoundDown - kvSStartIdx;
	725	+ } else {
	726	+ gmOffsetMaskRow = 0;
	727	+ gmOffsetMaskColumn = kvSStartIdx - triUp;
	728	+ maskColumn = columnNum;
	729	+ addMaskUbOffset = 0;
	730	+ }
	731	+ uint32_t maskColumnRound = RoundUp(maskColumn, BLOCK_SIZE_IN_BYTE);
	732	+
	733	+ int64_t offsetMask =
	734	+ layoutMask.GetOffset(MatrixCoord(gmOffsetMaskRow + maskOffsetThisSubBlock, gmOffsetMaskColumn));
	735	+ auto gMaskThisSubBlock = gMask[offsetMask];
	736	+
	737	+ uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound;
	738	+ uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE);
	739	+ rowNumTile = AscendC::Std::min(rowNumTile, FLOAT_VECTOR_SIZE);
	740	+ uint32_t rowLoopNum = CeilDiv(rowActualThisSubBlock, rowNumTile);
	741	+ uint32_t preLoad = 1;
	742	+
	743	+ if (rowActualThisSubBlock == 0) {
	744	+ Arch::CrossCoreWaitFlag(qkReady);
	745	+ return;
	746	+ }
	747	+
	748	+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) {
	749	+ if (rowLoopIdx < rowLoopNum) {
	750	+ uint32_t pingpongFlag = rowLoopIdx % 2;
	751	+ uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile;
	752	+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
	753	+ uint32_t rowNumCurLoop =
	754	+ (rowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
	755	+ // loop 0 mask load before cross core sync
	756	+ if (rowLoopIdx == 0) {
	757	+ // the token idx of the start token of the prologue part
	758	+ uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock;
	759	+ // the token num of the prologue part
	760	+ uint32_t proTokenNum =
	761	+ Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % tokenNumPerHeadThisSubBlock;
	762	+ // the token num of the epilogue part
	763	+ uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock;
	764	+ // the number of integral heads within a cycle
	765	+ uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock;
	766	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
	767	+ CopyMaskGmToUb(
	768	+ gMaskThisSubBlock, maskColumn, maskColumnRound, maskStride, tokenNumPerHeadThisSubBlock,
	769	+ proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum);
	770	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
	771	+ Arch::CrossCoreWaitFlag(qkReady);
	772	+ }
	773	+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
	774	+ auto gInputCurLoop = gInput[offsetInput];
	775	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(pingpongFlag);
	776	+ CopySGmToUb(
	777	+ gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad);
	778	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
	779	+ }
	780	+ if (rowLoopIdx >= preLoad) {
	781	+ uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad;
	782	+ uint32_t pingpongFlag = delayedRowLoopIdx % 2;
	783	+ uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile;
	784	+ uint32_t rowNumCurLoop =
	785	+ (delayedRowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
	786	+
	787	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
	788	+ UpCastMask<half, ElementMask>(maskUbTensor16, maskUbTensor, rowNumCurLoop, columnNumRound);
	789	+ UpCastMask<float, half>(maskUbTensor32, maskUbTensor16, rowNumCurLoop, columnNumRound);
	790	+
	791	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
	792	+ ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound);
	793	+ ApplyMask(
	794	+ (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, maskColumnRound,
	795	+ addMaskUbOffset);
	796	+ // online softmax vectorized compute
	797	+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
	798	+ int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
	799	+ auto gOutputCurLoop = gOutput[offsetOutput];
	800	+ auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum));
	801	+ SubCoreCompute<true>(
	802	+ gOutputCurLoop, layoutOutputCurLoop, rowOffsetCurLoop, isFirstStackTile, 0, delayedRowLoopIdx == 0,
	803	+ delayedRowLoopIdx == rowLoopNum - 1, columnNumRound, pingpongFlag, curStackTileMod);
	804	+ // next loop mask load
	805	+ if (rowLoopIdx < rowLoopNum) {
	806	+ uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile;
	807	+ uint32_t rowNumCurLoop =
	808	+ (rowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
	809	+ // the token idx of the start token of the prologue part
	810	+ uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock;
	811	+ // the token num of the prologue part
	812	+ uint32_t proTokenNum =
	813	+ Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % tokenNumPerHeadThisSubBlock;
	814	+ // the number of integral heads within a cycle
	815	+ uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock;
	816	+ // the token num of the epilogue part
	817	+ uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock;
	818	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
	819	+ CopyMaskGmToUb(
	820	+ gMaskThisSubBlock, maskColumn, maskColumnRound, maskStride, tokenNumPerHeadThisSubBlock,
	821	+ proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum);
	822	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
	823	+ }
	824	+ }
	825	+ }
	826	+ }
	827	+
	828	+private:
	829	+ float scaleValue;
	830	+ AscendC::LocalTensor<float> lsUbTensor;
	831	+ AscendC::LocalTensor<ElementOutput> lpUbTensor;
	832	+ AscendC::LocalTensor<ElementMask> maskUbTensor;
	833	+ AscendC::LocalTensor<half> maskUbTensor16;
	834	+ AscendC::LocalTensor<float> maskUbTensor32;
	835	+ AscendC::LocalTensor<float> lmUbTensor;
	836	+ AscendC::LocalTensor<float> hmUbTensor;
	837	+ AscendC::LocalTensor<float> gmUbTensor;
	838	+ AscendC::LocalTensor<float> dmUbTensor;
	839	+ AscendC::LocalTensor<float> llUbTensor;
	840	+ AscendC::LocalTensor<float> tvUbTensor;
	841	+ AscendC::LocalTensor<float> glUbTensor;
	842	+};
	843	+
	844	+// Ascend950 specialization - same implementation as AtlasA2 since Ascend C APIs are compatible
	845	+template <class OutputType_, class InputType_, class MaskType_, LseMode LSE_MODE_>
	846	+class BlockEpilogue<EpilogueAscend950OnlineSoftmax_FD<LSE_MODE_>, OutputType_, InputType_, MaskType_>
	847	+ : public BlockEpilogue<EpilogueAtlasA2OnlineSoftmax_FD<LSE_MODE_>, OutputType_, InputType_, MaskType_> {
	848	+public:
	849	+ using Base = BlockEpilogue<EpilogueAtlasA2OnlineSoftmax_FD<LSE_MODE_>, OutputType_, InputType_, MaskType_>;
	850	+ using DispatchPolicy = EpilogueAscend950OnlineSoftmax_FD<LSE_MODE_>;
	851	+ using ArchTag = typename DispatchPolicy::ArchTag;
	852	+ using Base::Base;
	853	+};
	854	+} // namespace Catlass::Epilogue::Block
	855	+
	856	+#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_FD_HPP

Ainclude/catlass/epilogue/block/block_epilogue_online_softmax_copy_glm.hpp+743-0

@@ -0,0 +1,743 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_COPY_SUM_MAX_HPP
	13	+#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_COPY_SUM_MAX_HPP
	14	+
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/arch/cross_core_sync.hpp"
	17	+#include "catlass/arch/resource.hpp"
	18	+#include "catlass/epilogue/dispatch_policy.hpp"
	19	+#include "catlass/epilogue/tile/tile_copy.hpp"
	20	+#include "catlass/gemm_coord.hpp"
	21	+#include "catlass/matrix_coord.hpp"
	22	+
	23	+namespace Catlass::Epilogue::Block {
	24	+
	25	+template <class OutputType_, class InputType_, class MaskType_>
	26	+class BlockEpilogue<EpilogueAtlasA2OnlineSoftmaxCopySumMax, OutputType_, InputType_, MaskType_> {
	27	+public:
	28	+ using DispatchPolicy = EpilogueAtlasA2OnlineSoftmaxCopySumMax;
	29	+ using ArchTag = typename DispatchPolicy::ArchTag;
	30	+ using ElementOutput = typename OutputType_::Element;
	31	+ using ElementInput = typename InputType_::Element;
	32	+ using ElementMask = typename MaskType_::Element;
	33	+ using ElementUpdate = typename InputType_::Element;
	34	+
	35	+ using LayoutOutput = typename OutputType_::Layout;
	36	+ using LayoutInput = typename InputType_::Layout;
	37	+ using LayoutMask = typename MaskType_::Layout;
	38	+
	39	+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
	40	+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
	41	+ static constexpr uint32_t HALF_VECTOR_SIZE = 128;
	42	+ static constexpr uint32_t BLOCK_SIZE = 16;
	43	+ static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024;
	44	+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384;
	45	+ static constexpr uint32_t VECTOR_SIZE = 128;
	46	+ static constexpr uint32_t MAX_UB_S_ELEM_NUM = 8192;
	47	+ static constexpr uint32_t ROW_SUM_PINGPONG_OFFSET = 64 * 8;
	48	+ static constexpr uint32_t REDUCE_UB_SIZE = 1024;
	49	+ static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32;
	50	+ static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4;
	51	+ static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 128;
	52	+ static constexpr int64_t UB_FLOAT_LINE_SIZE = 64;
	53	+ enum class MaskCategory
	54	+ {
	55	+ NO_MASK = 0,
	56	+ CAUSAL_MASK = 1
	57	+ };
	58	+ CATLASS_DEVICE
	59	+ BlockEpilogue(Arch::Resource<ArchTag>& resource, float scaleValue_)
	60	+ {
	61	+ // Allocate UB space
	62	+ constexpr uint32_t LS_UB_TENSOR_OFFSET = 0;
	63	+ constexpr uint32_t LP_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE;
	64	+ constexpr uint32_t MASK32_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE;
	65	+
	66	+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE;
	67	+ constexpr uint32_t LM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 8 * UB_UINT8_VECTOR_SIZE;
	68	+
	69	+ constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE;
	70	+ constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE;
	71	+ constexpr uint32_t LL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 11 * UB_UINT8_VECTOR_SIZE;
	72	+ constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE;
	73	+ constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE;
	74	+
	75	+ constexpr uint32_t MASK_UB_TENSOR_OFFSET = 11 * UB_UINT8_BLOCK_SIZE;
	76	+
	77	+ scaleValue = scaleValue_;
	78	+ lsUbTensor = resource.ubBuf.template GetBufferByByte<float>(LS_UB_TENSOR_OFFSET);
	79	+ lpUbTensor = resource.ubBuf.template GetBufferByByte<ElementOutput>(LP_UB_TENSOR_OFFSET);
	80	+ maskUbTensor = resource.ubBuf.template GetBufferByByte<ElementMask>(MASK_UB_TENSOR_OFFSET);
	81	+ maskUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(MASK32_UB_TENSOR_OFFSET);
	82	+ lmUbTensor = resource.ubBuf.template GetBufferByByte<float>(LM_UB_TENSOR_OFFSET);
	83	+ hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);
	84	+ gmUbTensor = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
	85	+ dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
	86	+ llUbTensor = resource.ubBuf.template GetBufferByByte<float>(LL_UB_TENSOR_OFFSET);
	87	+ tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
	88	+ glUbTensor = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
	89	+ }
	90	+
	91	+ CATLASS_DEVICE
	92	+ ~BlockEpilogue()
	93	+ {}
	94	+
	95	+ template <typename T>
	96	+ CATLASS_DEVICE T Min(T a, T b)
	97	+ {
	98	+ return (a > b) ? b : a;
	99	+ }
	100	+
	101	+ CATLASS_DEVICE
	102	+ void SetVecMask(int32_t len)
	103	+ {
	104	+ uint64_t mask = 0;
	105	+ uint64_t one = 1;
	106	+ uint64_t temp = len % FLOAT_VECTOR_SIZE;
	107	+ for (int64_t i = 0; i < temp; i++) {
	108	+ mask \|= one << i;
	109	+ }
	110	+
	111	+ if (len == VECTOR_SIZE \|\| len == 0) {
	112	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	113	+ } else if (len >= FLOAT_VECTOR_SIZE) {
	114	+ AscendC::SetVectorMask<int8_t>(mask, (uint64_t)-1);
	115	+ } else {
	116	+ AscendC::SetVectorMask<int8_t>(0x0, mask);
	117	+ }
	118	+ }
	119	+
	120	+ CATLASS_DEVICE
	121	+ void SetBlockReduceMask(int32_t len)
	122	+ {
	123	+ if (len > 8 \|\| len < 1) {
	124	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	125	+ return;
	126	+ }
	127	+ uint64_t subMask = ((uint64_t)1 << len) - 1;
	128	+ uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask + (subMask << 56) +
	129	+ (subMask << 40) + (subMask << 24) + (subMask << 8);
	130	+ AscendC::SetVectorMask<int8_t>(maskValue, maskValue);
	131	+ }
	132	+
	133	+ CATLASS_DEVICE
	134	+ void RowsumSPECTILE512(
	135	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowsumUb,
	136	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	137	+ uint32_t numElemsAligned)
	138	+ {
	139	+ AscendC::BlockReduceSum<float, false>(
	140	+ tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	141	+ AscendC::PipeBarrier<PIPE_V>();
	142	+
	143	+ AscendC::BlockReduceSum<float, false>(
	144	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
	145	+ numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	146	+ AscendC::PipeBarrier<PIPE_V>();
	147	+ AscendC::BlockReduceSum<float, false>(
	148	+ rowsumUb, tvUbTensor[REDUCE_UB_SIZE],
	149	+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	150	+ AscendC::PipeBarrier<PIPE_V>();
	151	+ }
	152	+
	153	+ CATLASS_DEVICE
	154	+ void RowsumSPECTILE256(
	155	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowsumUb,
	156	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	157	+ uint32_t numElemsAligned)
	158	+ {
	159	+ AscendC::BlockReduceSum<float, false>(
	160	+ tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	161	+ AscendC::PipeBarrier<PIPE_V>();
	162	+ SetVecMask(ROW_OPS_SPEC_MASK_32);
	163	+ AscendC::BlockReduceSum<float, false>(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4);
	164	+ AscendC::PipeBarrier<PIPE_V>();
	165	+ SetBlockReduceMask(ROW_OPS_SPEC_MASK_4);
	166	+ AscendC::BlockReduceSum<float, false>(
	167	+ rowsumUb, tvUbTensor[REDUCE_UB_SIZE],
	168	+ (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	169	+ AscendC::PipeBarrier<PIPE_V>();
	170	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	171	+ }
	172	+
	173	+ CATLASS_DEVICE
	174	+ void RowsumTAILTILE(
	175	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowsumUb,
	176	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	177	+ uint32_t numElemsAligned)
	178	+ {
	179	+ if (numElems >= FLOAT_VECTOR_SIZE) {
	180	+ AscendC::BlockReduceSum<float, false>(
	181	+ tvUbTensor, srcUb, numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
	182	+ AscendC::PipeBarrier<PIPE_V>();
	183	+ AscendC::BlockReduceSum<float, false>(
	184	+ rowsumUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0,
	185	+ 1, 1, 8);
	186	+ AscendC::PipeBarrier<PIPE_V>();
	187	+ for (uint64_t rowSumIdx = 1; rowSumIdx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowSumIdx) {
	188	+ AscendC::BlockReduceSum<float, false>(
	189	+ tvUbTensor, srcUb[rowSumIdx * FLOAT_VECTOR_SIZE], numRowsRound, 0, 1, 1,
	190	+ numElemsAligned / FLOAT_BLOCK_SIZE);
	191	+ AscendC::PipeBarrier<PIPE_V>();
	192	+ AscendC::BlockReduceSum<float, false>(
	193	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
	194	+ (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	195	+ AscendC::PipeBarrier<PIPE_V>();
	196	+ SetVecMask(numRowsRound);
	197	+ AscendC::Add<float, false>(
	198	+ rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
	199	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	200	+ AscendC::PipeBarrier<PIPE_V>();
	201	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	202	+ }
	203	+ }
	204	+ if (numElems % FLOAT_VECTOR_SIZE > 0) {
	205	+ SetVecMask(numElems % FLOAT_VECTOR_SIZE);
	206	+ AscendC::BlockReduceSum<float, false>(
	207	+ tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], numRowsRound, 0, 1, 1,
	208	+ numElemsAligned / FLOAT_BLOCK_SIZE);
	209	+ AscendC::PipeBarrier<PIPE_V>();
	210	+ SetBlockReduceMask((numElems % FLOAT_VECTOR_SIZE + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE);
	211	+ if (numElems < FLOAT_VECTOR_SIZE) {
	212	+ AscendC::BlockReduceSum<float, false>(
	213	+ rowsumUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	214	+ 0, 1, 1, 8);
	215	+ AscendC::PipeBarrier<PIPE_V>();
	216	+ } else {
	217	+ AscendC::BlockReduceSum<float, false>(
	218	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
	219	+ (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	220	+ AscendC::PipeBarrier<PIPE_V>();
	221	+ SetVecMask(numRowsRound);
	222	+ AscendC::Add<float, false>(
	223	+ rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
	224	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	225	+ AscendC::PipeBarrier<PIPE_V>();
	226	+ }
	227	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	228	+ }
	229	+ }
	230	+
	231	+ CATLASS_DEVICE
	232	+ void RowmaxSPECTILE512(
	233	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowmaxUb,
	234	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	235	+ uint32_t numElemsAligned)
	236	+ {
	237	+ AscendC::BlockReduceMax<float, false>(
	238	+ tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	239	+ AscendC::PipeBarrier<PIPE_V>();
	240	+ AscendC::BlockReduceMax<float, false>(
	241	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
	242	+ numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	243	+ AscendC::PipeBarrier<PIPE_V>();
	244	+ AscendC::BlockReduceMax<float, false>(
	245	+ rowmaxUb, tvUbTensor[REDUCE_UB_SIZE],
	246	+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	247	+ AscendC::PipeBarrier<PIPE_V>();
	248	+ }
	249	+
	250	+ CATLASS_DEVICE
	251	+ void RowmaxSPECTILE256(
	252	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowmaxUb,
	253	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	254	+ uint32_t numElemsAligned)
	255	+ {
	256	+ AscendC::BlockReduceMax<float, false>(
	257	+ tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	258	+ AscendC::PipeBarrier<PIPE_V>();
	259	+ SetVecMask(ROW_OPS_SPEC_MASK_32);
	260	+ AscendC::BlockReduceMax<float, false>(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4);
	261	+ AscendC::PipeBarrier<PIPE_V>();
	262	+ SetBlockReduceMask(ROW_OPS_SPEC_MASK_4);
	263	+ AscendC::BlockReduceMax<float, false>(
	264	+ rowmaxUb, tvUbTensor[REDUCE_UB_SIZE],
	265	+ (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	266	+ AscendC::PipeBarrier<PIPE_V>();
	267	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	268	+ }
	269	+
	270	+ CATLASS_DEVICE
	271	+ void RowmaxTAILTILE(
	272	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowmaxUb,
	273	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	274	+ uint32_t numElemsAligned)
	275	+ {
	276	+ if (numElems >= FLOAT_VECTOR_SIZE) {
	277	+ AscendC::BlockReduceMax<float, false>(
	278	+ tvUbTensor, srcUb, numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
	279	+ AscendC::PipeBarrier<PIPE_V>();
	280	+ AscendC::BlockReduceMax<float, false>(
	281	+ rowmaxUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0,
	282	+ 1, 1, 8);
	283	+ AscendC::PipeBarrier<PIPE_V>();
	284	+ for (uint64_t rowmax_idx = 1; rowmax_idx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowmax_idx) {
	285	+ AscendC::BlockReduceMax<float, false>(
	286	+ tvUbTensor, srcUb[rowmax_idx * FLOAT_VECTOR_SIZE], numRowsRound, 0, 1, 1,
	287	+ numElemsAligned / FLOAT_BLOCK_SIZE);
	288	+ AscendC::PipeBarrier<PIPE_V>();
	289	+ AscendC::BlockReduceMax<float, false>(
	290	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
	291	+ (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	292	+ AscendC::PipeBarrier<PIPE_V>();
	293	+ SetVecMask(numRowsRound);
	294	+ AscendC::Max<float, false>(
	295	+ rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
	296	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	297	+ AscendC::PipeBarrier<PIPE_V>();
	298	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	299	+ }
	300	+ }
	301	+ if (numElems % FLOAT_VECTOR_SIZE > 0) {
	302	+ SetVecMask(numElems % FLOAT_VECTOR_SIZE);
	303	+ AscendC::BlockReduceMax<float, false>(
	304	+ tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], numRowsRound, 0, 1, 1,
	305	+ numElemsAligned / FLOAT_BLOCK_SIZE);
	306	+ AscendC::PipeBarrier<PIPE_V>();
	307	+ SetBlockReduceMask((numElems % FLOAT_VECTOR_SIZE + FLOAT_BLOCK_SIZE - 1) / FLOAT_BLOCK_SIZE);
	308	+ if (numElems < FLOAT_VECTOR_SIZE) {
	309	+ AscendC::BlockReduceMax<float, false>(
	310	+ rowmaxUb, tvUbTensor, (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	311	+ 0, 1, 1, 8);
	312	+ AscendC::PipeBarrier<PIPE_V>();
	313	+ } else {
	314	+ AscendC::BlockReduceMax<float, false>(
	315	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
	316	+ (numRowsRound * FLOAT_BLOCK_SIZE + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	317	+ AscendC::PipeBarrier<PIPE_V>();
	318	+ SetVecMask(numRowsRound);
	319	+ AscendC::Max<float, false>(
	320	+ rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
	321	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	322	+ AscendC::PipeBarrier<PIPE_V>();
	323	+ }
	324	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	325	+ }
	326	+ }
	327	+
	328	+ CATLASS_DEVICE
	329	+ void CopySGmToUb(
	330	+ AscendC::GlobalTensor<ElementInput> gInput, uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound,
	331	+ uint32_t columnNumPad)
	332	+ {
	333	+ // input S
	334	+ AscendC::DataCopy(
	335	+ lsUbTensor[sUbOffset], gInput,
	336	+ AscendC::DataCopyParams(
	337	+ rowNumCurLoop, columnNumRound / FLOAT_BLOCK_SIZE, (columnNumPad - columnNumRound) / FLOAT_BLOCK_SIZE,
	338	+ 0));
	339	+ }
	340	+
	341	+ CATLASS_DEVICE
	342	+ void CopyMaskGmToUb(
	343	+ AscendC::GlobalTensor<ElementMask> gMask, uint32_t columnNum, uint32_t columnNumRound, uint32_t maskStride,
	344	+ uint32_t qSBlockSize, uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t integralHeadNum,
	345	+ uint32_t epiTokenNum)
	346	+ {
	347	+ uint32_t innerUbRowOffset = 0;
	348	+ if (proTokenNum != 0) {
	349	+ AscendC::DataCopyPad(
	350	+ maskUbTensor[innerUbRowOffset], gMask[proTokenIdx * maskStride],
	351	+ AscendC::DataCopyExtParams(proTokenNum, columnNum * 2, (maskStride - columnNum) * 2, 0, 0),
	352	+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
	353	+ innerUbRowOffset += proTokenNum * columnNumRound;
	354	+ }
	355	+ for (uint32_t headIdx = 0; headIdx < integralHeadNum; headIdx++) {
	356	+ AscendC::DataCopyPad(
	357	+ maskUbTensor[innerUbRowOffset], gMask,
	358	+ AscendC::DataCopyExtParams(qSBlockSize, columnNum * 2, (maskStride - columnNum) * 2, 0, 0),
	359	+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
	360	+ innerUbRowOffset += qSBlockSize * columnNumRound;
	361	+ }
	362	+ if (epiTokenNum != 0) {
	363	+ AscendC::DataCopyPad(
	364	+ maskUbTensor[innerUbRowOffset], gMask,
	365	+ AscendC::DataCopyExtParams(epiTokenNum, columnNum * 2, (maskStride - columnNum) * 2, 0, 0),
	366	+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
	367	+ }
	368	+ }
	369	+
	370	+ CATLASS_DEVICE
	371	+ void ScaleS(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound)
	372	+ {
	373	+ // *** ls = scaleValue * ls
	374	+ AscendC::Muls<float, false>(
	375	+ lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], scaleValue, (uint64_t)0,
	376	+ (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	377	+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
	378	+
	379	+ AscendC::PipeBarrier<PIPE_V>();
	380	+ }
	381	+
	382	+ CATLASS_DEVICE
	383	+ void UpCastMask(uint32_t rowNumCurLoop, uint32_t columnNumRound)
	384	+ {
	385	+ AscendC::Cast<float, ElementMask, false>(
	386	+ maskUbTensor32, maskUbTensor, AscendC::RoundMode::CAST_NONE, (uint64_t)0,
	387	+ (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	388	+ AscendC::UnaryRepeatParams(1, 1, 8, 4));
	389	+ AscendC::PipeBarrier<PIPE_V>();
	390	+ }
	391	+
	392	+ CATLASS_DEVICE
	393	+ void ApplyMask(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound)
	394	+ {
	395	+ AscendC::Muls<float, false>(
	396	+ maskUbTensor32, maskUbTensor32, (float)-3e38, (uint64_t)0,
	397	+ (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	398	+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
	399	+ AscendC::PipeBarrier<PIPE_V>();
	400	+ AscendC::Add<float, false>(
	401	+ lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], maskUbTensor32, (uint64_t)0,
	402	+ (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	403	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	404	+ AscendC::PipeBarrier<PIPE_V>();
	405	+ }
	406	+
	407	+ CATLASS_DEVICE
	408	+ void CalcLocalRowMax(
	409	+ uint32_t sUbOffset, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound,
	410	+ uint32_t rowOffset)
	411	+ {
	412	+ if (columnNum == 512) {
	413	+ RowmaxSPECTILE512(
	414	+ lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	415	+ columnNumRound);
	416	+ } else if (columnNum == 256) {
	417	+ RowmaxSPECTILE256(
	418	+ lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	419	+ columnNumRound);
	420	+ } else {
	421	+ RowmaxTAILTILE(
	422	+ lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	423	+ columnNumRound);
	424	+ }
	425	+ }
	426	+
	427	+ CATLASS_DEVICE
	428	+ void UpdateGlobalRowMax(
	429	+ uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound,
	430	+ uint32_t dmUbOffsetCurCycle, uint32_t rowOffset, uint32_t isFirstStackTile)
	431	+ {
	432	+ if (isFirstStackTile) {
	433	+ AscendC::DataCopy(
	434	+ hmUbTensor[rowOffset], lmUbTensor[rowOffset],
	435	+ AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0));
	436	+ AscendC::PipeBarrier<PIPE_V>();
	437	+ } else {
	438	+ SetVecMask(rowNumCurLoop);
	439	+ // *** hm = vmax(lm, gm)
	440	+ AscendC::Max<float, false>(
	441	+ hmUbTensor[rowOffset], lmUbTensor[rowOffset], gmUbTensor[rowOffset], (uint64_t)0, 1,
	442	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	443	+
	444	+ AscendC::PipeBarrier<PIPE_V>();
	445	+ // *** dm = gm - hm
	446	+ AscendC::Sub<float, false>(
	447	+ dmUbTensor[dmUbOffsetCurCycle], gmUbTensor[rowOffset], hmUbTensor[rowOffset], (uint64_t)0, 1,
	448	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	449	+
	450	+ AscendC::PipeBarrier<PIPE_V>();
	451	+ // *** dm = exp(dm)
	452	+ AscendC::Exp<float, false>(
	453	+ dmUbTensor[dmUbOffsetCurCycle], dmUbTensor[dmUbOffsetCurCycle], (uint64_t)0, 1,
	454	+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
	455	+ }
	456	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	457	+ AscendC::PipeBarrier<PIPE_V>();
	458	+ // *** gm = hm
	459	+ AscendC::DataCopy(
	460	+ gmUbTensor[rowOffset], hmUbTensor[rowOffset],
	461	+ AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0));
	462	+ AscendC::PipeBarrier<PIPE_V>();
	463	+ }
	464	+
	465	+ CATLASS_DEVICE
	466	+ void CalcExp(
	467	+ uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum,
	468	+ uint32_t columnNumRound, uint32_t rowOffset)
	469	+ {
	470	+ // *** hm_block = expand_to_block(hm), 存放于 tv
	471	+ AscendC::Brcb(
	472	+ tvUbTensor.template ReinterpretCast<uint32_t>(), hmUbTensor[rowOffset].template ReinterpretCast<uint32_t>(),
	473	+ rowNumCurLoopRound / FLOAT_BLOCK_SIZE, AscendC::BrcbRepeatParams(1, 8));
	474	+ AscendC::PipeBarrier<PIPE_V>();
	475	+ // *** ls = ls - hm_block
	476	+ for (uint32_t subIdx = 0; subIdx < columnNum / FLOAT_VECTOR_SIZE; ++subIdx) {
	477	+ AscendC::Sub<float, false>(
	478	+ lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE],
	479	+ tvUbTensor, (uint64_t)0, rowNumCurLoop,
	480	+ AscendC::BinaryRepeatParams(
	481	+ 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1));
	482	+ }
	483	+ if (columnNum % FLOAT_VECTOR_SIZE > 0) {
	484	+ SetVecMask(columnNum % FLOAT_VECTOR_SIZE);
	485	+ AscendC::Sub<float, false>(
	486	+ lsUbTensor[sUbOffset][columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
	487	+ lsUbTensor[sUbOffset][columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0,
	488	+ rowNumCurLoop,
	489	+ AscendC::BinaryRepeatParams(
	490	+ 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1));
	491	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	492	+ }
	493	+ AscendC::PipeBarrier<PIPE_V>();
	494	+ // *** ls = exp(ls)
	495	+ AscendC::Exp<float, false>(
	496	+ lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], (uint64_t)0,
	497	+ (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	498	+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
	499	+ AscendC::PipeBarrier<PIPE_V>();
	500	+ }
	501	+
	502	+ CATLASS_DEVICE
	503	+ void CalcLocalRowSum(
	504	+ uint32_t sUbOffset, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound,
	505	+ uint32_t rowOffset)
	506	+ {
	507	+ // *** ll = rowsum(ls32)
	508	+ if (columnNum == 512) {
	509	+ RowsumSPECTILE512(
	510	+ lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	511	+ columnNumRound);
	512	+ } else if (columnNum == 256) {
	513	+ RowsumSPECTILE256(
	514	+ lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	515	+ columnNumRound);
	516	+ } else {
	517	+ RowsumTAILTILE(
	518	+ lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	519	+ columnNumRound);
	520	+ }
	521	+ }
	522	+
	523	+ CATLASS_DEVICE
	524	+ void UpdateGlobalRowSum(
	525	+ uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t dmUbOffsetCurCycle,
	526	+ uint32_t rowOffset, uint32_t isFirstStackTile)
	527	+ {
	528	+ if (isFirstStackTile) {
	529	+ // *** gl = ll
	530	+ AscendC::DataCopy(
	531	+ glUbTensor[rowOffset], llUbTensor[rowOffset],
	532	+ AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0));
	533	+ AscendC::PipeBarrier<PIPE_V>();
	534	+ } else {
	535	+ SetVecMask(rowNumCurLoop);
	536	+ // *** gl = dm * gl
	537	+ AscendC::Mul<float, false>(
	538	+ glUbTensor[rowOffset], dmUbTensor[dmUbOffsetCurCycle], glUbTensor[rowOffset], (uint64_t)0, 1,
	539	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	540	+ AscendC::PipeBarrier<PIPE_V>();
	541	+ // *** gl = ll + gl
	542	+ AscendC::Add<float, false>(
	543	+ glUbTensor[rowOffset], glUbTensor[rowOffset], llUbTensor[rowOffset], (uint64_t)0, 1,
	544	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	545	+ AscendC::PipeBarrier<PIPE_V>();
	546	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	547	+ }
	548	+ }
	549	+
	550	+ CATLASS_DEVICE
	551	+ void DownCastP(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound)
	552	+ {
	553	+ // *** lp = castfp32to16(ls)
	554	+ if (std::is_same<ElementOutput, bfloat16_t>::value) {
	555	+ AscendC::Cast<ElementOutput, float, false>(
	556	+ lpUbTensor[sUbOffset], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_RINT, (uint64_t)0,
	557	+ (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	558	+ AscendC::UnaryRepeatParams(1, 1, 4, 8));
	559	+ } else {
	560	+ AscendC::Cast<ElementOutput, float, false>(
	561	+ lpUbTensor[sUbOffset], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_NONE, (uint64_t)0,
	562	+ (rowNumCurLoop * columnNumRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	563	+ AscendC::UnaryRepeatParams(1, 1, 4, 8));
	564	+ }
	565	+ }
	566	+
	567	+ CATLASS_DEVICE
	568	+ void CopyPUbToGm(
	569	+ AscendC::GlobalTensor<ElementOutput> gOutput, uint32_t sUbOffset, uint32_t rowNumCurLoop,
	570	+ uint32_t columnNumRound, uint32_t columnNumPad)
	571	+ {
	572	+ if (columnNumRound == columnNumPad) {
	573	+ AscendC::DataCopy(
	574	+ gOutput, lpUbTensor[sUbOffset],
	575	+ AscendC::DataCopyParams(1, rowNumCurLoop * columnNumRound / BLOCK_SIZE, 0, 0));
	576	+ } else {
	577	+ AscendC::DataCopy(
	578	+ gOutput, lpUbTensor[sUbOffset],
	579	+ AscendC::DataCopyParams(
	580	+ rowNumCurLoop, columnNumRound / BLOCK_SIZE, 0, (columnNumPad - columnNumRound) / BLOCK_SIZE));
	581	+ }
	582	+ }
	583	+
	584	+ template <MaskCategory maskCat>
	585	+ CATLASS_DEVICE void SubCoreCompute(
	586	+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementUpdate> gSharedMax,
	587	+ AscendC::GlobalTensor<ElementUpdate> gSharedSum, const LayoutOutput& layoutOutput, uint32_t rowOffset,
	588	+ uint32_t isFirstStackTile, uint32_t columnNumRound, uint32_t pingpongFlag, uint32_t curStackTileMod,
	589	+ uint32_t isLastStackTile, uint32_t headNum)
	590	+ {
	591	+ uint32_t rowNumCurLoop = layoutOutput.shape(0);
	592	+ uint32_t rowNumCurLoopRound = RoundUp(rowNumCurLoop, FLOAT_BLOCK_SIZE);
	593	+ uint32_t columnNum = layoutOutput.shape(1);
	594	+ // Align colNum to 16, for both float&half compute&copy
	595	+ uint32_t columnNumPad = layoutOutput.stride(0);
	596	+ uint32_t sUbOffset = pingpongFlag * MAX_UB_S_ELEM_NUM;
	597	+ uint32_t dmUbOffsetCurCycle = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffset;
	598	+
	599	+ CalcLocalRowMax(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
	600	+ UpdateGlobalRowMax(
	601	+ rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, dmUbOffsetCurCycle, rowOffset,
	602	+ isFirstStackTile);
	603	+ if (isLastStackTile) {
	604	+ // 保证gSharedMax经过了 Outer的偏移传入
	605	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(6 + pingpongFlag);
	606	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(6 + pingpongFlag);
	607	+ AscendC::LocalTensor<float> sharedGmUbTensor = maskUbTensor.template ReinterpretCast<
	608	+ float>()[pingpongFlag * ROW_SUM_PINGPONG_OFFSET + rowOffset * FLOAT_BLOCK_SIZE];
	609	+ AscendC::Brcb(
	610	+ sharedGmUbTensor, gmUbTensor[rowOffset], rowNumCurLoopRound / FLOAT_BLOCK_SIZE,
	611	+ AscendC::BrcbRepeatParams(1, 8));
	612	+ // AscendC::PipeBarrier<PIPE_V>();
	613	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
	614	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
	615	+ AscendC::DataCopy(gSharedMax, sharedGmUbTensor, AscendC::DataCopyParams(rowNumCurLoop, 1, 0, headNum - 1));
	616	+ }
	617	+
	618	+ CalcExp(sUbOffset, rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
	619	+ if constexpr (maskCat == MaskCategory::NO_MASK) {
	620	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(pingpongFlag);
	621	+ }
	622	+
	623	+ DownCastP(sUbOffset, rowNumCurLoop, columnNumRound);
	624	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
	625	+
	626	+ CalcLocalRowSum(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
	627	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(pingpongFlag);
	628	+
	629	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
	630	+ CopyPUbToGm(gOutput, sUbOffset, rowNumCurLoop, columnNumRound, columnNumPad);
	631	+ if constexpr (maskCat == MaskCategory::NO_MASK) {
	632	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(pingpongFlag);
	633	+ }
	634	+ UpdateGlobalRowSum(
	635	+ sUbOffset, rowNumCurLoop, rowNumCurLoopRound, dmUbOffsetCurCycle, rowOffset, isFirstStackTile);
	636	+ if (isLastStackTile) {
	637	+ // AscendC::LocalTensor<float> copyTempTensor = maskUbTensor32.template
	638	+ // ReinterpretCast<float>()[pingpongFlag * 128];
	639	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(pingpongFlag);
	640	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(pingpongFlag);
	641	+ AscendC::LocalTensor<float> sharedGlUbTensor = maskUbTensor.template ReinterpretCast<
	642	+ float>()[pingpongFlag * ROW_SUM_PINGPONG_OFFSET + 2048 + rowOffset * FLOAT_BLOCK_SIZE];
	643	+ AscendC::Brcb(
	644	+ sharedGlUbTensor, glUbTensor[rowOffset], rowNumCurLoopRound / FLOAT_BLOCK_SIZE,
	645	+ AscendC::BrcbRepeatParams(1, 8));
	646	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
	647	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
	648	+ AscendC::DataCopy(gSharedSum, sharedGlUbTensor, AscendC::DataCopyParams(rowNumCurLoop, 1, 0, headNum - 1));
	649	+ }
	650	+ }
	651	+
	652	+ CATLASS_DEVICE
	653	+ void operator()(
	654	+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
	655	+ AscendC::GlobalTensor<ElementUpdate> gSharedMax, AscendC::GlobalTensor<ElementUpdate> gSharedSum,
	656	+ const LayoutOutput& layoutOutput, const LayoutInput& layoutInput, GemmCoord actualBlockShape,
	657	+ uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t qSBlockSize, uint32_t qNBlockSize,
	658	+ uint32_t curStackTileMod, uint32_t headNum)
	659	+ {
	660	+ uint32_t rowNum = actualBlockShape.m();
	661	+ uint32_t columnNum = actualBlockShape.n();
	662	+ uint32_t columnNumRound = RoundUp(columnNum, BLOCK_SIZE);
	663	+ uint32_t columnNumPad = layoutInput.stride(0);
	664	+
	665	+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
	666	+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
	667	+
	668	+ uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum;
	669	+ uint32_t rowSplitSubBlock = (qNBlockSize == 1) ? (qSBlockSize / 2) : (qSBlockSize * qNSplitSubBlock);
	670	+ uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? (rowNum - rowSplitSubBlock) : rowSplitSubBlock;
	671	+ uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock;
	672	+ uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound; // 8192 / 128 = 64
	673	+ uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE); // 64
	674	+ uint32_t rowLoopNum = CeilDiv(rowActualThisSubBlock, rowNumTile);
	675	+ uint32_t preLoad = 1;
	676	+
	677	+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) {
	678	+ if (rowLoopIdx < rowLoopNum) {
	679	+ uint32_t pingpongFlag = rowLoopIdx % 2;
	680	+ uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile;
	681	+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
	682	+ uint32_t rowNumCurLoop =
	683	+ (rowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
	684	+
	685	+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
	686	+ auto gInputCurLoop = gInput[offsetInput];
	687	+
	688	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(pingpongFlag);
	689	+ CopySGmToUb(
	690	+ gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad);
	691	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
	692	+ }
	693	+ if (rowLoopIdx >= preLoad) {
	694	+ uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad;
	695	+ uint32_t pingpongFlag = delayedRowLoopIdx % 2;
	696	+ uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile;
	697	+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
	698	+ uint32_t rowNumCurLoop =
	699	+ (delayedRowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
	700	+
	701	+ int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
	702	+ auto gOutputCurLoop = gOutput[offsetOutput];
	703	+ uint32_t rowOffsetSumMax = rowOffsetIoGm * headNum * FLOAT_BLOCK_SIZE;
	704	+ auto gSharedMaxCurLoop = gSharedMax[rowOffsetSumMax];
	705	+ auto gSharedSumCurLoop = gSharedSum[rowOffsetSumMax];
	706	+ auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum));
	707	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
	708	+ ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound);
	709	+ SubCoreCompute<MaskCategory::NO_MASK>(
	710	+ gOutputCurLoop, gSharedMaxCurLoop, gSharedSumCurLoop, layoutOutputCurLoop, rowOffsetCurLoop,
	711	+ isFirstStackTile, columnNumRound, pingpongFlag, curStackTileMod, isLastStackTile, headNum);
	712	+ }
	713	+ }
	714	+ }
	715	+
	716	+private:
	717	+ float scaleValue;
	718	+ AscendC::LocalTensor<float> lsUbTensor;
	719	+ AscendC::LocalTensor<ElementOutput> lpUbTensor;
	720	+ AscendC::LocalTensor<ElementMask> maskUbTensor;
	721	+ AscendC::LocalTensor<float> maskUbTensor32;
	722	+ AscendC::LocalTensor<float> lmUbTensor;
	723	+ AscendC::LocalTensor<float> hmUbTensor;
	724	+ AscendC::LocalTensor<float> gmUbTensor;
	725	+ AscendC::LocalTensor<float> dmUbTensor;
	726	+ AscendC::LocalTensor<float> llUbTensor;
	727	+ AscendC::LocalTensor<float> tvUbTensor;
	728	+ AscendC::LocalTensor<float> glUbTensor;
	729	+};
	730	+
	731	+// Ascend950 specialization - same implementation as AtlasA2 since Ascend C APIs are compatible
	732	+template <class OutputType_, class InputType_, class MaskType_>
	733	+class BlockEpilogue<EpilogueAscend950OnlineSoftmaxCopySumMax, OutputType_, InputType_, MaskType_>
	734	+ : public BlockEpilogue<EpilogueAtlasA2OnlineSoftmaxCopySumMax, OutputType_, InputType_, MaskType_> {
	735	+public:
	736	+ using Base = BlockEpilogue<EpilogueAtlasA2OnlineSoftmaxCopySumMax, OutputType_, InputType_, MaskType_>;
	737	+ using DispatchPolicy = EpilogueAscend950OnlineSoftmaxCopySumMax;
	738	+ using ArchTag = typename DispatchPolicy::ArchTag;
	739	+ using Base::Base;
	740	+};
	741	+} // namespace Catlass::Epilogue::Block
	742	+
	743	+#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_COPY_SUM_MAX_HPP

Ainclude/catlass/epilogue/block/block_epilogue_rescale_o_FD.hpp+513-0

@@ -0,0 +1,513 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_HPP_FD
	13	+#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_HPP_FD
	14	+
	15	+#include "catlass/arch/resource.hpp"
	16	+#include "catlass/epilogue/dispatch_policy.hpp"
	17	+#include "catlass/epilogue/tile/tile_copy.hpp"
	18	+#include "catlass/gemm_coord.hpp"
	19	+#include "catlass/matrix_coord.hpp"
	20	+
	21	+namespace Catlass::Epilogue::Block {
	22	+
	23	+template <class OutputType_, class InputType_, class UpdateType_, class LseType_, LseMode LSE_MODE_>
	24	+class BlockEpilogue<EpilogueAtlasA2RescaleO_FD<LSE_MODE_>, OutputType_, InputType_, UpdateType_, LseType_> {
	25	+public:
	26	+ // Type aliases
	27	+ using DispatchPolicy = EpilogueAtlasA2RescaleO_FD<LSE_MODE_>;
	28	+ using ArchTag = typename DispatchPolicy::ArchTag;
	29	+
	30	+ using ElementOutput = typename OutputType_::Element;
	31	+ using ElementInput = typename InputType_::Element;
	32	+ using ElementUpdate = typename UpdateType_::Element;
	33	+ using ElementLse = typename LseType_::Element;
	34	+
	35	+ using LayoutOutput = typename OutputType_::Layout;
	36	+ using LayoutInput = typename InputType_::Layout;
	37	+ using LayoutUpdate = typename UpdateType_::Layout;
	38	+ using LayoutLse = typename LseType_::Layout;
	39	+
	40	+ static constexpr LseMode LSE_MODE = DispatchPolicy::LSE_MODE;
	41	+
	42	+ static constexpr uint32_t HALF_ELENUM_PER_BLK = 16;
	43	+ static constexpr uint32_t BLOCK_SIZE = 16;
	44	+ static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128;
	45	+ static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64;
	46	+ static constexpr uint32_t HALF_ELENUM_PER_LINE = 256;
	47	+ static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128;
	48	+ static constexpr uint32_t MULTIPLIER = 2;
	49	+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
	50	+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
	51	+ static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024;
	52	+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384;
	53	+ static constexpr uint32_t HALF_DM_UB_SIZE = 64;
	54	+ static constexpr uint32_t HALF_LL_UB_SIZE = 256;
	55	+ static constexpr uint32_t VECTOR_SIZE = 128;
	56	+ static constexpr uint32_t NUM4 = 4;
	57	+ static constexpr uint32_t MAX_UB_O_ELEM_NUM = 8192;
	58	+ static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 256;
	59	+ static constexpr uint32_t SIZE_OF_16BIT = 2;
	60	+
	61	+ __aicore__ inline BlockEpilogue(Arch::Resource<ArchTag>& resource)
	62	+ {
	63	+ // Allocate UB space
	64	+ constexpr uint32_t LO_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE;
	65	+ constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE;
	66	+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE;
	67	+
	68	+ constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE;
	69	+ constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE;
	70	+ constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE;
	71	+ constexpr uint32_t LSE_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE;
	72	+ constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE;
	73	+
	74	+ loUbTensor = resource.ubBuf.template GetBufferByByte<float>(LO_UB_TENSOR_OFFSET);
	75	+ dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
	76	+ glUbTensor = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
	77	+ tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
	78	+ goUbTensor16 = resource.ubBuf.template GetBufferByByte<ElementOutput>(GO_UB_TENSOR_OFFSET);
	79	+ goUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(GO_UB_TENSOR_OFFSET);
	80	+ hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);
	81	+ gmUbTensor = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
	82	+ lse32_ubuf_tensor = resource.ubBuf.template GetBufferByByte<float>(LSE_UB_TENSOR_OFFSET);
	83	+ }
	84	+
	85	+ __aicore__ inline ~BlockEpilogue()
	86	+ {}
	87	+
	88	+ __aicore__ inline void SetMask(int32_t len)
	89	+ {
	90	+ uint64_t mask = 0;
	91	+ uint64_t one = 1;
	92	+ uint64_t temp = static_cast<uint64_t>(len) % static_cast<uint64_t>(FLOAT_VECTOR_SIZE);
	93	+ for (uint64_t i = 0; i < temp; i++) {
	94	+ mask \|= one << i;
	95	+ }
	96	+
	97	+ if (len == VECTOR_SIZE) {
	98	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	99	+ } else if (len >= FLOAT_VECTOR_SIZE) {
	100	+ AscendC::SetVectorMask<int8_t>(mask, (uint64_t)-1);
	101	+ } else {
	102	+ AscendC::SetVectorMask<int8_t>(0x0, mask);
	103	+ }
	104	+ }
	105	+
	106	+ __aicore__ inline void CopyOToGm(
	107	+ AscendC::GlobalTensor<ElementOutput> gOutput, uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t epiTokenNum,
	108	+ uint32_t integralHeadNum, uint32_t qSThisSubBlock, uint32_t embed, uint32_t oHiddenSize)
	109	+ {
	110	+ uint32_t innerOGmOffset = 0;
	111	+ uint32_t innerGOUbOffset = 0;
	112	+ if (proTokenNum != 0U) {
	113	+ AscendC::DataCopyPad(
	114	+ gOutput[innerOGmOffset + proTokenIdx * oHiddenSize], goUbTensor16[innerGOUbOffset],
	115	+ AscendC::DataCopyExtParams(
	116	+ proTokenNum, embed * SIZE_OF_16BIT, 0, (oHiddenSize - embed) * SIZE_OF_16BIT, 0));
	117	+ innerOGmOffset += embed;
	118	+ innerGOUbOffset += proTokenNum * embed;
	119	+ }
	120	+ for (uint32_t qN_idx = 0; qN_idx < integralHeadNum; qN_idx++) {
	121	+ AscendC::DataCopyPad(
	122	+ gOutput[innerOGmOffset], goUbTensor16[innerGOUbOffset],
	123	+ AscendC::DataCopyExtParams(
	124	+ qSThisSubBlock, embed * SIZE_OF_16BIT, 0, (oHiddenSize - embed) * SIZE_OF_16BIT, 0));
	125	+ innerOGmOffset += embed;
	126	+ innerGOUbOffset += qSThisSubBlock * embed;
	127	+ }
	128	+ if (epiTokenNum != 0U) {
	129	+ AscendC::DataCopyPad(
	130	+ gOutput[innerOGmOffset], goUbTensor16[innerGOUbOffset],
	131	+ AscendC::DataCopyExtParams(
	132	+ epiTokenNum, embed * SIZE_OF_16BIT, 0, (oHiddenSize - embed) * SIZE_OF_16BIT, 0));
	133	+ }
	134	+ }
	135	+
	136	+ __aicore__ inline void CopyOToGmFp32(
	137	+ AscendC::GlobalTensor<float> gOutput, uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t epiTokenNum,
	138	+ uint32_t integralHeadNum, uint32_t qSThisSubBlock, uint32_t embed, uint32_t oHiddenSize,
	139	+ uint32_t oHiddenSize_gmlo)
	140	+ {
	141	+ uint32_t innerOGmOffset = 0;
	142	+ uint32_t innerGOUbOffset = 0;
	143	+
	144	+ // 前缀 token
	145	+ if (proTokenNum != 0U) {
	146	+ AscendC::DataCopyPad(
	147	+ gOutput[innerOGmOffset + proTokenIdx * oHiddenSize], goUbTensor32[innerGOUbOffset],
	148	+ AscendC::DataCopyExtParams(
	149	+ proTokenNum,
	150	+ embed * sizeof(float), // 每行有效 embed 列，float
	151	+ 0,
	152	+ (oHiddenSize_gmlo - embed) * sizeof(float), // 每行 padding
	153	+ 0));
	154	+ innerOGmOffset += embed;
	155	+ innerGOUbOffset += proTokenNum * embed;
	156	+ }
	157	+
	158	+ // 中间整块部分
	159	+ for (uint32_t qN_idx = 0; qN_idx < integralHeadNum; qN_idx++) {
	160	+ AscendC::DataCopyPad(
	161	+ gOutput[innerOGmOffset], goUbTensor32[innerGOUbOffset],
	162	+ AscendC::DataCopyExtParams(
	163	+ qSThisSubBlock, embed * sizeof(float), 0, (oHiddenSize_gmlo - embed) * sizeof(float), 0));
	164	+ innerOGmOffset += embed;
	165	+ innerGOUbOffset += qSThisSubBlock * embed;
	166	+ }
	167	+
	168	+ // 尾部 token
	169	+ if (epiTokenNum != 0U) {
	170	+ AscendC::DataCopyPad(
	171	+ gOutput[innerOGmOffset], goUbTensor32[innerGOUbOffset],
	172	+ AscendC::DataCopyExtParams(
	173	+ epiTokenNum, embed * sizeof(float), 0, (oHiddenSize_gmlo - embed) * sizeof(float), 0));
	174	+ }
	175	+ }
	176	+
	177	+ __aicore__ inline void SubCoreCompute(
	178	+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
	179	+ AscendC::GlobalTensor<ElementUpdate> gUpdate, AscendC::GlobalTensor<ElementLse> gLse,
	180	+ AscendC::GlobalTensor<ElementLse> gCombineLse, AscendC::GlobalTensor<ElementLse> gCombineo,
	181	+ const LayoutOutput& layoutOutput, const LayoutInput& layoutInput, const LayoutUpdate& layoutUpdate,
	182	+ const LayoutLse& layoutLse, uint32_t qNThisSubBlock, uint32_t qSThisSubBlock, uint32_t totalRowNum,
	183	+ uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod, uint32_t needRowLoop,
	184	+ uint32_t isLastRowLoop, uint32_t rowOffsetLoop, uint32_t proTokenIdx, uint32_t proTokenNum,
	185	+ uint32_t epiTokenNum, uint32_t integralHeadNum, bool isSplitkv, const LayoutOutput& layoutgmLse,
	186	+ const LayoutInput& layoutgmLo)
	187	+ {
	188	+ uint32_t curRowNum = layoutInput.shape(0);
	189	+ uint32_t embed = layoutInput.shape(1);
	190	+ uint32_t embedRound = layoutInput.stride(0);
	191	+ uint32_t curRowNumRound = RoundUp(curRowNum, FLOAT_BLOCK_SIZE);
	192	+ uint32_t qSBlockSize = layoutOutput.shape(0);
	193	+ uint32_t oHiddenSize = layoutOutput.shape(1);
	194	+ uint32_t qHeads = layoutLse.shape(1);
	195	+ uint32_t dmUbOffsetCurStackTile = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffsetLoop;
	196	+
	197	+ uint32_t oHiddenSize_gmlo = layoutgmLo.shape(1);
	198	+ uint32_t qHeads_gmlse = layoutgmLse.shape(1);
	199	+
	200	+ if (!isFirstStackTile) {
	201	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
	202	+ AscendC::DataCopy(
	203	+ loUbTensor, gInput, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
	204	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	205	+ }
	206	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID6);
	207	+ if (!isFirstStackTile) {
	208	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	209	+ AscendC::Brcb(
	210	+ tvUbTensor.ReinterpretCast<uint32_t>(), dmUbTensor[dmUbOffsetCurStackTile].ReinterpretCast<uint32_t>(),
	211	+ curRowNumRound / FLOAT_BLOCK_SIZE, AscendC::BrcbRepeatParams(1, 8));
	212	+ AscendC::PipeBarrier<PIPE_V>();
	213	+ if (needRowLoop) {
	214	+ AscendC::DataCopy(
	215	+ goUbTensor32, gUpdate, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
	216	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1);
	217	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1);
	218	+ }
	219	+ // *** go = go * dm_block
	220	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	221	+ for (uint32_t vmul_idx = 0; vmul_idx < embed / FLOAT_VECTOR_SIZE; ++vmul_idx) {
	222	+ AscendC::Mul<float, false>(
	223	+ goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE], goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE], tvUbTensor,
	224	+ (uint64_t)0, curRowNum,
	225	+ AscendC::BinaryRepeatParams(
	226	+ 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1));
	227	+ }
	228	+ if (embed % FLOAT_VECTOR_SIZE > 0) {
	229	+ SetMask(embed % FLOAT_VECTOR_SIZE);
	230	+ AscendC::Mul<float, false>(
	231	+ goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
	232	+ goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0, curRowNum,
	233	+ AscendC::BinaryRepeatParams(
	234	+ 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1));
	235	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	236	+ }
	237	+ AscendC::PipeBarrier<PIPE_V>();
	238	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	239	+ // *** go = lo + go
	240	+ AscendC::Add<float, false>(
	241	+ goUbTensor32, goUbTensor32, loUbTensor, (uint64_t)0,
	242	+ (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	243	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	244	+ AscendC::PipeBarrier<PIPE_V>();
	245	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
	246	+ } else {
	247	+ // *** go = lo
	248	+ AscendC::DataCopy(
	249	+ goUbTensor32, gInput, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
	250	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	251	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	252	+ }
	253	+
	254	+ if (isLastStackTile) {
	255	+ // *** gl_block = expand_to_block(gl), 存放于 tv
	256	+ AscendC::Brcb(
	257	+ tvUbTensor.ReinterpretCast<uint32_t>(), glUbTensor.ReinterpretCast<uint32_t>()[rowOffsetLoop],
	258	+ curRowNumRound / FLOAT_BLOCK_SIZE, AscendC::BrcbRepeatParams(1, 8));
	259	+ AscendC::PipeBarrier<PIPE_V>();
	260	+ // *** go = go / gl_block
	261	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	262	+ for (uint32_t vdiv_idx = 0; vdiv_idx < embed / FLOAT_VECTOR_SIZE; ++vdiv_idx) {
	263	+ AscendC::Div<float, false>(
	264	+ goUbTensor32[vdiv_idx * FLOAT_VECTOR_SIZE], goUbTensor32[vdiv_idx * FLOAT_VECTOR_SIZE], tvUbTensor,
	265	+ (uint64_t)0, curRowNum,
	266	+ AscendC::BinaryRepeatParams(
	267	+ 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1));
	268	+ }
	269	+ if (embed % FLOAT_VECTOR_SIZE > 0) {
	270	+ SetMask(embed % FLOAT_VECTOR_SIZE);
	271	+ AscendC::Div<float, false>(
	272	+ goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
	273	+ goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0, curRowNum,
	274	+ AscendC::BinaryRepeatParams(
	275	+ 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1));
	276	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	277	+ }
	278	+ AscendC::PipeBarrier<PIPE_V>();
	279	+ // *** go = castfp32to16(go)
	280	+ if (!isSplitkv) {
	281	+ if (std::is_same<ElementOutput, bfloat16_t>::value) {
	282	+ AscendC::Cast<ElementOutput, float, false>(
	283	+ goUbTensor16, goUbTensor32, AscendC::RoundMode::CAST_RINT, (uint64_t)0,
	284	+ (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	285	+ AscendC::UnaryRepeatParams(1, 1, 4, 8));
	286	+ } else {
	287	+ AscendC::Cast<ElementOutput, float, false>(
	288	+ goUbTensor16, goUbTensor32, AscendC::RoundMode::CAST_NONE, (uint64_t)0,
	289	+ (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	290	+ AscendC::UnaryRepeatParams(1, 1, 4, 8));
	291	+ }
	292	+ }
	293	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
	294	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
	295	+ //***move O to GM
	296	+ if (isSplitkv) {
	297	+ CopyOToGmFp32(
	298	+ gCombineo, // 注意：这里的 gCombineo 需要是 GlobalTensor<float>
	299	+ proTokenIdx, proTokenNum, epiTokenNum, integralHeadNum, qSThisSubBlock, embed, oHiddenSize,
	300	+ oHiddenSize_gmlo);
	301	+ } else {
	302	+ CopyOToGm(
	303	+ gOutput, proTokenIdx, proTokenNum, epiTokenNum, integralHeadNum, qSThisSubBlock, embed,
	304	+ oHiddenSize);
	305	+ }
	306	+
	307	+ if constexpr (LSE_MODE_ == LseMode::LSE_OUT) {
	308	+ if (isLastRowLoop) {
	309	+ AscendC::PipeBarrier<PIPE_V>();
	310	+ AscendC::Ln<float, false>(
	311	+ lse32_ubuf_tensor, glUbTensor, (uint64_t)0, CeilDiv(totalRowNum, FLOAT_VECTOR_SIZE),
	312	+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
	313	+
	314	+ AscendC::PipeBarrier<PIPE_V>();
	315	+ AscendC::Add<float, false>(
	316	+ lse32_ubuf_tensor, lse32_ubuf_tensor, gmUbTensor, (uint64_t)0,
	317	+ CeilDiv(totalRowNum, FLOAT_VECTOR_SIZE), AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	318	+ AscendC::PipeBarrier<PIPE_V>();
	319	+
	320	+ // *** lse_block = expand_to_block(lse), 存放于 tv
	321	+ AscendC::Brcb(
	322	+ tvUbTensor.ReinterpretCast<uint32_t>(), lse32_ubuf_tensor.ReinterpretCast<uint32_t>(),
	323	+ CeilDiv(totalRowNum, FLOAT_BLOCK_SIZE), AscendC::BrcbRepeatParams(1, 8));
	324	+ AscendC::PipeBarrier<PIPE_V>();
	325	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID4);
	326	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID4);
	327	+
	328	+ if (qNThisSubBlock == 0U) {
	329	+ if (isSplitkv) {
	330	+ AscendC::DataCopyPad(
	331	+ gCombineLse, tvUbTensor,
	332	+ AscendC::DataCopyExtParams(
	333	+ totalRowNum, sizeof(float), 0, (qHeads_gmlse - 1) * sizeof(float), 0));
	334	+ }
	335	+ AscendC::DataCopyPad(
	336	+ gLse, tvUbTensor,
	337	+ AscendC::DataCopyExtParams(totalRowNum, sizeof(float), 0, (qHeads - 1) * sizeof(float), 0));
	338	+ } else {
	339	+ for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) {
	340	+ if (isSplitkv) {
	341	+ AscendC::DataCopyPad(
	342	+ gCombineLse[qNIdx], tvUbTensor[qNIdx * qSBlockSize * FLOAT_BLOCK_SIZE],
	343	+ AscendC::DataCopyExtParams(
	344	+ qSBlockSize, sizeof(float), 0, (qHeads_gmlse - 1) * sizeof(float), 0));
	345	+ }
	346	+ AscendC::DataCopyPad(
	347	+ gLse[qNIdx], tvUbTensor[qNIdx * qSBlockSize * FLOAT_BLOCK_SIZE],
	348	+ AscendC::DataCopyExtParams(
	349	+ qSBlockSize, sizeof(float), 0, (qHeads - 1) * sizeof(float), 0));
	350	+ }
	351	+ }
	352	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
	353	+ }
	354	+ } else {
	355	+ if (isSplitkv) {
	356	+ if (isLastRowLoop) {
	357	+ AscendC::PipeBarrier<PIPE_V>();
	358	+ AscendC::Ln<float, false>(
	359	+ lse32_ubuf_tensor, glUbTensor, (uint64_t)0, CeilDiv(totalRowNum, FLOAT_VECTOR_SIZE),
	360	+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
	361	+
	362	+ AscendC::PipeBarrier<PIPE_V>();
	363	+ AscendC::Add<float, false>(
	364	+ lse32_ubuf_tensor, lse32_ubuf_tensor, gmUbTensor, (uint64_t)0,
	365	+ CeilDiv(totalRowNum, FLOAT_VECTOR_SIZE), AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	366	+ AscendC::PipeBarrier<PIPE_V>();
	367	+
	368	+ // *** lse_block = expand_to_block(lse), 存放于 tv
	369	+ AscendC::Brcb(
	370	+ tvUbTensor.ReinterpretCast<uint32_t>(), lse32_ubuf_tensor.ReinterpretCast<uint32_t>(),
	371	+ CeilDiv(totalRowNum, FLOAT_BLOCK_SIZE), AscendC::BrcbRepeatParams(1, 8));
	372	+ AscendC::PipeBarrier<PIPE_V>();
	373	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID4);
	374	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID4);
	375	+
	376	+ if (qNThisSubBlock == 0U) {
	377	+ AscendC::DataCopyPad(
	378	+ gCombineLse, tvUbTensor,
	379	+ AscendC::DataCopyExtParams(
	380	+ totalRowNum, sizeof(float), 0, (qHeads_gmlse - 1) * sizeof(float), 0));
	381	+ } else {
	382	+ for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) {
	383	+ AscendC::DataCopyPad(
	384	+ gCombineLse[qNIdx], tvUbTensor[qNIdx * qSBlockSize * FLOAT_BLOCK_SIZE],
	385	+ AscendC::DataCopyExtParams(
	386	+ qSBlockSize, sizeof(float), 0, (qHeads_gmlse - 1) * sizeof(float), 0));
	387	+ }
	388	+ }
	389	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
	390	+ }
	391	+ } else if (isLastRowLoop) {
	392	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
	393	+ }
	394	+ }
	395	+ } else if (needRowLoop) {
	396	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID5);
	397	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID5);
	398	+ AscendC::DataCopy(
	399	+ gUpdate, goUbTensor32, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
	400	+ }
	401	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID6);
	402	+ }
	403	+
	404	+ __aicore__ inline void operator()(
	405	+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
	406	+ AscendC::GlobalTensor<ElementUpdate> gUpdate, AscendC::GlobalTensor<ElementLse> gLse,
	407	+ AscendC::GlobalTensor<ElementLse> gCombineLse, AscendC::GlobalTensor<ElementLse> gCombineo,
	408	+ const LayoutOutput& layoutOutput, const LayoutInput& layoutInput, const LayoutUpdate& layoutUpdate,
	409	+ const LayoutLse& layoutLse, GemmCoord actualBlockShape, uint32_t qSBlockSize, uint32_t qNBlockSize,
	410	+ uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod, bool isSplitkv,
	411	+ const LayoutLse& layoutgmLse, const LayoutInput& layoutgmLo)
	412	+ {
	413	+ uint32_t rowNum = actualBlockShape.m();
	414	+ uint32_t embed = actualBlockShape.n();
	415	+ uint32_t maxRowNumPerLoop = MAX_UB_O_ELEM_NUM / embed;
	416	+ uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE);
	417	+
	418	+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
	419	+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
	420	+
	421	+ uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum;
	422	+ uint32_t qNThisSubBlock = (qNBlockSize == 1U) ? 0 :
	423	+ (subBlockIdx == 1U) ? (qNBlockSize - qNSplitSubBlock) :
	424	+ qNSplitSubBlock;
	425	+ uint32_t inRowSplitSubBlock =
	426	+ (qNBlockSize == 1U) ? (qSBlockSize / subBlockNum) : (qSBlockSize * qNSplitSubBlock);
	427	+ uint32_t inRowActualThisSubBlock = (subBlockIdx == 1U) ? (rowNum - inRowSplitSubBlock) : inRowSplitSubBlock;
	428	+ uint32_t inRowOffsetThisSubBlock = subBlockIdx * inRowSplitSubBlock;
	429	+ uint32_t outRowOffsetThisSubBlock = (qNBlockSize == 1U) ? inRowOffsetThisSubBlock : 0;
	430	+ uint32_t outColOffsetThisSubBlock = (qNBlockSize == 1U) ? 0 : subBlockIdx * qNSplitSubBlock * embed;
	431	+ uint32_t qSThisSubBlock = (qNBlockSize == 1U) ? inRowActualThisSubBlock : qSBlockSize;
	432	+ int64_t outOffsetSubBlock =
	433	+ layoutOutput.GetOffset(MatrixCoord(outRowOffsetThisSubBlock, outColOffsetThisSubBlock));
	434	+
	435	+ int64_t gmlooutOffsetSubBlock =
	436	+ layoutgmLo.GetOffset(MatrixCoord(outRowOffsetThisSubBlock, outColOffsetThisSubBlock));
	437	+
	438	+ uint32_t outLseRowOffsetThisSubBlock = (qNBlockSize == 1U) ? inRowOffsetThisSubBlock : 0;
	439	+ uint32_t outLseColOffsetThisSubBlock = (qNBlockSize == 1U) ? 0 : subBlockIdx * qNSplitSubBlock;
	440	+ int64_t offsetLse = layoutLse.GetOffset(MatrixCoord(outLseRowOffsetThisSubBlock, outLseColOffsetThisSubBlock));
	441	+ auto gLseThisSubBlock = gLse[offsetLse];
	442	+
	443	+ auto layoutOutLseThisSubBlock = layoutLse;
	444	+
	445	+ int64_t gmLseoffsetLse =
	446	+ layoutgmLse.GetOffset(MatrixCoord(outLseRowOffsetThisSubBlock, outLseColOffsetThisSubBlock));
	447	+
	448	+ auto gCombineLseSubBlock = gCombineLse[gmLseoffsetLse];
	449	+
	450	+ if (inRowActualThisSubBlock > 0U) {
	451	+ uint32_t rowLoop = CeilDiv(inRowActualThisSubBlock, rowNumTile);
	452	+ uint32_t needRowLoop = (rowLoop > 1U) ? 1 : 0;
	453	+
	454	+ // The rows of each cycle consist of multiple heads with several tokens.
	455	+ // There are several integral heads, one prologue head, one epilogue head.
	456	+ uint32_t proTokenIdx = 0; // the token idx of the start token of the prologue part
	457	+ uint32_t proTokenNum = 0; // the token num of the prologue part
	458	+ uint32_t epiTokenNum = 0; // the token num of the epilogue part
	459	+ uint32_t integralHeadNum = 0; // the number of integral heads within a cycle
	460	+
	461	+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoop; rowLoopIdx++) {
	462	+ uint32_t rowOffsetLoop = rowLoopIdx * rowNumTile;
	463	+ uint32_t rowOffsetCurLoop = inRowOffsetThisSubBlock + rowOffsetLoop;
	464	+ uint32_t rowActualCurLoop =
	465	+ (rowLoopIdx == (rowLoop - 1U)) ? inRowActualThisSubBlock - rowLoopIdx * rowNumTile : rowNumTile;
	466	+
	467	+ int64_t offsetOutput =
	468	+ static_cast<int64_t>(rowLoopIdx * rowNumTile / qSThisSubBlock * embed) + outOffsetSubBlock;
	469	+
	470	+ int64_t gmloffset =
	471	+ static_cast<int64_t>(rowLoopIdx * rowNumTile / qSThisSubBlock * embed) + gmlooutOffsetSubBlock;
	472	+
	473	+ auto gCombineoSubBlock = gCombineo[gmloffset];
	474	+
	475	+ auto gOutputCurLoop = gOutput[offsetOutput];
	476	+ auto layoutOutputCurLoop = layoutOutput;
	477	+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetCurLoop, 0));
	478	+ auto gInputCurLoop = gInput[offsetInput];
	479	+ auto layoutInputCurLoop = layoutInput.GetTileLayout(MatrixCoord(rowActualCurLoop, embed));
	480	+
	481	+ int64_t offsetUpdate = layoutUpdate.GetOffset(MatrixCoord(rowOffsetCurLoop, 0));
	482	+ auto gUpdateCurLoop = gUpdate[offsetUpdate];
	483	+ auto layoutUpdateCurLoop = layoutUpdate.GetTileLayout(MatrixCoord(rowActualCurLoop, embed));
	484	+
	485	+ proTokenIdx = rowOffsetLoop % qSThisSubBlock;
	486	+ proTokenNum = AscendC::Std::min(rowActualCurLoop, (qSThisSubBlock - proTokenIdx)) % qSThisSubBlock;
	487	+ integralHeadNum = (rowActualCurLoop - proTokenNum) / qSThisSubBlock;
	488	+ epiTokenNum = rowActualCurLoop - proTokenNum - integralHeadNum * qSThisSubBlock;
	489	+
	490	+ SubCoreCompute(
	491	+ gOutputCurLoop, gInputCurLoop, gUpdateCurLoop, gLseThisSubBlock, gCombineLseSubBlock,
	492	+ gCombineoSubBlock, layoutOutputCurLoop, layoutInputCurLoop, layoutUpdateCurLoop,
	493	+ layoutOutLseThisSubBlock, qNThisSubBlock, qSThisSubBlock, inRowActualThisSubBlock, isFirstStackTile,
	494	+ isLastStackTile, curStackTileMod, needRowLoop, (rowLoopIdx == rowLoop - 1U), rowOffsetLoop,
	495	+ proTokenIdx, proTokenNum, epiTokenNum, integralHeadNum, isSplitkv, layoutgmLse, layoutgmLo);
	496	+ }
	497	+ }
	498	+ }
	499	+
	500	+private:
	501	+ AscendC::LocalTensor<float> loUbTensor;
	502	+ AscendC::LocalTensor<float> dmUbTensor;
	503	+ AscendC::LocalTensor<float> hmUbTensor;
	504	+ AscendC::LocalTensor<float> glUbTensor;
	505	+ AscendC::LocalTensor<float> tvUbTensor;
	506	+ AscendC::LocalTensor<ElementOutput> goUbTensor16;
	507	+ AscendC::LocalTensor<float> goUbTensor32;
	508	+ AscendC::LocalTensor<float> gmUbTensor;
	509	+ AscendC::LocalTensor<float> lse32_ubuf_tensor;
	510	+};
	511	+} // namespace Catlass::Epilogue::Block
	512	+
	513	+#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_HPP_FD

Ainclude/catlass/epilogue/block/block_epilogue_rescale_o_no_div_rowsum.hpp+261-0

@@ -0,0 +1,261 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_NO_DIV_ROWSUM_HPP
	13	+#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_NO_DIV_ROWSUM_HPP
	14	+
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/arch/resource.hpp"
	17	+#include "catlass/epilogue/dispatch_policy.hpp"
	18	+#include "catlass/epilogue/tile/tile_copy.hpp"
	19	+#include "catlass/gemm_coord.hpp"
	20	+#include "catlass/matrix_coord.hpp"
	21	+
	22	+namespace Catlass::Epilogue::Block {
	23	+
	24	+template <class OutputType_, class InputType_, class UpdateType_>
	25	+class BlockEpilogue<
	26	+ EpilogueAtlasA2RescaleOWithoutDivSum,
	27	+ OutputType_, // fp16
	28	+ InputType_, // fp32
	29	+ UpdateType_> // fp32
	30	+{
	31	+public:
	32	+ // Type aliases
	33	+ using DispatchPolicy = EpilogueAtlasA2RescaleOWithoutDivSum;
	34	+ using ArchTag = typename DispatchPolicy::ArchTag;
	35	+
	36	+ using ElementOutput = typename OutputType_::Element;
	37	+ using ElementInput = typename InputType_::Element;
	38	+ using ElementUpdate = typename UpdateType_::Element;
	39	+
	40	+ using LayoutOutput = typename OutputType_::Layout;
	41	+ using LayoutInput = typename InputType_::Layout;
	42	+ using LayoutUpdate = typename UpdateType_::Layout;
	43	+
	44	+ static constexpr uint32_t HALF_ELENUM_PER_BLK = 16;
	45	+ static constexpr uint32_t BLOCK_SIZE = 16;
	46	+ static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128;
	47	+ static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64;
	48	+ static constexpr uint32_t HALF_ELENUM_PER_LINE = 256;
	49	+ static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128;
	50	+ static constexpr uint32_t MULTIPLIER = 2;
	51	+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
	52	+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
	53	+ static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024;
	54	+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384;
	55	+ static constexpr uint32_t HALF_DM_UB_SIZE = 64;
	56	+ static constexpr uint32_t HALF_LL_UB_SIZE = 256;
	57	+ static constexpr uint32_t VECTOR_SIZE = 128;
	58	+ static constexpr uint32_t NUM4 = 4;
	59	+ static constexpr uint32_t MAX_UB_O_ELEM_NUM = 4096;
	60	+ static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 128;
	61	+
	62	+ CATLASS_DEVICE
	63	+ BlockEpilogue(Arch::Resource<ArchTag>& resource)
	64	+ {
	65	+ // Allocate UB space
	66	+ constexpr uint32_t LO_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE;
	67	+ constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE;
	68	+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE;
	69	+
	70	+ constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE;
	71	+ constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE;
	72	+ constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE;
	73	+
	74	+ loUbTensor = resource.ubBuf.template GetBufferByByte<float>(LO_UB_TENSOR_OFFSET);
	75	+ dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
	76	+ glUbTensor = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
	77	+ tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
	78	+ goUbTensor16 = resource.ubBuf.template GetBufferByByte<ElementOutput>(GO_UB_TENSOR_OFFSET);
	79	+ goUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(GO_UB_TENSOR_OFFSET);
	80	+ hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);
	81	+ }
	82	+
	83	+ CATLASS_DEVICE
	84	+ ~BlockEpilogue()
	85	+ {}
	86	+
	87	+ CATLASS_DEVICE
	88	+ void SetMask(int32_t len)
	89	+ {
	90	+ uint64_t mask = 0;
	91	+ uint64_t one = 1;
	92	+ uint64_t temp = len % FLOAT_VECTOR_SIZE;
	93	+ for (int64_t i = 0; i < temp; i++) {
	94	+ mask \|= one << i;
	95	+ }
	96	+
	97	+ if (len == VECTOR_SIZE) {
	98	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	99	+ } else if (len >= FLOAT_VECTOR_SIZE) {
	100	+ AscendC::SetVectorMask<int8_t>(mask, (uint64_t)-1);
	101	+ } else {
	102	+ AscendC::SetVectorMask<int8_t>(0x0, mask);
	103	+ }
	104	+ }
	105	+
	106	+ CATLASS_DEVICE
	107	+ void CopyOToGm(
	108	+ AscendC::GlobalTensor<ElementOutput> gOutput, uint32_t curRowNum, uint32_t qSBlockSize, uint32_t embed,
	109	+ uint32_t embedRound, uint32_t qNThisSubBlock, uint32_t oHiddenSize)
	110	+ {
	111	+ if (qNThisSubBlock == 0) {
	112	+ AscendC::DataCopyPad(
	113	+ gOutput, goUbTensor16,
	114	+ AscendC::DataCopyExtParams(curRowNum, embed * 2, 0, (oHiddenSize - embed) * 2, 0));
	115	+ } else {
	116	+ for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) {
	117	+ AscendC::DataCopyPad(
	118	+ gOutput[qNIdx * embed], goUbTensor16[qNIdx * embedRound * qSBlockSize],
	119	+ AscendC::DataCopyExtParams(qSBlockSize, embed * 2, 0, (oHiddenSize - embed) * 2, 0));
	120	+ }
	121	+ }
	122	+ }
	123	+
	124	+ CATLASS_DEVICE
	125	+ void CopyFloatOToGm(
	126	+ AscendC::GlobalTensor<float> gSharedOut, uint32_t curRowNum, uint32_t qSBlockSize, uint32_t embed,
	127	+ uint32_t embedRound, uint32_t qNThisSubBlock, uint32_t oHiddenSize)
	128	+ {
	129	+ if (qNThisSubBlock == 0) {
	130	+ AscendC::DataCopyPad(
	131	+ gSharedOut, goUbTensor32,
	132	+ AscendC::DataCopyExtParams(curRowNum, embed * 4, 0, (oHiddenSize - embed) * 4, 0));
	133	+ } else {
	134	+ for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) {
	135	+ AscendC::DataCopyPad(
	136	+ gSharedOut[qNIdx * embed], goUbTensor32[qNIdx * embedRound * qSBlockSize],
	137	+ AscendC::DataCopyExtParams(qSBlockSize, embed * 4, 0, (oHiddenSize - embed) * 4, 0));
	138	+ }
	139	+ }
	140	+ }
	141	+
	142	+ CATLASS_DEVICE
	143	+ void SubCoreCompute(
	144	+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
	145	+ AscendC::GlobalTensor<float> gSharedOut, const LayoutOutput& layoutOutput, const LayoutInput& layoutInput,
	146	+ uint32_t qNThisSubBlock, uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod)
	147	+ {
	148	+ uint32_t curRowNum = layoutInput.shape(0);
	149	+ uint32_t embed = layoutInput.shape(1);
	150	+ uint32_t embedRound = layoutInput.stride(0);
	151	+ uint32_t curRowNumRound = RoundUp(curRowNum, FLOAT_BLOCK_SIZE);
	152	+ uint32_t qSBlockSize = layoutOutput.shape(0);
	153	+ uint32_t oHiddenSize = layoutOutput.shape(1);
	154	+ uint32_t dmUbOffsetCurStackTile = curStackTileMod * MAX_ROW_NUM_SUB_CORE;
	155	+
	156	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
	157	+ if (!isFirstStackTile) {
	158	+ AscendC::DataCopy(
	159	+ loUbTensor, gInput, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
	160	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	161	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	162	+
	163	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	164	+ AscendC::Brcb(
	165	+ tvUbTensor.ReinterpretCast<uint32_t>(), dmUbTensor[dmUbOffsetCurStackTile].ReinterpretCast<uint32_t>(),
	166	+ curRowNumRound / FLOAT_BLOCK_SIZE, AscendC::BrcbRepeatParams(1, 8));
	167	+ AscendC::PipeBarrier<PIPE_V>();
	168	+ // *** go = go * dm_block
	169	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	170	+ for (uint32_t vmul_idx = 0; vmul_idx < embed / FLOAT_VECTOR_SIZE; ++vmul_idx) {
	171	+ AscendC::Mul<float, false>(
	172	+ goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE], goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE], tvUbTensor,
	173	+ (uint64_t)0, curRowNum,
	174	+ AscendC::BinaryRepeatParams(
	175	+ 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1));
	176	+ }
	177	+ if (embed % FLOAT_VECTOR_SIZE > 0) {
	178	+ SetMask(embed % FLOAT_VECTOR_SIZE);
	179	+ AscendC::Mul<float, false>(
	180	+ goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
	181	+ goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0, curRowNum,
	182	+ AscendC::BinaryRepeatParams(
	183	+ 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1));
	184	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	185	+ }
	186	+ AscendC::PipeBarrier<PIPE_V>();
	187	+ // *** go = lo + go
	188	+ AscendC::Add<float, false>(
	189	+ goUbTensor32, goUbTensor32, loUbTensor, (uint64_t)0,
	190	+ (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	191	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	192	+ AscendC::PipeBarrier<PIPE_V>();
	193	+ } else {
	194	+ // *** go = lo
	195	+ AscendC::DataCopy(
	196	+ goUbTensor32, gInput, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
	197	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	198	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	199	+ }
	200	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
	201	+
	202	+ if (isLastStackTile) {
	203	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
	204	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
	205	+ CopyFloatOToGm(gSharedOut, curRowNum, qSBlockSize, embed, embedRound, qNThisSubBlock, oHiddenSize);
	206	+ }
	207	+ }
	208	+
	209	+ CATLASS_DEVICE
	210	+ void operator()(
	211	+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
	212	+ AscendC::GlobalTensor<float> gSharedOut, const LayoutOutput& layoutOutput, const LayoutInput& layoutInput,
	213	+ GemmCoord actualBlockShape, uint32_t qSBlockSize, uint32_t qNBlockSize, uint32_t isFirstStackTile,
	214	+ uint32_t isLastStackTile, uint32_t curStackTileMod)
	215	+ {
	216	+ uint32_t rowNum = actualBlockShape.m();
	217	+ uint32_t embed = actualBlockShape.n();
	218	+
	219	+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
	220	+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
	221	+
	222	+ uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum;
	223	+ uint32_t qNThisSubBlock = (qNBlockSize == 1) ? 0 :
	224	+ (subBlockIdx == 1) ? (qNBlockSize - qNSplitSubBlock) :
	225	+ qNSplitSubBlock;
	226	+ uint32_t inRowSplitSubBlock =
	227	+ (qNBlockSize == 1) ? (qSBlockSize / subBlockNum) : (qSBlockSize * qNSplitSubBlock);
	228	+ uint32_t inRowActualThisSubBlock = (subBlockIdx == 1) ? (rowNum - inRowSplitSubBlock) : inRowSplitSubBlock;
	229	+ uint32_t inRowOffsetThisSubBlock = subBlockIdx * inRowSplitSubBlock;
	230	+ uint32_t outRowOffsetThisSubBlock = (qNBlockSize == 1) ? inRowOffsetThisSubBlock : 0;
	231	+ uint32_t outColOffsetThisSubBlock = (qNBlockSize == 1) ? 0 : subBlockIdx * qNSplitSubBlock * embed;
	232	+
	233	+ if (inRowActualThisSubBlock > 0) {
	234	+ int64_t offsetOutput =
	235	+ layoutOutput.GetOffset(MatrixCoord(outRowOffsetThisSubBlock, outColOffsetThisSubBlock));
	236	+ auto gOutputThisSubBlock = gOutput[offsetOutput];
	237	+ auto gSharedOutThisSubBlock = gSharedOut[offsetOutput];
	238	+ auto layoutOutputThisSubBlock = layoutOutput;
	239	+
	240	+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(inRowOffsetThisSubBlock, 0));
	241	+ auto gInputThisSubBlock = gInput[offsetInput];
	242	+ auto layoutInputThisSubBlock = layoutInput.GetTileLayout(MatrixCoord(inRowActualThisSubBlock, embed));
	243	+
	244	+ SubCoreCompute(
	245	+ gOutputThisSubBlock, gInputThisSubBlock, gSharedOutThisSubBlock, layoutOutputThisSubBlock,
	246	+ layoutInputThisSubBlock, qNThisSubBlock, isFirstStackTile, isLastStackTile, curStackTileMod);
	247	+ }
	248	+ }
	249	+
	250	+private:
	251	+ AscendC::LocalTensor<float> loUbTensor;
	252	+ AscendC::LocalTensor<float> dmUbTensor;
	253	+ AscendC::LocalTensor<float> hmUbTensor;
	254	+ AscendC::LocalTensor<float> glUbTensor;
	255	+ AscendC::LocalTensor<float> tvUbTensor;
	256	+ AscendC::LocalTensor<ElementOutput> goUbTensor16;
	257	+ AscendC::LocalTensor<float> goUbTensor32;
	258	+};
	259	+} // namespace Catlass::Epilogue::Block
	260	+
	261	+#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_NO_DIV_ROWSUM_HPP

Ainclude/catlass/epilogue/block/block_epilogue_xa_combine_scale_ascend950.hpp+255-0

@@ -0,0 +1,255 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_XA_COMBINE_SCALE_ASCEND950_HPP
	13	+#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_XA_COMBINE_SCALE_ASCEND950_HPP
	14	+
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/arch/resource.hpp"
	17	+#include "catlass/epilogue/dispatch_policy.hpp"
	18	+#include "catlass/epilogue/tile/tile_copy.hpp"
	19	+#include "catlass/gemm_coord.hpp"
	20	+#include "catlass/matrix_coord.hpp"
	21	+
	22	+namespace Catlass::Epilogue::Block {
	23	+
	24	+template <class OutputType_, class InputType_>
	25	+class BlockEpilogue<EpilogueAscend950XACombineScale, OutputType_, InputType_> {
	26	+public:
	27	+ // Type aliases
	28	+ using DispatchPolicy = EpilogueAscend950XACombineScale;
	29	+ using ArchTag = typename DispatchPolicy::ArchTag;
	30	+
	31	+ using ElementOutput = typename OutputType_::Element;
	32	+ using ElementInput = typename InputType_::Element;
	33	+
	34	+ using LayoutOutput = typename OutputType_::Layout;
	35	+ using LayoutInput = typename InputType_::Layout;
	36	+
	37	+ CATLASS_DEVICE
	38	+ BlockEpilogue(
	39	+ Arch::Resource<ArchTag>& resource, uint32_t& ubBufAddrStart, uint32_t rowNumPerLoop_, uint32_t headDim_)
	40	+ {
	41	+ rowNumPerLoop = rowNumPerLoop_;
	42	+ headDim = headDim_;
	43	+ uint32_t reduceUbSize = rowNumPerLoop * sizeof(ElementInput);
	44	+ uint32_t inputUbSize = rowNumPerLoop * headDim * sizeof(ElementInput);
	45	+
	46	+ for (int i = 0; i < 2; i++) {
	47	+ sharedAttnUbTensorList[i] = resource.ubBuf.template GetBufferByByte<ElementInput>(ubBufAddrStart);
	48	+ ubBufAddrStart += inputUbSize;
	49	+ sharedGmUbTensorList[i] = resource.ubBuf.template GetBufferByByte<ElementInput>(ubBufAddrStart);
	50	+ ubBufAddrStart += reduceUbSize;
	51	+ sharedGlUbTensorList[i] = resource.ubBuf.template GetBufferByByte<ElementInput>(ubBufAddrStart);
	52	+ ubBufAddrStart += reduceUbSize;
	53	+ unsharedAttnUbTensorList[i] = resource.ubBuf.template GetBufferByByte<ElementInput>(ubBufAddrStart);
	54	+ ubBufAddrStart += inputUbSize;
	55	+ unsharedGmUbTensorList[i] = resource.ubBuf.template GetBufferByByte<ElementInput>(ubBufAddrStart);
	56	+ ubBufAddrStart += reduceUbSize;
	57	+ unsharedGlUbTensorList[i] = resource.ubBuf.template GetBufferByByte<ElementInput>(ubBufAddrStart);
	58	+ ubBufAddrStart += reduceUbSize;
	59	+ finalAttnUbTensorList[i] = resource.ubBuf.template GetBufferByByte<ElementInput>(ubBufAddrStart);
	60	+ ubBufAddrStart += inputUbSize;
	61	+ eventAttnVMTE2[i] = i;
	62	+ eventGmGlVMTE2[i] = i + 2;
	63	+ eventOutMTE3V[i] = i;
	64	+ }
	65	+
	66	+ finalGlUbTensor = resource.ubBuf.template GetBufferByByte<ElementInput>(ubBufAddrStart);
	67	+ ubBufAddrStart += reduceUbSize;
	68	+ expMaxSharedUbTensor = resource.ubBuf.template GetBufferByByte<ElementInput>(ubBufAddrStart);
	69	+ ubBufAddrStart += reduceUbSize;
	70	+ expMaxUnSharedUbTensor = resource.ubBuf.template GetBufferByByte<ElementInput>(ubBufAddrStart);
	71	+ ubBufAddrStart += reduceUbSize;
	72	+
	73	+ for (int i = 0; i < 2; i++) {
	74	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(eventAttnVMTE2[i]);
	75	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(eventGmGlVMTE2[i]);
	76	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventOutMTE3V[i]);
	77	+ }
	78	+ }
	79	+
	80	+ CATLASS_DEVICE
	81	+ ~BlockEpilogue()
	82	+ {
	83	+ for (int i = 0; i < 2; i++) {
	84	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(eventAttnVMTE2[i]);
	85	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(eventGmGlVMTE2[i]);
	86	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventOutMTE3V[i]);
	87	+ }
	88	+ }
	89	+
	90	+ CATLASS_DEVICE
	91	+ void operator()(
	92	+ AscendC::GlobalTensor<ElementInput> gSharedGm, AscendC::GlobalTensor<ElementInput> gUnsharedGm,
	93	+ AscendC::GlobalTensor<ElementInput> gSharedGl, AscendC::GlobalTensor<ElementInput> gUnsharedGl,
	94	+ AscendC::GlobalTensor<ElementInput> gSharedOut, AscendC::GlobalTensor<ElementInput> gUnsharedOut,
	95	+ AscendC::GlobalTensor<ElementOutput> gFinalOutput, uint32_t m, int8_t& taskId)
	96	+ {
	97	+ uint32_t attnCount = m * headDim;
	98	+ uint32_t gmglCount = m;
	99	+ // copyIn
	100	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(eventAttnVMTE2[taskId]);
	101	+ AscendC::DataCopy(sharedAttnUbTensorList[taskId], gSharedOut, attnCount);
	102	+ AscendC::DataCopy(unsharedAttnUbTensorList[taskId], gUnsharedOut, attnCount);
	103	+
	104	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(eventGmGlVMTE2[taskId]);
	105	+ AscendC::DataCopy(sharedGmUbTensorList[taskId], gSharedGm, gmglCount);
	106	+ AscendC::DataCopy(sharedGlUbTensorList[taskId], gSharedGl, gmglCount);
	107	+ AscendC::DataCopy(unsharedGmUbTensorList[taskId], gUnsharedGm, gmglCount);
	108	+ AscendC::DataCopy(unsharedGlUbTensorList[taskId], gUnsharedGl, gmglCount);
	109	+
	110	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(taskId);
	111	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(taskId);
	112	+
	113	+ __ubuf__ ElementInput* sharedAttnUbAddr = (__ubuf__ ElementInput*)sharedAttnUbTensorList[taskId].GetPhyAddr();
	114	+ __ubuf__ ElementInput* unsharedAttnUbAddr =
	115	+ (__ubuf__ ElementInput*)unsharedAttnUbTensorList[taskId].GetPhyAddr();
	116	+ __ubuf__ ElementInput* sharedGmUbAddr = (__ubuf__ ElementInput*)sharedGmUbTensorList[taskId].GetPhyAddr();
	117	+ __ubuf__ ElementInput* sharedGlUbAddr = (__ubuf__ ElementInput*)sharedGlUbTensorList[taskId].GetPhyAddr();
	118	+ __ubuf__ ElementInput* unsharedGmUbAddr = (__ubuf__ ElementInput*)unsharedGmUbTensorList[taskId].GetPhyAddr();
	119	+ __ubuf__ ElementInput* unsharedGlUbAddr = (__ubuf__ ElementInput*)unsharedGlUbTensorList[taskId].GetPhyAddr();
	120	+ __ubuf__ ElementInput* finalGlUbAddr = (__ubuf__ ElementInput*)finalGlUbTensor.GetPhyAddr();
	121	+ __ubuf__ ElementInput* expMaxSharedUbAddr = (__ubuf__ ElementInput*)expMaxSharedUbTensor.GetPhyAddr();
	122	+ __ubuf__ ElementInput* expMaxUnSharedUbAddr = (__ubuf__ ElementInput*)expMaxUnSharedUbTensor.GetPhyAddr();
	123	+
	124	+ __ubuf__ ElementInput* finalAttnFloatUbAddr =
	125	+ (__ubuf__ ElementInput*)finalAttnUbTensorList[taskId].GetPhyAddr();
	126	+
	127	+ ComputeExpSumAndExpMax(
	128	+ sharedGmUbAddr, sharedGlUbAddr, unsharedGmUbAddr, unsharedGlUbAddr, finalGlUbAddr, expMaxSharedUbAddr,
	129	+ expMaxUnSharedUbAddr, m);
	130	+
	131	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(eventGmGlVMTE2[taskId]);
	132	+
	133	+ AscendC::PipeBarrier<PIPE_V>();
	134	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventOutMTE3V[taskId]);
	135	+ constexpr int16_t vlSize = static_cast<int16_t>(AscendC::GetVecLen() / sizeof(ElementInput));
	136	+ int16_t nLoops = AscendC::CeilDivision(headDim, vlSize) - 1;
	137	+ uint32_t tailN = (headDim - 1) % vlSize + 1;
	138	+ ComputeFinalAttn(
	139	+ sharedAttnUbAddr, unsharedAttnUbAddr, expMaxSharedUbAddr, expMaxUnSharedUbAddr, finalGlUbAddr,
	140	+ finalAttnFloatUbAddr, static_cast<uint16_t>(m), headDim, nLoops, tailN);
	141	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(eventAttnVMTE2[taskId]);
	142	+
	143	+ AscendC::PipeBarrier<PIPE_V>();
	144	+ AscendC::LocalTensor<ElementOutput> finalAttnUb =
	145	+ finalAttnUbTensorList[taskId].template ReinterpretCast<ElementOutput>();
	146	+ AscendC::Cast(finalAttnUb, finalAttnUbTensorList[taskId], AscendC::RoundMode::CAST_ROUND, attnCount);
	147	+
	148	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(taskId);
	149	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(taskId);
	150	+ AscendC::DataCopy(gFinalOutput, finalAttnUb, attnCount);
	151	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventOutMTE3V[taskId]);
	152	+ taskId = 1 - taskId;
	153	+ }
	154	+
	155	+private:
	156	+ template <typename ElementInput>
	157	+ __simd_vf__ inline void ComputeExpSumAndExpMax(
	158	+ __ubuf__ ElementInput* sharedGm, __ubuf__ ElementInput* sharedGl, __ubuf__ ElementInput* unsharedGm,
	159	+ __ubuf__ ElementInput* unsharedGl, __ubuf__ ElementInput* finalGl, __ubuf__ ElementInput* expMaxShared,
	160	+ __ubuf__ ElementInput* expMaxUnShared, uint32_t tailM)
	161	+ {
	162	+ using namespace AscendC::Reg;
	163	+ RegTensor<ElementInput> sharedMaxVreg;
	164	+ RegTensor<ElementInput> sharedSumVreg;
	165	+ RegTensor<ElementInput> unsharedMaxVreg;
	166	+ RegTensor<ElementInput> unsharedSumVreg;
	167	+ RegTensor<ElementInput> finalGmVreg;
	168	+ RegTensor<ElementInput> finalGlVreg;
	169	+ RegTensor<ElementInput> expMaxSharedVreg;
	170	+ RegTensor<ElementInput> expMaxUnSharedVreg;
	171	+ MaskReg pregTailM = UpdateMask<ElementInput>(tailM);
	172	+
	173	+ LoadAlign(sharedMaxVreg, sharedGm);
	174	+ LoadAlign(unsharedMaxVreg, unsharedGm);
	175	+ Max(finalGmVreg, sharedMaxVreg, unsharedMaxVreg, pregTailM);
	176	+ ExpSub(expMaxSharedVreg, sharedMaxVreg, finalGmVreg, pregTailM);
	177	+ ExpSub(expMaxUnSharedVreg, unsharedMaxVreg, finalGmVreg, pregTailM);
	178	+ StoreAlign<ElementInput, StoreDist::DIST_NORM_B32>(expMaxShared, expMaxSharedVreg, pregTailM);
	179	+ StoreAlign<ElementInput, StoreDist::DIST_NORM_B32>(expMaxUnShared, expMaxUnSharedVreg, pregTailM);
	180	+
	181	+ LoadAlign(sharedSumVreg, sharedGl);
	182	+ LoadAlign(unsharedSumVreg, unsharedGl);
	183	+ Mul(finalGlVreg, sharedSumVreg, expMaxSharedVreg, pregTailM);
	184	+ MulAddDst(finalGlVreg, unsharedSumVreg, expMaxUnSharedVreg, pregTailM);
	185	+ StoreAlign<ElementInput, StoreDist::DIST_NORM_B32>(finalGl, finalGlVreg, pregTailM);
	186	+ }
	187	+
	188	+ template <class ElementInput>
	189	+ __simd_vf__ inline void ComputeFinalAttn(
	190	+ __ubuf__ ElementInput* sharedO, __ubuf__ ElementInput* unsharedO, __ubuf__ ElementInput* expMaxShared,
	191	+ __ubuf__ ElementInput* expMaxUnShared, __ubuf__ ElementInput* finalGl, __ubuf__ ElementInput* finalAttn,
	192	+ uint16_t m, uint32_t headDim, uint16_t nLoops, uint32_t tailN)
	193	+ {
	194	+ using namespace AscendC::Reg;
	195	+ RegTensor<ElementInput> sharedOVreg;
	196	+ RegTensor<ElementInput> unsharedOVreg;
	197	+ RegTensor<ElementInput> expMaxSharedVreg;
	198	+ RegTensor<ElementInput> expMaxUnsharedVreg;
	199	+ RegTensor<ElementInput> finalGlVreg;
	200	+
	201	+ RegTensor<ElementInput> finalAttnVreg;
	202	+
	203	+ MaskReg pregFull = CreateMask<ElementInput, MaskPattern::ALL>();
	204	+ MaskReg pregTailN = UpdateMask<ElementInput>(tailN);
	205	+
	206	+ for (uint16_t i = 0; i < m; ++i) {
	207	+ LoadAlign<ElementInput, LoadDist::DIST_BRC_B32>(expMaxSharedVreg, expMaxShared + i);
	208	+ LoadAlign<ElementInput, LoadDist::DIST_BRC_B32>(expMaxUnsharedVreg, expMaxUnShared + i);
	209	+ LoadAlign<ElementInput, LoadDist::DIST_BRC_B32>(finalGlVreg, finalGl + i);
	210	+ for (uint16_t j = 0; j < nLoops; ++j) {
	211	+ LoadAlign(sharedOVreg, sharedO + i * headDim + j * FLOAT_REP_SIZE);
	212	+ LoadAlign(unsharedOVreg, unsharedO + i * headDim + j * FLOAT_REP_SIZE);
	213	+ Mul(sharedOVreg, sharedOVreg, expMaxSharedVreg, pregFull);
	214	+ Mul(unsharedOVreg, unsharedOVreg, expMaxUnsharedVreg, pregFull);
	215	+ Add(finalAttnVreg, sharedOVreg, unsharedOVreg, pregFull);
	216	+ Div(finalAttnVreg, finalAttnVreg, finalGlVreg, pregFull);
	217	+ StoreAlign<ElementInput, StoreDist::DIST_NORM_B32>(
	218	+ finalAttn + i * headDim + j * FLOAT_REP_SIZE, finalAttnVreg, pregFull);
	219	+ }
	220	+ LoadAlign(sharedOVreg, sharedO + i * headDim + nLoops * FLOAT_REP_SIZE);
	221	+ LoadAlign(unsharedOVreg, unsharedO + i * headDim + nLoops * FLOAT_REP_SIZE);
	222	+ Mul(sharedOVreg, sharedOVreg, expMaxSharedVreg, pregTailN);
	223	+ Mul(unsharedOVreg, unsharedOVreg, expMaxUnsharedVreg, pregTailN);
	224	+ Add(finalAttnVreg, sharedOVreg, unsharedOVreg, pregTailN);
	225	+ Div(finalAttnVreg, finalAttnVreg, finalGlVreg, pregTailN);
	226	+ StoreAlign<ElementInput, StoreDist::DIST_NORM_B32>(
	227	+ finalAttn + i * headDim + nLoops * FLOAT_REP_SIZE, finalAttnVreg, pregTailN);
	228	+ }
	229	+ }
	230	+
	231	+private:
	232	+ static constexpr uint32_t FLOAT_REP_SIZE = 64;
	233	+ AscendC::LocalTensor<ElementInput> sharedAttnUbTensorList[2];
	234	+ AscendC::LocalTensor<ElementInput> sharedGmUbTensorList[2];
	235	+ AscendC::LocalTensor<ElementInput> sharedGlUbTensorList[2];
	236	+ AscendC::LocalTensor<ElementInput> unsharedAttnUbTensorList[2];
	237	+ AscendC::LocalTensor<ElementInput> unsharedGmUbTensorList[2];
	238	+ AscendC::LocalTensor<ElementInput> unsharedGlUbTensorList[2];
	239	+
	240	+ // tmp buffer
	241	+ AscendC::LocalTensor<ElementInput> finalGlUbTensor;
	242	+ AscendC::LocalTensor<ElementInput> expMaxSharedUbTensor;
	243	+ AscendC::LocalTensor<ElementInput> expMaxUnSharedUbTensor;
	244	+
	245	+ AscendC::LocalTensor<ElementInput> finalAttnUbTensorList[2];
	246	+ int32_t eventAttnVMTE2[2];
	247	+ int32_t eventGmGlVMTE2[2];
	248	+ int32_t eventOutMTE3V[2];
	249	+ uint32_t headDim;
	250	+ uint32_t rowNumPerLoop;
	251	+};
	252	+
	253	+} // namespace Catlass::Epilogue::Block
	254	+
	255	+#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_XA_COMBINE_SCALE_ASCEND950_HPP

Ainclude/catlass/epilogue/block/block_epilogue_xa_shared_rescale_ascend950.hpp+132-0

@@ -0,0 +1,132 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_XA_SHARED_RESCALE_O_ASCEND950
	13	+#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_XA_SHARED_RESCALE_O_ASCEND950
	14	+
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/arch/resource.hpp"
	17	+#include "catlass/epilogue/dispatch_policy.hpp"
	18	+#include "catlass/gemm_coord.hpp"
	19	+#include "catlass/matrix_coord.hpp"
	20	+#include "tla/tensor.hpp"
	21	+#include "tla/layout.hpp"
	22	+
	23	+namespace Catlass::Epilogue::Block {
	24	+
	25	+template <class L1TileShape_, class OTmpType_>
	26	+class BlockEpilogue<EpilogueAscend950XASharedRescaleO, L1TileShape_, OTmpType_> {
	27	+public:
	28	+ using DispatchPolicy = EpilogueAscend950XASharedRescaleO;
	29	+ using ArchTag = typename DispatchPolicy::ArchTag;
	30	+
	31	+ using ElementOTmp = typename OTmpType_::Element;
	32	+ using LayoutTagOTmp = typename OTmpType_::Layout;
	33	+ using L1TileShape = L1TileShape_;
	34	+
	35	+ static constexpr uint32_t S1_BASE_SIZE = tla::get<0>(L1TileShape{});
	36	+ static constexpr uint32_t D_BASE_SIZE = tla::get<2>(L1TileShape{});
	37	+ static constexpr uint32_t HALF_S1_BASE_SIZE = S1_BASE_SIZE / 2;
	38	+ static constexpr uint32_t VEC2_UB_SIZE = HALF_S1_BASE_SIZE * D_BASE_SIZE * sizeof(ElementOTmp);
	39	+
	40	+ CATLASS_DEVICE
	41	+ BlockEpilogue(Arch::Resource<ArchTag>& resource, uint32_t& ubBufAddrStart)
	42	+ {
	43	+ attnTmpBuf = resource.ubBuf.template GetBufferByByte<ElementOTmp>(ubBufAddrStart);
	44	+ ubBufAddrStart += VEC2_UB_SIZE;
	45	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventOMTE3V);
	46	+ }
	47	+
	48	+ CATLASS_DEVICE
	49	+ ~BlockEpilogue()
	50	+ {
	51	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventOMTE3V);
	52	+ }
	53	+
	54	+ template <class TensorDst, class TensorSrc>
	55	+ CATLASS_DEVICE void operator()(
	56	+ TensorDst& attenOutGm, const AscendC::LocalTensor<ElementOTmp>& expMaxUb, TensorSrc& pvRes, bool isFirstKv,
	57	+ bool isLastKv, uint16_t PV_RELEASE_FLAG)
	58	+ {
	59	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventOMTE3V);
	60	+ uint32_t m = tla::get<0>(pvRes.shape());
	61	+ uint32_t n = tla::get<1>(pvRes.shape());
	62	+ constexpr int16_t vlSize = static_cast<int16_t>(AscendC::GetVecLen() / sizeof(ElementOTmp));
	63	+ int16_t nLoops = AscendC::CeilDivision(n, vlSize) - 1;
	64	+ uint32_t tailN = (n - 1) % vlSize + 1;
	65	+
	66	+ __ubuf__ ElementOTmp* pvResUbAddr = (__ubuf__ ElementOTmp*)pvRes.data().GetPhyAddr();
	67	+ __ubuf__ ElementOTmp* attnTmpAddr = (__ubuf__ ElementOTmp*)attnTmpBuf.GetPhyAddr();
	68	+ __ubuf__ ElementOTmp* expMaxUbAddr = (__ubuf__ ElementOTmp*)expMaxUb.GetPhyAddr();
	69	+
	70	+ if (isFirstKv) {
	71	+ DataCopy(attnTmpBuf, pvRes.data(), m * n);
	72	+ } else {
	73	+ FlashUpdateNew<ElementOTmp, D_BASE_SIZE>(
	74	+ attnTmpAddr, pvResUbAddr, expMaxUbAddr, static_cast<uint16_t>(m), nLoops, tailN);
	75	+ }
	76	+
	77	+ if (isLastKv) {
	78	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(eventOVMTE3);
	79	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(eventOVMTE3);
	80	+ auto layoutUb = tla::MakeLayout(tla::MakeShape(m, n), tla::MakeStride(D_BASE_SIZE, tla::Int<1>{}));
	81	+ auto attenOutUb = tla::MakeTensor(attnTmpBuf, layoutUb, Arch::PositionUB{});
	82	+ using CopyUbToGmO = Tile::CopyUb2GmTla<ArchTag, decltype(attenOutUb), TensorDst>;
	83	+ CopyUbToGmO copyUbToGmO;
	84	+ copyUbToGmO(attenOutGm, attenOutUb);
	85	+ }
	86	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventOMTE3V);
	87	+ AscendC::CrossCoreSetFlag<SYNC_MODE, PIPE_V>(PV_RELEASE_FLAG);
	88	+ }
	89	+
	90	+private:
	91	+ AscendC::LocalTensor<ElementOTmp> attnTmpBuf;
	92	+ static constexpr int32_t SYNC_MODE = 4;
	93	+ static constexpr uint16_t FLOAT_REP_SIZE = 64;
	94	+ static constexpr int32_t eventOVMTE3 = 3;
	95	+ static constexpr int32_t eventOMTE3V = 3;
	96	+
	97	+ template <class T, uint16_t DBaseSize>
	98	+ __simd_vf__ inline void FlashUpdateNew(
	99	+ __ubuf__ T* updateUb, __ubuf__ T* curUb, __ubuf__ T* expMaxUb, uint16_t m, uint16_t nLoops, uint32_t tailN)
	100	+ {
	101	+ using namespace AscendC::Reg;
	102	+ RegTensor<float> expMaxVreg;
	103	+ RegTensor<float> preSrcVreg;
	104	+ RegTensor<float> curSrcVreg;
	105	+ RegTensor<float> mulVreg;
	106	+ RegTensor<float> addVreg;
	107	+
	108	+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
	109	+ MaskReg pregTailN = UpdateMask<float>(tailN);
	110	+
	111	+ for (uint16_t i = 0; i < m; ++i) {
	112	+ LoadAlign<T, LoadDist::DIST_BRC_B32>(expMaxVreg, expMaxUb + i);
	113	+ for (uint16_t j = 0; j < nLoops; ++j) {
	114	+ LoadAlign(preSrcVreg, updateUb + i * DBaseSize + j * FLOAT_REP_SIZE);
	115	+ LoadAlign(curSrcVreg, curUb + i * DBaseSize + j * FLOAT_REP_SIZE);
	116	+ Mul(mulVreg, expMaxVreg, preSrcVreg, pregFull);
	117	+ Add(addVreg, mulVreg, curSrcVreg, pregFull);
	118	+ StoreAlign<T, StoreDist::DIST_NORM_B32>(
	119	+ updateUb + i * DBaseSize + j * FLOAT_REP_SIZE, addVreg, pregFull);
	120	+ }
	121	+ LoadAlign(preSrcVreg, updateUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE);
	122	+ LoadAlign(curSrcVreg, curUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE);
	123	+ Mul(mulVreg, expMaxVreg, preSrcVreg, pregTailN);
	124	+ Add(addVreg, mulVreg, curSrcVreg, pregTailN);
	125	+ StoreAlign<T, StoreDist::DIST_NORM_B32>(
	126	+ updateUb + i * DBaseSize + nLoops * FLOAT_REP_SIZE, addVreg, pregTailN);
	127	+ }
	128	+ }
	129	+};
	130	+} // namespace Catlass::Epilogue::Block
	131	+
	132	+#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_ASCEND950

Ainclude/catlass/epilogue/block/block_epilogue_xa_shared_softmax_ascend950.hpp+359-0

@@ -0,0 +1,359 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_XA_SHARED_SOFTMAX_ASCEND950
	13	+#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_XA_SHARED_SOFTMAX_ASCEND950
	14	+
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/arch/cross_core_sync.hpp"
	17	+#include "catlass/arch/resource.hpp"
	18	+#include "catlass/epilogue/dispatch_policy.hpp"
	19	+#include "catlass/gemm_coord.hpp"
	20	+#include "catlass/matrix_coord.hpp"
	21	+#include "tla/tensor.hpp"
	22	+#include "tla/layout.hpp"
	23	+
	24	+namespace Catlass::Epilogue::Block {
	25	+
	26	+template <class L1TileShape_, class PType_, class SType_>
	27	+class BlockEpilogue<EpilogueAscend950XASharedSoftmax, L1TileShape_, PType_, SType_> {
	28	+public:
	29	+ using DispatchPolicy = EpilogueAscend950XASharedSoftmax;
	30	+ using ArchTag = typename DispatchPolicy::ArchTag;
	31	+ using L1TileShape = L1TileShape_;
	32	+ using ElementP = typename PType_::Element;
	33	+ using ElementS = typename SType_::Element;
	34	+ using LayoutTagP = typename PType_::Layout;
	35	+
	36	+ static constexpr uint32_t S1_BASE_SIZE = tla::get<0>(L1TileShape{});
	37	+ static constexpr uint32_t S2_BASE_SIZE = tla::get<1>(L1TileShape{});
	38	+ static constexpr int32_t HALF_S1_BASE_SIZE = S1_BASE_SIZE >> 1;
	39	+ static constexpr int32_t HALF_VEC_SIZE = HALF_S1_BASE_SIZE * sizeof(ElementS);
	40	+ static constexpr int32_t HALF_SCM_BLOCK_SIZE = HALF_S1_BASE_SIZE * S2_BASE_SIZE * sizeof(ElementP);
	41	+ static constexpr int32_t HALF_BRCB_SIZE = S1_BASE_SIZE * 32;
	42	+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementP);
	43	+
	44	+ CATLASS_DEVICE
	45	+ BlockEpilogue(Arch::Resource<ArchTag>& resource, uint32_t& ubBufAddrStart, ElementS scaleValue_, uint32_t qHeads_)
	46	+ {
	47	+ int32_t eventMTE3V = 0;
	48	+ int32_t eventVMTE3 = 0;
	49	+ scaleValue = scaleValue_;
	50	+ qHeads = qHeads_;
	51	+ for (int32_t i = 0; i < TASK_NUM2; i++) {
	52	+ pNzOutTensorList[i] = resource.ubBuf.template GetBufferByByte<ElementP>(ubBufAddrStart);
	53	+ ubBufAddrStart += HALF_SCM_BLOCK_SIZE;
	54	+ }
	55	+
	56	+ for (int32_t i = 0; i < TASK_NUM3; i++) {
	57	+ eventUbPMTE3VList[i] = eventMTE3V++;
	58	+ eventUbPVMTE3List[i] = eventVMTE3++;
	59	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventUbPMTE3VList[i]);
	60	+ }
	61	+
	62	+ maxBrcbTensor = resource.ubBuf.template GetBufferByByte<ElementS>(ubBufAddrStart);
	63	+ ubBufAddrStart += HALF_BRCB_SIZE;
	64	+ sumBrcbTensor = resource.ubBuf.template GetBufferByByte<ElementS>(ubBufAddrStart);
	65	+ ubBufAddrStart += HALF_BRCB_SIZE;
	66	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID7);
	67	+ }
	68	+
	69	+ CATLASS_DEVICE
	70	+ ~BlockEpilogue()
	71	+ {
	72	+ for (uint32_t i = 0; i < TASK_NUM3; i++) {
	73	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventUbPMTE3VList[i]);
	74	+ }
	75	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID7);
	76	+ }
	77	+
	78	+ template <class TensorDst, class TensorSrc>
	79	+ CATLASS_DEVICE void operator()(
	80	+ TensorDst& pOutL1, TensorSrc& qkRes, AscendC::LocalTensor<ElementS>& lastExpSum,
	81	+ AscendC::LocalTensor<ElementS>& nowExpSum, AscendC::LocalTensor<ElementS>& nowExpMax,
	82	+ AscendC::LocalTensor<ElementS>& lastMax, AscendC::LocalTensor<ElementS>& nowMax,
	83	+ AscendC::GlobalTensor<ElementS>& sharedMaxGm, AscendC::GlobalTensor<ElementS>& sharedSumGm, bool isUpdate,
	84	+ bool isLastKv, uint16_t SYNC_QK_READY_FLAG, uint16_t SYNC_SOFTMAX_READY_FLAG, uint16_t QK_UB_RELEASE_FLAG,
	85	+ int8_t taskIdMod2, int8_t taskIdMod3)
	86	+ {
	87	+ uint32_t m = tla::get<0>(qkRes.shape());
	88	+ uint32_t n = tla::get<1>(qkRes.shape());
	89	+ uint32_t blockStride = AscendC::CeilDivision(m, ELE_NUM_PER_C0) * ELE_NUM_PER_C0;
	90	+ auto pNzUbLayout = tla::MakeLayout(
	91	+ tla::MakeShape(
	92	+ tla::MakeShape(m, tla::Int<1>{}),
	93	+ tla::MakeShape(
	94	+ tla::Int<ELE_NUM_PER_C0>{}, AscendC::CeilDivision(S2_BASE_SIZE, tla::Int<ELE_NUM_PER_C0>{}))),
	95	+ tla::MakeStride(
	96	+ tla::MakeStride(tla::Int<ELE_NUM_PER_C0>{}, tla::Int<ELE_NUM_PER_C0>{} * m),
	97	+ tla::MakeStride(tla::Int<1>{}, blockStride * ELE_NUM_PER_C0)));
	98	+ auto pOutUb = tla::MakeTensor(pNzOutTensorList[taskIdMod2], pNzUbLayout, Arch::PositionUB{});
	99	+
	100	+ constexpr int16_t vlSize = static_cast<int16_t>(AscendC::GetVecLen() / sizeof(ElementS));
	101	+ uint32_t tailN = (n - 1) % vlSize + 1;
	102	+ uint32_t tailM = m;
	103	+
	104	+ __ubuf__ ElementP* outputAddr = (__ubuf__ ElementP*)pOutUb.data().GetPhyAddr();
	105	+ __ubuf__ ElementS* inputAddr = (__ubuf__ ElementS*)qkRes.data().GetPhyAddr();
	106	+ __ubuf__ ElementS* lastMaxUbAddr = (__ubuf__ ElementS*)lastMax.GetPhyAddr();
	107	+ __ubuf__ ElementS* nowMaxUbAddr = (__ubuf__ ElementS*)nowMax.GetPhyAddr();
	108	+ __ubuf__ ElementS* lastExpSumUbAddr = (__ubuf__ ElementS*)lastExpSum.GetPhyAddr();
	109	+ __ubuf__ ElementS* nowExpSumUbAddr = (__ubuf__ ElementS*)nowExpSum.GetPhyAddr();
	110	+ __ubuf__ ElementS* nowExpMaxUbAddr = (__ubuf__ ElementS*)nowExpMax.GetPhyAddr();
	111	+
	112	+ AscendC::CrossCoreWaitFlag<SYNC_MODE, PIPE_V>(SYNC_QK_READY_FLAG);
	113	+
	114	+ if (n == 128) {
	115	+ ComputeScaleandMax<ElementS, S2_BASE_SIZE, NRangeIndex::N128>(
	116	+ inputAddr, nowMaxUbAddr, static_cast<uint16_t>(m), tailN, scaleValue);
	117	+ } else if (n <= 64) {
	118	+ ComputeScaleandMax<ElementS, S2_BASE_SIZE, NRangeIndex::N0_64>(
	119	+ inputAddr, nowMaxUbAddr, static_cast<uint16_t>(m), tailN, scaleValue);
	120	+ } else {
	121	+ ComputeScaleandMax<ElementS, S2_BASE_SIZE, NRangeIndex::N65_127>(
	122	+ inputAddr, nowMaxUbAddr, static_cast<uint16_t>(m), tailN, scaleValue);
	123	+ }
	124	+
	125	+ if (likely(isUpdate)) {
	126	+ UpdateMax<ElementS>(nowMaxUbAddr, lastMaxUbAddr, tailM);
	127	+ }
	128	+
	129	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(taskIdMod3);
	130	+ if (unlikely(n > 64)) {
	131	+ ComputeExpSubSum<ElementP, ElementS, S2_BASE_SIZE, NRangeIndex::N128>(
	132	+ outputAddr, inputAddr, nowMaxUbAddr, nowExpSumUbAddr, static_cast<uint16_t>(m), blockStride);
	133	+ } else {
	134	+ ComputeExpSubSum<ElementP, ElementS, S2_BASE_SIZE, NRangeIndex::N0_64>(
	135	+ outputAddr, inputAddr, nowMaxUbAddr, nowExpSumUbAddr, static_cast<uint16_t>(m), blockStride);
	136	+ }
	137	+
	138	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(taskIdMod2);
	139	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(taskIdMod2);
	140	+ AscendC::CrossCoreSetFlag<SYNC_MODE, PIPE_V>(QK_UB_RELEASE_FLAG);
	141	+
	142	+ if (likely(m != 0)) {
	143	+ using CopyUbToL1P = Tile::CopyUb2L1Tla<ArchTag, decltype(pOutUb), TensorDst>;
	144	+ CopyUbToL1P copyUb2L1P;
	145	+ copyUb2L1P(pOutL1, pOutUb);
	146	+ }
	147	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(taskIdMod3);
	148	+ AscendC::CrossCoreSetFlag<SYNC_MODE, PIPE_MTE3>(SYNC_SOFTMAX_READY_FLAG);
	149	+
	150	+ if (likely(isUpdate)) {
	151	+ UpdateExpSumAndExpMax<ElementS>(
	152	+ lastExpSumUbAddr, nowExpMaxUbAddr, lastMaxUbAddr, nowExpSumUbAddr, nowMaxUbAddr, tailM);
	153	+ }
	154	+
	155	+ if (unlikely(isLastKv)) {
	156	+ CopyOutMaxAndSum(sharedMaxGm, sharedSumGm, nowMax, nowExpSum, m);
	157	+ }
	158	+ }
	159	+
	160	+private:
	161	+ static constexpr uint32_t TASK_NUM2 = 2;
	162	+ static constexpr uint32_t TASK_NUM3 = 3;
	163	+ static constexpr int32_t SYNC_MODE = 4;
	164	+ static constexpr int64_t BLOCK_BYTES = 32;
	165	+ static constexpr uint32_t REPEAT_STRIDE = 1;
	166	+ static constexpr float MIN_VALUE = -3e38f;
	167	+ static constexpr uint16_t FLOAT_REP_SIZE = 64;
	168	+ static constexpr uint16_t DOUBLE_FLOAT_REP_SIZE = 128;
	169	+ static constexpr uint32_t FLOATS_PER_BLOCK = 32 / sizeof(ElementS);
	170	+
	171	+ int32_t eventUbPMTE3VList[TASK_NUM3];
	172	+ int32_t eventUbPVMTE3List[TASK_NUM3];
	173	+ uint32_t qHeads;
	174	+
	175	+ ElementS scaleValue;
	176	+ AscendC::LocalTensor<ElementP> pNzOutTensorList[TASK_NUM2];
	177	+ AscendC::LocalTensor<ElementS> maxBrcbTensor;
	178	+ AscendC::LocalTensor<ElementS> sumBrcbTensor;
	179	+
	180	+ enum class NRangeIndex
	181	+ {
	182	+ N0_64 = 0,
	183	+ N65_127,
	184	+ N128,
	185	+ N128_INF
	186	+ };
	187	+
	188	+ template <typename ElementS, uint16_t S2BaseSize, NRangeIndex NRange>
	189	+ __simd_vf__ inline void ComputeScaleandMax(
	190	+ __ubuf__ ElementS* srcUb, __ubuf__ ElementS* newMaxUb, uint16_t m, uint32_t tailN, ElementS dScale)
	191	+ {
	192	+ using namespace AscendC::Reg;
	193	+ RegTensor<ElementS> minVreg;
	194	+ RegTensor<ElementS> srcVreg0;
	195	+ RegTensor<ElementS> srcVreg1;
	196	+ RegTensor<ElementS> maxVreg;
	197	+ RegTensor<ElementS> maxTmpVreg;
	198	+ UnalignReg maxUreg;
	199	+ MaskReg pregFull = CreateMask<ElementS, MaskPattern::ALL>();
	200	+ MaskReg pregTailN = UpdateMask<ElementS>(tailN);
	201	+
	202	+ Duplicate(minVreg, MIN_VALUE);
	203	+
	204	+ for (uint16_t i = 0; i < m; ++i) {
	205	+ if constexpr (NRange > NRangeIndex::N0_64) {
	206	+ LoadAlign(srcVreg0, srcUb + i * S2BaseSize);
	207	+ LoadAlign(srcVreg1, srcUb + i * S2BaseSize + FLOAT_REP_SIZE);
	208	+ Muls(srcVreg0, srcVreg0, dScale, pregFull);
	209	+ Muls(srcVreg1, srcVreg1, dScale, pregTailN);
	210	+ if constexpr (NRange < NRangeIndex::N128) {
	211	+ Select(srcVreg1, srcVreg1, minVreg, pregTailN);
	212	+ }
	213	+ StoreAlign<ElementS, StoreDist::DIST_NORM_B32>(srcUb + i * S2BaseSize, srcVreg0, pregFull);
	214	+ StoreAlign<ElementS, StoreDist::DIST_NORM_B32>(
	215	+ srcUb + i * S2BaseSize + FLOAT_REP_SIZE, srcVreg1, pregFull);
	216	+ Max(maxTmpVreg, srcVreg0, srcVreg1, pregFull);
	217	+ ReduceMax(maxVreg, maxTmpVreg, pregFull);
	218	+ } else {
	219	+ LoadAlign(srcVreg0, srcUb + i * S2BaseSize);
	220	+ Muls(srcVreg0, srcVreg0, dScale, pregTailN);
	221	+ Select(srcVreg0, srcVreg0, minVreg, pregTailN);
	222	+ StoreAlign<ElementS, StoreDist::DIST_NORM_B32>(srcUb + i * S2BaseSize, srcVreg0, pregFull);
	223	+ ReduceMax(maxVreg, srcVreg0, pregFull);
	224	+ }
	225	+ StoreUnAlign<ElementS, PostLiteral::POST_MODE_UPDATE>(newMaxUb, maxVreg, maxUreg, 1);
	226	+ }
	227	+ StoreUnAlignPost(newMaxUb, maxUreg, 0);
	228	+ }
	229	+
	230	+ template <typename ElementS>
	231	+ __simd_vf__ inline void UpdateMax(__ubuf__ ElementS* nowMaxUb, __ubuf__ ElementS* lastMaxUb, uint32_t tailM)
	232	+ {
	233	+ using namespace AscendC::Reg;
	234	+ RegTensor<ElementS> nowMaxVreg;
	235	+ RegTensor<ElementS> lastMaxVreg;
	236	+ RegTensor<ElementS> maxVreg;
	237	+
	238	+ MaskReg pregTailM = UpdateMask<ElementS>(tailM);
	239	+ LoadAlign(lastMaxVreg, lastMaxUb);
	240	+ LoadAlign(nowMaxVreg, nowMaxUb);
	241	+ Max(maxVreg, nowMaxVreg, lastMaxVreg, pregTailM);
	242	+ StoreAlign<ElementS, StoreDist::DIST_NORM_B32>(nowMaxUb, maxVreg, pregTailM);
	243	+ }
	244	+
	245	+ template <typename ElementP, typename ElementS, uint16_t S2BaseSize, NRangeIndex NRange>
	246	+ __simd_vf__ inline void ComputeExpSubSum(
	247	+ __ubuf__ ElementP* expUb, __ubuf__ ElementS* srcUb, __ubuf__ ElementS* nowMaxUb, __ubuf__ ElementS* expSumUb,
	248	+ uint16_t m, uint32_t blockStride)
	249	+ {
	250	+ using namespace AscendC::Reg;
	251	+ constexpr static CastTrait castTraitZero = {
	252	+ RegLayout::ZERO,
	253	+ SatMode::SAT,
	254	+ MaskMergeMode::ZEROING,
	255	+ AscendC::RoundMode::CAST_ROUND,
	256	+ };
	257	+
	258	+ constexpr static CastTrait castTraitOne = {
	259	+ RegLayout::ONE,
	260	+ SatMode::SAT,
	261	+ MaskMergeMode::ZEROING,
	262	+ AscendC::RoundMode::CAST_ROUND,
	263	+ };
	264	+ RegTensor<ElementS> srcVreg0;
	265	+ RegTensor<ElementS> srcVreg1;
	266	+ RegTensor<ElementS> expVreg;
	267	+ RegTensor<ElementS> expVreg0;
	268	+ RegTensor<ElementS> expVreg1;
	269	+ RegTensor<ElementS> expSumVreg;
	270	+ RegTensor<ElementS> maxVreg;
	271	+
	272	+ RegTensor<ElementP> expDstVreg0;
	273	+ RegTensor<ElementP> expDstVreg1;
	274	+ RegTensor<ElementP> expDstVreg;
	275	+
	276	+ UnalignReg expSumUreg;
	277	+
	278	+ MaskReg pregFull = CreateMask<ElementS, MaskPattern::ALL>();
	279	+ MaskReg pregFull16 = CreateMask<uint16_t, MaskPattern::ALL>();
	280	+ for (uint16_t i = 0; i < m; ++i) {
	281	+ LoadAlign<ElementS, LoadDist::DIST_BRC_B32>(maxVreg, nowMaxUb + i);
	282	+ if constexpr (NRange > NRangeIndex::N0_64) {
	283	+ LoadAlign<ElementS, LoadDist::DIST_DINTLV_B32>(srcVreg0, srcVreg1, srcUb + i * S2BaseSize);
	284	+ ExpSub(expVreg0, srcVreg0, maxVreg, pregFull);
	285	+ ExpSub(expVreg1, srcVreg1, maxVreg, pregFull);
	286	+ Add(expVreg, expVreg0, expVreg1, pregFull);
	287	+ Cast<ElementP, ElementS, castTraitZero>(expDstVreg0, expVreg0, pregFull);
	288	+ Cast<ElementP, ElementS, castTraitOne>(expDstVreg1, expVreg1, pregFull);
	289	+ Or((RegTensor<uint16_t>&)expDstVreg, (RegTensor<uint16_t>&)expDstVreg0,
	290	+ (RegTensor<uint16_t>&)expDstVreg1, pregFull16);
	291	+ StoreAlign<ElementP, DataCopyMode::DATA_BLOCK_COPY, PostLiteral::POST_MODE_UPDATE>(
	292	+ expUb, expDstVreg, blockStride, REPEAT_STRIDE, pregFull16);
	293	+ } else {
	294	+ LoadAlign(srcVreg0, srcUb + i * S2BaseSize);
	295	+ ExpSub(expVreg, srcVreg0, maxVreg, pregFull);
	296	+ Cast<ElementP, ElementS, castTraitZero>(expDstVreg, expVreg, pregFull);
	297	+ DeInterleave(expDstVreg0, expDstVreg1, expDstVreg, expDstVreg);
	298	+ StoreAlign<ElementP, DataCopyMode::DATA_BLOCK_COPY, PostLiteral::POST_MODE_UPDATE>(
	299	+ expUb, expDstVreg0, blockStride, REPEAT_STRIDE, pregFull16);
	300	+ }
	301	+ ReduceSum(expSumVreg, expVreg, pregFull);
	302	+ StoreUnAlign<ElementS, PostLiteral::POST_MODE_UPDATE>(expSumUb, expSumVreg, expSumUreg, 1);
	303	+ }
	304	+ StoreUnAlignPost(expSumUb, expSumUreg, 0);
	305	+ }
	306	+
	307	+ template <typename ElementS>
	308	+ __simd_vf__ inline void UpdateExpSumAndExpMax(
	309	+ __ubuf__ ElementS* sumUb, __ubuf__ ElementS* expMaxUb, __ubuf__ ElementS* maxUb, __ubuf__ ElementS* expSumUb,
	310	+ __ubuf__ ElementS* nowMaxUb, uint32_t tailM)
	311	+ {
	312	+ using namespace AscendC::Reg;
	313	+ RegTensor<ElementS> nowMaxVreg;
	314	+ RegTensor<ElementS> lastMaxVreg;
	315	+ RegTensor<ElementS> expMaxVreg;
	316	+ RegTensor<ElementS> lastExpSumVreg;
	317	+ RegTensor<ElementS> brcExpSumVreg;
	318	+ MaskReg pregTailM = UpdateMask<ElementS>(tailM);
	319	+ LoadAlign(lastMaxVreg, maxUb);
	320	+ LoadAlign(nowMaxVreg, nowMaxUb);
	321	+ ExpSub(expMaxVreg, lastMaxVreg, nowMaxVreg, pregTailM);
	322	+ StoreAlign<ElementS, StoreDist::DIST_NORM_B32>(expMaxUb, expMaxVreg, pregTailM);
	323	+ StoreAlign<ElementS, StoreDist::DIST_NORM_B32>(maxUb, nowMaxVreg, pregTailM);
	324	+
	325	+ LoadAlign(lastExpSumVreg, sumUb);
	326	+ LoadAlign(brcExpSumVreg, expSumUb);
	327	+ MulDstAdd(expMaxVreg, lastExpSumVreg, brcExpSumVreg, pregTailM);
	328	+ StoreAlign<ElementS, StoreDist::DIST_NORM_B32>(expSumUb, expMaxVreg, pregTailM);
	329	+ }
	330	+
	331	+ CATLASS_DEVICE
	332	+ void CopyOutMaxAndSum(
	333	+ AscendC::GlobalTensor<ElementS> maxGm, AscendC::GlobalTensor<ElementS> sumGm,
	334	+ AscendC::LocalTensor<ElementS> maxUb, AscendC::LocalTensor<ElementS> sumUb, uint16_t tailM)
	335	+ {
	336	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID7);
	337	+
	338	+ uint8_t repeatTimes = (tailM + 7) / 8;
	339	+ AscendC::Brcb(maxBrcbTensor, maxUb, repeatTimes, AscendC::BrcbRepeatParams(1, FLOATS_PER_BLOCK));
	340	+
	341	+ AscendC::Brcb(sumBrcbTensor, sumUb, repeatTimes, AscendC::BrcbRepeatParams(1, FLOATS_PER_BLOCK));
	342	+
	343	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID7);
	344	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID7);
	345	+
	346	+ AscendC::DataCopyExtParams copyOutParams(
	347	+ static_cast<uint16_t>(tailM), static_cast<uint32_t>(sizeof(ElementS)), 0,
	348	+ static_cast<uint32_t>((qHeads - 1) * sizeof(ElementS)), 0);
	349	+
	350	+ AscendC::DataCopyPad(maxGm, maxBrcbTensor, copyOutParams);
	351	+ AscendC::DataCopyPad(sumGm, sumBrcbTensor, copyOutParams);
	352	+
	353	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID7);
	354	+ }
	355	+};
	356	+
	357	+} // namespace Catlass::Epilogue::Block
	358	+
	359	+#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_XA_SHARED_SOFTMAX_ASCEND950

Ainclude/catlass/epilogue/block/block_epilogue_xa_unshared_softmax_ascend950.hpp+356-0

@@ -0,0 +1,356 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_XA_UNSHARED_SOFTMAX_ASCEND950
	13	+#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_XA_UNSHARED_SOFTMAX_ASCEND950
	14	+
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/arch/cross_core_sync.hpp"
	17	+#include "catlass/arch/resource.hpp"
	18	+#include "catlass/epilogue/dispatch_policy.hpp"
	19	+#include "catlass/gemm_coord.hpp"
	20	+#include "catlass/matrix_coord.hpp"
	21	+#include "tla/tensor.hpp"
	22	+#include "tla/layout.hpp"
	23	+
	24	+namespace Catlass::Epilogue::Block {
	25	+
	26	+template <class L1TileShape_, class PType_, class SType_>
	27	+class BlockEpilogue<EpilogueAscend950XAUnsharedSoftmax, L1TileShape_, PType_, SType_> {
	28	+public:
	29	+ using DispatchPolicy = EpilogueAscend950XAUnsharedSoftmax;
	30	+ using ArchTag = typename DispatchPolicy::ArchTag;
	31	+ using L1TileShape = L1TileShape_;
	32	+ using ElementP = typename PType_::Element;
	33	+ using ElementS = typename SType_::Element;
	34	+ using LayoutTagP = typename PType_::Layout;
	35	+
	36	+ static constexpr uint32_t S1_BASE_SIZE = tla::get<0>(L1TileShape{});
	37	+ static constexpr uint32_t S2_BASE_SIZE = tla::get<1>(L1TileShape{});
	38	+ static constexpr int32_t HALF_S1_BASE_SIZE = S1_BASE_SIZE >> 1;
	39	+ static constexpr int32_t HALF_VEC_SIZE = HALF_S1_BASE_SIZE * sizeof(ElementS);
	40	+ static constexpr int32_t HALF_SCM_BLOCK_SIZE = HALF_S1_BASE_SIZE * S2_BASE_SIZE * sizeof(ElementP);
	41	+ static constexpr int32_t HALF_MASK_BLOCK_SIZE = HALF_S1_BASE_SIZE * S2_BASE_SIZE * sizeof(uint8_t);
	42	+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementP);
	43	+
	44	+ CATLASS_DEVICE
	45	+ BlockEpilogue(
	46	+ Arch::Resource<ArchTag>& resource, uint32_t& ubBufAddrStart, ElementS scaleValue_, uint32_t unsharedKvLen_,
	47	+ uint32_t maxDecodeStep_, uint32_t groupCountPerLoop_, uint32_t groupSize_)
	48	+ {
	49	+ int32_t eventMTE3V = 0;
	50	+ scaleValue = scaleValue_;
	51	+ unsharedKvLen = unsharedKvLen_;
	52	+ maxDecodeStep = maxDecodeStep_;
	53	+ groupCountPerLoop = groupCountPerLoop_;
	54	+ groupSize = groupSize_;
	55	+
	56	+ for (int32_t i = 0; i < TASK_NUM2; i++) {
	57	+ pNzOutTensorList[i] = resource.ubBuf.template GetBufferByByte<ElementP>(ubBufAddrStart);
	58	+ ubBufAddrStart += HALF_SCM_BLOCK_SIZE;
	59	+ }
	60	+ maskUbTensor = resource.ubBuf.template GetBufferByByte<uint8_t>(ubBufAddrStart);
	61	+ ubBufAddrStart += HALF_MASK_BLOCK_SIZE;
	62	+ maxUbTensor = resource.ubBuf.template GetBufferByByte<ElementS>(ubBufAddrStart);
	63	+ ubBufAddrStart += HALF_VEC_SIZE;
	64	+ sumUbTensor = resource.ubBuf.template GetBufferByByte<ElementS>(ubBufAddrStart);
	65	+ ubBufAddrStart += HALF_VEC_SIZE;
	66	+
	67	+ InitUnsharedMask(groupCountPerLoop, groupSize, unsharedKvLen);
	68	+ AscendC::PipeBarrier<PIPE_V>();
	69	+
	70	+ for (int32_t i = 0; i < TASK_NUM3; i++) {
	71	+ eventUbPMTE3VList[i] = eventMTE3V++;
	72	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(eventUbPMTE3VList[i]);
	73	+ }
	74	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID7);
	75	+ }
	76	+
	77	+ CATLASS_DEVICE
	78	+ ~BlockEpilogue()
	79	+ {
	80	+ for (uint32_t i = 0; i < TASK_NUM3; i++) {
	81	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(eventUbPMTE3VList[i]);
	82	+ }
	83	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID7);
	84	+ }
	85	+
	86	+ template <class TensorDst, class TensorSrc>
	87	+ CATLASS_DEVICE void operator()(
	88	+ TensorDst& pOutL1, TensorSrc& qkRes, AscendC::GlobalTensor<ElementS>& unSharedMaxGm,
	89	+ AscendC::GlobalTensor<ElementS>& unSharedSumGm, uint16_t SYNC_QK_READY_FLAG, uint16_t SYNC_SOFTMAX_READY_FLAG,
	90	+ uint16_t QK_UB_RELEASE_FLAG, int8_t taskIdMod2, int8_t taskIdMod3)
	91	+ {
	92	+ uint32_t m = tla::get<0>(qkRes.shape());
	93	+ uint32_t n = tla::get<1>(qkRes.shape());
	94	+ uint32_t blockStride = AscendC::CeilDivision(m, ELE_NUM_PER_C0) * ELE_NUM_PER_C0;
	95	+ auto pNzUbLayout = tla::MakeLayout(
	96	+ tla::MakeShape(
	97	+ tla::MakeShape(m, tla::Int<1>{}),
	98	+ tla::MakeShape(
	99	+ tla::Int<ELE_NUM_PER_C0>{}, AscendC::CeilDivision(S2_BASE_SIZE, tla::Int<ELE_NUM_PER_C0>{}))),
	100	+ tla::MakeStride(
	101	+ tla::MakeStride(tla::Int<ELE_NUM_PER_C0>{}, tla::Int<ELE_NUM_PER_C0>{} * m),
	102	+ tla::MakeStride(tla::Int<1>{}, blockStride * ELE_NUM_PER_C0)));
	103	+ auto pOutUb = tla::MakeTensor(pNzOutTensorList[taskIdMod2], pNzUbLayout, Arch::PositionUB{});
	104	+
	105	+ constexpr int16_t vlSize = static_cast<int16_t>(AscendC::GetVecLen() / sizeof(ElementS));
	106	+ uint32_t tailN = (n - 1) % vlSize + 1;
	107	+
	108	+ __ubuf__ ElementP* outputAddr = (__ubuf__ ElementP*)pOutUb.data().GetPhyAddr();
	109	+ __ubuf__ ElementS* inputAddr = (__ubuf__ ElementS*)qkRes.data().GetPhyAddr();
	110	+ __ubuf__ ElementS* maxUbAddr = (__ubuf__ ElementS*)maxUbTensor.GetPhyAddr();
	111	+ __ubuf__ ElementS* sumUbAddr = (__ubuf__ ElementS*)sumUbTensor.GetPhyAddr();
	112	+ __ubuf__ uint8_t* maskUbUnRoll1 = (__ubuf__ uint8_t*)maskUbTensor.GetPhyAddr();
	113	+ __ubuf__ uint8_t* maskUbUnRoll2 = maskUbUnRoll1 + FLOAT_REP_SIZE;
	114	+
	115	+ AscendC::CrossCoreWaitFlag<SYNC_MODE, PIPE_V>(SYNC_QK_READY_FLAG);
	116	+
	117	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID7);
	118	+ if (n == 128) {
	119	+ ComputeMaskandScale<ElementS, S2_BASE_SIZE, NRangeIndex::N128>(
	120	+ inputAddr, maskUbUnRoll1, maskUbUnRoll2, maxUbAddr, static_cast<uint16_t>(m), tailN, scaleValue);
	121	+ } else if (n <= 64) {
	122	+ ComputeMaskandScale<ElementS, S2_BASE_SIZE, NRangeIndex::N0_64>(
	123	+ inputAddr, maskUbUnRoll1, maskUbUnRoll2, maxUbAddr, static_cast<uint16_t>(m), tailN, scaleValue);
	124	+ } else {
	125	+ ComputeMaskandScale<ElementS, S2_BASE_SIZE, NRangeIndex::N65_127>(
	126	+ inputAddr, maskUbUnRoll1, maskUbUnRoll2, maxUbAddr, static_cast<uint16_t>(m), tailN, scaleValue);
	127	+ }
	128	+
	129	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(taskIdMod3);
	130	+ if (unlikely(n > 64)) {
	131	+ ComputeExpSubSum<ElementP, ElementS, S2_BASE_SIZE, NRangeIndex::N128>(
	132	+ outputAddr, inputAddr, maxUbAddr, sumUbAddr, static_cast<uint16_t>(m), blockStride);
	133	+ } else {
	134	+ ComputeExpSubSum<ElementP, ElementS, S2_BASE_SIZE, NRangeIndex::N0_64>(
	135	+ outputAddr, inputAddr, maxUbAddr, sumUbAddr, static_cast<uint16_t>(m), blockStride);
	136	+ }
	137	+
	138	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(taskIdMod2);
	139	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(taskIdMod2);
	140	+ AscendC::CrossCoreSetFlag<SYNC_MODE, PIPE_V>(QK_UB_RELEASE_FLAG);
	141	+
	142	+ if (likely(m != 0)) {
	143	+ using CopyUbToL1P = Tile::CopyUb2L1Tla<ArchTag, decltype(pOutUb), TensorDst>;
	144	+ CopyUbToL1P copyUb2L1P;
	145	+ copyUb2L1P(pOutL1, pOutUb);
	146	+ }
	147	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(taskIdMod3);
	148	+ AscendC::CrossCoreSetFlag<SYNC_MODE, PIPE_MTE3>(SYNC_SOFTMAX_READY_FLAG);
	149	+
	150	+ if (likely(m != 0)) {
	151	+ CopyOutMaxAndSum(unSharedMaxGm, unSharedSumGm, m);
	152	+ }
	153	+
	154	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID7);
	155	+ }
	156	+
	157	+private:
	158	+ static constexpr uint32_t TASK_NUM2 = 2;
	159	+ static constexpr uint32_t TASK_NUM3 = 3;
	160	+ static constexpr int32_t SYNC_MODE = 4;
	161	+ static constexpr int64_t BLOCK_BYTES = 32;
	162	+ static constexpr uint32_t REPEAT_STRIDE = 1;
	163	+ static constexpr float MIN_VALUE = -3e38f;
	164	+ static constexpr uint16_t FLOAT_REP_SIZE = 64;
	165	+ static constexpr uint16_t DOUBLE_FLOAT_REP_SIZE = 128;
	166	+ static constexpr uint32_t UINT8_BLOCK_SIZE = 256;
	167	+
	168	+ int32_t eventUbPMTE3VList[TASK_NUM3];
	169	+
	170	+ uint32_t groupSize;
	171	+ uint32_t groupCountPerLoop;
	172	+ uint32_t unsharedKvLen;
	173	+ uint32_t maxDecodeStep;
	174	+
	175	+ ElementS scaleValue;
	176	+ AscendC::LocalTensor<ElementP> pNzOutTensorList[TASK_NUM2];
	177	+ AscendC::LocalTensor<uint8_t> maskUbTensor;
	178	+ AscendC::LocalTensor<ElementS> maxUbTensor;
	179	+ AscendC::LocalTensor<ElementS> sumUbTensor;
	180	+
	181	+ enum class NRangeIndex
	182	+ {
	183	+ N0_64 = 0,
	184	+ N65_127,
	185	+ N128,
	186	+ N128_INF
	187	+ };
	188	+
	189	+ CATLASS_DEVICE
	190	+ void InitUnsharedMask(uint32_t groupCountPerLoop, uint32_t groupSize, uint32_t unsharedKvSeqLen)
	191	+ {
	192	+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
	193	+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
	194	+ uint32_t curGroupCount = (groupCountPerLoop + subBlockNum - 1) / subBlockNum;
	195	+ uint32_t curGroupThisSubBlock = (subBlockIdx == 0) ? curGroupCount : (groupCountPerLoop - curGroupCount);
	196	+ uint32_t groupOffset = (subBlockIdx == 0) ? 0 : curGroupCount;
	197	+ // init mask
	198	+
	199	+ uint8_t repeatStride = S2_BASE_SIZE * sizeof(uint8_t);
	200	+
	201	+ AscendC::Duplicate(maskUbTensor, uint8_t(0), HALF_MASK_BLOCK_SIZE);
	202	+ AscendC::PipeBarrier<PIPE_V>();
	203	+
	204	+ for (uint32_t round = 0; round < curGroupThisSubBlock; ++round) {
	205	+ uint32_t colOffset = (groupOffset + round) * maxDecodeStep;
	206	+ uint32_t colOffsetFloor = colOffset / UINT8_BLOCK_SIZE * UINT8_BLOCK_SIZE;
	207	+ uint32_t floorSub = colOffset - colOffsetFloor;
	208	+ uint64_t rowOffset = round * groupSize * S2_BASE_SIZE;
	209	+ uint64_t totalOffset = rowOffset + colOffsetFloor;
	210	+ auto totalDupLen = unsharedKvSeqLen + floorSub;
	211	+ for (int groupId = 0; groupId < groupSize; groupId++) {
	212	+ AscendC::Duplicate(maskUbTensor[totalOffset + groupId * repeatStride], uint8_t(1), totalDupLen);
	213	+ }
	214	+
	215	+ AscendC::PipeBarrier<PIPE_V>();
	216	+ if (floorSub > 0) {
	217	+ for (int groupId = 0; groupId < groupSize; groupId++) {
	218	+ AscendC::Duplicate(maskUbTensor[totalOffset + groupId * repeatStride], uint8_t(0), floorSub);
	219	+ }
	220	+ AscendC::PipeBarrier<PIPE_V>();
	221	+ }
	222	+ }
	223	+ }
	224	+
	225	+ template <typename ElementS, uint16_t S2BaseSize, NRangeIndex NRange>
	226	+ __simd_vf__ inline void ComputeMaskandScale(
	227	+ __ubuf__ ElementS* srcUb, __ubuf__ uint8_t* maskUbUnRoll1, __ubuf__ uint8_t* maskUbUnRoll2,
	228	+ __ubuf__ ElementS* maxUb, uint16_t m, uint32_t tailN, ElementS scale)
	229	+ {
	230	+ using namespace AscendC::Reg;
	231	+ static_assert(NRange <= NRangeIndex::N128, "ComputeMaskandScale only supports N <= 128");
	232	+ RegTensor<ElementS> minVreg;
	233	+ RegTensor<ElementS> srcVreg0;
	234	+ RegTensor<ElementS> srcVreg1;
	235	+ RegTensor<ElementS> maxVreg;
	236	+ RegTensor<ElementS> maxTmpVreg;
	237	+ UnalignReg maxUreg;
	238	+ MaskReg pregCompare0;
	239	+ MaskReg pregCompare1;
	240	+ MaskReg pregFull = CreateMask<ElementS, MaskPattern::ALL>();
	241	+ MaskReg pregTailN = UpdateMask<ElementS>(tailN);
	242	+
	243	+ Duplicate(minVreg, MIN_VALUE);
	244	+ for (uint16_t i = 0; i < m; ++i) {
	245	+ if constexpr (NRange == NRangeIndex::N0_64) {
	246	+ LoadAlign(srcVreg0, srcUb + i * S2BaseSize);
	247	+ Muls(srcVreg0, srcVreg0, scale, pregTailN);
	248	+ LoadAlign<uint32_t, PostLiteral::POST_MODE_UPDATE, MaskDist::DIST_DS>(
	249	+ pregCompare0, (__ubuf__ uint32_t*&)maskUbUnRoll1, S2BaseSize);
	250	+ Select(srcVreg0, srcVreg0, minVreg, pregCompare0);
	251	+ Select(srcVreg0, srcVreg0, minVreg, pregTailN);
	252	+ StoreAlign<ElementS, StoreDist::DIST_NORM_B32>(srcUb + i * S2BaseSize, srcVreg0, pregFull);
	253	+ ReduceMax(maxVreg, srcVreg0, pregFull);
	254	+ } else {
	255	+ LoadAlign(srcVreg0, srcUb + i * S2BaseSize);
	256	+ LoadAlign(srcVreg1, srcUb + i * S2BaseSize + FLOAT_REP_SIZE);
	257	+ Muls(srcVreg0, srcVreg0, scale, pregFull);
	258	+ Muls(srcVreg1, srcVreg1, scale, NRange == NRangeIndex::N128 ? pregFull : pregTailN);
	259	+ LoadAlign<uint32_t, PostLiteral::POST_MODE_UPDATE, MaskDist::DIST_DS>(
	260	+ pregCompare0, (__ubuf__ uint32_t*&)maskUbUnRoll1, S2BaseSize);
	261	+ LoadAlign<uint32_t, PostLiteral::POST_MODE_UPDATE, MaskDist::DIST_DS>(
	262	+ pregCompare1, (__ubuf__ uint32_t*&)maskUbUnRoll2, S2BaseSize);
	263	+ Select(srcVreg0, srcVreg0, minVreg, pregCompare0);
	264	+ Select(srcVreg1, srcVreg1, minVreg, pregCompare1);
	265	+ if constexpr (NRange == NRangeIndex::N65_127) {
	266	+ Select(srcVreg1, srcVreg1, minVreg, pregTailN);
	267	+ }
	268	+ StoreAlign<ElementS, StoreDist::DIST_NORM_B32>(srcUb + i * S2BaseSize, srcVreg0, pregFull);
	269	+ StoreAlign<ElementS, StoreDist::DIST_NORM_B32>(
	270	+ srcUb + i * S2BaseSize + FLOAT_REP_SIZE, srcVreg1, pregFull);
	271	+ Max(maxTmpVreg, srcVreg0, srcVreg1, pregFull);
	272	+ ReduceMax(maxVreg, maxTmpVreg, pregFull);
	273	+ }
	274	+ StoreUnAlign<ElementS, PostLiteral::POST_MODE_UPDATE>(maxUb, maxVreg, maxUreg, 1);
	275	+ }
	276	+ StoreUnAlignPost(maxUb, maxUreg, 0);
	277	+ }
	278	+
	279	+ template <typename ElementP, typename ElementS, uint16_t S2BaseSize, NRangeIndex NRange>
	280	+ __simd_vf__ inline void ComputeExpSubSum(
	281	+ __ubuf__ ElementP* expUb, __ubuf__ ElementS* srcUb, __ubuf__ ElementS* nowMaxUb, __ubuf__ ElementS* expSumUb,
	282	+ uint16_t m, uint32_t blockStride)
	283	+ {
	284	+ using namespace AscendC::Reg;
	285	+ constexpr static CastTrait castTraitZero = {
	286	+ RegLayout::ZERO,
	287	+ SatMode::SAT,
	288	+ MaskMergeMode::ZEROING,
	289	+ AscendC::RoundMode::CAST_ROUND,
	290	+ };
	291	+
	292	+ constexpr static CastTrait castTraitOne = {
	293	+ RegLayout::ONE,
	294	+ SatMode::SAT,
	295	+ MaskMergeMode::ZEROING,
	296	+ AscendC::RoundMode::CAST_ROUND,
	297	+ };
	298	+ RegTensor<ElementS> srcVreg0;
	299	+ RegTensor<ElementS> srcVreg1;
	300	+ RegTensor<ElementS> expVreg;
	301	+ RegTensor<ElementS> expVreg0;
	302	+ RegTensor<ElementS> expVreg1;
	303	+ RegTensor<ElementS> expSumVreg;
	304	+ RegTensor<ElementS> maxVreg;
	305	+
	306	+ RegTensor<ElementP> expDstVreg0;
	307	+ RegTensor<ElementP> expDstVreg1;
	308	+ RegTensor<ElementP> expDstVreg;
	309	+
	310	+ UnalignReg expSumUreg;
	311	+
	312	+ MaskReg pregFull = CreateMask<ElementS, MaskPattern::ALL>();
	313	+ MaskReg pregFull16 = CreateMask<uint16_t, MaskPattern::ALL>();
	314	+ for (uint16_t i = 0; i < m; ++i) {
	315	+ LoadAlign<ElementS, LoadDist::DIST_BRC_B32>(maxVreg, nowMaxUb + i);
	316	+ if constexpr (NRange > NRangeIndex::N0_64) {
	317	+ LoadAlign<ElementS, LoadDist::DIST_DINTLV_B32>(srcVreg0, srcVreg1, srcUb + i * S2BaseSize);
	318	+ ExpSub(expVreg0, srcVreg0, maxVreg, pregFull);
	319	+ ExpSub(expVreg1, srcVreg1, maxVreg, pregFull);
	320	+ Add(expVreg, expVreg0, expVreg1, pregFull);
	321	+ Cast<ElementP, ElementS, castTraitZero>(expDstVreg0, expVreg0, pregFull);
	322	+ Cast<ElementP, ElementS, castTraitOne>(expDstVreg1, expVreg1, pregFull);
	323	+ Or((RegTensor<uint16_t>&)expDstVreg, (RegTensor<uint16_t>&)expDstVreg0,
	324	+ (RegTensor<uint16_t>&)expDstVreg1, pregFull16);
	325	+ StoreAlign<ElementP, DataCopyMode::DATA_BLOCK_COPY, PostLiteral::POST_MODE_UPDATE>(
	326	+ expUb, expDstVreg, blockStride, REPEAT_STRIDE, pregFull16);
	327	+ } else {
	328	+ LoadAlign(srcVreg0, srcUb + i * S2BaseSize);
	329	+ ExpSub(expVreg, srcVreg0, maxVreg, pregFull);
	330	+ Cast<ElementP, ElementS, castTraitZero>(expDstVreg, expVreg, pregFull);
	331	+ DeInterleave(expDstVreg0, expDstVreg1, expDstVreg, expDstVreg);
	332	+ StoreAlign<ElementP, DataCopyMode::DATA_BLOCK_COPY, PostLiteral::POST_MODE_UPDATE>(
	333	+ expUb, expDstVreg0, blockStride, REPEAT_STRIDE, pregFull16);
	334	+ }
	335	+ ReduceSum(expSumVreg, expVreg, pregFull);
	336	+ StoreUnAlign<ElementS, PostLiteral::POST_MODE_UPDATE>(expSumUb, expSumVreg, expSumUreg, 1);
	337	+ }
	338	+ StoreUnAlignPost(expSumUb, expSumUreg, 0);
	339	+ }
	340	+
	341	+ template <typename ElementS>
	342	+ CATLASS_DEVICE void CopyOutMaxAndSum(
	343	+ AscendC::GlobalTensor<ElementS> maxGm, AscendC::GlobalTensor<ElementS> sumGm, uint16_t tailM)
	344	+ {
	345	+ AscendC::DataCopyExtParams copyOutParams(1, static_cast<uint32_t>(tailM) * sizeof(ElementS), 0, 0, 0);
	346	+
	347	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID7);
	348	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID7);
	349	+ AscendC::DataCopyPad(maxGm, maxUbTensor, copyOutParams);
	350	+ AscendC::DataCopyPad(sumGm, sumUbTensor, copyOutParams);
	351	+ }
	352	+};
	353	+
	354	+} // namespace Catlass::Epilogue::Block
	355	+
	356	+#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_XA_SHARED_SOFTMAX_ASCEND950

Ainclude/catlass/epilogue/block/block_epilogue_xfai_combine_scale.hpp+283-0

@@ -0,0 +1,283 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_EPILOGUE_BLOCK_XFAI_COMBINE_SCALE_HPP
	13	+#define CATLASS_EPILOGUE_BLOCK_XFAI_COMBINE_SCALE_HPP
	14	+
	15	+#include <limits>
	16	+#include "catlass/arch/resource.hpp"
	17	+
	18	+namespace Catlass::Epilogue::Block {
	19	+
	20	+template <class OutputType_, class LseType_>
	21	+class CombineScale {
	22	+public:
	23	+ using ElementOutput = typename OutputType_::Element;
	24	+ using ElementLse = typename LseType_::Element;
	25	+ using ArchTag = Arch::AtlasA2;
	26	+
	27	+ // Placeholder constants matching original logic
	28	+ static constexpr uint32_t STAGE2_UB_UINT8_BLOCK_SIZE = 6144; // 24 * 64 * 4
	29	+ static constexpr uint32_t UB_UINT8_LINE_SIZE = 32768; // 1 * 64 * 128 * 4
	30	+ static constexpr uint32_t BASIC_BLOCK = 64 * 128;
	31	+
	32	+ __aicore__ inline CombineScale(Arch::Resource<ArchTag>& resource)
	33	+ {
	34	+ // UB Memory Allocation
	35	+ constexpr uint32_t LL_UB_OFFSET = 0; // splitnum_align * (q * h)_algin
	36	+ constexpr uint32_t LM_UB_OFFSET = 1 * STAGE2_UB_UINT8_BLOCK_SIZE; // 1 * (q * h)_algin
	37	+ constexpr uint32_t BROADCAST_OFFSET = 2 * STAGE2_UB_UINT8_BLOCK_SIZE; // splitnum_align * (q * h)_algin
	38	+ constexpr uint32_t TL_UB_OFFSET = 3 * STAGE2_UB_UINT8_BLOCK_SIZE; // splitnum_align * (q * h)_algin
	39	+ constexpr uint32_t RS_UB_OFFSET = 4 * STAGE2_UB_UINT8_BLOCK_SIZE; // 1 * (q * h)_algin
	40	+ constexpr uint32_t TS_UB_OFFSET = 5 * STAGE2_UB_UINT8_BLOCK_SIZE; // 1 * (q * h)_algin
	41	+ constexpr uint32_t BROADCASTSCALE_OFFSET = 6 * STAGE2_UB_UINT8_BLOCK_SIZE; // splitnum_align * (q * h)_algin
	42	+ constexpr uint32_t GL_UB_OFFSET = 7 * STAGE2_UB_UINT8_BLOCK_SIZE; // splitnum_align * (q * h)_algin
	43	+ constexpr uint32_t BROADCASTO_OFFSET = 8 * STAGE2_UB_UINT8_BLOCK_SIZE; // splitnum_align * (q * h)_algin * v
	44	+ constexpr uint32_t GO16_UB_OFFSET =
	45	+ 8 * STAGE2_UB_UINT8_BLOCK_SIZE + 2 * UB_UINT8_LINE_SIZE; // (q * h)_algin * v
	46	+ constexpr uint32_t tempReduceMax_OFFSET = 9 * STAGE2_UB_UINT8_BLOCK_SIZE; // splitnum_align * (q * h)_algin * v
	47	+ constexpr uint32_t tempReduceSum_OFFSET = 10 * STAGE2_UB_UINT8_BLOCK_SIZE; // splitnum_align * (q * h)_algin * v
	48	+ constexpr uint32_t GO_UB_OFFSET = 0; // (q * h)_algin * v
	49	+
	50	+ // Buffer Init
	51	+ llUbTensor = resource.ubBuf.template GetBufferByByte<float>(LL_UB_OFFSET);
	52	+ lmUbTensor = resource.ubBuf.template GetBufferByByte<float>(LM_UB_OFFSET);
	53	+ broadCastTensor = resource.ubBuf.template GetBufferByByte<float>(BROADCAST_OFFSET);
	54	+ tlUbTensor = resource.ubBuf.template GetBufferByByte<float>(TL_UB_OFFSET);
	55	+ rsUbTensor = resource.ubBuf.template GetBufferByByte<float>(RS_UB_OFFSET);
	56	+ tsUbTensor = resource.ubBuf.template GetBufferByByte<float>(TS_UB_OFFSET);
	57	+ broadCastScaleTensor = resource.ubBuf.template GetBufferByByte<float>(BROADCASTSCALE_OFFSET);
	58	+ glUbTensor = resource.ubBuf.template GetBufferByByte<float>(GL_UB_OFFSET);
	59	+ broadCastOTensor = resource.ubBuf.template GetBufferByByte<float>(BROADCASTO_OFFSET);
	60	+ toUbTensor = resource.ubBuf.template GetBufferByByte<float>(BROADCASTO_OFFSET);
	61	+ goUbTensor = resource.ubBuf.template GetBufferByByte<float>(GO_UB_OFFSET);
	62	+ loFloatUbTensor = resource.ubBuf.template GetBufferByByte<float>(GO16_UB_OFFSET);
	63	+ go16UbTensor = resource.ubBuf.template GetBufferByByte<ElementOutput>(GO_UB_OFFSET);
	64	+
	65	+ tempReduceMax = resource.ubBuf.template GetBufferByByte<uint8_t>(tempReduceMax_OFFSET);
	66	+ tempReduceSum = resource.ubBuf.template GetBufferByByte<uint8_t>(tempReduceSum_OFFSET); // Offset example
	67	+ }
	68	+
	69	+ __aicore__ inline void operator()(
	70	+ uint32_t qHeads, uint32_t kvSplitCoreNum, uint32_t headSizeV, __gm__ SplitKvExtraInfo* extraInfo,
	71	+ AscendC::GlobalTensor<ElementLse> lGmTensor, AscendC::GlobalTensor<ElementLse> oCoreTmpGmTensor,
	72	+ AscendC::GlobalTensor<ElementOutput> oGmTensor, AscendC::GlobalTensor<int32_t> gActualQseqlen,
	73	+ bool inputLayoutTND = true)
	74	+ {
	75	+ int64_t vectorsubBlockID = AscendC::GetSubBlockIdx();
	76	+ int64_t subBlockNum = AscendC::GetBlockNum() * 2;
	77	+ int64_t subBlockID = AscendC::GetBlockIdx();
	78	+
	79	+ for (uint32_t process = subBlockID; process < kvSplitCoreNum * 2; process += subBlockNum) {
	80	+ uint32_t batchIdx = extraInfo->splitInfo[process / 2].batchIdx;
	81	+ uint32_t headStartIndx = extraInfo->splitInfo[process / 2].headStartIdx;
	82	+ uint32_t headEndIndx = extraInfo->splitInfo[process / 2].headEndIdx;
	83	+ uint32_t qStartIndx = extraInfo->splitInfo[process / 2].qStartIdx;
	84	+ uint32_t qEndIndx = extraInfo->splitInfo[process / 2].qEndIdx;
	85	+ uint32_t pingpongFlag = 0;
	86	+
	87	+ uint32_t q_len = (qEndIndx - qStartIndx);
	88	+ uint32_t n_len = (headEndIndx - headStartIndx);
	89	+
	90	+ uint32_t sum = q_len * n_len;
	91	+ uint32_t sum_former = q_len == 1 ? sum / 2 : (q_len / 2) * n_len;
	92	+
	93	+ uint32_t addrLOffset = vectorsubBlockID == 0 ? extraInfo->splitInfo[process / 2].lseTaskOffset :
	94	+ extraInfo->splitInfo[process / 2].lseTaskOffset + sum_former;
	95	+ uint32_t addrOOffset = vectorsubBlockID == 0 ?
	96	+ extraInfo->splitInfo[process / 2].oTaskOffset :
	97	+ extraInfo->splitInfo[process / 2].oTaskOffset + sum_former * headSizeV;
	98	+
	99	+ uint32_t prevQSeqlenSum = 0;
	100	+ if (inputLayoutTND) {
	101	+ prevQSeqlenSum = (batchIdx == 0) ? 0 : static_cast<uint32_t>(gActualQseqlen.GetValue(batchIdx - 1));
	102	+ }
	103	+ uint32_t baseGmOffset =
	104	+ prevQSeqlenSum * qHeads * headSizeV + qStartIndx * qHeads * headSizeV + headStartIndx * headSizeV;
	105	+ uint32_t gmOScalar = 0;
	106	+ if (q_len == 1) {
	107	+ gmOScalar = vectorsubBlockID == 0 ? baseGmOffset : baseGmOffset + sum_former * headSizeV;
	108	+ } else {
	109	+ uint32_t q_half = q_len / 2;
	110	+ gmOScalar = vectorsubBlockID == 0 ? baseGmOffset : baseGmOffset + q_half * qHeads * headSizeV;
	111	+ }
	112	+
	113	+ uint32_t splitNum = extraInfo->splitInfo[process / 2].splitNum;
	114	+
	115	+ uint32_t splitNumAlign = (splitNum + 7) / 8 * 8; // 32b align
	116	+ uint32_t lseBlock = vectorsubBlockID == 0 ? sum_former : sum - sum_former;
	117	+ uint32_t lseBlockAlign = (lseBlock + 7) / 8 * 8; // 32b align
	118	+ int32_t count = splitNum * lseBlockAlign;
	119	+ int32_t lnCount = 1 * lseBlockAlign;
	120	+ // Initialize LSE UB space
	121	+ int32_t calcLen = splitNumAlign * lseBlockAlign;
	122	+ int32_t oCount = lseBlock * headSizeV;
	123	+ int32_t oCount_vector = sum * headSizeV;
	124	+
	125	+ AscendC::Duplicate(llUbTensor, std::numeric_limits<float>::lowest(), calcLen);
	126	+ AscendC::Duplicate(tlUbTensor, 0.0f, calcLen);
	127	+
	128	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
	129	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
	130	+
	131	+ // Copy LSE from GM to UB
	132	+ uint32_t srcStride = vectorsubBlockID == 0 ? sum - sum_former : sum_former;
	133	+ AscendC::DataCopyPad(
	134	+ llUbTensor, lGmTensor[addrLOffset],
	135	+ AscendC::DataCopyExtParams(splitNum, lseBlock * sizeof(float), srcStride * sizeof(float), 0, 0),
	136	+ AscendC::DataCopyPadExtParams<float>(false, 0, lseBlockAlign - lseBlock, 0));
	137	+
	138	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	139	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	140	+
	141	+ // // Preload O
	142	+ AscendC::DataCopyPad(
	143	+ loFloatUbTensor[pingpongFlag * BASIC_BLOCK], oCoreTmpGmTensor[addrOOffset],
	144	+ AscendC::DataCopyExtParams(1, oCount * sizeof(float), 0, 0, 0),
	145	+ AscendC::DataCopyPadExtParams<float>(true, 0, 0, 0));
	146	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1 + pingpongFlag);
	147	+
	148	+ // ReduceMax
	149	+ uint32_t reduceMaxShape[] = {splitNumAlign, lseBlockAlign};
	150	+ AscendC::ReduceMax<float, AscendC::Pattern::Reduce::RA, false>(
	151	+ lmUbTensor, llUbTensor, tempReduceMax, reduceMaxShape, true);
	152	+ AscendC::PipeBarrier<PIPE_V>();
	153	+
	154	+ // Broadcast Max
	155	+ uint32_t dstShapeBroadcast[] = {splitNum, lseBlockAlign};
	156	+ uint32_t srcShapeBroadcast[] = {1, lseBlockAlign};
	157	+ AscendC::BroadCast<float, 2, 0>(
	158	+ broadCastTensor, lmUbTensor, dstShapeBroadcast, srcShapeBroadcast, tempReduceSum);
	159	+ AscendC::PipeBarrier<PIPE_V>();
	160	+
	161	+ AscendC::Sub(tlUbTensor, llUbTensor, broadCastTensor, count);
	162	+ AscendC::PipeBarrier<PIPE_V>();
	163	+
	164	+ // expf
	165	+ AscendC::Exp(tlUbTensor, tlUbTensor, count);
	166	+ AscendC::PipeBarrier<PIPE_V>();
	167	+
	168	+ // ReduceSum
	169	+ uint32_t reduceSumShape[] = {splitNumAlign, lseBlockAlign};
	170	+ AscendC::ReduceSum<float, AscendC::Pattern::Reduce::RA, false>(
	171	+ rsUbTensor, tlUbTensor, tempReduceSum, reduceSumShape, true);
	172	+ AscendC::PipeBarrier<PIPE_V>();
	173	+
	174	+ // Ln
	175	+ AscendC::Ln(rsUbTensor, rsUbTensor, lnCount);
	176	+ AscendC::PipeBarrier<PIPE_V>();
	177	+
	178	+ // logf(lse_sum) + lse_max
	179	+ AscendC::Add(tsUbTensor, rsUbTensor, lmUbTensor, lnCount);
	180	+ AscendC::PipeBarrier<PIPE_V>();
	181	+
	182	+ // Broadcast scale
	183	+ AscendC::BroadCast<float, 2, 0>(
	184	+ broadCastScaleTensor, tsUbTensor, dstShapeBroadcast, srcShapeBroadcast, tempReduceSum);
	185	+ AscendC::PipeBarrier<PIPE_V>();
	186	+
	187	+ AscendC::Sub(glUbTensor, llUbTensor, broadCastScaleTensor, count);
	188	+ AscendC::PipeBarrier<PIPE_V>();
	189	+
	190	+ AscendC::Exp(glUbTensor, glUbTensor, count);
	191	+ AscendC::PipeBarrier<PIPE_V>();
	192	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
	193	+ for (uint32_t nIdx = 0; nIdx < splitNum; nIdx++) {
	194	+ if (nIdx != 0) {
	195	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2 + pingpongFlag);
	196	+ AscendC::DataCopyPad(
	197	+ loFloatUbTensor[pingpongFlag * BASIC_BLOCK],
	198	+ oCoreTmpGmTensor[addrOOffset + nIdx * oCount_vector],
	199	+ AscendC::DataCopyExtParams(1, oCount * sizeof(float), 0, 0, 0),
	200	+ AscendC::DataCopyPadExtParams<float>(true, 0, 0, 0));
	201	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1 + pingpongFlag);
	202	+ }
	203	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1 + pingpongFlag);
	204	+
	205	+ uint32_t dstShapeO[2] = {lseBlockAlign, headSizeV};
	206	+ uint32_t srcShapeO[2] = {lseBlockAlign, 1};
	207	+ AscendC::BroadCast<float, 2, 1>(
	208	+ broadCastOTensor[pingpongFlag * BASIC_BLOCK], glUbTensor[nIdx * lseBlockAlign], dstShapeO,
	209	+ srcShapeO, tempReduceSum);
	210	+ AscendC::PipeBarrier<PIPE_V>();
	211	+
	212	+ AscendC::Mul(
	213	+ toUbTensor[pingpongFlag * BASIC_BLOCK], loFloatUbTensor[pingpongFlag * BASIC_BLOCK],
	214	+ broadCastOTensor[pingpongFlag * BASIC_BLOCK], oCount); // toUbTensor和broadCastOTensor共用一块空间
	215	+ AscendC::PipeBarrier<PIPE_V>();
	216	+
	217	+ if (nIdx == 0) {
	218	+ AscendC::Adds(goUbTensor, toUbTensor[pingpongFlag * BASIC_BLOCK], 0.0f, oCount);
	219	+ AscendC::PipeBarrier<PIPE_V>();
	220	+ } else {
	221	+ AscendC::Add(goUbTensor, toUbTensor[pingpongFlag * BASIC_BLOCK], goUbTensor, oCount);
	222	+ AscendC::PipeBarrier<PIPE_V>();
	223	+ }
	224	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2 + pingpongFlag);
	225	+ pingpongFlag = 1 - pingpongFlag;
	226	+ }
	227	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);
	228	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
	229	+ AscendC::PipeBarrier<PIPE_V>();
	230	+ // Cast and move out
	231	+ if (std::is_same<ElementOutput, bfloat16_t>::value) {
	232	+ AscendC::Cast(go16UbTensor, goUbTensor, AscendC::RoundMode::CAST_RINT, oCount);
	233	+ } else {
	234	+ AscendC::Cast(go16UbTensor, goUbTensor, AscendC::RoundMode::CAST_NONE, oCount);
	235	+ }
	236	+
	237	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID1);
	238	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID1);
	239	+
	240	+ if (q_len == 1) {
	241	+ AscendC::DataCopyPad(
	242	+ oGmTensor[gmOScalar], go16UbTensor,
	243	+ AscendC::DataCopyExtParams(1, oCount * sizeof(ElementOutput), 0, 0, 0));
	244	+ } else {
	245	+ uint32_t q_half = q_len / 2;
	246	+ if (vectorsubBlockID == 0) {
	247	+ AscendC::DataCopyPad(
	248	+ oGmTensor[gmOScalar], go16UbTensor,
	249	+ AscendC::DataCopyExtParams(
	250	+ q_half, (headEndIndx - headStartIndx) * headSizeV * sizeof(ElementOutput), 0,
	251	+ (qHeads - (headEndIndx - headStartIndx)) * headSizeV * sizeof(ElementOutput), 0));
	252	+ } else {
	253	+ AscendC::DataCopyPad(
	254	+ oGmTensor[gmOScalar], go16UbTensor,
	255	+ AscendC::DataCopyExtParams(
	256	+ q_len - q_half, (headEndIndx - headStartIndx) * headSizeV * sizeof(ElementOutput), 0,
	257	+ (qHeads - (headEndIndx - headStartIndx)) * headSizeV * sizeof(ElementOutput), 0));
	258	+ }
	259	+ }
	260	+ }
	261	+ }
	262	+
	263	+private:
	264	+ AscendC::LocalTensor<float> llUbTensor;
	265	+ AscendC::LocalTensor<float> lmUbTensor;
	266	+ AscendC::LocalTensor<float> tlUbTensor;
	267	+ AscendC::LocalTensor<float> rsUbTensor;
	268	+ AscendC::LocalTensor<float> tsUbTensor;
	269	+ AscendC::LocalTensor<float> glUbTensor;
	270	+ AscendC::LocalTensor<float> toUbTensor;
	271	+ AscendC::LocalTensor<float> goUbTensor;
	272	+ AscendC::LocalTensor<ElementOutput> go16UbTensor;
	273	+ AscendC::LocalTensor<float> loFloatUbTensor;
	274	+
	275	+ AscendC::LocalTensor<uint8_t> tempReduceMax;
	276	+ AscendC::LocalTensor<uint8_t> tempReduceSum;
	277	+ AscendC::LocalTensor<float> broadCastTensor;
	278	+ AscendC::LocalTensor<float> broadCastScaleTensor;
	279	+ AscendC::LocalTensor<float> broadCastOTensor;
	280	+};
	281	+} // namespace Catlass::Epilogue::Block
	282	+
	283	+#endif // CATLASS_EPILOGUE_BLOCK_XFAI_COMBINE_SCALE_HPP

Ainclude/catlass/epilogue/block/block_epilogue_xfai_online_softmax.hpp+841-0

@@ -0,0 +1,841 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_HPP
	13	+#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_HPP
	14	+
	15	+#include "catlass/arch/cross_core_sync.hpp"
	16	+#include "catlass/arch/resource.hpp"
	17	+#include "catlass/epilogue/dispatch_policy.hpp"
	18	+#include "catlass/epilogue/tile/tile_copy.hpp"
	19	+#include "catlass/gemm_coord.hpp"
	20	+#include "catlass/matrix_coord.hpp"
	21	+
	22	+namespace Catlass::Epilogue::Block {
	23	+
	24	+template <class OutputType_, class InputType_, class MaskType_, LseMode LSE_MODE_>
	25	+class BlockEpilogue<EpilogueAtlasA2XFAIOnlineSoftmax<LSE_MODE_>, OutputType_, InputType_, MaskType_> {
	26	+public:
	27	+ using DispatchPolicy = EpilogueAtlasA2XFAIOnlineSoftmax<LSE_MODE_>;
	28	+ using ArchTag = typename DispatchPolicy::ArchTag;
	29	+ using ElementOutput = typename OutputType_::Element;
	30	+ using ElementInput = typename InputType_::Element;
	31	+ using ElementMask = typename MaskType_::Element;
	32	+
	33	+ using LayoutOutput = typename OutputType_::Layout;
	34	+ using LayoutInput = typename InputType_::Layout;
	35	+ using LayoutMask = typename MaskType_::Layout;
	36	+
	37	+ static constexpr LseMode LSE_MODE = DispatchPolicy::LSE_MODE;
	38	+
	39	+ static constexpr uint32_t BLOCK_SIZE_IN_BYTE = 32;
	40	+ static constexpr uint32_t REPEAT_SIZE_IN_BYTE = 256;
	41	+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
	42	+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
	43	+ static constexpr uint32_t HALF_VECTOR_SIZE = 128;
	44	+ static constexpr uint32_t BLOCK_SIZE = 16;
	45	+ static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024;
	46	+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384;
	47	+ static constexpr uint32_t VECTOR_SIZE = 128;
	48	+ static constexpr uint32_t MAX_UB_S_ELEM_NUM = 8192;
	49	+
	50	+ static constexpr uint32_t REDUCE_UB_SIZE = 1024;
	51	+ static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32;
	52	+ static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4;
	53	+ static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 256;
	54	+ static constexpr int64_t UB_FLOAT_LINE_SIZE = 64;
	55	+
	56	+ __aicore__ inline BlockEpilogue(Arch::Resource<ArchTag>& resource, float scaleValue_)
	57	+ {
	58	+ // Allocate UB space
	59	+ constexpr uint32_t LS_UB_TENSOR_OFFSET = 0;
	60	+ constexpr uint32_t LP_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE;
	61	+ constexpr uint32_t MASK_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE;
	62	+ constexpr uint32_t MASK32_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE;
	63	+
	64	+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE;
	65	+ constexpr uint32_t LM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 8 * UB_UINT8_VECTOR_SIZE;
	66	+
	67	+ constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE;
	68	+ constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE;
	69	+ constexpr uint32_t LL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 11 * UB_UINT8_VECTOR_SIZE;
	70	+ constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE;
	71	+ constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE;
	72	+
	73	+ constexpr uint32_t MASK16_UB_TENSOR_OFFSET = 11 * UB_UINT8_BLOCK_SIZE;
	74	+
	75	+ scaleValue = scaleValue_;
	76	+ lsUbTensor = resource.ubBuf.template GetBufferByByte<float>(LS_UB_TENSOR_OFFSET);
	77	+ lpUbTensor = resource.ubBuf.template GetBufferByByte<ElementOutput>(LP_UB_TENSOR_OFFSET);
	78	+ maskUbTensor = resource.ubBuf.template GetBufferByByte<ElementMask>(MASK_UB_TENSOR_OFFSET);
	79	+ maskUbTensor16 = resource.ubBuf.template GetBufferByByte<half>(MASK16_UB_TENSOR_OFFSET);
	80	+ maskUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(MASK32_UB_TENSOR_OFFSET);
	81	+ lmUbTensor = resource.ubBuf.template GetBufferByByte<float>(LM_UB_TENSOR_OFFSET);
	82	+ hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);
	83	+ gmUbTensor = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
	84	+ dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
	85	+ llUbTensor = resource.ubBuf.template GetBufferByByte<float>(LL_UB_TENSOR_OFFSET);
	86	+ tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
	87	+ glUbTensor = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
	88	+ }
	89	+
	90	+ __aicore__ inline ~BlockEpilogue()
	91	+ {}
	92	+
	93	+ template <typename T>
	94	+ __aicore__ inline T Min(T a, T b)
	95	+ {
	96	+ return (a > b) ? b : a;
	97	+ }
	98	+
	99	+ __aicore__ inline void SetVecMask(int32_t len)
	100	+ {
	101	+ uint64_t mask = 0;
	102	+ uint64_t one = 1;
	103	+ uint64_t temp = len % FLOAT_VECTOR_SIZE;
	104	+ for (int64_t i = 0; i < temp; i++) {
	105	+ mask \|= one << i;
	106	+ }
	107	+
	108	+ if (len == VECTOR_SIZE \|\| len == 0) {
	109	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	110	+ } else if (len >= FLOAT_VECTOR_SIZE) {
	111	+ AscendC::SetVectorMask<int8_t>(mask, (uint64_t)-1);
	112	+ } else {
	113	+ AscendC::SetVectorMask<int8_t>(0x0, mask);
	114	+ }
	115	+ }
	116	+
	117	+ __aicore__ inline void SetBlockReduceMask(int32_t len)
	118	+ {
	119	+ if (len > 8 \|\| len < 1) {
	120	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	121	+ return;
	122	+ }
	123	+ uint64_t subMask = ((uint64_t)1 << len) - 1;
	124	+ uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask + (subMask << 56) +
	125	+ (subMask << 40) + (subMask << 24) + (subMask << 8);
	126	+ AscendC::SetVectorMask<int8_t>(maskValue, maskValue);
	127	+ }
	128	+
	129	+ __aicore__ inline void RowsumSPECTILE512(
	130	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowsumUb,
	131	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	132	+ uint32_t numElemsAligned)
	133	+ {
	134	+ AscendC::BlockReduceSum<float, false>(
	135	+ tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	136	+ AscendC::PipeBarrier<PIPE_V>();
	137	+
	138	+ AscendC::BlockReduceSum<float, false>(
	139	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
	140	+ numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	141	+ AscendC::PipeBarrier<PIPE_V>();
	142	+ AscendC::BlockReduceSum<float, false>(
	143	+ rowsumUb, tvUbTensor[REDUCE_UB_SIZE],
	144	+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	145	+ AscendC::PipeBarrier<PIPE_V>();
	146	+ }
	147	+
	148	+ __aicore__ inline void RowsumSPECTILE256(
	149	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowsumUb,
	150	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	151	+ uint32_t numElemsAligned)
	152	+ {
	153	+ AscendC::BlockReduceSum<float, false>(
	154	+ tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	155	+ AscendC::PipeBarrier<PIPE_V>();
	156	+ SetVecMask(ROW_OPS_SPEC_MASK_32);
	157	+ AscendC::BlockReduceSum<float, false>(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4);
	158	+ AscendC::PipeBarrier<PIPE_V>();
	159	+ SetBlockReduceMask(ROW_OPS_SPEC_MASK_4);
	160	+ AscendC::BlockReduceSum<float, false>(
	161	+ rowsumUb, tvUbTensor[REDUCE_UB_SIZE], CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), 0, 1, 1,
	162	+ 8);
	163	+ AscendC::PipeBarrier<PIPE_V>();
	164	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	165	+ }
	166	+
	167	+ __aicore__ inline void RowsumTAILTILE(
	168	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowsumUb,
	169	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	170	+ uint32_t numElemsAligned)
	171	+ {
	172	+ if (numElems >= FLOAT_VECTOR_SIZE) {
	173	+ AscendC::BlockReduceSum<float, false>(
	174	+ tvUbTensor, srcUb, numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
	175	+ AscendC::PipeBarrier<PIPE_V>();
	176	+ AscendC::BlockReduceSum<float, false>(
	177	+ rowsumUb, tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), 0, 1, 1, 8);
	178	+ AscendC::PipeBarrier<PIPE_V>();
	179	+ for (uint64_t rowSumIdx = 1; rowSumIdx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowSumIdx) {
	180	+ AscendC::BlockReduceSum<float, false>(
	181	+ tvUbTensor, srcUb[rowSumIdx * FLOAT_VECTOR_SIZE], numRowsRound, 0, 1, 1,
	182	+ numElemsAligned / FLOAT_BLOCK_SIZE);
	183	+ AscendC::PipeBarrier<PIPE_V>();
	184	+ AscendC::BlockReduceSum<float, false>(
	185	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
	186	+ 0, 1, 1, 8);
	187	+ AscendC::PipeBarrier<PIPE_V>();
	188	+ SetVecMask(numRowsRound);
	189	+ AscendC::Add<float, false>(
	190	+ rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
	191	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	192	+ AscendC::PipeBarrier<PIPE_V>();
	193	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	194	+ }
	195	+ }
	196	+ if (numElems % FLOAT_VECTOR_SIZE > 0) {
	197	+ SetVecMask(numElems % FLOAT_VECTOR_SIZE);
	198	+ AscendC::BlockReduceSum<float, false>(
	199	+ tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], numRowsRound, 0, 1, 1,
	200	+ numElemsAligned / FLOAT_BLOCK_SIZE);
	201	+ AscendC::PipeBarrier<PIPE_V>();
	202	+ SetBlockReduceMask(CeilDiv(numElems % FLOAT_VECTOR_SIZE, FLOAT_BLOCK_SIZE));
	203	+ if (numElems < FLOAT_VECTOR_SIZE) {
	204	+ AscendC::BlockReduceSum<float, false>(
	205	+ rowsumUb, tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), 0, 1, 1, 8);
	206	+ AscendC::PipeBarrier<PIPE_V>();
	207	+ } else {
	208	+ AscendC::BlockReduceSum<float, false>(
	209	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
	210	+ 0, 1, 1, 8);
	211	+ AscendC::PipeBarrier<PIPE_V>();
	212	+ SetVecMask(numRowsRound);
	213	+ AscendC::Add<float, false>(
	214	+ rowsumUb, rowsumUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
	215	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	216	+ AscendC::PipeBarrier<PIPE_V>();
	217	+ }
	218	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	219	+ }
	220	+ }
	221	+
	222	+ __aicore__ inline void RowmaxSPECTILE512(
	223	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowmaxUb,
	224	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	225	+ uint32_t numElemsAligned)
	226	+ {
	227	+ AscendC::BlockReduceMax<float, false>(
	228	+ tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	229	+ AscendC::PipeBarrier<PIPE_V>();
	230	+ AscendC::BlockReduceMax<float, false>(
	231	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor,
	232	+ numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	233	+ AscendC::PipeBarrier<PIPE_V>();
	234	+ AscendC::BlockReduceMax<float, false>(
	235	+ rowmaxUb, tvUbTensor[REDUCE_UB_SIZE],
	236	+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	237	+ AscendC::PipeBarrier<PIPE_V>();
	238	+ }
	239	+
	240	+ __aicore__ inline void RowmaxSPECTILE256(
	241	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowmaxUb,
	242	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	243	+ uint32_t numElemsAligned)
	244	+ {
	245	+ AscendC::BlockReduceMax<float, false>(
	246	+ tvUbTensor, srcUb, numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, 0, 1, 1, 8);
	247	+ AscendC::PipeBarrier<PIPE_V>();
	248	+ SetVecMask(ROW_OPS_SPEC_MASK_32);
	249	+ AscendC::BlockReduceMax<float, false>(tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, numRowsRound, 0, 1, 1, 4);
	250	+ AscendC::PipeBarrier<PIPE_V>();
	251	+ SetBlockReduceMask(ROW_OPS_SPEC_MASK_4);
	252	+ AscendC::BlockReduceMax<float, false>(
	253	+ rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), 0, 1, 1,
	254	+ 8);
	255	+ AscendC::PipeBarrier<PIPE_V>();
	256	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	257	+ }
	258	+
	259	+ __aicore__ inline void RowmaxTAILTILE(
	260	+ const AscendC::LocalTensor<float>& srcUb, const AscendC::LocalTensor<float>& rowmaxUb,
	261	+ const AscendC::LocalTensor<float>& tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
	262	+ uint32_t numElemsAligned)
	263	+ {
	264	+ if (numElems >= FLOAT_VECTOR_SIZE) {
	265	+ AscendC::BlockReduceMax<float, false>(
	266	+ tvUbTensor, srcUb, numRowsRound, 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
	267	+ AscendC::PipeBarrier<PIPE_V>();
	268	+ AscendC::BlockReduceMax<float, false>(
	269	+ rowmaxUb, tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), 0, 1, 1, 8);
	270	+ AscendC::PipeBarrier<PIPE_V>();
	271	+ for (uint64_t rowmax_idx = 1; rowmax_idx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowmax_idx) {
	272	+ AscendC::BlockReduceMax<float, false>(
	273	+ tvUbTensor, srcUb[rowmax_idx * FLOAT_VECTOR_SIZE], numRowsRound, 0, 1, 1,
	274	+ numElemsAligned / FLOAT_BLOCK_SIZE);
	275	+ AscendC::PipeBarrier<PIPE_V>();
	276	+ AscendC::BlockReduceMax<float, false>(
	277	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
	278	+ 0, 1, 1, 8);
	279	+ AscendC::PipeBarrier<PIPE_V>();
	280	+ SetVecMask(numRowsRound);
	281	+ AscendC::Max<float, false>(
	282	+ rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
	283	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	284	+ AscendC::PipeBarrier<PIPE_V>();
	285	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	286	+ }
	287	+ }
	288	+ if (numElems % FLOAT_VECTOR_SIZE > 0) {
	289	+ SetVecMask(numElems % FLOAT_VECTOR_SIZE);
	290	+ AscendC::BlockReduceMax<float, false>(
	291	+ tvUbTensor, srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], numRowsRound, 0, 1, 1,
	292	+ numElemsAligned / FLOAT_BLOCK_SIZE);
	293	+ AscendC::PipeBarrier<PIPE_V>();
	294	+ SetBlockReduceMask(CeilDiv(numElems % FLOAT_VECTOR_SIZE, FLOAT_BLOCK_SIZE));
	295	+ if (numElems < FLOAT_VECTOR_SIZE) {
	296	+ AscendC::BlockReduceMax<float, false>(
	297	+ rowmaxUb, tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), 0, 1, 1, 8);
	298	+ AscendC::PipeBarrier<PIPE_V>();
	299	+ } else {
	300	+ AscendC::BlockReduceMax<float, false>(
	301	+ tvUbTensor[REDUCE_UB_SIZE], tvUbTensor, CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
	302	+ 0, 1, 1, 8);
	303	+ AscendC::PipeBarrier<PIPE_V>();
	304	+ SetVecMask(numRowsRound);
	305	+ AscendC::Max<float, false>(
	306	+ rowmaxUb, rowmaxUb, tvUbTensor[REDUCE_UB_SIZE], (uint64_t)0, 1,
	307	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	308	+ AscendC::PipeBarrier<PIPE_V>();
	309	+ }
	310	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	311	+ }
	312	+ }
	313	+
	314	+ __aicore__ inline void CopySGmToUb(
	315	+ AscendC::GlobalTensor<ElementInput> gInput, uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound,
	316	+ uint32_t columnNumPad)
	317	+ {
	318	+ AscendC::DataCopy(
	319	+ lsUbTensor[sUbOffset], gInput,
	320	+ AscendC::DataCopyParams(
	321	+ rowNumCurLoop, columnNumRound / FLOAT_BLOCK_SIZE, (columnNumPad - columnNumRound) / FLOAT_BLOCK_SIZE,
	322	+ 0));
	323	+ }
	324	+
	325	+ __aicore__ inline void CopyMaskGmToUb(
	326	+ AscendC::GlobalTensor<ElementMask> gMask, uint32_t columnNum, uint32_t columnNumRound, uint32_t maskStride,
	327	+ uint32_t tokenNumPerHead, uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t integralHeadNum,
	328	+ uint32_t epiTokenNum)
	329	+ {
	330	+ uint32_t innerUbRowOffset = 0;
	331	+ if (proTokenNum != 0) {
	332	+ AscendC::DataCopyPad(
	333	+ maskUbTensor[innerUbRowOffset], gMask[proTokenIdx * maskStride],
	334	+ AscendC::DataCopyExtParams(
	335	+ proTokenNum, columnNum * sizeof(ElementMask), (maskStride - columnNum) * sizeof(ElementMask), 0, 0),
	336	+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
	337	+ innerUbRowOffset += proTokenNum * columnNumRound;
	338	+ }
	339	+ for (uint32_t headIdx = 0; headIdx < integralHeadNum; headIdx++) {
	340	+ AscendC::DataCopyPad(
	341	+ maskUbTensor[innerUbRowOffset], gMask,
	342	+ AscendC::DataCopyExtParams(
	343	+ tokenNumPerHead, columnNum * sizeof(ElementMask), (maskStride - columnNum) * sizeof(ElementMask), 0,
	344	+ 0),
	345	+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
	346	+ innerUbRowOffset += tokenNumPerHead * columnNumRound;
	347	+ }
	348	+ if (epiTokenNum != 0) {
	349	+ AscendC::DataCopyPad(
	350	+ maskUbTensor[innerUbRowOffset], gMask,
	351	+ AscendC::DataCopyExtParams(
	352	+ epiTokenNum, columnNum * sizeof(ElementMask), (maskStride - columnNum) * sizeof(ElementMask), 0, 0),
	353	+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
	354	+ }
	355	+ }
	356	+
	357	+ __aicore__ inline void ScaleS(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound)
	358	+ {
	359	+ AscendC::Muls<float, false>(
	360	+ lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], scaleValue, (uint64_t)0,
	361	+ CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE), AscendC::UnaryRepeatParams(1, 1, 8, 8));
	362	+
	363	+ AscendC::PipeBarrier<PIPE_V>();
	364	+ }
	365	+
	366	+ template <typename ElementMaskDst, typename ElementMaskSrc>
	367	+ __aicore__ inline void UpCastMask(
	368	+ const AscendC::LocalTensor<ElementMaskDst>& maskUbTensorDst,
	369	+ const AscendC::LocalTensor<ElementMaskSrc>& maskUbTensorSrc, uint32_t rowNumCurLoop, uint32_t columnNumRound)
	370	+ {
	371	+ AscendC::Cast<ElementMaskDst, ElementMaskSrc, false>(
	372	+ maskUbTensorDst, maskUbTensorSrc, AscendC::RoundMode::CAST_NONE, (uint64_t)0,
	373	+ CeilDiv(rowNumCurLoop * columnNumRound, (uint32_t)(REPEAT_SIZE_IN_BYTE / sizeof(ElementMaskDst))),
	374	+ AscendC::UnaryRepeatParams(1, 1, 8, 4));
	375	+ AscendC::PipeBarrier<PIPE_V>();
	376	+ }
	377	+
	378	+ __aicore__ inline void ApplyMask(
	379	+ uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound, uint32_t maskColumnRound,
	380	+ uint32_t addMaskUbOffset)
	381	+ {
	382	+ AscendC::Muls<float, false>(
	383	+ maskUbTensor32, maskUbTensor32, (float)-3e38, (uint64_t)0,
	384	+ CeilDiv(rowNumCurLoop * maskColumnRound, FLOAT_VECTOR_SIZE), AscendC::UnaryRepeatParams(1, 1, 8, 8));
	385	+ AscendC::PipeBarrier<PIPE_V>();
	386	+ if (maskColumnRound == columnNumRound) {
	387	+ AscendC::Add<float, false>(
	388	+ lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], maskUbTensor32, (uint64_t)0,
	389	+ CeilDiv(rowNumCurLoop * maskColumnRound, FLOAT_VECTOR_SIZE),
	390	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	391	+ } else {
	392	+ uint32_t loop = maskColumnRound / FLOAT_VECTOR_SIZE;
	393	+ for (uint32_t i = 0; i < loop; i++) {
	394	+ AscendC::Add<float, false>(
	395	+ lsUbTensor[sUbOffset][addMaskUbOffset + i * FLOAT_VECTOR_SIZE],
	396	+ lsUbTensor[sUbOffset][addMaskUbOffset + i * FLOAT_VECTOR_SIZE],
	397	+ maskUbTensor32[i * FLOAT_VECTOR_SIZE], (uint64_t)0, rowNumCurLoop,
	398	+ AscendC::BinaryRepeatParams(
	399	+ 1, 1, 1, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE,
	400	+ maskColumnRound / FLOAT_BLOCK_SIZE));
	401	+ }
	402	+ if (maskColumnRound % FLOAT_VECTOR_SIZE > 0) {
	403	+ SetVecMask(maskColumnRound % FLOAT_VECTOR_SIZE);
	404	+ AscendC::Add<float, false>(
	405	+ lsUbTensor[sUbOffset][addMaskUbOffset + loop * FLOAT_VECTOR_SIZE],
	406	+ lsUbTensor[sUbOffset][addMaskUbOffset + loop * FLOAT_VECTOR_SIZE],
	407	+ maskUbTensor32[loop * FLOAT_VECTOR_SIZE], (uint64_t)0, rowNumCurLoop,
	408	+ AscendC::BinaryRepeatParams(
	409	+ 1, 1, 1, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE,
	410	+ maskColumnRound / FLOAT_BLOCK_SIZE));
	411	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	412	+ }
	413	+ }
	414	+ AscendC::PipeBarrier<PIPE_V>();
	415	+ }
	416	+
	417	+ __aicore__ inline void CalcLocalRowMax(
	418	+ uint32_t sUbOffset, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound,
	419	+ uint32_t rowOffset)
	420	+ {
	421	+ if (columnNum == 512) {
	422	+ RowmaxSPECTILE512(
	423	+ lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	424	+ columnNumRound);
	425	+ } else if (columnNum == 256) {
	426	+ RowmaxSPECTILE256(
	427	+ lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	428	+ columnNumRound);
	429	+ } else {
	430	+ RowmaxTAILTILE(
	431	+ lsUbTensor[sUbOffset], lmUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	432	+ columnNumRound);
	433	+ }
	434	+ }
	435	+
	436	+ __aicore__ inline void UpdateGlobalRowMax(
	437	+ uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound,
	438	+ uint32_t dmUbOffsetCurCycle, uint32_t rowOffset, uint32_t isFirstStackTile)
	439	+ {
	440	+ if (isFirstStackTile) {
	441	+ AscendC::DataCopy(
	442	+ hmUbTensor[rowOffset], lmUbTensor[rowOffset],
	443	+ AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0));
	444	+ AscendC::PipeBarrier<PIPE_V>();
	445	+ } else {
	446	+ SetVecMask(rowNumCurLoop);
	447	+ // *** hm = vmax(lm, gm)
	448	+ AscendC::Max<float, false>(
	449	+ hmUbTensor[rowOffset], lmUbTensor[rowOffset], gmUbTensor[rowOffset], (uint64_t)0, 1,
	450	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	451	+ AscendC::PipeBarrier<PIPE_V>();
	452	+ // *** dm = gm - hm
	453	+ AscendC::Sub<float, false>(
	454	+ dmUbTensor[dmUbOffsetCurCycle], gmUbTensor[rowOffset], hmUbTensor[rowOffset], (uint64_t)0, 1,
	455	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	456	+ AscendC::PipeBarrier<PIPE_V>();
	457	+ // *** dm = exp(dm)
	458	+ AscendC::Exp<float, false>(
	459	+ dmUbTensor[dmUbOffsetCurCycle], dmUbTensor[dmUbOffsetCurCycle], (uint64_t)0, 1,
	460	+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
	461	+ }
	462	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	463	+ AscendC::PipeBarrier<PIPE_V>();
	464	+ // *** gm = hm
	465	+ AscendC::DataCopy(
	466	+ gmUbTensor[rowOffset], hmUbTensor[rowOffset],
	467	+ AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0));
	468	+ AscendC::PipeBarrier<PIPE_V>();
	469	+ }
	470	+
	471	+ __aicore__ inline void CalcExp(
	472	+ uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum,
	473	+ uint32_t columnNumRound, uint32_t rowOffset)
	474	+ {
	475	+ // *** hm_block = expand_to_block(hm), 存放于 tv
	476	+ AscendC::Brcb(
	477	+ tvUbTensor.template ReinterpretCast<uint32_t>(), hmUbTensor[rowOffset].template ReinterpretCast<uint32_t>(),
	478	+ rowNumCurLoopRound / FLOAT_BLOCK_SIZE, AscendC::BrcbRepeatParams(1, 8));
	479	+ AscendC::PipeBarrier<PIPE_V>();
	480	+ // *** ls = ls - hm_block
	481	+ for (uint32_t subIdx = 0; subIdx < columnNum / FLOAT_VECTOR_SIZE; ++subIdx) {
	482	+ AscendC::Sub<float, false>(
	483	+ lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE],
	484	+ tvUbTensor, (uint64_t)0, rowNumCurLoop,
	485	+ AscendC::BinaryRepeatParams(
	486	+ 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1));
	487	+ }
	488	+ if (columnNum % FLOAT_VECTOR_SIZE > 0) {
	489	+ SetVecMask(columnNum % FLOAT_VECTOR_SIZE);
	490	+ AscendC::Sub<float, false>(
	491	+ lsUbTensor[sUbOffset][columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
	492	+ lsUbTensor[sUbOffset][columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0,
	493	+ rowNumCurLoop,
	494	+ AscendC::BinaryRepeatParams(
	495	+ 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1));
	496	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	497	+ }
	498	+ AscendC::PipeBarrier<PIPE_V>();
	499	+ // *** ls = exp(ls)
	500	+ AscendC::Exp<float, false>(
	501	+ lsUbTensor[sUbOffset], lsUbTensor[sUbOffset], (uint64_t)0,
	502	+ CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE), AscendC::UnaryRepeatParams(1, 1, 8, 8));
	503	+ AscendC::PipeBarrier<PIPE_V>();
	504	+ }
	505	+
	506	+ __aicore__ inline void CalcLocalRowSum(
	507	+ uint32_t sUbOffset, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound,
	508	+ uint32_t rowOffset)
	509	+ {
	510	+ // *** ll = rowsum(ls32)
	511	+ if (columnNum == 512) {
	512	+ RowsumSPECTILE512(
	513	+ lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	514	+ columnNumRound);
	515	+ } else if (columnNum == 256) {
	516	+ RowsumSPECTILE256(
	517	+ lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	518	+ columnNumRound);
	519	+ } else {
	520	+ RowsumTAILTILE(
	521	+ lsUbTensor[sUbOffset], llUbTensor[rowOffset], tvUbTensor, rowNumCurLoopRound, columnNum,
	522	+ columnNumRound);
	523	+ }
	524	+ }
	525	+
	526	+ __aicore__ inline void UpdateGlobalRowSum(
	527	+ uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t dmUbOffsetCurCycle,
	528	+ uint32_t rowOffset, uint32_t isFirstStackTile)
	529	+ {
	530	+ if (isFirstStackTile) {
	531	+ // *** gl = ll
	532	+ AscendC::DataCopy(
	533	+ glUbTensor[rowOffset], llUbTensor[rowOffset],
	534	+ AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0));
	535	+ AscendC::PipeBarrier<PIPE_V>();
	536	+ } else {
	537	+ SetVecMask(rowNumCurLoop);
	538	+ // *** gl = dm * gl
	539	+ AscendC::Mul<float, false>(
	540	+ glUbTensor[rowOffset], dmUbTensor[dmUbOffsetCurCycle], glUbTensor[rowOffset], (uint64_t)0, 1,
	541	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	542	+ AscendC::PipeBarrier<PIPE_V>();
	543	+ // *** gl = ll + gl
	544	+ AscendC::Add<float, false>(
	545	+ glUbTensor[rowOffset], glUbTensor[rowOffset], llUbTensor[rowOffset], (uint64_t)0, 1,
	546	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	547	+ AscendC::PipeBarrier<PIPE_V>();
	548	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	549	+ }
	550	+ }
	551	+
	552	+ __aicore__ inline void DownCastP(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound)
	553	+ {
	554	+ // *** lp = castfp32to16(ls)
	555	+ if (std::is_same<ElementOutput, bfloat16_t>::value) {
	556	+ AscendC::Cast<ElementOutput, float, false>(
	557	+ lpUbTensor[sUbOffset], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_RINT, (uint64_t)0,
	558	+ CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE), AscendC::UnaryRepeatParams(1, 1, 4, 8));
	559	+ } else {
	560	+ AscendC::Cast<ElementOutput, float, false>(
	561	+ lpUbTensor[sUbOffset], lsUbTensor[sUbOffset], AscendC::RoundMode::CAST_NONE, (uint64_t)0,
	562	+ CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE), AscendC::UnaryRepeatParams(1, 1, 4, 8));
	563	+ }
	564	+ }
	565	+
	566	+ __aicore__ inline void CopyPUbToGm(
	567	+ AscendC::GlobalTensor<ElementOutput> gOutput, uint32_t sUbOffset, uint32_t rowNumCurLoop,
	568	+ uint32_t columnNumRound, uint32_t columnNumPad)
	569	+ {
	570	+ AscendC::DataCopy(
	571	+ gOutput, lpUbTensor[sUbOffset],
	572	+ AscendC::DataCopyParams(
	573	+ rowNumCurLoop, columnNumRound / BLOCK_SIZE, 0, (columnNumPad - columnNumRound) / BLOCK_SIZE));
	574	+ }
	575	+
	576	+ template <bool doTriUMask>
	577	+ __aicore__ inline void SubCoreCompute(
	578	+ AscendC::GlobalTensor<ElementOutput> gOutput, const LayoutOutput& layoutOutput, uint32_t rowOffset,
	579	+ uint32_t isFirstStackTile, uint32_t isLastNoMaskStackTile, uint32_t isFirstRowLoop, uint32_t isLastRowLoop,
	580	+ uint32_t columnNumRound, uint32_t pingpongFlag, uint32_t curStackTileMod)
	581	+ {
	582	+ uint32_t rowNumCurLoop = layoutOutput.shape(0);
	583	+ uint32_t rowNumCurLoopRound = RoundUp(rowNumCurLoop, FLOAT_BLOCK_SIZE);
	584	+ uint32_t columnNum = layoutOutput.shape(1);
	585	+ uint32_t columnNumPad = layoutOutput.stride(0);
	586	+ uint32_t sUbOffset = pingpongFlag * MAX_UB_S_ELEM_NUM;
	587	+ uint32_t dmUbOffsetCurCycle = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffset;
	588	+
	589	+ if constexpr (LSE_MODE_ == LseMode::LSE_OUT) {
	590	+ // In lse out-only mode, tv is used in the last stack tile to transport lse
	591	+ if (isFirstStackTile && isFirstRowLoop) {
	592	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
	593	+ }
	594	+ }
	595	+ CalcLocalRowMax(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
	596	+ UpdateGlobalRowMax(
	597	+ rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, dmUbOffsetCurCycle, rowOffset,
	598	+ isFirstStackTile);
	599	+
	600	+ CalcExp(sUbOffset, rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
	601	+ if constexpr (!doTriUMask) {
	602	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(pingpongFlag);
	603	+ }
	604	+
	605	+ DownCastP(sUbOffset, rowNumCurLoop, columnNumRound);
	606	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
	607	+
	608	+ CalcLocalRowSum(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
	609	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(pingpongFlag);
	610	+
	611	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
	612	+ CopyPUbToGm(gOutput, sUbOffset, rowNumCurLoop, columnNumRound, columnNumPad);
	613	+ if constexpr (!doTriUMask) {
	614	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(pingpongFlag);
	615	+ if (isLastNoMaskStackTile && isLastRowLoop) {
	616	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
	617	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
	618	+ }
	619	+ } else {
	620	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
	621	+ }
	622	+ UpdateGlobalRowSum(
	623	+ sUbOffset, rowNumCurLoop, rowNumCurLoopRound, dmUbOffsetCurCycle, rowOffset, isFirstStackTile);
	624	+ }
	625	+
	626	+ __aicore__ inline void operator()(
	627	+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
	628	+ const LayoutOutput& layoutOutput, const LayoutInput& layoutInput, GemmCoord actualBlockShape,
	629	+ uint32_t isFirstStackTile, uint32_t isLastNoMaskStackTile, uint32_t qSBlockSize, uint32_t qNBlockSize,
	630	+ uint32_t curStackTileMod)
	631	+ {
	632	+ uint32_t rowNum = actualBlockShape.m();
	633	+ uint32_t columnNum = actualBlockShape.n();
	634	+ uint32_t columnNumRound = RoundUp(columnNum, BLOCK_SIZE);
	635	+ uint32_t columnNumPad = layoutInput.stride(0);
	636	+
	637	+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
	638	+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
	639	+
	640	+ uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum;
	641	+ uint32_t rowSplitSubBlock = (qNBlockSize == 1) ? (qSBlockSize / 2) : (qSBlockSize * qNSplitSubBlock);
	642	+ uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? (rowNum - rowSplitSubBlock) : rowSplitSubBlock;
	643	+ uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock;
	644	+ uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound;
	645	+ uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE);
	646	+ rowNumTile = AscendC::Std::min(rowNumTile, FLOAT_VECTOR_SIZE);
	647	+ uint32_t rowLoopNum = CeilDiv(rowActualThisSubBlock, rowNumTile);
	648	+ uint32_t preLoad = 1;
	649	+
	650	+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) {
	651	+ if (rowLoopIdx < rowLoopNum) {
	652	+ uint32_t pingpongFlag = rowLoopIdx % 2;
	653	+ uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile;
	654	+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
	655	+ uint32_t rowNumCurLoop =
	656	+ (rowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
	657	+
	658	+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
	659	+ auto gInputCurLoop = gInput[offsetInput];
	660	+
	661	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(pingpongFlag);
	662	+ CopySGmToUb(
	663	+ gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad);
	664	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
	665	+ }
	666	+ if (rowLoopIdx >= preLoad) {
	667	+ uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad;
	668	+ uint32_t pingpongFlag = delayedRowLoopIdx % 2;
	669	+ uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile;
	670	+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
	671	+ uint32_t rowNumCurLoop =
	672	+ (delayedRowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
	673	+
	674	+ int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
	675	+ auto gOutputCurLoop = gOutput[offsetOutput];
	676	+ auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum));
	677	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
	678	+ ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound);
	679	+ SubCoreCompute<false>(
	680	+ gOutputCurLoop, layoutOutputCurLoop, rowOffsetCurLoop, isFirstStackTile, isLastNoMaskStackTile,
	681	+ delayedRowLoopIdx == 0, delayedRowLoopIdx == rowLoopNum - 1, columnNumRound, pingpongFlag,
	682	+ curStackTileMod);
	683	+ }
	684	+ }
	685	+ }
	686	+
	687	+ __aicore__ inline void operator()(
	688	+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
	689	+ AscendC::GlobalTensor<ElementMask> gMask, const LayoutOutput& layoutOutput, const LayoutInput& layoutInput,
	690	+ const LayoutInput& layoutMask, GemmCoord actualBlockShape, uint32_t isFirstStackTile, uint32_t qSBlockSize,
	691	+ uint32_t qNBlockSize, uint32_t curStackTileMod, Arch::CrossCoreFlag qkReady, uint32_t triUp, uint32_t triDown,
	692	+ uint32_t kvSStartIdx, uint32_t kvSEndIdx)
	693	+ {
	694	+ uint32_t rowNum = actualBlockShape.m();
	695	+ uint32_t columnNum = actualBlockShape.n();
	696	+ uint32_t columnNumRound = RoundUp(columnNum, BLOCK_SIZE_IN_BYTE);
	697	+ uint32_t columnNumPad = layoutInput.stride(0);
	698	+ uint32_t maskStride = layoutMask.stride(0);
	699	+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
	700	+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
	701	+
	702	+ uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum;
	703	+ uint32_t rowSplitSubBlock = (qNBlockSize == 1) ? (qSBlockSize / 2) : (qSBlockSize * qNSplitSubBlock);
	704	+ uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? (rowNum - rowSplitSubBlock) : rowSplitSubBlock;
	705	+ uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock;
	706	+
	707	+ uint32_t tokenNumPerHeadThisSubBlock = Min(qSBlockSize, rowActualThisSubBlock);
	708	+ uint32_t maskOffsetThisSubBlock = (qNBlockSize == 1) ? rowOffsetThisSubBlock : 0;
	709	+
	710	+ // calc mask shift in gm
	711	+ uint32_t gmOffsetMaskRow;
	712	+ uint32_t gmOffsetMaskColumn;
	713	+ uint32_t maskColumn;
	714	+ uint32_t addMaskUbOffset;
	715	+ if (triUp >= kvSStartIdx) {
	716	+ uint32_t triUpRoundDown = RoundDown(triUp, BLOCK_SIZE_IN_BYTE);
	717	+ gmOffsetMaskRow = triUp - triUpRoundDown;
	718	+ gmOffsetMaskColumn = 0;
	719	+ maskColumn = kvSEndIdx - triUpRoundDown;
	720	+ addMaskUbOffset = triUpRoundDown - kvSStartIdx;
	721	+ } else {
	722	+ gmOffsetMaskRow = 0;
	723	+ gmOffsetMaskColumn = kvSStartIdx - triUp;
	724	+ maskColumn = columnNum;
	725	+ addMaskUbOffset = 0;
	726	+ }
	727	+ uint32_t maskColumnRound = RoundUp(maskColumn, BLOCK_SIZE_IN_BYTE);
	728	+
	729	+ int64_t offsetMask =
	730	+ layoutMask.GetOffset(MatrixCoord(gmOffsetMaskRow + maskOffsetThisSubBlock, gmOffsetMaskColumn));
	731	+ auto gMaskThisSubBlock = gMask[offsetMask];
	732	+
	733	+ uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound;
	734	+ uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE);
	735	+ rowNumTile = AscendC::Std::min(rowNumTile, FLOAT_VECTOR_SIZE);
	736	+ uint32_t rowLoopNum = CeilDiv(rowActualThisSubBlock, rowNumTile);
	737	+ uint32_t preLoad = 1;
	738	+
	739	+ if (rowActualThisSubBlock == 0) {
	740	+ Arch::CrossCoreWaitFlag(qkReady);
	741	+ return;
	742	+ }
	743	+
	744	+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) {
	745	+ if (rowLoopIdx < rowLoopNum) {
	746	+ uint32_t pingpongFlag = rowLoopIdx % 2;
	747	+ uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile;
	748	+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
	749	+ uint32_t rowNumCurLoop =
	750	+ (rowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
	751	+ // loop 0 mask load before cross core sync
	752	+ if (rowLoopIdx == 0) {
	753	+ // the token idx of the start token of the prologue part
	754	+ uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock;
	755	+ // the token num of the prologue part
	756	+ uint32_t proTokenNum =
	757	+ Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % tokenNumPerHeadThisSubBlock;
	758	+ // the token num of the epilogue part
	759	+ uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock;
	760	+ // the number of integral heads within a cycle
	761	+ uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock;
	762	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
	763	+ CopyMaskGmToUb(
	764	+ gMaskThisSubBlock, maskColumn, maskColumnRound, maskStride, tokenNumPerHeadThisSubBlock,
	765	+ proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum);
	766	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
	767	+ Arch::CrossCoreWaitFlag(qkReady);
	768	+ }
	769	+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
	770	+ auto gInputCurLoop = gInput[offsetInput];
	771	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(pingpongFlag);
	772	+ CopySGmToUb(
	773	+ gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad);
	774	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
	775	+ }
	776	+ if (rowLoopIdx >= preLoad) {
	777	+ uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad;
	778	+ uint32_t pingpongFlag = delayedRowLoopIdx % 2;
	779	+ uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile;
	780	+ uint32_t rowNumCurLoop =
	781	+ (delayedRowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
	782	+
	783	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
	784	+ UpCastMask<half, ElementMask>(maskUbTensor16, maskUbTensor, rowNumCurLoop, columnNumRound);
	785	+ UpCastMask<float, half>(maskUbTensor32, maskUbTensor16, rowNumCurLoop, columnNumRound);
	786	+
	787	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
	788	+ ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound);
	789	+ ApplyMask(
	790	+ (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, maskColumnRound,
	791	+ addMaskUbOffset);
	792	+ // next loop mask load
	793	+ if (rowLoopIdx < rowLoopNum) {
	794	+ uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile;
	795	+ uint32_t rowNumCurLoop =
	796	+ (rowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
	797	+ // the token idx of the start token of the prologue part
	798	+ uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock;
	799	+ // the token num of the prologue part
	800	+ uint32_t proTokenNum =
	801	+ Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % tokenNumPerHeadThisSubBlock;
	802	+ // the number of integral heads within a cycle
	803	+ uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock;
	804	+ // the token num of the epilogue part
	805	+ uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock;
	806	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
	807	+ CopyMaskGmToUb(
	808	+ gMaskThisSubBlock, maskColumn, maskColumnRound, maskStride, tokenNumPerHeadThisSubBlock,
	809	+ proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum);
	810	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
	811	+ }
	812	+ // online softmax vectorized compute
	813	+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
	814	+ int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
	815	+ auto gOutputCurLoop = gOutput[offsetOutput];
	816	+ auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum));
	817	+ SubCoreCompute<true>(
	818	+ gOutputCurLoop, layoutOutputCurLoop, rowOffsetCurLoop, isFirstStackTile, 0, delayedRowLoopIdx == 0,
	819	+ delayedRowLoopIdx == rowLoopNum - 1, columnNumRound, pingpongFlag, curStackTileMod);
	820	+ }
	821	+ }
	822	+ }
	823	+
	824	+private:
	825	+ float scaleValue;
	826	+ AscendC::LocalTensor<float> lsUbTensor;
	827	+ AscendC::LocalTensor<ElementOutput> lpUbTensor;
	828	+ AscendC::LocalTensor<ElementMask> maskUbTensor;
	829	+ AscendC::LocalTensor<half> maskUbTensor16;
	830	+ AscendC::LocalTensor<float> maskUbTensor32;
	831	+ AscendC::LocalTensor<float> lmUbTensor;
	832	+ AscendC::LocalTensor<float> hmUbTensor;
	833	+ AscendC::LocalTensor<float> gmUbTensor;
	834	+ AscendC::LocalTensor<float> dmUbTensor;
	835	+ AscendC::LocalTensor<float> llUbTensor;
	836	+ AscendC::LocalTensor<float> tvUbTensor;
	837	+ AscendC::LocalTensor<float> glUbTensor;
	838	+};
	839	+} // namespace Catlass::Epilogue::Block
	840	+
	841	+#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_HPP

Ainclude/catlass/epilogue/block/block_epilogue_xfai_rescale_o.hpp+380-0

@@ -0,0 +1,380 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_XFAI_RESCALE_O_HPP
	13	+#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_XFAI_RESCALE_O_HPP
	14	+
	15	+#include "catlass/arch/resource.hpp"
	16	+#include "catlass/epilogue/dispatch_policy.hpp"
	17	+#include "catlass/epilogue/tile/tile_copy.hpp"
	18	+#include "catlass/gemm_coord.hpp"
	19	+#include "catlass/matrix_coord.hpp"
	20	+
	21	+namespace Catlass::Epilogue::Block {
	22	+
	23	+template <class OutputType_, class InputType_, class UpdateType_, class LseType_, LseMode LSE_MODE_>
	24	+class BlockEpilogue<EpilogueAtlasA2XFAIRescaleO<LSE_MODE_>, OutputType_, InputType_, UpdateType_, LseType_> {
	25	+public:
	26	+ // Type aliases
	27	+ using DispatchPolicy = EpilogueAtlasA2XFAIRescaleO<LSE_MODE_>;
	28	+ using ArchTag = typename DispatchPolicy::ArchTag;
	29	+
	30	+ using ElementOutput = typename OutputType_::Element;
	31	+ using ElementInput = typename InputType_::Element;
	32	+ using ElementUpdate = typename UpdateType_::Element;
	33	+ using ElementLse = typename LseType_::Element;
	34	+
	35	+ using LayoutOutput = typename OutputType_::Layout;
	36	+ using LayoutInput = typename InputType_::Layout;
	37	+ using LayoutUpdate = typename UpdateType_::Layout;
	38	+ using LayoutLse = typename LseType_::Layout;
	39	+
	40	+ static constexpr LseMode LSE_MODE = DispatchPolicy::LSE_MODE;
	41	+
	42	+ static constexpr uint32_t HALF_ELENUM_PER_BLK = 16;
	43	+ static constexpr uint32_t BLOCK_SIZE = 16;
	44	+ static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128;
	45	+ static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64;
	46	+ static constexpr uint32_t HALF_ELENUM_PER_LINE = 256;
	47	+ static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128;
	48	+ static constexpr uint32_t MULTIPLIER = 2;
	49	+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
	50	+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
	51	+ static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024;
	52	+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384;
	53	+ static constexpr uint32_t HALF_DM_UB_SIZE = 64;
	54	+ static constexpr uint32_t HALF_LL_UB_SIZE = 256;
	55	+ static constexpr uint32_t VECTOR_SIZE = 128;
	56	+ static constexpr uint32_t NUM4 = 4;
	57	+ static constexpr uint32_t MAX_UB_O_ELEM_NUM = 8192;
	58	+ static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 256;
	59	+
	60	+ __aicore__ inline BlockEpilogue(Arch::Resource<ArchTag>& resource)
	61	+ {
	62	+ // Allocate UB space
	63	+ constexpr uint32_t LO_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE;
	64	+ constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE;
	65	+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE;
	66	+
	67	+ constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE;
	68	+ constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE;
	69	+ constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE;
	70	+ constexpr uint32_t LSE_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE;
	71	+ constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE;
	72	+
	73	+ loUbTensor = resource.ubBuf.template GetBufferByByte<float>(LO_UB_TENSOR_OFFSET);
	74	+ dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
	75	+ glUbTensor = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
	76	+ tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
	77	+ goUbTensor16 = resource.ubBuf.template GetBufferByByte<ElementOutput>(GO_UB_TENSOR_OFFSET);
	78	+ goUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(GO_UB_TENSOR_OFFSET);
	79	+ hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);
	80	+ gmUbTensor = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
	81	+ lse32_ubuf_tensor = resource.ubBuf.template GetBufferByByte<float>(LSE_UB_TENSOR_OFFSET);
	82	+ }
	83	+
	84	+ __aicore__ inline ~BlockEpilogue()
	85	+ {}
	86	+
	87	+ __aicore__ inline void SetMask(int32_t len)
	88	+ {
	89	+ uint64_t mask = 0;
	90	+ uint64_t one = 1;
	91	+ uint64_t temp = len % FLOAT_VECTOR_SIZE;
	92	+ for (int64_t i = 0; i < temp; i++) {
	93	+ mask \|= one << i;
	94	+ }
	95	+
	96	+ if (len == VECTOR_SIZE) {
	97	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	98	+ } else if (len >= FLOAT_VECTOR_SIZE) {
	99	+ AscendC::SetVectorMask<int8_t>(mask, (uint64_t)-1);
	100	+ } else {
	101	+ AscendC::SetVectorMask<int8_t>(0x0, mask);
	102	+ }
	103	+ }
	104	+
	105	+ __aicore__ inline void CopyOToGm(
	106	+ AscendC::GlobalTensor<ElementOutput> gOutput, uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t epiTokenNum,
	107	+ uint32_t integralHeadNum, uint32_t qSThisSubBlock, uint32_t embed, uint32_t oHiddenSize)
	108	+ {
	109	+ uint32_t innerOGmOffset = 0;
	110	+ uint32_t innerGOUbOffset = 0;
	111	+ if (proTokenNum != 0) {
	112	+ AscendC::DataCopyPad(
	113	+ gOutput[innerOGmOffset + proTokenIdx * oHiddenSize], goUbTensor16[innerGOUbOffset],
	114	+ AscendC::DataCopyExtParams(proTokenNum, embed * 2, 0, (oHiddenSize - embed) * 2, 0));
	115	+ innerOGmOffset += embed;
	116	+ innerGOUbOffset += proTokenNum * embed;
	117	+ }
	118	+ for (uint32_t qN_idx = 0; qN_idx < integralHeadNum; qN_idx++) {
	119	+ AscendC::DataCopyPad(
	120	+ gOutput[innerOGmOffset], goUbTensor16[innerGOUbOffset],
	121	+ AscendC::DataCopyExtParams(qSThisSubBlock, embed * 2, 0, (oHiddenSize - embed) * 2, 0));
	122	+ innerOGmOffset += embed;
	123	+ innerGOUbOffset += qSThisSubBlock * embed;
	124	+ }
	125	+ if (epiTokenNum != 0) {
	126	+ AscendC::DataCopyPad(
	127	+ gOutput[innerOGmOffset], goUbTensor16[innerGOUbOffset],
	128	+ AscendC::DataCopyExtParams(epiTokenNum, embed * 2, 0, (oHiddenSize - embed) * 2, 0));
	129	+ }
	130	+ }
	131	+
	132	+ __aicore__ inline void SubCoreCompute(
	133	+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
	134	+ AscendC::GlobalTensor<ElementUpdate> gUpdate, AscendC::GlobalTensor<ElementLse> gLse,
	135	+ const LayoutOutput& layoutOutput, const LayoutInput& layoutInput, const LayoutUpdate& layoutUpdate,
	136	+ const LayoutLse& layoutLse, uint32_t qNThisSubBlock, uint32_t isFirstStackTile, uint32_t isLastStackTile,
	137	+ uint32_t curStackTileMod, uint32_t needRowLoop, uint32_t isLastRowLoop, uint32_t rowOffsetLoop,
	138	+ uint32_t qSThisSubBlock, uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t epiTokenNum,
	139	+ uint32_t integralHeadNum)
	140	+ {
	141	+ uint32_t curRowNum = layoutInput.shape(0);
	142	+ uint32_t embed = layoutInput.shape(1);
	143	+ uint32_t embedRound = layoutInput.stride(0);
	144	+ uint32_t curRowNumRound = RoundUp(curRowNum, FLOAT_BLOCK_SIZE);
	145	+ uint32_t qSBlockSize = layoutOutput.shape(0);
	146	+ uint32_t oHiddenSize = layoutOutput.shape(1);
	147	+ uint32_t qHeads = layoutLse.shape(1);
	148	+ uint32_t dmUbOffsetCurStackTile = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffsetLoop;
	149	+
	150	+ if (!isFirstStackTile) {
	151	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
	152	+ AscendC::DataCopy(
	153	+ loUbTensor, gInput, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
	154	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	155	+ }
	156	+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID6);
	157	+ if (!isFirstStackTile) {
	158	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	159	+ AscendC::Brcb(
	160	+ tvUbTensor.ReinterpretCast<uint32_t>(), dmUbTensor[dmUbOffsetCurStackTile].ReinterpretCast<uint32_t>(),
	161	+ curRowNumRound / FLOAT_BLOCK_SIZE, AscendC::BrcbRepeatParams(1, 8));
	162	+ AscendC::PipeBarrier<PIPE_V>();
	163	+ if (needRowLoop) {
	164	+ AscendC::DataCopy(
	165	+ goUbTensor32, gUpdate, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
	166	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1);
	167	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1);
	168	+ }
	169	+ // *** go = go * dm_block
	170	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	171	+ for (uint32_t vmul_idx = 0; vmul_idx < embed / FLOAT_VECTOR_SIZE; ++vmul_idx) {
	172	+ AscendC::Mul<float, false>(
	173	+ goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE], goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE], tvUbTensor,
	174	+ (uint64_t)0, curRowNum,
	175	+ AscendC::BinaryRepeatParams(
	176	+ 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1));
	177	+ }
	178	+ if (embed % FLOAT_VECTOR_SIZE > 0) {
	179	+ SetMask(embed % FLOAT_VECTOR_SIZE);
	180	+ AscendC::Mul<float, false>(
	181	+ goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
	182	+ goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0, curRowNum,
	183	+ AscendC::BinaryRepeatParams(
	184	+ 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1));
	185	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	186	+ }
	187	+ AscendC::PipeBarrier<PIPE_V>();
	188	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	189	+ // *** go = lo + go
	190	+ AscendC::Add<float, false>(
	191	+ goUbTensor32, goUbTensor32, loUbTensor, (uint64_t)0,
	192	+ (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	193	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	194	+ AscendC::PipeBarrier<PIPE_V>();
	195	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
	196	+ } else {
	197	+ // *** go = lo
	198	+ AscendC::DataCopy(
	199	+ goUbTensor32, gInput, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
	200	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	201	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
	202	+ }
	203	+
	204	+ if (isLastStackTile) {
	205	+ // *** gl_block = expand_to_block(gl), 存放于 tv
	206	+ AscendC::Brcb(
	207	+ tvUbTensor.ReinterpretCast<uint32_t>(), glUbTensor.ReinterpretCast<uint32_t>()[rowOffsetLoop],
	208	+ curRowNumRound / FLOAT_BLOCK_SIZE, AscendC::BrcbRepeatParams(1, 8));
	209	+ AscendC::PipeBarrier<PIPE_V>();
	210	+ // *** go = go / gl_block
	211	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	212	+ for (uint32_t vdiv_idx = 0; vdiv_idx < embed / FLOAT_VECTOR_SIZE; ++vdiv_idx) {
	213	+ AscendC::Div<float, false>(
	214	+ goUbTensor32[vdiv_idx * FLOAT_VECTOR_SIZE], goUbTensor32[vdiv_idx * FLOAT_VECTOR_SIZE], tvUbTensor,
	215	+ (uint64_t)0, curRowNum,
	216	+ AscendC::BinaryRepeatParams(
	217	+ 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1));
	218	+ }
	219	+ if (embed % FLOAT_VECTOR_SIZE > 0) {
	220	+ SetMask(embed % FLOAT_VECTOR_SIZE);
	221	+ AscendC::Div<float, false>(
	222	+ goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
	223	+ goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], tvUbTensor, (uint64_t)0, curRowNum,
	224	+ AscendC::BinaryRepeatParams(
	225	+ 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1));
	226	+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
	227	+ }
	228	+ AscendC::PipeBarrier<PIPE_V>();
	229	+
	230	+ // *** go = castfp32to16(go)
	231	+ if (std::is_same<ElementOutput, bfloat16_t>::value) {
	232	+ AscendC::Cast<ElementOutput, float, false>(
	233	+ goUbTensor16, goUbTensor32, AscendC::RoundMode::CAST_RINT, (uint64_t)0,
	234	+ (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	235	+ AscendC::UnaryRepeatParams(1, 1, 4, 8));
	236	+ } else {
	237	+ AscendC::Cast<ElementOutput, float, false>(
	238	+ goUbTensor16, goUbTensor32, AscendC::RoundMode::CAST_NONE, (uint64_t)0,
	239	+ (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
	240	+ AscendC::UnaryRepeatParams(1, 1, 4, 8));
	241	+ }
	242	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
	243	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
	244	+
	245	+ // ***move O to GM
	246	+ CopyOToGm(
	247	+ gOutput, proTokenIdx, proTokenNum, epiTokenNum, integralHeadNum, qSThisSubBlock, embed, oHiddenSize);
	248	+ if constexpr (LSE_MODE_ == LseMode::LSE_OUT) {
	249	+ if (isLastRowLoop) {
	250	+ AscendC::PipeBarrier<PIPE_V>();
	251	+ uint32_t lenBrust = sizeof(float);
	252	+ AscendC::Ln<float, false>(
	253	+ lse32_ubuf_tensor, glUbTensor, (uint64_t)0, 1, AscendC::UnaryRepeatParams(1, 1, 8, 8));
	254	+
	255	+ AscendC::PipeBarrier<PIPE_V>();
	256	+ AscendC::Add<float, false>(
	257	+ lse32_ubuf_tensor, lse32_ubuf_tensor, gmUbTensor, (uint64_t)0, 1,
	258	+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
	259	+ AscendC::PipeBarrier<PIPE_V>();
	260	+
	261	+ // *** lse_block = expand_to_block(lse), 存放于 tv
	262	+ AscendC::Brcb(
	263	+ tvUbTensor.ReinterpretCast<uint32_t>(), lse32_ubuf_tensor.ReinterpretCast<uint32_t>(),
	264	+ curRowNumRound / FLOAT_BLOCK_SIZE, AscendC::BrcbRepeatParams(1, 8));
	265	+ AscendC::PipeBarrier<PIPE_V>();
	266	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID4);
	267	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID4);
	268	+
	269	+ if (qNThisSubBlock == 0) {
	270	+ AscendC::DataCopyPad(
	271	+ gLse, tvUbTensor,
	272	+ AscendC::DataCopyExtParams(curRowNum, lenBrust, 0, (qHeads - 1) * lenBrust, 0));
	273	+ } else {
	274	+ for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) {
	275	+ AscendC::DataCopyPad(
	276	+ gLse[qNIdx], tvUbTensor[qNIdx * qSBlockSize * FLOAT_BLOCK_SIZE],
	277	+ AscendC::DataCopyExtParams(qSBlockSize, lenBrust, 0, (qHeads - 1) * lenBrust, 0));
	278	+ }
	279	+ }
	280	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
	281	+ }
	282	+ }
	283	+ } else if (needRowLoop) {
	284	+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID5);
	285	+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID5);
	286	+ AscendC::DataCopy(
	287	+ gUpdate, goUbTensor32, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
	288	+ }
	289	+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID6);
	290	+ }
	291	+
	292	+ __aicore__ inline void operator()(
	293	+ AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
	294	+ AscendC::GlobalTensor<ElementUpdate> gUpdate, AscendC::GlobalTensor<ElementLse> gLse,
	295	+ const LayoutOutput& layoutOutput, const LayoutInput& layoutInput, const LayoutUpdate& layoutUpdate,
	296	+ const LayoutLse& layoutLse, GemmCoord actualBlockShape, uint32_t qSBlockSize, uint32_t qNBlockSize,
	297	+ uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod)
	298	+ {
	299	+ uint32_t rowNum = actualBlockShape.m();
	300	+ uint32_t embed = actualBlockShape.n();
	301	+ uint32_t rowTile = MAX_UB_O_ELEM_NUM / embed;
	302	+
	303	+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
	304	+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
	305	+
	306	+ uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum;
	307	+ uint32_t qNThisSubBlock = (qNBlockSize == 1) ? 0 :
	308	+ (subBlockIdx == 1) ? (qNBlockSize - qNSplitSubBlock) :
	309	+ qNSplitSubBlock;
	310	+ uint32_t inRowSplitSubBlock =
	311	+ (qNBlockSize == 1) ? (qSBlockSize / subBlockNum) : (qSBlockSize * qNSplitSubBlock);
	312	+ uint32_t inRowActualThisSubBlock = (subBlockIdx == 1) ? (rowNum - inRowSplitSubBlock) : inRowSplitSubBlock;
	313	+ uint32_t inRowOffsetThisSubBlock = subBlockIdx * inRowSplitSubBlock;
	314	+ uint32_t outRowOffsetThisSubBlock = (qNBlockSize == 1) ? inRowOffsetThisSubBlock : 0;
	315	+ uint32_t outColOffsetThisSubBlock = (qNBlockSize == 1) ? 0 : subBlockIdx * qNSplitSubBlock * embed;
	316	+ uint32_t qSThisSubBlock = (qNBlockSize == 1) ? inRowActualThisSubBlock : qSBlockSize;
	317	+ int64_t outOffsetSubBlock =
	318	+ layoutOutput.GetOffset(MatrixCoord(outRowOffsetThisSubBlock, outColOffsetThisSubBlock));
	319	+
	320	+ uint32_t outLseRowOffsetThisSubBlock = (qNBlockSize == 1) ? inRowOffsetThisSubBlock : 0;
	321	+ uint32_t outLseColOffsetThisSubBlock = (qNBlockSize == 1) ? 0 : subBlockIdx * qNSplitSubBlock;
	322	+ int64_t offsetLse = layoutLse.GetOffset(MatrixCoord(outLseRowOffsetThisSubBlock, outLseColOffsetThisSubBlock));
	323	+ auto gLseThisSubBlock = gLse[offsetLse];
	324	+ auto layoutOutLseThisSubBlock = layoutLse;
	325	+
	326	+ if (inRowActualThisSubBlock > 0) {
	327	+ uint32_t rowLoop = CeilDiv(inRowActualThisSubBlock, rowTile);
	328	+ uint32_t needRowLoop = (rowLoop > 1) ? 1 : 0;
	329	+
	330	+ // The rows of each cycle consist of multiple heads with several tokens.
	331	+ // There are several integral heads, one prologue head, one epilogue head.
	332	+ uint32_t proTokenIdx = 0; // the token idx of the start token of the prologue part
	333	+ uint32_t proTokenNum = 0; // the token num of the prologue part
	334	+ uint32_t epiTokenNum = 0; // the token num of the epilogue part
	335	+ uint32_t integralHeadNum = 0; // the number of integral heads within a cycle
	336	+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoop; rowLoopIdx++) {
	337	+ uint32_t rowOffsetLoop = rowLoopIdx * rowTile;
	338	+ uint32_t rowOffsetCurLoop = inRowOffsetThisSubBlock + rowOffsetLoop;
	339	+ uint32_t rowActualCurLoop =
	340	+ (rowLoopIdx == (rowLoop - 1)) ? inRowActualThisSubBlock - rowLoopIdx * rowTile : rowTile;
	341	+
	342	+ int64_t offsetOutput = rowLoopIdx * rowTile / qSThisSubBlock * embed + outOffsetSubBlock;
	343	+ auto gOutputCurLoop = gOutput[offsetOutput];
	344	+ auto layoutOutputCurLoop = layoutOutput;
	345	+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetCurLoop, 0));
	346	+ auto gInputCurLoop = gInput[offsetInput];
	347	+ auto layoutInputCurLoop = layoutInput.GetTileLayout(MatrixCoord(rowActualCurLoop, embed));
	348	+
	349	+ int64_t offsetUpdate = layoutUpdate.GetOffset(MatrixCoord(rowOffsetCurLoop, 0));
	350	+ auto gUpdateCurLoop = gUpdate[offsetUpdate];
	351	+ auto layoutUpdateCurLoop = layoutUpdate.GetTileLayout(MatrixCoord(rowActualCurLoop, embed));
	352	+
	353	+ proTokenIdx = rowOffsetLoop % qSThisSubBlock;
	354	+ proTokenNum = AscendC::Std::min(rowActualCurLoop, (qSThisSubBlock - proTokenIdx)) % qSThisSubBlock;
	355	+ integralHeadNum = (rowActualCurLoop - proTokenNum) / qSThisSubBlock;
	356	+ epiTokenNum = rowActualCurLoop - proTokenNum - integralHeadNum * qSThisSubBlock;
	357	+
	358	+ SubCoreCompute(
	359	+ gOutputCurLoop, gInputCurLoop, gUpdateCurLoop, gLseThisSubBlock, layoutOutputCurLoop,
	360	+ layoutInputCurLoop, layoutUpdateCurLoop, layoutOutLseThisSubBlock, qNThisSubBlock, isFirstStackTile,
	361	+ isLastStackTile, curStackTileMod, needRowLoop, (rowLoopIdx == rowLoop - 1), rowOffsetLoop,
	362	+ qSThisSubBlock, proTokenIdx, proTokenNum, epiTokenNum, integralHeadNum);
	363	+ }
	364	+ }
	365	+ }
	366	+
	367	+private:
	368	+ AscendC::LocalTensor<float> loUbTensor;
	369	+ AscendC::LocalTensor<float> dmUbTensor;
	370	+ AscendC::LocalTensor<float> hmUbTensor;
	371	+ AscendC::LocalTensor<float> glUbTensor;
	372	+ AscendC::LocalTensor<float> tvUbTensor;
	373	+ AscendC::LocalTensor<ElementOutput> goUbTensor16;
	374	+ AscendC::LocalTensor<float> goUbTensor32;
	375	+ AscendC::LocalTensor<float> gmUbTensor;
	376	+ AscendC::LocalTensor<float> lse32_ubuf_tensor;
	377	+};
	378	+} // namespace Catlass::Epilogue::Block
	379	+
	380	+#endif // CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_XFAI_RESCALE_O_HPP

Minclude/catlass/epilogue/dispatch_policy.hpp+85-0

	@@ -16,6 +16,40 @@
16		16
17	namespace Catlass::Epilogue {	17	namespace Catlass::Epilogue {
18		18
		19	+// For AtlasA2, FA Infer RescaleO without div rowsum
		20	+struct EpilogueAtlasA2RescaleOWithoutDivSum {
		21	+ using ArchTag = Arch::AtlasA2;
		22	+};
		23	+
		24	+// For AtlasA2, FA Infer OnlineSoftmaxCopySumMax
		25	+struct EpilogueAtlasA2OnlineSoftmaxCopySumMax {
		26	+ using ArchTag = Arch::AtlasA2;
		27	+};
		28	+
		29	+// For Ascend950, FA Infer OnlineSoftmaxCopySumMax
		30	+struct EpilogueAscend950OnlineSoftmaxCopySumMax {
		31	+ using ArchTag = Arch::Ascend950;
		32	+};
		33	+
		34	+// For AtlasA2, FA Unshared Softmax
		35	+struct EpilogueAtlasA2FAUnsharedSoftmax {
		36	+ using ArchTag = Arch::AtlasA2;
		37	+};
		38	+
		39	+// For Ascend950, FA Unshared Softmax
		40	+struct EpilogueAscend950FAUnsharedSoftmax {
		41	+ using ArchTag = Arch::Ascend950;
		42	+};
		43	+
		44	+// For AtlasA2, FA Infer CombineScale
		45	+struct EpilogueAtlasA2CombineScale {
		46	+ using ArchTag = Arch::AtlasA2;
		47	+};
		48	+
		49	+// For Ascend950, FA Infer CombineScale
		50	+struct EpilogueAscend950CombineScale {
		51	+ using ArchTag = Arch::Ascend950;
		52	+};
19	// For AtlasA2, an element wise epilogue of the form D = C + X, where X is an additional source	53	// For AtlasA2, an element wise epilogue of the form D = C + X, where X is an additional source
20	struct EpilogueAtlasA2ElemWiseOneSource {	54	struct EpilogueAtlasA2ElemWiseOneSource {
21	using ArchTag = Arch::AtlasA2;	55	using ArchTag = Arch::AtlasA2;
	@@ -67,6 +101,26 @@ struct EpilogueAscend950FARescaleO {
67	using ArchTag = Arch::Ascend950;	101	using ArchTag = Arch::Ascend950;
68	};	102	};
69		103
		104	+// For Ascend950, XA Shared Infer online Softmax
		105	+struct EpilogueAscend950XASharedSoftmax {
		106	+ using ArchTag = Arch::Ascend950;
		107	+};
		108	+
		109	+// For Ascend950, XA unshared KV infer softmax
		110	+struct EpilogueAscend950XAUnsharedSoftmax {
		111	+ using ArchTag = Arch::Ascend950;
		112	+};
		113	+
		114	+// For Ascend950, XA Shared Infer RescaleO
		115	+struct EpilogueAscend950XASharedRescaleO {
		116	+ using ArchTag = Arch::Ascend950;
		117	+};
		118	+
		119	+// For Ascend950, XA Combine scale
		120	+struct EpilogueAscend950XACombineScale {
		121	+ using ArchTag = Arch::Ascend950;
		122	+};
		123	+
70	// For AtlasA2, MLA RescaleO	124	// For AtlasA2, MLA RescaleO
71	struct EpilogueAtlasA2MLARescaleO {	125	struct EpilogueAtlasA2MLARescaleO {
72	using ArchTag = Arch::AtlasA2;	126	using ArchTag = Arch::AtlasA2;
	@@ -225,6 +279,37 @@ struct EpilogueElemWiseNoSourceFromUB {
225	static constexpr uint32_t OPERANDS_NUM = 2;	279	static constexpr uint32_t OPERANDS_NUM = 2;
226	static constexpr uint32_t UB_STAGES = 2;	280	static constexpr uint32_t UB_STAGES = 2;
227	};	281	};
		282	+
		283	+template <LseMode LSE_MODE_>
		284	+struct EpilogueAtlasA2XFAIOnlineSoftmax {
		285	+ using ArchTag = Arch::AtlasA2;
		286	+ static constexpr LseMode LSE_MODE = LSE_MODE_;
		287	+};
		288	+
		289	+template <LseMode LSE_MODE_>
		290	+struct EpilogueAtlasA2OnlineSoftmax_FD {
		291	+ using ArchTag = Arch::AtlasA2;
		292	+ static constexpr LseMode LSE_MODE = LSE_MODE_;
		293	+};
		294	+
		295	+// For Ascend950, FA Infer online Softmax FD
		296	+template <LseMode LSE_MODE_>
		297	+struct EpilogueAscend950OnlineSoftmax_FD {
		298	+ using ArchTag = Arch::Ascend950;
		299	+ static constexpr LseMode LSE_MODE = LSE_MODE_;
		300	+};
		301	+
		302	+template <LseMode LSE_MODE_>
		303	+struct EpilogueAtlasA2XFAIRescaleO {
		304	+ using ArchTag = Arch::AtlasA2;
		305	+ static constexpr LseMode LSE_MODE = LSE_MODE_;
		306	+};
		307	+
		308	+template <LseMode LSE_MODE_>
		309	+struct EpilogueAtlasA2RescaleO_FD {
		310	+ using ArchTag = Arch::AtlasA2;
		311	+ static constexpr LseMode LSE_MODE = LSE_MODE_;
		312	+};
228	} // namespace Catlass::Epilogue	313	} // namespace Catlass::Epilogue
229		314
230	#endif // CATLASS_EPILOGUE_DISPATCH_POLICY_HPP	315	#endif // CATLASS_EPILOGUE_DISPATCH_POLICY_HPP

Minclude/catlass/gemm/block/block_mmad.hpp+11-0

	@@ -89,6 +89,11 @@ struct BlockPrologue {
89	} // namespace Catlass::Gemm::Block	89	} // namespace Catlass::Gemm::Block
90		90
91	#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201)	91	#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201)
		92	+#include "catlass/gemm/block/block_mmad_unshared_fa_pv.hpp"
		93	+#include "catlass/gemm/block/block_mmad_unshared_fa_qk.hpp"
		94	+#include "catlass/gemm/block/block_mmad_pingpong.hpp"
		95	+#include "catlass/gemm/block/block_mmad_fai_qk_split_row.hpp"
		96	+#include "catlass/gemm/block/block_mmad_fai_pv_split_row.hpp"
92	#include "catlass/gemm/block/block_mmad_fa_qk.hpp"	97	#include "catlass/gemm/block/block_mmad_fa_qk.hpp"
93	#include "catlass/gemm/block/block_mmad_fa_pv.hpp"	98	#include "catlass/gemm/block/block_mmad_fa_pv.hpp"
94	#include "catlass/gemm/block/block_mmad_mla_qk.hpp"	99	#include "catlass/gemm/block/block_mmad_mla_qk.hpp"
	@@ -159,5 +164,11 @@ struct BlockPrologue {
159	#include "catlass/gemm/block/block_mmad_flash_attention_qk.hpp"	164	#include "catlass/gemm/block/block_mmad_flash_attention_qk.hpp"
160	#include "catlass/gemm/block/block_mmad_flash_attention_qk_DN.hpp"	165	#include "catlass/gemm/block/block_mmad_flash_attention_qk_DN.hpp"
161	#include "catlass/gemm/block/block_mmad_preload_async_with_callback_tla_l0c_to_ub.hpp"	166	#include "catlass/gemm/block/block_mmad_preload_async_with_callback_tla_l0c_to_ub.hpp"
		167	+#include "catlass/gemm/block/block_mmad_xa_shared_qk_tla.hpp"
		168	+#include "catlass/gemm/block/block_mmad_xa_unshared_qk_tla.hpp"
		169	+#include "catlass/gemm/block/block_mmad_xa_shared_pv_tla.hpp"
		170	+#include "catlass/gemm/block/block_mmad_xa_unshared_pv_tla.hpp"
162	#endif	171	#endif
		172	+#include "catlass/gemm/block/block_mmad_xfai_qk.hpp"
		173	+#include "catlass/gemm/block/block_mmad_xfai_pv.hpp"
163	#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_HPP	174	#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_HPP

Ainclude/catlass/gemm/block/block_mmad_fai_pv_split_row.hpp+226-0

@@ -0,0 +1,226 @@
	1	+/*
	2	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+ * This file is a part of the CANN Open Software.
	4	+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
	5	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+ * See LICENSE in the root of the software repository for the full text of the License.
	9	+ */
	10	+
	11	+#ifndef CATLASS_GEMM_BLOCK_MMAD_PV_SPLIT_ROW_HPP
	12	+#define CATLASS_GEMM_BLOCK_MMAD_PV_SPLIT_ROW_HPP
	13	+
	14	+#include "catlass/catlass.hpp"
	15	+#include "catlass/arch/resource.hpp"
	16	+#include "catlass/coord.hpp"
	17	+#include "catlass/gemm/dispatch_policy.hpp"
	18	+#include "catlass/gemm/helper.hpp"
	19	+#include "catlass/gemm_coord.hpp"
	20	+#include "catlass/gemm/tile/tile_copy.hpp"
	21	+#include "catlass/gemm/tile/tile_mmad.hpp"
	22	+
	23	+////////////////////////////////////////////////////////////////////
	24	+
	25	+namespace Catlass::Gemm::Block {
	26	+////////////////////////////////////////////////////////////////////
	27	+
	28	+template <
	29	+ bool PAGED_CACHE_FLAG_, bool ENABLE_UNIT_FLAG_, class L1TileShape_, class L0TileShape_, class AType_, class BType_,
	30	+ class CType_, class BiasType_, class TileCopy_, class TileMmad_>
	31	+struct BlockMmad<
	32	+ MmadAtlasA2FAIPVSplitRow<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>, L1TileShape_, L0TileShape_, AType_, BType_, CType_,
	33	+ BiasType_, TileCopy_, TileMmad_> {
	34	+public:
	35	+ // Type Aliases
	36	+ using DispatchPolicy = MmadAtlasA2FAIPVSplitRow<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>;
	37	+ using ArchTag = typename DispatchPolicy::ArchTag;
	38	+ using L1TileShape = L1TileShape_;
	39	+ using L0TileShape = L0TileShape_;
	40	+ using ElementA = typename AType_::Element;
	41	+ using LayoutA = typename AType_::Layout;
	42	+ using ElementB = typename BType_::Element;
	43	+ using LayoutB = typename BType_::Layout;
	44	+ using ElementC = typename CType_::Element;
	45	+ using LayoutC = typename CType_::Layout;
	46	+ using TileMmad = TileMmad_;
	47	+ using CopyGmToL1A = typename TileCopy_::CopyGmToL1A;
	48	+ using CopyGmToL1B = typename TileCopy_::CopyGmToL1B;
	49	+ using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A;
	50	+ using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B;
	51	+ using CopyL0CToGm = typename TileCopy_::CopyL0CToGm;
	52	+ using ElementAccumulator =
	53	+ typename Gemm::helper::ElementAccumulatorSelector<ElementA, ElementB>::ElementAccumulator;
	54	+ using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc;
	55	+ using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc;
	56	+ using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst;
	57	+ using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst;
	58	+ using LayoutCInL0 = layout::zN;
	59	+
	60	+ using L1AAlignHelper = Gemm::helper::L1AlignHelper<ElementA, LayoutA>;
	61	+ using L1BAlignHelper = Gemm::helper::L1AlignHelper<ElementB, LayoutB>;
	62	+
	63	+ static constexpr uint32_t STAGES = DispatchPolicy::STAGES;
	64	+ static constexpr uint32_t L1A_SIZE = 32768 * sizeof(ElementA);
	65	+ static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB);
	66	+ static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE;
	67	+ static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE;
	68	+ static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE;
	69	+ static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES;
	70	+ static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES;
	71	+ static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES;
	72	+
	73	+ // Check LayoutC
	74	+ static_assert(std::is_same_v<LayoutC, layout::RowMajor>, "LayoutC only support RowMajor yet!");
	75	+
	76	+ static_assert(L1TileShape::M * L1TileShape::K <= 32768, "L1TileShape::M * L1TileShape::K must be less than 32768");
	77	+
	78	+ /// Construct
	79	+ CATLASS_DEVICE
	80	+ BlockMmad(Arch::Resource<ArchTag>& resource, uint32_t l1BufAddrStart = 0)
	81	+ {
	82	+ // Allocate L1 memory space
	83	+ l1BTensor = resource.l1Buf.template GetBufferByByte<ElementB>(l1BufAddrStart + L1A_SIZE * 2);
	84	+ for (uint32_t i = 0; i < STAGES; i++) {
	85	+ l1ATensor[i] = resource.l1Buf.template GetBufferByByte<ElementA>(l1BufAddrStart + L1A_SIZE * i);
	86	+ l0ATensor[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(L0A_PINGPONG_BUF_SIZE * i);
	87	+ l0BTensor[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(L0B_PINGPONG_BUF_SIZE * i);
	88	+ l0CTensor[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(L0C_PINGPONG_BUF_SIZE * i);
	89	+ }
	90	+ }
	91	+
	92	+ /// Destructor
	93	+ CATLASS_DEVICE
	94	+ ~BlockMmad()
	95	+ {}
	96	+
	97	+ CATLASS_DEVICE
	98	+ void getKVOffset(
	99	+ AscendC::GlobalTensor<int32_t>& gBlockTable, uint32_t& kOffset, uint32_t& nowNIdx, uint32_t& strideKV,
	100	+ uint32_t& blockSize)
	101	+ {
	102	+ if constexpr (PAGED_CACHE_FLAG_) {
	103	+ uint32_t blockTableId = gBlockTable.GetValue(nowNIdx);
	104	+ kOffset = blockTableId * blockSize * strideKV;
	105	+ } else {
	106	+ kOffset = nowNIdx * blockSize * strideKV;
	107	+ }
	108	+ }
	109	+
	110	+ CATLASS_DEVICE
	111	+ void operator()(
	112	+ AscendC::GlobalTensor<ElementA> gA, AscendC::GlobalTensor<ElementB> gB, AscendC::GlobalTensor<ElementC> gC,
	113	+ AscendC::GlobalTensor<int32_t> gBlockTable, LayoutA layoutA, LayoutB layoutB, LayoutC layoutC,
	114	+ GemmCoord actualOriShape, uint32_t& nIdx, uint32_t& blockSize, uint32_t strideKV,
	115	+ Arch::CrossCoreFlag softmaxFlag, uint32_t& l1KPPingPongFlag, uint32_t& l0ABPingPongFlag,
	116	+ uint32_t& l0CPingPongFlag)
	117	+ {
	118	+ uint32_t rowNum = actualOriShape[0];
	119	+ uint32_t embed = actualOriShape[1];
	120	+ uint32_t stackSeqTile = actualOriShape[2];
	121	+ uint32_t gBOffset = 0;
	122	+
	123	+ // load V
	124	+ LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout<ElementB>(stackSeqTile, embed);
	125	+ uint32_t kLoop = CeilDiv(stackSeqTile, blockSize);
	126	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID2);
	127	+ for (uint32_t blockStackIdx = 0; blockStackIdx < kLoop; blockStackIdx++) {
	128	+ uint32_t nowNIdx = nIdx + blockStackIdx;
	129	+ uint32_t kActual = AscendC::Std::min(stackSeqTile - blockStackIdx * blockSize, blockSize);
	130	+ getKVOffset(gBlockTable, gBOffset, nowNIdx, strideKV, blockSize);
	131	+
	132	+ auto layoutBTile = layoutB.GetTileLayout(MakeCoord(kActual, embed));
	133	+ MatrixCoord l1BTileCoord{blockStackIdx * blockSize, 0};
	134	+ auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BTileCoord)];
	135	+ copyGmToL1B(l1BTile, gB[gBOffset], layoutBInL1, layoutBTile);
	136	+ }
	137	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID2);
	138	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID2);
	139	+
	140	+ Arch::CrossCoreWaitFlag(softmaxFlag);
	141	+
	142	+ uint32_t mL1Loop = CeilDiv<L1TileShape::M>(rowNum);
	143	+ uint32_t kL1Loop = CeilDiv<L1TileShape::K>(stackSeqTile);
	144	+ for (uint32_t mL1Idx = 0; mL1Idx < mL1Loop; mL1Idx++) {
	145	+ uint32_t mL1Actual = (mL1Idx < mL1Loop - 1) ? L1TileShape::M : (rowNum - mL1Idx * L1TileShape::M);
	146	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CPingPongFlag);
	147	+ for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) {
	148	+ uint32_t kL1Actual = (kL1Idx < kL1Loop - 1) ? L1TileShape::K : (stackSeqTile - kL1Idx * L1TileShape::K);
	149	+
	150	+ // load P
	151	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1KPPingPongFlag);
	152	+ MatrixCoord gmATileCoord{mL1Idx * L1TileShape::M, kL1Idx * L1TileShape::K};
	153	+ auto gmTileA = gA[layoutA.GetOffset(gmATileCoord)];
	154	+ auto layoutTileA = layoutA.GetTileLayout(MakeCoord(mL1Actual, kL1Actual));
	155	+ LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout<ElementA>(mL1Actual, kL1Actual);
	156	+ copyGmToL1A(l1ATensor[l1KPPingPongFlag], gmTileA, layoutAInL1, layoutTileA);
	157	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1KPPingPongFlag);
	158	+
	159	+ uint32_t kL0Loop = CeilDiv<L0TileShape::K>(kL1Actual);
	160	+ for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) {
	161	+ uint32_t kL0Actual =
	162	+ (kL0Idx < kL0Loop - 1) ? L0TileShape::K : (kL1Actual - kL0Idx * L0TileShape::K);
	163	+
	164	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag);
	165	+ LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout<ElementB>(kL0Actual, embed);
	166	+ MatrixCoord l1BTileCoord{kL1Idx * L1TileShape::K + kL0Idx * L0TileShape::K, 0};
	167	+ auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BTileCoord)];
	168	+
	169	+ copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTile, layoutBInL0, layoutBInL1);
	170	+
	171	+ LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout<ElementA>(mL1Actual, kL0Actual);
	172	+ MatrixCoord l1ATileCoord{0, kL0Idx * L0TileShape::K};
	173	+ auto l1ATile = l1ATensor[l1KPPingPongFlag][layoutAInL1.GetOffset(l1ATileCoord)];
	174	+
	175	+ if (kL0Idx == 0) {
	176	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1KPPingPongFlag);
	177	+ }
	178	+ copyL1ToL0A(l0ATensor[l0ABPingPongFlag], l1ATile, layoutAInL0, layoutAInL1);
	179	+ if (kL0Idx == kL0Loop - 1) {
	180	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1KPPingPongFlag);
	181	+ }
	182	+
	183	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0ABPingPongFlag);
	184	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0ABPingPongFlag);
	185	+ bool initMmad = kL1Idx == 0 && kL0Idx == 0;
	186	+ tileMmad(
	187	+ l0CTensor[l0CPingPongFlag], l0ATensor[l0ABPingPongFlag], l0BTensor[l0ABPingPongFlag], mL1Actual,
	188	+ embed, kL0Actual, initMmad);
	189	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag);
	190	+ l0ABPingPongFlag = 1 - l0ABPingPongFlag;
	191	+ }
	192	+ l1KPPingPongFlag = 1 - l1KPPingPongFlag;
	193	+ }
	194	+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(l0CPingPongFlag);
	195	+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(l0CPingPongFlag);
	196	+ MatrixCoord gmCTileCoord{mL1Idx * L0TileShape::M, 0};
	197	+ LayoutC layoutCTile = layoutC.GetTileLayout(MakeCoord(mL1Actual, embed));
	198	+ auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mL1Actual, embed));
	199	+ copyL0CToGm(gC[layoutC.GetOffset(gmCTileCoord)], l0CTensor[l0CPingPongFlag], layoutCTile, layoutInL0C);
	200	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CPingPongFlag);
	201	+ l0CPingPongFlag = 1 - l0CPingPongFlag;
	202	+ }
	203	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID2);
	204	+ }
	205	+
	206	+protected:
	207	+ /// Data members
	208	+ AscendC::LocalTensor<ElementA> l1ATensor[STAGES];
	209	+ AscendC::LocalTensor<ElementB> l1BTensor;
	210	+ AscendC::LocalTensor<ElementA> l0ATensor[STAGES];
	211	+ AscendC::LocalTensor<ElementB> l0BTensor[STAGES];
	212	+ AscendC::LocalTensor<ElementAccumulator> l0CTensor[STAGES];
	213	+
	214	+ TileMmad tileMmad;
	215	+ CopyGmToL1A copyGmToL1A;
	216	+ CopyGmToL1B copyGmToL1B;
	217	+ CopyL1ToL0A copyL1ToL0A;
	218	+ CopyL1ToL0B copyL1ToL0B;
	219	+ CopyL0CToGm copyL0CToGm;
	220	+};
	221	+
	222	+////////////////////////////////////////////////////////////////////
	223	+
	224	+} // namespace Catlass::Gemm::Block
	225	+
	226	+#endif // CATLASS_GEMM_BLOCK_MMAD_PV_SPLIT_ROW_HPP

Ainclude/catlass/gemm/block/block_mmad_fai_qk_split_row.hpp+237-0

@@ -0,0 +1,237 @@
	1	+/*
	2	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	3	+ * This file is a part of the CANN Open Software.
	4	+ * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
	5	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	6	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
	7	+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
	8	+ * See LICENSE in the root of the software repository for the full text of the License.
	9	+ */
	10	+
	11	+#ifndef CATLASS_GEMM_BLOCK_MMAD_QK_SPLIT_ROW_HPP
	12	+#define CATLASS_GEMM_BLOCK_MMAD_QK_SPLIT_ROW_HPP
	13	+
	14	+#include "catlass/catlass.hpp"
	15	+#include "catlass/arch/resource.hpp"
	16	+#include "catlass/coord.hpp"
	17	+#include "catlass/gemm/dispatch_policy.hpp"
	18	+#include "catlass/gemm/helper.hpp"
	19	+#include "catlass/gemm_coord.hpp"
	20	+#include "catlass/gemm/tile/tile_copy.hpp"
	21	+#include "catlass/gemm/tile/tile_mmad.hpp"
	22	+
	23	+////////////////////////////////////////////////////////////////////
	24	+
	25	+namespace Catlass::Gemm::Block {
	26	+////////////////////////////////////////////////////////////////////
	27	+
	28	+template <
	29	+ bool PAGED_CACHE_FLAG_, bool ENABLE_UNIT_FLAG_, class L1TileShape_, class L0TileShape_, class AType_, class BType_,
	30	+ class CType_, class BiasType_, class TileCopy_, class TileMmad_>
	31	+struct BlockMmad<
	32	+ MmadAtlasA2FAIQKSplitRow<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>, L1TileShape_, L0TileShape_, AType_, BType_, CType_,
	33	+ BiasType_, TileCopy_, TileMmad_> {
	34	+public:
	35	+ // Type Aliases
	36	+ using DispatchPolicy = MmadAtlasA2FAIQKSplitRow<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>;
	37	+ using ArchTag = typename DispatchPolicy::ArchTag;
	38	+ using L1TileShape = L1TileShape_;
	39	+ using L0TileShape = L0TileShape_;
	40	+ using ElementA = typename AType_::Element;
	41	+ using LayoutA = typename AType_::Layout;
	42	+ using ElementB = typename BType_::Element;
	43	+ using LayoutB = typename BType_::Layout;
	44	+ using ElementC = typename CType_::Element;
	45	+ using LayoutC = typename CType_::Layout;
	46	+ using TileMmad = TileMmad_;
	47	+ using CopyGmToL1A = typename TileCopy_::CopyGmToL1A;
	48	+ using CopyGmToL1B = typename TileCopy_::CopyGmToL1B;
	49	+ using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A;
	50	+ using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B;
	51	+ using CopyL0CToGm = typename TileCopy_::CopyL0CToGm;
	52	+ using ElementAccumulator =
	53	+ typename Gemm::helper::ElementAccumulatorSelector<ElementA, ElementB>::ElementAccumulator;
	54	+ using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc;
	55	+ using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc;
	56	+ using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst;
	57	+ using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst;
	58	+ using LayoutCInL0 = layout::zN;
	59	+
	60	+ using L1AAlignHelper = Gemm::helper::L1AlignHelper<ElementA, LayoutA>;
	61	+ using L1BAlignHelper = Gemm::helper::L1AlignHelper<ElementB, LayoutB>;
	62	+
	63	+ static constexpr uint32_t STAGES = DispatchPolicy::STAGES;
	64	+ static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA);
	65	+ static constexpr uint32_t L1B_SIZE = 32768 * sizeof(ElementB);
	66	+ static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE;
	67	+ static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE;
	68	+ static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE;
	69	+ static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES;
	70	+ static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES;
	71	+ static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES;
	72	+ static constexpr uint32_t BLOCK_SIZE = 16;
	73	+
	74	+ static_assert(std::is_same_v<LayoutC, layout::RowMajor>, "LayoutC only support RowMajor yet!");
	75	+
	76	+ static_assert(L1TileShape::N * L1TileShape::K <= 32768, "L1TileShape::N * L1TileShape::K must be less than 32768");
	77	+
	78	+ CATLASS_DEVICE
	79	+ BlockMmad(Arch::Resource<ArchTag>& resource, uint32_t l1BufAddrStart = 0)
	80	+ {
	81	+ // Allocate L1 memory space
	82	+ l1ATensor = resource.l1Buf.template GetBufferByByte<ElementA>(l1BufAddrStart);
	83	+ for (uint32_t i = 0; i < STAGES; i++) {
	84	+ l1BTensor[i] = resource.l1Buf.template GetBufferByByte<ElementB>(l1BufAddrStart + L1A_SIZE + L1B_SIZE * i);
	85	+ l0ATensor[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(L0A_PINGPONG_BUF_SIZE * i);
	86	+ l0BTensor[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(L0B_PINGPONG_BUF_SIZE * i);
	87	+ l0CTensor[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(L0C_PINGPONG_BUF_SIZE * i);
	88	+ }
	89	+ }
	90	+
	91	+ CATLASS_DEVICE
	92	+ ~BlockMmad()
	93	+ {}
	94	+
	95	+ CATLASS_DEVICE
	96	+ void loadQGM(
	97	+ AscendC::GlobalTensor<ElementA> gA, LayoutA layoutA, uint32_t rowNum, uint32_t& singleGroupHeads,
	98	+ uint32_t& qHeads)
	99	+ {
	100	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID3);
	101	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID3);
	102	+ uint32_t embed = layoutA.shape(1);
	103	+ uint32_t rowNumRound = RoundUp<L1AAlignHelper::M_ALIGNED>(rowNum);
	104	+ uint32_t tokenNumPerGroup = rowNum / singleGroupHeads;
	105	+ auto layoutSingleANd = layoutA.GetTileLayout(MakeCoord(singleGroupHeads, embed));
	106	+ LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout<ElementA>(rowNum, embed);
	107	+ copyGmToL1A(
	108	+ l1ATensor, gA, layoutAInL1, layoutSingleANd, tokenNumPerGroup, qHeads * embed, tokenNumPerGroup, BLOCK_SIZE,
	109	+ rowNumRound);
	110	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID3);
	111	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID3);
	112	+ }
	113	+
	114	+ CATLASS_DEVICE
	115	+ void getBlockShape(GemmCoord& actualShape, uint32_t nL1Idx, uint32_t nL1Loop, uint32_t stackSeqTile)
	116	+ {
	117	+ uint32_t nSplitSize = L1TileShape::N;
	118	+ if (nL1Idx == nL1Loop - 1) {
	119	+ nSplitSize = stackSeqTile - nL1Idx * L1TileShape::N;
	120	+ }
	121	+ actualShape[1] = nSplitSize;
	122	+ }
	123	+
	124	+ CATLASS_DEVICE
	125	+ void getKVOffset(
	126	+ AscendC::GlobalTensor<int32_t>& gBlockTable, uint32_t& kOffset, uint32_t nowNIdx, uint32_t nL1Idx,
	127	+ uint32_t strideKV, uint32_t blockSize)
	128	+ {
	129	+ if constexpr (PAGED_CACHE_FLAG_) {
	130	+ uint32_t blockTableId = gBlockTable.GetValue(nowNIdx);
	131	+ kOffset = blockTableId * blockSize * strideKV + nL1Idx * L1TileShape::N * strideKV;
	132	+ } else {
	133	+ kOffset = nowNIdx * blockSize * strideKV + nL1Idx * L1TileShape::N * strideKV;
	134	+ }
	135	+ }
	136	+
	137	+ CATLASS_DEVICE
	138	+ void operator()(
	139	+ AscendC::GlobalTensor<ElementA> gA, AscendC::GlobalTensor<ElementB> gB, AscendC::GlobalTensor<ElementC> gC,
	140	+ AscendC::GlobalTensor<int32_t> gBlockTable, LayoutA layoutA, LayoutB layoutB, LayoutC layoutC,
	141	+ GemmCoord actualOriShape, uint32_t& nIdx, uint32_t& blockSize, uint32_t strideKV, uint32_t& l1KPPingPongFlag,
	142	+ uint32_t& l0ABPingPongFlag, uint32_t& l0CPingPongFlag)
	143	+ {
	144	+ uint32_t rowNum = actualOriShape[0];
	145	+ uint32_t stackSeqTile = actualOriShape[1];
	146	+ uint32_t embed = actualOriShape[2];
	147	+
	148	+ GemmCoord actualShape{rowNum, 0, embed};
	149	+ uint32_t gBOffset = 0;
	150	+
	151	+ LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout<ElementA>(rowNum, embed);
	152	+
	153	+ uint32_t tileNNumPerPaged = blockSize / L1TileShape::N;
	154	+ uint32_t nL1Loop = CeilDiv<L1TileShape::N>(stackSeqTile);
	155	+ for (uint32_t nL1Idx = 0; nL1Idx < nL1Loop; ++nL1Idx) {
	156	+ uint32_t nowNIdx = nIdx + nL1Idx / tileNNumPerPaged;
	157	+ getBlockShape(actualShape, nL1Idx, nL1Loop, stackSeqTile);
	158	+ getKVOffset(gBlockTable, gBOffset, nowNIdx, nL1Idx % tileNNumPerPaged, strideKV, blockSize);
	159	+
	160	+ uint32_t mActual = actualShape.m();
	161	+ uint32_t kActual = actualShape.k();
	162	+ uint32_t nActual = actualShape.n();
	163	+ LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout<ElementB>(kActual, nActual);
	164	+
	165	+ auto layoutBTile = layoutB.GetTileLayout(MakeCoord(kActual, nActual));
	166	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1KPPingPongFlag);
	167	+ copyGmToL1B(l1BTensor[l1KPPingPongFlag], gB[gBOffset], layoutBInL1, layoutBTile);
	168	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1KPPingPongFlag);
	169	+
	170	+ uint32_t mL0Loop = CeilDiv<L0TileShape::M>(mActual);
	171	+ uint32_t kL0Loop = CeilDiv<L0TileShape::K>(kActual);
	172	+ for (uint32_t mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) {
	173	+ uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0TileShape::M : (mActual - mL0Idx * L0TileShape::M);
	174	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CPingPongFlag);
	175	+ for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) {
	176	+ uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0TileShape::K : (kActual - kL0Idx * L0TileShape::K);
	177	+
	178	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag);
	179	+ LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout<ElementA>(mL0Actual, kL0Actual);
	180	+ MatrixCoord l1ATileCoord{mL0Idx * L0TileShape::M, kL0Idx * L0TileShape::K};
	181	+ auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ATileCoord)];
	182	+ copyL1ToL0A(l0ATensor[l0ABPingPongFlag], l1ATile, layoutAInL0, layoutAInL1);
	183	+
	184	+ LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout<ElementB>(kL0Actual, nActual);
	185	+ MatrixCoord l1BTileCoord{kL0Idx * L0TileShape::K, 0};
	186	+ auto l1BTile = l1BTensor[l1KPPingPongFlag][layoutBInL1.GetOffset(l1BTileCoord)];
	187	+ if ((mL0Idx == 0) && (kL0Idx == 0)) {
	188	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1KPPingPongFlag);
	189	+ }
	190	+ copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTile, layoutBInL0, layoutBInL1);
	191	+ if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) {
	192	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1KPPingPongFlag);
	193	+ }
	194	+
	195	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0ABPingPongFlag);
	196	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0ABPingPongFlag);
	197	+ bool initMmad = kL0Idx == 0;
	198	+ tileMmad(
	199	+ l0CTensor[l0CPingPongFlag], l0ATensor[l0ABPingPongFlag], l0BTensor[l0ABPingPongFlag], mL0Actual,
	200	+ nActual, kL0Actual, initMmad);
	201	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag);
	202	+ l0ABPingPongFlag = 1 - l0ABPingPongFlag;
	203	+ }
	204	+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(l0CPingPongFlag);
	205	+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(l0CPingPongFlag);
	206	+ MatrixCoord gmCTileCoord{mL0Idx * L0TileShape::M, nL1Idx * L1TileShape::N};
	207	+ LayoutC layoutCTile = layoutC.GetTileLayout(MakeCoord(mL0Actual, nActual));
	208	+ auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mL0Actual, nActual));
	209	+ copyL0CToGm(gC[layoutC.GetOffset(gmCTileCoord)], l0CTensor[l0CPingPongFlag], layoutCTile, layoutInL0C);
	210	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CPingPongFlag);
	211	+ l0CPingPongFlag = 1 - l0CPingPongFlag;
	212	+ }
	213	+ l1KPPingPongFlag = 1 - l1KPPingPongFlag;
	214	+ }
	215	+ }
	216	+
	217	+protected:
	218	+ /// Data members
	219	+ AscendC::LocalTensor<ElementA> l1ATensor;
	220	+ AscendC::LocalTensor<ElementB> l1BTensor[STAGES];
	221	+ AscendC::LocalTensor<ElementA> l0ATensor[STAGES];
	222	+ AscendC::LocalTensor<ElementB> l0BTensor[STAGES];
	223	+ AscendC::LocalTensor<ElementAccumulator> l0CTensor[STAGES];
	224	+
	225	+ TileMmad tileMmad;
	226	+ CopyGmToL1A copyGmToL1A;
	227	+ CopyGmToL1B copyGmToL1B;
	228	+ CopyL1ToL0A copyL1ToL0A;
	229	+ CopyL1ToL0B copyL1ToL0B;
	230	+ CopyL0CToGm copyL0CToGm;
	231	+};
	232	+
	233	+////////////////////////////////////////////////////////////////////
	234	+
	235	+} // namespace Catlass::Gemm::Block
	236	+
	237	+#endif // CATLASS_GEMM_BLOCK_MMAD_QK_SPLIT_ROW_HPP

Ainclude/catlass/gemm/block/block_mmad_unshared_fa_pv.hpp+167-0

@@ -0,0 +1,167 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_UNSHARED_FA_PV_HPP
	13	+#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_UNSHARED_FA_PV_HPP
	14	+
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/arch/cross_core_sync.hpp"
	17	+#include "catlass/arch/resource.hpp"
	18	+#include "catlass/coord.hpp"
	19	+#include "catlass/gemm/dispatch_policy.hpp"
	20	+#include "catlass/gemm/helper.hpp"
	21	+#include "catlass/gemm_coord.hpp"
	22	+#include "catlass/gemm/tile/tile_copy.hpp"
	23	+#include "catlass/gemm/tile/tile_mmad.hpp"
	24	+
	25	+////////////////////////////////////////////////////////////////////
	26	+
	27	+namespace Catlass::Gemm::Block {
	28	+////////////////////////////////////////////////////////////////////
	29	+
	30	+template <
	31	+ class L1TileShape_, class L0TileShape_, class AType_, class BType_, class CType_, class BiasType_, class TileCopy_,
	32	+ class TileMmad_>
	33	+struct BlockMmad<
	34	+ MmadAtlasA2UnsharedFAPV, L1TileShape_, L0TileShape_, AType_, BType_, CType_, BiasType_, TileCopy_, TileMmad_> {
	35	+public:
	36	+ // Type Aliases
	37	+ using DispatchPolicy = MmadAtlasA2UnsharedFAPV;
	38	+ using ArchTag = typename DispatchPolicy::ArchTag;
	39	+ using L1TileShape = L1TileShape_;
	40	+ using L0TileShape = L0TileShape_;
	41	+ using ElementA = typename AType_::Element;
	42	+ using LayoutA = typename AType_::Layout;
	43	+ using ElementB = typename BType_::Element;
	44	+ using LayoutB = typename BType_::Layout;
	45	+ using ElementC = typename CType_::Element;
	46	+ using LayoutC = typename CType_::Layout;
	47	+ using TileMmad = TileMmad_;
	48	+ using CopyGmToL1A = typename TileCopy_::CopyGmToL1A;
	49	+ using CopyGmToL1B = typename TileCopy_::CopyGmToL1B;
	50	+ using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A;
	51	+ using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B;
	52	+ using CopyL0CToGm = typename TileCopy_::CopyL0CToGm;
	53	+ using ElementAccumulator =
	54	+ typename Gemm::helper::ElementAccumulatorSelector<ElementA, ElementB>::ElementAccumulator;
	55	+ using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc;
	56	+ using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc;
	57	+ using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst;
	58	+ using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst;
	59	+ using LayoutCInL0 = layout::zN;
	60	+
	61	+ using L1AAlignHelper = Gemm::helper::L1AlignHelper<ElementA, LayoutA>;
	62	+ using L1BAlignHelper = Gemm::helper::L1AlignHelper<ElementB, LayoutB>;
	63	+
	64	+ static constexpr uint32_t STAGES = DispatchPolicy::STAGES;
	65	+ static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA);
	66	+ static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB);
	67	+ static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE;
	68	+ static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE;
	69	+ static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE;
	70	+ static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES;
	71	+ static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES;
	72	+ static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES;
	73	+
	74	+ // Check LayoutC
	75	+ static_assert(std::is_same_v<LayoutC, layout::RowMajor>, "LayoutC only support RowMajor yet!");
	76	+
	77	+ /// Construct
	78	+ CATLASS_DEVICE
	79	+ BlockMmad(Arch::Resource<ArchTag>& resource, uint32_t l1BufAddrStart = 0)
	80	+ {
	81	+ l1ATensor = resource.l1Buf.template GetBufferByByte<ElementA>(l1BufAddrStart);
	82	+ l1BTensor = resource.l1Buf.template GetBufferByByte<ElementB>(l1BufAddrStart + L1A_SIZE);
	83	+ l0ATensor = resource.l0ABuf.template GetBufferByByte<ElementA>(0);
	84	+ l0BTensor = resource.l0BBuf.template GetBufferByByte<ElementB>(0);
	85	+ l0CTensor = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(0);
	86	+
	87	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID2);
	88	+ }
	89	+
	90	+ /// Destructor
	91	+ CATLASS_DEVICE
	92	+ ~BlockMmad()
	93	+ {
	94	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID2);
	95	+ }
	96	+
	97	+ /// Perform a block-scoped matrix multiply-accumulate
	98	+ CATLASS_DEVICE
	99	+ void operator()(
	100	+ AscendC::GlobalTensor<ElementA> gA, AscendC::GlobalTensor<ElementB> gB, AscendC::GlobalTensor<ElementC> gC,
	101	+ LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, GemmCoord actualShape, Arch::CrossCoreFlag softmaxReady)
	102	+ {
	103	+ LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout<ElementA>(L1TileShape::M, L1TileShape::K);
	104	+ LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout<ElementB>(L1TileShape::K, L1TileShape::N);
	105	+
	106	+ uint32_t mRound = RoundUp<L1AAlignHelper::M_ALIGNED>(actualShape.m());
	107	+ uint32_t nRound = RoundUp<L1BAlignHelper::N_ALIGNED>(actualShape.n());
	108	+ uint32_t kRound = RoundUp<L1BAlignHelper::K_ALIGNED>(actualShape.k());
	109	+
	110	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID2);
	111	+ auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShape.k(), actualShape.n()));
	112	+ copyGmToL1B(l1BTensor, gB, layoutBInL1, layoutTileB);
	113	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
	114	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
	115	+
	116	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
	117	+ LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout<ElementB>(kRound, nRound);
	118	+ copyL1ToL0B(l0BTensor, l1BTensor, layoutBInL0, layoutBInL1);
	119	+
	120	+ Arch::CrossCoreWaitFlag(softmaxReady);
	121	+ auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), actualShape.k()));
	122	+ copyGmToL1A(l1ATensor, gA, layoutAInL1, layoutTileA);
	123	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
	124	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
	125	+ LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout<ElementA>(mRound, kRound);
	126	+ copyL1ToL0A(l0ATensor, l1ATensor, layoutAInL0, layoutAInL1);
	127	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID2);
	128	+
	129	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
	130	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
	131	+
	132	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
	133	+ tileMmad(l0CTensor, l0ATensor, l0BTensor, mRound, nRound, actualShape.k());
	134	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
	135	+
	136	+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
	137	+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
	138	+
	139	+ // copy block out
	140	+ auto blockShape = MakeCoord(actualShape.m(), actualShape.n());
	141	+ auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(blockShape);
	142	+
	143	+ copyL0CToGm(gC, l0CTensor, layoutC, layoutInL0C);
	144	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
	145	+ }
	146	+
	147	+protected:
	148	+ /// Data members
	149	+ AscendC::LocalTensor<ElementA> l1ATensor;
	150	+ AscendC::LocalTensor<ElementB> l1BTensor;
	151	+ AscendC::LocalTensor<ElementA> l0ATensor;
	152	+ AscendC::LocalTensor<ElementB> l0BTensor;
	153	+ AscendC::LocalTensor<ElementAccumulator> l0CTensor;
	154	+
	155	+ TileMmad tileMmad;
	156	+ CopyGmToL1A copyGmToL1A;
	157	+ CopyGmToL1B copyGmToL1B;
	158	+ CopyL1ToL0A copyL1ToL0A;
	159	+ CopyL1ToL0B copyL1ToL0B;
	160	+ CopyL0CToGm copyL0CToGm;
	161	+};
	162	+
	163	+////////////////////////////////////////////////////////////////////
	164	+
	165	+} // namespace Catlass::Gemm::Block
	166	+
	167	+#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_FA_PV_HPP

Ainclude/catlass/gemm/block/block_mmad_unshared_fa_qk.hpp+170-0

@@ -0,0 +1,170 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_UNSHARED_FA_QK_HPP
	13	+#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_UNSHARED_FA_QK_HPP
	14	+
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/arch/resource.hpp"
	17	+#include "catlass/coord.hpp"
	18	+#include "catlass/gemm/dispatch_policy.hpp"
	19	+#include "catlass/gemm/helper.hpp"
	20	+#include "catlass/gemm_coord.hpp"
	21	+#include "catlass/gemm/tile/tile_copy.hpp"
	22	+#include "catlass/gemm/tile/tile_mmad.hpp"
	23	+
	24	+////////////////////////////////////////////////////////////////////
	25	+
	26	+namespace Catlass::Gemm::Block {
	27	+////////////////////////////////////////////////////////////////////
	28	+
	29	+template <
	30	+ class L1TileShape_, class L0TileShape_, class AType_, class BType_, class CType_, class BiasType_, class TileCopy_,
	31	+ class TileMmad_>
	32	+struct BlockMmad<
	33	+ MmadAtlasA2UnsharedFAQK, L1TileShape_, L0TileShape_, AType_, BType_, CType_, BiasType_, TileCopy_, TileMmad_> {
	34	+public:
	35	+ // Type Aliases
	36	+ using DispatchPolicy = MmadAtlasA2UnsharedFAQK;
	37	+ using ArchTag = typename DispatchPolicy::ArchTag;
	38	+ using L1TileShape = L1TileShape_;
	39	+ using L0TileShape = L0TileShape_;
	40	+ using ElementA = typename AType_::Element;
	41	+ using LayoutA = typename AType_::Layout;
	42	+ using ElementB = typename BType_::Element;
	43	+ using LayoutB = typename BType_::Layout;
	44	+ using ElementC = typename CType_::Element;
	45	+ using LayoutC = typename CType_::Layout;
	46	+ using TileMmad = TileMmad_;
	47	+ using CopyGmToL1A = typename TileCopy_::CopyGmToL1A;
	48	+ using CopyGmToL1B = typename TileCopy_::CopyGmToL1B;
	49	+ using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A;
	50	+ using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B;
	51	+ using CopyL0CToGm = typename TileCopy_::CopyL0CToGm;
	52	+ using ElementAccumulator =
	53	+ typename Gemm::helper::ElementAccumulatorSelector<ElementA, ElementB>::ElementAccumulator;
	54	+ using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc;
	55	+ using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc;
	56	+ using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst;
	57	+ using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst;
	58	+ using LayoutCInL0 = layout::zN;
	59	+
	60	+ using L1AAlignHelper = Gemm::helper::L1AlignHelper<ElementA, LayoutA>;
	61	+ using L1BAlignHelper = Gemm::helper::L1AlignHelper<ElementB, LayoutB>;
	62	+
	63	+ static constexpr uint32_t STAGES = DispatchPolicy::STAGES;
	64	+ static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA);
	65	+ static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB);
	66	+ static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE;
	67	+ static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE;
	68	+ static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE;
	69	+ static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES;
	70	+ static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES;
	71	+ static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES;
	72	+
	73	+ // Check LayoutC
	74	+ static_assert(std::is_same_v<LayoutC, layout::RowMajor>, "LayoutC only support RowMajor yet!");
	75	+
	76	+ /// Construct
	77	+ CATLASS_DEVICE
	78	+ BlockMmad(Arch::Resource<ArchTag>& resource, uint32_t l1BufAddrStart = 0)
	79	+ {
	80	+ l1ATensor = resource.l1Buf.template GetBufferByByte<ElementA>(l1BufAddrStart);
	81	+ l1BTensor = resource.l1Buf.template GetBufferByByte<ElementB>(l1BufAddrStart + L1A_SIZE);
	82	+ l0ATensor = resource.l0ABuf.template GetBufferByByte<ElementA>(0);
	83	+ l0BTensor = resource.l0BBuf.template GetBufferByByte<ElementB>(0);
	84	+ l0CTensor = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(0);
	85	+
	86	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
	87	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
	88	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
	89	+ }
	90	+
	91	+ /// Destructor
	92	+ CATLASS_DEVICE
	93	+ ~BlockMmad()
	94	+ {
	95	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
	96	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
	97	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
	98	+ }
	99	+
	100	+ /// Perform a block-scoped matrix multiply-accumulate
	101	+ CATLASS_DEVICE
	102	+ void operator()(
	103	+ AscendC::GlobalTensor<ElementA> gA, AscendC::GlobalTensor<ElementB> gB, AscendC::GlobalTensor<ElementC> gC,
	104	+ LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, GemmCoord actualShape)
	105	+ {
	106	+ LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout<ElementA>(L1TileShape::M, L1TileShape::K);
	107	+ LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout<ElementB>(L1TileShape::K, L1TileShape::N);
	108	+
	109	+ uint32_t mRound = RoundUp<L1AAlignHelper::M_ALIGNED>(actualShape.m());
	110	+ uint32_t nRound = RoundUp<L1BAlignHelper::N_ALIGNED>(actualShape.n());
	111	+ uint32_t kRound = RoundUp<L1BAlignHelper::K_ALIGNED>(actualShape.k());
	112	+
	113	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
	114	+ auto layoutTileA = layoutA.GetTileLayout(MakeCoord(actualShape.m(), actualShape.k()));
	115	+ copyGmToL1A(l1ATensor, gA, layoutAInL1, layoutTileA);
	116	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
	117	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
	118	+
	119	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
	120	+ LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout<ElementA>(mRound, kRound);
	121	+ copyL1ToL0A(l0ATensor, l1ATensor, layoutAInL0, layoutAInL1);
	122	+
	123	+ auto layoutTileB = layoutB.GetTileLayout(MakeCoord(actualShape.k(), actualShape.n()));
	124	+ copyGmToL1B(l1BTensor, gB, layoutBInL1, layoutTileB);
	125	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
	126	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
	127	+ LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout<ElementB>(kRound, nRound);
	128	+ copyL1ToL0B(l0BTensor, l1BTensor, layoutBInL0, layoutBInL1);
	129	+
	130	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
	131	+
	132	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
	133	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
	134	+
	135	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
	136	+ tileMmad(l0CTensor, l0ATensor, l0BTensor, mRound, nRound, actualShape.k());
	137	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
	138	+
	139	+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
	140	+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
	141	+
	142	+ // copy block out
	143	+ auto blockShape = MakeCoord(actualShape.m(), actualShape.n());
	144	+ auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(blockShape);
	145	+
	146	+ copyL0CToGm(gC, l0CTensor, layoutC, layoutInL0C);
	147	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
	148	+ }
	149	+
	150	+protected:
	151	+ /// Data members
	152	+ AscendC::LocalTensor<ElementA> l1ATensor;
	153	+ AscendC::LocalTensor<ElementB> l1BTensor;
	154	+ AscendC::LocalTensor<ElementA> l0ATensor;
	155	+ AscendC::LocalTensor<ElementB> l0BTensor;
	156	+ AscendC::LocalTensor<ElementAccumulator> l0CTensor;
	157	+
	158	+ TileMmad tileMmad;
	159	+ CopyGmToL1A copyGmToL1A;
	160	+ CopyGmToL1B copyGmToL1B;
	161	+ CopyL1ToL0A copyL1ToL0A;
	162	+ CopyL1ToL0B copyL1ToL0B;
	163	+ CopyL0CToGm copyL0CToGm;
	164	+};
	165	+
	166	+////////////////////////////////////////////////////////////////////
	167	+
	168	+} // namespace Catlass::Gemm::Block
	169	+
	170	+#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_FA_QK_HPP

Ainclude/catlass/gemm/block/block_mmad_xa_shared_pv_tla.hpp+221-0

@@ -0,0 +1,221 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_XA_SHARED_PV_TLA_HPP
	13	+#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_XA_SHARED_PV_TLA_HPP
	14	+
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/arch/resource.hpp"
	17	+#include "catlass/coord.hpp"
	18	+#include "catlass/gemm/dispatch_policy.hpp"
	19	+#include "catlass/gemm/helper.hpp"
	20	+#include "catlass/gemm_coord.hpp"
	21	+#include "catlass/gemm/tile/tile_copy.hpp"
	22	+#include "catlass/gemm/tile/tile_mmad.hpp"
	23	+#include "tla/tensor.hpp"
	24	+#include "tla/layout.hpp"
	25	+
	26	+////////////////////////////////////////////////////////////////////
	27	+
	28	+namespace Catlass::Gemm::Block {
	29	+////////////////////////////////////////////////////////////////////
	30	+
	31	+template <
	32	+ class L1TileShape_, class L0TileShape_, class ElementA_, class ElementB_, class ElementC_, class ElementBias_,
	33	+ class TileCopy_, class TileMmad_>
	34	+struct BlockMmadTla<
	35	+ MmadXASharedPV<Arch::Ascend950>, L1TileShape_, L0TileShape_, ElementA_, ElementB_, ElementC_, ElementBias_,
	36	+ TileCopy_, TileMmad_> {
	37	+public:
	38	+ using DispatchPolicy = MmadXASharedPV<Arch::Ascend950>;
	39	+ using ArchTag = typename DispatchPolicy::ArchTag;
	40	+ using L1TileShape = L1TileShape_;
	41	+ using L0TileShape = L0TileShape_;
	42	+ using TileCopy = TileCopy_;
	43	+ using TileMmad = TileMmad_;
	44	+
	45	+ using ElementA = ElementA_;
	46	+ using ElementB = ElementB_;
	47	+ using ElementC = ElementC_;
	48	+ using ElementAccumulator = typename TileCopy::ElementAccumulator;
	49	+
	50	+ using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A;
	51	+ using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B;
	52	+
	53	+ using LayoutTagA = typename TileCopy::LayoutTagA;
	54	+ using LayoutTagB = typename TileCopy::LayoutTagB;
	55	+ using LayoutTagC = typename TileCopy::LayoutTagC;
	56	+
	57	+ using LayoutTagL1A = typename TileCopy::LayoutTagL1A;
	58	+ using LayoutTagL1B = typename TileCopy::LayoutTagL1B;
	59	+ using LayoutTagL0A = typename TileCopy::LayoutTagL0A;
	60	+ using LayoutTagL0B = typename TileCopy::LayoutTagL0B;
	61	+
	62	+ static constexpr uint32_t STAGES = DispatchPolicy::STAGES;
	63	+ // L1 tile shape
	64	+ static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); // s1
	65	+ static constexpr uint32_t L1_TILE_K = tla::get<1>(L1TileShape{}); // s2
	66	+ static constexpr uint32_t L1_TILE_N = tla::get<2>(L1TileShape{}); // d
	67	+ // L0 tile shape
	68	+ static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); // s1
	69	+ static constexpr uint32_t L0_TILE_K = tla::get<1>(L0TileShape{}); // s2
	70	+ static constexpr uint32_t L0_TILE_N = tla::get<2>(L0TileShape{}); // d
	71	+ // L1 tile size
	72	+ static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA);
	73	+ static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB);
	74	+ // L0 tile size
	75	+ static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA);
	76	+ static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB);
	77	+ static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator);
	78	+
	79	+ static constexpr uint32_t BLOCK_L1_SIZE = L1B_TILE_SIZE * STAGES;
	80	+ static constexpr uint32_t BLOCK_L0C_SIZE = L0C_TILE_SIZE * STAGES;
	81	+
	82	+ // Check L1/L0TileShape
	83	+ static_assert(
	84	+ L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N && L1_TILE_K == L0_TILE_K,
	85	+ "The situation where the basic blocks of L1 and L0 differ on the m, n, k axes is not supported yet");
	86	+ static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!");
	87	+ static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!");
	88	+
	89	+public:
	90	+ /// Construct
	91	+ CATLASS_DEVICE
	92	+ BlockMmadTla(Arch::Resource<ArchTag>& resource, uint32_t& l1BufAddrStart, uint32_t& l0CBufAddrStart)
	93	+ {
	94	+ for (uint32_t i = 0; i < STAGES; i++) {
	95	+ // Assign L1/L0A/L0B space for each stages
	96	+ l1BTensorList_[i] = resource.l1Buf.template GetBufferByByte<ElementB>(l1BufAddrStart);
	97	+ l1BufAddrStart += L1B_TILE_SIZE;
	98	+ l0ATensorList_[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(L0A_TILE_SIZE * i);
	99	+ l0BTensorList_[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(L0B_TILE_SIZE * i);
	100	+ l0CTensorList_[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(l0CBufAddrStart);
	101	+ l0CBufAddrStart += L0C_TILE_SIZE;
	102	+
	103	+ // Assign event ID for each stages
	104	+ l1BEventList_[i] = BLOCK_EVENT_ID + i + STAGES;
	105	+ l0AEventList_[i] = i;
	106	+ l0BEventList_[i] = i + STAGES;
	107	+ l0CEventList_[i] = BLOCK_EVENT_ID + i;
	108	+
	109	+ // The event id that needs to be set before the loop
	110	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[i]);
	111	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[i]);
	112	+ }
	113	+ }
	114	+
	115	+ /// Destructor
	116	+ CATLASS_DEVICE
	117	+ ~BlockMmadTla()
	118	+ {
	119	+ for (uint32_t i = 0; i < STAGES; i++) {
	120	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[i]);
	121	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[i]);
	122	+ }
	123	+ }
	124	+
	125	+public:
	126	+ template <class TensorA, class TensorB, class TensorC, class Shape>
	127	+ CATLASS_DEVICE void operator()(
	128	+ TensorA& tensorA, TensorB& tensorB, TensorC& tensorC, Shape& actualShape, uint32_t& taskIdL0A,
	129	+ uint32_t& taskIdL0B, uint16_t PV_RELEASE_FLAG)
	130	+ {
	131	+ uint32_t blockM = tla::get<0>(actualShape);
	132	+ uint32_t blockN = tla::get<1>(actualShape);
	133	+ uint32_t blockK = tla::get<2>(actualShape);
	134	+
	135	+ // Matrix A is already resident in L1. Only matrix B needs to be loaded from GM.
	136	+ auto layoutBInL1 = tla::MakeLayout<ElementB, LayoutTagL1B>(blockK, blockN);
	137	+ auto tensorBInL1 = tla::MakeTensor(l1BTensorList_[l1BListId_], layoutBInL1, Arch::PositionL1{});
	138	+ using CopyGmToL1B = typename TileCopy::template CopyGmToL1B<TensorB>;
	139	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[l1BListId_]);
	140	+ CopyGmToL1B copyGmToL1B;
	141	+ copyGmToL1B(tensorBInL1, tensorB);
	142	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList_[l1BListId_]);
	143	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList_[l1BListId_]);
	144	+
	145	+ auto layoutAInL0 = tla::MakeLayout<ElementA, LayoutTagL0A>(blockM, blockK);
	146	+ auto tensorAInL0 = tla::MakeTensor(l0ATensorList_[taskIdL0A], layoutAInL0, Arch::PositionL0A{});
	147	+ auto tensorAInL1 = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(blockM, blockK));
	148	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList_[taskIdL0A]);
	149	+ copyL1ToL0A(tensorAInL0, tensorAInL1);
	150	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0AEventList_[taskIdL0A]);
	151	+
	152	+ auto layoutBInL0 = tla::MakeLayout<ElementB, LayoutTagL0B>(blockK, blockN);
	153	+ auto tensorBInL0 = tla::MakeTensor(l0BTensorList_[taskIdL0B], layoutBInL0, Arch::PositionL0B{});
	154	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventList_[taskIdL0B]);
	155	+ auto tensorBL1Tile = GetTile(tensorBInL1, tla::MakeCoord(0, 0), tla::MakeShape(blockK, blockN));
	156	+ copyL1ToL0B(tensorBInL0, tensorBL1Tile);
	157	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0BEventList_[taskIdL0B]);
	158	+
	159	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0AEventList_[taskIdL0A]);
	160	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0BEventList_[taskIdL0B]);
	161	+
	162	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[l1BListId_]);
	163	+ l1BListId_ = 1 - l1BListId_;
	164	+
	165	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[l0CListId_]);
	166	+ auto layoutInL0C = tla::MakeLayoutL0C(blockM, blockN);
	167	+ auto tensorInL0C = tla::MakeTensor(l0CTensorList_[l0CListId_], layoutInL0C, Arch::PositionL0C{});
	168	+
	169	+ tileMmad(tensorInL0C, tensorAInL0, tensorBInL0, blockM, blockN, blockK, true);
	170	+
	171	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList_[taskIdL0A]);
	172	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventList_[taskIdL0B]);
	173	+ taskIdL0A = 1 - taskIdL0A;
	174	+ taskIdL0B = 1 - taskIdL0B;
	175	+
	176	+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(l0CEventList_[l0CListId_]);
	177	+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(l0CEventList_[l0CListId_]);
	178	+
	179	+ AscendC::CrossCoreWaitFlag<SYNC_MODE, PIPE_FIX>(PV_RELEASE_FLAG);
	180	+ AscendC::CrossCoreWaitFlag<SYNC_MODE, PIPE_FIX>(PV_RELEASE_FLAG + AIV1_EVENT_OFFSET);
	181	+
	182	+ using CopyL0CToDst = typename TileCopy::template CopyL0CToDst<TensorC>;
	183	+ CopyL0CToDst copyL0CToDst;
	184	+ copyL0CToDst(tensorC, tensorInL0C);
	185	+
	186	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[l0CListId_]);
	187	+ l0CListId_ = 1 - l0CListId_;
	188	+ }
	189	+
	190	+private:
	191	+ static constexpr uint16_t M_ALIGN = 16;
	192	+ static constexpr uint32_t SYNC_MODE = 4;
	193	+ static constexpr uint32_t AIV1_EVENT_OFFSET = 16;
	194	+ static constexpr uint32_t BLOCK_EVENT_ID = 4;
	195	+
	196	+ // Multi-stage tensors list
	197	+ AscendC::LocalTensor<ElementA> l1ATensorList_[STAGES];
	198	+ AscendC::LocalTensor<ElementB> l1BTensorList_[STAGES];
	199	+ AscendC::LocalTensor<ElementA> l0ATensorList_[STAGES];
	200	+ AscendC::LocalTensor<ElementB> l0BTensorList_[STAGES];
	201	+ AscendC::LocalTensor<ElementAccumulator> l0CTensorList_[STAGES];
	202	+
	203	+ // Multi-stage event id list
	204	+ int32_t l1AEventList_[STAGES];
	205	+ int32_t l1BEventList_[STAGES];
	206	+ int32_t l0AEventList_[STAGES];
	207	+ int32_t l0BEventList_[STAGES];
	208	+ int32_t l0CEventList_[STAGES];
	209	+
	210	+ uint32_t l1AListId_{0};
	211	+ uint32_t l1BListId_{0};
	212	+ uint32_t l0AListId_{0};
	213	+ uint32_t l0BListId_{0};
	214	+ uint32_t l0CListId_{0};
	215	+
	216	+ CopyL1ToL0A copyL1ToL0A;
	217	+ CopyL1ToL0B copyL1ToL0B;
	218	+ TileMmad tileMmad;
	219	+};
	220	+} // namespace Catlass::Gemm::Block
	221	+#endif

Ainclude/catlass/gemm/block/block_mmad_xa_shared_qk_tla.hpp+244-0

@@ -0,0 +1,244 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_XA_SHARED_QK_TLA_HPP
	13	+#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_XA_SHARED_QK_TLA_HPP
	14	+
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/arch/resource.hpp"
	17	+#include "catlass/coord.hpp"
	18	+#include "catlass/gemm/dispatch_policy.hpp"
	19	+#include "catlass/gemm/helper.hpp"
	20	+#include "catlass/gemm_coord.hpp"
	21	+#include "catlass/gemm/tile/tile_copy.hpp"
	22	+#include "catlass/gemm/tile/tile_mmad.hpp"
	23	+#include "tla/tensor.hpp"
	24	+#include "tla/layout.hpp"
	25	+
	26	+////////////////////////////////////////////////////////////////////
	27	+
	28	+namespace Catlass::Gemm::Block {
	29	+////////////////////////////////////////////////////////////////////
	30	+
	31	+template <
	32	+ class L1TileShape_, class L0TileShape_, class ElementA_, class ElementB_, class ElementC_, class ElementBias_,
	33	+ class TileCopy_, class TileMmad_>
	34	+struct BlockMmadTla<
	35	+ MmadXASharedQK<Arch::Ascend950>, L1TileShape_, L0TileShape_, ElementA_, ElementB_, ElementC_, ElementBias_,
	36	+ TileCopy_, TileMmad_> {
	37	+public:
	38	+ using DispatchPolicy = MmadXASharedQK<Arch::Ascend950>;
	39	+ using ArchTag = typename DispatchPolicy::ArchTag;
	40	+ using L1TileShape = L1TileShape_;
	41	+ using L0TileShape = L0TileShape_;
	42	+ using TileCopy = TileCopy_;
	43	+ using TileMmad = TileMmad_;
	44	+
	45	+ using ElementA = ElementA_;
	46	+ using ElementB = ElementB_;
	47	+ using ElementC = ElementC_;
	48	+ using ElementAccumulator = typename TileCopy::ElementAccumulator;
	49	+
	50	+ using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A;
	51	+ using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B;
	52	+
	53	+ using LayoutTagA = typename TileCopy::LayoutTagA;
	54	+ using LayoutTagB = typename TileCopy::LayoutTagB;
	55	+ using LayoutTagC = typename TileCopy::LayoutTagC;
	56	+
	57	+ using LayoutTagL1A = typename TileCopy::LayoutTagL1A;
	58	+ using LayoutTagL1B = typename TileCopy::LayoutTagL1B;
	59	+ using LayoutTagL0A = typename TileCopy::LayoutTagL0A;
	60	+ using LayoutTagL0B = typename TileCopy::LayoutTagL0B;
	61	+
	62	+ static constexpr uint32_t STAGES = DispatchPolicy::STAGES;
	63	+ // L1 tile shape
	64	+ static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); // s1
	65	+ static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); // s2
	66	+ static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); // d
	67	+ // L0 tile shape
	68	+ static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); // s1
	69	+ static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); // s2
	70	+ static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); // d
	71	+ // L1 tile size
	72	+ static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA);
	73	+ static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB);
	74	+ // L0 tile size
	75	+ static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA);
	76	+ static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB);
	77	+ static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator);
	78	+
	79	+ static constexpr uint32_t BLOCK_L1_SIZE = (L1A_TILE_SIZE + L1B_TILE_SIZE) * STAGES;
	80	+ static constexpr uint32_t BLOCK_L0C_SIZE = L0C_TILE_SIZE * STAGES;
	81	+
	82	+ // Check L1/L0TileShape
	83	+ static_assert(
	84	+ L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N && L1_TILE_K == L0_TILE_K,
	85	+ "The situation where the basic blocks of L1 and L0 differ on the m, n, k axes is not supported yet");
	86	+ static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!");
	87	+ static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!");
	88	+
	89	+public:
	90	+ /// Construct
	91	+ CATLASS_DEVICE
	92	+ BlockMmadTla(Arch::Resource<ArchTag>& resource, uint32_t& l1BufAddrStart, uint32_t& l0CBufAddrStart)
	93	+ {
	94	+ for (uint32_t i = 0; i < STAGES; i++) {
	95	+ // Assign L1/L0A/L0B space for each stages
	96	+ l1ATensorList_[i] = resource.l1Buf.template GetBufferByByte<ElementA>(l1BufAddrStart);
	97	+ l1BufAddrStart += L1A_TILE_SIZE;
	98	+ l1BTensorList_[i] = resource.l1Buf.template GetBufferByByte<ElementB>(l1BufAddrStart);
	99	+ l1BufAddrStart += L1B_TILE_SIZE;
	100	+ l0ATensorList_[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(L0A_TILE_SIZE * i);
	101	+ l0BTensorList_[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(L0B_TILE_SIZE * i);
	102	+ l0CTensorList_[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(l0CBufAddrStart);
	103	+ l0CBufAddrStart += L0C_TILE_SIZE;
	104	+
	105	+ // Assign event ID for each stages
	106	+ l1AEventList_[i] = BLOCK_EVENT_ID + i;
	107	+ l1BEventList_[i] = BLOCK_EVENT_ID + i + STAGES;
	108	+ l0AEventList_[i] = i;
	109	+ l0BEventList_[i] = i + STAGES;
	110	+ l0CEventList_[i] = BLOCK_EVENT_ID + i;
	111	+
	112	+ // The event id that needs to be set before the loop
	113	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList_[i]);
	114	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[i]);
	115	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[i]);
	116	+ }
	117	+ }
	118	+
	119	+ /// Destructor
	120	+ CATLASS_DEVICE
	121	+ ~BlockMmadTla()
	122	+ {
	123	+ for (uint32_t i = 0; i < STAGES; i++) {
	124	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList_[i]);
	125	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[i]);
	126	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[i]);
	127	+ }
	128	+ }
	129	+
	130	+public:
	131	+ template <class TensorA, class TensorB, class TensorC, class Shape>
	132	+ CATLASS_DEVICE void operator()(
	133	+ TensorA& tensorA, TensorB& tensorB, TensorC& tensorC, Shape& actualShape, uint64_t QK_UB_RELEASE_FLAG,
	134	+ bool isFirstKv, bool isLastKv, uint32_t& taskIdL0A, uint32_t& taskIdL0B)
	135	+ {
	136	+ uint32_t qActual = tla::get<0>(actualShape);
	137	+ uint32_t kActual = tla::get<1>(actualShape);
	138	+ uint32_t dimActual = tla::get<2>(actualShape);
	139	+
	140	+ auto layoutAInL1 = tla::MakeLayout<ElementA, LayoutTagL1A>(qActual, dimActual);
	141	+ auto tensorAInL1 = tla::MakeTensor(l1ATensorList_[l1AListId_], layoutAInL1, Arch::PositionL1{});
	142	+
	143	+ if (unlikely(isFirstKv)) {
	144	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList_[l1AListId_]);
	145	+ using CopyGmToL1A = typename TileCopy::template CopyGmToL1A<TensorA>;
	146	+ CopyGmToL1A copyGmToL1A;
	147	+ copyGmToL1A(tensorAInL1, tensorA);
	148	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList_[l1AListId_]);
	149	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList_[l1AListId_]);
	150	+ }
	151	+
	152	+ auto layoutBInL1 = tla::MakeLayout<ElementB, LayoutTagL1B>(dimActual, kActual);
	153	+ auto tensorBInL1 = tla::MakeTensor(l1BTensorList_[l1BListId_], layoutBInL1, Arch::PositionL1{});
	154	+ using CopyGmToL1B = typename TileCopy::template CopyGmToL1B<TensorB>;
	155	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[l1BListId_]);
	156	+ CopyGmToL1B copyGmToL1B;
	157	+ copyGmToL1B(tensorBInL1, tensorB);
	158	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList_[l1BListId_]);
	159	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList_[l1BListId_]);
	160	+
	161	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[l0CListId_]);
	162	+ auto layoutInL0C = tla::MakeLayoutL0C(qActual, kActual);
	163	+ auto tensorInL0C = tla::MakeTensor(l0CTensorList_[l0CListId_], layoutInL0C, Arch::PositionL0C{});
	164	+
	165	+ auto layoutAInL0 = tla::MakeLayout<ElementA, LayoutTagL0A>(qActual, dimActual);
	166	+ auto tensorAInL0 = tla::MakeTensor(l0ATensorList_[taskIdL0A], layoutAInL0, Arch::PositionL0A{});
	167	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList_[taskIdL0A]);
	168	+
	169	+ auto tensorAL1Tile = GetTile(tensorAInL1, tla::MakeCoord(0, 0), tla::MakeShape(qActual, dimActual));
	170	+ copyL1ToL0A(tensorAInL0, tensorAL1Tile);
	171	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0AEventList_[taskIdL0A]);
	172	+
	173	+ auto layoutBInL0 = tla::MakeLayout<ElementB, LayoutTagL0B>(dimActual, kActual);
	174	+ auto tensorBInL0 = tla::MakeTensor(l0BTensorList_[taskIdL0B], layoutBInL0, Arch::PositionL0B{});
	175	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventList_[taskIdL0B]);
	176	+ auto tensorBL1Tile = GetTile(tensorBInL1, tla::MakeCoord(0, 0), tla::MakeShape(dimActual, kActual));
	177	+ copyL1ToL0B(tensorBInL0, tensorBL1Tile);
	178	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0BEventList_[taskIdL0B]);
	179	+
	180	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0AEventList_[taskIdL0A]);
	181	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0BEventList_[taskIdL0B]);
	182	+
	183	+ // MTE1 has consumed the L1 tiles, so their buffers can be reused by MTE2.
	184	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[l1BListId_]);
	185	+ l1BListId_ = 1 - l1BListId_;
	186	+
	187	+ if (unlikely(isLastKv)) {
	188	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList_[l1AListId_]);
	189	+ l1AListId_ = 1 - l1AListId_;
	190	+ }
	191	+
	192	+ tileMmad(tensorInL0C, tensorAInL0, tensorBInL0, qActual, kActual, dimActual, true);
	193	+
	194	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList_[taskIdL0A]);
	195	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventList_[taskIdL0B]);
	196	+ taskIdL0A = 1 - taskIdL0A;
	197	+ taskIdL0B = 1 - taskIdL0B;
	198	+
	199	+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(l0CEventList_[l0CListId_]);
	200	+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(l0CEventList_[l0CListId_]);
	201	+
	202	+ // The two AIVs must both release the shared UB tile before FixPipe overwrites it.
	203	+ AscendC::CrossCoreWaitFlag<SYNC_MODE, PIPE_FIX>(QK_UB_RELEASE_FLAG);
	204	+ AscendC::CrossCoreWaitFlag<SYNC_MODE, PIPE_FIX>(QK_UB_RELEASE_FLAG + AIV1_EVENT_OFFSET);
	205	+
	206	+ using CopyL0CToDst = typename TileCopy::template CopyL0CToDst<TensorC>;
	207	+ CopyL0CToDst copyL0CToDst;
	208	+ copyL0CToDst(tensorC, tensorInL0C);
	209	+
	210	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[l0CListId_]);
	211	+ l0CListId_ = 1 - l0CListId_;
	212	+ }
	213	+
	214	+private:
	215	+ static constexpr uint16_t M_ALIGN = 16;
	216	+ static constexpr uint32_t SYNC_MODE = 4;
	217	+ static constexpr uint32_t AIV1_EVENT_OFFSET = 16;
	218	+ static constexpr uint32_t BLOCK_EVENT_ID = 0;
	219	+
	220	+ // Multi-stage tensors list
	221	+ AscendC::LocalTensor<ElementA> l1ATensorList_[STAGES];
	222	+ AscendC::LocalTensor<ElementB> l1BTensorList_[STAGES];
	223	+ AscendC::LocalTensor<ElementA> l0ATensorList_[STAGES];
	224	+ AscendC::LocalTensor<ElementB> l0BTensorList_[STAGES];
	225	+ AscendC::LocalTensor<ElementAccumulator> l0CTensorList_[STAGES];
	226	+
	227	+ // Multi-stage event id list
	228	+ int32_t l1AEventList_[STAGES];
	229	+ int32_t l1BEventList_[STAGES];
	230	+ int32_t l0AEventList_[STAGES];
	231	+ int32_t l0BEventList_[STAGES];
	232	+ int32_t l0CEventList_[STAGES];
	233	+
	234	+ uint32_t l1AListId_{0};
	235	+ uint32_t l1BListId_{0};
	236	+ uint32_t l0ListId_{0};
	237	+ uint32_t l0CListId_{0};
	238	+
	239	+ CopyL1ToL0A copyL1ToL0A;
	240	+ CopyL1ToL0B copyL1ToL0B;
	241	+ TileMmad tileMmad;
	242	+};
	243	+} // namespace Catlass::Gemm::Block
	244	+#endif

Ainclude/catlass/gemm/block/block_mmad_xa_unshared_pv_tla.hpp+205-0

@@ -0,0 +1,205 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_XA_UNSHARED_PV_TLA_HPP
	13	+#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_XA_UNSHARED_PV_TLA_HPP
	14	+
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/arch/resource.hpp"
	17	+#include "catlass/coord.hpp"
	18	+#include "catlass/gemm/dispatch_policy.hpp"
	19	+#include "catlass/gemm/helper.hpp"
	20	+#include "catlass/gemm_coord.hpp"
	21	+#include "catlass/gemm/tile/tile_copy.hpp"
	22	+#include "catlass/gemm/tile/tile_mmad.hpp"
	23	+#include "tla/layout.hpp"
	24	+#include "tla/tensor.hpp"
	25	+
	26	+namespace Catlass::Gemm::Block {
	27	+
	28	+template <
	29	+ class L1TileShape_, class L0TileShape_, class ElementA_, class ElementB_, class ElementC_, class ElementBias_,
	30	+ class TileCopy_, class TileMmad_>
	31	+struct BlockMmadTla<
	32	+ MmadXAUnsharedPV<Arch::Ascend950>, L1TileShape_, L0TileShape_, ElementA_, ElementB_, ElementC_, ElementBias_,
	33	+ TileCopy_, TileMmad_> {
	34	+public:
	35	+ using DispatchPolicy = MmadXAUnsharedPV<Arch::Ascend950>;
	36	+ using ArchTag = typename DispatchPolicy::ArchTag;
	37	+ using L1TileShape = L1TileShape_;
	38	+ using L0TileShape = L0TileShape_;
	39	+ using TileCopy = TileCopy_;
	40	+ using TileMmad = TileMmad_;
	41	+
	42	+ using ElementA = ElementA_;
	43	+ using ElementB = ElementB_;
	44	+ using ElementC = ElementC_;
	45	+ using ElementAccumulator = typename TileCopy::ElementAccumulator;
	46	+
	47	+ using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A;
	48	+ using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B;
	49	+
	50	+ using LayoutTagA = typename TileCopy::LayoutTagA;
	51	+ using LayoutTagB = typename TileCopy::LayoutTagB;
	52	+ using LayoutTagC = typename TileCopy::LayoutTagC;
	53	+ using LayoutTagL1A = typename TileCopy::LayoutTagL1A;
	54	+ using LayoutTagL1B = typename TileCopy::LayoutTagL1B;
	55	+ using LayoutTagL0A = typename TileCopy::LayoutTagL0A;
	56	+ using LayoutTagL0B = typename TileCopy::LayoutTagL0B;
	57	+
	58	+ static constexpr uint32_t STAGES = DispatchPolicy::STAGES;
	59	+ // PV shape is M x N x K. QK's N axis maps to PV's K axis.
	60	+ static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{});
	61	+ static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{});
	62	+ static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{});
	63	+ static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{});
	64	+ static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{});
	65	+ static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{});
	66	+
	67	+ static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_K * L1_TILE_N * sizeof(ElementB);
	68	+ static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA);
	69	+ static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB);
	70	+ static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L0_TILE_N * sizeof(ElementAccumulator);
	71	+ // QK owns a 128 x 256 L0C stage. PV reuses its first 128 x 128 result region.
	72	+ static constexpr uint32_t SHARED_L0C_STAGE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementAccumulator);
	73	+
	74	+ static constexpr uint32_t BLOCK_L1_SIZE = L1B_TILE_SIZE * STAGES;
	75	+ static constexpr uint32_t BLOCK_L0C_SIZE = SHARED_L0C_STAGE_SIZE * STAGES;
	76	+
	77	+ static_assert(
	78	+ L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, "L1 and L0 tile shapes must match on the M and N axes");
	79	+ static_assert(
	80	+ L1_TILE_K >= L0_TILE_K && L1_TILE_K % L0_TILE_K == 0,
	81	+ "The L1 K tile must be an integer multiple of the L0 K tile");
	82	+ static_assert(BLOCK_L1_SIZE <= ArchTag::L1_SIZE, "L1TileShape exceeding the L1 space!");
	83	+ static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!");
	84	+ static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!");
	85	+ static_assert(BLOCK_L0C_SIZE <= ArchTag::L0C_SIZE, "Shared L0C buffers exceeding the L0C space!");
	86	+
	87	+public:
	88	+ CATLASS_DEVICE
	89	+ BlockMmadTla(Arch::Resource<ArchTag>& resource, uint32_t& l1BufAddrStart, uint32_t& l0CBufAddrStart)
	90	+ {
	91	+ for (uint32_t i = 0; i < STAGES; ++i) {
	92	+ l1BTensorList_[i] = resource.l1Buf.template GetBufferByByte<ElementB>(l1BufAddrStart);
	93	+ l1BufAddrStart += L1B_TILE_SIZE;
	94	+ l0ATensorList_[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(L0A_TILE_SIZE * i);
	95	+ l0BTensorList_[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(L0B_TILE_SIZE * i);
	96	+ l0CTensorList_[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(
	97	+ l0CBufAddrStart + SHARED_L0C_STAGE_SIZE * i);
	98	+
	99	+ l1BEventList_[i] = BLOCK_EVENT_ID + i + STAGES;
	100	+ l0AEventList_[i] = i;
	101	+ l0BEventList_[i] = i + STAGES;
	102	+ l0CEventList_[i] = i;
	103	+
	104	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[i]);
	105	+ }
	106	+ }
	107	+
	108	+ CATLASS_DEVICE
	109	+ ~BlockMmadTla()
	110	+ {
	111	+ for (uint32_t i = 0; i < STAGES; ++i) {
	112	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[i]);
	113	+ }
	114	+ }
	115	+
	116	+ template <class TensorA, class TensorB, class TensorC, class Shape>
	117	+ CATLASS_DEVICE void operator()(
	118	+ TensorA& tensorA, TensorB& tensorB, TensorC& tensorC, Shape& actualShape, uint32_t& taskIdL0A,
	119	+ uint32_t& taskIdL0B, uint32_t& taskIdL0C)
	120	+ {
	121	+ uint32_t mActual = tla::get<0>(actualShape);
	122	+ uint32_t nActual = tla::get<1>(actualShape);
	123	+ uint32_t kActual = tla::get<2>(actualShape);
	124	+
	125	+ auto layoutBInL1 = tla::MakeLayout<ElementB, LayoutTagL1B>(kActual, nActual);
	126	+ auto tensorBInL1 = tla::MakeTensor(l1BTensorList_[l1BListId_], layoutBInL1, Arch::PositionL1{});
	127	+ using CopyGmToL1B = typename TileCopy::template CopyGmToL1B<TensorB>;
	128	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[l1BListId_]);
	129	+ CopyGmToL1B copyGmToL1B;
	130	+ copyGmToL1B(tensorBInL1, tensorB);
	131	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList_[l1BListId_]);
	132	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList_[l1BListId_]);
	133	+
	134	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[taskIdL0C]);
	135	+ auto layoutInL0C = tla::MakeLayoutL0C(mActual, nActual);
	136	+ auto tensorInL0C = tla::MakeTensor(l0CTensorList_[taskIdL0C], layoutInL0C, Arch::PositionL0C{});
	137	+
	138	+ uint32_t kLoops = (kActual + L0_TILE_K - 1) / L0_TILE_K;
	139	+ for (uint32_t kIdx = 0; kIdx < kLoops; ++kIdx) {
	140	+ uint32_t kOffset = kIdx * L0_TILE_K;
	141	+ uint32_t tileK = kIdx + 1 == kLoops ? kActual - kOffset : L0_TILE_K;
	142	+
	143	+ auto tensorATileInL1 = GetTile(tensorA, tla::MakeCoord(0, kOffset), tla::MakeShape(mActual, tileK));
	144	+ auto layoutAInL0 = tla::MakeLayout<ElementA, LayoutTagL0A>(mActual, tileK);
	145	+ auto tensorAInL0 = tla::MakeTensor(l0ATensorList_[taskIdL0A], layoutAInL0, Arch::PositionL0A{});
	146	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList_[taskIdL0A]);
	147	+ copyL1ToL0A_(tensorAInL0, tensorATileInL1);
	148	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0AEventList_[taskIdL0A]);
	149	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0AEventList_[taskIdL0A]);
	150	+
	151	+ auto tensorBTileInL1 = GetTile(tensorBInL1, tla::MakeCoord(kOffset, 0), tla::MakeShape(tileK, nActual));
	152	+ auto layoutBInL0 = tla::MakeLayout<ElementB, LayoutTagL0B>(tileK, nActual);
	153	+ auto tensorBInL0 = tla::MakeTensor(l0BTensorList_[taskIdL0B], layoutBInL0, Arch::PositionL0B{});
	154	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventList_[taskIdL0B]);
	155	+ copyL1ToL0B_(tensorBInL0, tensorBTileInL1);
	156	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0BEventList_[taskIdL0B]);
	157	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0BEventList_[taskIdL0B]);
	158	+
	159	+ tileMmad_(tensorInL0C, tensorAInL0, tensorBInL0, mActual, nActual, tileK, kIdx == 0);
	160	+
	161	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList_[taskIdL0A]);
	162	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventList_[taskIdL0B]);
	163	+ taskIdL0A = 1 - taskIdL0A;
	164	+ taskIdL0B = 1 - taskIdL0B;
	165	+ }
	166	+
	167	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[l1BListId_]);
	168	+ l1BListId_ = 1 - l1BListId_;
	169	+
	170	+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(l0CEventList_[taskIdL0C]);
	171	+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(l0CEventList_[taskIdL0C]);
	172	+
	173	+ using CopyL0CToDst = typename TileCopy::template CopyL0CToDst<TensorC>;
	174	+ CopyL0CToDst copyL0CToDst;
	175	+ copyL0CToDst(tensorC, tensorInL0C);
	176	+
	177	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[taskIdL0C]);
	178	+ taskIdL0C = 1 - taskIdL0C;
	179	+ }
	180	+
	181	+private:
	182	+ static constexpr uint32_t SYNC_MODE = 4;
	183	+ static constexpr uint32_t AIV1_EVENT_OFFSET = 16;
	184	+ static constexpr uint32_t BLOCK_EVENT_ID = 4;
	185	+
	186	+ AscendC::LocalTensor<ElementB> l1BTensorList_[STAGES];
	187	+ AscendC::LocalTensor<ElementA> l0ATensorList_[STAGES];
	188	+ AscendC::LocalTensor<ElementB> l0BTensorList_[STAGES];
	189	+ AscendC::LocalTensor<ElementAccumulator> l0CTensorList_[STAGES];
	190	+
	191	+ int32_t l1BEventList_[STAGES];
	192	+ int32_t l0AEventList_[STAGES];
	193	+ int32_t l0BEventList_[STAGES];
	194	+ int32_t l0CEventList_[STAGES];
	195	+
	196	+ uint32_t l1BListId_{0};
	197	+
	198	+ CopyL1ToL0A copyL1ToL0A_;
	199	+ CopyL1ToL0B copyL1ToL0B_;
	200	+ TileMmad tileMmad_;
	201	+};
	202	+
	203	+} // namespace Catlass::Gemm::Block
	204	+
	205	+#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_XA_UNSHARED_PV_TLA_HPP

Ainclude/catlass/gemm/block/block_mmad_xa_unshared_qk_tla.hpp+229-0

@@ -0,0 +1,229 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_XA_UNSHARED_QK_TLA_HPP
	13	+#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_XA_UNSHARED_QK_TLA_HPP
	14	+
	15	+#include "catlass/catlass.hpp"
	16	+#include "catlass/arch/resource.hpp"
	17	+#include "catlass/coord.hpp"
	18	+#include "catlass/gemm/dispatch_policy.hpp"
	19	+#include "catlass/gemm/helper.hpp"
	20	+#include "catlass/gemm_coord.hpp"
	21	+#include "catlass/gemm/tile/tile_copy.hpp"
	22	+#include "catlass/gemm/tile/tile_mmad.hpp"
	23	+#include "tla/layout.hpp"
	24	+#include "tla/tensor.hpp"
	25	+
	26	+namespace Catlass::Gemm::Block {
	27	+
	28	+template <
	29	+ class L1TileShape_, class L0TileShape_, class ElementA_, class ElementB_, class ElementC_, class ElementBias_,
	30	+ class TileCopy_, class TileMmad_>
	31	+struct BlockMmadTla<
	32	+ MmadXAUnsharedQK<Arch::Ascend950>, L1TileShape_, L0TileShape_, ElementA_, ElementB_, ElementC_, ElementBias_,
	33	+ TileCopy_, TileMmad_> {
	34	+public:
	35	+ using DispatchPolicy = MmadXAUnsharedQK<Arch::Ascend950>;
	36	+ using ArchTag = typename DispatchPolicy::ArchTag;
	37	+ using L1TileShape = L1TileShape_;
	38	+ using L0TileShape = L0TileShape_;
	39	+ using TileCopy = TileCopy_;
	40	+ using TileMmad = TileMmad_;
	41	+
	42	+ using ElementA = ElementA_;
	43	+ using ElementB = ElementB_;
	44	+ using ElementC = ElementC_;
	45	+ using ElementAccumulator = typename TileCopy::ElementAccumulator;
	46	+
	47	+ using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A;
	48	+ using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B;
	49	+
	50	+ using LayoutTagA = typename TileCopy::LayoutTagA;
	51	+ using LayoutTagB = typename TileCopy::LayoutTagB;
	52	+ using LayoutTagC = typename TileCopy::LayoutTagC;
	53	+ using LayoutTagL1A = typename TileCopy::LayoutTagL1A;
	54	+ using LayoutTagL1B = typename TileCopy::LayoutTagL1B;
	55	+ using LayoutTagL0A = typename TileCopy::LayoutTagL0A;
	56	+ using LayoutTagL0B = typename TileCopy::LayoutTagL0B;
	57	+
	58	+ static constexpr uint32_t STAGES = DispatchPolicy::STAGES;
	59	+ static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{});
	60	+ static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{});
	61	+ static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{});
	62	+ static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{});
	63	+ static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{});
	64	+ static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{});
	65	+
	66	+ static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA);
	67	+ static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB);
	68	+ static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA);
	69	+ static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB);
	70	+ static constexpr uint32_t L0C_TILE_SIZE = L0_TILE_M * L1_TILE_N * sizeof(ElementAccumulator);
	71	+
	72	+ static constexpr uint32_t BLOCK_L1_SIZE = (L1A_TILE_SIZE + L1B_TILE_SIZE) * STAGES;
	73	+ static constexpr uint32_t BLOCK_L0C_SIZE = L0C_TILE_SIZE * STAGES;
	74	+
	75	+ static_assert(
	76	+ L1_TILE_M == L0_TILE_M && L1_TILE_K == L0_TILE_K, "L1 and L0 tile shapes must match on the M and K axes");
	77	+ static_assert(
	78	+ L1_TILE_N >= L0_TILE_N && L1_TILE_N % L0_TILE_N == 0,
	79	+ "The L1 N tile must be an integer multiple of the L0 N tile");
	80	+ static_assert(BLOCK_L1_SIZE <= ArchTag::L1_SIZE, "L1TileShape exceeding the L1 space!");
	81	+ static_assert(L0A_TILE_SIZE * STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!");
	82	+ static_assert(L0B_TILE_SIZE * STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!");
	83	+ static_assert(BLOCK_L0C_SIZE <= ArchTag::L0C_SIZE, "L0C tile buffers exceeding the L0C space!");
	84	+
	85	+public:
	86	+ CATLASS_DEVICE
	87	+ BlockMmadTla(Arch::Resource<ArchTag>& resource, uint32_t& l1BufAddrStart, uint32_t& l0CBufAddrStart)
	88	+ {
	89	+ for (uint32_t i = 0; i < STAGES; ++i) {
	90	+ l1ATensorList_[i] = resource.l1Buf.template GetBufferByByte<ElementA>(l1BufAddrStart);
	91	+ l1BufAddrStart += L1A_TILE_SIZE;
	92	+ l1BTensorList_[i] = resource.l1Buf.template GetBufferByByte<ElementB>(l1BufAddrStart);
	93	+ l1BufAddrStart += L1B_TILE_SIZE;
	94	+ l0ATensorList_[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(L0A_TILE_SIZE * i);
	95	+ l0BTensorList_[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(L0B_TILE_SIZE * i);
	96	+ l0CTensorList_[i] =
	97	+ resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(l0CBufAddrStart + L0C_TILE_SIZE * i);
	98	+
	99	+ l1AEventList_[i] = BLOCK_EVENT_ID + i;
	100	+ l1BEventList_[i] = BLOCK_EVENT_ID + i + STAGES;
	101	+ l0AEventList_[i] = i;
	102	+ l0BEventList_[i] = i + STAGES;
	103	+ l0CEventList_[i] = BLOCK_EVENT_ID + i;
	104	+
	105	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList_[i]);
	106	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[i]);
	107	+ }
	108	+ }
	109	+
	110	+ CATLASS_DEVICE
	111	+ ~BlockMmadTla()
	112	+ {
	113	+ for (uint32_t i = 0; i < STAGES; ++i) {
	114	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList_[i]);
	115	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[i]);
	116	+ }
	117	+ }
	118	+
	119	+ template <class TensorA, class TensorB, class TensorC, class Shape>
	120	+ CATLASS_DEVICE void operator()(
	121	+ TensorA& tensorA, TensorB& tensorB, TensorC& tensorC, Shape& actualShape, uint64_t QK_UB_RELEASE_FLAG,
	122	+ uint32_t& taskIdL0A, uint32_t& taskIdL0B, uint32_t& taskIdL0C)
	123	+ {
	124	+ uint32_t qActual = tla::get<0>(actualShape);
	125	+ uint32_t kActual = tla::get<1>(actualShape);
	126	+ uint32_t dimActual = tla::get<2>(actualShape);
	127	+
	128	+ auto layoutAInL1 = tla::MakeLayout<ElementA, LayoutTagL1A>(qActual, dimActual);
	129	+ auto tensorAInL1 = tla::MakeTensor(l1ATensorList_[l1AListId_], layoutAInL1, Arch::PositionL1{});
	130	+ using CopyGmToL1A = typename TileCopy::template CopyGmToL1A<TensorA>;
	131	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList_[l1AListId_]);
	132	+ CopyGmToL1A copyGmToL1A;
	133	+ copyGmToL1A(tensorAInL1, tensorA);
	134	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList_[l1AListId_]);
	135	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventList_[l1AListId_]);
	136	+
	137	+ auto layoutBInL1 = tla::MakeLayout<ElementB, LayoutTagL1B>(dimActual, kActual);
	138	+ auto tensorBInL1 = tla::MakeTensor(l1BTensorList_[l1BListId_], layoutBInL1, Arch::PositionL1{});
	139	+ using CopyGmToL1B = typename TileCopy::template CopyGmToL1B<TensorB>;
	140	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[l1BListId_]);
	141	+ CopyGmToL1B copyGmToL1B;
	142	+ copyGmToL1B(tensorBInL1, tensorB);
	143	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList_[l1BListId_]);
	144	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList_[l1BListId_]);
	145	+
	146	+ auto layoutAInL0 = tla::MakeLayout<ElementA, LayoutTagL0A>(qActual, dimActual);
	147	+ auto tensorAInL0 = tla::MakeTensor(l0ATensorList_[taskIdL0A], layoutAInL0, Arch::PositionL0A{});
	148	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList_[taskIdL0A]);
	149	+ auto tensorAL1Tile = GetTile(tensorAInL1, tla::MakeCoord(0, 0), tla::MakeShape(qActual, dimActual));
	150	+
	151	+ copyL1ToL0A_(tensorAInL0, tensorAL1Tile);
	152	+
	153	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventList_[l1AListId_]);
	154	+ l1AListId_ = 1 - l1AListId_;
	155	+
	156	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0AEventList_[taskIdL0A]);
	157	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0AEventList_[taskIdL0A]);
	158	+
	159	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[taskIdL0C]);
	160	+ auto layoutInL0C = tla::MakeLayoutL0C(qActual, kActual);
	161	+ auto tensorInL0C = tla::MakeTensor(l0CTensorList_[taskIdL0C], layoutInL0C, Arch::PositionL0C{});
	162	+
	163	+ uint32_t nLoops = (kActual + L0_TILE_N - 1) / L0_TILE_N;
	164	+ for (uint32_t nIdx = 0; nIdx < nLoops; ++nIdx) {
	165	+ uint32_t nOffset = nIdx * L0_TILE_N;
	166	+ uint32_t tileN = nIdx + 1 == nLoops ? kActual - nOffset : L0_TILE_N;
	167	+
	168	+ auto layoutBInL0 = tla::MakeLayout<ElementB, LayoutTagL0B>(dimActual, tileN);
	169	+ auto tensorBInL0 = tla::MakeTensor(l0BTensorList_[taskIdL0B], layoutBInL0, Arch::PositionL0B{});
	170	+ auto tensorBTileInL1 = GetTile(tensorBInL1, tla::MakeCoord(0, nOffset), tla::MakeShape(dimActual, tileN));
	171	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventList_[taskIdL0B]);
	172	+ copyL1ToL0B_(tensorBInL0, tensorBTileInL1);
	173	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0BEventList_[taskIdL0B]);
	174	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0BEventList_[taskIdL0B]);
	175	+
	176	+ auto tensorTileInL0C = GetTile(tensorInL0C, tla::MakeCoord(0, nOffset), tla::MakeShape(qActual, tileN));
	177	+ tileMmad_(tensorTileInL0C, tensorAInL0, tensorBInL0, qActual, tileN, dimActual, true);
	178	+
	179	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventList_[taskIdL0B]);
	180	+ taskIdL0B = 1 - taskIdL0B;
	181	+ }
	182	+
	183	+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(l0CEventList_[taskIdL0C]);
	184	+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(l0CEventList_[taskIdL0C]);
	185	+ AscendC::CrossCoreWaitFlag<SYNC_MODE, PIPE_FIX>(QK_UB_RELEASE_FLAG);
	186	+ AscendC::CrossCoreWaitFlag<SYNC_MODE, PIPE_FIX>(QK_UB_RELEASE_FLAG + 16);
	187	+
	188	+ using CopyL0CToDst = typename TileCopy::template CopyL0CToDst<TensorC>;
	189	+ CopyL0CToDst copyL0CToDst;
	190	+ copyL0CToDst(tensorC, tensorInL0C);
	191	+
	192	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[taskIdL0C]);
	193	+ taskIdL0C = 1 - taskIdL0C;
	194	+
	195	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList_[taskIdL0A]);
	196	+ taskIdL0A = 1 - taskIdL0A;
	197	+
	198	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[l1BListId_]);
	199	+ l1BListId_ = 1 - l1BListId_;
	200	+ }
	201	+
	202	+private:
	203	+ static constexpr uint32_t SYNC_MODE = 4;
	204	+ static constexpr uint32_t AIV1_EVENT_OFFSET = 16;
	205	+ static constexpr uint32_t BLOCK_EVENT_ID = 0;
	206	+
	207	+ AscendC::LocalTensor<ElementA> l1ATensorList_[STAGES];
	208	+ AscendC::LocalTensor<ElementB> l1BTensorList_[STAGES];
	209	+ AscendC::LocalTensor<ElementA> l0ATensorList_[STAGES];
	210	+ AscendC::LocalTensor<ElementB> l0BTensorList_[STAGES];
	211	+ AscendC::LocalTensor<ElementAccumulator> l0CTensorList_[STAGES];
	212	+
	213	+ int32_t l1AEventList_[STAGES];
	214	+ int32_t l1BEventList_[STAGES];
	215	+ int32_t l0AEventList_[STAGES];
	216	+ int32_t l0BEventList_[STAGES];
	217	+ int32_t l0CEventList_[STAGES];
	218	+
	219	+ uint32_t l1AListId_{0};
	220	+ uint32_t l1BListId_{0};
	221	+
	222	+ CopyL1ToL0A copyL1ToL0A_;
	223	+ CopyL1ToL0B copyL1ToL0B_;
	224	+ TileMmad tileMmad_;
	225	+};
	226	+
	227	+} // namespace Catlass::Gemm::Block
	228	+
	229	+#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_XA_UNSHARED_QK_TLA_HPP

Ainclude/catlass/gemm/block/block_mmad_xfai_pv.hpp+254-0

@@ -0,0 +1,254 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef GEMM_BLOCK_MMAD_XFAI_PV_HPP
	13	+#define GEMM_BLOCK_MMAD_XFAI_PV_HPP
	14	+
	15	+#include "catlass/arch/resource.hpp"
	16	+#include "catlass/coord.hpp"
	17	+#include "catlass/arch/cross_core_sync.hpp"
	18	+#include "catlass/gemm/dispatch_policy.hpp"
	19	+#include "catlass/gemm/helper.hpp"
	20	+#include "catlass/gemm_coord.hpp"
	21	+#include "catlass/gemm/tile/tile_copy.hpp"
	22	+#include "catlass/gemm/tile/tile_mmad.hpp"
	23	+
	24	+////////////////////////////////////////////////////////////////////
	25	+
	26	+namespace Catlass::Gemm::Block {
	27	+////////////////////////////////////////////////////////////////////
	28	+
	29	+template <
	30	+ bool PAGED_CACHE_FLAG_, bool ENABLE_UNIT_FLAG_, class L1TileShape_, class L0TileShape_, class AType_, class BType_,
	31	+ class CType_, class BiasType_, class TileCopy_, class TileMmad_>
	32	+struct BlockMmad<
	33	+ MmadAtlasA2XFAIPV<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>, L1TileShape_, L0TileShape_, AType_, BType_, CType_,
	34	+ BiasType_, TileCopy_, TileMmad_> {
	35	+public:
	36	+ // Type Aliases
	37	+ using DispatchPolicy = MmadAtlasA2XFAIPV<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>;
	38	+ using ArchTag = typename DispatchPolicy::ArchTag;
	39	+ using L1TileShape = L1TileShape_;
	40	+ using L0TileShape = L0TileShape_;
	41	+ using ElementA = typename AType_::Element;
	42	+ using LayoutA = typename AType_::Layout;
	43	+ using ElementB = typename BType_::Element;
	44	+ using LayoutB = typename BType_::Layout;
	45	+ using ElementC = typename CType_::Element;
	46	+ using LayoutC = typename CType_::Layout;
	47	+ using TileMmad = TileMmad_;
	48	+ using CopyGmToL1A = typename TileCopy_::CopyGmToL1A;
	49	+ using CopyGmToL1B = typename TileCopy_::CopyGmToL1B;
	50	+ using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A;
	51	+ using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B;
	52	+ using CopyL0CToGm = typename TileCopy_::CopyL0CToGm;
	53	+ using ElementAccumulator =
	54	+ typename Gemm::helper::ElementAccumulatorSelector<ElementA, ElementB>::ElementAccumulator;
	55	+ using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc;
	56	+ using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc;
	57	+ using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst;
	58	+ using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst;
	59	+ using LayoutCInL0 = layout::zN;
	60	+
	61	+ using L1AAlignHelper = Gemm::helper::L1AlignHelper<ElementA, LayoutA>;
	62	+ using L1BAlignHelper = Gemm::helper::L1AlignHelper<ElementB, LayoutB>;
	63	+
	64	+ static constexpr uint32_t STAGES = DispatchPolicy::STAGES;
	65	+ static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA);
	66	+ static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB);
	67	+ static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE;
	68	+ static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE;
	69	+ static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE;
	70	+ static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES;
	71	+ static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES;
	72	+ static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES;
	73	+ static constexpr uint32_t BLOCK_SIZE = 16;
	74	+ static constexpr uint32_t EMBED_SPLIT_SIZE = 128;
	75	+ static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4;
	76	+ static constexpr uint32_t KV_BASE_BLOCK = 512;
	77	+ static constexpr uint32_t KV_SPLIT_SIZE = 128;
	78	+ static constexpr uint32_t LOAB_BLOCK = 1;
	79	+
	80	+ // Check LayoutC
	81	+ static_assert(std::is_same_v<LayoutC, layout::RowMajor>, "LayoutC only support RowMajor yet!");
	82	+
	83	+ /// Construct
	84	+ __aicore__ inline BlockMmad(
	85	+ Arch::Resource<ArchTag>& resource, uint32_t nDyn, uint32_t kDyn, uint32_t l1BufAddrStart = 0)
	86	+ {
	87	+ // Allocate L1 memory space
	88	+ l1BTensor = resource.l1Buf.template GetBufferByByte<ElementB>(
	89	+ l1BufAddrStart + L1TileShape::M * kDyn * sizeof(ElementA) * 2);
	90	+ for (uint32_t i = 0; i < STAGES; i++) {
	91	+ l1ATensor[i] = resource.l1Buf.template GetBufferByByte<ElementA>(
	92	+ l1BufAddrStart + L1TileShape::M * kDyn * sizeof(ElementA) * i);
	93	+ l0ATensor[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(L0A_PINGPONG_BUF_SIZE * i);
	94	+ l0BTensor[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(L0B_PINGPONG_BUF_SIZE * i);
	95	+ l0CTensor[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(L0C_PINGPONG_BUF_SIZE * i);
	96	+ }
	97	+ l1NDynamic = nDyn;
	98	+ l1KDynamic = kDyn;
	99	+ }
	100	+
	101	+ /// Destructor
	102	+ __aicore__ inline ~BlockMmad()
	103	+ {}
	104	+
	105	+ __aicore__ inline void getBlockShape(
	106	+ GemmCoord& actualShape, uint32_t& nowNIdx, uint32_t& nLoop, uint32_t& kvSeqlen, uint32_t& blockSize)
	107	+ {
	108	+ uint32_t nSplitSize = blockSize;
	109	+ if (nowNIdx == nLoop - 1) {
	110	+ nSplitSize = kvSeqlen - nowNIdx * blockSize;
	111	+ }
	112	+ actualShape[2] = nSplitSize;
	113	+ }
	114	+
	115	+ __aicore__ inline void getKVOffset(
	116	+ AscendC::GlobalTensor<int32_t>& gBlockTable, uint32_t& kOffset, uint32_t& nowNIdx, uint32_t& strideKV,
	117	+ uint32_t& blockSize)
	118	+ {
	119	+ if constexpr (PAGED_CACHE_FLAG_) {
	120	+ uint32_t blockTableId = gBlockTable.GetValue(nowNIdx);
	121	+ kOffset = blockTableId * blockSize * strideKV;
	122	+ } else {
	123	+ kOffset = nowNIdx * blockSize * strideKV;
	124	+ }
	125	+ }
	126	+
	127	+ __aicore__ inline void operator()(
	128	+ AscendC::GlobalTensor<ElementA> gA, AscendC::GlobalTensor<ElementB> gB, AscendC::GlobalTensor<ElementC> gC,
	129	+ AscendC::GlobalTensor<int32_t> gBlockTable, LayoutA layoutA, LayoutB layoutB, LayoutC layoutC,
	130	+ GemmCoord actualOriShape, uint32_t& nIdx, uint32_t& nLoop, uint32_t& blockSize, uint32_t kvSeqlen,
	131	+ uint32_t strideKV, uint32_t blockStackNum, Arch::CrossCoreFlag softmaxFlag)
	132	+ {
	133	+ uint32_t rowNum = actualOriShape[0];
	134	+ uint32_t embed = actualOriShape[1];
	135	+ uint32_t stackSeqTile = actualOriShape[2];
	136	+ GemmCoord actualShape{rowNum, embed, 0};
	137	+ uint32_t gBOffset = 0;
	138	+
	139	+ LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout<ElementB>(stackSeqTile, embed);
	140	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID4);
	141	+ for (uint32_t blockStackIdx = 0; (blockStackIdx < blockStackNum) && ((nIdx + blockStackIdx) < nLoop);
	142	+ blockStackIdx++) {
	143	+ uint32_t nowNIdx = nIdx + blockStackIdx;
	144	+ getBlockShape(actualShape, nowNIdx, nLoop, kvSeqlen, blockSize);
	145	+ getKVOffset(gBlockTable, gBOffset, nowNIdx, strideKV, blockSize);
	146	+ auto layoutBTile = layoutB.GetTileLayout(MakeCoord(actualShape.k(), actualShape.n()));
	147	+ MatrixCoord l1BTileCoord{blockStackIdx * blockSize, 0};
	148	+ auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BTileCoord)];
	149	+ copyGmToL1B(l1BTile, gB[gBOffset], layoutBInL1, layoutBTile);
	150	+ }
	151	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
	152	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
	153	+ Arch::CrossCoreWaitFlag(softmaxFlag);
	154	+
	155	+ uint32_t mL1Loop = CeilDiv(rowNum, L1TileShape::M);
	156	+ uint32_t kL1Loop = CeilDiv(stackSeqTile, l1KDynamic);
	157	+ uint32_t nL1Loop = CeilDiv(embed, L0TileShape::N);
	158	+
	159	+ for (uint32_t nL1Idx = 0; nL1Idx < nL1Loop; nL1Idx++) {
	160	+ uint32_t nL1Actual = (nL1Idx < nL1Loop - 1) ? L0TileShape::N : (embed - nL1Idx * L0TileShape::N);
	161	+ for (uint32_t mL1Idx = 0; mL1Idx < mL1Loop; mL1Idx++) {
	162	+ uint32_t mL1Actual = (mL1Idx < mL1Loop - 1) ? L1TileShape::M : (rowNum - mL1Idx * L1TileShape::M);
	163	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CPingPongFlag);
	164	+ for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) {
	165	+ uint32_t kL1Actual = (kL1Idx < kL1Loop - 1) ? l1KDynamic : (stackSeqTile - kL1Idx * l1KDynamic);
	166	+ // load P
	167	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1PPingPongFlag);
	168	+ MatrixCoord gmATileCoord{mL1Idx * L1TileShape::M, kL1Idx * l1KDynamic};
	169	+ auto gmTileA = gA[layoutA.GetOffset(gmATileCoord)];
	170	+ auto layoutTileA = layoutA.GetTileLayout(MakeCoord(mL1Actual, kL1Actual));
	171	+ LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout<ElementA>(mL1Actual, kL1Actual);
	172	+ copyGmToL1A(l1ATensor[l1PPingPongFlag], gmTileA, layoutAInL1, layoutTileA);
	173	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1PPingPongFlag);
	174	+
	175	+ uint32_t kL0Loop = CeilDiv(kL1Actual, L0TileShape::K);
	176	+ for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) {
	177	+ uint32_t kL0Actual =
	178	+ (kL0Idx < kL0Loop - 1) ? L0TileShape::K : (kL1Actual - kL0Idx * L0TileShape::K);
	179	+ LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout<ElementA>(mL1Actual, kL0Actual);
	180	+ MatrixCoord l1ATileCoord{0, kL0Idx * L0TileShape::K};
	181	+ auto l1ATile = l1ATensor[l1PPingPongFlag][layoutAInL1.GetOffset(l1ATileCoord)];
	182	+
	183	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag);
	184	+ if (kL0Idx == 0) {
	185	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1PPingPongFlag);
	186	+ }
	187	+ copyL1ToL0A(l0ATensor[l0ABPingPongFlag], l1ATile, layoutAInL0, layoutAInL1);
	188	+ if (kL0Idx == kL0Loop - 1) {
	189	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1PPingPongFlag);
	190	+ }
	191	+
	192	+ LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout<ElementB>(kL0Actual, nL1Actual);
	193	+ MatrixCoord l1BTileCoord{
	194	+ kL1Idx * l1KDynamic + kL0Idx * L0TileShape::K, L0TileShape::N * nL1Idx};
	195	+ auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BTileCoord)];
	196	+
	197	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag + 2);
	198	+ copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTile, layoutBInL0, layoutBInL1);
	199	+
	200	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
	201	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
	202	+ bool initMmad = kL1Idx == 0 && kL0Idx == 0;
	203	+ uint32_t mL0Align = (mL1Actual + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE;
	204	+ tileMmad(
	205	+ l0CTensor[l0CPingPongFlag], l0ATensor[l0ABPingPongFlag], l0BTensor[l0ABPingPongFlag],
	206	+ mL0Align, nL1Actual, kL0Actual, initMmad);
	207	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag);
	208	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag + 2);
	209	+ l0ABPingPongFlag = 1 - l0ABPingPongFlag;
	210	+ }
	211	+ l1PPingPongFlag = 1 - l1PPingPongFlag;
	212	+ }
	213	+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
	214	+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
	215	+ MatrixCoord gmCTileCoord{mL1Idx * L0TileShape::M, L0TileShape::N * nL1Idx};
	216	+ LayoutC layoutCTile = layoutC.GetTileLayout(MakeCoord(mL1Actual, nL1Actual));
	217	+ auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mL1Actual, nL1Actual));
	218	+ copyL0CToGm(gC[layoutC.GetOffset(gmCTileCoord)], l0CTensor[l0CPingPongFlag], layoutCTile, layoutInL0C);
	219	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CPingPongFlag);
	220	+ l0CPingPongFlag = 1 - l0CPingPongFlag;
	221	+ }
	222	+ }
	223	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID4);
	224	+ }
	225	+
	226	+protected:
	227	+ /// Data members
	228	+ AscendC::LocalTensor<ElementA> l1ATensor[STAGES];
	229	+ AscendC::LocalTensor<ElementB> l1BTensor;
	230	+ AscendC::LocalTensor<ElementA> l0ATensor[STAGES];
	231	+ AscendC::LocalTensor<ElementB> l0BTensor[STAGES];
	232	+ AscendC::LocalTensor<ElementAccumulator> l0CTensor[STAGES];
	233	+
	234	+ TileMmad tileMmad;
	235	+ CopyGmToL1A copyGmToL1A;
	236	+ CopyGmToL1B copyGmToL1B;
	237	+ CopyL1ToL0A copyL1ToL0A;
	238	+ CopyL1ToL0B copyL1ToL0B;
	239	+ CopyL0CToGm copyL0CToGm;
	240	+
	241	+ uint32_t l1PPingPongFlag = 0;
	242	+ uint32_t l0CPingPongFlag = 0;
	243	+ uint32_t l0ABPingPongFlag = 0;
	244	+
	245	+ uint32_t l1MDynamic = 0;
	246	+ uint32_t l1NDynamic = 0;
	247	+ uint32_t l1KDynamic = 0;
	248	+};
	249	+
	250	+////////////////////////////////////////////////////////////////////
	251	+
	252	+} // namespace Catlass::Gemm::Block
	253	+
	254	+#endif // GEMM_BLOCK_MMAD_XFAI_PV_HPP

Ainclude/catlass/gemm/block/block_mmad_xfai_qk.hpp+246-0

@@ -0,0 +1,246 @@
	1	+/**
	2	+ * This program is free software, you can redistribute it and/or modify.
	3	+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
	4	+ * This file is a part of the CANN Open Software.
	5	+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
	6	+ * Please refer to the License for details. You may not use this file except in compliance with the License.
	7	+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING
	8	+ * BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. See LICENSE in the root of
	9	+ * the software repository for the full text of the License.
	10	+ */
	11	+
	12	+#ifndef GEMM_BLOCK_MMAD_XFAI_QK_HPP
	13	+#define GEMM_BLOCK_MMAD_XFAI_QK_HPP
	14	+
	15	+#include "catlass/arch/resource.hpp"
	16	+#include "catlass/coord.hpp"
	17	+#include "catlass/gemm/dispatch_policy.hpp"
	18	+#include "catlass/gemm/helper.hpp"
	19	+#include "catlass/gemm_coord.hpp"
	20	+#include "catlass/gemm/tile/tile_copy.hpp"
	21	+#include "catlass/gemm/tile/tile_mmad.hpp"
	22	+
	23	+////////////////////////////////////////////////////////////////////
	24	+
	25	+namespace Catlass::Gemm::Block {
	26	+////////////////////////////////////////////////////////////////////
	27	+
	28	+template <
	29	+ bool PAGED_CACHE_FLAG_, bool ENABLE_UNIT_FLAG_, class L1TileShape_, class L0TileShape_, class AType_, class BType_,
	30	+ class CType_, class BiasType_, class TileCopy_, class TileMmad_>
	31	+struct BlockMmad<
	32	+ MmadAtlasA2XFAIQK<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>, L1TileShape_, L0TileShape_, AType_, BType_, CType_,
	33	+ BiasType_, TileCopy_, TileMmad_> {
	34	+public:
	35	+ // Type Aliases
	36	+ using DispatchPolicy = MmadAtlasA2XFAIQK<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>;
	37	+ using ArchTag = typename DispatchPolicy::ArchTag;
	38	+ using L1TileShape = L1TileShape_;
	39	+ using L0TileShape = L0TileShape_;
	40	+ using ElementA = typename AType_::Element;
	41	+ using LayoutA = typename AType_::Layout;
	42	+ using ElementB = typename BType_::Element;
	43	+ using LayoutB = typename BType_::Layout;
	44	+ using ElementC = typename CType_::Element;
	45	+ using LayoutC = typename CType_::Layout;
	46	+ using TileMmad = TileMmad_;
	47	+ using CopyGmToL1A = typename TileCopy_::CopyGmToL1A;
	48	+ using CopyGmToL1B = typename TileCopy_::CopyGmToL1B;
	49	+ using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A;
	50	+ using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B;
	51	+ using CopyL0CToGm = typename TileCopy_::CopyL0CToGm;
	52	+ using ElementAccumulator =
	53	+ typename Gemm::helper::ElementAccumulatorSelector<ElementA, ElementB>::ElementAccumulator;
	54	+ using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc;
	55	+ using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc;
	56	+ using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst;
	57	+ using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst;
	58	+ using LayoutCInL0 = layout::zN;
	59	+
	60	+ using L1AAlignHelper = Gemm::helper::L1AlignHelper<ElementA, LayoutA>;
	61	+ using L1BAlignHelper = Gemm::helper::L1AlignHelper<ElementB, LayoutB>;
	62	+
	63	+ static constexpr uint32_t STAGES = DispatchPolicy::STAGES;
	64	+ static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA);
	65	+ static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB);
	66	+ static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE;
	67	+ static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE;
	68	+ static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE;
	69	+ static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES;
	70	+ static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES;
	71	+ static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES;
	72	+ static constexpr uint32_t BLOCK_SIZE = 16;
	73	+ static constexpr uint32_t EMBED_SPLIT_SIZE = 128;
	74	+ static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4;
	75	+ static constexpr uint32_t KV_BASE_BLOCK = 512;
	76	+ static constexpr uint32_t KV_SPLIT_SIZE = 128;
	77	+
	78	+ static_assert(std::is_same_v<LayoutC, layout::RowMajor>, "LayoutC only support RowMajor yet!");
	79	+
	80	+ __aicore__ inline BlockMmad(
	81	+ Arch::Resource<ArchTag>& resource, uint32_t nDyn, uint32_t kDyn, uint32_t l1BufAddrStart = 0)
	82	+ {
	83	+ // Allocate L1 memory space
	84	+ l1ATensor = resource.l1Buf.template GetBufferByByte<ElementA>(l1BufAddrStart);
	85	+ for (uint32_t i = 0; i < STAGES; i++) {
	86	+ l1BTensor[i] = resource.l1Buf.template GetBufferByByte<ElementB>(
	87	+ l1BufAddrStart + L1TileShape::M * kDyn * sizeof(ElementA) + nDyn * kDyn * sizeof(ElementB) * i);
	88	+ l0ATensor[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(L0A_PINGPONG_BUF_SIZE * i);
	89	+ l0BTensor[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(L0B_PINGPONG_BUF_SIZE * i);
	90	+ l0CTensor[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(L0C_PINGPONG_BUF_SIZE * i);
	91	+ }
	92	+ l1NDynamic = nDyn;
	93	+ l1KDynamic = kDyn;
	94	+ }
	95	+
	96	+ __aicore__ inline ~BlockMmad()
	97	+ {}
	98	+
	99	+ __aicore__ inline void loadQGM(
	100	+ AscendC::GlobalTensor<ElementA> gA, LayoutA layoutA, uint32_t rowNum, uint32_t& singleGroupHeads,
	101	+ uint32_t& qHeads)
	102	+ {
	103	+ uint32_t embed = layoutA.shape(1);
	104	+ uint32_t rowNumRound = RoundUp(rowNum, L1AAlignHelper::M_ALIGNED);
	105	+ uint32_t tokenNumPerGroup = rowNum / singleGroupHeads;
	106	+ auto layoutSingleANd = layoutA.GetTileLayout(MakeCoord(singleGroupHeads, embed));
	107	+ LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout<ElementA>(rowNum, embed);
	108	+ copyGmToL1A(
	109	+ l1ATensor, gA, layoutAInL1, layoutSingleANd, tokenNumPerGroup, qHeads * embed, tokenNumPerGroup, BLOCK_SIZE,
	110	+ rowNumRound);
	111	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID3);
	112	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID3);
	113	+ }
	114	+
	115	+ __aicore__ inline void getBlockShape(
	116	+ GemmCoord& actualShape, uint32_t nL1Idx, uint32_t nL1Loop, uint32_t stackSeqTile)
	117	+ {
	118	+ uint32_t nSplitSize = l1NDynamic;
	119	+ if (nL1Idx == nL1Loop - 1) {
	120	+ nSplitSize = stackSeqTile - nL1Idx * l1NDynamic;
	121	+ }
	122	+ actualShape[1] = nSplitSize;
	123	+ }
	124	+
	125	+ __aicore__ inline void getKVOffset(
	126	+ AscendC::GlobalTensor<int32_t>& gBlockTable, uint32_t& kOffset, uint32_t nowNIdx, uint32_t nL1Idx,
	127	+ uint32_t strideKV, uint32_t blockSize)
	128	+ {
	129	+ if constexpr (PAGED_CACHE_FLAG_) {
	130	+ uint32_t blockTableId = gBlockTable.GetValue(nowNIdx);
	131	+ kOffset = blockTableId * blockSize * strideKV + nL1Idx * l1NDynamic * strideKV;
	132	+ } else {
	133	+ kOffset = nowNIdx * blockSize * strideKV + nL1Idx * l1NDynamic * strideKV;
	134	+ }
	135	+ }
	136	+
	137	+ __aicore__ inline void operator()(
	138	+ AscendC::GlobalTensor<ElementA> gA, AscendC::GlobalTensor<ElementB> gB, AscendC::GlobalTensor<ElementC> gC,
	139	+ AscendC::GlobalTensor<int32_t> gBlockTable, LayoutA layoutA, LayoutB layoutB, LayoutC layoutC,
	140	+ GemmCoord actualOriShape, uint32_t nIdx, uint32_t nLoop, uint32_t blockSize, uint32_t strideKV)
	141	+ {
	142	+ uint32_t rowNum = actualOriShape[0];
	143	+ uint32_t stackSeqTile = actualOriShape[1];
	144	+ uint32_t embed = actualOriShape[2];
	145	+
	146	+ GemmCoord actualShape{rowNum, 0, embed};
	147	+ uint32_t gBOffset = 0;
	148	+
	149	+ LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout<ElementA>(rowNum, embed);
	150	+
	151	+ uint32_t tileNNumPerBaseBlock = blockSize / l1NDynamic;
	152	+ uint32_t nL1Loop = CeilDiv(stackSeqTile, l1NDynamic);
	153	+ for (uint32_t nL1Idx = 0; nL1Idx < nL1Loop; ++nL1Idx) {
	154	+ uint32_t nowNIdx = nIdx + nL1Idx / tileNNumPerBaseBlock;
	155	+ getBlockShape(actualShape, nL1Idx, nL1Loop, stackSeqTile);
	156	+ getKVOffset(gBlockTable, gBOffset, nowNIdx, nL1Idx % tileNNumPerBaseBlock, strideKV, blockSize);
	157	+ uint32_t mActual = actualShape.m();
	158	+ uint32_t kActual = actualShape.k();
	159	+ uint32_t nActual = actualShape.n();
	160	+ LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout<ElementB>(kActual, nActual);
	161	+
	162	+ auto layoutBTile = layoutB.GetTileLayout(MakeCoord(kActual, nActual));
	163	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1KvPingPongFlag);
	164	+ copyGmToL1B(l1BTensor[l1KvPingPongFlag], gB[gBOffset], layoutBInL1, layoutBTile);
	165	+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1KvPingPongFlag);
	166	+
	167	+ uint32_t mL0Loop = CeilDiv(mActual, L0TileShape::M);
	168	+ uint32_t kL0Loop = CeilDiv(kActual, L0TileShape::K);
	169	+ for (uint32_t mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) {
	170	+ uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0TileShape::M : (mActual - mL0Idx * L0TileShape::M);
	171	+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CPingPongFlag);
	172	+ for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) {
	173	+ uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0TileShape::K : (kActual - kL0Idx * L0TileShape::K);
	174	+
	175	+ LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout<ElementA>(mL0Actual, kL0Actual);
	176	+ MatrixCoord l1ATileCoord{mL0Idx * L0TileShape::M, kL0Idx * L0TileShape::K};
	177	+ auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ATileCoord)];
	178	+
	179	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag);
	180	+ copyL1ToL0A(l0ATensor[l0ABPingPongFlag], l1ATile, layoutAInL0, layoutAInL1);
	181	+
	182	+ LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout<ElementB>(kL0Actual, nActual);
	183	+ MatrixCoord l1BTileCoord{kL0Idx * L0TileShape::K, 0};
	184	+ auto l1BTile = l1BTensor[l1KvPingPongFlag][layoutBInL1.GetOffset(l1BTileCoord)];
	185	+ if ((mL0Idx == 0) && (kL0Idx == 0)) {
	186	+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1KvPingPongFlag);
	187	+ }
	188	+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag + 2);
	189	+ copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTile, layoutBInL0, layoutBInL1);
	190	+ if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) {
	191	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1KvPingPongFlag);
	192	+ }
	193	+
	194	+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
	195	+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
	196	+ bool initMmad = kL0Idx == 0;
	197	+ uint32_t mL0Align = (mL0Actual + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE;
	198	+ tileMmad(
	199	+ l0CTensor[l0CPingPongFlag], l0ATensor[l0ABPingPongFlag], l0BTensor[l0ABPingPongFlag], mL0Align,
	200	+ nActual, kL0Actual, initMmad);
	201	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag);
	202	+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag + 2);
	203	+ l0ABPingPongFlag = 1 - l0ABPingPongFlag;
	204	+ }
	205	+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
	206	+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
	207	+ MatrixCoord gmCTileCoord{mL0Idx * L0TileShape::M, nL1Idx * l1NDynamic};
	208	+ LayoutC layoutCTile = layoutC.GetTileLayout(MakeCoord(mL0Actual, nActual));
	209	+ auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mL0Actual, nActual));
	210	+ copyL0CToGm(gC[layoutC.GetOffset(gmCTileCoord)], l0CTensor[l0CPingPongFlag], layoutCTile, layoutInL0C);
	211	+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CPingPongFlag);
	212	+ l0CPingPongFlag = 1 - l0CPingPongFlag;
	213	+ }
	214	+ l1KvPingPongFlag = 1 - l1KvPingPongFlag;
	215	+ }
	216	+ }
	217	+
	218	+protected:
	219	+ /// Data members
	220	+ AscendC::LocalTensor<ElementA> l1ATensor;
	221	+ AscendC::LocalTensor<ElementB> l1BTensor[STAGES];
	222	+ AscendC::LocalTensor<ElementA> l0ATensor[STAGES];
	223	+ AscendC::LocalTensor<ElementB> l0BTensor[STAGES];
	224	+ AscendC::LocalTensor<ElementAccumulator> l0CTensor[STAGES];
	225	+
	226	+ TileMmad tileMmad;
	227	+ CopyGmToL1A copyGmToL1A;
	228	+ CopyGmToL1B copyGmToL1B;
	229	+ CopyL1ToL0A copyL1ToL0A;
	230	+ CopyL1ToL0B copyL1ToL0B;
	231	+ CopyL0CToGm copyL0CToGm;
	232	+
	233	+ uint32_t l1KvPingPongFlag = 0;
	234	+ uint32_t l0CPingPongFlag = 0;
	235	+ uint32_t l0ABPingPongFlag = 0;
	236	+
	237	+ uint32_t l1MDynamic = 0;
	238	+ uint32_t l1NDynamic = 0;
	239	+ uint32_t l1KDynamic = 0;
	240	+};
	241	+
	242	+////////////////////////////////////////////////////////////////////
	243	+
	244	+} // namespace Catlass::Gemm::Block
	245	+
	246	+#endif // GEMM_BLOCK_MMAD_XFAI_QK_HPP

Minclude/catlass/gemm/dispatch_policy.hpp+56-0

	@@ -63,6 +63,14 @@ struct MmadAtlasA2Preload : public MmadAtlasA2 {
63	static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_;	63	static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_;
64	};	64	};
65		65
		66	+struct MmadAtlasA2UnsharedFAQK : public MmadAtlasA2 {
		67	+ static constexpr uint32_t STAGES = 2;
		68	+};
		69	+
		70	+struct MmadAtlasA2UnsharedFAPV : public MmadAtlasA2 {
		71	+ static constexpr uint32_t STAGES = 2;
		72	+};
		73	+
66	struct MmadAtlasA2FAQK : public MmadAtlasA2 {	74	struct MmadAtlasA2FAQK : public MmadAtlasA2 {
67	static constexpr uint32_t STAGES = 2;	75	static constexpr uint32_t STAGES = 2;
68	};	76	};
	@@ -164,6 +172,20 @@ struct MmadAtlasA2FAIQK : public MmadAtlasA2 {
164	static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;	172	static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
165	};	173	};
166		174
		175	+template <bool PAGED_CACHE_FLAG_ = false, bool ENABLE_UNIT_FLAG_ = false>
		176	+struct MmadAtlasA2XFAIQK : public MmadAtlasA2 {
		177	+ static constexpr uint32_t STAGES = 2;
		178	+ static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_;
		179	+ static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
		180	+};
		181	+
		182	+template <bool PAGED_CACHE_FLAG_ = false, bool ENABLE_UNIT_FLAG_ = false>
		183	+struct MmadAtlasA2XFAIPV : public MmadAtlasA2 {
		184	+ static constexpr uint32_t STAGES = 2;
		185	+ static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_;
		186	+ static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
		187	+};
		188	+
167	template <bool PAGED_CACHE_FLAG_ = false, bool ENABLE_UNIT_FLAG_ = false>	189	template <bool PAGED_CACHE_FLAG_ = false, bool ENABLE_UNIT_FLAG_ = false>
168	struct MmadAtlasA2FAIPV : public MmadAtlasA2 {	190	struct MmadAtlasA2FAIPV : public MmadAtlasA2 {
169	static constexpr uint32_t STAGES = 2;	191	static constexpr uint32_t STAGES = 2;
	@@ -185,6 +207,20 @@ struct MmadAtlasA2FAITailPV : public MmadAtlasA2 {
185	static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;	207	static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
186	};	208	};
187		209
		210	+template <bool PAGED_CACHE_FLAG_ = false, bool ENABLE_UNIT_FLAG_ = false>
		211	+struct MmadAtlasA2FAIQKSplitRow : public MmadAtlasA2 {
		212	+ static constexpr uint32_t STAGES = 2;
		213	+ static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_;
		214	+ static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
		215	+};
		216	+
		217	+template <bool PAGED_CACHE_FLAG_ = false, bool ENABLE_UNIT_FLAG_ = false>
		218	+struct MmadAtlasA2FAIPVSplitRow : public MmadAtlasA2 {
		219	+ static constexpr uint32_t STAGES = 2;
		220	+ static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_;
		221	+ static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
		222	+};
		223	+
188	template <bool ENABLE_UNIT_FLAG_ = false>	224	template <bool ENABLE_UNIT_FLAG_ = false>
189	struct MmadAtlasA2FullLoadA : public MmadAtlasA2 {	225	struct MmadAtlasA2FullLoadA : public MmadAtlasA2 {
190	static constexpr uint32_t STAGES = 2;	226	static constexpr uint32_t STAGES = 2;
	@@ -567,6 +603,26 @@ struct MmadPlanarComplexFused : public MmadBase<ArchTag_, false> {
567	static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_;	603	static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_;
568	};	604	};
569		605
		606	+template <class ArchTag_>
		607	+struct MmadXASharedQK : public MmadBase<ArchTag_, false> {
		608	+ static constexpr uint32_t STAGES = 2;
		609	+};
		610	+
		611	+template <class ArchTag_>
		612	+struct MmadXAUnsharedQK : public MmadBase<ArchTag_, false> {
		613	+ static constexpr uint32_t STAGES = 2;
		614	+};
		615	+
		616	+template <class ArchTag_>
		617	+struct MmadXAUnsharedPV : public MmadBase<ArchTag_, false> {
		618	+ static constexpr uint32_t STAGES = 2;
		619	+};
		620	+
		621	+template <class ArchTag_>
		622	+struct MmadXASharedPV : public MmadBase<ArchTag_, false> {
		623	+ static constexpr uint32_t STAGES = 2;
		624	+};
		625	+
570	} // namespace Catlass::Gemm	626	} // namespace Catlass::Gemm
571		627
572	#endif // CATLASS_GEMM_DISPATCH_POLICY_HPP	628	#endif // CATLASS_GEMM_DISPATCH_POLICY_HPP