feat: add sdpa_math operator (Ascend Triton + P800 exact-P/A1) by TruNcat3 · Pull Request #10 · FLAG-OP/FlagOS-OP

TruNcat3 · 2026-09-30T06:57:52Z

概述

新增算子 ops/sdpa_math —— aten::_scaled_dot_product_attention_math（SDPA 的 math 后端，返回 (out, attn_probs) 双输出），A1 路线接管，与既有 ops/sdpa 目录独立、代码零依赖。为"必须拿到注意力概率图"的消费方（蒸馏 / 熵正则 / rollout / 对齐分析）提供一个快而正确的实现。

实现

三级实现：reference.py（纯 PyTorch 数学权威）· kernel/torch_level.py（ATen 组合，CPU 交付）· kernel/triton_level.py（Triton 两段式，NPU 交付）
A1 注册：register_a1() 成对注册 AutogradPrivateUse1+PrivateUse1 / AutogradCPU+CPU；每次注册绑定独立 _Fn 闭包（同进程多 impl 不互相覆盖）；grad 走 autograd.Function（dout+dprobs 双路梯度），no_grad/inference 直通
语义固化：显式/随机 dropout 双规则、causal+mask 互斥、全遮蔽行 P=O=0、GQA 不整除报错 —— 均有 probes/native_semantics.py（6 节实测证据）+ 黄金互验 + 测试三重锁定

验证（ascend910 与 cpu 双 profile）

项	结果
黄金回归	reference / torch / triton 175/175；原生对照 131/131 + 44 组 bool 按有意分歧跳过
kernel 层	52 组（3 dtype + 非整倍数 shape + 哨兵 + dropout + 错误路径），最差 fp32 1.79e-7 / fp16 1.95e-3 / bf16 7.81e-3
op 层	21 项：拦截 3822 次、A1 包装=直调逐位、6 组 fp64 gradcheck、`F.sdpa(MATH)` 命中
应用层	8 项：mini-decoder logits L∞ 1.86e-8、top1 1.000、全权重梯度 L∞ 1.75e-10、贪心序列 1.000
性能门禁	`perf_compare` FAIL 0 · WARN 0，基线 `triton 0.504ms / 8.524 TFLOPS`

python3 example.py ascend910                     # 三层一键（也支持 cpu）
python3 script/gen_golden.py && python3 script/check_accuracy.py --impl triton --device npu:0
python3 ../../scripts/perf_compare.py --device ascend910
python3 probes/native_semantics.py               # native 语义证据表

性能（NPU fp16，基线 = 同样返回 `(out,P)` 的原生 math 后端）

形状	自研	原生	speedup
prefill 1k D64	0.515ms	0.626ms	1.22x
prefill 1k D128	0.538ms	0.695ms	1.29x
prefill 2k D128	1.369ms	2.776ms	2.03x
GQA 1k D128	0.810ms	1.351ms	1.67x
decode S64 / tail S100	0.217 / 0.219ms	0.255 / 0.253ms	1.17x / 1.16x

CPU fp32 六形状 1.05–1.40x。对照不返回 P 的 F.sdpa 融合路径约慢 4x —— 物化概率图的契约代价，判卷只用同契约基线。

有意分歧与遗留

bool attn_mask 直调按 -inf 遮蔽实现（对齐 F.sdpa），与 native 直调的 0/1 加性怪癖有意分歧 → 44 组黄金在 --impl native 下跳过并注明
decode/tail 的 A1 包装路径 0.86x（直调经单 kernel 融合已 1.16–1.17x 反超）：autograd.Function 包装 ~0.078ms 属 torch 图节点开销，非 kernel 可控
dropout_p>0 的 A1 路径 out=P@V 走 ATen matmul（dropout 非性能路径，register.py 已注明）

交付物

README + REPORT + 报告四件套（development / test-report / accuracy / performance）+ 175 组黄金 + 4 脚本 + 3 层测试 + native 语义 probe；已登记 ops/README.md 索引、common/perf_registry.py 与 docs/known-issues.md #19。

2026-10-01 update: P800 / Kunlunxin second platform

1e69c53 adds the P800 profile and backend path:

production implementation: kernel/torch_level.py ATen composition;
A1 registration: AutogradCUDA + CUDA paired keys;
framework/application consumer runs inside the FlagOS operator stack with flag_gems.only_enable(["gelu"]);
benchmark and perf-gate selection are platform-aware.

P800 validation

item	result
kernel	52/52
golden (`torch @ cuda:1`)	175/175, worst 9.766e-4
A1 op layer	21/21, including 6 fp64 gradchecks
FlagOS application	9/9; logits / probability-map / gradient diffs all 0
P800 perf gate	FAIL 0 · WARN 0 · NEW 0

P800 fp16 performance

Native baseline is the same private math op and therefore has the same (output, attention) contract.

shape	direct ours	native	speedup	A1	A1 speedup
prefill 1k D64	1.077ms	1.197ms	1.111x	0.820ms	1.459x
prefill 1k D128	0.804ms	1.019ms	1.267x	0.834ms	1.222x
prefill 2k D128	2.392ms	2.999ms	1.253x	2.420ms	1.239x
GQA 1k D128	1.363ms	1.727ms	1.266x	1.404ms	1.230x
decode D128	0.370ms	0.382ms	1.031x	0.407ms	0.939x
tail100 D64	0.328ms	0.383ms	1.166x	0.375ms	1.019x

Raw data: ops/sdpa_math/reports/perf_p800-kunlunxin.json.

2026-10-01 update: 单 kernel 融合（小形状）

e0b9c87 + 09ce990：_probs_kernel 增加 HAS_PV 模式 —— QK→softmax→P→PV 在一个 kernel 内完成（P 块只在寄存器里、O 直写，probs 仍按契约物化），causal 尾块补零也收进 kernel（去掉独立 torch.zeros 启动）。

选择判据：按 causal 截断后实际处理的 score tile 数，阈值 16 tiles（实测交叉点 15~21：15 tiles 0.98x、21 tiles 1.07x；同面积 causal 512²=36 tiles 融合慢 13%、128×2048=3 tiles 反而快 15%，故不能用面积判）；dropout_p>0 恒走两段式（掩码须先作用在 P）。SDPA_MATH_FUSED=on/off/auto 可强制（仅实验）。
收益：decode 0.287→0.217ms（0.92x→1.17x）、tail100 0.285→0.219ms（0.97x→1.16x）；融合路径 launch 3→1（本栈单次启动地板 ~97µs）。大形状按阈值走两段式，与融合前持平。
回归：黄金 175/175（worst 3.906e-3 不变，两路径都覆盖）· example.py ascend910 与 cpu 三层全绿 · 门禁 FAIL 0 · WARN 0（triton 0.543 vs 基线 0.504，+7.8%，阈值内）。
18 组阈值实测表与栈级差距拆解（GEMM 5.2x / UB 192KB / P 往返 / launch）见 reports/performance.md §5–§6。

2026-10-02 update: 微优化 + 融合阈值重定标（16→56）+ 端到端

7885074 + 6ce4b59 + 83482a1（报告新增 §8/§9，8 个有效探针入库 probes/）：

微优化：_d1() 仅要求 stride(-1)==1（transposed 输入免 contiguous，单次 ~14µs）；P+O 合并为一次 empty + 两个 view（省一次设备启动）。直调口径 decode 195.2→149.5µs（-23%）、prefill 288.2→223.8µs（-22%）。
融合阈值 16→56 tiles：e2e decode_cache 形状（kv 1025→1152 = 17-18 tiles）原被阈值 16 排除 → 强制融合实测 -15%（331→273ms）；decode_re（causal 36→55 tiles）强制融合 -9%。64+ tiles 与 prefill 136 tiles 两段式仍胜（0.684 vs 0.522ms）→ 定标 56。单次同步的 latency 口径在 32-55 tiles 会低估融合（两段式多 1-2 次启动 ~97µs，流水下才摊薄），故以 e2e 实测为准；门禁形状（1/3/136 tiles）两侧均不受影响。
端到端（probes/e2e_probe.py，4 层 D512 H8 + 每层熵正则消费 P，fp16 注册路径）：prefill 2.6 vs 2.9ms（-10%）、decode_re 297 vs 304ms（反超 native）、decode_cache 291 vs 241ms（+21%——差距构成见 §8：attn 调用 host 发射/注册包装 ~120µs·call⁻¹，设备 kernel 时间已接近 native）。
decode 差距归因（§8，含测量陷阱）：register_a1() 返回的 lib 列表不被持有会随 GC 失效（早期多支探针 "ours" 实为 native，已作废）；修正后同进程 A/B 显示 prefill ours 反超（382.8 vs 489.7µs）、decode 落后 ~39µs、设备时间 native 189.2 vs ours 直调 196.5µs → 差距主要在 host 侧（triton launch 83.7 + 分配 19.4 + python ~34）。
回归：黄金 175/175（worst 3.906e-3 不变）· example.py 双端三层全绿 · Ascend 门禁 FAIL 0 · WARN 0（triton 0.544 vs 基线 0.504，+7.9%）。

2026-10-02 update: P800 exact-P fast path

88a3b57 replaces the P800 production selection with kernel/p800_fast_level.py, while preserving the operator's exact (output, attention_probs) contract:

aten::_scaled_dot_product_efficient_attention
    -> vendor output + row-wise LSE
P = exp(scale * QK^T - LSE)
    -> materialize the full probability map
return (output, P)

The fast path applies to fp16/bf16 no-mask/no-dropout calls. fp32, masks, dropout, and direct autograd construction fall back to torch_level; A1 continues to use the custom mathematical backward.

Updated P800 validation

item	result
golden (`p800 @ cuda:1`)	175/175, worst 3.906e-3 (bf16, below 2e-2 tolerance)
kernel layer	52/52
A1 op layer	21/21, including 6 fp64 gradchecks
FlagOS application layer	9/9, with FlagGems GELU enabled around the target A1 op
P800 perf gate	FAIL 0 / WARN 0 / NEW 0; fast 0.489ms vs native 1.039ms

Updated P800 fp16 performance

shape	fast direct	native	direct speedup	A1	A1 speedup
prefill 1k D64	0.509ms	1.021ms	2.004x	0.497ms	2.054x
prefill 1k D128	0.494ms	1.036ms	2.096x	0.499ms	2.075x
prefill 2k D128	1.429ms	3.013ms	2.109x	1.442ms	2.089x
GQA 1k D128	0.783ms	1.726ms	2.206x	0.796ms	2.168x
decode D128	0.290ms	0.474ms	1.637x	0.311ms	1.526x
tail100 D64	0.278ms	0.405ms	1.458x	0.310ms	1.307x

A non-equivalent no-P F.sdpa reference remains 1.98–4.46x faster because it does not materialize P; this is documented separately rather than used as a same-contract speedup. No same-contract A100 private-math measurement is available yet.

2026-10-02 update: FlagOS end-to-end test

944b050 adds ops/sdpa_math/script/e2e_flagos.py and packages the operator as ops.sdpa_math.

Workload

4-layer Llama-style mini-LLM
D=128, GQA heads 8/2, FFN=256, vocab=1024
Each attention block calls aten::_scaled_dot_product_attention_math
The returned probability map is consumed by an attention-entropy regularizer
Native and plugin phases both run with FlagGems GELU enabled
Plugin phase registers ops.sdpa_math.register_a1 on AutogradCUDA+CUDA
Separate subprocesses avoid irreversible in-process torch.library registration
Covers no-grad forward, CE + entropy backward, and 8-step greedy generation

E2E correctness

check	result
A1 interceptions	856
logits L∞	0.015625 (bf16 threshold 0.02)
top-1 agreement	98.75%
attention entropy abs diff	7.71e-4
greedy sequence	exact match
embedding / q0 / head grad L∞	1.14e-5 / 3.81e-6 / 2.44e-4

E2E latency

Shared-machine repeated runs show:

workload	native/plugin speedup range
forward	1.02–1.09x
generate	1.06–1.12x
training with CE + `dprobs` entropy path	0.89–1.00x

The model-level gain is diluted by GEMM/MLP work. Training can consume the fast-forward gain through the custom mathematical backward, especially the probability-map gradient path; this is recorded as the next optimization target rather than hidden.

Raw report: ops/sdpa_math/reports/e2e_flagos_p800-kunlunxin.json.

2026-10-07 update: P800 fused backward and E2E A/B

4671389 extends the P800 fast path into A1 backward.

Implementation

The P800 forward now optionally returns vendor auxiliary state:
- output
- probability map
- log-sumexp
- philox seed/offset
- effective scale
A1 saves that state for eligible fp16/bf16, no-mask, no-dropout calls.
dout uses aten::_scaled_dot_product_efficient_attention_backward.
dprobs computes the exact softmax Jacobian from the already materialized probability map.
GQA Q/K gradients are combined and reduced back to KV heads.
Mask, dropout, fp32, and other fallback cases retain the generic mathematical backward.
SDPA_MATH_P800_FUSED_BWD=0 can disable the new branch for attribution.

Correctness

New op-level P800 case: fp16 causal GQA with both output and probability gradients.
A1 fused backward vs native maximum gradient error: 3.91e-3.
P800 one-button remains green:
- kernel 52/52
- A1 op 22/22
- FlagOS application 9/9
- golden 175/175, worst 3.906e-3

Updated FlagOS E2E

The mini-LLM E2E now separates ordinary CE training from CE plus probability-map entropy training.

workload	native	A1	speedup
forward B1 S256	7.683ms	7.134ms	1.077x
forward B2 S128	7.659ms	7.134ms	1.074x
CE training	36.225ms	24.941ms	1.452x
CE + attention-entropy training	42.733ms	28.793ms	1.484x
8-step generation	66.201ms	60.835ms	1.088x

Fused-backward off/on A/B:

training workload	generic backward	fused backward	off/on
CE only	37.376ms	24.941ms	1.499x
CE + probability entropy	41.723ms	28.793ms	1.449x

E2E correctness remains green: logits L∞ 0.015625, top-1 agreement 99.375%, greedy sequence exact match, and selected gradient L∞ values remain below 2e-4.

TruNcat3 · 2026-10-01T06:51:07Z

Added P800/Kunlunxin as the second sdpa_math platform in 1e69c53.

production path: torch_level ATen composition
A1: paired AutogradCUDA + CUDA registration
framework test runs with FlagGems GELU active
validation: kernel 52/52, golden 175/175, A1 21/21, FlagOS application 9/9
P800 fp16 direct performance: 1.03-1.27x native; A1 large shapes 1.02-1.24x
P800 perf gate: FAIL 0 · WARN 0 · NEW 0

Raw data and repro commands are in ops/sdpa_math/reports/perf_p800-kunlunxin.json and reports/performance.md.

- _probs_kernel 增 HAS_PV：QK→softmax→P 落地→PV 同 tile 累加，1 次启动 - causal 尾块补零收进 kernel，P 一律 empty（去掉独立 torch.zeros 启动） - 选择判据按实际处理 score tile 数（causal 截断后），阈值 16： 15 tiles 0.98x · 21 tiles 1.07x · 64 tiles 1.16x（同面积 128x2048 反而 0.85x） - dropout_p>0 恒走两段式（掩码须先作用在 P） - 实测：decode 0.287→0.217ms(0.92x→1.17x)、tail 0.285→0.219ms(0.97x→1.16x)，大形状按阈值走两段式持平；门禁 FAIL 0（triton +7.8%，阈值内） - 回归：黄金 175/175 worst 3.906e-3 不变；三层 ascend910/cpu 全绿

- performance.md §5 补栈级差距拆解（GEMM 5.2x / UB 192KB / P 往返 / launch 97µs）与 JIT 口径，新增 §6 单 kernel 融合实验（tiles 阈值表 + 18 组实测） - development/test-report/REPORT/README: 6 形状性能表更新为融合后数字， decode/tail 遗留项改记 A1 包装 0.86x，后续优化项标记已完成 - ops/README.md 索引行 1.20-2.05x → 1.16-2.03x

- _d1(): 仅要求 stride(-1)==1，transposed 输入免物化（单次 ~14µs） - P+O 合并为一次 empty + 两个 view（省一次分配/启动） - 直调口径: decode 195.2→149.5µs (-23%)、prefill 288.2→223.8µs (-22%) - 融合阈值 16→56: e2e decode_cache (17-18 tiles) 强制融合 -15%、 decode_re (36-55 tiles) -9%；64+/prefill 136 tiles 仍走两段式。依据为 e2e 流水口径实测（latency 单次同步口径在 32-55 tiles 低估融合） - 回归: 175/175 worst 3.906e-3 不变、example 双端全绿、门禁 FAIL 0

- performance.md §8: decode 差距归因（GC 丢注册测量陷阱、同进程 A/B、 host/device 分离、e2e 分段、布局排除） - performance.md §9: e2e 三场景结果 + 本轮三项改动分解（跳 contiguous、 P+O 合并分配、融合阈值 16→56）与口径分歧说明 - §2/§4 数字随 2026-10-02 复测更新（环境漂移同向，门禁 FAIL 0） - §6/development.md/REPORT.md 阈值引用 16→56 同步 - probes/ 入库 8 个有效探针（去掉绝对路径，相对 repo 定位）

- §4: triton +7.9% / torch -0.4% / reference -0.8%（rebase 后 FAIL 0），撤回窗口性 +13.9% 漂移的临时读数 - §2/§5: 跨日绝对值标注共享机 ±5% 波动，判定看同运行 speedup/趋势

TruNcat3 · 2026-10-03T02:47:49Z

2026-10-02 更新总结（最终态）

本轮改动（7885074 perf + 6ce4b59 docs + 83482a1 docs，均已过 CI）：

微优化：_d1() 仅要求 stride(-1)==1（transposed 输入免 contiguous）+ P+O 合并为一次 empty → 直调口径 decode 195.2→149.5µs（-23%）、prefill 288.2→223.8µs（-22%）
融合阈值 16→56 tiles（e2e 流水口径定标；latency 单次同步口径在 32-55 tiles 会低估融合）→ e2e decode_re -9%、decode_cache -12%；64+ 与 prefill 136 tiles 仍走两段式，门禁形状（1/3/136）不受影响
报告 §8/§9 + 8 个探针入库：decode 差距归因（含 register_a1() 返回 lib 不持有会随 GC 失效的测量陷阱）与优化分解

总体性能（vs 同契约原生 math，NPU fp16，2026-10-02）：

形状	ours	native	speedup
prefill 2k D128	1.370ms	2.800ms	2.04x
GQA 1k D128	0.838ms	1.415ms	1.69x
prefill 1k D128	0.548ms	0.694ms	1.27x
prefill 1k D64	0.549ms	0.629ms	1.15x
decode D128	0.243ms	0.270ms	1.11x
tail100 D64	0.246ms	0.268ms	1.09x

端到端（4 层 D512 H8 + 熵正则消费 P）：prefill 2.6 vs 2.9ms（-10%） · decode_re 297 vs 304ms（反超） · decode_cache 291 vs 241ms（+21%，归因：设备 kernel 时间已 ≈ native，残余为 host 发射 + 注册包装 ~120µs·call⁻¹，见 §8）

验证：黄金 175/175（worst 3.906e-3 不变）· example.py 双端三层全绿 · Ascend 门禁 FAIL 0（triton +7.9%）· CPU 1.05-1.40x · P800 exact-P 门禁 OK

天花板：vs F.sdpa 仍差 ~5.2x 为栈级差距（P 按契约必须物化 + 同栈 GEMM 5.2x + UB 192KB tile 封死），单算子内不可追平；消费方不需要概率图时应优先用 F.sdpa。

TruNcat3 · 2026-10-07T04:10:16Z

📦 后续维护提示（合入后补记，#16 审计整改关联）：

本 PR 合入后 ops/ 做了全库审计整改（#16），与你相关的变化：

ops/sdpa_math 现在有 PLATFORM.md——你的双平台绑定项（ascend910 Triton 64×64 tile 上限、p800 vendor O/LSE + exact-P）已沉淀；新平台接入参考其中 §4
合并前自查：templates/operator/README 新增「交付自查清单」（核心八件/三层测试/性能对照/黄金/包化回归/多平台文档/署名七项），后续平台扩展合并前照单打勾
包化回归：python3 ops/test_packaging.py 现覆盖全 16 算子导入扫描，是你 PR 里的 __init__.py 之外的最后一道门

感谢按 #14/#15 新范式交付（本 PR 是三个带包化入口的算子之一）。

TruNcat3 · 2026-10-07T11:44:29Z

按“后续维护提示”做作者自查时发现并修复一个当前平台问题：多算子 perf registry 加载后，ops.sdpa P800 backend 的裸 reference 导入会被 sdpa_math.reference 遮蔽，导致 ops.sdpa.p800.forward SKIP；同时旧 perf JSON 会让 compare 假绿。修复与复测在 #22。

TruNcat3 changed the title ~~feat: 新增算子 ops/sdpa_math（aten::_scaled_dot_product_attention_math · Triton 级 · A1 · ascend910）~~ feat: add sdpa_math operator (Ascend Triton + P800 torch/A1) Oct 1, 2026

TruNcat3 force-pushed the feat/op-sdpa-math branch from 1ea7fee to 88a3b57 Compare October 2, 2026 03:23

TruNcat3 changed the title ~~feat: add sdpa_math operator (Ascend Triton + P800 torch/A1)~~ feat: add sdpa_math operator (Ascend Triton + P800 exact-P/A1) Oct 2, 2026

FlagOS Templates and others added 5 commits October 2, 2026 11:56

feat: add P800 backend to sdpa_math

672020b

perf: add P800 exact-P fast path to sdpa_math

ae62050

test: add FlagOS sdpa_math end-to-end benchmark

944b050

TruNcat3 force-pushed the feat/op-sdpa-math branch from 88a3b57 to 944b050 Compare October 2, 2026 03:59

TruNcat3 added 3 commits October 2, 2026 13:40

perf: reuse vendor attention backward on P800

4671389

TruNcat3 merged commit f7b22a2 into main Oct 7, 2026
2 checks passed

TruNcat3 mentioned this pull request Oct 7, 2026

fix: guard P800 sdpa perf loading #22

Merged

3 of 7 tasks

Navigation Menu

Repository navigation

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

feat: add sdpa_math operator (Ascend Triton + P800 exact-P/A1)#10
TruNcat3 merged 9 commits into
mainfrom
feat/op-sdpa-math

TruNcat3 commented Sep 30, 2026 •

edited

Loading

Uh oh!

TruNcat3 commented Oct 1, 2026

Uh oh!

TruNcat3 commented Oct 3, 2026

Uh oh!

Uh oh!

TruNcat3 commented Oct 7, 2026

Uh oh!

TruNcat3 commented Oct 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

TruNcat3 commented Sep 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

概述

实现

验证（ascend910 与 cpu 双 profile）

性能（NPU fp16，基线 = 同样返回 (out,P) 的原生 math 后端）

有意分歧与遗留

交付物

2026-10-01 update: P800 / Kunlunxin second platform

P800 validation

P800 fp16 performance

2026-10-01 update: 单 kernel 融合（小形状）

2026-10-02 update: 微优化 + 融合阈值重定标（16→56）+ 端到端

2026-10-02 update: P800 exact-P fast path

Updated P800 validation

Updated P800 fp16 performance

2026-10-02 update: FlagOS end-to-end test

Workload

E2E correctness

E2E latency

2026-10-07 update: P800 fused backward and E2E A/B

Implementation

Correctness

Updated FlagOS E2E

Uh oh!

TruNcat3 commented Oct 1, 2026

Uh oh!

TruNcat3 commented Oct 3, 2026

2026-10-02 更新总结（最终态）

Uh oh!

Uh oh!

TruNcat3 commented Oct 7, 2026

Uh oh!

TruNcat3 commented Oct 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

TruNcat3 commented Sep 30, 2026 •

edited

Loading

性能（NPU fp16，基线 = 同样返回 `(out,P)` 的原生 math 后端）