support flash_attention_encoder custom op · Pull Request !37 · MindSpore/ms_custom_ops

# FlashAttention Encoder（Self-Attention）算子说明

## 背景简介
FlashAttention 是一种高性能的自注意力实现，主要通过分块/重计算、mask压缩与更优的访存/算子融合来降低显存占用和提升吞吐。本项目中的 `flash_attention_encoder` 为自定义算子接口，统一对接 Ascend 910B（ND 内核）与 310P（NZ 内核）的 Unpad Flash Attention 系列内核：
- 910B: `UnpadFlashAttentionOperation`（ND 路径，支持高精度 FP32 模式）
- 310P: `UnpadFlashAttentionNzOperation`（NZ 路径，支持 SWA/Alibi 压缩等变体）

参考（ATB 文档，SelfAttentionOperation 参数与特性说明）：
- [SelfAttentionOperation 参数列表](https://www.hiascend.com/document/detail/zh/canncommercial/82RC1/API/ascendtbapi/ascendtb_01_0278.html)

## 计算公式
给定 Q/K/V 与缩放系数 `qkScale` 与可选掩码 M：

- 令 Q ∈ R^(N_q × H × D), K ∈ R^(N_k × H_kv × D), V ∈ R^(N_k × H_kv × D_v)。H 为注意力头数，D 为 head_dim。
- 计算 logits：
  ```
  S = (Q · K^T) / scale + M          # 其中 scale = 1 / qkScale
  ```
- 归一化：
  ```
  P = softmax(S)
  ```
- 输出：
  ```
  O = P · V
  ```

说明：本版本未开放 LOGN 变体与量化相关的叠加，仅对接常规 scale 和常见掩码类型；不同平台（ND/NZ）内部实现差异对用户透明。

## 接口与输入输出
### 名称
- 算子名：`flash_attention_encoder`

### 输入参数

| Name            | DType                     | Shape                     | Optional | Format | Description |
|-----------------|---------------------------|---------------------------|----------|--------|-------------|
| query           | Tensor[float16/bfloat16]  | TH/TND                    | No       | ND/NZ | 查询向量 Q |
| key             | Tensor[float16/bfloat16]  | 同 query                  | No       | ND/NZ | 键向量 K |
| value           | Tensor[float16/bfloat16]  | 同 query                  | No       | ND/NZ | 值向量 V |
| layer_id        | Tensor[int32]             | (1,) 或按实现要求         | Yes      | ND    | 层索引（保留，可选） |
| mask            | Tensor[float16/bfloat16]  | 见不同 mask 类型要求      | Yes      | ND/NZ | 注意力掩码，支持 NORM/ALIBI/SWA 等 |
| alibi_coeff     | Tensor[float32]           | 见 ALIBI 需求             | Yes      | ND    | ALIBI 相关（可选） |
| deq_scale_qk    | Tensor[float32]           | (head_num,)               | Yes      | ND    | 量化相关（保留） |
| deq_offset_qk   | Tensor[float32]           | (head_num,)               | Yes      | ND    | 量化相关（保留） |
| deq_scale_pv    | Tensor[float32]           | (head_num,)               | Yes      | ND    | 量化相关（保留） |
| deq_offset_pv   | Tensor[float32]           | (head_num,)               | Yes      | ND    | 量化相关（保留） |
| quant_p         | Tensor[float32]           | 按需                      | Yes      | ND    | 量化相关（保留） |
| logN            | Tensor[float32]           | (1,) 或按实现要求         | Yes      | ND    | 预留（本版未使用） |
| q_seq_len       | Tensor[int32]             | (batch,)                  | Cond     | ND    | TH/TND 布局必需：每 batch 的 Q 序列长度 |
| kv_seq_len      | Tensor[int32]             | (batch,)                  | Cond     | ND    | TH/TND 布局必需：每 batch 的 KV 序列长度 |
| head_num        | int                       | -                         | Yes      | -     | 注意力头数（H） |
| scale_value     | float                     | -                         | Yes      | -     | QK 缩放系数 `qkScale`（通常为 1/sqrt(head_dim)） |
| kv_head_num     | int                       | -                         | Yes      | -     | KV 头数（GQA 场景）；0 表示与 head_num 对齐 |
| mask_type       | int                       | -                         | Yes      | -     | 掩码类型：UNDEFINED/NORM/ALIBI/SWA_* |
| kernel_type     | int                       | -                         | Yes      | -     | 内核精度：0-默认半精度，1-高精度（FP32 BMM1） |
| window_size     | int                       | -                         | Yes      | -     | SWA 窗口大小（SWA 场景） |
| cache_type      | int                       | -                         | Yes      | -     | 缓存类型：0-NORM，1-SWA（SWA 优化） |
| cache_mode      | int                       | -                         | Yes      | -     | 后端选择：0-910B（ND 路径），1-310P（NZ 路径） |

注：当前版本未接线量化/online-offline QKV/clamp/ring/prefix 等高级特性，相关张量为占位，可保持为 None。

### 输出参数

| Name          | DType                    | Shape           | Description |
|---------------|--------------------------|-----------------|-------------|
| attention_out | Tensor[float16/bfloat16] | 与 query 对齐   | 注意力输出   |

### 形状与布局
- 支持布局：TH / TND（以及 BNSD 的等价场景）
  - TH：展平 token 流；`query [sum(q_seq), H*D]`，`key/value [sum(kv_seq), H_kv*D]` → 输出同 Q
  - TND：可按每 batch 的真实序列构造 token 流（建议采用 TH 形式输入以避免额外填充）
- 动态长度（TH/TND）需提供 `q_seq_len` 与 `kv_seq_len`（int32/CPU）。
- 910B（ND）与 310P（NZ）会在内部完成必要的 ND↔NZ 格式与参数设置，用户接口一致。

### Python 使用示例

```python
import numpy as np
from mindspore import context, Tensor
import ms_custom_ops

np.random.seed(0)
context.set_context(device_target="Ascend", mode=context.PYNATIVE_MODE)

head_num = 8
kv_head_num = 4
head_dim = 128
scale_value = float(1.0 / np.sqrt(head_dim))

q_seq = np.array([16, 20], dtype=np.int32)
kv_seq = np.array([16, 20], dtype=np.int32)
q_tokens = int(q_seq.sum())
kv_tokens = int(kv_seq.sum())

q = Tensor(np.random.uniform(-1, 1, size=(q_tokens, head_num, head_dim)).astype(np.float16))
k = Tensor(np.random.uniform(-1, 1, size=(kv_tokens, kv_head_num, head_dim)).astype(np.float16))
v = Tensor(np.random.uniform(-1, 1, size=(kv_tokens, kv_head_num, head_dim)).astype(np.float16))

q_seq_len = Tensor(q_seq).move_to("CPU")
kv_seq_len = Tensor(kv_seq).move_to("CPU")

out = ms_custom_ops.flash_attention_encoder(
    q, k, v,
    layer_id=None, mask=None, alibi_coeff=None,
    deq_scale_qk=None, deq_offset_qk=None, deq_scale_pv=None, deq_offset_pv=None,
    quant_p=None, logN=None,
    q_seq_len=q_seq_len, kv_seq_len=kv_seq_len,
    head_num=head_num, scale_value=scale_value, kv_head_num=kv_head_num,
    kernel_type=0, mask_type=0, window_size=0, cache_type=0
)

print(out.shape)  # 期望: (q_tokens, head_num, head_dim)
```

## 支持的参数与调用顺序
- 固定设置：`calcType = PA_ENCODER`
- 标量参数（与 ATB 对应关系）：
  - head_num → atb.headNum（>0）
  - scale_value → atb.qkScale（通常 1/sqrt(head_dim)）
  - kv_head_num → atb.kvHeadNum（0 表示与 head_num 对齐；GQA 约束：head_num % kv_head_num == 0）
  - mask_type → atb.MaskType（支持 NORM、ALIBI、SWA_*，见下）
  - kernel_type → atb.KernelType（0: 默认半精度；1: 高精度/910B FP32 BMM1）
  - window_size → atb.windowSize（SWA 窗口）
  - cache_type → atb.CacheType（0: NORM；1: SWA）
- 动态长度（仅 TH/TND 需要）：`q_seq_len`、`kv_seq_len`（int32，CPU Tensor）

参考：ATB SelfAttention 参数定义（官方文档）
链接: https://www.hiascend.com/document/detail/zh/canncommercial/82RC1/API/ascendtbapi/ascendtb_01_0278.html

## 特性支持矩阵（本版）
- 平台/内核
  - 910B（ND）：`UnpadFlashAttentionOperation`
    - 掩码：NORM、ALIBI、SWA_NORM、SWA_COMPRESS（内部会将部分压缩类型映射到基础类型）
    - 高精度：通过 `kernel_type=1` 启用（FP32 BMM1）；`kernel_type=0` 走半精度高吞吐
    - LOGN：未开放
  - 310P（NZ）：`UnpadFlashAttentionNzOperation`
    - 掩码：NORM、ALIBI、SWA_NORM、SWA_COMPRESS（含压缩/左对齐等变体的映射）
    - LOGN：未开放于本版接口
    - SWA：按 `mask_type` 选择；`window_size`/`cache_type` 推荐联动配置

- DType
  - Q/K/V：fp16、bf16
  - mask：与 Q/K/V 一致（fp16 或 bf16）
  - alibi_coeff：当前未使用（可留空）

- GQA
  - 支持 `kv_head_num > 0`，需满足 `head_num % kv_head_num == 0`

- 动态序列长度
  - TH/TND：必须提供 `q_seq_len/kv_seq_len`（int32/CPU）

- 暂未支持（保留/默认关闭）
  - 量化（在线/离线 INT8）、clamp、LOGN、prefix/decoder、ring/send-recv 等高级特性

## 性能与精度建议
- 910B：对精度敏感场景设置 `kernel_type=1`（FP32 BMM1），默认 `0` 走半精度高吞吐。
- 310P：长序列场景开启 SWA，并合理设置 `window_size` 与 `cache_type` 以控内存与吞吐。

## 版本与兼容性
- 当前文档对应的实现为 `flash_attention_encoder` 的首个版本，接口已与 ATB `SelfAttentionParam` 对齐，后续将逐步开放量化/decoder/prefix/ring 等特性，保持 YAML 与参数结构不变，向后兼容。

类型	指派人员	状态
审查	fangwenyi	已完成 (0/0人)
测试	fangwenyi	已完成 (0/0人)

MindSpore/ms_custom_ops .gitee-modal { width: 500px !important; }

搜索帮助

MindSpore/ms_custom_ops