# SmallQ FlashAttention (Vector)

面向推测解码（Speculative Decoding）/ MTP 场景的注意力算子 AscendC Vector 实现，针对 qLen 很小（典型 1–8 token）、kvLen 较长的工作负载优化。运行在 Ascend 910B3 / Atlas 800I A2 (ascend910b) 上。

## 接口

| 张量 | Shape | Dtype | 说明 |
|---|---|---|---|
| Q (input) | `[numHeads, qLen, headDim]` | fp16 | Query |
| K (input) | `[numHeads, kvLen, headDim]` | fp16 | Key |
| V (input) | `[numHeads, kvLen, headDim]` | fp16 | Value |
| O (output) | `[numHeads, qLen, headDim]` | fp16 | 输出 |

```
O = softmax(Q · Kᵀ / √headDim) · V
```

## 支持范围

| 参数 | 范围 | 备注 |
|---|---|---|
| numHeads | 任意 ≥ 1 | 每 head 一个 AI Core，blockDim = numHeads |
| qLen | 任意 ≥ 1（典型 1–8）| Q 常驻 UB，UB 占用与 qLen 成正比 |
| headDim | **任意** ≥ 1 | 内部对齐到 16；非对齐尾部用 DataCopyPad 写出 |
| kvLen | **任意** ≥ 1 | 内部按 blockK 分块，最后一块自动处理尾部 |
| dtype | fp16 | 内部累加用 fp32 |
| 硬件 | ascend910b | CANN 8.5.0+ |

实测稳定通过的真实模型 shape（在 7 张 910B3 上并行 512 用例 × 0 失败）：

- Llama 3.1 8B / Qwen3-8B：nh=32, hd=128
- Llama 3 70B / Qwen3-32B：nh=64, hd=128
- Llama 3.1 405B：nh=128, hd=128
- Qwen3-235B-A22B：nh=64, hd=128
- DeepSeek V3 (MLA K-side hd=192, V-side hd=128)
- 各种 qLen ∈ {1, 2, 4, 8}，kvLen 从 1 到 16384（含 1023/1025/1500/3000 非 2ⁿ 长度）

## 算法

**K-outer + Q 驻留 UB**：blockK 默认 64，UB 装不下时按 ÷2 自动缩减到下限 4。Q 整体一次性搬入并驻留，K/V 按 blockK 分块流式读入，双缓冲（BUFFER_NUM=2）预取下一块 K。

**Online Softmax**：每处理一个 K 分块，对每行维护 running max 和 running sum。新分块到来时计算 α = exp(m_old − m_new) 修正历史累加器 oAcc。

**精度策略**：
- QK 点积：标量 fp32 累加（`for d in [0, headDim): dot += float(q[d]) * float(k[d])`）。对任意 headDim 都正确，且彻底避开 fp16 累加误差。
- ReduceMax/ReduceSum：fp32 + 掩码（mask 长度等于 blockRows，处理 blockK 不整除 kvLen 的尾块）。
- PV 累加：标量 fp32（`oAcc[d] += weight * float(v[d])`）。
- 输出：fp32 oAcc × (1/rowSum) 再 cast 回 fp16。

**任意 headDim 支持**：
- UB 内部以 `hdPad = ceil(headDim/16)*16` 对齐，DataCopy GM↔UB 全部走 16 对齐快速路径。
- 输出回 GM 时压紧到 `headDim` stride。压紧后的 `qLen × headDim` 元素如果不是 16 的整数倍，**最后一段用 `DataCopyPad` UB→GM 写出**——底层硬件指令 `copy_ubuf_to_gm_align_b16` 使用字节级写使能，多核同时写同一 cache line 的不同字节是安全的（这是修复多 head 共享 cache line 竞争的关键）。

**动态 blockK**：Tiling 阶段根据 UB 实际空闲容量从 64 开始按 ÷2 缩小，下限 4。`(headDim=2048, qLen=7)` 这样的极端 UB 紧张 case 也能跑。

## 文件结构

```
smallq_flash_attention/
├── op_kernel/
│   ├── smallq_flash_attention.cpp         # Kernel 入口
│   └── smallq_flash_attention_impl.h      # 计算实现 + DataCopyPad tail
├── op_host/
│   ├── smallq_flash_attention_def.cpp     # 算子定义
│   ├── smallq_flash_attention_proto.cpp   # InferShape
│   ├── smallq_flash_attention_tiling.h    # TilingData 结构
│   └── smallq_flash_attention_tiling.cpp  # 动态 blockK + UB 预算
├── CMakeLists.txt
└── README.md
```

## 编译部署

```bash
# 编译
cd cann-recipes-infer/ops/ascendc
bash build.sh -n "smallq_flash_attention" -c "ascend910b"

# 部署（一定要装到 opp 路径下，opp 优先级高于 vendors）
bash output/CANN-custom_ops-none-linux.aarch64.run \
     --quiet --install-path=$ASCEND_HOME_PATH/opp/vendors/customize

# 安装脚本会在路径下嵌套一层 vendors/customize/，把内层文件拷到外层即可
yes | cp -rf $ASCEND_HOME_PATH/opp/vendors/customize/vendors/customize/. \
             $ASCEND_HOME_PATH/opp/vendors/customize/
```

## aclnn API 调用

```cpp
#include "aclnn_smallq_flash_attention.h"

// 1) 准备 [numHeads, qLen, headDim] / [numHeads, kvLen, headDim] 的 fp16 aclTensor
aclTensor *qTensor, *kTensor, *vTensor, *oTensor;
// ... aclCreateTensor 略

// 2) 获取 workspace 大小 + 创建执行器
uint64_t workspaceSize = 0;
aclOpExecutor* executor = nullptr;
aclnnSmallqFlashAttentionGetWorkspaceSize(
    qTensor, kTensor, vTensor, oTensor, &workspaceSize, &executor);

// 3) 分配 workspace 并执行
void* workspace = nullptr;
if (workspaceSize > 0) aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
aclnnSmallqFlashAttention(workspace, workspaceSize, executor, stream);
aclrtSynchronizeStream(stream);
```

## 测试

`test_aclnn/` 目录下有 `test_aclnn.cpp` 测试程序和 `run_model_tests.py` 多卡并行测试驱动。后者基于 Llama 3 / Qwen 3 / DeepSeek V3 / Mistral 真实 config 生成测试 shape，560 用例 × 7 卡并行，约 9 分钟跑完，全 PASS。