feat: add sdpa_math operator (Ascend Triton + P800 exact-P/A1)

CHANGELOG.md

@@ -34,6 +34,19 @@
 ## [Unreleased]
 
 ### Added
+- **sdpa_math 第二平台: p800-kunlunxin**（vendor O/LSE + exact-P + A1）
+  - `_profile` / benchmark / framework 消费方接入 P800；A1 使用
+    `AutogradCUDA+CUDA` 成对注册；fast path 由 vendor efficient attention
+    产出 `O/LSE`，再物化 `P=exp(scale·QKᵀ-LSE)`，fp32/mask/dropout 回退
+    `torch_level`
+  - A1 backward P800 分支复用 vendor efficient backward，`dprobs` 用已
+    物化 P 单独求 Jacobian；fp16 GQA dO+dP 对照 native 通过
+  - kernel 52/52、黄金 175/175、op 22 项、FlagGems 应用层 9 项、
+    P800 perf gate 全绿；direct 1.46-2.21x native
+  - 新增 FlagOS mini-LLM 端到端：FlagGems GELU + A1 sdpa_math +
+    概率图可选熵正则，forward 1.07x、train-out 1.45x、
+    train+probs 1.48x、generate 1.09x；fused backward off/on 1.45-1.50x
+  - `ops.sdpa_math` 包化，支持 `from ops.sdpa_math import register_a1`
 - **新算子: [ops/embedding](ops/embedding/)**（`aten::embedding`，
   A1 路线，p800-kunlunxin）
   - forward 复用 XMLIR native `index_select` row-gather；任意 rank indices、

docs/known-issues.md

@@ -195,7 +195,7 @@ N=5120 误差 0.6。原 softmax 样例的测试形状恰好全是整倍数，漏
 平台结论与复现命令见
 [ops/sdpa/reports/mlu590.md](../ops/sdpa/reports/mlu590.md)。
 
-### 19. SDPA math 后端的六类语义/注册分化（sdpa_math 实测发现）
+### 19. SDPA math 后端的七类语义/注册/平台分化（sdpa_math 实测发现）
 
 1. **torch_npu 的 `F.sdpa` 从不调用 math 后端**: NPU 上
    `F.sdpa(sdpa_kernel(MATH))` 对本 op 命中 **0 次**（整条走
@@ -223,6 +223,10 @@ N=5120 误差 0.6。原 softmax 样例的测试形状恰好全是整倍数，漏
 6. **`causal` 与 `attn_mask` 互斥**: 同时给（bool/float 皆然）直接
    raise `Explicit attn_mask should not be set when is_causal=True`；
    全 `-inf` 行返回 P=O=0 而非 NaN。
+7. **P800 不能直接移植 Ascend exact-P Triton 双 pass**: 源码级
+   row-stat + P-write 结构会触发 XMLIR pointer-state rewrite 失败；
+   既有 no-P custom schedule 也慢于 vendor efficient attention。生产路径
+   改为 vendor `O/LSE` + `P=exp(scale·QKᵀ-LSE)`，保持双输出契约。
 
 复现: `python3 ops/sdpa_math/probes/native_semantics.py`（6 节证据表）·
 黄金互验: `ops/sdpa_math/script/gen_golden.py`。

ops/README.md

@@ -23,12 +23,12 @@ grep -rl my_op . | xargs sed -i 's/my_op/<你的算子名>/g'
 
 ## 算子索引
 
-> 硬件平台合计 **3 个**：`ascend910`（sdpa / sdpa_math）· `p800-kunlunxin`（sdpa / embedding）· `cambricon` / `mlu590`（MLU590，sdpa + 13 个算子）。 算子级明细见下表「硬件平台」列与各算子的 `REPORT.md` 表头。
+> 硬件平台合计 **3 个**：`ascend910`（sdpa / sdpa_math）· `p800-kunlunxin`（sdpa / sdpa_math / embedding）· `cambricon` / `mlu590`（MLU590，sdpa + 13 个算子）。 算子级明细见下表「硬件平台」列与各算子的 `REPORT.md` 表头。
 
 | 算子 | 硬件平台 | 开发级别 | 路线 | 状态 | 交付报告 | 备注 |
 |---|---|---|---|---|---|---|
 | [sdpa](sdpa/)（scaled_dot_product_attention） | **3** · ascend910 / p800-kunlunxin / mlu590 | Triton / 厂商委托 | A1（aten 拦截） | ✅ 定稿 | [REPORT](sdpa/REPORT.md) | ascend910 + p800-kunlunxin + mlu590 · 三平台黄金 397/397（Ascend 历史 265/265） · [平台绑定](sdpa/PLATFORM.md) / [多平台合并](sdpa/MERGE.md) |
-| [sdpa_math](sdpa_math/)（aten::_scaled_dot_product_attention_math） | 1 · ascend910 | Triton 级 | A1（aten 拦截） | ✅ 定稿 | [REPORT](sdpa_math/REPORT.md) | ascend910 · SDPA math 后端（双输出 out+概率图）· 黄金 175/175 · 原生对照 131/131（44 组 bool 按有意分歧跳过）· NPU fp16 1.20-2.05x |
+| [sdpa_math](sdpa_math/)（aten::_scaled_dot_product_attention_math） | 2 · ascend910 / p800-kunlunxin | Triton / vendor O/LSE + exact-P | A1（aten 拦截） | ✅ 定稿 | [REPORT](sdpa_math/REPORT.md) | SDPA math 后端（双输出 out+概率图）· 三平台黄金 175/175 · NPU fp16 1.16-2.03x · P800 direct 1.46-2.21x native |
 | [embedding](embedding/)（aten::embedding） | 2 · p800-kunlunxin / ascend910 | 厂商委托 / Triton 探针 | A1（aten 拦截） | ✅ 定稿 | [REPORT](embedding/REPORT.md) | p800-kunlunxin + ascend910 · 稠密查表/反向 · 黄金 117/117（Ascend 三层+350 PASS）|
 | [type_as](type_as/) | 1 · cambricon | Triton 级 | A1 aten | 三层全绿 | [REPORT.md](type_as/REPORT.md) | 黄金 111/111；位级一致 |
 | [clone](clone/) | 1 · cambricon | Triton 级 | A1 aten | 三层全绿 | [REPORT.md](clone/REPORT.md) | 黄金 48/48；存储独立 |
@@ -60,7 +60,8 @@ grep -rl my_op . | xargs sed -i 's/my_op/<你的算子名>/g'
 - **目录名 = 算子名**（小写下划线，如 `gelu_and_mul`）
 - **包化导入（多算子同进程必用）**: 目录含 `__init__.py`，
   `from ops.sdpa import register_a1` / `from ops.embedding import
-  register_a1` 顶层名带包前缀，互不遮蔽（根因与验证见
+  register_a1` / `from ops.sdpa_math import register_a1` 顶层名带包前缀，
+  互不遮蔽（根因与验证见
   sdpa/reports/e2e_mini_llm.md + ops/test_packaging.py）；
   单算子脚本模式（sys.path 注入目录）继续兼容
 - 未实现的级别**置空 + 说明**（`kernel/hardware_level/README.md`），

ops/sdpa_math/README.md

@@ -8,18 +8,20 @@
 | 项 | 值 |
 |---|---|
 | 算子 | `aten::_scaled_dot_product_attention_math`（torch SDPA 的 **math 后端**，`CompositeImplicitAutograd`） |
-| 语义 | `softmax(QKᵀ·scale + mask)` 双输出：`(out, attn_probs)` · GQA · causal · bool/float mask · dropout 双规则 · fp32 内部（fp64 输入保持 fp64，供 gradcheck） |
+| 语义 | `softmax(QKᵀ·scale + mask)` 双输出：`(out, attn_probs)` · GQA · causal · bool/float mask · dropout 双规则 · fp32 内部（fp64 保持；P800 半精度 fast path 见下） |
 | 路线 | **A1** aten 拦截（成对注册 `Autograd*` + 纯设备键），旧业务代码零改动 |
-| 平台 | **ascend910**: 自研 Triton（两段式 probs + PV）；**cpu**: ATen 组合（对照/梯度兜底） |
-| 验证 | kernel 52 组 ×2 profile · 黄金 **175/175** · 原生对照 **131/131**（44 组 bool 跳过）· op 21 项（6 gradcheck）· 应用层 8 项 ✅ |
-| NPU 性能（fp16，vs 同为"返回 out+P"的原生 math） | prefill1k D64 **1.20x** · 1k D128 **1.32x** · 2k D128 **2.05x** · GQA **1.74x** · decode 0.92x |
+| 平台 | **ascend910**: 自研 Triton（两段式 probs + PV，小形状单 kernel 融合）；**p800-kunlunxin**: vendor O/LSE + exact-P + A1；**cpu**: ATen 组合（对照/梯度兜底） |
+| 验证 | kernel 52 组 ×3 profile · 黄金 **175/175** · 原生对照 **131/131**（44 组 bool 跳过）· op 22 项（6 gradcheck + fused backward）· 应用层 8-9 项 · FlagOS E2E ✅ |
+| NPU 性能（fp16，vs 同为"返回 out+P"的原生 math） | prefill1k D64 **1.22x** · 1k D128 **1.29x** · 2k D128 **2.03x** · GQA **1.67x** · decode **1.17x** · tail100 **1.16x** |
 | CPU 性能（fp32，同口径） | 1.05-1.40x（6 形状全过） |
+| P800 性能（fp16，同口径） | direct **1.46-2.21x** native；A1 **1.31-2.17x** |
 
 ## 快速开始
 
 ```bash
 python3 example.py ascend910          # 三层一键（kernel/op/framework）
 python3 example.py cpu
+python3 example.py p800-kunlunxin
 
 python3 test/kernel_level.py ascend910
 python3 test/op_level.py ascend910
@@ -28,11 +30,17 @@ python3 test/framework_level.py ascend910
 # 黄金（CPU 生成一次，跨平台复用）
 python3 script/gen_golden.py
 python3 script/check_accuracy.py --impl triton   --device npu:0
+python3 script/check_accuracy.py --impl p800     --device cuda:1
 python3 script/check_accuracy.py --impl reference --device cpu
 python3 script/check_accuracy.py --impl native   --device npu:0   # 原生对照
 
 python3 script/bench_perf.py --device npu:0 --register \
   --json-out reports/perf_ascend910.json
+python3 script/bench_perf.py --device cuda:1 --register \
+  --json-out reports/perf_p800-kunlunxin.json
+python3 script/bench_f_context.py --device cuda:1 \
+  --json-out reports/perf_f_context_p800-kunlunxin.json  # no-P 参考列
+python3 script/e2e_flagos.py --device cuda:1 --dtype bfloat16
 python3 probes/native_semantics.py ascend910     # native 语义证据表
 ```
 
@@ -64,19 +72,63 @@ python3 probes/native_semantics.py ascend910     # native 语义证据表
 
 ```
 sdpa_math/
+├── __init__.py              # 包化入口：from ops.sdpa_math import register_a1
 ├── reference.py              # CPU 语义参考（判卷标准，含 dropout 双规则顶注）
 ├── register.py               # A1 成对注册 + autograd.Function（dout/dprobs 两路）
-├── _profile.py               # ascend910 / cpu 本地 profile
+├── _profile.py               # ascend910 / p800-kunlunxin / cpu 本地 profile
 ├── kernel/
 │   ├── triton_level.py       # 两段式 _probs_kernel + _pv_kernel（NPU）
+│   ├── p800_fast_level.py    # vendor O/LSE + exact-P（P800 fast path）
 │   └── torch_level.py        # ATen 组合（独立写法，非 reference 转发）
 ├── test/                     # kernel / op / framework 三层
 ├── goldendata/               # 声明式 175 组黄金（inputs_spec.yaml）
-├── script/                   # gen_golden / check_accuracy / bench_perf
+├── script/                   # gen_golden / accuracy / perf / F-context / FlagOS E2E
 ├── probes/native_semantics.py# native 语义证据（schema/参数形态/dropout 表/注册键）
 └── reports/                  # development / test-report / accuracy / performance
 ```
 
+## P800 / Kunlunxin 路线
+
+P800 生产实现选择 `p800_fast_level.sdpa_math_p800_fast`：
+
+```text
+A1 AutogradCUDA + CUDA 成对注册
+    ↓
+aten::_scaled_dot_product_efficient_attention → output + LSE
+QKᵀ → exp(scale·QKᵀ - LSE) → full P
+return (output, P)
+```
+
+本算子必须 materialize 第二输出 `P`，不能直接替换为 efficient attention。
+fast path 只在 fp16/bf16、无 mask、无 dropout 时启用；fp32、mask、dropout
+与 direct-autograd 场景回退 `torch_level`。P800 实测 direct
+**1.46-2.21x** native，A1 **1.31-2.17x**。
+
+`script/bench_f_context.py` 另提供 no-P `F.sdpa` 参考列：P800 fp16 下
+exact-contract math 约慢 **1.98-4.46x**。这不是同输出结论，因为 `F.sdpa`
+不返回 `P`、也不物化概率图；它说明后续 P800 内部融合有机会继续压缩
+QK/GQA expansion/mask/P 写出，但受“必须写全量 P”限制，不应以
+no-P FlashAttention 作为可达目标。源码级 exact-P Triton 双 pass 目前会触发
+XMLIR pointer-state rewrite 限制，故先采用 vendor LSE 组合。
+
+框架层运行在 FlagOS 算子栈内：
+
+```python
+import flag_gems
+flag_gems.only_enable(include=["gelu"])
+```
+
+锁定镜像上全量 `flag_gems.enable()` 非确定，因此选择稳定且真实被消费方
+调用的 GELU 作为 surrounding op；目标 `sdpa_math` 走 A1。
+
+端到端入口 `script/e2e_flagos.py` 自建 4 层 Llama 风格 mini-LLM
+（D128 · GQA 8/2 · vocab1024），每层直调本算子并把 `P` 送入注意力熵
+正则，覆盖 forward、backward、8 步贪心生成与 FlagGems GELU。实测
+推理 forward **1.07x**、仅 CE 训练 **1.45x**、CE+概率图熵正则训练
+**1.48x**，多步生成 **1.09x**；fused backward off/on A/B **1.45-1.50x**。
+原始数据见
+[reports/e2e_flagos_p800-kunlunxin.json](reports/e2e_flagos_p800-kunlunxin.json)。
+
 ## 硬约束（本栈已内建）
 
 1. Triton 启动必须包 `flag_gems.runtime.torch_device_fn.device` 上下文（#11）