diff --git a/alto/kernels/dispatch/attention.py b/alto/kernels/dispatch/attention.py index 3372762..c5d0ae3 100644 --- a/alto/kernels/dispatch/attention.py +++ b/alto/kernels/dispatch/attention.py @@ -6,6 +6,7 @@ from torchtitan.models.common.attention import (ScaledDotProductAttentionWrapper) from alto.kernels.fp4.mxfp4.triton_flash_attention_mxfp4 import triton_attention_mxfp4 +from alto.kernels.mxfp8.triton_flash_attention_mxfp8 import triton_attention_mxfp8 from .config import TrainingOpConfig __all__ = ["LPScaledDotProductAttentionWrapper"] @@ -20,6 +21,8 @@ def __init__(self, config: TrainingOpConfig): if isinstance(config, TrainingOpConfig) and config.precision == "mxfp4": self.attn_func = triton_attention_mxfp4 + elif isinstance(config, TrainingOpConfig) and config.precision == "mxfp8_e4m3": + self.attn_func = triton_attention_mxfp8 else: raise ValueError(f"Unsupported SDPA config: {config}") diff --git a/alto/kernels/mxfp8/MXFP8_ATTENTION_PLAN.md b/alto/kernels/mxfp8/MXFP8_ATTENTION_PLAN.md new file mode 100644 index 0000000..a0a2219 --- /dev/null +++ b/alto/kernels/mxfp8/MXFP8_ATTENTION_PLAN.md @@ -0,0 +1,273 @@ +# MXFP8 E4M3 Flash-Attention（forward + backward） + +在 **AMD MI350 (CDNA4 / gfx950)** 上支撑 mxfp8 训练的 flash-attention kernel，前向与反向（dQ / dK / dV）均已实现。 + +**状态：已在 CDNA4 真机验收通过。** §7 的六项验收标准全部达成；测试 75 项全绿。本文档描述**已实现的现状**，末尾 §9 记录关键决策的理由。 + +--- + +## 1. 范围与设计决定 + +1. **forward + backward 都做。** forward 由 mxfp4 attention 机械改写，backward 由 `blockwise_fp8` attention backward 机械 port（三 kernel 结构同构）。 +2. **全 e4m3。** forward 的 Q/K/V/P 与 backward 的 dO/dP/dS 全部 e4m3，kernel 无 dtype 分发。 + ⚠️ **e5m2 没有参数化预留。** 格式以 `E4M3_TARGET_MAX_POW2 / E4M3_MBITS / E4M3_FORMAT_ID` 三个模块级 constexpr 硬编码，公开 API 也没有格式开关。若将来要切 e5m2，**必须改 kernel**，不是切参数。§3 的数值实测表明目前不需要切。 +3. **仅 CDNA4，无 fallback。** kernel 只走 `tl.dot_scaled`（e4m3×e4m3），不写 CDNA3 dequant fallback、无 `USE_DOT_SCALED` 开关。MI300 / CDNA3 不在范围内。数值 ground-truth 由 PyTorch 侧模拟 reference 提供（§6）。 + +## 2. 参照实现来源 + +| 部分 | 来源 | +|---|---| +| forward kernel | `alto/kernels/fp4/mxfp4/triton_flash_attention_mxfp4.py`（FA v2 + mxfp4，forward-only） | +| backward 三 kernel | `alto/kernels/blockwise_fp8/triton_flash_attention_fp8_block.py`（`_bwd_preprocess` + `_bwd_kernel_dkdv` + `_bwd_kernel_dq`）。**mxfp4 没有 backward，不是这部分的来源** | +| 量化原语 | 本目录 `mxfp8_quantization.py`（`convert_to_mxfp8` / `_calculate_scales` / `_quantize_fp8`） | +| 流程与验收范式 | 本目录 `MXFP8_GROUPED_GEMM_PLAN.md`（「机械改写 + 分层校验 + 真机验证」） | + +**forward 相对 mxfp4 的改动**：删掉全部 head_dim packing（mxfp4 一 byte 装两个 e2m1，mxfp8 一 byte 一元素）——包括 `head_size_q/k/v *= 2`、`HALF_BLOCK_DMODEL_*`、`offs_d_*_pack`、`tl.dot_scaled` 的 `*_k_pack` 参数；`e2m1` → `e4m3`；`_pack_fp4` → `_quantize_fp8`；`convert_to_mxfp4` → `convert_to_mxfp8`。与量化正交的 FA v2 骨架（causal masking、GQA/MQA、LSE 写回、padded head、online-softmax、全 0 块 early-exit、`triton_op` + `autograd.Function` 三段式）原样保留。 + +## 3. 数值格式：全 e4m3 + +attention 的四个 forward operand 都偏「动态范围小、单元素精度重要」，e4m3 天然合适： + +| operand | 分布特征 | e4m3 | +|---|---|---| +| Q / K | LayerNorm 后激活，分布集中（±几十） | ✅ 单元素精度更重要 | +| V | 同上 | ✅ | +| P（softmax 概率） | ∈ [0,1]，行和为 1；长尾但有界，被 online-softmax 逐行重归一化 | ✅ ~6% 相对误差可接受 | + +**backward 全 e4m3 的顾虑已实测排除。** grad 是长尾 + 大动态范围分布（`dP`/`dS` 尤甚），理论上是 e5m2 的场景，曾担心小尾部 underflow。实测（stage-2 全量化参照 vs SDPA autograd）：dQ/dK SNR ≈ 23 dB、cossim ≈ 0.9976；dV SNR ≈ 25~26 dB、cossim ≈ 0.9985。**未出现崩盘**，故不切 e5m2。 + +## 4. 接口契约 + +### 4.1 用户 API + +```python +def triton_attention_mxfp8( + q: torch.Tensor, # [batch, nheads_q, seqlen_q, head_dim_qk]（bhsd） + k: torch.Tensor, + v: torch.Tensor, + alibi_slopes: torch.Tensor | None, + bias: torch.Tensor | None, + sm_scale: float, + dropout_p: float, + cu_seqlens_q: int, + cu_seqlens_k: int, + max_seqlens_q: int, + max_seqlens_k: int, + causal: bool, + return_scores: bool, + use_exp2: bool, + layout: str, # 仅支持 "bhsd"，见 §4.3 +) -> Tuple[Tensor, Tensor, Tensor]: # (o, softmax_lse, exp_scores) +``` + +与 `triton_attention_mxfp4` **完全同形**，dispatch 可直接替换。无格式开关、无 `use_dot_scaled`。反向经 `autograd.Function` 自动触发，用户不直接调 backward kernel。 + +`dispatch/attention.py` 的接入分支： + +```python +elif isinstance(config, TrainingOpConfig) and config.precision == "mxfp8_e4m3": + self.attn_func = triton_attention_mxfp8 +``` + +### 4.2 scale 布局 + +- **Q / K scale**：沿 head_dim（QK 的 reduction 维）量化。K 的 scale 在非 reduction 维（seqlen_k）上 major。 +- **V scale**：PV 的 reduction 维是 seqlen_k，V scale 沿 seqlen_k。 +- **P scale**：kernel 内动态生成，沿 `BLOCK_N`。 +- forward 把 Q/K/V 存成紧凑 2D-block scale `[.., seqlen/32, head_dim/32]`，backward 复用（§5.3）。 + +### 4.3 三类硬约束（入口断言） + +| 断言 | 位置 | 理由 | +|---|---|---| +| `layout == "bhsd"` | forward op + backward op | `convert_to_mxfp8(is_2d_block=True)` 按数据张量 `shape[-2] × shape[-1]` 分 32×32 块，只有 bhsd 时 `shape[-2]` 才是 seqlen。`bshd`/`thd` 下它按 **nheads** 分组（nheads 恰为 32 倍数时 `torch._check` 也拦不住），而 kernel 的 scale 指针数学假定按 seqlen 分组 → **静默读错 scale** | +| `not causal or seqlen_q == seqlen_k` | forward op + backward op | kernel 掩码用 bottom-right，PyTorch 参照与 `F.sdpa` 用 top-left，两者只在方阵下等价。V1 生产只跑 self-attention，故不统一约定（那是解决不存在的问题），改为断言拒绝。两侧都拦——一个 op 允许、另一个禁止同一种形状，调用者只能靠踩坑才知道边界 | +| backward 不支持 alibi / dropout | `autograd.Function.backward` | 上游 `blockwise_fp8` forward 支持 dropout 但 backward 无对应逻辑且不检查（开 dropout 训练即静默错梯度），alibi 同理。此处改为当场报错 | + +维度约束：head_dim 与 seqlen_k 必须被 `QUANT_BLOCK_SIZE=32` 整除（由 `convert_to_mxfp8` 内的 `torch._check` 保证，不在下游重复检查）；head_dim < 64 因 `tl.dot_scaled` 限制不支持（mxfp4 亦然）。 + +断言的有效性单独验过：`causal=True, sq=64/sk=128` 走 backward 按预期抛 `AssertionError`；`causal=False` 同形状正常通过（非 causal 无对齐问题，不该拦）；`layout="bshd"` 走 forward 按预期抛 `AssertionError`。forward 侧配 `test_causal_forward_rejects_non_square_shape` 做回归。 + +## 5. 量化方案 + +### 5.1 Forward 的两个 dot + +| dot | 计算 | reduction 维 | LHS quant axis | RHS quant axis | 跨几个 scale group | +|---|---|---|---|---|---| +| QK | `S = Q @ Kᵀ` | head_dim | Q: head_dim | K: head_dim | head_dim/32（=4 @ dim128） | +| PV | `O += P @ V` | seqlen_k (BLOCK_N) | P: BLOCK_N | V: seqlen_k | BLOCK_N/32（=2 @ BLOCK_N=64） | + +> 立项时把「QK 一次 `dot_scaled` 跨 head_dim/32 个 32-wide scale group」列为头号数值风险，退路是沿 head_dim 分块累加。**真机实测未触发**：kernel vs 参照 53~63 dB，误差全部来自量化本身而非跨 group 累加，退路未启用。 + +### 5.2 Backward 的七个 dot + +| dot | kernel | 计算 | reduction 轴 | LHS 量化轴 | RHS 量化轴 | +|---|---|---|---|---|---| +| a | dkdv | `qk = q@kᵀ` | head_dim | q 沿 head_dim | k 沿 head_dim | +| b | dkdv | `dp = do@vᵀ` | head_dim_v | do 沿 head_dim_v | v 沿 head_dim_v | +| c | dkdv | `dv += pᵀ@do` | seqlen_q (BLOCK_M) | p 沿 seqlen_q | do 沿 seqlen_q | +| d | dkdv | `dk += dsᵀ@q` | seqlen_q (BLOCK_M) | ds 沿 seqlen_q | q 沿 seqlen_q | +| e | dq | `qk = q@kᵀ` | head_dim | 同 a | 同 a | +| f | dq | `dp = do@vᵀ` | head_dim_v | 同 b | 同 b | +| g | dq | `dq += ds@k` | seqlen_k (BLOCK_N) | ds 沿 seqlen_k | k 沿 seqlen_k | + +同一个 operand 被不同 dot 沿不同轴归约，因此需要多套 scale：q 沿 head_dim（a/e）+ 沿 seqlen_q（d）；k 沿 head_dim（a/e）+ 沿 seqlen_k（g）；do 沿 head_dim_v（b/f）+ 沿 seqlen_q（c）；v 只需 head_dim_v。这是整个 backward 的核心工作量。 + +标准 FA v2 backward 数学（`delta = rowsum(dO ∘ O)`；`dV += Pᵀ@dO`；`dP = dO@Vᵀ`；`dS = P∘(dP−delta)·sm_scale`；`dK += dSᵀ@Q`；`dQ += dS@K`）与上游一致，不赘述。 + +### 5.3 scale 复用方案（A1） + +**q/k/v 直接复用 forward 存的 e4m3 值 + 2D-block scale，backward 不重量化。** 每个 dot 用指针索引把紧凑的 2D scale 广播成 `tl.dot_scaled` 要的 `[outer, reduction/32]`： + +- `_load_scale_hd`（head_dim 收缩，dot a/b/e/f）：`scale[outer//32, dgroup]`，与 forward QK 的 `qs_ptrs` 广播同款。 +- `_load_scale_sq`（seqlen 收缩，dot d/g）：**转置对称复用**——同一个 32×32 块只有一个 scale，两轴共用，换个轴索引成 `[head_dim, seqlen_block/32]`。 + +**dO / P / dS 是 backward 新产生的，forward 没存过**，仍现场 1D per-row 沿各自 reduction 轴量化（`_mx_quant`，`IS_2D_BLOCK=False`），scale 直接匹配 `dot_scaled`。 + +这套方案 correct-by-construction：与 forward 逐位一致、零重量化、无双重量化。演进过程见 §9.1。 + +## 6. 测试组织 + +三层校验，归因清晰。文件：`tests/unittest/mxfp8/test_mxfp8_attention.py`（基线）、`test_mxfp8_attention_reference.py`（分层）、`utils.py`（黄金参照）。 + +| 层 | 对比 | 隔离出什么 | 硬件 | +|---|---|---|---| +| 1 | 纯 PyTorch 黄金参照 vs bf16 SDPA | 算法与量化放置是否正确（不含 Triton） | CPU / 任意设备 | +| 2 | kernel vs 黄金参照 | Triton 移植 bug（masking / online-softmax / LSE / strides / scale 广播），排除量化误差本身 | CDNA4 | +| 3 | kernel vs bf16 SDPA | mxfp8 量化的端到端总误差 | CDNA4 | + +黄金参照（`utils.py`）三个：`mxfp8_attention_forward_reference`、`mxfp8_attention_backward_reference`（stage-1，FA 数学基准）、`mxfp8_attention_backward_reference_stage2`（A1 全量化，复刻 §5.3）。全部纯 fp32 matmul、不含 `tl.dot_scaled`。 + +**参数网格**：`test_cases` 八组（causal × GQA × head_dim{128,192} × seqlen{1024,2048}，含 GQA 与非对称 head_dim 的交叉点）× `causal ∈ {True, False}`；`non_causal_cases` 两组覆盖 `seqlen_q != seqlen_k`（causal 断言使其只在非 causal 下合法）；`reference_cases` 三组小形状供 CPU 层。共 **75 项**。 + +> 测试分层曾经是反的：三个纯 PyTorch 参照测试跑 `causal=[True, False]`，而每一个真正启动 kernel 的测试都钉在 `[True]`——不可能有移植 bug 的那层测两个值，会有移植 bug 的那层只测一个。骨架抄自 mxfp4 时把严格性一起抄丢了。已全部改为两个值。 + +## 7. 验收标准与实测结果 + +| # | 标准 | 结果 | +|---|---|---| +| 1 | forward kernel 在 CDNA4 跑通 | ✅ | +| 2 | backward（dQ/dK/dV）在 CDNA4 跑通 | ✅ | +| 3 | forward vs bf16 SDPA：cossim > 0.99、SNR 硬断言 | ✅ | +| 4 | backward dQ/dK/dV vs bf16 SDPA autograd：硬断言 | ✅ | +| 5 | 覆盖 causal × GQA × head_dim{128,192} × seqlen{1024,2048} | ✅ | +| 6 | dispatch 层 `mxfp8_e4m3` 分支可路由 | ✅ | + +实测数值（gfx950）： + +| 检查 | 结果 | +|---|---| +| backward kernel vs A1 参照 | dQ/dK **53~60 dB**、dV **59~63 dB**，cossim ≈ 1.0 | +| forward kernel vs 黄金参照 | ✅ 全网格通过 | +| 公开 autograd 端到端 vs bf16 SDPA | dQ/dK ≈ 23.3 dB·cossim ≈ 0.9976、dV ≈ 26 dB·cossim ≈ 0.9988 | +| 全量测试 | **75 passed**（约 11 秒） | + +kernel vs 参照 53~63 dB ⇒ 移植与 scale 广播无 bug；端到端 23 dB 与参照预测一致 ⇒ **误差全部来自 mxfp8 量化本身，不是 kernel 实现**。 + +> **相对 mxfp4 attention 的主动加严**：mxfp4 的 `test_mxfp_attention.py` 只 print、无 assert 且无 backward 测试。mxfp8 把前向精度做成硬断言，并新增完整 backward 梯度校验。 + +## 8. 不做的事（明确划线） + +- ❌ CDNA3 / MI300 支持与 fallback 路径 +- ❌ 混合格式（e5m2）——见 §1 第 2 条，切换需改 kernel +- ❌ head_dim < 64（`tl.dot_scaled` 限制） +- ❌ 2D-block P 量化（P 沿 `BLOCK_N` 一维即可） +- ❌ autotune 扩展（`BLOCK_M = BLOCK_N = 64, PRE_LOAD_V=False` 单 config） +- ❌ TMA / async copy / pipelining 调优 +- ❌ 沿 head_dim 分块 QK（§5.1 风险未触发） +- ❌ 非方阵 causal（§4.3 断言拒绝，放开条件见 §9.2） +- ❌ FSDP/TP 集成测试 + +--- + +## 9. 关键决策记录 + +### 9.1 backward 量化源：option A → A1 + +曾实现过 **option A**：backward 入口 `convert_from_mxfp8` 把 saved e4m3 q/k/v 反量化回 bf16，kernel 内每个 dot 沿其 reduction 轴 1D 重量化。它能跑，参照实测 SNR ≈ 23 dB。 + +2026-07-16 翻案改 **A1**（复用 forward 的 e4m3 + 2D scale，§5.3），并**把 A 的代码与参照全部删除**，不留作备用。理由：A 有双重量化（dequant 后再量化），A1 与 forward 逐位一致、correct-by-construction；留 A 当回退是冗余代码。代价是 kernel 内多两个 scale 广播辅助（`_load_scale_hd` / `_load_scale_sq`）。 + +A1 落地时采用「先 kernel 后参照」的顺序，意味着 kernel 写完后有一段时间没有任何数值信号，`_load_scale_sq`（seqlen 轴 re-index）全仓无先例、是当时最大盲区。CDNA4 验收 53~63 dB 后该盲区解除。 + +### 9.2 backward causal 循环边界：判定不修 + +两个 backward kernel 的**掩码**按 bottom-right 对齐（`col_offset = N_CTX_Q - N_CTX_K`），但决定**循环范围**的两处把 `col_offset` 丢了，等价于硬编码 `col_offset == 0`： + +| 位置 | 现状 | 非方阵下的后果 | +|---|---|---| +| `_bwd_kernel_dkdv` 的 `lo` | `(start_n*BLOCK_N - BLOCK_M + 1) // BLOCK_M * BLOCK_M` | `seqlen_k − seqlen_q > BLOCK_M` 时跳过必须计算的 query 块 → dK/dV 漏贡献 | +| `_bwd_kernel_dq` 的 `hi` | `BLOCK_M // BLOCK_N * (start_m+1) * BLOCK_N` | `seqlen_k > seqlen_q` 时差 1 就漏 key 列 → dQ 漏贡献 | + +一度改成带 `col_offset` 的正确形式，**最终回退**。理由按硬度排： + +1. **这条路不可达，所以它不是 bug。** §4.3 两条断言合起来把 `col_offset != 0` 完全堵死：`causal → seqlen_q == seqlen_k` 拦住定长，`layout == "bhsd"` 顺带杀掉 varlen（thd）——这点关键，varlen 下每条序列实际长度不同，光断言 max_seqlen 拦不住。而 `col_offset == 0` 时新旧公式只差一块全掩块，输出必然逐位一致（实测撤回改动前后 21 个 SNR 数字一致到小数点后四位：`p = exp(-inf) = 0`，量化成 e4m3 仍是 0，加进 fp32 累加器是精确的 0）。 +2. **它从来没有独立可达过。** 非方阵 + causal 本就因 kernel（bottom-right）与参照/`F.sdpa`（top-left）约定不一致而算不对，边界错只是被一个更大的破坏盖住。修边界不能让这条路可用，必须先统一约定。 +3. **「负 `lo` 越界读内存」这个理由是假的，已实测推翻。** Triton 的整数 `//` 对负数**向零截断**（不是 Python 的下取整），实测 `BLOCK_M=BLOCK_N=64` 时原式在 `start_n=0..3` 给出 `[0, 0, 64, 128]`，floor 语义才会给 `[-64, 0, 64, 128]`。原式不会产生负 `lo`、不会越界；`tl.maximum(..., 0)` 是新公式自己引入的需求。 +4. **唯一真实代价是性能。** 方阵下每个 key block 多跑一个全掩 query block（Triton 循环内无法提前退出，那块的 4 次 dot 实打实算完再被掩成 0）。多出比例 `2(N−1) / (N(N+1))`，N = seqlen/64：seqlen 1024 → dkdv 内层迭代 +11.0%；2048 → +5.9%；4096 → +3.0%。序列越长越不值钱。 + +结论：这是**性能与可读性**改动，不是正确性修复；按 fix 提交会误导 reviewer 去找线上不存在的 bug。 + +同两行还隐含另一个约束：**`lo` / `hi` 要求 `BLOCK_M == BLOCK_N`**。若把 `BLOCK_N` 调成 128，`BLOCK_M // BLOCK_N` 整除成 0，`hi` 直接为 0 ⇒ **dQ 全 0 且不报错**。当前两者硬编码 64，无断言保护。调 BLOCK 尺寸前必须先处理这里。 + +**若未来要支持非方阵 causal**（cross-attention / prefix），三件事顺序不能颠倒：① 先在 kernel 与参照之间统一 top-left 或 bottom-right 约定；② 再把 `lo`/`hi` 补上 `col_offset`（正确形式见本节，届时才是真修复）；③ 最后才放开 §4.3 的断言。 + +上游 `blockwise_fp8` 有一字不差的两行（1297、1665 行），且**没有** `seqlen_q == seqlen_k` 断言，但同样不可达：唯一生产入口 `blockwise_fa.py` 第 219 行 assert 后把同一个 `seqlen` 传给 q 和 k。同为潜伏、未暴露。 + +### 9.3 port 时丢失的 `num_stages=1` + +上游把这个启动参数藏在一个**只有一个空 config 的 `@triton.autotune`** 里（`triton.Config({}, num_stages=1, num_warps=4)`），不调任何 BLOCK 尺寸，唯一作用就是钉住 `num_stages=1`。port 时它被当作「可选的调优基础设施」删掉，于是静默继承 Triton 默认值（3.6.0 实测 `num_warps=4, num_stages=2`——`num_warps` 恰好撞对，`num_stages` 翻倍）。 + +`num_stages=2` 给内层循环加两级软流水，多出的预取缓冲挤占寄存器；head_dim=192（padded 256）时 fp32 累加器与 scale tile 最多，退化最严重： + +| 形状（batch 4, causal, gfx950） | num_stages=2 | num_stages=1 | 提速 | +|---|---|---|---| +| s1024, hq32/hkv8, d128 | 1.741 ms | 1.505 ms | **1.16x** | +| s2048, hq32/hkv8, d128 | 6.201 ms | 5.411 ms | **1.15x** | +| s2048, hq16/hkv16, d192 | 7.015 ms | 4.630 ms | **1.51x** | + +（`triton.testing.do_bench`，warmup 50 / rep 200，跑两遍复现，偏差 < 0.3%。另试 `num_warps=8`：2.084 / 7.538 / 9.091 ms，明显更差 ⇒ 上游的 4 是对的。） + +修法是给两个 launch 各加显式 `num_warps=4, num_stages=1`，**不照搬那个空 autotune**——参数只有一个取值就不该披着 autotune 的皮，那正是它被丢掉的原因。 + +**同类缺陷全仓排查无第二例**：扫了 `alto/kernels` 全部 20 个含 `@triton.jit` 的文件，循环密集型（attention / grouped GEMM）都有显式 launch 配置，其余是单趟访存受限的量化/elementwise kernel，`num_stages` 对它们无意义。 + +### 9.4 修掉的真 bug：dkdv 组循环用错 dO 的 head stride + +`_bwd_kernel_dkdv` 的 GQA 组循环里，推进到下一个 query head 时把 dO 的指针按 **q 的** head stride 前进： + +```python +q_offset += stride_qh +do_offset += stride_qh # ← 应为 stride_doh +``` + +`stride_doh` 本就传进了 kernel、初始化 `do_offset` 时也用对了，只有这一步用错。q 的最后一维是 `head_dim_qk`、dO 的是 `head_dim_v`，bhsd 连续布局下两者**仅在 `head_dim_qk == head_dim_v` 时相等**。 + +触发要两个条件**同时**成立：`GROUP_SIZE > 1`（GQA）**且** `head_dim_qk != head_dim_v`。原 `test_cases` 七组恰好把这两个轴**各自覆盖但从未交叉**，所以当时 41 passed 完全盖不住。实测（`head_dim_qk=192, head_dim_v=128`）修前 dK/dV 全 nan（dQ 正常——`_bwd_kernel_dq` 每个 query head 一个 program，没有这个组循环），GROUP_SIZE 与 seqlen 再大些直接 HIP 内存访问错误、进程 abort；修后 57~63 dB。 + +`test_cases` 已补上交叉配置（`num_head_q=32, num_head_kv=8, head_dim_qk=192, head_dim_v=128`），并验证过撤回修复后它确实会崩。修复对原有配置数值零影响。 + +**上游 `blockwise_fp8` 第 1378 行是一字不差的同一行**，其测试的 10 组用例有完全相同的覆盖盲区 ⇒ 那份代码在 GQA + 非对称 head_dim 下同样会崩，且无断言拦截。 + +### 9.5 相对上游 blockwise_fp8 更好的四处（有意保持） + +- **`tl.dot_scaled` 消掉了一整套 scale 记账。** 上游是「普通 `tl.dot` + 事后乘标量 descale」，于是 `p_scale` / `log_p_scale` / `acc_descale` 要在 forward inner、dkdv、dq 全程穿（softmax 后 `p *= p_scale`、epilogue `acc *= acc_descale`、backward 重算 p 补 `+ log_p_scale * RCP_LN2`、最后 `dq *= sm_scale / p_scale` 除回去）。换成硬件指令后这些全部消失——不是简化，是特殊情况不存在了。 +- **砍掉上游传了但没读的 kernel 参数**（给 dkdv 传 `Out`/`DO`/`DQ`、给 dq 传 `Out`/`DK`/`DV`）。 +- **把上游的静默失效变成当场报错**（§4.3 第三条）。 +- **无 fallback**：上游留 `use_fp8` 开关让同一套 kernel 兼跑 bf16，代价是每个 kernel 里 `if USE_FP8` 分叉。 + +### 9.6 已排除的虚警（记录以免重复排查） + +| 初查疑点 | 排除依据 | +|---|---| +| 上游 dkdv 的 exp2 分支对 `l_i` 双乘 `RCP_LN2` | 误读。`exp2(qk - l_i[:,None] + log_p_scale * RCP_LN2)` 里 `* RCP_LN2` 作用在 `log_p_scale` 上，`l_i` 只乘一次。两边都对 | +| 上游 autograd backward 返回 18 个梯度但 forward 只有 16 个输入 | 实测最小 `autograd.Function`：PyTorch 容忍尾部多余的 `None`。mxfp8 这边是 15 对 15 精确匹配 | +| mxfp8 未 assert `o` / `softmax_lse` 连续（上游有） | `get_strides_from_layout` 与 `softmax_lse.stride()` 均直读真实 stride，非连续也读得对；`delta = empty_like(softmax_lse)` 继承同布局 | +| mxfp8 缺 head_dim 整除检查（上游 assert `>=32` 且 `%2==0`） | mxfp8 需要更强的 `%32==0`，已由 `convert_to_mxfp8` 的 `torch._check` 保证。上游保证的不在下游重复检查。这也解释了为何 layout 那条必须自己加：`convert_to_mxfp8` 检的是 `shape[-2] % 32`，bshd 下那是 nheads，32 头刚好整除，拦不住 | +| mxfp4 测试里 `causal=[True, False]` 被注释掉，疑似 `causal=False` 有问题 | 那句注释属于整个 `test_attention_fp8_with_sparse_do` 函数被注掉，不是有人发现 `causal=False` 挂掉才关的。mxfp8 放开两个值后实测全过 | + +### 9.7 两处 Triton 语言坑（修复方式记录） + +- **fp8 `tl.load` 的 `other` 类型**：`other=0`（int32）无法 cast 为 `fp8e4nv`，报 `cannot cast int32[...] to fp8e4nv`。改为 `other=0.0`。 +- **模块级 constexpr 必须用实例化写法**：`X: tl.constexpr = 8` 这种注解式写法在 `@jit` 内不可见，报 `NameError: Cannot access global variable`。改为 `X = tl.constexpr(8)`。forward 的 `E4M3_*` 与 backward 的 `RCP_LN2` 都踩过。 + diff --git a/alto/kernels/mxfp8/__init__.py b/alto/kernels/mxfp8/__init__.py index 85eac52..f8c68c7 100644 --- a/alto/kernels/mxfp8/__init__.py +++ b/alto/kernels/mxfp8/__init__.py @@ -1,3 +1,7 @@ # Copyright (c) 2026 Advanced Micro Devices, Inc. # # SPDX-License-Identifier: MIT + +from .triton_flash_attention_mxfp8 import triton_attention_mxfp8 + +__all__ = ("triton_attention_mxfp8", ) diff --git a/alto/kernels/mxfp8/triton_flash_attention_mxfp8.py b/alto/kernels/mxfp8/triton_flash_attention_mxfp8.py new file mode 100644 index 0000000..33828f1 --- /dev/null +++ b/alto/kernels/mxfp8/triton_flash_attention_mxfp8.py @@ -0,0 +1,2248 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +""" +Fused Attention (MXFP8 / e4m3) +============================== + +Triton implementation of the Flash Attention v2 algorithm from Tri Dao +(https://tridao.me/publications/flash2/flash2.pdf) with MXFP8 (e4m3) block +scaled operands. + +This is a mechanical port of ``triton_flash_attention_mxfp4.py``: + * FP4 packs two elements per byte along head_dim; FP8 stores one element per + byte, so all ``// 2`` head-dim halving is removed. + * ``tl.dot_scaled`` dtype string ``"e2m1"`` -> ``"e4m3"``; the ``lhs_k_pack`` + / ``rhs_k_pack`` arguments are dropped (they only exist for packed FP4). + * The on-the-fly softmax-probability quantization uses the MXFP8 helpers + ``_calculate_scales`` / ``_quantize_fp8`` instead of the FP4 pack path. + +The scale layout is identical to the FP4 kernel (uint8, ``[.., dim/32]``, 2D +block), so all scale pointer arithmetic carries over unchanged. + +Forward is a mechanical port of the FP4 kernel; backward is implemented from +the blockwise-fp8 backward structure with per-reduction-axis MXFP8 scales +(stage-2, ``tl.dot_scaled``; see the plan §9.5 and the backward section below). +""" +from typing import Tuple, Optional +import os +import torch +import triton +import triton.language as tl +from torch._library import triton_op, wrap_triton + +from .mxfp8_quantization import ( + BLOCK_SIZE_DEFAULT, + is_cdna4, + _calculate_scales, + _quantize_fp8, +) + +fwd_torch_dtype: tl.constexpr = torch.bfloat16 +bwd_torch_dtype: tl.constexpr = torch.float32 +# Seed the RNG so we get reproducible results for testing. +philox_seed: tl.constexpr = 0x1BF52 +philox_offset: tl.constexpr = 0x1D4B42 + +# e4m3 quantization constants (see mxfp8_quantization.FORMAT_TO_*). +E4M3_TARGET_MAX_POW2 = tl.constexpr(8) +E4M3_MBITS = tl.constexpr(3) +E4M3_FORMAT_ID = tl.constexpr(0) + +AUTOTUNE = os.environ.get('FLASH_ATTENTION_TRITON_AMD_AUTOTUNE', '0').lower() in ('1', 'true', 'yes') +DEBUG = os.environ.get('FLASH_ATTENTION_TRITON_AMD_DEBUG', '0').lower() in ('1', 'true', 'yes') +PERF = os.environ.get('FLASH_ATTENTION_TRITON_AMD_PERF', '0').lower() in ('1', 'true', 'yes') + + +def get_shape_from_layout(q, k, v, layout, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=None, max_seqlen_k=None): + if layout == 'bhsd': + batch_q, nheads_q, max_seqlen_q, head_size_q = q.shape + batch_k, nheads_k, max_seqlen_k, head_size_k = k.shape + batch_v, nheads_v, max_seqlen_v, head_size_v = v.shape + elif layout == 'bshd': + batch_q, max_seqlen_q, nheads_q, head_size_q = q.shape + batch_k, max_seqlen_k, nheads_k, head_size_k = k.shape + batch_v, max_seqlen_v, nheads_v, head_size_v = v.shape + elif layout == 'thd': + batch_q, max_seqlen_q, nheads_q, head_size_q = len(cu_seqlens_q) - 1, max_seqlen_q, q.shape[1], q.shape[2] + batch_k, max_seqlen_k, nheads_k, head_size_k = len(cu_seqlens_k) - 1, max_seqlen_k, k.shape[1], k.shape[2] + batch_v, max_seqlen_v, nheads_v, head_size_v = len(cu_seqlens_k) - 1, max_seqlen_k, v.shape[1], v.shape[2] + else: + assert False, "Got unsupported layout." + + # FP8 stores one element per byte, so head_dim is not packed. + + # assert + assert batch_q == batch_k + assert head_size_q == head_size_k + + return batch_q, nheads_q, nheads_k, head_size_q, head_size_v, max_seqlen_q, max_seqlen_k + + +def get_strides_from_layout(q, layout): + if layout == 'thd': + q_strides = (0, q.stride(1), q.stride(0), q.stride(2)) + elif layout == 'bhsd': + q_strides = (q.stride(0), q.stride(1), q.stride(2), q.stride(3)) + elif layout == 'bshd': + q_strides = (q.stride(0), q.stride(2), q.stride(1), q.stride(3)) + else: + assert False, 'Got unsupported layout.' + return q_strides + + +def get_padded_headsize(size): + # Get closest power of 2 over or equal to 32. + padded_d_model = 1 << (size - 1).bit_length() + # Smallest head_dim supported is 16. If smaller, the tile in the + # kernel is padded - there is no padding in memory for any dims. + padded_d_model = max(padded_d_model, 16) + return padded_d_model + + +@triton.jit +def cdiv_fn(x, y): + return (x + y - 1) // y + + +@triton.jit +def dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, stride): + ms = tl.arange(0, m) + ns = tl.arange(0, n) + return philox_offset + ms[:, None] * stride + ns[None, :] + + +@triton.jit +def dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride): + rng_offsets = dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, stride).to(tl.uint32) + # TODO: use tl.randint for better performance + return tl.rand(philox_seed, rng_offsets) + + +@triton.jit +def dropout_mask(philox_seed, philox_offset, dropout_p, m, n, stride): + rng_output = dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride) + rng_keep = rng_output > dropout_p + return rng_keep + + +# Convenience function to load with optional boundary checks. +# "First" is the major dim, "second" is the minor dim. +@triton.jit +def load_fn(ptrs, offset_first, offset_second, boundary_first, boundary_second, other=0.0): + if offset_first is not None and offset_second is not None: + mask = (offset_first[:, None] < boundary_first) & \ + (offset_second[None, :] < boundary_second) + tensor = tl.load(ptrs, mask=mask, other=other) + elif offset_first is not None: + mask = offset_first[:, None] < boundary_first + tensor = tl.load(ptrs, mask=mask, other=other) + elif offset_second is not None: + mask = offset_second[None, :] < boundary_second + tensor = tl.load(ptrs, mask=mask, other=other) + else: + tensor = tl.load(ptrs) + return tensor + + +@triton.jit +def compute_alibi_block(alibi_slope, seqlen_q, seqlen_k, offs_m, offs_n, transpose=False): + # when seqlen_k and seqlen_q are different we want the diagonal to stick to the bottom right of the attention matrix + # for casual mask we want something like this where (1 is kept and 0 is masked) + # seqlen_q = 2 and seqlen_k = 5 + # 1 1 1 1 0 + # 1 1 1 1 1 + # seqlen_q = 5 and seqlen_k = 2 + # 0 0 + # 0 0 + # 0 0 + # 1 0 + # 1 1 + # for alibi the diagonal is 0 indicating no penalty for attending to that spot and increasing penalty for attending further from the diagonal + relative_pos_block = offs_m[:, None] + seqlen_k - seqlen_q - offs_n[None, :] + alibi_block = -1 * alibi_slope * tl.abs(relative_pos_block) + if transpose: + return alibi_block.T + else: + return alibi_block + + +@triton.jit +def _attn_fwd_inner( + acc, + l_i, + m_i, + q, + qs, + k_ptrs, + v_ptrs, + ks_ptrs, + vs_ptrs, + bias_ptrs, + stride_kn, + stride_vk, + stride_bn, + stride_kscale_n, + stride_vscale_k, + start_m, + actual_seqlen_k, + actual_seqlen_k_scale, + actual_seqlen_q, + dropout_p, + philox_seed, + batch_philox_offset, + exp_scores_ptrs, + block_min, + block_max, + offs_n_causal, + masked_blocks, + n_extra_tokens, + alibi_slope, + score_ptrs, + scores_scaled_shifted_ptrs, + IS_CAUSAL: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL_QK: tl.constexpr, + SCALE_BLOCK_DMODEL_QK: tl.constexpr, + BLOCK_DMODEL_V: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_N_SCALE: tl.constexpr, + OFFS_M: tl.constexpr, + OFFS_N: tl.constexpr, + PRE_LOAD_V: tl.constexpr, + MASK_STEPS: tl.constexpr, + ENABLE_DROPOUT: tl.constexpr, + PADDED_HEAD_QK: tl.constexpr, + PADDED_HEAD_V: tl.constexpr, + ACTUAL_BLOCK_DMODEL_QK: tl.constexpr, + SCALE_ACTUAL_BLOCK_DMODEL_QK: tl.constexpr, + ACTUAL_BLOCK_DMODEL_V: tl.constexpr, + SM_SCALE: tl.constexpr, + USE_EXP2: tl.constexpr, + RETURN_SCORES: tl.constexpr, + QUANT_BLOCK_SIZE: tl.constexpr, + USE_ASM: tl.constexpr, +): + if USE_EXP2: + RCP_LN2: tl.constexpr = 1.4426950408889634 + + # loop over k, v, and update accumulator + for start_n in range(block_min, block_max, BLOCK_N): + # For padded blocks, we will overrun the tensor size if + # we load all BLOCK_N. For others, the blocks are all within range. + if MASK_STEPS: + k_offs_n = start_n + tl.arange(0, BLOCK_N) + vs_offs_n = start_n // QUANT_BLOCK_SIZE + tl.arange(0, BLOCK_N_SCALE) + else: + k_offs_n = None + vs_offs_n = None + if PADDED_HEAD_QK: + k_offs_k = tl.arange(0, BLOCK_DMODEL_QK) + ks_offs_k = tl.arange(0, SCALE_BLOCK_DMODEL_QK) + else: + k_offs_k = None + ks_offs_k = None + k = load_fn(k_ptrs, k_offs_k, k_offs_n, ACTUAL_BLOCK_DMODEL_QK, actual_seqlen_k) + ks = load_fn(ks_ptrs, k_offs_n, ks_offs_k, actual_seqlen_k, SCALE_ACTUAL_BLOCK_DMODEL_QK, other=1) + + if PADDED_HEAD_V: + v_offs_k = tl.arange(0, BLOCK_DMODEL_V) + vs_offs_k = tl.arange(0, BLOCK_DMODEL_V) + else: + v_offs_k = None + vs_offs_k = None + if PRE_LOAD_V: + # We can use the same offsets as k, just with dims transposed. + v = load_fn(v_ptrs, k_offs_n, v_offs_k, actual_seqlen_k, ACTUAL_BLOCK_DMODEL_V) + vs = load_fn(vs_ptrs, vs_offs_k, vs_offs_n, ACTUAL_BLOCK_DMODEL_V, actual_seqlen_k_scale, other=1) + qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) + # We start from end of seqlen_k so only the first iteration would need + # to be checked for padding if it is not a multiple of block_n + if MASK_STEPS: + # If this is the last block / iteration, we want to + # mask if the sequence length is not a multiple of block size + if (start_n + BLOCK_N == block_max) and (n_extra_tokens != 0): + boundary_m = tl.full([BLOCK_M], actual_seqlen_k, dtype=tl.int32) + size_n = start_n + OFFS_N[None, :] + mask = size_n < boundary_m[:, None] + qk = tl.where(mask, qk, float("-inf")) + + # -- compute qk ---- + qk += tl.dot_scaled(q, qs, "e4m3", k, ks, "e4m3", out_dtype=tl.float32) + + qk_scaled = qk * SM_SCALE + + if RETURN_SCORES: + score_mask = (OFFS_M[:, None] < actual_seqlen_q) & ( + (start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k) + tl.store(score_ptrs, qk_scaled, mask=score_mask) + + if IS_CAUSAL: + causal_boundary = start_n + offs_n_causal + causal_mask = OFFS_M[:, None] >= causal_boundary[None, :] + qk_scaled = tl.where(causal_mask, qk_scaled, float("-inf")) + if bias_ptrs is not None: + bias_offs_n = start_n + tl.arange(0, BLOCK_N) if MASK_STEPS else None + bias = load_fn(bias_ptrs, OFFS_M, bias_offs_n, actual_seqlen_q, actual_seqlen_k) + qk_scaled += bias + + if alibi_slope is not None: + # Compute the global position of each token within the sequence + global_m_positions = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + global_n_positions = start_n + tl.arange(0, BLOCK_N) + alibi_block = compute_alibi_block(alibi_slope, actual_seqlen_q, actual_seqlen_k, global_m_positions, + global_n_positions) + qk_scaled += alibi_block + # get max scores so far + m_ij = tl.maximum(m_i, tl.max(qk_scaled, 1)) + + # scale and subtract max + q_shifted = qk_scaled - m_ij[:, None] + if RETURN_SCORES: + scores_scaled_shifted_mask = (OFFS_M[:, None] < actual_seqlen_q) & ( + (start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k) + tl.store(scores_scaled_shifted_ptrs, q_shifted, mask=scores_scaled_shifted_mask) + + # Compute scaled QK and softmax probabilities + if USE_EXP2: + p = tl.math.exp2(q_shifted * RCP_LN2) + else: + p = tl.math.exp(q_shifted) + + # CAVEAT: Must update l_ij before applying dropout + l_ij = tl.sum(p, 1) + if ENABLE_DROPOUT: + philox_offset = batch_philox_offset + start_m * BLOCK_M * actual_seqlen_k + start_n - BLOCK_N + keep = dropout_mask(philox_seed, philox_offset, dropout_p, BLOCK_M, BLOCK_N, actual_seqlen_k) + if RETURN_SCORES: + exp_score_mask = (OFFS_M[:, None] < actual_seqlen_q) & ( + (start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k) + tl.store(exp_scores_ptrs, tl.where(keep, p, -p), mask=exp_score_mask) + p = tl.where(keep, p, 0.0) + elif RETURN_SCORES: + exp_score_mask = (OFFS_M[:, None] < actual_seqlen_q) & ( + (start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k) + tl.store(exp_scores_ptrs, p, mask=exp_score_mask) + + # -- update output accumulator -- + # alpha is an adjustment factor for acc and li as we loop and find new maxes + m_diff = m_i - m_ij + if USE_EXP2: + alpha = tl.math.exp2(m_diff * RCP_LN2) + else: + alpha = tl.math.exp(m_diff) + acc = acc * alpha[:, None] + if not PRE_LOAD_V: + v = load_fn(v_ptrs, k_offs_n, v_offs_k, actual_seqlen_k, ACTUAL_BLOCK_DMODEL_V) + vs = load_fn(vs_ptrs, vs_offs_k, vs_offs_n, ACTUAL_BLOCK_DMODEL_V, actual_seqlen_k_scale, other=1) + # -- update m_i and l_i + l_i = l_i * alpha + l_ij + # update m_i and l_i + m_i = m_ij + + ps = _calculate_scales( + p, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + QUANT_BLOCK_SIZE=QUANT_BLOCK_SIZE, + target_max_pow2=E4M3_TARGET_MAX_POW2, + mbits=E4M3_MBITS, + IS_2D_BLOCK=False, + ) + p_fp8 = _quantize_fp8( + p, + ps, + philox_seed, + batch_philox_offset, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + QUANT_BLOCK_SIZE=QUANT_BLOCK_SIZE, + FP8_FORMAT=E4M3_FORMAT_ID, + IS_2D_BLOCK=False, + USE_ASM=USE_ASM, + USE_SR=False, + ) + + acc += tl.dot_scaled(p_fp8, ps, "e4m3", v, vs, "e4m3", out_dtype=tl.float32) + + k_ptrs += BLOCK_N * stride_kn + v_ptrs += BLOCK_N * stride_vk + ks_ptrs += BLOCK_N_SCALE * stride_kscale_n + vs_ptrs += BLOCK_N_SCALE * stride_vscale_k + if bias_ptrs is not None: + bias_ptrs += BLOCK_N * stride_bn + if RETURN_SCORES: + score_ptrs += BLOCK_N + scores_scaled_shifted_ptrs += BLOCK_N + exp_scores_ptrs += BLOCK_N + return acc, l_i, m_i + + +def get_autotune_fwd_configs(): + return [ + triton.Config( + { + "PRE_LOAD_V": False, + }, + num_stages=1, + num_warps=4, + ), + ], [ + "IS_CAUSAL", "dropout_p", "MAX_SEQLENS_Q", "MAX_SEQLENS_K", "ACTUAL_BLOCK_DMODEL_QK", "ACTUAL_BLOCK_DMODEL_V", + "VARLEN", "HQ", "HK" + ] + + +autotune_fwd_configs, autotune_fwd_keys = get_autotune_fwd_configs() + + +@triton.autotune( + configs=autotune_fwd_configs, + key=autotune_fwd_keys, +) +@triton.jit +def attn_fwd( + Q, + K, + V, + bias, + q_scale_ptr, + k_scale_ptr, + v_scale_ptr, + SM_SCALE: tl.constexpr, + LSE, + Out, + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_vz, + stride_vh, + stride_vk, + stride_vn, + stride_oz, + stride_oh, + stride_om, + stride_on, + stride_bz, + stride_bh, + stride_bm, + stride_bn, + stride_az, + stride_ah, + stride_sz, + stride_sh, + stride_sm, + stride_sn, + stride_lse_z, + stride_lse_h, + stride_lse_m, + stride_qscale_z, + stride_qscale_h, + stride_qscale_m, + stride_qscale_k, + stride_kscale_z, + stride_kscale_h, + stride_kscale_n, + stride_kscale_k, + stride_vscale_z, + stride_vscale_h, + stride_vscale_k, + stride_vscale_n, + cu_seqlens_q, + cu_seqlens_k, + dropout_p, + philox_seed, + philox_offset_base, + scores, + scores_scaled_shifted, + exp_scores, + alibi_slopes, + HQ: tl.constexpr, + HK: tl.constexpr, + ACTUAL_BLOCK_DMODEL_QK: tl.constexpr, + ACTUAL_BLOCK_DMODEL_V: tl.constexpr, + MAX_SEQLENS_Q: tl.constexpr, + MAX_SEQLENS_K: tl.constexpr, + VARLEN: tl.constexpr, + IS_CAUSAL: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL_QK: tl.constexpr, + BLOCK_DMODEL_V: tl.constexpr, + BLOCK_N: tl.constexpr, + PRE_LOAD_V: tl.constexpr, + USE_BIAS: tl.constexpr, + ENABLE_DROPOUT: tl.constexpr, + RETURN_SCORES: tl.constexpr, + USE_ALIBI: tl.constexpr, + USE_EXP2: tl.constexpr, + QUANT_BLOCK_SIZE: tl.constexpr, + USE_ASM: tl.constexpr, +): + start_m = tl.program_id(0) + off_h_q = tl.program_id(1) + off_z = tl.program_id(2) + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + + SCALE_BLOCK_DMODEL_QK: tl.constexpr = BLOCK_DMODEL_QK // QUANT_BLOCK_SIZE + SCALE_ACTUAL_BLOCK_DMODEL_QK: tl.constexpr = ACTUAL_BLOCK_DMODEL_QK // QUANT_BLOCK_SIZE + BLOCK_N_SCALE: tl.constexpr = BLOCK_N // QUANT_BLOCK_SIZE + offs_d_qk = tl.arange(0, BLOCK_DMODEL_QK) + offs_d_qk_scale = tl.arange(0, SCALE_BLOCK_DMODEL_QK) + offs_d_v = tl.arange(0, BLOCK_DMODEL_V) + offs_n_scale = tl.arange(0, BLOCK_N_SCALE) + + # If MQA / GQA, set the K and V head offsets appropriately. + GROUP_SIZE: tl.constexpr = HQ // HK + if GROUP_SIZE != 1: + off_h_k = off_h_q // GROUP_SIZE + else: + off_h_k = off_h_q + + PADDED_HEAD_QK: tl.constexpr = (ACTUAL_BLOCK_DMODEL_QK != BLOCK_DMODEL_QK) + PADDED_HEAD_V: tl.constexpr = (ACTUAL_BLOCK_DMODEL_V != BLOCK_DMODEL_V) + + if VARLEN: + cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z) + cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1) + + seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start + # We have a one-size-fits-all grid in id(0). Some seqlens might be too + # small for all start_m so for those we return early. + if start_m * BLOCK_M > seqlen_q: + return + cu_seqlens_k_start = tl.load(cu_seqlens_k + off_z) + cu_seqlens_k_end = tl.load(cu_seqlens_k + off_z + 1) + seqlen_k = cu_seqlens_k_end - cu_seqlens_k_start + else: + cu_seqlens_q_start = 0 + cu_seqlens_k_start = 0 + seqlen_q = MAX_SEQLENS_Q + seqlen_k = MAX_SEQLENS_K + + cu_seqlens_q_start_scale = cu_seqlens_q_start // QUANT_BLOCK_SIZE + cu_seqlens_k_start_scale = cu_seqlens_k_start // QUANT_BLOCK_SIZE + seqlen_k_scale = seqlen_k // QUANT_BLOCK_SIZE + + # Now we compute whether we need to exit early due to causal masking. + n_blocks = cdiv_fn(seqlen_k, BLOCK_N) + if IS_CAUSAL: + n_blocks_seqlen = cdiv_fn((start_m + 1) * BLOCK_M + seqlen_k - seqlen_q, BLOCK_N) + n_blocks = min(n_blocks, n_blocks_seqlen) + # If we have no blocks after adjusting for seqlen deltas, this WG is part of + # the blocks that are all 0. We exit early. + if n_blocks <= 0: + o_offset = Out + off_z * stride_oz + off_h_q * stride_oh + cu_seqlens_q_start * stride_om + o_ptrs = o_offset + offs_m[:, None] * stride_om + offs_d_v[None, :] * stride_on + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_V], dtype=Out.type.element_ty) + o_ptrs_mask = offs_m[:, None] < seqlen_q + if PADDED_HEAD_V: + o_ptrs_mask = o_ptrs_mask & (offs_d_v[None, :] < ACTUAL_BLOCK_DMODEL_V) + # We still need to write 0s to the result + tl.store(o_ptrs, acc, mask=o_ptrs_mask) + l_offset = LSE + off_z * stride_lse_z + off_h_q * stride_lse_h + cu_seqlens_q_start * stride_lse_m + l_ptrs = l_offset + offs_m * stride_lse_m + + l = tl.full([BLOCK_M], value=0.0, dtype=tl.float32) + + l_ptrs_mask = offs_m < MAX_SEQLENS_Q + tl.store(l_ptrs, l, mask=l_ptrs_mask) + return + + n_extra_tokens = 0 + if seqlen_k < BLOCK_N: + n_extra_tokens = BLOCK_N - seqlen_k + elif seqlen_k % BLOCK_N: + n_extra_tokens = seqlen_k % BLOCK_N + + # Compute pointers for all the tensors used in this kernel. + q_offset = Q + off_z * stride_qz + off_h_q * stride_qh + cu_seqlens_q_start * stride_qm + q_ptrs = q_offset + offs_m[:, None] * stride_qm + offs_d_qk[None, :] * stride_qk + k_offset = K + off_z * stride_kz + off_h_k * stride_kh + cu_seqlens_k_start * stride_kn + k_ptrs = k_offset + offs_d_qk[:, None] * stride_kk + offs_n[None, :] * stride_kn + v_offset = V + off_z * stride_vz + off_h_k * stride_vh + cu_seqlens_k_start * stride_vk + v_ptrs = v_offset + offs_n[:, None] * stride_vk + offs_d_v[None, :] * stride_vn + qs_offset = (q_scale_ptr + off_z * stride_qscale_z + off_h_q * stride_qscale_h + + cu_seqlens_q_start_scale * stride_qscale_m) + qs_ptrs = (qs_offset + (offs_m[:, None] // QUANT_BLOCK_SIZE) * stride_qscale_m + + offs_d_qk_scale[None, :] * stride_qscale_k) + ks_offset = (k_scale_ptr + off_z * stride_kscale_z + off_h_k * stride_kscale_h + + cu_seqlens_k_start_scale * stride_kscale_n) + # k scale is N*K even though k is K*N, this is required by tl.dot_scaled + ks_ptrs = (ks_offset + (offs_n[:, None] // QUANT_BLOCK_SIZE) * stride_kscale_n + + offs_d_qk_scale[None, :] * stride_kscale_k) + vs_offset = (v_scale_ptr + off_z * stride_vscale_z + off_h_k * stride_vscale_h + + cu_seqlens_k_start_scale * stride_vscale_k) + vs_ptrs = (vs_offset + (offs_d_v[:, None] // QUANT_BLOCK_SIZE) * stride_vscale_n + + offs_n_scale[None, :] * stride_vscale_k) + if USE_BIAS: + # Note: this might get large enough to overflow on some configs + bias_offset = off_h_q * stride_bh + bias_ptrs = bias + bias_offset + offs_m[:, None] * stride_bm + offs_n[None, :] * stride_bn + else: + bias_ptrs = None + + if USE_ALIBI: + a_offset = off_z * stride_az + off_h_q * stride_ah + alibi_slope = tl.load(alibi_slopes + a_offset) + else: + alibi_slope = None + + if RETURN_SCORES: + scores_offset = scores + off_z * stride_sz + off_h_q * stride_sh + cu_seqlens_q_start * stride_sm + score_ptrs = scores_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn + + scores_scaled_shifted_offset = scores_scaled_shifted + off_z * stride_sz + off_h_q * stride_sh + cu_seqlens_q_start * stride_sm + scores_scaled_shifted_ptrs = scores_scaled_shifted_offset + offs_m[:, None] * stride_sm + offs_n[ + None, :] * stride_sn + + exp_scores_offset = exp_scores + off_z * stride_sz + off_h_q * stride_sh + cu_seqlens_q_start * stride_sm + exp_scores_ptrs = exp_scores_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn + else: + score_ptrs = None + scores_scaled_shifted_ptrs = None + exp_scores_ptrs = None + + if ENABLE_DROPOUT: + off_hz = off_z * HQ + off_h_q + batch_philox_offset = philox_offset_base + off_hz * seqlen_q * seqlen_k + else: + batch_philox_offset = 0 + # initialize pointer to m and l + m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) + l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_V], dtype=tl.float32) + # Q is loaded once at the beginning and shared by all N blocks. + q_ptrs_mask = offs_m[:, None] < seqlen_q + qs_ptrs_mask = q_ptrs_mask + if PADDED_HEAD_QK: + q_ptrs_mask = q_ptrs_mask & (offs_d_qk[None, :] < ACTUAL_BLOCK_DMODEL_QK) + qs_ptrs_mask = qs_ptrs_mask & (offs_d_qk_scale[None, :] < SCALE_ACTUAL_BLOCK_DMODEL_QK) + + q = tl.load(q_ptrs, mask=q_ptrs_mask, other=0.0) + qs = tl.load(qs_ptrs, mask=qs_ptrs_mask, other=1) + + # Here we compute how many full and masked blocks we have. + padded_block_k = n_extra_tokens != 0 + is_modulo_mn = not padded_block_k and (seqlen_q % BLOCK_M == 0) + if IS_CAUSAL: + # There are always at least BLOCK_M // BLOCK_N masked blocks. + masked_blocks = BLOCK_M // BLOCK_N + (not is_modulo_mn) + else: + # Padding on Q does not need to be masked in the FA loop. + masked_blocks = padded_block_k + + masked_blocks = min(masked_blocks, n_blocks) + n_full_blocks = n_blocks - masked_blocks + block_min = 0 + block_max = n_blocks * BLOCK_N + # Compute for full blocks. Here we set causal to false regardless of its actual + # value because there is no masking. Similarly we do not need padding. + + if n_full_blocks > 0: + block_max = (n_blocks - masked_blocks) * BLOCK_N + acc, l_i, m_i = _attn_fwd_inner( + acc, + l_i, + m_i, + q, + qs, + k_ptrs, + v_ptrs, + ks_ptrs, + vs_ptrs, + bias_ptrs, + stride_kn, + stride_vk, + stride_bn, + stride_kscale_n, + stride_vscale_k, + start_m, + seqlen_k, + seqlen_k_scale, + seqlen_q, + dropout_p, + philox_seed, + batch_philox_offset, + exp_scores_ptrs, + # _, _, offs_n_causal, masked_blocks, n_extra_tokens, _ + block_min, + block_max, + 0, + 0, + 0, + alibi_slope, + score_ptrs, + scores_scaled_shifted_ptrs, + # IS_CAUSAL, .... + False, + BLOCK_M, + BLOCK_DMODEL_QK, + SCALE_BLOCK_DMODEL_QK, + BLOCK_DMODEL_V, + BLOCK_N, + BLOCK_N_SCALE, + offs_m, + offs_n, + # _, MASK_STEPS, ... + PRE_LOAD_V, + False, + ENABLE_DROPOUT, + PADDED_HEAD_QK, + PADDED_HEAD_V, + ACTUAL_BLOCK_DMODEL_QK, + SCALE_ACTUAL_BLOCK_DMODEL_QK, + ACTUAL_BLOCK_DMODEL_V, + SM_SCALE, + USE_EXP2=USE_EXP2, + RETURN_SCORES=RETURN_SCORES, + QUANT_BLOCK_SIZE=QUANT_BLOCK_SIZE, + USE_ASM=USE_ASM, + ) + block_min = block_max + block_max = n_blocks * BLOCK_N + + # Remaining blocks, if any, are full / not masked. + if (masked_blocks > 0): + if IS_CAUSAL: + offs_n_causal = offs_n + (seqlen_q - seqlen_k) + else: + offs_n_causal = 0 + k_ptrs += n_full_blocks * BLOCK_N * stride_kn + v_ptrs += n_full_blocks * BLOCK_N * stride_vk + ks_ptrs += n_full_blocks * BLOCK_N_SCALE * stride_kscale_n + vs_ptrs += n_full_blocks * BLOCK_N_SCALE * stride_vscale_k + if USE_BIAS: + bias_ptrs += n_full_blocks * BLOCK_N * stride_bn + if RETURN_SCORES: + score_ptrs += n_full_blocks * BLOCK_N + scores_scaled_shifted_ptrs += n_full_blocks * BLOCK_N + exp_scores_ptrs += n_full_blocks * BLOCK_N + + acc, l_i, m_i = _attn_fwd_inner( + acc, + l_i, + m_i, + q, + qs, + k_ptrs, + v_ptrs, + ks_ptrs, + vs_ptrs, + bias_ptrs, + stride_kn, + stride_vk, + stride_bn, + stride_kscale_n, + stride_vscale_k, + start_m, + seqlen_k, + seqlen_k_scale, + seqlen_q, + dropout_p, + philox_seed, + batch_philox_offset, + exp_scores_ptrs, + block_min, + block_max, + offs_n_causal, + masked_blocks, + n_extra_tokens, + alibi_slope, + score_ptrs, + scores_scaled_shifted_ptrs, + IS_CAUSAL, + BLOCK_M, + BLOCK_DMODEL_QK, + SCALE_BLOCK_DMODEL_QK, + BLOCK_DMODEL_V, + BLOCK_N, + BLOCK_N_SCALE, + offs_m, + offs_n, + # _, MASK_STEPS, ... + PRE_LOAD_V, + True, + ENABLE_DROPOUT, + PADDED_HEAD_QK, + PADDED_HEAD_V, + ACTUAL_BLOCK_DMODEL_QK, + SCALE_ACTUAL_BLOCK_DMODEL_QK, + ACTUAL_BLOCK_DMODEL_V, + SM_SCALE, + USE_EXP2=USE_EXP2, + RETURN_SCORES=RETURN_SCORES, + QUANT_BLOCK_SIZE=QUANT_BLOCK_SIZE, + USE_ASM=USE_ASM, + ) + + # epilogue + l_recip = 1 / l_i[:, None] + acc = acc * l_recip + if ENABLE_DROPOUT: + acc = acc / (1 - dropout_p) + # If seqlen_q > seqlen_k but the delta is not a multiple of BLOCK_M, + # then we have one block with a row of all NaNs which come from computing + # softmax over a row of all -infs (-inf - inf = NaN). We check for that here + # and store 0s where there are NaNs as these rows should've been zeroed out. + end_m_idx = (start_m + 1) * BLOCK_M + start_m_idx = start_m * BLOCK_M + causal_start_idx = seqlen_q - seqlen_k + + acc = acc.to(Out.type.element_ty) + if IS_CAUSAL: + if causal_start_idx > start_m_idx and causal_start_idx < end_m_idx: + out_mask_boundary = tl.full((BLOCK_DMODEL_V,), causal_start_idx, dtype=tl.int32) + mask_m_offsets = start_m_idx + tl.arange(0, BLOCK_M) + out_ptrs_mask = mask_m_offsets[:, None] >= out_mask_boundary[None, :] + z = 0.0 + acc = tl.where(out_ptrs_mask, acc, z.to(acc.dtype)) + + # write back LSE(Log Sum Exponents), the log of the normalization constant + l_offset = LSE + off_z * stride_lse_z + off_h_q * stride_lse_h + cu_seqlens_q_start * stride_lse_m + offs_l_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + l_ptrs = l_offset + offs_l_m * stride_lse_m + if USE_EXP2: + RCP_LN2: tl.constexpr = 1.4426950408889634 + LN2: tl.constexpr = 0.6931471824645996 + # compute log-sum-exp in base 2 units + mi_base2 = m_i * RCP_LN2 + softmax_lse = mi_base2 + tl.math.log2(l_i) + # convert back to natural units + softmax_lse *= LN2 + else: + softmax_lse = m_i + tl.math.log(l_i) + + if IS_CAUSAL: + # zero out nans caused by -infs when doing causal + lse_mask = (start_m_idx + tl.arange(0, BLOCK_M)) < causal_start_idx + softmax_lse = tl.where(lse_mask, 0.0, softmax_lse) + + # If seqlen_q not multiple of BLOCK_M, we need to mask out the last few rows. + # This is only true for the last M block. For others, overflow_size will be -ve + overflow_size = end_m_idx - seqlen_q + if overflow_size > 0: + boundary = tl.full((BLOCK_M,), BLOCK_M - overflow_size, dtype=tl.int32) + l_ptrs_mask = tl.arange(0, BLOCK_M) < boundary + tl.store(l_ptrs, softmax_lse, mask=l_ptrs_mask) # the log of the normalization constant + else: + tl.store(l_ptrs, softmax_lse) # the log of the normalization constant + + # write back O + o_offset = Out + off_z * stride_oz + off_h_q * stride_oh + cu_seqlens_q_start * stride_om + o_ptrs = o_offset + offs_m[:, None] * stride_om + offs_d_v[None, :] * stride_on + o_ptrs_mask = tl.full([BLOCK_M, BLOCK_DMODEL_V], 1, dtype=tl.int1) + if overflow_size > 0: + o_ptrs_mask = o_ptrs_mask & (offs_m[:, None] < seqlen_q) + if PADDED_HEAD_V: + o_ptrs_mask = o_ptrs_mask & (offs_d_v[None, :] < ACTUAL_BLOCK_DMODEL_V) + tl.store(o_ptrs, acc.to(Out.type.element_ty), mask=o_ptrs_mask) + + +def get_padded_head_dim(head_size: int): + # Get closest power of 2 over or equal to 32. + padded_d_model = 1 << (head_size - 1).bit_length() + # Smallest head_dim supported is 16. If smaller, the tile in the + # kernel is padded - there is no padding in memory for any dims. + padded_d_model = max(padded_d_model, 16) + return padded_d_model + + +@triton_op("alto::attention_mxfp8_forward_triton_impl", mutates_args=()) +def attention_mxfp8_forward_triton_impl( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_scale: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + sm_scale: float, + alibi_slopes: Optional[torch.Tensor], + causal: bool, + bias: Optional[torch.Tensor], + dropout_p: float, + layout: str, + cu_seqlens_q: Optional[int], + cu_seqlens_k: Optional[int], + max_seqlens_q: Optional[int], + max_seqlens_k: Optional[int], + return_scores: bool, + use_exp2: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if DEBUG: + print() + print("attention_mxfp8_forward_triton_impl") + print("q:", q, q.shape) + print("k:", k, k.shape) + print("v:", v, v.shape) + print("sm_scale:", sm_scale) + print("alibi_slopes:", alibi_slopes) + print("causal:", causal) + print("bias:", bias) + print("dropout_p:", dropout_p) + print("layout:", layout) + print("cu_seqlens_q:", cu_seqlens_q) + print("cu_seqlens_k:", cu_seqlens_k) + print("max_seqlens_q:", max_seqlens_q) + print("max_seqlens_k:", max_seqlens_k) + print("return_scores:", return_scores) + print("use_exp2:", use_exp2) + + assert q.is_contiguous() + assert k.is_contiguous() + assert v.is_contiguous() + assert q_scale.is_contiguous() + assert k_scale.is_contiguous() + assert v_scale.is_contiguous() + assert layout == "bhsd", ( + f"MXFP8 attention requires layout='bhsd', got {layout!r}: the 2D-block scale groups the " + "data tensor's last two dims, which are (seqlen, head_dim) only for bhsd. Any other " + "layout groups across heads and the scale pointer math below reads the wrong scale.") + + # check if varlen + is_varlen = layout == "thd" + + # NOTE: a large bias tensor leads to overflow during pointer arithmetic + if (bias is not None): + assert (bias.numel() < 2**31) + + batch, nheads_q, nheads_k, head_size_qk, head_size_v, seqlen_q, seqlen_k = get_shape_from_layout( + q, k, v, layout, cu_seqlens_q, cu_seqlens_k, max_seqlens_q, max_seqlens_k) + assert not causal or seqlen_q == seqlen_k, ( + f"causal forward requires seqlen_q == seqlen_k, got {seqlen_q} vs {seqlen_k}: the " + "kernel uses bottom-right causal masking while PyTorch SDPA/reference tests use top-left " + "masking for non-square shapes.") + o_shape = (*q.shape[:-1], head_size_v) + o = torch.empty( + o_shape, + device=q.device, + dtype=fwd_torch_dtype, + requires_grad=True, + ) + + q_strides = get_strides_from_layout(q, layout) + k_strides = get_strides_from_layout(k, layout) + v_strides = get_strides_from_layout(v, layout) + o_strides = get_strides_from_layout(o, layout) + + # Get closest power of 2 over or equal to 32. + padded_d_model_qk = get_padded_head_dim(head_size_qk) + padded_d_model_v = get_padded_head_dim(head_size_v) + + grid = lambda META: (triton.cdiv(max_seqlens_q, META["BLOCK_M"]), nheads_q, batch) + + if return_scores: + scores = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device, dtype=torch.float32) + scores_scaled_shifted = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), + device=q.device, + dtype=torch.float32) + scores_strides = (scores.stride(0), scores.stride(1), scores.stride(2), scores.stride(3)) + else: + scores = torch.empty([], device=q.device, dtype=torch.float32) + scores_scaled_shifted = None + scores_strides = (0, 0, 0, 0) + + # exp_scores is used to validate dropout behavior vs the PyTorch SDPA math backend reference. + if return_scores: + exp_scores = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device, dtype=torch.float32) + else: + exp_scores = torch.empty([], device=q.device, dtype=torch.float32) + + # stores LSE the log of the normalization constant / sum of exponential score (unnormalized probabilities) + if is_varlen: + softmax_lse = torch.empty((q.shape[0], nheads_q), device=q.device, dtype=torch.float32) + stride_lse_m, stride_lse_h = softmax_lse.stride() + stride_lse_z = 0 + else: + softmax_lse = torch.empty((batch, nheads_q, max_seqlens_q), device=q.device, dtype=torch.float32) + stride_lse_z, stride_lse_h, stride_lse_m = softmax_lse.stride() + + if bias is not None: + bias_strides = (bias.stride(0), bias.stride(1), bias.stride(2), bias.stride(3)) + else: + bias_strides = (0, 0, 0, 0) + + if alibi_slopes is not None: + alibi_strides = (alibi_slopes.stride(0), alibi_slopes.stride(1)) + else: + alibi_strides = (0, 0) + + qs_strides = get_strides_from_layout(q_scale, layout) + ks_strides = get_strides_from_layout(k_scale, layout) + vs_strides = get_strides_from_layout(v_scale, layout) + + wrap_triton(attn_fwd)[grid]( + q, + k, + v, + bias, + q_scale, + k_scale, + v_scale, + sm_scale, + softmax_lse, + o, + *q_strides, + *k_strides, + *v_strides, + *o_strides, + *bias_strides, + *alibi_strides, + *scores_strides, + stride_lse_z, + stride_lse_h, + stride_lse_m, + *qs_strides, + *ks_strides, + *vs_strides, + cu_seqlens_q, + cu_seqlens_k, + dropout_p=dropout_p, + philox_seed=philox_seed, + philox_offset_base=philox_offset, + scores=scores, + scores_scaled_shifted=scores_scaled_shifted, + exp_scores=exp_scores, + alibi_slopes=alibi_slopes, + HQ=nheads_q, + HK=nheads_k, + ACTUAL_BLOCK_DMODEL_QK=head_size_qk, + ACTUAL_BLOCK_DMODEL_V=head_size_v, + MAX_SEQLENS_Q=max_seqlens_q, + MAX_SEQLENS_K=max_seqlens_k, + IS_CAUSAL=causal, + VARLEN=is_varlen, + BLOCK_DMODEL_QK=padded_d_model_qk, + BLOCK_DMODEL_V=padded_d_model_v, + USE_BIAS=False if bias is None else True, + USE_ALIBI=False if alibi_slopes is None else True, + ENABLE_DROPOUT=dropout_p > 0.0, + USE_EXP2=use_exp2, + RETURN_SCORES=return_scores, + BLOCK_M=64, + BLOCK_N=64, + QUANT_BLOCK_SIZE=BLOCK_SIZE_DEFAULT, + USE_ASM=is_cdna4(), + ) + return o, softmax_lse, exp_scores + + +RCP_LN2 = tl.constexpr(1.4426950408889634) + + +@triton.jit +def _mx_quant( + x, + BM: tl.constexpr, + BN: tl.constexpr, + QUANT_BLOCK_SIZE: tl.constexpr, + IS_2D_BLOCK: tl.constexpr, + USE_ASM: tl.constexpr, +): + """Quantize a [BM, BN] tile to e4m3 with MX block scales along the last axis. + + Wraps ``_calculate_scales`` + ``_quantize_fp8`` with the e4m3 constants, so + each backward dot can quantize an operand along its reduction axis (which the + caller places last, transposing the tile when needed). ``IS_2D_BLOCK=True`` + groups both axes (head_dim reduction, like fwd Q/K/V); ``False`` is a 1D block + along the last axis (seqlen reduction, like fwd P). Non-SR path, so philox is + unused. + """ + scales = _calculate_scales( + x, + BLOCK_M=BM, + BLOCK_N=BN, + QUANT_BLOCK_SIZE=QUANT_BLOCK_SIZE, + target_max_pow2=E4M3_TARGET_MAX_POW2, + mbits=E4M3_MBITS, + IS_2D_BLOCK=IS_2D_BLOCK, + ) + xq = _quantize_fp8( + x, + scales, + 0, + 0, + BLOCK_M=BM, + BLOCK_N=BN, + QUANT_BLOCK_SIZE=QUANT_BLOCK_SIZE, + FP8_FORMAT=E4M3_FORMAT_ID, + IS_2D_BLOCK=IS_2D_BLOCK, + USE_ASM=USE_ASM, + USE_SR=False, + ) + return xq, scales + + +# Backward operands split into two families (A1, plan §9.5 / AB decision): +# +# * Q/K/V come from the forward pass as saved e4m3 + a compact 2D-block scale +# ``[.., seqlen/32, head_dim/32]``. A1 **reuses** that exact quantization — +# no dequant, no re-quantization. Each dot rebuilds the scale tile the +# ``tl.dot_scaled`` layout wants (``[outer, reduction/32]``) directly from the +# saved compact scale via pointer index math (the same broadcast the forward +# kernel uses for QK). Because a 32x32 block shares one scale, the same saved +# e4m3 values + tile serve both the head_dim-reduction dots (a/b/e/f) and the +# seqlen-reduction dots (d/g); only the broadcast axis differs. +# * dO/P/dS are backward-only tensors the forward never saw, so they are +# quantized fresh in-kernel with 1D per-row scales (``IS_2D_BLOCK=False``) +# along their reduction axis — the canonical layout ``tl.dot_scaled`` eats +# with no broadcast. +_MX_2D = tl.constexpr(False) + + +@triton.jit +def _load_scale_hd( + scale_base, + offs_outer, + stride_sg, + stride_dg, + N_CTX, + SCALE_D: tl.constexpr, + SCALE_ACTUAL_D: tl.constexpr, + QUANT_BLOCK_SIZE: tl.constexpr, +): + """Rebuild a ``[len(offs_outer), SCALE_D]`` scale tile for a **head_dim**- + reduction dot from the saved compact 2D-block scale. + + ``offs_outer`` indexes the outer (token/seqlen) rows; each maps to compact + seqlen-group ``offs_outer // QUANT_BLOCK_SIZE``. ``SCALE_D = head_dim/32`` + columns index the reduction (head_dim) groups directly. This mirrors the + forward kernel's ``qs_ptrs`` broadcast (a 32-row block shares one scale). + Out-of-range rows (past ``N_CTX``) and padded head groups load a neutral + scale; they only ever multiply masked-to-zero e4m3 values. + """ + offs_dg = tl.arange(0, SCALE_D) + mask = (offs_outer[:, None] < N_CTX) & (offs_dg[None, :] < SCALE_ACTUAL_D) + ptrs = (scale_base + (offs_outer[:, None] // QUANT_BLOCK_SIZE) * stride_sg + + offs_dg[None, :] * stride_dg) + return tl.load(ptrs, mask=mask, other=127) + + +@triton.jit +def _load_scale_sq( + scale_base, + start_outer, + offs_d, + stride_sg, + stride_dg, + N_CTX, + BLOCK_OUTER: tl.constexpr, + SCALE_ACTUAL_D: tl.constexpr, + QUANT_BLOCK_SIZE: tl.constexpr, +): + """Rebuild a ``[len(offs_d), BLOCK_OUTER/32]`` scale tile for a **seqlen**- + reduction dot from the same saved compact 2D-block scale. + + This is the transpose-symmetric reuse (AB decision A1): the reduction axis is + now the seqlen block ``[start_outer, start_outer+BLOCK_OUTER)``, grouped by 32 + into ``BLOCK_OUTER/32`` columns, while ``offs_d`` (head_dim rows) map to + compact head_dim-group ``offs_d // QUANT_BLOCK_SIZE`` and are broadcast across + the 32 rows of each group. Element ``[d, sg] = scale[start_outer/32 + sg, + d/32]`` — the same 32x32 block scale the head_dim path reads, indexed for the + other axis. + """ + n_sg: tl.constexpr = BLOCK_OUTER // QUANT_BLOCK_SIZE + offs_sg = tl.arange(0, n_sg) + seq_group = start_outer // QUANT_BLOCK_SIZE + offs_sg + mask = (offs_d[:, None] < SCALE_ACTUAL_D * QUANT_BLOCK_SIZE) & \ + ((seq_group[None, :] * QUANT_BLOCK_SIZE) < N_CTX) + ptrs = (scale_base + seq_group[None, :] * stride_sg + + (offs_d[:, None] // QUANT_BLOCK_SIZE) * stride_dg) + return tl.load(ptrs, mask=mask, other=127) + + +@triton.jit +def _bwd_preprocess( + Out, + DO, + Delta, + stride_oz, + stride_oh, + stride_om, + stride_ok, + stride_doz, + stride_doh, + stride_dom, + stride_dok, + stride_deltaz, + stride_deltah, + stride_deltam, + cu_seqlens_q, + max_seqlen_q, + BLOCK_M: tl.constexpr, + BLOCK_DMODEL_V: tl.constexpr, + ACTUAL_BLOCK_DMODEL_V: tl.constexpr, + HQ: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """delta = rowsum(o * do), one value per query row. Feeds ds = p * (dp - delta).""" + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + off_z = pid_bh // HQ + off_h = pid_bh % HQ + + if IS_VARLEN: + q_start = tl.load(cu_seqlens_q + off_z) + q_end = tl.load(cu_seqlens_q + off_z + 1) + N_CTX_Q = q_end - q_start + else: + q_start = 0 + N_CTX_Q = max_seqlen_q + + off_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + off_d_v = tl.arange(0, BLOCK_DMODEL_V) + mask_m = off_m < N_CTX_Q + mask_o = mask_m[:, None] & (off_d_v[None, :] < ACTUAL_BLOCK_DMODEL_V) + + o_offset = Out + off_z * stride_oz + off_h * stride_oh + q_start * stride_om + do_offset = DO + off_z * stride_doz + off_h * stride_doh + q_start * stride_dom + out_ptrs = o_offset + off_m[:, None] * stride_om + off_d_v[None, :] * stride_ok + do_ptrs = do_offset + off_m[:, None] * stride_dom + off_d_v[None, :] * stride_dok + + o = tl.load(out_ptrs, mask=mask_o, other=0.0).to(tl.float32) + do = tl.load(do_ptrs, mask=mask_o, other=0.0).to(tl.float32) + delta = tl.sum(o * do, axis=1) + + delta_offset = Delta + off_z * stride_deltaz + off_h * stride_deltah + q_start * stride_deltam + tl.store(delta_offset + off_m * stride_deltam, delta, mask=mask_m) + + +@triton.jit +def _attn_bwd_dkdv_inner( + kt_fp8, + ks_hd, + vt_fp8, + vs_hd, + dk, + dv, + q_scale_base, + stride_qsm, + stride_qsk, + offs_d_qk, + offs_d_v, + offs_n, + mask_d_qk, + mask_d_v, + q_offset, + do_offset, + stride_qm, + stride_qk, + stride_dom, + stride_dok, + l_offset, + d_offset, + stride_ldm, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DMODEL_QK: tl.constexpr, + BLOCK_DMODEL_V: tl.constexpr, + SCALE_BLOCK_DMODEL_QK: tl.constexpr, + SCALE_ACTUAL_BLOCK_DMODEL_QK: tl.constexpr, + sm_scale: tl.constexpr, + lo, + num_block_m: tl.constexpr, + USE_EXP2: tl.constexpr, + N_CTX_Q: tl.constexpr, + N_CTX_K: tl.constexpr, + CAUSAL: tl.constexpr, + QUANT_BLOCK_SIZE: tl.constexpr, + USE_ASM: tl.constexpr, +): + """Accumulate dk, dv over query blocks for one key block (A1, dot_scaled). + + ``kt_fp8`` / ``vt_fp8`` are the saved-e4m3 key/value tiles transposed to + ``[head_dim, BLOCK_N]`` (fixed across the m-loop); ``ks_hd`` / ``vs_hd`` are + their ``[BLOCK_N, head_dim/32]`` scales rebuilt from the forward's compact 2D + scale by the caller. ``q`` is likewise the **saved e4m3** and is reused with + no re-quantization: dot a reads its head_dim-grouped scale, dot d reads the + same 2D block scale re-indexed along seqlen. Only dO/P/dS are quantized fresh + (plan §9.5 dots a/b/c/d). + """ + for start_m in range(lo, num_block_m * BLOCK_M, BLOCK_M): + offs_m = start_m + tl.arange(0, BLOCK_M) + q_ptrs = q_offset + offs_m[:, None] * stride_qm + offs_d_qk[None, :] * stride_qk + do_ptrs = do_offset + offs_m[:, None] * stride_dom + offs_d_v[None, :] * stride_dok + + mask_m = offs_m < N_CTX_Q + q_mask = mask_m[:, None] & mask_d_qk[None, :] + do_mask = mask_m[:, None] & mask_d_v[None, :] + + # Saved e4m3 q — reused directly, no re-quantization (A1). + q = tl.load(q_ptrs, mask=q_mask, other=0.0) + # dot a: qk = q @ kᵀ, reduction = head_dim. Reuse q's saved head_dim scale. + qs_hd = _load_scale_hd(q_scale_base, offs_m, stride_qsm, stride_qsk, N_CTX_Q, + SCALE_BLOCK_DMODEL_QK, SCALE_ACTUAL_BLOCK_DMODEL_QK, QUANT_BLOCK_SIZE) + qk = tl.dot_scaled(q, qs_hd, "e4m3", kt_fp8, ks_hd, "e4m3", out_dtype=tl.float32) + + if CAUSAL: + # Bottom-right causal, matches the forward kernel (offs_n_causal = offs_n + N_CTX_Q - N_CTX_K). + col_offset = N_CTX_Q - N_CTX_K + causal_mask = offs_m[:, None] >= (col_offset + offs_n[None, :]) + qk = tl.where(causal_mask, qk, float("-inf")) + + l_ptrs = l_offset + offs_m * stride_ldm + l_i = tl.load(l_ptrs, mask=mask_m, other=0.0) + if USE_EXP2: + qk *= sm_scale * RCP_LN2 + p = tl.math.exp2(qk - l_i[:, None] * RCP_LN2) + else: + qk *= sm_scale + p = tl.math.exp(qk - l_i[:, None]) + + do = tl.load(do_ptrs, mask=do_mask, other=0.0) + # dot b: dp = do @ vᵀ, reduction = head_dim_v. dO is fresh -> quantize here. + do_fp8_hd, dos_hd = _mx_quant(do, BLOCK_M, BLOCK_DMODEL_V, QUANT_BLOCK_SIZE, _MX_2D, USE_ASM) + dp = tl.dot_scaled(do_fp8_hd, dos_hd, "e4m3", vt_fp8, vs_hd, "e4m3", out_dtype=tl.float32) + + d_ptrs = d_offset + offs_m * stride_ldm + Di = tl.load(d_ptrs, mask=mask_m, other=0.0) + ds = p * (dp - Di[:, None]) + + # dot c: dv += pᵀ @ do, reduction = seqlen_q (BLOCK_M). P/dO fresh -> quantize + # 1D-block along BLOCK_M (their last axis after transpose). + pt_fp8, ps_c = _mx_quant(tl.trans(p), BLOCK_N, BLOCK_M, QUANT_BLOCK_SIZE, _MX_2D, USE_ASM) + do_t_fp8, dos_c = _mx_quant(tl.trans(do), BLOCK_DMODEL_V, BLOCK_M, QUANT_BLOCK_SIZE, _MX_2D, USE_ASM) + dv += tl.dot_scaled(pt_fp8, ps_c, "e4m3", tl.trans(do_t_fp8), dos_c, "e4m3", out_dtype=tl.float32) + + # dot d: dk += dsᵀ @ q, reduction = seqlen_q (BLOCK_M). dS fresh (1D along + # BLOCK_M); q reuses the saved 2D block scale re-indexed along seqlen. + dst_fp8, dss_d = _mx_quant(tl.trans(ds), BLOCK_N, BLOCK_M, QUANT_BLOCK_SIZE, _MX_2D, USE_ASM) + qs_sq = _load_scale_sq(q_scale_base, start_m, offs_d_qk, stride_qsm, stride_qsk, N_CTX_Q, + BLOCK_M, SCALE_ACTUAL_BLOCK_DMODEL_QK, QUANT_BLOCK_SIZE) + dk += tl.dot_scaled(dst_fp8, dss_d, "e4m3", q, qs_sq, "e4m3", out_dtype=tl.float32) + + return dk, dv + + +@triton.jit +def _bwd_kernel_dkdv( + Q, + K, + V, + Q_scale, + K_scale, + V_scale, + sm_scale: tl.constexpr, + DO, + DK, + DV, + LSE, + Delta, + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_vz, + stride_vh, + stride_vn, + stride_vk, + stride_doz, + stride_doh, + stride_dom, + stride_dok, + stride_qsz, + stride_qsh, + stride_qsm, + stride_qsk, + stride_ksz, + stride_ksh, + stride_ksn, + stride_ksk, + stride_vsz, + stride_vsh, + stride_vsn, + stride_vsk, + stride_ldz, + stride_ldh, + stride_ldm, + Z, + HQ: tl.constexpr, + HK: tl.constexpr, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + num_block_m: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DMODEL_QK: tl.constexpr, + BLOCK_DMODEL_V: tl.constexpr, + ACTUAL_BLOCK_DMODEL_QK: tl.constexpr, + ACTUAL_BLOCK_DMODEL_V: tl.constexpr, + CAUSAL: tl.constexpr, + USE_EXP2: tl.constexpr, + IS_VARLEN: tl.constexpr, + QUANT_BLOCK_SIZE: tl.constexpr, + USE_ASM: tl.constexpr, +): + """One program per (batch*head_k, key block). Parallelizes dk/dv over keys. + + A1: q/k/v arrive as the forward's saved e4m3 with compact 2D-block scales; + the kernel reuses them (no re-quant), rebuilding each dot's scale tile from + the saved scale. Only dO/P/dS are quantized fresh (see ``_attn_bwd_dkdv_inner``). + """ + SCALE_BLOCK_DMODEL_QK: tl.constexpr = BLOCK_DMODEL_QK // QUANT_BLOCK_SIZE + SCALE_BLOCK_DMODEL_V: tl.constexpr = BLOCK_DMODEL_V // QUANT_BLOCK_SIZE + SCALE_ACTUAL_BLOCK_DMODEL_QK: tl.constexpr = ACTUAL_BLOCK_DMODEL_QK // QUANT_BLOCK_SIZE + SCALE_ACTUAL_BLOCK_DMODEL_V: tl.constexpr = ACTUAL_BLOCK_DMODEL_V // QUANT_BLOCK_SIZE + + off_hz = tl.program_id(0) + start_n = tl.program_id(1) + off_z = off_hz // HK + off_h_k = off_hz % HK + + GROUP_SIZE: tl.constexpr = HQ // HK + off_h_q = off_h_k * GROUP_SIZE if GROUP_SIZE != 1 else off_h_k + + if IS_VARLEN: + q_start = tl.load(cu_seqlens_q + off_z) + k_start = tl.load(cu_seqlens_k + off_z) + N_CTX_Q = tl.load(cu_seqlens_q + off_z + 1) - q_start + N_CTX_K = tl.load(cu_seqlens_k + off_z + 1) - k_start + else: + q_start = 0 + k_start = 0 + N_CTX_Q = max_seqlen_q + N_CTX_K = max_seqlen_k + + q_scale_start = q_start // QUANT_BLOCK_SIZE + k_scale_start = k_start // QUANT_BLOCK_SIZE + + q_offset = Q + off_z * stride_qz + off_h_q * stride_qh + q_start * stride_qm + k_offset = K + off_z * stride_kz + off_h_k * stride_kh + k_start * stride_kn + v_offset = V + off_z * stride_vz + off_h_k * stride_vh + k_start * stride_vn + do_offset = DO + off_z * stride_doz + off_h_q * stride_doh + q_start * stride_dom + q_scale_base = Q_scale + off_z * stride_qsz + off_h_q * stride_qsh + q_scale_start * stride_qsm + k_scale_base = K_scale + off_z * stride_ksz + off_h_k * stride_ksh + k_scale_start * stride_ksn + v_scale_base = V_scale + off_z * stride_vsz + off_h_k * stride_vsh + k_scale_start * stride_vsn + adj_delta = off_z * stride_ldz + off_h_q * stride_ldh + q_start * stride_ldm + l_offset = LSE + adj_delta + d_offset = Delta + adj_delta + dk_offset = DK + off_z * stride_kz + off_h_k * stride_kh + k_start * stride_kn + dv_offset = DV + off_z * stride_vz + off_h_k * stride_vh + k_start * stride_vn + + if CAUSAL: + causal_boundary = start_n * BLOCK_N - BLOCK_M + lo = (causal_boundary + 1) // BLOCK_M * BLOCK_M + else: + lo = 0 + + offs_d_qk = tl.arange(0, BLOCK_DMODEL_QK) + offs_d_v = tl.arange(0, BLOCK_DMODEL_V) + offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N) + + mask_n = offs_n < N_CTX_K + mask_d_qk = offs_d_qk < ACTUAL_BLOCK_DMODEL_QK + mask_d_v = offs_d_v < ACTUAL_BLOCK_DMODEL_V + + k_ptrs = k_offset + offs_n[:, None] * stride_kn + offs_d_qk[None, :] * stride_kk + v_ptrs = v_offset + offs_n[:, None] * stride_vn + offs_d_v[None, :] * stride_vk + # Saved e4m3 k/v reused directly; transpose to [head_dim, BLOCK_N] for the + # QK / dP dots (a/b). Scales are rebuilt from the saved compact 2D scale. + k_fp8 = tl.load(k_ptrs, mask=mask_n[:, None] & mask_d_qk[None, :], other=0.0) + v_fp8 = tl.load(v_ptrs, mask=mask_n[:, None] & mask_d_v[None, :], other=0.0) + ks_hd = _load_scale_hd(k_scale_base, offs_n, stride_ksn, stride_ksk, N_CTX_K, + SCALE_BLOCK_DMODEL_QK, SCALE_ACTUAL_BLOCK_DMODEL_QK, QUANT_BLOCK_SIZE) + vs_hd = _load_scale_hd(v_scale_base, offs_n, stride_vsn, stride_vsk, N_CTX_K, + SCALE_BLOCK_DMODEL_V, SCALE_ACTUAL_BLOCK_DMODEL_V, QUANT_BLOCK_SIZE) + kt_fp8 = tl.trans(k_fp8) + vt_fp8 = tl.trans(v_fp8) + + dk = tl.zeros([BLOCK_N, BLOCK_DMODEL_QK], dtype=tl.float32) + dv = tl.zeros([BLOCK_N, BLOCK_DMODEL_V], dtype=tl.float32) + + for _ in range(GROUP_SIZE): + dk, dv = _attn_bwd_dkdv_inner( + kt_fp8, + ks_hd, + vt_fp8, + vs_hd, + dk, + dv, + q_scale_base, + stride_qsm, + stride_qsk, + offs_d_qk, + offs_d_v, + offs_n, + mask_d_qk, + mask_d_v, + q_offset, + do_offset, + stride_qm, + stride_qk, + stride_dom, + stride_dok, + l_offset, + d_offset, + stride_ldm, + BLOCK_M, + BLOCK_N, + BLOCK_DMODEL_QK, + BLOCK_DMODEL_V, + SCALE_BLOCK_DMODEL_QK, + SCALE_ACTUAL_BLOCK_DMODEL_QK, + sm_scale, + lo, + num_block_m, + USE_EXP2, + N_CTX_Q, + N_CTX_K, + CAUSAL, + QUANT_BLOCK_SIZE, + USE_ASM, + ) + q_offset += stride_qh + do_offset += stride_doh + q_scale_base += stride_qsh + l_offset += stride_ldh + d_offset += stride_ldh + + dk *= sm_scale + + tl.store(dk_offset + offs_n[:, None] * stride_kn + offs_d_qk[None, :] * stride_kk, dk, + mask=mask_n[:, None] & mask_d_qk[None, :]) + tl.store(dv_offset + offs_n[:, None] * stride_vn + offs_d_v[None, :] * stride_vk, dv, + mask=mask_n[:, None] & mask_d_v[None, :]) + + +@triton.jit +def _attn_bwd_dq_inner( + dq, + q_fp8_hd, + qs_hd, + do_fp8_hd, + dos_hd, + k_scale_base, + v_scale_base, + stride_ksn, + stride_ksk, + stride_vsn, + stride_vsk, + offs_d_qk, + offs_d_v, + offs_m, + l_i, + Di, + mask_d_qk, + mask_d_v, + k_offset, + v_offset, + stride_kn, + stride_kk, + stride_vn, + stride_vk, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DMODEL_QK: tl.constexpr, + BLOCK_DMODEL_V: tl.constexpr, + SCALE_BLOCK_DMODEL_QK: tl.constexpr, + SCALE_BLOCK_DMODEL_V: tl.constexpr, + SCALE_ACTUAL_BLOCK_DMODEL_QK: tl.constexpr, + SCALE_ACTUAL_BLOCK_DMODEL_V: tl.constexpr, + sm_scale: tl.constexpr, + hi, + USE_EXP2: tl.constexpr, + N_CTX_Q: tl.constexpr, + N_CTX_K: tl.constexpr, + CAUSAL: tl.constexpr, + QUANT_BLOCK_SIZE: tl.constexpr, + USE_ASM: tl.constexpr, +): + """Accumulate dq over key blocks for one query block (A1, dot_scaled). + + ``q_fp8_hd`` / ``qs_hd`` are the saved-e4m3 q + its head_dim scale (fixed + across the n-loop, built once by the caller); ``do_fp8_hd`` / ``dos_hd`` are + the fresh-quantized dO. k/v are the saved e4m3 reused with no re-quant: dots + e/f read their head_dim scale, dot g reads k's 2D block scale re-indexed along + seqlen. Only dS is quantized fresh (plan §9.5 dots e/f/g). + """ + if USE_EXP2: + l_i *= RCP_LN2 + + for start_n in range(0, hi, BLOCK_N): + offs_n = start_n + tl.arange(0, BLOCK_N) + mask_n = offs_n < N_CTX_K + mask_k = mask_n[:, None] & mask_d_qk[None, :] + mask_v = mask_n[:, None] & mask_d_v[None, :] + + k_ptrs = k_offset + offs_n[:, None] * stride_kn + offs_d_qk[None, :] * stride_kk + v_ptrs = v_offset + offs_n[:, None] * stride_vn + offs_d_v[None, :] * stride_vk + # Saved e4m3 k/v, reused directly. + k = tl.load(k_ptrs, mask=mask_k, other=0.0) + v = tl.load(v_ptrs, mask=mask_v, other=0.0) + + # dot e: qk = q @ kᵀ, reduction = head_dim. Reuse k's saved head_dim scale. + ks_e = _load_scale_hd(k_scale_base, offs_n, stride_ksn, stride_ksk, N_CTX_K, + SCALE_BLOCK_DMODEL_QK, SCALE_ACTUAL_BLOCK_DMODEL_QK, QUANT_BLOCK_SIZE) + qk = tl.dot_scaled(q_fp8_hd, qs_hd, "e4m3", tl.trans(k), ks_e, "e4m3", out_dtype=tl.float32) + + if CAUSAL: + col_offset = N_CTX_Q - N_CTX_K + causal_mask = offs_m[:, None] >= (col_offset + offs_n[None, :]) + qk = tl.where(causal_mask, qk, float("-inf")) + + if USE_EXP2: + qk *= sm_scale * RCP_LN2 + p = tl.math.exp2(qk - l_i[:, None]) + else: + qk *= sm_scale + p = tl.math.exp(qk - l_i[:, None]) + + # dot f: dp = do @ vᵀ, reduction = head_dim_v. Reuse v's saved head_dim scale. + vs_f = _load_scale_hd(v_scale_base, offs_n, stride_vsn, stride_vsk, N_CTX_K, + SCALE_BLOCK_DMODEL_V, SCALE_ACTUAL_BLOCK_DMODEL_V, QUANT_BLOCK_SIZE) + dp = tl.dot_scaled(do_fp8_hd, dos_hd, "e4m3", tl.trans(v), vs_f, "e4m3", out_dtype=tl.float32) + ds = p * (dp - Di[:, None]) + + # dot g: dq += ds @ k, reduction = seqlen_k (BLOCK_N). dS fresh (1D along + # BLOCK_N); k reuses the saved 2D block scale re-indexed along seqlen. + ds_fp8, dss_g = _mx_quant(ds, BLOCK_M, BLOCK_N, QUANT_BLOCK_SIZE, _MX_2D, USE_ASM) + ks_sk = _load_scale_sq(k_scale_base, start_n, offs_d_qk, stride_ksn, stride_ksk, N_CTX_K, + BLOCK_N, SCALE_ACTUAL_BLOCK_DMODEL_QK, QUANT_BLOCK_SIZE) + dq += tl.dot_scaled(ds_fp8, dss_g, "e4m3", k, ks_sk, "e4m3", out_dtype=tl.float32) + + return dq + + +@triton.jit +def _bwd_kernel_dq( + Q, + K, + V, + Q_scale, + K_scale, + V_scale, + sm_scale: tl.constexpr, + DO, + DQ, + LSE, + Delta, + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_vz, + stride_vh, + stride_vn, + stride_vk, + stride_doz, + stride_doh, + stride_dom, + stride_dok, + stride_qsz, + stride_qsh, + stride_qsm, + stride_qsk, + stride_ksz, + stride_ksh, + stride_ksn, + stride_ksk, + stride_vsz, + stride_vsh, + stride_vsn, + stride_vsk, + stride_ldz, + stride_ldh, + stride_ldm, + Z, + HQ: tl.constexpr, + HK: tl.constexpr, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + num_block_n: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_DMODEL_QK: tl.constexpr, + BLOCK_DMODEL_V: tl.constexpr, + ACTUAL_BLOCK_DMODEL_QK: tl.constexpr, + ACTUAL_BLOCK_DMODEL_V: tl.constexpr, + CAUSAL: tl.constexpr, + USE_EXP2: tl.constexpr, + IS_VARLEN: tl.constexpr, + QUANT_BLOCK_SIZE: tl.constexpr, + USE_ASM: tl.constexpr, +): + """One program per (batch*head_q, query block). Parallelizes dq over queries. + + A1: q/k/v are the forward's saved e4m3 + compact 2D-block scales, reused with + no re-quant. Only dO (fresh input) and dS are quantized in-kernel. + """ + SCALE_BLOCK_DMODEL_QK: tl.constexpr = BLOCK_DMODEL_QK // QUANT_BLOCK_SIZE + SCALE_BLOCK_DMODEL_V: tl.constexpr = BLOCK_DMODEL_V // QUANT_BLOCK_SIZE + SCALE_ACTUAL_BLOCK_DMODEL_QK: tl.constexpr = ACTUAL_BLOCK_DMODEL_QK // QUANT_BLOCK_SIZE + SCALE_ACTUAL_BLOCK_DMODEL_V: tl.constexpr = ACTUAL_BLOCK_DMODEL_V // QUANT_BLOCK_SIZE + + off_hz = tl.program_id(0) + start_m = tl.program_id(1) + off_z = off_hz // HQ + off_h_q = off_hz % HQ + + GROUP_SIZE: tl.constexpr = HQ // HK + off_h_k = off_h_q // GROUP_SIZE if GROUP_SIZE != 1 else off_h_q + + if IS_VARLEN: + q_start = tl.load(cu_seqlens_q + off_z) + k_start = tl.load(cu_seqlens_k + off_z) + N_CTX_Q = tl.load(cu_seqlens_q + off_z + 1) - q_start + N_CTX_K = tl.load(cu_seqlens_k + off_z + 1) - k_start + else: + q_start = 0 + k_start = 0 + N_CTX_Q = max_seqlen_q + N_CTX_K = max_seqlen_k + + q_scale_start = q_start // QUANT_BLOCK_SIZE + k_scale_start = k_start // QUANT_BLOCK_SIZE + + q_offset = Q + off_z * stride_qz + off_h_q * stride_qh + q_start * stride_qm + k_offset = K + off_z * stride_kz + off_h_k * stride_kh + k_start * stride_kn + v_offset = V + off_z * stride_vz + off_h_k * stride_vh + k_start * stride_vn + do_offset = DO + off_z * stride_doz + off_h_q * stride_doh + q_start * stride_dom + q_scale_base = Q_scale + off_z * stride_qsz + off_h_q * stride_qsh + q_scale_start * stride_qsm + k_scale_base = K_scale + off_z * stride_ksz + off_h_k * stride_ksh + k_scale_start * stride_ksn + v_scale_base = V_scale + off_z * stride_vsz + off_h_k * stride_vsh + k_scale_start * stride_vsn + adj_delta = off_z * stride_ldz + off_h_q * stride_ldh + q_start * stride_ldm + l_offset = LSE + adj_delta + d_offset = Delta + adj_delta + dq_offset = DQ + off_z * stride_qz + off_h_q * stride_qh + q_start * stride_qm + + if CAUSAL: + hi = tl.minimum(BLOCK_M // BLOCK_N * (start_m + 1), num_block_n) * BLOCK_N + else: + hi = num_block_n * BLOCK_N + + offs_d_qk = tl.arange(0, BLOCK_DMODEL_QK) + offs_d_v = tl.arange(0, BLOCK_DMODEL_V) + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + + mask_m = offs_m < N_CTX_Q + mask_d_qk = offs_d_qk < ACTUAL_BLOCK_DMODEL_QK + mask_d_v = offs_d_v < ACTUAL_BLOCK_DMODEL_V + + q_ptrs = q_offset + offs_m[:, None] * stride_qm + offs_d_qk[None, :] * stride_qk + do_ptrs = do_offset + offs_m[:, None] * stride_dom + offs_d_v[None, :] * stride_dok + # Saved e4m3 q reused directly; dO is a fresh input and is quantized here. + q_fp8_hd = tl.load(q_ptrs, mask=mask_m[:, None] & mask_d_qk[None, :], other=0.0) + do = tl.load(do_ptrs, mask=mask_m[:, None] & mask_d_v[None, :], other=0.0) + + l_i = tl.load(l_offset + offs_m * stride_ldm, mask=mask_m, other=0.0) + Di = tl.load(d_offset + offs_m * stride_ldm, mask=mask_m, other=0.0) + + # q's saved head_dim scale (dot e), fixed across the n-loop; dO quantized fresh. + qs_hd = _load_scale_hd(q_scale_base, offs_m, stride_qsm, stride_qsk, N_CTX_Q, + SCALE_BLOCK_DMODEL_QK, SCALE_ACTUAL_BLOCK_DMODEL_QK, QUANT_BLOCK_SIZE) + do_fp8_hd, dos_hd = _mx_quant(do, BLOCK_M, BLOCK_DMODEL_V, QUANT_BLOCK_SIZE, _MX_2D, USE_ASM) + + dq = tl.zeros([BLOCK_M, BLOCK_DMODEL_QK], dtype=tl.float32) + dq = _attn_bwd_dq_inner( + dq, + q_fp8_hd, + qs_hd, + do_fp8_hd, + dos_hd, + k_scale_base, + v_scale_base, + stride_ksn, + stride_ksk, + stride_vsn, + stride_vsk, + offs_d_qk, + offs_d_v, + offs_m, + l_i, + Di, + mask_d_qk, + mask_d_v, + k_offset, + v_offset, + stride_kn, + stride_kk, + stride_vn, + stride_vk, + BLOCK_M, + BLOCK_N, + BLOCK_DMODEL_QK, + BLOCK_DMODEL_V, + SCALE_BLOCK_DMODEL_QK, + SCALE_BLOCK_DMODEL_V, + SCALE_ACTUAL_BLOCK_DMODEL_QK, + SCALE_ACTUAL_BLOCK_DMODEL_V, + sm_scale, + hi, + USE_EXP2, + N_CTX_Q, + N_CTX_K, + CAUSAL, + QUANT_BLOCK_SIZE, + USE_ASM, + ) + + dq *= sm_scale + tl.store(dq_offset + offs_m[:, None] * stride_qm + offs_d_qk[None, :] * stride_qk, dq, + mask=mask_m[:, None] & mask_d_qk[None, :]) + + +@triton_op("alto::attention_mxfp8_backward_triton_impl", mutates_args=()) +def attention_mxfp8_backward_triton_impl( + do: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + softmax_lse: torch.Tensor, + q_scale: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + sm_scale: float, + causal: bool, + layout: str, + cu_seqlens_q: Optional[int], + cu_seqlens_k: Optional[int], + max_seqlen_q: Optional[int], + max_seqlen_k: Optional[int], + use_exp2: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """MXFP8 flash-attention backward (stage-2 / A1, ``tl.dot_scaled``). + + A1 operand source: the forward-saved e4m3 q/k/v and their compact 2D-block + scales are passed straight in and reused in *every* dot (no entry dequant, no + re-quantization of q/k/v). For head_dim-reduction dots the 2D scale already + groups along the reduction axis; for seqlen-reduction dots the same compact + 2D scale is re-indexed (``_load_scale_sq``) to broadcast per reduction row. + Only the backward-only tensors dO/P/dS are quantized in-kernel (1D per-row + along each dot's reduction axis), then fed to ``tl.dot_scaled`` (plan §9.5). + Matches ``mxfp8_attention_backward_reference_stage2``. + + Requires CDNA4 (native ``tl.dot_scaled``). + """ + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + q_scale = q_scale.contiguous() + k_scale = k_scale.contiguous() + v_scale = v_scale.contiguous() + + if not do.is_contiguous(): + do = do.contiguous() + + assert layout == "bhsd", ( + f"MXFP8 attention requires layout='bhsd', got {layout!r}: the 2D-block scale groups the " + "data tensor's last two dims, which are (seqlen, head_dim) only for bhsd. Any other " + "layout groups across heads and the scale pointer math below reads the wrong scale.") + + batch, nheads_q, nheads_k, head_size_qk, head_size_v, max_seqlen_q, max_seqlen_k = get_shape_from_layout( + q, k, v, layout, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k) + assert not causal or max_seqlen_q == max_seqlen_k, ( + f"causal backward requires seqlen_q == seqlen_k, got {max_seqlen_q} vs {max_seqlen_k}: the " + "kernel masks bottom-right aligned while the PyTorch references and F.sdpa mask top-left, " + "and the two conventions only agree on square shapes.") + stride_qz, stride_qh, stride_qm, stride_qk = get_strides_from_layout(q, layout) + stride_kz, stride_kh, stride_kn, stride_kk = get_strides_from_layout(k, layout) + stride_vz, stride_vh, stride_vn, stride_vk = get_strides_from_layout(v, layout) + stride_oz, stride_oh, stride_om, stride_ok = get_strides_from_layout(o, layout) + stride_doz, stride_doh, stride_dom, stride_dok = get_strides_from_layout(do, layout) + # Compact 2D-block scale strides ([.., seqlen/32, head_dim/32]); reused by the + # kernels for both head_dim- and seqlen-reduction dots (A1). + stride_qsz, stride_qsh, stride_qsm, stride_qsk = get_strides_from_layout(q_scale, layout) + stride_ksz, stride_ksh, stride_ksn, stride_ksk = get_strides_from_layout(k_scale, layout) + stride_vsz, stride_vsh, stride_vsn, stride_vsk = get_strides_from_layout(v_scale, layout) + is_varlen = layout == "thd" + + padded_d_model_qk = get_padded_head_dim(head_size_qk) + padded_d_model_v = get_padded_head_dim(head_size_v) + + dq = torch.zeros_like(q, dtype=bwd_torch_dtype) + dk = torch.zeros_like(k, dtype=bwd_torch_dtype) + dv = torch.zeros_like(v, dtype=bwd_torch_dtype) + + delta = torch.empty_like(softmax_lse) + if is_varlen: + stride_lse_m, stride_lse_h = softmax_lse.stride() + stride_lse_z = 0 + else: + stride_lse_z, stride_lse_h, stride_lse_m = softmax_lse.stride() + + BLOCK_M = 64 + BLOCK_N = 64 + num_block_m = triton.cdiv(max_seqlen_q, BLOCK_M) + num_block_n = triton.cdiv(max_seqlen_k, BLOCK_N) + + wrap_triton(_bwd_preprocess)[(num_block_m, batch * nheads_q)]( + o, + do, + delta, + stride_oz, + stride_oh, + stride_om, + stride_ok, + stride_doz, + stride_doh, + stride_dom, + stride_dok, + stride_lse_z, + stride_lse_h, + stride_lse_m, + cu_seqlens_q, + max_seqlen_q, + BLOCK_M=BLOCK_M, + BLOCK_DMODEL_V=padded_d_model_v, + ACTUAL_BLOCK_DMODEL_V=head_size_v, + HQ=nheads_q, + IS_VARLEN=is_varlen, + ) + + wrap_triton(_bwd_kernel_dq)[(batch * nheads_q, num_block_m)]( + q, + k, + v, + q_scale, + k_scale, + v_scale, + sm_scale, + do, + dq, + softmax_lse, + delta, + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_vz, + stride_vh, + stride_vn, + stride_vk, + stride_doz, + stride_doh, + stride_dom, + stride_dok, + stride_qsz, + stride_qsh, + stride_qsm, + stride_qsk, + stride_ksz, + stride_ksh, + stride_ksn, + stride_ksk, + stride_vsz, + stride_vsh, + stride_vsn, + stride_vsk, + stride_lse_z, + stride_lse_h, + stride_lse_m, + batch, + nheads_q, + nheads_k, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + num_block_n=num_block_n, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_DMODEL_QK=padded_d_model_qk, + BLOCK_DMODEL_V=padded_d_model_v, + ACTUAL_BLOCK_DMODEL_QK=head_size_qk, + ACTUAL_BLOCK_DMODEL_V=head_size_v, + CAUSAL=causal, + USE_EXP2=use_exp2, + IS_VARLEN=is_varlen, + QUANT_BLOCK_SIZE=BLOCK_SIZE_DEFAULT, + USE_ASM=is_cdna4(), + num_warps=4, + num_stages=1, + ) + + wrap_triton(_bwd_kernel_dkdv)[(batch * nheads_k, num_block_n)]( + q, + k, + v, + q_scale, + k_scale, + v_scale, + sm_scale, + do, + dk, + dv, + softmax_lse, + delta, + stride_qz, + stride_qh, + stride_qm, + stride_qk, + stride_kz, + stride_kh, + stride_kn, + stride_kk, + stride_vz, + stride_vh, + stride_vn, + stride_vk, + stride_doz, + stride_doh, + stride_dom, + stride_dok, + stride_qsz, + stride_qsh, + stride_qsm, + stride_qsk, + stride_ksz, + stride_ksh, + stride_ksn, + stride_ksk, + stride_vsz, + stride_vsh, + stride_vsn, + stride_vsk, + stride_lse_z, + stride_lse_h, + stride_lse_m, + batch, + nheads_q, + nheads_k, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + num_block_m=num_block_m, + BLOCK_M=BLOCK_M, + BLOCK_N=BLOCK_N, + BLOCK_DMODEL_QK=padded_d_model_qk, + BLOCK_DMODEL_V=padded_d_model_v, + ACTUAL_BLOCK_DMODEL_QK=head_size_qk, + ACTUAL_BLOCK_DMODEL_V=head_size_v, + CAUSAL=causal, + USE_EXP2=use_exp2, + IS_VARLEN=is_varlen, + QUANT_BLOCK_SIZE=BLOCK_SIZE_DEFAULT, + USE_ASM=is_cdna4(), + num_warps=4, + num_stages=1, + ) + + return dq, dk, dv + + +@attention_mxfp8_backward_triton_impl.register_fake +def fake_attention_mxfp8_backward_triton_impl( + do: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + softmax_lse: torch.Tensor, + q_scale: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + sm_scale: float, + causal: bool, + layout: str, + cu_seqlens_q: Optional[int], + cu_seqlens_k: Optional[int], + max_seqlen_q: Optional[int], + max_seqlen_k: Optional[int], + use_exp2: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + dq = torch.empty_like(q, dtype=bwd_torch_dtype) + dk = torch.empty_like(k, dtype=bwd_torch_dtype) + dv = torch.empty_like(v, dtype=bwd_torch_dtype) + return dq, dk, dv + + +@torch.compiler.allow_in_graph +class _triton_attention_mxfp8(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + alibi_slopes: torch.Tensor | None, + bias: torch.Tensor | None, + sm_scale: float, + dropout_p: float, + cu_seqlens_q: int, + cu_seqlens_k: int, + max_seqlens_q: int, + max_seqlens_k: int, + causal: bool, + return_scores: bool, + use_exp2: bool, + layout: str, + ): + q, q_scale = torch.ops.alto.convert_to_mxfp8(q, mxfp_format="e4m3", axis=-1, is_2d_block=True) + k, k_scale = torch.ops.alto.convert_to_mxfp8(k, mxfp_format="e4m3", axis=-1, is_2d_block=True) + v, v_scale = torch.ops.alto.convert_to_mxfp8(v, mxfp_format="e4m3", axis=-1, is_2d_block=True) + + output, softmax_lse, exp_scores = torch.ops.alto.attention_mxfp8_forward_triton_impl( + q, + k, + v, + q_scale, + k_scale, + v_scale, + sm_scale=sm_scale, + alibi_slopes=alibi_slopes, + causal=causal, + bias=bias, + dropout_p=dropout_p, + layout=layout, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlens_q=max_seqlens_q, + max_seqlens_k=max_seqlens_k, + return_scores=return_scores, + use_exp2=use_exp2, + ) + + ctx.save_for_backward(q, k, v, output, softmax_lse, alibi_slopes, bias, q_scale, k_scale, v_scale) + ctx.sm_scale = sm_scale + ctx.causal = causal + ctx.dropout_p = dropout_p + ctx.layout = layout + ctx.use_exp2 = use_exp2 + ctx.cu_seqlens_q = cu_seqlens_q + ctx.cu_seqlens_k = cu_seqlens_k + ctx.max_seqlens_q = max_seqlens_q + ctx.max_seqlens_k = max_seqlens_k + + return output, softmax_lse, exp_scores + + @staticmethod + def backward(ctx, *grad_outputs): + do = grad_outputs[0] + q, k, v, o, softmax_lse, alibi_slopes, bias, q_scale, k_scale, v_scale = ctx.saved_tensors + assert bias is None, "MXFP8 attention backward does not support bias yet." + assert alibi_slopes is None, "MXFP8 attention backward does not support alibi yet." + assert ctx.dropout_p == 0.0, "MXFP8 attention backward does not support dropout yet." + + dq, dk, dv = torch.ops.alto.attention_mxfp8_backward_triton_impl( + do, + q, + k, + v, + o, + softmax_lse, + q_scale, + k_scale, + v_scale, + sm_scale=ctx.sm_scale, + causal=ctx.causal, + layout=ctx.layout, + cu_seqlens_q=ctx.cu_seqlens_q, + cu_seqlens_k=ctx.cu_seqlens_k, + max_seqlen_q=ctx.max_seqlens_q, + max_seqlen_k=ctx.max_seqlens_k, + use_exp2=ctx.use_exp2, + ) + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None + + +@attention_mxfp8_forward_triton_impl.register_fake +def fake_attention_mxfp8_forward_triton_impl( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_scale: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + sm_scale: float, + alibi_slopes: Optional[torch.Tensor], + causal: bool, + bias: Optional[torch.Tensor], + dropout_p: float, + layout: str, + cu_seqlens_q: Optional[int], + cu_seqlens_k: Optional[int], + max_seqlens_q: Optional[int], + max_seqlens_k: Optional[int], + return_scores: bool, + use_exp2: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + o_shape = list(q.shape) + o_shape[-1] = v.shape[-1] # output shape should match v's head dim + o = torch.empty( + o_shape, + device=q.device, + dtype=fwd_torch_dtype, + requires_grad=True, + ) + + # check if varlen + is_varlen = layout == "thd" + + batch, nheads_q, nheads_k, head_size_qk, head_size_v, seqlen_q, seqlen_k = get_shape_from_layout( + q, k, v, layout, cu_seqlens_q, cu_seqlens_k, max_seqlens_q, max_seqlens_k) + + if return_scores: + scores = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device, dtype=torch.float32) + else: + scores = torch.empty([], device=q.device, dtype=torch.float32) + + if return_scores: + exp_scores = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device, dtype=torch.float32) + else: + exp_scores = torch.empty([], device=q.device, dtype=torch.float32) + + if is_varlen: + softmax_lse = torch.empty((q.shape[0], nheads_q), device=q.device, dtype=torch.float32) + else: + softmax_lse = torch.empty((batch, nheads_q, max_seqlens_q), device=q.device, dtype=torch.float32) + return o, softmax_lse, exp_scores + + +def triton_attention_mxfp8( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + alibi_slopes: torch.Tensor | None, + bias: torch.Tensor | None, + sm_scale: float, + dropout_p: float, + cu_seqlens_q: int, + cu_seqlens_k: int, + max_seqlens_q: int, + max_seqlens_k: int, + causal: bool, + return_scores: bool, + use_exp2: bool, + layout: str, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return _triton_attention_mxfp8.apply( + q, + k, + v, + alibi_slopes, + bias, + sm_scale, + dropout_p, + cu_seqlens_q, + cu_seqlens_k, + max_seqlens_q, + max_seqlens_k, + causal, + return_scores, + use_exp2, + layout, + ) diff --git a/tests/unittest/mxfp8/test_mxfp8_attention.py b/tests/unittest/mxfp8/test_mxfp8_attention.py new file mode 100644 index 0000000..0b7cf14 --- /dev/null +++ b/tests/unittest/mxfp8/test_mxfp8_attention.py @@ -0,0 +1,127 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Forward numerical tests for the MXFP8 (e4m3) flash attention kernel. + +Mirrors ``tests/unittest/mxfp4/test_mxfp_attention.py`` for the forward +kernel-vs-bf16 SDPA check. Backward and kernel-vs-golden reference coverage live +in ``test_mxfp8_attention_reference.py``. +""" + +import pytest +from tabulate import tabulate +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA device is required.") + +from alto.kernels.mxfp8.triton_flash_attention_mxfp8 import triton_attention_mxfp8 +from .utils import calc_snr, calc_cossim + + +def attention_vanilla_forward_pytorch_ref_impl(q, k, v, sm_scale, causal, layout="bshd"): + """Compute reference output using PyTorch's built-in SDPA.""" + if layout == "bshd": + num_heads = q.shape[2] + n_kv_heads = k.shape[2] + n_rep = num_heads // n_kv_heads + + q = q.transpose(1, 2).contiguous() + k = k.transpose(1, 2).contiguous() + v = v.transpose(1, 2).contiguous() + else: + raise ValueError(f"Unknown layout {layout}") + + o_ref = torch.nn.functional.scaled_dot_product_attention(q, + k, + v, + is_causal=causal, + scale=sm_scale, + enable_gqa=n_rep > 1) + if layout == "bshd": + o_ref = o_ref.transpose(1, 2) + return o_ref + + +class AttnConfig: + + def __init__(self, seqlen_q, seqlen_kv, num_head_q, num_head_kv, head_dim_qk, head_dim_v): + self.seqlen_q = seqlen_q + self.seqlen_kv = seqlen_kv + self.num_head_q = num_head_q + self.num_head_kv = num_head_kv + self.head_dim_qk = head_dim_qk + self.head_dim_v = head_dim_v + + +test_cases = [ + AttnConfig(seqlen_q=1024, seqlen_kv=1024, num_head_q=32, num_head_kv=32, head_dim_qk=128, head_dim_v=128), + AttnConfig(seqlen_q=1024, seqlen_kv=1024, num_head_q=64, num_head_kv=8, head_dim_qk=128, head_dim_v=128), + AttnConfig(seqlen_q=1024, seqlen_kv=1024, num_head_q=32, num_head_kv=8, head_dim_qk=128, head_dim_v=128), + AttnConfig(seqlen_q=1024, seqlen_kv=1024, num_head_q=16, num_head_kv=16, head_dim_qk=192, head_dim_v=128), + AttnConfig(seqlen_q=1024, seqlen_kv=1024, num_head_q=128, num_head_kv=128, head_dim_qk=192, head_dim_v=128), + AttnConfig(seqlen_q=1024, seqlen_kv=1024, num_head_q=48, num_head_kv=8, head_dim_qk=128, head_dim_v=128), + AttnConfig(seqlen_q=2048, seqlen_kv=2048, num_head_q=64, num_head_kv=8, head_dim_qk=128, head_dim_v=128), + # GQA crossed with head_dim_qk != head_dim_v. The cases above cover those two + # axes only in isolation, which leaves the dO head-stride advance in the + # backward dkdv group loop untested (it produced NaN dK/dV). + AttnConfig(seqlen_q=1024, seqlen_kv=1024, num_head_q=32, num_head_kv=8, head_dim_qk=192, head_dim_v=128), +] + + +@pytest.mark.parametrize("batch", [4]) +@pytest.mark.parametrize("config", test_cases) +@pytest.mark.parametrize("causal", [True, False]) +def test_attention(batch, config, causal): + device = "cuda" + dtype = torch.bfloat16 + seqlen_q, seqlen_kv, num_head_q, num_head_kv, head_dim_qk, head_dim_v = ( + config.seqlen_q, + config.seqlen_kv, + config.num_head_q, + config.num_head_kv, + config.head_dim_qk, + config.head_dim_v, + ) + q_layout = (batch, seqlen_q, num_head_q, head_dim_qk) + k_layout = (batch, seqlen_kv, num_head_kv, head_dim_qk) + v_layout = (batch, seqlen_kv, num_head_kv, head_dim_v) + + torch.manual_seed(1234) + + query = torch.randn(q_layout, device=device, dtype=dtype, requires_grad=True) + key = torch.randn(k_layout, device=device, dtype=dtype, requires_grad=True) + value = torch.randn(v_layout, device=device, dtype=dtype, requires_grad=True) + query_ref = query.clone().detach().requires_grad_() + key_ref = key.clone().detach().requires_grad_() + value_ref = value.clone().detach().requires_grad_() + + sm_scale = query.shape[-1]**(-0.5) + o_ref = attention_vanilla_forward_pytorch_ref_impl(query_ref, key_ref, value_ref, sm_scale, causal) + + o = triton_attention_mxfp8( + query.transpose(1, 2).contiguous(), + key.transpose(1, 2).contiguous(), + value.transpose(1, 2).contiguous(), + bias=None, + alibi_slopes=None, + sm_scale=sm_scale, + dropout_p=0.0, + cu_seqlens_q=0, + cu_seqlens_k=0, + max_seqlens_q=seqlen_q, + max_seqlens_k=seqlen_kv, + causal=causal, + return_scores=False, + use_exp2=True, + layout="bhsd", + )[0].transpose(1, 2) + + output_snr = calc_snr(o_ref, o) + output_sim = calc_cossim(o_ref, o) + print() + print(tabulate([ + ["O", output_snr, output_sim], + ], headers=["Tensor", "SNR", "Cosine Sim"], tablefmt="github")) + + assert output_sim > 0.99, f"output cosine-sim too low: {output_sim}" + assert output_snr > 20, f"output SNR too low: {output_snr}" diff --git a/tests/unittest/mxfp8/test_mxfp8_attention_reference.py b/tests/unittest/mxfp8/test_mxfp8_attention_reference.py new file mode 100644 index 0000000..f3fa969 --- /dev/null +++ b/tests/unittest/mxfp8/test_mxfp8_attention_reference.py @@ -0,0 +1,435 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT +"""Golden-reference validation for the MXFP8 (e4m3) flash attention kernel. + +Additive to ``test_mxfp8_attention.py`` (which mirrors the mxfp4 attention test: +kernel vs bf16 SDPA). This file adds the two layers that the pure-PyTorch golden +reference enables — a strengthening over the mxfp4 reference, which has neither a +golden reference nor asserts: + + * Layer 1 ``test_reference_matches_bf16_sdpa`` — golden reference vs bf16 SDPA. + Pure PyTorch, **runs on CPU / any device without CDNA4**; validates the + algorithm and mxfp8 quantization placement *before* the hardware arrives. + * Layer 2 ``test_kernel_matches_reference`` — kernel vs golden reference. + Isolates Triton port bugs (masking / online-softmax / LSE) from the mxfp8 + quantization error itself. Requires CDNA4 (native ``tl.dot_scaled``). + +(Layer 3 — kernel vs bf16 SDPA, total end-to-end error — is the committed +``test_mxfp8_attention.py::test_attention``; not duplicated here.) +""" + +import pytest +from tabulate import tabulate +import torch + +from .utils import ( + calc_snr, + calc_cossim, + mxfp8_attention_forward_reference, + mxfp8_attention_backward_reference, + mxfp8_attention_backward_reference_stage2, +) +# Reuse the committed shape grid so Layer 2 stays in lock-step with Layer 3. +from .test_mxfp8_attention import AttnConfig, test_cases + +cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA/ROCm device is required.") + + +def _sdpa_bf16_bhsd(q, k, v, sm_scale, causal): + """bf16 SDPA reference in bhsd layout (the native SDPA head layout).""" + n_rep = q.shape[1] // k.shape[1] + return torch.nn.functional.scaled_dot_product_attention( + q, k, v, is_causal=causal, scale=sm_scale, enable_gqa=n_rep > 1) + + +def _make_qkv_bhsd(batch, config, device, dtype): + torch.manual_seed(1234) + q = torch.randn((batch, config.num_head_q, config.seqlen_q, config.head_dim_qk), device=device, dtype=dtype) + k = torch.randn((batch, config.num_head_kv, config.seqlen_kv, config.head_dim_qk), device=device, dtype=dtype) + v = torch.randn((batch, config.num_head_kv, config.seqlen_kv, config.head_dim_v), device=device, dtype=dtype) + return q, k, v + + +# --------------------------------------------------------------------------- +# Layer 1 — golden reference vs bf16 SDPA (pure PyTorch, runs anywhere, no CDNA4) +# --------------------------------------------------------------------------- + +# Small shapes: the reference loops over key blocks in Python, keep it cheap for CPU. +reference_cases = [ + AttnConfig(seqlen_q=128, seqlen_kv=128, num_head_q=4, num_head_kv=4, head_dim_qk=64, head_dim_v=64), + AttnConfig(seqlen_q=256, seqlen_kv=256, num_head_q=8, num_head_kv=2, head_dim_qk=128, head_dim_v=128), # GQA + AttnConfig(seqlen_q=128, seqlen_kv=256, num_head_q=4, num_head_kv=4, head_dim_qk=128, head_dim_v=128), # seqlen_kv > seqlen_q +] + + +@pytest.mark.parametrize("config", reference_cases) +@pytest.mark.parametrize("causal", [True, False]) +def test_reference_matches_bf16_sdpa(config, causal): + """The pure-PyTorch golden reference must track bf16 SDPA within mxfp8 error. + + Runnable on CPU without CDNA4 — the check we can do *before* the hardware + arrives, to de-risk the algorithm and quantization placement. + """ + device = "cuda" if torch.cuda.is_available() else "cpu" + dtype = torch.float32 if device == "cpu" else torch.bfloat16 + + q, k, v = _make_qkv_bhsd(2, config, device, dtype) + sm_scale = config.head_dim_qk**(-0.5) + + o_ref = _sdpa_bf16_bhsd(q, k, v, sm_scale, causal) + o_golden, _ = mxfp8_attention_forward_reference(q, k, v, sm_scale, causal) + + snr = calc_snr(o_ref, o_golden) + sim = calc_cossim(o_ref, o_golden) + print() + print(tabulate([["O (golden vs bf16)", snr, sim]], headers=["Tensor", "SNR", "Cosine Sim"], tablefmt="github")) + + assert sim > 0.99, f"golden reference cosine-sim vs bf16 SDPA too low: {sim}" + assert snr > 15, f"golden reference SNR vs bf16 SDPA too low: {snr}" + + +# --------------------------------------------------------------------------- +# Layer 2 — kernel vs golden reference (requires CDNA4 native tl.dot_scaled) +# --------------------------------------------------------------------------- + +# Both ops assert seqlen_q == seqlen_k when causal, so non-square shapes are only +# reachable without causal masking. They stay out of the shared `test_cases` grid +# (which must remain valid for both causal values) and get their own coverage. +non_causal_cases = [ + AttnConfig(seqlen_q=256, seqlen_kv=512, num_head_q=8, num_head_kv=2, head_dim_qk=128, head_dim_v=128), + AttnConfig(seqlen_q=512, seqlen_kv=256, num_head_q=8, num_head_kv=2, head_dim_qk=192, head_dim_v=128), +] + + +def _check_forward_kernel_vs_reference(config, causal, batch=4): + from alto.kernels.mxfp8.triton_flash_attention_mxfp8 import triton_attention_mxfp8 + + device = "cuda" + dtype = torch.bfloat16 + q, k, v = _make_qkv_bhsd(batch, config, device, dtype) + sm_scale = config.head_dim_qk**(-0.5) + + o_kernel = triton_attention_mxfp8( + q.contiguous(), + k.contiguous(), + v.contiguous(), + bias=None, + alibi_slopes=None, + sm_scale=sm_scale, + dropout_p=0.0, + cu_seqlens_q=0, + cu_seqlens_k=0, + max_seqlens_q=config.seqlen_q, + max_seqlens_k=config.seqlen_kv, + causal=causal, + return_scores=False, + use_exp2=True, + layout="bhsd", + )[0] + o_golden, _ = mxfp8_attention_forward_reference(q, k, v, sm_scale, causal) + + snr = calc_snr(o_golden, o_kernel) + sim = calc_cossim(o_golden, o_kernel) + print() + print(tabulate([["O (kernel vs golden)", snr, sim]], headers=["Tensor", "SNR", "Cosine Sim"], tablefmt="github")) + + # Threshold placeholder — calibrate on the first CDNA4 run. + assert sim > 0.99, f"kernel vs golden cosine-sim too low: {sim}" + assert snr > 30, f"kernel vs golden SNR too low: {snr}" + + +@cuda_only +@pytest.mark.parametrize("config", test_cases) +@pytest.mark.parametrize("causal", [True, False]) +def test_kernel_matches_reference(config, causal): + """Kernel vs golden reference — isolates Triton port bugs from quant error. + + Both sides apply the same mxfp8 quantization, so a large gap here points at a + Triton port bug (masking / online-softmax / LSE / strides) rather than + quantization error. + """ + _check_forward_kernel_vs_reference(config, causal) + + +@cuda_only +@pytest.mark.parametrize("config", non_causal_cases) +def test_kernel_matches_reference_non_square(config): + """Forward kernel on seqlen_q != seqlen_k, the only shape family causal rejects.""" + _check_forward_kernel_vs_reference(config, causal=False, batch=2) + + +@cuda_only +def test_causal_forward_rejects_non_square_shape(): + """Non-square causal masking is intentionally blocked until its semantics are decided.""" + from alto.kernels.mxfp8.triton_flash_attention_mxfp8 import triton_attention_mxfp8 + + config = AttnConfig(seqlen_q=128, seqlen_kv=256, num_head_q=4, num_head_kv=4, head_dim_qk=128, head_dim_v=128) + q, k, v = _make_qkv_bhsd(1, config, "cuda", torch.bfloat16) + + with pytest.raises(AssertionError, match="causal forward requires seqlen_q == seqlen_k"): + triton_attention_mxfp8( + q.contiguous(), + k.contiguous(), + v.contiguous(), + bias=None, + alibi_slopes=None, + sm_scale=config.head_dim_qk**(-0.5), + dropout_p=0.0, + cu_seqlens_q=0, + cu_seqlens_k=0, + max_seqlens_q=config.seqlen_q, + max_seqlens_k=config.seqlen_kv, + causal=True, + return_scores=False, + use_exp2=True, + layout="bhsd", + ) + + +# --------------------------------------------------------------------------- +# Backward Layer 1 — golden reference vs autograd fp32 SDPA (pure PyTorch, CPU) +# --------------------------------------------------------------------------- + +def _make_do_bhsd(batch, config, device, dtype): + torch.manual_seed(4321) + return torch.randn((batch, config.num_head_q, config.seqlen_q, config.head_dim_v), device=device, dtype=dtype) + + +@pytest.mark.parametrize("config", reference_cases) +@pytest.mark.parametrize("causal", [True, False]) +def test_backward_reference_matches_sdpa(config, causal): + """Backward golden reference (mxfp8 quant) must track autograd fp32 SDPA grads. + + Both sides use top-left causal, so this is valid for non-square shapes too. + The gap is the mxfp8 quantization error in the gradients; runs on CPU. + """ + device = "cuda" if torch.cuda.is_available() else "cpu" + dtype = torch.float32 if device == "cpu" else torch.bfloat16 + n_rep = config.num_head_q // config.num_head_kv + + q, k, v = _make_qkv_bhsd(2, config, device, dtype) + do = _make_do_bhsd(2, config, device, dtype) + sm_scale = config.head_dim_qk**(-0.5) + + # Autograd fp32 SDPA backward — the "ideal" (unquantized) gradient. + qg = q.clone().requires_grad_(True) + kg = k.clone().requires_grad_(True) + vg = v.clone().requires_grad_(True) + o_auto = torch.nn.functional.scaled_dot_product_attention( + qg, kg, vg, is_causal=causal, scale=sm_scale, enable_gqa=n_rep > 1) + o_auto.backward(do) + + # Golden reference backward — quantized operands, forward-ref o / lse. + o_ref, lse_ref = mxfp8_attention_forward_reference(q, k, v, sm_scale, causal) + dq, dk, dv = mxfp8_attention_backward_reference(q, k, v, do, o_ref, lse_ref, sm_scale, causal) + + rows = [] + for name, ref, got in [("dQ", qg.grad, dq), ("dK", kg.grad, dk), ("dV", vg.grad, dv)]: + rows.append([f"{name} (golden vs sdpa)", calc_snr(ref, got), calc_cossim(ref, got)]) + print() + print(tabulate(rows, headers=["Tensor", "SNR", "Cosine Sim"], tablefmt="github")) + + for name, ref, got in [("dQ", qg.grad, dq), ("dK", kg.grad, dk), ("dV", vg.grad, dv)]: + sim = calc_cossim(ref, got) + snr = calc_snr(ref, got) + assert sim > 0.97, f"{name} golden-vs-sdpa cosine-sim too low: {sim}" + assert snr > 8, f"{name} golden-vs-sdpa SNR too low: {snr}" + + +# --------------------------------------------------------------------------- +# Backward Stage-2 (A1) Layer 1 — full-quant reference vs autograd fp32 SDPA +# Adds dO/P/dS quantization on top of Stage-1 (a numerical model of the +# kernel's tl.dot_scaled). q/k/v reuse the forward 2D-block dequant (A1, no +# double quant). The gap vs SDPA is the *full* mxfp8 backward quant error — +# the signal that all-e4m3 backward is viable. Pure PyTorch, runs on CPU/MI250. +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("config", reference_cases) +@pytest.mark.parametrize("causal", [True, False]) +def test_backward_stage2_reference_matches_sdpa(config, causal): + """Stage-2 (A1, all-operand-quantized) backward reference vs autograd fp32 SDPA.""" + device = "cuda" if torch.cuda.is_available() else "cpu" + dtype = torch.float32 if device == "cpu" else torch.bfloat16 + n_rep = config.num_head_q // config.num_head_kv + + q, k, v = _make_qkv_bhsd(2, config, device, dtype) + do = _make_do_bhsd(2, config, device, dtype) + sm_scale = config.head_dim_qk**(-0.5) + + qg = q.clone().requires_grad_(True) + kg = k.clone().requires_grad_(True) + vg = v.clone().requires_grad_(True) + o_auto = torch.nn.functional.scaled_dot_product_attention( + qg, kg, vg, is_causal=causal, scale=sm_scale, enable_gqa=n_rep > 1) + o_auto.backward(do) + + o_ref, lse_ref = mxfp8_attention_forward_reference(q, k, v, sm_scale, causal) + dq, dk, dv = mxfp8_attention_backward_reference_stage2(q, k, v, do, o_ref, lse_ref, sm_scale, causal) + + rows = [] + for name, ref, got in [("dQ", qg.grad, dq), ("dK", kg.grad, dk), ("dV", vg.grad, dv)]: + rows.append([f"{name} (stage2 vs sdpa)", calc_snr(ref, got), calc_cossim(ref, got)]) + print() + print(tabulate(rows, headers=["Tensor", "SNR", "Cosine Sim"], tablefmt="github")) + + # A1 has no double-quant, so it should land at/above the retired option-A + # baseline (cossim~0.998 / SNR~23-26 dB). Thresholds keep margin. + for name, ref, got in [("dQ", qg.grad, dq), ("dK", kg.grad, dk), ("dV", vg.grad, dv)]: + sim = calc_cossim(ref, got) + snr = calc_snr(ref, got) + assert sim > 0.995, f"{name} stage2-vs-sdpa cosine-sim too low: {sim}" + assert snr > 18, f"{name} stage2-vs-sdpa SNR too low: {snr}" + + +# --------------------------------------------------------------------------- +# Public autograd path — user-facing forward + backward wrapper +# --------------------------------------------------------------------------- + +public_autograd_cases = [ + AttnConfig(seqlen_q=128, seqlen_kv=128, num_head_q=4, num_head_kv=4, head_dim_qk=128, head_dim_v=128), + AttnConfig(seqlen_q=128, seqlen_kv=128, num_head_q=8, num_head_kv=2, head_dim_qk=128, head_dim_v=128), # GQA +] + + +@cuda_only +@pytest.mark.parametrize("config", public_autograd_cases) +@pytest.mark.parametrize("causal", [True, False]) +def test_public_autograd_matches_sdpa(config, causal): + """Public ``triton_attention_mxfp8`` forward/backward must wire gradients correctly. + + The lower-level backward op is tested across the large shape grid below. This + smaller test covers the user-facing autograd wrapper: saved tensors, ctx + metadata, backward op dispatch, and returned gradient positions. + """ + from alto.kernels.mxfp8.triton_flash_attention_mxfp8 import triton_attention_mxfp8 + + device = "cuda" + dtype = torch.bfloat16 + n_rep = config.num_head_q // config.num_head_kv + + q, k, v = _make_qkv_bhsd(1, config, device, dtype) + do = _make_do_bhsd(1, config, device, dtype) + sm_scale = config.head_dim_qk**(-0.5) + + q_ref = q.clone().detach().requires_grad_(True) + k_ref = k.clone().detach().requires_grad_(True) + v_ref = v.clone().detach().requires_grad_(True) + o_ref = torch.nn.functional.scaled_dot_product_attention( + q_ref, k_ref, v_ref, is_causal=causal, scale=sm_scale, enable_gqa=n_rep > 1) + o_ref.backward(do) + + q_kernel = q.clone().detach().requires_grad_(True) + k_kernel = k.clone().detach().requires_grad_(True) + v_kernel = v.clone().detach().requires_grad_(True) + o_kernel = triton_attention_mxfp8( + q_kernel.contiguous(), + k_kernel.contiguous(), + v_kernel.contiguous(), + bias=None, + alibi_slopes=None, + sm_scale=sm_scale, + dropout_p=0.0, + cu_seqlens_q=0, + cu_seqlens_k=0, + max_seqlens_q=config.seqlen_q, + max_seqlens_k=config.seqlen_kv, + causal=causal, + return_scores=False, + use_exp2=True, + layout="bhsd", + )[0] + o_kernel.backward(do) + + pairs = [ + ("O", o_ref, o_kernel, 0.99, 20), + ("dQ", q_ref.grad, q_kernel.grad, 0.995, 18), + ("dK", k_ref.grad, k_kernel.grad, 0.995, 18), + ("dV", v_ref.grad, v_kernel.grad, 0.995, 18), + ] + rows = [] + for name, ref, got, _, _ in pairs: + rows.append([f"{name} (public autograd vs sdpa)", calc_snr(ref, got), calc_cossim(ref, got)]) + print() + print(tabulate(rows, headers=["Tensor", "SNR", "Cosine Sim"], tablefmt="github")) + + for name, ref, got, sim_threshold, snr_threshold in pairs: + sim = calc_cossim(ref, got) + snr = calc_snr(ref, got) + assert sim > sim_threshold, f"{name} public-autograd-vs-sdpa cosine-sim too low: {sim}" + assert snr > snr_threshold, f"{name} public-autograd-vs-sdpa SNR too low: {snr}" + + +# --------------------------------------------------------------------------- +# Backward Layer 2 — kernel vs Stage-2 (A1) golden reference +# The A1 backward kernel runs tl.dot_scaled in-kernel, so this requires CDNA4 +# (like the forward Layer 2). Calls the backward op directly with the +# forward-reference o/lse (bypassing the forward kernel) and compares to the +# Stage-2 reference (same full quantization both sides -> gap = port bugs). +# --------------------------------------------------------------------------- + +def _check_backward_kernel_vs_reference(config, causal, batch=4): + from alto.kernels.mxfp8.mxfp8_quantization import convert_to_mxfp8 + + device = "cuda" + dtype = torch.bfloat16 + q, k, v = _make_qkv_bhsd(batch, config, device, dtype) + do = _make_do_bhsd(batch, config, device, dtype) + sm_scale = config.head_dim_qk**(-0.5) + + o_ref, lse_ref = mxfp8_attention_forward_reference(q, k, v, sm_scale, causal) + + # Quantize operands the way the forward autograd does, then run the backward op. + q8, q_scale = convert_to_mxfp8(q, mxfp_format="e4m3", axis=-1, is_2d_block=True) + k8, k_scale = convert_to_mxfp8(k, mxfp_format="e4m3", axis=-1, is_2d_block=True) + v8, v_scale = convert_to_mxfp8(v, mxfp_format="e4m3", axis=-1, is_2d_block=True) + + dq_k, dk_k, dv_k = torch.ops.alto.attention_mxfp8_backward_triton_impl( + do.contiguous(), + q8.contiguous(), + k8.contiguous(), + v8.contiguous(), + o_ref.contiguous(), + lse_ref.contiguous(), + q_scale, + k_scale, + v_scale, + sm_scale=sm_scale, + causal=causal, + layout="bhsd", + cu_seqlens_q=0, + cu_seqlens_k=0, + max_seqlen_q=config.seqlen_q, + max_seqlen_k=config.seqlen_kv, + use_exp2=True, + ) + dq_r, dk_r, dv_r = mxfp8_attention_backward_reference_stage2(q, k, v, do, o_ref, lse_ref, sm_scale, causal) + + rows = [] + pairs = [("dQ", dq_r, dq_k), ("dK", dk_r, dk_k), ("dV", dv_r, dv_k)] + for name, ref, got in pairs: + rows.append([f"{name} (kernel vs golden)", calc_snr(ref, got), calc_cossim(ref, got)]) + print() + print(tabulate(rows, headers=["Tensor", "SNR", "Cosine Sim"], tablefmt="github")) + + # Threshold placeholder — calibrate on the first CDNA4 run. + for name, ref, got in pairs: + sim = calc_cossim(ref, got) + snr = calc_snr(ref, got) + assert sim > 0.99, f"{name} kernel-vs-golden cosine-sim too low: {sim}" + assert snr > 25, f"{name} kernel-vs-golden SNR too low: {snr}" + + +@cuda_only +@pytest.mark.parametrize("config", test_cases) +@pytest.mark.parametrize("causal", [True, False]) +def test_backward_kernel_matches_reference(config, causal): + """A1 backward kernel vs Stage-2 golden reference — isolates port bugs. CDNA4-only.""" + _check_backward_kernel_vs_reference(config, causal) + + +@cuda_only +@pytest.mark.parametrize("config", non_causal_cases) +def test_backward_kernel_matches_reference_non_square(config): + """Backward kernel on seqlen_q != seqlen_k, the only shape family causal rejects.""" + _check_backward_kernel_vs_reference(config, causal=False, batch=2) diff --git a/tests/unittest/mxfp8/utils.py b/tests/unittest/mxfp8/utils.py index 0cc5039..750d876 100644 --- a/tests/unittest/mxfp8/utils.py +++ b/tests/unittest/mxfp8/utils.py @@ -174,3 +174,307 @@ def convert_from_mxfp8_pytorch( data_hp = data_hp.to(torch.bfloat16) return data_hp.transpose(axis, -1) + + +def _mxfp8_qdq(x: torch.Tensor, axis: int, is_2d_block: bool, block_size: int = BLOCK_SIZE_DEFAULT) -> torch.Tensor: + """Round-trip a tensor through pure-PyTorch MXFP8 e4m3 quantize+dequantize. + + Returns the value an mxfp8 kernel operand would actually see. + """ + lp, s = convert_to_mxfp8_pytorch(x, block_size=block_size, mxfp_format="e4m3", axis=axis, is_2d_block=is_2d_block) + return convert_from_mxfp8_pytorch(lp, s, output_dtype=x.dtype, block_size=block_size, axis=axis, + is_2d_block=is_2d_block) + + +def mxfp8_attention_forward_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + sm_scale: float, + causal: bool, + block_n: int = 64, + block_size: int = BLOCK_SIZE_DEFAULT, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Pure-PyTorch golden reference for the MXFP8 (e4m3) flash-attention forward. + + Faithfully replicates what ``triton_flash_attention_mxfp8.attn_fwd`` computes, + so a kernel-vs-this-reference gap isolates Triton port bugs (masking / + online-softmax / LSE) from mxfp8 quantization error itself: + + * Q/K/V are quantized 2D-block along head_dim (``is_2d_block=True``, + ``axis=-1``), matching the autograd forward. + * The softmax probabilities ``P`` are quantized **per key block** exactly + like the kernel: within the online-softmax loop, using the *running* row + max (not the global max), 1D-block along the key axis with + ``block_size`` groups. + * Contains no ``tl.dot_scaled``: every matmul is a plain fp32 matmul on the + dequantized operands, so this runs on CPU / any device (no CDNA4). + + Args: + q, k, v: ``[batch, nheads, seqlen, head_dim]`` (bhsd), fp32/bf16. GQA is + supported (``nheads_k`` may be < ``nheads_q``). + sm_scale: softmax scale applied to ``Q @ K^T``. + causal: top-left causal mask ``key_j <= query_i`` (matches + ``F.sdpa(is_causal=True)`` on bhsd tensors in PyTorch 2.x). + block_n: key-block width, must match the kernel ``BLOCK_N`` (default 64). + block_size: MXFP8 quant block, must match ``QUANT_BLOCK_SIZE`` (default 32). + + Returns: + (o, softmax_lse): output ``[batch, nheads_q, seqlen_q, head_dim_v]`` and + log-sum-exp ``[batch, nheads_q, seqlen_q]``. + """ + assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4 + b, hq, sq, dqk = q.shape + _, hk, sk, _ = k.shape + dv = v.shape[-1] + assert hq % hk == 0, f"nheads_q ({hq}) must be a multiple of nheads_k ({hk})" + n_rep = hq // hk + + # Quantize Q/K/V (2D block along head_dim) then dequantize -> operands the kernel sees. + q_dq = _mxfp8_qdq(q, axis=-1, is_2d_block=True).to(torch.float32) + k_dq = _mxfp8_qdq(k, axis=-1, is_2d_block=True).to(torch.float32) + v_dq = _mxfp8_qdq(v, axis=-1, is_2d_block=True).to(torch.float32) + + # Expand K/V heads for GQA (head_q h -> head_k h // n_rep, matching off_h_k). + if n_rep > 1: + k_dq = k_dq.repeat_interleave(n_rep, dim=1) + v_dq = v_dq.repeat_interleave(n_rep, dim=1) + + device = q.device + # Real -inf (not finfo.min): masked scores must dequantize to exp(-inf)=0. + # A finite sentinel would make a fully-masked block compute exp(0)=1. + neg_inf = float("-inf") + + m_i = torch.full((b, hq, sq), float("-inf"), dtype=torch.float32, device=device) + l_i = torch.zeros((b, hq, sq), dtype=torch.float32, device=device) + acc = torch.zeros((b, hq, sq, dv), dtype=torch.float32, device=device) + + q_pos = torch.arange(sq, device=device) + + for j0 in range(0, sk, block_n): + j1 = min(j0 + block_n, sk) + k_blk = k_dq[:, :, j0:j1, :] # [b, hq, bn, d] + v_blk = v_dq[:, :, j0:j1, :] + + s = torch.matmul(q_dq, k_blk.transpose(-1, -2)) * sm_scale # [b, hq, sq, bn] + + if causal: + key_pos = torch.arange(j0, j1, device=device) + allowed = key_pos[None, :] <= q_pos[:, None] # [sq, bn] + s = torch.where(allowed[None, None, :, :], s, torch.full_like(s, neg_inf)) + + m_ij = torch.maximum(m_i, s.max(dim=-1).values) # [b, hq, sq] + p = torch.exp(s - m_ij[..., None]) # unnormalized, running max + # Rows fully masked in this block yield exp(neg_inf - (-inf)) issues; zero them. + p = torch.nan_to_num(p, nan=0.0, posinf=0.0, neginf=0.0) + l_ij = p.sum(dim=-1) # [b, hq, sq] + + # Quantize P per key block exactly like the kernel (1D block along key axis). + p_dq = _mxfp8_qdq(p.to(torch.float32), axis=-1, is_2d_block=False, block_size=block_size) + + alpha = torch.exp(m_i - m_ij) + alpha = torch.nan_to_num(alpha, nan=0.0, posinf=0.0, neginf=0.0) + acc = acc * alpha[..., None] + torch.matmul(p_dq, v_blk) + l_i = l_i * alpha + l_ij + m_i = m_ij + + l_safe = torch.where(l_i > 0, l_i, torch.ones_like(l_i)) + o = acc / l_safe[..., None] + softmax_lse = m_i + torch.log(l_safe) + # Fully-masked query rows (can happen with causal when sk < sq): zero output. + fully_masked = ~torch.isfinite(m_i) | (l_i <= 0) + o = torch.where(fully_masked[..., None], torch.zeros_like(o), o) + softmax_lse = torch.where(fully_masked, torch.zeros_like(softmax_lse), softmax_lse) + + return o.to(q.dtype), softmax_lse + + +def mxfp8_attention_backward_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + do: torch.Tensor, + o: torch.Tensor, + softmax_lse: torch.Tensor, + sm_scale: float, + causal: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pure-PyTorch golden reference for the MXFP8 (e4m3) flash-attention backward. + + Mirrors the Stage-1 backward kernel exactly so a kernel-vs-this-reference gap + isolates Triton port bugs from mxfp8 quantization error: + + * Q/K/V are quantized 2D-block along head_dim then dequantized — the same + operands the (Stage-1) kernel runs on. + * The softmax probabilities ``P`` are recomputed in fp32 from the *saved* + ``softmax_lse`` (``P = exp(S - lse)``); backward does **not** re-quantize + ``P`` (matching the kernel). + * ``o`` / ``softmax_lse`` must be the values the forward produced (feed the + outputs of ``mxfp8_attention_forward_reference``), so ``delta = rowsum(o * + do)`` and ``P`` are consistent with the forward pass. + + Uses the same top-left causal mask as the forward reference, so causal tests + must keep ``seqlen_q == seqlen_k`` (top-left == bottom-right there). + + Standard FlashAttention-v2 backward:: + + dV = Pᵀ @ dO + dP = dO @ Vᵀ + delta = rowsum(O * dO) + dS = P * (dP - delta) + dQ = sm_scale * dS @ K + dK = sm_scale * dSᵀ @ Q + + Args: + q, k, v: ``[batch, nheads, seqlen, head_dim]`` (bhsd). GQA supported. + do: upstream gradient, same shape/layout as ``o``. + o, softmax_lse: forward outputs (see above). + sm_scale: softmax scale. + causal: top-left causal mask. + + Returns: + (dq, dk, dv) matching the shapes of q, k, v. + """ + b, hq, sq, dqk = q.shape + _, hk, sk, dvv = v.shape + assert hq % hk == 0, f"nheads_q ({hq}) must be a multiple of nheads_k ({hk})" + n_rep = hq // hk + + q_dq = _mxfp8_qdq(q, axis=-1, is_2d_block=True).to(torch.float32) + k_dq = _mxfp8_qdq(k, axis=-1, is_2d_block=True).to(torch.float32) + v_dq = _mxfp8_qdq(v, axis=-1, is_2d_block=True).to(torch.float32) + if n_rep > 1: + k_dq = k_dq.repeat_interleave(n_rep, dim=1) + v_dq = v_dq.repeat_interleave(n_rep, dim=1) + + do_f = do.to(torch.float32) + o_f = o.to(torch.float32) + lse = softmax_lse.to(torch.float32) + + s = torch.matmul(q_dq, k_dq.transpose(-1, -2)) * sm_scale # [b, hq, sq, sk] + if causal: + q_pos = torch.arange(sq, device=q.device) + k_pos = torch.arange(sk, device=q.device) + allowed = k_pos[None, :] <= q_pos[:, None] # [sq, sk], top-left + s = torch.where(allowed[None, None, :, :], s, torch.full_like(s, float("-inf"))) + + p = torch.exp(s - lse[..., None]) # softmax probabilities + p = torch.nan_to_num(p, nan=0.0, posinf=0.0, neginf=0.0) + + delta = (o_f * do_f).sum(dim=-1) # [b, hq, sq] + dp = torch.matmul(do_f, v_dq.transpose(-1, -2)) # [b, hq, sq, sk] + ds = p * (dp - delta[..., None]) # [b, hq, sq, sk] + + dv_full = torch.matmul(p.transpose(-1, -2), do_f) # [b, hq, sk, dv] + dq = sm_scale * torch.matmul(ds, k_dq) # [b, hq, sq, dqk] + dk_full = sm_scale * torch.matmul(ds.transpose(-1, -2), q_dq) # [b, hq, sk, dqk] + + if n_rep > 1: + dk = dk_full.view(b, hk, n_rep, sk, dqk).sum(dim=2) + dv = dv_full.view(b, hk, n_rep, sk, dvv).sum(dim=2) + else: + dk = dk_full + dv = dv_full + + return dq, dk, dv + + +def mxfp8_attention_backward_reference_stage2( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + do: torch.Tensor, + o: torch.Tensor, + softmax_lse: torch.Tensor, + sm_scale: float, + causal: bool, + block_size: int = BLOCK_SIZE_DEFAULT, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pure-PyTorch golden reference for the **Stage-2 (A1)** MXFP8 backward. + + Where Stage-1 (``mxfp8_attention_backward_reference``) runs every backward + matmul in fp32 on the dequantized Q/K/V, Stage-2 additionally quantizes the + backward-only tensors (dO / P / dS) along each dot's reduction axis — a + numerical model of what ``tl.dot_scaled`` does in the kernel. A kernel-vs-this + gap isolates Triton port bugs from mxfp8 error; a this-vs-bf16-SDPA gap + measures the full mxfp8 backward quantization error. No ``tl.dot_scaled`` + here (plain fp32 matmul on dequantized operands), so it runs on CPU / MI250. + + **A1 operand source (2026-07-16 decision):** the kernel reuses the forward's + saved e4m3 q/k/v + their 2D-block scale directly in *every* dot — including + the seqlen-reduction dots — with no re-quantization. So this reference uses a + single 2D-block dequant ``q_dq``/``k_dq``/``v_dq`` everywhere (no double + quantization along seqlen, unlike the retired option-A reference). Only + dO/P/dS are quantized fresh (1D per-row along their reduction axis), since the + forward never saw them. + + Quantization axis per dot (plan §9.5):: + + S = Q @ Kᵀ reduction head_dim : Q,K = 2D-block dequant (reuse) + dP = dO @ Vᵀ reduction head_dim_v : dO 1D axis=-1 ; V 2D reuse + dV = Pᵀ @ dO reduction seqlen_q : P,dO 1D along sq + dK = dSᵀ @ Q reduction seqlen_q : dS 1D along sq ; Q 2D reuse + dQ = dS @ K reduction seqlen_k : dS 1D along sk ; K 2D reuse + + Same top-left causal mask as the other references (matches SDPA on bhsd), so + causal tests must keep ``seqlen_q == seqlen_k``. Args mirror + ``mxfp8_attention_backward_reference``; ``o`` / ``softmax_lse`` must be the + forward-reference outputs. Returns ``(dq, dk, dv)``. + """ + b, hq, sq, dqk = q.shape + _, hk, sk, dvv = v.shape + assert hq % hk == 0, f"nheads_q ({hq}) must be a multiple of nheads_k ({hk})" + n_rep = hq // hk + + def qdq(x, axis): + return _mxfp8_qdq(x.to(torch.float32), axis=axis, is_2d_block=False, block_size=block_size) + + # q/k/v: single 2D-block dequant, reused in every dot (A1, no re-quant). + q_dq = _mxfp8_qdq(q, axis=-1, is_2d_block=True, block_size=block_size).to(torch.float32) + k_dq = _mxfp8_qdq(k, axis=-1, is_2d_block=True, block_size=block_size).to(torch.float32) + v_dq = _mxfp8_qdq(v, axis=-1, is_2d_block=True, block_size=block_size).to(torch.float32) + if n_rep > 1: + k_dq = k_dq.repeat_interleave(n_rep, dim=1) + v_dq = v_dq.repeat_interleave(n_rep, dim=1) + + do_f = do.to(torch.float32) + o_f = o.to(torch.float32) + lse = softmax_lse.to(torch.float32) + + # S -> P (recompute from saved lse; P feeding dS is elementwise, not quantized). + s = torch.matmul(q_dq, k_dq.transpose(-1, -2)) * sm_scale + if causal: + q_pos = torch.arange(sq, device=q.device) + k_pos = torch.arange(sk, device=q.device) + allowed = k_pos[None, :] <= q_pos[:, None] # top-left + s = torch.where(allowed[None, None, :, :], s, torch.full_like(s, float("-inf"))) + p = torch.exp(s - lse[..., None]) + p = torch.nan_to_num(p, nan=0.0, posinf=0.0, neginf=0.0) + + # dP = dO @ Vᵀ (reduction head_dim_v): dO 1D along head_dim_v, V 2D reuse. + do_hd = qdq(do_f, axis=-1) + dp = torch.matmul(do_hd, v_dq.transpose(-1, -2)) + delta = (o_f * do_f).sum(dim=-1) + ds = p * (dp - delta[..., None]) # fp32 elementwise + + # dV = Pᵀ @ dO (reduction seqlen_q): P and dO 1D along sq. + p_sq = qdq(p, axis=-2) + do_sq = qdq(do_f, axis=-2) + dv_full = torch.matmul(p_sq.transpose(-1, -2), do_sq) + + # dQ = sm_scale * dS @ K (reduction seqlen_k): dS 1D along sk, K 2D reuse. + ds_sk = qdq(ds, axis=-1) + dq = sm_scale * torch.matmul(ds_sk, k_dq) + + # dK = sm_scale * dSᵀ @ Q (reduction seqlen_q): dS 1D along sq, Q 2D reuse. + ds_sq = qdq(ds, axis=-2) + dk_full = sm_scale * torch.matmul(ds_sq.transpose(-1, -2), q_dq) + + if n_rep > 1: + dk = dk_full.view(b, hk, n_rep, sk, dqk).sum(dim=2) + dv = dv_full.view(b, hk, n_rep, sk, dvv).sum(dim=2) + else: + dk = dk_full + dv = dv_full + + return dq, dk, dv