ffpa-cuda-understand

## 使用说明（先读本节） - **本文件主体**是《ffpa-attn CUDA Backend 特性支持现状技术报告》（下称"报告"）： §1-§8 特性现状、§9 优化方向概览、§10 附录（环境变量 / bench / 构建 / 辅助 skill）、§11 核心技术点数学原理（推导级）。 - **references/**： - `rfc-future-optimizations.md`——工程级未来优化 RFC（功能完备性 > 性能优化 双轨，含**完成状态清单**、实施路线图、已证伪附录）。做优化前先读其附录 A； 做完一项勾选清单 + 更新总览表状态列。 - `ffpa_split_d_m4n2_design.md` / `ffpa_split_d_m4n4_analysis.md`——大 D 两大 核心技术之一 M4N2 TiledMMA<4,2,1> 的设计推导（配合报告 §11.4）。 - `sage-attention-principles.md`——SageAttention 1/2/2++/3 核心技术原理 （ffpa 量化数学的上游出处）。 - `flash-attention-principles.md`——FlashAttention 1-4 核心技术原理 （online softmax / tiling / WS / lazy rescale 的上游出处）。 - **论文原文**：`references/papers/`（SA1/SA2/SA2++/SA3 + FA-1/FA-2/FA-3/FA-4 文本版：`SageAttention{1,2,2++,3}.txt`、`FlashAttention-{1,2,3,4}.txt`）。 - **辅助 skill**（kernel 级优化/调试时搭配）：`cuda-auto-tune`（NCU 驱动迭代 优化）、`cuda-cpp-kernel`（CUDA/PTX 编写与 profiling）、`cutlass-cpp-kernel` （CUTLASS/CuTe C++ 模板）。 - **代码路径约定**：报告与 RFC 中所有代码路径均相对于 **ffpa-attn 仓库根目录**。 ## RFC实现规范（⚠️ 强制约束） - 1. 每项动手前：先在 plan 模式（Copilot下要切换到plan agent）完成实施规划（改动面 / 注入点 / 验证矩阵），规划好再动手 (自动模式下可以按照规划继续实施操作)。 - 2. 每做完一项：勾选对应条目（`- [ ]` → `- [x]`），并同步更新上方总览表状态列。注意，要同时更新 SKILL.md 中的技术报告和RFC 文档的总览表状态列，确保两处状态一致。 --- # ffpa-attn CUDA Backend 特性支持现状技术报告 > 基准代码：ffpa-attn `dev` 分支（strided-NHD 零拷贝已合入，对应 commit `02d49ab` 附近； > fp16 家族 per-tag variant TU 拆分 + 默认编译集裁剪已合入，对应 `0b90ddf` 附近）。 > 本报告梳理 CUDA backend 的全部 kernel 路径、特性矩阵、限制、核心技术与未来优化方向。 > 所有行号/行为以当前工作树为准。 --- ## 目录 1. [总体架构与分发链路](#1-总体架构与分发链路) 2. [后端实现枚举与 Python 入口门禁](#2-后端实现枚举与-python-入口门禁) 3. [Native 路径（NATIVE / TMA hint）](#3-native-路径) 4. [CUTE fp16/bf16 路径](#4-cute-fp16bf16-路径) 5. [CUTE FP8 路径](#5-cute-fp8-路径) 6. [CUTE FP4 (NVFP4) 路径](#6-cute-fp4-nvfp4-路径) 7. [横向特性矩阵](#7-横向特性矩阵) 8. [split-D / M4N2 相对 persist-D 的功能缺口](#8-split-d--m4n2-相对-persist-d-的功能缺口) 9. [未来优化方向](#9-未来优化方向) 10. [附录：环境变量 / CLI / 构建备忘](#10-附录) 11. [核心技术点数学原理](#11-核心技术点数学原理) --- ## 1. 总体架构与分发链路 ### 1.1 调用链 ```mermaid flowchart TD A["ffpa_attn_func(q,k,v,..., forward_backend=CUDABackend(...))"] --> B{推理 fast path 条件?} B -- 满足 --> C["_ffpa_attn_forward (functional.py)<br/>跳过 meta/autograd, ~20us"] B -- 不满足 --> D["FFPAAttnMeta 全链<br/>(autograd Function + 校验)"] C --> E{meta.fallback ?} D --> E E -- D≤256 显式关闭 small-D env / D>1024 /<br/>8≤Nq<512 / Nkv<512 / grad 等 --> F["SDPA fallback<br/>(torch._C._nn.sdpa)"] E -- 通过 --> G["_fwd_cuda torch op<br/>(cuda/__init__.py)"] G --> H["ffpa_api.cc<br/>dtype/acc/pad 分发 → generated per-headdim 入口"] H --> I["launch_ffpa_attn_fwd_template<br/>(csrc/cuffpa/launch/router.cuh)"] I --> J1["Native cp.async (sm80+)<br/>+ split-KV decode"] I --> J2["Native TMA (sm90/100 WS,<br/>sm120 non-WS)"] I --> J3["CUTE_TMA fp16 家族 (sm120)<br/>persist-D / split-D / M4N2"] I --> J4["CUTE_TMA_FP8<br/>persist-D / split-D / M4N2"] I --> J5["CUTE_TMA_FP4<br/>persist-D / split-D / M4N2"] I --> J6["CUTE hint cp.async<br/>(sm80+, 非 TMA)<br/>persist-D / split-D / M4N2"] ``` ### 1.2 关键文件 | 文件 | 职责 | |---|---| | `src/ffpa_attn/functional.py` | `CUDABackend` dataclass（全部量化 knob）、fast path decline 判定、NHD gate | | `src/ffpa_attn/ffpa_attn_interface.py` | SDPA 对齐签名入口、meta 校验、SDPA fallback 路由 | | `src/ffpa_attn/cuda/__init__.py` | `_fwd_cuda` torch op 注册、O/lse 分配、NHD permute 归一化 | | `csrc/cuffpa/ffpa_api.cc` | dtype×acc 分发、head_dim pad（O-only）、generated dispatcher | | `csrc/cuffpa/launch/router.cuh` | **顶层 dispatcher**：impl hint → 路径；hybrid stage-1 编排；NHD/strided 物化决策 | | `csrc/cuffpa/launch/{cute_fp16,cute_fp8,cute_fp4}.cuh` | CUTE 家族 launcher **伞形头**（per-impl 头 + 尾部 extern 表）：以 `cute_fp16.cuh` 为例 = common include + `cute_fp16_{persist_d,split_d,split_d_m4n2}.cuh`（单源 `ffpa::fp16_{impl}_bias_plan` + `_v` 变体体）+ 3 个 tag-dispatch wrapper（签名不变，lambda `launch_variant[_with_dropout]` 按 runtime bias/dropout 选 `(kBiasOn,kBiasPlanMode,kBias4BytesPerElem,kHasDropout)` tag 发 `_v`）+ 尾部 include `generated/fwd_cute_fp16_variants.cuh`（extern 表，variant TU 负责实例化） | | `csrc/cuffpa/launch/common.cuh` | CUTE 家族共享 helper：布局 gate（`ffpa_is_nhd_view` / `ffpa_is_strided_nhd` / `ffpa_layout_of`）、bias tile plan | | `csrc/cuffpa/launch/native_fp16.cuh` | Native kernel 编译期 config（MMA atom、stage、smem 复用、decode split 数选择） | | `csrc/cuffpa/cute/{fp8,fp4}/` | 量化前处理 kernel（quantize / kv_mean / delta_s / smooth）与主 kernel | ### 1.3 C++ dispatch/launch 详细调用链（router 之后） ```mermaid flowchart TD API["ffpa_api.cc::ffpa_attn_fwd_*<br/>dtype×acc 分发 + head_dim pad"] --> GEN["generated dispatcher<br/>(per-headdim case 表, env.py 生成)"] GEN --> TU["generated family TU<br/>fwd_{fp16,bf16}_{cute_fp16,cute_fp8,cute_fp4}_hdim{D}[_s{stage}].cu / fwd_{fp16,bf16}_native_hdim{D}[_s{stage}].cu<br/>显式实例化 family 入口模板（native 文件名 drop f32 后缀，符号保留全 variant 名）"] TU --> VT["generated variant TU（量化族 + fp16 族）<br/>fwd_{token}_cute_{f}_{impl}_hdim{D}_s{S}_b{b}m{m}f{f}[r{r}].cu<br/>每 (bias_on, plan_mode, bias_4B, dropout) tag 一个 TU，实例化 `_v` kernel 表"] VT --> DISP["dispatch/{native_fp16,cute_fp16,cute_fp8,cute_fp4}.cuh<br/>family 入口: 参数校验 + D 路由 + hybrid 编排"] DISP --> LAUNCH["launch/{family}.cuh<br/>launcher _impl: hadamard → 布局探测 → bias<br/>→ quant knob 派生 → kBr/kBc/smem/stages clamp → traits"] LAUNCH --> PRE["前处理编排"] LAUNCH --> TMA["TMA descriptor 构建<br/>(make_tma_copy Q/K/V/O)"] TMA --> KERN["attention 主 kernel launch<br/>(<<<>>> 或 cute cooperative)"] subgraph PRE["前处理（stage 无关）"] P1["fp8: cute/fp8/prepare_inputs.cuh<br/>prepare_fp8_inputs&lt;dtype,kBr,kBc,D,kQKInt8&gt;<br/>分配 + smooth-K mean + Q/K quantize + V per-channel<br/>⚠ extern-template: 仅 fwd_cute_fp8_preprocess.cu 实例化"] P2["fp4: launch_fp4_quant_*&lt;D&gt; 量化链<br/>(launcher 模板不带 dtype)<br/>⚠ extern-template: 仅 fwd_cute_fp4_preprocess.cu 实例化"] P3["fp8/fp4 共享: launch_kv_mean_sm120&lt;dtype,D&gt;<br/>(cute/fp8/smooth_k.cuh)"] end ``` **family 入口的 D 路由决策**（dispatch/cute_fp8.cuh、cute_fp4.cuh、cute_fp16.cuh 同构交叉点）： | 家族 | D≤224 | 224<D<768 | D≥768 | 备注 | |---|---|---|---|---| | fp16 (cute_fp16) | persist-D | split-D M8N1 | split-D M4N2 | M4N2 交叉点 768（sm120 dense 全家 M4N2 / sm80 仅 causal）；D%32≠0 或 bias/dropout 路由级回退 native；sm80 路由见 §4.5 | | fp8 | persist-D (kBc=128/D≤128, else 64) | split-D M8N1 (128/128) | M4N2 (64/64) | kBr/kBc 被 env.py::_fp8_variant_blocks 镜像（extern 表） | | fp4 | persist-D | split-D | M4N2 | 量化 launcher 模板仅 &lt;kHeadDim&gt;（dtype 无关） | | native | sm80 split-D + split-KV decode (Nq==1) | TMA: sm90/100 WS / sm120 non-WS | 同左 | decode split 数 = select_decode_num_splits 波效率贪心 | **hybrid 双阶段**（dispatch/cute_hybrid.cuh::prepare_hybrid_stage1）：`fp8_hybrid/fp4_hybrid && Nq ≥ n_early` 时，前 `n_early` 行走 fp16 stage-1（`ffpa_fwd_fp16_stage1<dtype,D,stage,224>`，cute_fp16 TU 实例化），其余行 stage-2 以 `q_start_row` 偏移交回量化家族 launcher——两阶段各写自己行集的 O/lse 切片。 **编译期 TU 结构**（谁实例化什么）： ```mermaid flowchart LR subgraph FamilyTU["family TU ×(dtype,D,stage)"] A["dispatch/cute_fp8.cuh"] --> B["launch/cute_fp8.cuh<br/>(extern: 前处理不实例化)"] B --> C["cute/fp8/sm_120/*.cuh 主 kernel<br/>(stage 相关, 每 TU 独有)"] end subgraph SharedTU["fwd_cute_fp8/fp4_preprocess.cu (×2)"] D["generated/fwd_cute_fp8_preprocess.cuh<br/>define 模式: prepare_fp8_inputs 全组合"] E["generated/fwd_cute_fp4_preprocess.cuh<br/>define 模式: launch_fp4_quant_* 全组合"] end B -. "extern template 声明<br/>(preprocess.cuh extern 模式)" .- D ``` **默认编译集裁剪（构建期 flags，env.py `_enabled_variants()`）**：冷构建 770s/451 TU → **525s/326 TU（-32%）**。两条轴： | flag | 默认 | 效果 | |---|---|---| | `ENABLE_FFPA_CUDA_INPUT_FP16` | 0 | 只编 bf16 输入 TU（fp16 输入 kernel 不编译，`fwd_{fp16,bf16}_native_*` 文件名 drop f32 后缀、fp16f16 保留全名） | | `ENABLE_FFPA_CUDA_MASK_FP32` | 0 | 剔除所有 `f=1`（bias 4B/fp32 mask）variant TU；量化族与 fp16 族统一 | 运行时配套（`functional._cuda_input_guard_and_mask_downcast`）：fp16 输入 raise（提示 rebuild flag）；fp32 mask downcast 到 `q.dtype` 再调（数值等价，PC-0-5-pd 缓解选项 B 的默认形态；全开构建的 fp4 persist_d C++ launcher 侧已镜像同语义 downcast——PC-0-5-pd-B，覆盖直调绕过面）。headdim 128+512 子集下默认 97 TU / 全开 255 TU。 **单条完整链路示例**（fp8、D=512、bf16、stage=2、无 hybrid）： `CUDABackend(enable_fp8=True)` → `_fwd_cuda` → `ffpa_api.cc` pad→512 → `fwd_bf16_cute_fp8_hdim512_s2.cu`（实例化 `ffpa_fwd_fp8<bf16,512,2>`）→ `dispatch/cute_fp8.cuh`：D=512∈(224,768) → `launch_cute_fwd_split_d_fp8_sm120<bf16,512,2>` → `_impl`：hadamard off → `ffpa_layout_of`×3 → kBr/kBc=128/128 → `prepare_fp8_inputs<bf16,128,128,512,kQKInt8>`（**extern→preprocess TU**）→ TMA Q/K/V/O → `split_d_fwd_cute_fp8_sm120<...Traits<512,...,kStagesQK=2,...>>`。 --- ## 2. 后端实现枚举与 Python 入口门禁 ### 2.1 `CudaBackendImpl` 与 `CUDABackend` flag 映射 `CudaBackendImpl`（`backend.h`）：`AUTO(0) / NATIVE(1) / TMA(2) / CUTE(3) / CUTE_TMA(4) / CUTE_TMA_FP8(5) / CUTE_TMA_FP4(6)`。 `_apply_cuda_backend_hint`（functional.py）优先级： | CUDABackend 字段 | hint | |---|---| | `enable_fp4=True` | `CUTE_TMA_FP4` | | `enable_fp8=True` | `CUTE_TMA_FP8` | | `enable_tma and enable_cute` | `CUTE_TMA` | | `enable_tma` | `TMA` | | `enable_cute` | `CUTE` | | 都不开 | `NATIVE`（`AUTO` 视同 `NATIVE`） | 注意 hint 是**进程级全局原子**（`set_cuda_backend_impl`），每次调用前由 Python 侧刷新；直调 `_ffpa_attn_forward_cuda` 绕过 Python 时必须手动 set，否则静默走错路径（见 memory `tool-use-pitfalls`）。 ### 2.2 推理 fast path 的 decline 条件（返回 `None` → 全链/SDPA） `_ffpa_attn_forward` 逐条检查（全部命中才走 CUDA fast path）： 1. `_ffpa_attn_forward_cuda` 已编译（`ENABLE_FFPA_CUDA_IMPL=1`） 2. `torch.is_grad_enabled()` 为 False 3. 非 `bf16 + acc='f16'`（无 bf16-acc MMA PTX，meta 层也强制拒绝） 4. small-D 路由：`D ≤ 256` 且显式设 `FFPA_CUDA_ALLOW_SMALL_D=0` → decline（SDPA）；默认（未设或 `=1`）走 CUDA kernel 5. `D > 1024` → decline 6. `8 ≤ Nq < 512` 或 `Nkv < 512` → decline（短序列不划算；`Nq ∈ [1,8)` 不 decline，可走 native decode）