前言

在上篇文章 斯坦福大学 | CS336 | 从零开始构建语言模型 | Spring 2025 | 笔记 | Assignment 2: FlashAttention-2 中，我们已经了解了 FlashAttention-2 的作业要求，下面我们就一起来看看这些作业该如何实现，本篇文章记录 CS336 作业 Assignment 2: Systems 中的 FlashAttention-2 实现，仅供自己参考😄

Note：博主并未遵循 from-scratch 的宗旨，所有代码几乎均由 ChatGPT 完成

Assignment 2：https://github.com/stanford-cs336/assignment2-systems

reference：https://chatgpt.com/

1. Problem (pytorch_attention): 2 points

(a) 在不同规模下对你的注意力实现进行基准测试，请编写一个脚本，完成以下工作：

1. 将 batch size 固定为 8，并且 不使用多头注意力（即去掉 head 这一维度）

2. 对以下参数组合进行遍历（笛卡尔积）：

• head 的嵌入维度 $d_{\text{model}}\in [16,32,64,128]$
• 序列长度 $\in [256,1024,4096,8192,16384]$

3. 为对应尺寸生成随机输入 $Q,K,V$

4. 使用这些输入对注意力模块进行 100 次前向传播 并计时

5. 在反向传播开始之前，测量当前的显存使用情况，并对 100 次反向传播 进行计时

6. 确保在正式计时前进行 warm-up，并且在每一次前向 / 反向传播之前调用 torch.cuda.synchronize()

请报告在这些配置下得到的运行时间（或是否发生了显存溢出错误），在哪些规模下你会遇到 out-of-memory（OOM） 错误？

请对你发现的 最小一个发生 OOM 的配置，对注意力模块的显存使用进行理论分析（你可以使用 Assignment 1 中给出的 Transformer 显存占用公式）。反向传播所需的显存节省量会如何随序列长度变化？如果要 彻底消除这部分显存开销，你会采取什么方法？

Deliverable：一张包含运行时间的表格；你对注意力显存使用的推导计算以及一段 1-2 段的文字分析说明。

bench_pytorch_attention.py 脚本实现如下：

import torch
import argparse
from typing import Callable, List
from pathlib import Path

from cs336_basics.modules import scaled_dot_product_attention
from cs336_systems.utils import AttentionRow, AttentionBenchmarkReporter


def cuda_sync():
    torch.cuda.synchronize()


def causal_mask(seq_len: int, device: torch.device) -> torch.Tensor:
    # True = keep, False = masked out
    return torch.tril(torch.ones((seq_len, seq_len), device=device, dtype=torch.bool))


def time_forward(
    fn: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor],
    q: torch.Tensor, k: torch. Tensor, v: torch.Tensor,
    iters: int
) -> float:
    # Use CUDA events
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    cuda_sync()
    start.record()
    for _ in range(iters):
        cuda_sync()
        _ = fn(q, k, v)
        cuda_sync()
    end.record()
    cuda_sync()
    return start.elapsed_time(end) / iters


def time_backward(
    fn: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor],
    q: torch.Tensor, k: torch. Tensor, v: torch.Tensor,
    iters: int
) -> float:
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    total_ms = 0.0
    for _ in range(iters):
        cuda_sync()
        out = fn(q, k, v)
        loss = out.sum()
        cuda_sync()

        start.record()
        loss.backward()
        end.record()
        cuda_sync()
        total_ms += start.elapsed_time(end)

        # clear grads for next iter
        q.grad = None
        k.grad = None
        v.grad = None

    return total_ms / iters


def run_one(
    d_model: int, seq_len: int, batch: int, dtype: torch.dtype,
    warmup: int, iters: int, use_causal_mask: bool
) -> AttentionRow:
    device = torch.device("cuda")

    try:
        q = torch.randn(batch, seq_len, d_model, device=device, dtype=dtype, requires_grad=True)
        k = torch.randn(batch, seq_len, d_model, device=device, dtype=dtype, requires_grad=True)
        v = torch.randn(batch, seq_len, d_model, device=device, dtype=dtype, requires_grad=True)

        mask = causal_mask(seq_len, device) if use_causal_mask else None

        def attn(q, k, v):
            return scaled_dot_product_attention(q, k, v, mask=mask)

        # warmup
        for _ in range(warmup):
            cuda_sync()
            out = attn(q, k, v)
            cuda_sync()
            out.sum().backward()
            cuda_sync()
            q.grad = None
            k.grad = None
            v.grad = None
        
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
        cuda_sync()

        fwd_ms = time_forward(attn, q, k, v, iters)

        # memory snapshot
        mem_before_bwd_mb = torch.cuda.memory_allocated() / (1024 ** 2)

        bwd_ms = time_backward(attn, q, k, v, iters)

        return AttentionRow(d_model, seq_len, fwd_ms, bwd_ms, mem_before_bwd_mb, "ok")

    except RuntimeError as e:
        msg =  str(e).lower()
        if "out of memory" in msg:
            # Important: release references, then clear cache
            try:
                del q, k, v
            except Exception:
                pass
            try:
                del out, loss
            except Exception:
                pass
            torch.cuda.empty_cache()
            return AttentionRow(d_model, seq_len, None, None, None, "oom")
        return AttentionRow(d_model, seq_len, None, None, None, f"error:{type(e).__name__}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch", type=int, default=8)
    parser.add_argument("--warmup", type=int, default=10)
    parser.add_argument("--iters", type=int, default=100)
    parser.add_argument("--dtype", type=str, default="float32", choices=[ "float32", "float16", "bfloat16"])
    parser.add_argument("--no-causal", action="store_true", help="Disable causal mask (default: causal enabled)")
    parser.add_argument("--out-dir", type=str, default="runs/pytorch_attention")

    args = parser.parse_args()

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for this benchmark.")

    dtype_map = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}
    dtype = dtype_map[args.dtype]

    out_dir = Path(args.out_dir)
    reporter = AttentionBenchmarkReporter(
        jsonl_path=out_dir / "metrics.jsonl",
        md_path=out_dir / "table.md",
        title="#### Problem (pytorch_attention): naive PyTorch attention benchmark",
    )

    d_models = [16, 32, 64, 128]
    seq_lens = [256, 1024, 4096, 8192, 16384]

    rows: List[AttentionRow] = []
    for d in d_models:
        for s in seq_lens:
            torch.cuda.empty_cache()
            r = run_one(
                d_model=d,
                seq_len=s,
                batch=args.batch,
                dtype=dtype,
                warmup=args.warmup,
                iters=args.iters,
                use_causal_mask=(not args.no_causal),
            )
            print(f"d={d:4d}, s={s:5d} -> {r.status}"
                  + ("" if r.fwd_ms is None else f", fwd={r.fwd_ms:.3f}ms, bwd={r.bwd_ms:.3f}ms, mem={r.mem_before_bwd_mb:.1f}MB"))
            rows.append(r)
            reporter.append(r)

    reporter.write_markdown()


if __name__ == "__main__":
    main()

运行指令如下：

uv run cs336_systems/bench_pytorch_attention.py

执行运行指令后输出如下：

整个代码实现就是作业要求的使用不同嵌入维度、序列长度组合对注意力模块的前向传播和反向传播计时

代码实现比较简单，大家可以自己看看

Note：由于后续需要进行其他测试，因此 bench_pytorch_attention.py 的实现会不断修改，最新的代码请参考：https://github.com/Melody-Zhou/stanford-cs336-spring2025-assignments

测量结果如下所示：

d_model	seq_len	fwd_ms	bwd_ms	mem_before_bwd_mb	status
16	256	0.231	0.450	16.812	ok
16	1024	2.020	4.990	19.250	ok
16	4096	30.149	73.746	40.250	ok
16	8192	-	-	-	oom
16	16384	-	-	-	oom
32	256	0.195	0.500	17.312	ok
32	1024	2.043	5.023	21.250	ok
32	4096	30.603	74.197	48.250	ok
32	8192	-	-	-	oom
32	16384	-	-	-	oom
64	256	0.204	0.487	18.312	ok
64	1024	2.086	5.060	25.250	ok
64	4096	31.429	75.045	64.250	ok
64	8192	-	-	-	oom
64	16384	-	-	-	oom
128	256	0.213	0.517	20.312	ok
128	1024	2.216	5.234	33.250	ok
128	4096	33.046	76.831	96.250	ok
128	8192	-	-	-	oom
128	16384	-	-	-	oom

Note：这是博主在自己主机 RTX3060（12GB）上测量得到的结果

从表中我们可以观察到，当 序列长度达到 8192 及以上时，所有配置均发生 OOM（out-of-memory）错误，且该现象与嵌入维度大小无关，因此，最小发生 OOM 的配置为：$(B,S,d_{\text{model}})=(8,8192,16)$

在朴素的 scaled dot-product attention 实现中，显存开销的主要来源是显式构造并保存注意力矩阵 $L=Q{K}^{\top }\in {\mathbb{R}}^{B\times S\times S}$ 以及对应的 softmax 概率矩阵，其形状同样为 $B\times S\times S$

对于最小 OOM 配置 $B=8,S=8192$，注意力矩阵的元素数量为 $B\cdot S^2=2\times 8192^2\approx 5.37\times 10^8$，且该矩阵以 FP32（4 bytes） 形式存储，因此单个 $B\times S\times S$ 张量的显存占用约为 $5.37\times 10^8\times 4\approx 2.0\text{ GB}$

在训练过程中，前向传播阶段需要保留注意力 logits 或 softmax 输出以供反向传播使用，反向传播还会引入与其同阶的中间梯度或临时张量，因此，仅注意力相关的中间结果就可能占用多个 $B\times S\times S$ 级别的张量，使得总显存开销呈现出 $\Theta (B\cdot S^2)$ 的增长趋势，当 $S=8192$ 时，该二次增长的显存需求迅速超过 12GB 显存限制，从而触发 OOM 错误

反向传播阶段中，可节省的激活显存规模同样与 $S^2$ 成正比，因此序列长度越大，潜在的显存优化空间越大。若要彻底消除这部分显存开销，必须避免显式物化完整的注意力矩阵，一种有效的方法是采用 内存高效注意力算法（如 FlashAttention），通过分块计算和在线 softmax，将显存复杂度从 $\Theta (B\cdot S^2)$ 降低至近似 $\Theta (B\cdot S\cdot d)$，在保持数值等价的同时显著提升可扩展性。

2. Problem (torch_compile): 2 points

(a) 扩展你的 attention 基准测试脚本，使其包含 PyTorch attention 实现的编译版本，并在与上面 pytorch_attention 问题相同的配置下，将其性能与 未编译版本 进行对比

Deliverable：一张表格，对比编译版 attention 模块与 pytorch_attention 问题中未编译版本在前向和反向传播上的耗时。

编译版 attention 模块测试运行指令如下：

uv run cs336_systems/bench_pytorch_attention.py --compile

执行运行指令后输出如下：

编译版本和未编译版本的性能对比如下所示：

d_model	seq_len	fwd_ms	bwd_ms	mem_before_bwd_mb	status	impl
16	256	0.105	0.202	16.812	ok	compiled
16	256	0.231	0.450	16.812	ok	eager
16	1024	0.736	1.563	19.250	ok	compiled
16	1024	2.020	4.990	19.250	ok	eager
16	4096	11.319	27.147	40.250	ok	compiled
16	4096	30.149	73.746	40.250	ok	eager
16	8192	52.370	109.351	96.250	ok	compiled
16	8192	-	-	-	oom	eager
16	16384	-	-	-	oom	compiled
16	16384	-	-	-	oom	eager
32	256	0.140	0.215	17.312	ok	compiled
32	256	0.195	0.500	17.312	ok	eager
32	1024	0.851	1.956	21.250	ok	compiled
32	1024	2.043	5.023	21.250	ok	eager
32	4096	13.905	28.484	48.250	ok	compiled
32	4096	30.603	74.197	48.250	ok	eager
32	8192	55.870	112.285	112.250	ok	compiled
32	8192	-	-	-	oom	eager
32	16384	-	-	-	oom	compiled
32	16384	-	-	-	oom	eager
64	256	0.315	0.305	18.312	ok	compiled
64	256	0.204	0.487	18.312	ok	eager
64	1024	1.337	1.758	25.250	ok	compiled
64	1024	2.086	5.060	25.250	ok	eager
64	4096	11.201	28.288	64.250	ok	compiled
64	4096	31.429	75.045	64.250	ok	eager
64	8192	47.166	115.708	144.250	ok	compiled
64	8192	-	-	-	oom	eager
64	16384	-	-	-	oom	compiled
64	16384	-	-	-	oom	eager
128	256	0.336	0.342	20.312	ok	compiled
128	256	0.213	0.517	20.312	ok	eager
128	1024	1.501	1.928	33.250	ok	compiled
128	1024	2.216	5.234	33.250	ok	eager
128	4096	13.177	30.452	96.250	ok	compiled
128	4096	33.046	76.831	96.250	ok	eager
128	8192	53.704	122.807	208.250	ok	compiled
128	8192	-	-	-	oom	eager
128	16384	-	-	-	oom	compiled
128	16384	-	-	-	oom	eager

从上面的结果中我们可以观察到，torch.compile 在所有未发生 OOM 的配置先均显著提升了运行性能，且这种加速效果随着序列长度的增大而更加明显

• 在 小序列长度（S=256） 下，compiled 模式的前向和反向传播时间相较 eager 模式约有 1.5x-2x 的加速，但由于计算规模较小，绝对时间差异有限
• 当 序列长度增大到 1024 和 4096 时，compiled 模式在前向和反向传播上均取得了 约 2x-3x 的速度提升，尤其在反向传播阶段收益更为明显
• 在 S=8192 时，eager 模式已经因显存不足而无法运行，而 compiled 模式仍可完成前向与反向传播，表明 torch.compile 在一定程度上 推迟了 OOM 的出现

需要注意的是，在所有可运行的配置中，compiled 模式与 eager 模式在反向传播前的显存占用几乎一致，这表明 torch.compile 的主要收益来自于 算子融合、调度优化以及 kernel 级别的性能改进，而并未从根本上改变 attention 的显存复杂度。

(b) 接下来，在你的 端到端基准测试脚本 中，对 整个 Transformer 模型 进行编译。前向传播的性能发生了怎样的变化？前向 + 反向传播以及优化器 step 的组合性能又如何变化？

Deliverable：一张表格，对比原始（vanilla）Transformer 模型与编译后的 Transformer 模型的性能。

编译版 transformer 模块测试运行指令如下：

uv run python cs336_systems/benchmark.py \
    --out-jsonl runs/bench.jsonl \
    --out-md runs/bench.md \
    --write-md \
    --sweep \
    --sweep-contexts 128 \
    --compile

执行运行指令后输出如下：

编译版本和未编译版本的性能对比如下所示：

specs	model_size	context_length	batch_size	impl	mode	mean_ms	std_ms	tok_per_s
RTX3060	small	128	4	compiled	forward	3.896	0.016	131428
RTX3060	small	128	4	eager	forward	5.476	0.071	93492.8
RTX3060	small	128	4	compiled	backward	8.166	0.026	62696.8
RTX3060	small	128	4	eager	backward	11.768	0.258	43509.4
RTX3060	medium	128	4	compiled	forward	7.469	0.024	68548.5
RTX3060	medium	128	4	eager	forward	9.641	0.159	53104.2
RTX3060	medium	128	4	compiled	backward	15.446	0.032	33147.7
RTX3060	medium	128	4	eager	backward	21.295	0.164	24043.4
RTX3060	large	128	4	compiled	forward	12.81	0.026	39967.9
RTX3060	large	128	4	eager	forward	16.012	0.075	31977
RTX3060	large	128	4	compiled	backward	28.149	0.037	18189
RTX3060	large	128	4	eager	backward	37.124	0.165	13791.7
RTX3060	xl	128	4	compiled	forward	20.394	0.401	25105.3
RTX3060	xl	128	4	eager	forward	24.802	0.085	20643.6
RTX3060	xl	128	4	compiled	backward	44.694	0.651	11455.7
RTX3060	xl	128	4	eager	backward	57.553	0.259	8896.22
RTX3060	2.7b	128	4	compiled	forward	26.762	0.278	19132
RTX3060	2.7b	128	4	eager	forward	32.77	0.23	15624
RTX3060	2.7b	128	4	compiled	backward	58.605	0.274	8736.42
RTX3060	2.7b	128	4	eager	backward	75.727	0.215	6761.11

从上面的对比结果中我们可以发现 torch.compile 能够稳定提升 Transformer 端到端执行性能，由于我们这里固定了 context_length，因此在不同模型规模下 torch.compile 几乎表现出一致性的加速效果。

3. Problem (flash_forward): 15 points

3.1 基础实现

(a) 编写一个 纯 PyTorch（不使用 Triton） 的 autograd.Function，用于实现 FlashAttention-2 的前向传播，该实现会比常规的 PyTorch 注意力实现慢得多，但它将有助于你调试后续的 Triton kernel

你的实现应当接收输入 $\mathbf{Q},\mathbf{K},\mathbf{V}$ 以及一个标志位 is_causal，并输出结果 $\mathbf{O}$ 以及 logsumexp 值 $L$，在本题中，你可以忽略 is_causal 标志。autograd.Function 的 forward 方法随后应当通过 save_for_backward 保存 $L,\mathbf{Q},\mathbf{K},\mathbf{V},\mathbf{O}$，以供反向传播阶段使用

请注意，autograd.Function 的 forward 方法始终将 context（ctx） 作为第一个参数，任何 autograd.Function 类都需要实现一个 backward 方法，不过在当前阶段你可以让它直接抛出 NotImplementedError。如果你需要一个对照实现，可以在 PyTorch 中实现公式 (4) 到 (6) 以及 (12)，并将其输出与你的实现进行比较

该接口定义为：

def forward(ctx, Q, K, V, is_causal=False)

tile 的尺寸可以由你自行决定，但请确保 至少为 16x16，我们在测试中始终使用 维度为 2 的整数次幂且不小于 16 的输入，因此你无需担心越界访问的问题

Deliverable：一个 torch.autograd.Function 的子类，实现 FlashAttention-2 的前向传播，为了测试你的代码，请实现 [adapters.get_flashattention_autograd_function_pytorch]，然后运行：

uv run pytest -k test_flash_forward_pass_pytorch

并确保你的实现能够通过测试。

代码实现如下：

import math
import torch


class FlashAttention2Pytorch(torch.autograd.Function):
    @staticmethod
    def forward(ctx, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, is_causal: bool = False):
        # q: (B, Q, D), k/v: (B, K, D)
        if q.dim() < 3:
            raise ValueError("q must have shape (..., Q, D)")
        if k.dim() != q.dim() or v.dim() != q.dim():
            raise ValueError("q/k/v must have same rank")
        if q.shape[:-2] != k.shape[:-2] or q.shape[:-2] != v.shape[:-2]:
            raise ValueError("leading dims of q/k/v must match")
        if k.shape[-2] != v.shape[-2] or k.shape[-1] != v.shape[-1]:
            raise ValueError("k and v must have sanme shape in last 2 dims")
        if q.shape[-1] != k.shape[-1]:
            raise ValueError("q and k must have same D")
        
        B, Q, D = q.shape
        K = k.shape[-2]        

        # ignore is_causal
        scale = 1.0 / math.sqrt(D)

        # choose tile sizes (>=16)
        Bq = 32
        Bk = 32

        # output buffers
        o = torch.empty((B, Q, D), device=q.device, dtype=q.dtype)
        L = torch.empty((B, Q), device=q.device, dtype=torch.float32)

        for i in range(0, Q, Bq):
            q_i = q[:, i : i + Bq, :]  # (B, Bq, D)

            # running stats for this query tile
            m = torch.full((B, Bq), -float("inf"), device=q.device, dtype=torch.float32)
            l = torch.zeros((B, Bq), device=q.device, dtype=torch.float32)
            o_acc = torch.zeros((B, Bq, D), device=q.device, dtype=torch.float32)

            for j in range(0, K, Bk):
                k_j = k[:, j : j + Bk, :]   # (B, Bk, D)
                v_j = v[:, j : j + Bk, :]   # (B, Bk, D)

                # S = q @ k^T * scale -> (B, Bq, Bk)
                S = torch.matmul(q_i, k_j.transpose(-1, -2)) * scale

                # online softmax update
                m_new = torch.maximum(m, S.max(dim=-1).values)   # (B, Bq)

                # exp(S - m_new)
                P_tilde = torch.exp(S - m_new.unsqueeze(-1))   # (B, Bq, Bk)

                # l_new = exp(m - m_new) * l + rowsum(P_tilde)
                l_new = torch.exp(m - m_new) * l + P_tilde.sum(dim=-1)   # (B, Bq)

                # o_acc = exp(m - m_new) * o_acc + P_tilde @ v_j
                o_acc = (torch.exp(m - m_new)).unsqueeze(-1) * o_acc + torch.matmul(P_tilde, v_j)

                m, l = m_new, l_new
            
            # finalize: O = o_acc / l
            o_i = o_acc / l.unsqueeze(-1)
            o[:, i : i + Bq, :] = o_i.to(dtype=q.dtype)

            # L = logsumexp(S_row) = m + log(l)
            L[:, i : i + Bq] = m + torch.log(l)

        # save tensor for later backward stage
        ctx.save_for_backward(L, q, k, v, o)
        ctx.is_causal = is_causal

        # return output
        return o
    
    @staticmethod
    def backward(ctx, *grad_outputs):
        raise NotImplementedError

测试适配器 [adapters.get_flashattention_autograd_function_pytorch] 的实现如下：

def get_flashattention_autograd_function_pytorch() -> Type:
    """
    Returns a torch.autograd.Function subclass that implements FlashAttention2.
    The expectation is that this class will implement FlashAttention2
    using only standard PyTorch operations (no Triton!).

    Returns:
        A class object (not an instance of the class)
    """
    # For example: return MyFlashAttnAutogradFunctionClass
    from cs336_systems.flash_pytorch import FlashAttention2Pytorch
    return FlashAttention2Pytorch

执行 uv run pytest -k test_flash_forward_pass_pytorch 后输出如下：

整个代码的实现非常简单，就是将下面的 Algorithm 1 用代码进行实现。上面的 forward 函数就是在 PyTorch 里严格复现了 FlashAttention 的 online softmax 分块扫描：通过维护 (m, l, o_acc)，实现了数值稳定的 softmax 累积，并在不构造完整 attention 矩阵的情况下得到输出 o 和 pre-row 的 logsumexp L

Note：关于 FlashAttention 的原理讲解，大家感兴趣的可以看看：从Online Softmax到FlashAttention、Flash Attention原理讲解，这里博主就不再赘述了。

(b) 接下来，请按照 Algorithm 1 编写一个 Triton kernel，用于实现 FlashAttention-2 的前向传播。随后，请再编写一个继承自 torch.autograd.Function 的子类，在其 forward 方法中调用你刚刚实现的 融合（fused）Triton kernel，而不是再用 PyTorch 逐步计算结果

下面是一些针对该问题的调试与实现建议：

• 为了便于调试，我们建议将你在 Triton 中执行的每一步操作结果，与 (a) 部分中你实现的 tiled Pytorch 版本 逐一进行对比

• kernel 的 launch grid 应设置为 $(T_q,batch\_size)$，这意味着每一个 Triton program instance 只会处理 一个 batch 索引，并且只会读取和写入 一个 query tile 中对应的 $\mathbf{Q},\mathbf{O}$ 和 $L$

• kernel 内部应当 只包含一个循环，该循环沿着 key 维度遍历所有 key tiles，即 $1\le j\le T_k$

• 在循环结束时，记得 推进（advance）所有 block pointer，以指向下一个 tile

• 请使用下面给出的 函数声明模板（我们已经为你提供了部分 block pointer 的定义，其余指针的设置方式应当可以自行推导出来）：

  @triton.jit
  def flash_fwd_kernel(
      Q_ptr, K_ptr, V_ptr,
      O_ptr, L_ptr,
      stride_qb, stride_qq, stride_qd,
      stride_kb, stride_kk, stride_kd,
      stride_vb, stride_vk, stride_vd,
      stride_ob, stride_oq, stride_od,
      stride_lb, stride_lq,
      N_QUERIES, N_KEYS,
      scale,
      D: tl.constexpr,
      Q_TILE_SIZE: tl.constexpr,
      K_TILE_SIZE: tl.constexpr,
  ):
  

  其中，$\mathrm{scale}=\frac{1}{\sqrt{d}}$，而 Q_TILE_SIZE 和 K_TILE_SIZE 分别对应 $B_q$ 和 $B_k$，这些参数后续都可以根据性能需要进行调优

下面是一些额外的实现建议，可以帮助你避免数值精度方面的问题：

• 位于片上（on-chip）的缓冲区 $(\mathbf{O},l,m)$ 应当使用 tl.float32 作为数据类型，如果你在向输出缓冲区中进行累加，请使用 acc 参数，例如 acc = tl.dot(..., acc=acc)
• 在将 ${\tilde{\mathrm{P}}}_i^{(j)}$ 与 ${\mathbf{V}}^{(j)}$ 相乘之前，应当先将 ${\tilde{\mathrm{P}}}_i^{(j)}$ 转换为与 ${\mathbf{V}}^{(j)}$ 相同的数据类型；在将 $\mathbf{O}$ 写回全局内存之前，也应将其转换为合适的数据类型。类型转换可以通过 tensor.to 来完成，你可以通过 tensor.dtype 获取一个 tensor 的数据类型，而 block pointer 或普通指针的数据类型可以通过 *_block_ptr.type.element_ty 获取

Deliverable：实现一个继承自 torch.autograd.Function 的子类，在其 forward 中调用你编写的 Triton kernel，从而实现 FlashAttention-2 的前向传播，实现 [adapters.get_flash_autograd_function_triton] 后，请运行以下测试命令以验证正确性：

uv run pytest -k test_flash_forward_pass_triton

(c) 请在你的 autograd.Function 实现中，将 因果掩码（causal masking）作为最后一个参数 加入，该参数应为一个布尔类型标志，当其设置为 True 时，启用用于因果掩码的索引比较逻辑。你的 Triton kernel 需要有一个与之对应的额外参数 is_causal: tl.constexpr（这是类型注解所必需的）

在 Triton 中需要为 queries 和 keys 构造合适的索引向量，将它们与一个大小为 $B_q\times B_k$ 的方形掩码进行比较，对于被掩码的位置，在注意力得分矩阵 ${\mathbf{S}}_i^{(j)}$ 的对应元素上 加上常数值 -1e6。请务必在 forward 中保存该掩码标志，以便在反向传播阶段使用 ctx.is_causal = is_causal

Deliverable：为你的 torch.autograd.Function 子类增加一个可选的因果掩码标志，使其能够通过你实现的 Triton kernel 执行带因果掩码的 FlashAttention-2 前向传播。请确保该标志是可选参数，默认值为 False，以保证之前的所有测试仍然可以通过。

代码实现如下：

import math
import torch
import triton
import triton.language as tl


@triton.jit
def flash_fwd_kernel(
    Q_ptr, K_ptr, V_ptr,
    O_ptr, L_ptr,
    stride_qb, stride_qq, stride_qd,
    stride_kb, stride_kk, stride_kd,
    stride_vb, stride_vk, stride_vd,
    stride_ob, stride_oq, stride_od,
    stride_lb, stride_lq,
    N_QUERIES: tl.constexpr,
    N_KEYS: tl.constexpr,
    scale,
    D: tl.constexpr,
    Q_TILE_SIZE: tl.constexpr,
    K_TILE_SIZE: tl.constexpr,
    IS_CAUSAL: tl.constexpr
):
    # program ids
    pid_q = tl.program_id(0)  # query tile id
    pid_b = tl.program_id(1)  # batch id

    # offsets
    q_offsets = pid_q * Q_TILE_SIZE + tl.arange(0, Q_TILE_SIZE)  # [Bq]
    d_offsets = tl.arange(0, D)  # [D]

    # pointers for Q tile: (Bq, D)
    q_ptrs = Q_ptr + pid_b * stride_qb + q_offsets[:, None] * stride_qq + d_offsets[None, :] * stride_qd
    q = tl.load(q_ptrs, mask=(q_offsets[:, None] < N_QUERIES), other=0.0).to(tl.float32)

    # running state (on-chip)
    m = tl.full((Q_TILE_SIZE,), -float("inf"), tl.float32)  # [Bq]
    l = tl.zeros((Q_TILE_SIZE,), tl.float32)                # [Bq]
    acc = tl.zeros((Q_TILE_SIZE, D), tl.float32)            # [Bq, D]

    # loop over K tiles
    for kb in tl.static_range(0, N_KEYS, K_TILE_SIZE):
        k_offsets = kb + tl.arange(0, K_TILE_SIZE)  # [Bk]

        k_ptrs = K_ptr + pid_b * stride_kb + k_offsets[:, None] * stride_kk + d_offsets[None, :] * stride_kd
        v_ptrs = V_ptr + pid_b * stride_vb + k_offsets[:, None] * stride_vk + d_offsets[None, :] * stride_vd
    
        k = tl.load(k_ptrs, mask=(k_offsets[:, None] < N_KEYS), other=0.0).to(tl.float32)  # [Bk, D]
        v = tl.load(v_ptrs, mask=(k_offsets[:, None] < N_KEYS), other=0.0)                 # [Bk, D]

        # S = q @ k^T * scale -> [Bq, Bk]
        S = tl.dot(q, tl.trans(k)) * scale  # float32

        # causal mask: keep if q_idx >= k_idx else -1e-6
        if IS_CAUSAL:
            q_abs = q_offsets[:, None]  # [Bq, 1]
            k_abs = k_offsets[None, :]  # [1, Bk]
            causal = q_abs >= k_abs
            S = tl.where(causal, S, -1.0e6)
        
        # online softmax update
        m_new = tl.maximum(m, tl.max(S, axis=1))  # [Bq]
        p = tl.exp(S - m_new[:, None])            # [Bq, Bk]

        alpha = tl.exp(m - m_new)                 # [Bq]
        l_new = alpha * l + tl.sum(p, axis=1)     # [Bq]

        # acc = alpha * acc + p @ v
        # p needs to match v dtype before dot
        p = p.to(v.dtype)
        acc = alpha[:, None] * acc
        acc = tl.dot(p, v, acc=acc)

        m = m_new
        l = l_new
    
    # write O and L
    o = acc / l[:, None]
    o = o.to(tl.float32)

    o_ptrs = O_ptr + pid_b * stride_ob + q_offsets[:, None] * stride_oq + d_offsets[None, :] * stride_od
    tl.store(o_ptrs, o, mask=(q_offsets[:, None] < N_QUERIES))

    L_out = m + tl.log(l)  # [Bq]
    l_ptrs = L_ptr + pid_b * stride_lb + q_offsets * stride_lq
    tl.store(l_ptrs, L_out, mask=(q_offsets < N_QUERIES))


class FlashAttention2Triton(torch.autograd.Function):
    @staticmethod
    def forward(ctx, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, is_causal: bool = False):
        # Expect (B, Q, D), (B, K, D), (B, K, D)
        if not q.is_cuda:
            raise RuntimeError("Triton implementation requires CUDA tensors")
        if q.dim() != 3 or k.dim() != 3 or v.dim() != 3:
            raise ValueError("Expected q/k/v to be 3D: (B, N, D)")
        if q.shape[0] != k.shape[0] or q.shape[0] != v.shape[0]:
            raise ValueError("Batch size mismatch")
        if k.shape[1] != v.shape[1] or k.shape[2] != v.shape[2]:
            raise ValueError("k/v shape mismatch")
        if q.shape[2] != k.shape[2]:
            raise ValueError("q/k D mismatch")
        
        B, Q, D = q.shape
        K = k.shape[1]
        scale = 1.0 / math.sqrt(D)

        # tile sizes
        Bq = 32
        Bk = 32

        # outputs
        o = torch.empty((B, Q, D), device=q.device, dtype=q.dtype)
        L = torch.empty((B, Q), device=q.device, dtype=torch.float32)

        grid = (triton.cdiv(Q, Bq), B)

        flash_fwd_kernel[grid](
            q, k, v,
            o, L,
            q.stride(0), q.stride(1), q.stride(2),
            k.stride(0), k.stride(1), k.stride(2),
            v.stride(0), v.stride(1), v.stride(2),
            o.stride(0), o.stride(1), o.stride(2),
            L.stride(0), L.stride(1),
            N_QUERIES=Q,
            N_KEYS=K,
            scale=scale,
            D=D,
            Q_TILE_SIZE=Bq,
            K_TILE_SIZE=Bk,
            IS_CAUSAL=is_causal,
            num_warps=4
        )

        # save for backward
        ctx.save_for_backward(L, q, k, v, o)
        ctx.is_causal = is_causal
        return o

    @staticmethod
    def backward(ctx, *grad_outputs):
        raise NotImplementedError

测试适配器 [adapters.get_flash_autograd_function_triton] 的实现如下：

def get_flashattention_autograd_function_triton() -> Type:
    """
    Returns a torch.autograd.Function subclass that implements FlashAttention2
    using Triton kernels.
    The expectation is that this class will implement the same operations
    as the class you return in get_flashattention_autograd_function_pytorch(),
    but it should do so by invoking custom Triton kernels in the forward
    and backward passes.

    Returns:
        A class object (not an instance of the class)
    """
    # For example: return MyTritonFlashAttentionAutogradFunctionClass
    from cs336_systems.flash_triton import FlashAttention2Triton
    return FlashAttention2Triton

执行 uv run pytest -k test_flash_forward_pass_triton 后输出如下：

flash_fwd_kernel 内部的核心实现其实和 PyTorch 没有太多的区别，都是对照着 Algorithm 1 来完成的，Triton 版相比于 PyTorch 版的主要优化点在于：

1. 并行性提升：把外层循环（遍历 query tiles）从串行 Python 循环变成 GPU 上并行的 program grid

2. 内存流量下降：不显式实现 $\mathbf{S}$ / $\mathbf{P}$ 的全矩阵，尤其是 ${\mathbf{QK}}^{\top }$ 这个巨大的注意力权重矩阵，只维护 online softmax 的 m,l,acc

3. kernel fusion：在一个 kernel 里完成 score→softmax→加权求和 的整条链路，减少中间张量的全局内存写回与 kernel launch 开销

所以总体优化关键就是：外层循环并行化 + 内存循环融合 + 中间结果不落显存

下面我们来简单看下 Triton 里 offset/索引/指针的设计：

1) 先定每个 program 负责那一块：grid → tile 坐标

grid = (triton.cdiv(Q, Bq), B)

这意味着：

• pid_q = tl.program_id(0)：第几个 query tile
• pid_b = tl.program_id(1)：第几个 batch

一个 Triton program instance（相当于 CUDA 中的一个 block）负责：固定 batch=pid_b，固定 query tile=pid_q，计算输出 O 的一个 tile（BqxD）以及 L 的一个 tile（Bq）

2) 写 tile 内的坐标：tl.arange

在一个 program 内，我们通过两组 arange 来生成 tile 内行列坐标：

q_offsets = pid_q * Q_TILE_SIZE + tl.arange(0, Q_TILE_SIZE)  # shape [Bq]
d_offsets = tl.arange(0, D)                                  # shape [D]

其中：

• q_offsets[r] 是这个 tile 里第 r 行对应的 全局 query index，即 q_global = pid_q * Bq + r
• d_offsets[c] 是 embedding 维度的 列 index，即 d = c

我们可以利用它们 broadcasting 组成一个 2D 网格：

• q_offsets[:, None] → shape [Bq, 1]
• d_offsets[None, :] → shape 1, D
• 相加/线性组合后变成 shape [Bq, D]，表示 tile 里每个元素的全局坐标

这就相当于 CUDA block 里二维 threadIdx，但 Triton 是向量化 + 编译器映射到 threads/warps，我们不用直接写 threadIdx

3) 从坐标到内存地址：stride + 线性地址公式

我们的 query 是三维 (B, Q, D)，PyTorch 的张量是 “基址 + stride * index” 寻址，我们传进 kernel 的是：

stride_qb, stride_qq, stride_qd = q.stride(0), q.stride(1), q.stride(2)

Note：Pytorch stride 的单位是元素个数，不是字节，Triton 也按元素算，最后底层会乘 dtype bytes

那么 Q[b, q, d] 的地址就是：

$$\text{addr}(b,q,d)=Q\_ptr+b\cdot stride\_qb+q\cdot stride\_qq+d\cdot stride\_qd$$

对应代码：

q_ptrs = Q_ptr + pid_b * stride_qb + q_offsets[:, None] * stride_qq + d_offsets[None, :] * stride_qd

其中：

• pid_b * stride_qb：跳到第 pid_b 个 batch 的起始位置
• q_offsets[:,None] * stride_qq：tile 内每一行跳到对应 query 行
• d_offsets[None,:] * stride_qd：tile 内每一列跳到对应 embedding 维

最终 q_ptrs 是一个 [Bq, D] 的指针矩阵，tl.load(q_ptrs) 就一次性把整个 tile 读出来，同样逻辑适用于 K/V/O/L，只是用各自 stride

4) K/V 的 tile 索引：循环推进

我们内部循环变量是 kb：

for kb in tl.static_range(0, N_KEYS, K_TILE_SIZE):
    k_offsets = kb + tl.arange(0, K_TILE_SIZE)  # [Bk]

其中：

• kb 是当前 key tile 的起始全局 index（0, Bk, 2Bk, …）
• k_offsets[t] = kb + t 是 tile 内第 t 个 key 的全局 index

将它和 d_offsets 组合就可以得到 K tile 的地址：

k_ptrs = K_ptr + pid_b*stride_kb + k_offsets[:,None]*stride_kk + d_offsets[None,:]*stride_kd
v_ptrs = V_ptr + pid_b*stride_vb + k_offsets[:,None]*stride_vk + d_offsets[None,:]*stride_vd

这其实就是论文中的 Load ${\mathbf{K}}^{(j)},{\mathbf{V}}^{(j)}$

值得注意的是，K/V 不需要 program id，因为它们不是按 tile 并行分配的，而是同一个 program 内沿 key 维扫一遍

5) mask 的全局坐标计算

causal 的判断依赖全局序列位置，而不是 tile 内相对位置，所以我们用：

q_abs = q_offsets[:, None]   # [Bq,1] global query index
k_abs = k_offsets[None, :]   # [1,Bk] global key index
causal = q_abs >= k_abs

这就是标准的 causal 条件：query 位置 i 只能看 key ≤ i

6) 尾部边界的处理

当 Q/K 不能整除 tile size 时，最后一个 tile 会越界，所以我们在 load/stroe 时都写了 mask：

tl.load(q_ptrs, mask=(q_offsets[:, None] < N_QUERIES), other=0.0)
tl.load(..., mask=(k_offsets[:, None] < N_KEYS), other=0.0)
tl.store(o_ptrs, o, mask=(q_offsets[:, None] < N_QUERIES))
tl.store(l_ptrs, L_out, mask=(q_offsets < N_QUERIES))

对于尾部越界的元素在 load 时当成 0 来处理，并且不写回

3.2 优化 (block pointer)

值得注意的是，博主在完成后续 triton-backward 优化时突然想到作业中提到的可以使用 Triton 提供的 block pointer 抽象 即 tl.make_block_prt 来简化繁琐的指针操作，而这个优化在上面的 flash_fwd_kernel 中并未实现

下面我们就来简单的优化下 flash_fwd_kernel，通过 tl.make_block_prt 来简化索引计算等繁琐操作

实现代码如下：

@triton.jit
def flash_fwd_kernel_bp(
    Q_ptr, K_ptr, V_ptr,
    O_ptr, L_ptr,
    stride_qb, stride_qq, stride_qd,
    stride_kb, stride_kk, stride_kd,
    stride_vb, stride_vk, stride_vd,
    stride_ob, stride_oq, stride_od,
    stride_lb, stride_lq,
    N_QUERIES: tl.constexpr,
    N_KEYS: tl.constexpr,
    scale,
    D: tl.constexpr,
    Q_TILE_SIZE: tl.constexpr,
    K_TILE_SIZE: tl.constexpr,
    IS_CAUSAL: tl.constexpr,
):
    # program ids
    pid_q = tl.program_id(0)  # query tile id
    pid_b = tl.program_id(1)  # batch id

    # block pointers for Q/K/V/O/L tiles
    # block pointers encapsulate base, shape, strides, and OOB handling;
    # advancing them avoids re-materializing raw pointer arithmetic in the loop
    Qb = Q_ptr + pid_b * stride_qb
    Kb = K_ptr + pid_b * stride_kb
    Vb = V_ptr + pid_b * stride_vb
    Ob = O_ptr + pid_b * stride_ob
    Lb = L_ptr + pid_b * stride_lb

    Q_bp = tl.make_block_ptr(
        base=Qb,
        shape=(N_QUERIES, D),
        strides=(stride_qq, stride_qd),
        offsets=(pid_q * Q_TILE_SIZE, 0),
        block_shape=(Q_TILE_SIZE, D),
        order=(1, 0),
    )
    K_bp = tl.make_block_ptr(
        base=Kb,
        shape=(N_KEYS, D),
        strides=(stride_kk, stride_kd),
        offsets=(0, 0),
        block_shape=(K_TILE_SIZE, D),
        order=(1, 0),
    )
    V_bp = tl.make_block_ptr(
        base=Vb,
        shape=(N_KEYS, D),
        strides=(stride_vk, stride_vd),
        offsets=(0, 0),
        block_shape=(K_TILE_SIZE, D),
        order=(1, 0),
    )
    O_bp = tl.make_block_ptr(
        base=Ob,
        shape=(N_QUERIES, D),
        strides=(stride_oq, stride_od),
        offsets=(pid_q * Q_TILE_SIZE, 0),
        block_shape=(Q_TILE_SIZE, D),
        order=(1, 0),
    )
    L_bp = tl.make_block_ptr(
        base=Lb,
        shape=(N_QUERIES,),
        strides=(stride_lq,),
        offsets=(pid_q * Q_TILE_SIZE,),
        block_shape=(Q_TILE_SIZE,),
        order=(0,),
    )

    # load the Q tile
    # keep the original dtype for the final store cast
    q_raw = tl.load(Q_bp, boundary_check=(0, 1), padding_option="zero")
    q = q_raw.to(tl.float32)

    # running state (on-chip)
    m = tl.full((Q_TILE_SIZE,), -float("inf"), tl.float32)  # [Bq]
    l = tl.zeros((Q_TILE_SIZE,), tl.float32)                # [Bq]
    acc = tl.zeros((Q_TILE_SIZE, D), tl.float32)            # [Bq, D]

    # iterate over K/V tile by advancing block pointers (instead of re-building raw pointers)
    K_it = K_bp
    V_it = V_bp

    # absolute query indices used only for causal masking
    q_abs = pid_q * Q_TILE_SIZE + tl.arange(0, Q_TILE_SIZE)

    for kb in range(0, N_KEYS, K_TILE_SIZE):
        # load one (K_TILE_SIZE, D) tile of K and V
        k = tl.load(K_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bk, D]
        v = tl.load(V_it, boundary_check=(0, 1), padding_option="zero")                 # [Bk, D]

        # S = q @ k^T * scale -> [Bq, Bk]
        S = tl.dot(q, tl.trans(k)) * scale  # float32

        # causal mask: keep if q_idx >= k_idx else -1e-6
        if IS_CAUSAL:
            k_abs = kb + tl.arange(0, K_TILE_SIZE)
            S = tl.where(q_abs[:, None] >= k_abs[None, :], S, -1.0e6)

        # online softmax update
        m_new = tl.maximum(m, tl.max(S, axis=1))  # [Bq]
        p = tl.exp(S - m_new[:, None])            # [Bq, Bk]

        alpha = tl.exp(m - m_new)                 # [Bq]
        l_new = alpha * l + tl.sum(p, axis=1)     # [Bq]

        # acc = alpha * acc + p @ v
        # p needs to match v dtype before dot
        p = p.to(v.dtype)
        acc = alpha[:, None] * acc
        acc = tl.dot(p, v, acc=acc)

        m = m_new
        l = l_new

        # advance K/V block pointers to the next tile along the sequence dimension
        K_it = K_it.advance((K_TILE_SIZE, 0))
        V_it = V_it.advance((K_TILE_SIZE, 0))

    # write O and L
    # store output in the original input dtype
    o = (acc / l[:, None]).to(q_raw.dtype)

    tl.store(O_bp, o, boundary_check=(0, 1))

    L_out = m + tl.log(l)  # [Bq]
    tl.store(L_bp, L_out, boundary_check=(0,))

在实现 flash_fwd_kernel_bp 时，我们的计算逻辑（例如 online softmax 的 m/l/acc 更新、p @ v 累加等）与原始 flash_fwd_kernel_bp 保持一致，主要变化集中在 tile 的索引/指针表达方式：从 “手写指针 + mask” 切换为 Triton 的 block pointer 抽象（tl.make_block_ptr），并通过 tl.advance 在循环中推进 tile

主要优化点如下：

1) 核心改动：用 block pointer 表达 tile，而不是手拼指针矩阵

在旧实现中，每次加载一个 tile 都需要显式构造二维指针网格：

• Q tile：通过 q_offsets[:,None] 与 d_offsets[None,:] 广播，拼出 q_ptrs（shape [Bq, D]）
• K/V tile：每次循环生成 k_offsets，再拼出 k_ptrs / v_ptrs（shape [Bk, D]）
• O/L 写回：同样要拼 o_ptrs / l_ptrs

这种写法直观，但会让 kernel 内充斥大量 “地址计算表达式”，并且边界处理需要显式写 mask=。

而在 block pointer 版本中，我们把 “一个 tile 对应哪块内存区域” 封装在 tl.make_block_ptr 里：

Q_bp = tl.make_block_ptr(
    base=Qb,
    shape=(N_QUERIES, D),
    strides=(stride_qq, stride_qd),
    offsets=(pid_q * Q_TILE_SIZE, 0),
    block_shape=(Q_TILE_SIZE, D),
    order=(1, 0),
)

这里 shape/strides/offsets/block_shape 合在一起，描述的就是一个 可滑动的 2D 视窗：起点在 (pid_q * Bq, 0)，每次加载/存储都对应一个 (Bq, D) tile。K/V/O/L 也是同样的构造方式，只是 tile 的起点和 shape 不同

因此，从 “坐标 → 指针” 的那一步（stride address arithmetic）在 block pointer 版本里只需要写一次初始化，而不需要在每轮循环里反复拼接

2) 循环推进方式改变：从 “重新构造指针” 到 “advance 滑窗”

旧实现的 K/V tile 遍历依赖：

for kb in tl.static_range(0, N_KEYS, K_TILE_SIZE):
    k_offsets = kb + tl.arange(0, K_TILE_SIZE)
    k_ptrs = ...
    v_ptrs = ...
    k = tl.load(k_ptrs, mask=..., other=0.0)
    v = tl.load(v_ptrs, mask=..., other=0.0)

也就是说每一轮都要重新计算 k_ptrs / v_ptrs。

block pointer 版本中，K/V 的 tile 不再通过显式 pointer arithmetic 重建，而是把 K/V 的 block pointer 作为迭代变量：

K_it = K_bp
V_it = V_bp
for kb in range(0, N_KEYS, K_TILE_SIZE):
    k = tl.load(K_it, boundary_check=(0, 1), padding_option="zero")
    v = tl.load(V_it, boundary_check=(0, 1), padding_option="zero")
    ...
    K_it = K_it.advance((K_TILE_SIZE, 0))
    V_it = V_it.advance((K_TILE_SIZE, 0))

advance((K_TILE_SIZE, 0)) 的含义就是：沿着第 0 维（sequence 维）把 tile 起点向前移动一个 Bk。这样整个循环更像 “滑动窗口扫描 K/V”，而不是每轮都重新拼地址表达式。

3) 边界处理方式变化：从 mask= 到 boundary_check + padding_option

旧实现里，Q/K/V 的 load 都显式写 mask，并在越界处用 other=0.0：

tl.load(q_ptrs, mask=(q_offsets[:, None] < N_QUERIES), other=0.0)
tl.load(k_ptrs, mask=(k_offsets[:, None] < N_KEYS), other=0.0)

block pointer 版本把越界逻辑交给 tl.load 对 block pointer 的边界机制：

tl.load(Q_bp, boundary_check=(0, 1), padding_option="zero")
tl.load(K_it, boundary_check=(0, 1), padding_option="zero")

其中 boundary_check=(0,1) 表示对 tile 的两个维度都做边界检查；padding_option="zero" 等价于旧代码的 other=0.0，即越界元素按 0 填充，这让代码更简洁，并且避免反复写 mask 条件。

同理，写回 O/L 也从显式 mask= 变成 boundary_check=

综合来看，flash_fwd_kernel_bp 相比于 flash_fwd_kernel 的主要收益可以概括为三点：

1. 代码结构更清晰

把 “tile 对应的内存区域” 集中在 make_block_ptr 初始化处，循环体主要保留数学逻辑（online softmax + matmul），减少大量重复的地址拼接表达式

2. 边界处理更统一

boundary_check + padding_option 替代手写 mask，避免在每个 load/store 上重复写条件，同时语义更接近 “加载一个 tile，越界补 0”

3. 潜在性能更好

block pointer 为编译器提供了更明确的 “块状访存” 信息，配合 order=(1, 0) 指定更连续维度优先，有助于生成更好的内存访问与向量化代码；并且循环中通过 advance 推进视窗，减少了显式 pointer arithmetic 的指令与寄存器压力。博主在后续 benchmark 中也发现该版本的 forward 耗时明显更优

OK，以上就是 FlashAttention forward 部分的 Triton 实现了。

4. Problem (flash_backward): 5 points

Implementing the backward pass with recomputation

注意，与公式 (7)-(11) 中给出的标准反向传播不同，我们可以通过 重计算（recomputation）来避免在反向传播阶段执行 softmax 运算（如公式 (13)-(19) 所示），这意味着反向传播可以通过一个相对简单的 kernel 来完成，不需要任何在线 softmax 技巧。因此，在这一部分，你可以直接通过在一个普通的 PyTorch 函数（而不是 Triton kernel）上调用 torch.compile 来实现反向传播

请使用 PyTorch（而非 Triton） 和 torch.compile 为你的 FlashAttention-2 autograd.Function 实现反向传播。你的实现应当接收张量 $\mathbf{Q},\mathbf{K},\mathbf{V},\mathbf{O},\mathbf{dO}$ 和 $L$ 作为输入，并返回 $\mathbf{dQ},\mathbf{dK}$ 和 $\mathbf{dV}$。请注意在计算过程中需要显式地计算并使用向量 $D$，你可以按照公式 (13)-(19) 中给出的计算流程来实现反向传播

Deliverable：为了测试你的实现，请运行以下命令：

uv run pytest -k test_flash_backward

并确保你的实现能够顺利通过测试。

代码实现如下：

def flash_bwd_recompute_impl(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    o: torch.Tensor,
    do: torch.Tensor,
    L: torch.Tensor,
    is_causal: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Inputs:
      q,k,v,o,do: (B, Q/K, D)
      L: (B, Q)   where L = logsumexp(S, dim=-1)
    Returns:
      dq, dk, dv    
    """
    B, Q, D = q.shape
    K = k.shape[1]
    scale = 1.0 / math.sqrt(D)

    # S = QK^T * scale : (B, Q, K)
    S = torch.matmul(q, k.transpose(-1, -2)) * scale

    if is_causal:
        q_idx = torch.arange(Q, device=S.device)[:, None]
        k_idx = torch.arange(K, device=S.device)[None, :]
        casual = (q_idx >= k_idx)  # (Q, K)
        S = torch.where(casual[None, :, :], S, torch.full_like(S, -1.0e6))
    
    # P = exp(S - L) : (B, Q, K)
    P = torch.exp(S - L.unsqueeze(-1))

    # dV = P^T @ dO : (B, K, D)
    dv = torch.matmul(P.transpose(-1, -2), do)

    # dP = dO @ V^T : (B, Q, K)
    dP = torch.matmul(do, v.transpose(-1, -2))

    # Dvec = sum(dO * O, dim=-1) : (B, Q)
    Dvec = (do * o).sum(dim=-1)

    # dS = P * (dP - Dvec) * scale : (B, Q, K)
    dS = P * (dP - Dvec.unsqueeze(-1)) * scale

    # dQ = dS @ K : (B, Q, D)
    dq = torch.matmul(dS, k)

    # dK = dS^T @ Q : (B, K, D)
    dk = torch.matmul(dS.transpose(-1, -2), q)

    return dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype)  


class FlashAttention2Pytorch(torch.autograd.Function):
    ...
    @staticmethod
    def backward(ctx, do):
        (L, q, k, v, o) = ctx.saved_tensors
        dq, dk, dv = flash_bwd_recompute_impl(q, k, v, o, do, L, ctx.is_causal)
        return dq, dk, dv, None    

执行 uv run pytest -k test_flash_backward 后输出如下：

整个代码的实现非常简单，就是将下面的 Algorithm 2 用代码进行实现。它本质上是在 PyTorch 里实现了 重计算版反向：不用保存 softmax 概率矩阵 P，也不用在 backward 里做在线 softmax，直接用 L = logsumexp 把 P 重建出来，然后套公式把 dQ/dK/dV 算出来就行

5. Problem (flash_benchmarking): 5 points

接下来，我们来比较你 基于 Triton 的 FlashAttention-2 实现与使用 PyTorch 实现的标准 Attention 在性能上的差异

(a) 使用 triton.testing.do_bench 编写一个 性能基准测试脚本，比较以下两种实现的性能：

• 你实现的 Triton 版 FlashAttention-2 的前向和反向传播
• 普通的 PyTorch Attention 实现（即不使用 FlashAttention）

具体要求如下：

• 你需要给出一张结果表，包含 前向传播（forward）、反向传播（backward）以及端到端前向 + 反向传播（end-to-end forward-backward） 的延迟；
• 对 Triton 实现 和 PyTorch 实现 都需要分别测量上述三类延迟
• 在开始基准测试之前，需随机生成所有必要的输入；
• 基准测试需在 单张 H100 GPU 上运行；
• 始终使用 batch size = 1，并启用 因果掩码（causal masking）
• 在测试中遍历如下组合（笛卡儿积）：
  • 序列长度：从 128 到 65536，取 2 的幂
  • 嵌入维度：从 16 到 128，取 2 的幂
  • 数据精度：torch.bfloat16 和 torch.float32
• 你可能需要根据输入规模调整 tile size

Deliverable：提交一张结果表，对比你实现的 FlashAttention-2 与 PyTorch Attention 在上述设置下的性能，并分别报告前向延迟、反向延迟、端到端前向+反向延迟。

bench_flash_vs_pytorch.py 脚本实现如下：

import torch
import triton.testing
from typing import Tuple, Optional
import argparse

from cs336_basics.modules import scaled_dot_product_attention
from cs336_systems.flash_triton import FlashAttention2Triton

from cs336_systems.utils import FlashBenchRow, FlashBenchmarkReporter


def make_causal_mask(n: int, device: torch.device) -> torch.Tensor:
    idx = torch.arange(n, device=device)
    return (idx[:, None] >= idx[None, :])


def pow2_list(lo: int, hi: int) -> list[int]:
    out = []
    x = lo
    while x <= hi:
        out.append(x)
        x *= 2
    return out


@torch.no_grad()
def make_inputs(
    seq_len: int,
    d_model: int,
    dtype: torch.dtype,
    device: torch.device,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    q = torch.randn((1, seq_len, d_model), device=device, dtype=dtype)
    k = torch.randn((1, seq_len, d_model), device=device, dtype=dtype)
    v = torch.randn((1, seq_len, d_model), device=device, dtype=dtype)
    return q, k, v


def bench_one(
    impl: str,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    is_causal: bool,
    warmup: int,
    rep: int
) -> Tuple[Optional[float], Optional[float], Optional[float], str]:
    """
    Returns: (fwd_ms, bwd_ms, e2e_ms, status)
    """
    assert q.shape[0] == 1, "batch size must be 1"

    device = q.device
    mask = make_causal_mask(q.shape[1], device) if is_causal else None

    def fwd_fn():
        if impl == "baseline":
            return scaled_dot_product_attention(q, k, v, mask=mask)
        elif impl == "flash":
            return FlashAttention2Triton.apply(q, k, v, is_causal)
        else:
            raise ValueError(f"unknow impl: {impl}")
    
    try:
        # foward benchmark
        fwd_ms = float(triton.testing.do_bench(fwd_fn, warmup=warmup, rep=rep))

        # build one graph for backward-only
        q_ = q.detach().requires_grad_(True)
        k_ = k.detach().requires_grad_(True)
        v_ = v.detach().requires_grad_(True)

        if impl == "baseline":
            out = scaled_dot_product_attention(q_, k_, v_, mask=mask)
        else:
            out = FlashAttention2Triton.apply(q_, k_, v_, is_causal)

        do = torch.randn_like(out)

        def bwd_only():
            torch.autograd.grad(out, (q_, k_, v_), grad_outputs=do, retain_graph=True)
        
        bwd_ms = float(triton.testing.do_bench(bwd_only, warmup=warmup, rep=rep))

        # end-to-end benchmark (rebuild graph each rep)
        def e2e():
            qx = q.detach().requires_grad_(True)
            kx = k.detach().requires_grad_(True)
            vx = v.detach().requires_grad_(True)

            if impl == "baseline":
                oy = scaled_dot_product_attention(qx, kx, vx, mask=mask)
            else:
                oy = FlashAttention2Triton.apply(qx, kx, vx, is_causal)
            
            doy = torch.rand_like(oy)
            torch.autograd.grad(oy, (qx, kx, vx), grad_outputs=doy, retain_graph=False)

        e2e_ms = float(triton.testing.do_bench(e2e, warmup=warmup, rep=rep))

        return fwd_ms, bwd_ms, e2e_ms, "ok"
    
    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        return None, None, None, "oom"

    except Exception as e:
        print(e)
        return None, None, None, f"error:{type(e).__name__}"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out-jsonl", type=str, default="runs/flash_bench.jsonl")
    ap.add_argument("--out-md", type=str, default="runs/flash_bench.md")
    ap.add_argument("--device", type=str, default="cuda")

    ap.add_argument("--warmup", type=int, default=25)
    ap.add_argument("--rep", type=int, default=100)

    ap.add_argument("--seq-min", type=int, default=128)
    ap.add_argument("--seq-max", type=int, default=65536)
    ap.add_argument("--d-min", type=int, default=16)
    ap.add_argument("--d-max", type=int, default=128)

    args = ap.parse_args()

    if args.device.startswith("cuda") and not torch.cuda.is_available():
        raise RuntimeError("CUDA not available")
    
    device = torch.device(args.device)

    reporter = FlashBenchmarkReporter(
        args.out_jsonl,
        args.out_md,
        title="#### FlashAttention-2 (Triton) vs Baseline (PyTorch)  (batch=1, causal=True)"
    )

    seq_list = pow2_list(args.seq_min, args.seq_max)
    d_list = pow2_list(args.d_min, args.d_max)
    dtypes = [(torch.bfloat16, "bf16"), (torch.float32, "fp32")]

    for n in seq_list:
        for d in d_list:
            for dtype, dtype_name in dtypes:
                q, k, v = make_inputs(n, d, dtype=dtype, device=device)

                for impl in ["baseline", "flash"]:
                    fwd_ms, bwd_ms, e2e_ms, status = bench_one(
                        impl=impl,
                        q=q,
                        k=k,
                        v=v,
                        is_causal=True,
                        warmup=args.warmup,
                        rep=args.rep,
                    )

                    row = FlashBenchRow(
                        impl=impl,
                        dtype=dtype_name,
                        seq_len=n,
                        d_model=d,
                        fwd_ms=fwd_ms,
                        bwd_ms=bwd_ms,
                        e2e_ms=e2e_ms,
                        status=status,
                    )

                    tag = f"{impl}|{dtype_name}"
                    msg = f"[{tag:13s}] d={d:4d}, s={n:6d} -> {row.status}"

                    if row.fwd_ms is not None:
                        msg += (
                            f", fwd={row.fwd_ms:.3f}ms"
                            f", bwd={row.bwd_ms:.3f}ms"
                            f", e2e={row.e2e_ms:.3f}ms"
                        )
                    
                    print(msg)

                    reporter.append(row)

    reporter.write_markdown()
    print(reporter.render_markdown())
    print(f"\nSaved: {args.out_jsonl}  and  {args.out_md}")


if __name__ == "__main__":
    main()    

运行指令如下：

uv run cs336_systems/bench_flash_vs_pytorch.py

执行运行指令后输出如下：

整个代码实现就是作业要求的使用 triton.testing.do_bench 对不同序列长度、不同嵌入维度以及不同数据精度下的 Triton Attention 和 Pytorch Attention 实现的前向传播、反向传播和端到端延迟对比

代码实现比较简单，大家可以自己看看

测量结果如下所示：

impl	seq_len	d_model	dtype	fwd_ms	bwd_ms	e2e_ms	status
baseline	128	16	bf16	0.0470	0.0778	0.1161	ok
flash	128	16	bf16	0.0064	0.0677	0.0712	ok
baseline	128	16	fp32	0.0340	0.0691	0.0973	ok
flash	128	16	fp32	0.0068	0.0481	0.0547	ok
baseline	128	32	bf16	0.0443	0.0775	0.1162	ok
flash	128	32	bf16	0.0069	0.0685	0.0723	ok
baseline	128	32	fp32	0.0343	0.0698	0.0995	ok
flash	128	32	fp32	0.0088	0.0491	0.0569	ok
baseline	128	64	bf16	0.0455	0.0807	0.1216	ok
flash	128	64	bf16	0.0086	0.0744	0.0806	ok
baseline	128	64	fp32	0.0354	0.0723	0.1039	ok
flash	128	64	fp32	0.0108	0.0540	0.0653	ok
baseline	128	128	bf16	0.0461	0.0803	0.1223	ok
flash	128	128	bf16	0.0127	0.0740	0.0834	ok
baseline	128	128	fp32	0.0364	0.0724	0.1039	ok
flash	128	128	fp32	0.0195	0.0531	0.0723	ok
baseline	256	16	bf16	0.0486	0.0918	0.1352	ok
flash	256	16	bf16	0.0110	0.0790	0.0873	ok
baseline	256	16	fp32	0.0387	0.0842	0.1175	ok
flash	256	16	fp32	0.0110	0.0578	0.0697	ok
baseline	256	32	bf16	0.0497	0.0900	0.1355	ok
flash	256	32	bf16	0.0110	0.0773	0.0864	ok
baseline	256	32	fp32	0.0392	0.0840	0.1182	ok
flash	256	32	fp32	0.0140	0.0573	0.0720	ok
baseline	256	64	bf16	0.0534	0.0936	0.1437	ok
flash	256	64	bf16	0.0154	0.0862	0.0993	ok
baseline	256	64	fp32	0.0419	0.0867	0.1266	ok
flash	256	64	fp32	0.0197	0.0638	0.0852	ok
baseline	256	128	bf16	0.0570	0.1030	0.1591	ok
flash	256	128	bf16	0.0247	0.0958	0.1206	ok
baseline	256	128	fp32	0.0458	0.0952	0.1434	ok
flash	256	128	fp32	0.0348	0.0723	0.1096	ok
baseline	512	16	bf16	0.0810	0.1723	0.2540	ok
flash	512	16	bf16	0.0181	0.1224	0.1395	ok
baseline	512	16	fp32	0.0700	0.1649	0.2317	ok
flash	512	16	fp32	0.0186	0.0997	0.1225	ok
baseline	512	32	bf16	0.0834	0.1774	0.2594	ok
flash	512	32	bf16	0.0203	0.1366	0.1524	ok
baseline	512	32	fp32	0.0730	0.1681	0.2396	ok
flash	512	32	fp32	0.0269	0.1098	0.1394	ok
baseline	512	64	bf16	0.0921	0.1923	0.2832	ok
flash	512	64	bf16	0.0317	0.1479	0.1815	ok
baseline	512	64	fp32	0.0808	0.1834	0.2657	ok
flash	512	64	fp32	0.0375	0.1330	0.1661	ok
baseline	512	128	bf16	0.1038	0.2130	0.3167	ok
flash	512	128	bf16	0.0640	0.1945	0.2544	ok
baseline	512	128	fp32	0.0911	0.2002	0.2945	ok
flash	512	128	fp32	0.0736	0.1600	0.2303	ok
baseline	1024	16	bf16	0.2842	0.6228	0.9038	ok
flash	1024	16	bf16	0.0517	0.3874	0.4278	ok
baseline	1024	16	fp32	0.2710	0.6152	0.8855	ok
flash	1024	16	fp32	0.0527	0.3632	0.4097	ok
baseline	1024	32	bf16	0.2862	0.6233	0.9085	ok
flash	1024	32	bf16	0.0644	0.3943	0.4520	ok
baseline	1024	32	fp32	0.2733	0.6141	0.8880	ok
flash	1024	32	fp32	0.0753	0.3763	0.4416	ok
baseline	1024	64	bf16	0.3034	0.6503	0.9549	ok
flash	1024	64	bf16	0.0830	0.4410	0.5133	ok
baseline	1024	64	fp32	0.2906	0.6404	0.9327	ok
flash	1024	64	fp32	0.0919	0.4095	0.4973	ok
baseline	1024	128	bf16	0.3609	0.7621	1.1223	ok
flash	1024	128	bf16	0.1321	0.6001	0.7044	ok
baseline	1024	128	fp32	0.3471	0.7471	1.0955	ok
flash	1024	128	fp32	0.1688	0.5682	0.7121	ok
baseline	2048	16	bf16	0.9939	2.2163	3.2262	ok
flash	2048	16	bf16	0.1152	1.3024	1.4038	ok
baseline	2048	16	fp32	0.9842	2.2054	3.1881	ok
flash	2048	16	fp32	0.1207	1.2772	1.3818	ok
baseline	2048	32	bf16	1.0074	2.2351	3.2467	ok
flash	2048	32	bf16	0.1377	1.3093	1.4356	ok
baseline	2048	32	fp32	0.9968	2.2254	3.2227	ok
flash	2048	32	fp32	0.1585	1.3063	1.4493	ok
baseline	2048	64	bf16	1.0497	2.3069	3.3631	ok
flash	2048	64	bf16	0.2696	1.4501	1.6767	ok
baseline	2048	64	fp32	1.0352	2.2913	3.3340	ok
flash	2048	64	fp32	0.3756	1.3933	1.7418	ok
baseline	2048	128	bf16	1.1553	2.5286	3.6887	ok
flash	2048	128	bf16	0.5241	1.7828	2.2167	ok
baseline	2048	128	fp32	1.1310	2.5111	3.6418	ok
flash	2048	128	fp32	0.6448	1.6765	2.2805	ok
baseline	4096	16	bf16	3.8105	8.5301	12.2380	ok
flash	4096	16	bf16	0.2369	4.7537	4.9881	ok
baseline	4096	16	fp32	3.7892	8.5441	12.2156	ok
flash	4096	16	fp32	0.2589	4.7633	4.9782	ok
baseline	4096	32	bf16	3.8535	8.5646	12.3668	ok
flash	4096	32	bf16	0.4550	4.9138	5.2974	ok
baseline	4096	32	fp32	3.8393	8.4995	12.4445	ok
flash	4096	32	fp32	0.6185	4.9034	5.4253	ok
baseline	4096	64	bf16	4.0966	8.8949	13.1057	ok
flash	4096	64	bf16	0.6352	5.4448	5.9079	ok
baseline	4096	64	fp32	4.0730	8.8768	12.9402	ok
flash	4096	64	fp32	1.1182	5.3835	6.3312	ok
baseline	4096	128	bf16	4.4149	9.6891	14.1139	ok
flash	4096	128	bf16	1.5544	6.6057	7.8837	ok
baseline	4096	128	fp32	4.4130	9.6632	14.0451	ok
flash	4096	128	fp32	1.9256	6.3899	8.1513	ok
baseline	8192	16	bf16	14.9943	33.4495	48.4823	ok
flash	8192	16	bf16	0.9313	19.0371	19.8308	ok
baseline	8192	16	fp32	14.9944	33.8135	48.4260	ok
flash	8192	16	fp32	1.0169	19.1322	19.9160	ok
baseline	8192	32	bf16	15.1243	33.6280	48.7791	ok
flash	8192	32	bf16	2.3201	19.1790	21.4119	ok
baseline	8192	32	fp32	15.1032	33.7637	48.6656	ok
flash	8192	32	fp32	2.6002	19.3208	21.6093	ok
baseline	8192	64	bf16	16.2875	35.1386	51.4376	ok
flash	8192	64	bf16	2.8117	21.0159	23.6802	ok
baseline	8192	64	fp32	16.2350	35.2031	51.4150	ok
flash	8192	64	fp32	3.7871	21.2019	24.5579	ok
baseline	8192	128	bf16	17.6281	38.6216	56.2483	ok
flash	8192	128	bf16	5.1404	25.4475	30.4227	ok
baseline	8192	128	fp32	17.5904	38.6356	56.1777	ok
flash	8192	128	fp32	6.6048	25.3943	31.6969	ok
baseline	16384	16	bf16	-	-	-	oom
flash	16384	16	bf16	3.9171	74.6281	78.3968	ok

延迟对比如下图所示：

加速比如下图所示：

Note：这是博主在自己主机 RTX3060（12GB）上测量得到的结果

从上面的基准测试结果来看，基于 Triton 实现的 FlashAttention-2 在前向传播阶段较于 PyTorch 基线实现具有显著的性能优势

随着序列长度的增加，前向传播的加速比迅速提升，在小 head 维度、bf16 精度条件下，前向计算最高可达到 15-16x 的加速。这主要得益于 FlashAttention 采用的 分块（block-wise）softmax 计算方式，避免了显式构造完整的注意力矩阵，从而显著降低了全局内存访问量并提升了缓存友好性

需要注意的是，本实验中 反向传播并未使用 Triton 实现，而是直接复用了 PyTorch 提供的标准 backward 实现，因此，反向阶段的性能提升相对温和，加速比约为 1.3-1.8x，总体加速比还是随着序列长度的增加而提升的。由于反向传播本身的计算与内存开销较大，其耗时在整体训练过程中占据主导地位，这也直接限制了端到端性能提升的幅度

综合前向和反向过程来看，整体端到端加速比约为 1.5-2.5x，尽管前向传播获得了数量级上的加速，但由于 backward 仍然是当前实现中的主要性能瓶颈，整体提升被显著 “稀释”。这说明在训练场景下，仅优化前向传播不足以获得最大端到端收益，高效的自定义反向传播实现同样是实现进一步加速的关键因素

总体而言，即便在博主的中端 GPU（RTX3060）上，我们也能清晰的看到 FlashAttention 在前向计算阶段的显著优势，同时也能了解当前端到端延迟的主要瓶颈所在，为后续实现 Triton-based backward 提供了明确的优化方向。

6. FlashAttention-2 Leaderboard

Assignment 2 的排行榜将测试你所实现的 FlashAttention-2 的速度表现（包括前向和反向传播），我们鼓励你尽可能使用各种技巧来进一步提升实现的性能

需要注意的限制条件如下：

• 不允许改变函数的输入/输出接口
• 必须使用 Triton（而非）CUDA 实现
• 输入将在 BF16 精度 + causal masking 条件下进行测试
• 你的实现必须通过与普通实现完全一致的正确性测试
• 实现必须是你 原创的，不能使用已有的第三方实现

性能测试将在一块 H100 GPU 上进行，测试配置如下：

• batch size = 1
• 查询、键和值（Q/K/V）的序列长度均为 16,384
• $d_{\text{model}}=1024$
• head 数为 16

我们将验证前 5 名提交的正确性和性能，测试代码形式如下：

def test_timing_flash_forward_backward():
    n_heads = 16
    d_head = 64
    sequence_length = 16384
    q, k, v = torch.randn(
        3, n_heads, sequence_length, d_head,
        device='cuda', dtype=torch.bfloat16, requires_grad=True
    )

    flash = torch.compile(FlashAttention2.apply)

    def flash_forward_backward():
        o = flash(q, k, v, True)
        loss = o.sum()
        loss.backward()

    results = triton.testing.do_bench(
        flash_forward_backward,
        rep=10000,
        warmup=1000
    )
    print(results)

在你本地测试时，可以适当缩短 repetition 和 warmup 时间（单位为毫秒）

你可以考虑以下方向来进一步提升性能：

• 调整 kernel 的 tile size（推荐使用 Triton autotune）
• 微调更多 Triton kernel 的配置参数
• 使用 Triton 实现反向传播，而不仅仅依赖 torch.compile（见 Algorithm 2）
• 在反向传播中对输入执行 两次遍历：一次用于计算 $\mathbf{dQ}$，另一次用于计算 $\mathbf{dK}$ 和 $\mathbf{dV}$，从而避免 block 之间的原子操作或同步
• 在 casual masking 场景下，提前终止 某些 program instance，跳过那些必然全为零的 tile
• 将 非 mask 的 tile 与对角线 tile 分开处理：前者完全不需要索引比较，后者只需一次比较
• 在 H100 上使用 TMA（Tensor Memory Accelerator） 功能，可参考相关教程 [tutorial] 采用类似模式

Note：对于打榜大家如果有想法的话可以自己试试，这里博主就跳过了

博主在当前硬件资源的参数配置下尝试了作业推荐的一些优化，主要包括：

• autotune 调整 kernel 的 tile size 及其他配置参数
• Triton 实现反向传播且对输入执行两次遍历
• causal mask 的特殊处理

下面我们就来看看这些优化对 FlashAttention-2 的性能提升表现如何

我们先来看下之前实现的 FlashAttention-2 的性能表现，博主写了一个基准测试脚本 bench_flash_leaderboard.py 来测试 FlashAttention-2 的性能，代码如下：

import torch
import triton.testing
from pathlib import Path

from cs336_systems.flash_triton import FlashAttention2Triton
from cs336_systems.utils import LeaderboardBenchRow, LeaderboardBenchmarkReporter


def bench_once(fn, warmup=200, rep=1000):
    torch.cuda.synchronize()
    return float(triton.testing.do_bench(fn, warmup=warmup, rep=rep))


def main(
    out_jsonl="runs/leaderboard_ablation.jsonl", 
    out_md="runs/leaderboard_ablation.md",
    warmup=200,
    rep=1000,
    *,
    variant="baseline"
):
    out_jsonl = Path(out_jsonl)
    out_md = Path(out_md)
    out_jsonl.parent.mkdir(parents=True, exist_ok=True)

    B = 1
    S = 4096
    H = 16
    Dh = 64
    dtype = torch.bfloat16
    device = "cuda"

    q = torch.randn(B * H, S, Dh, device=device, dtype=dtype, requires_grad=True)
    k = torch.randn(B * H, S, Dh, device=device, dtype=dtype, requires_grad=True)
    v = torch.randn(B * H, S, Dh, device=device, dtype=dtype, requires_grad=True)    

    def flash(q, k, v):
        return FlashAttention2Triton.apply(q, k, v, True)

    try:
        # fwd
        def fwd():
            return flash(q, k, v)

        fwd_ms = bench_once(fwd, warmup=warmup, rep=rep)

        # bwd (fixed graph)
        o = flash(q, k, v)
        loss = o.sum()

        def bwd():
            q.grad = None
            k.grad = None
            v.grad = None
            loss.backward(retain_graph=True)

        bwd_ms = bench_once(bwd, warmup=warmup, rep=rep)

        # e2e (rebuild graph)
        def e2e():
            qx = q.detach().requires_grad_(True)
            kx = k.detach().requires_grad_(True)
            vx = v.detach().requires_grad_(True)
            oy = flash(qx, kx, vx)
            oy.sum().backward()

        e2e_ms = bench_once(e2e, warmup=warmup, rep=rep)

        status = "ok"

    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        fwd_ms = bwd_ms = e2e_ms = None
        status = "oom"

    except Exception as e:
        print(e)
        fwd_ms = bwd_ms = e2e_ms = None
        status = f"error:{type(e).__name__}"

    reporter = LeaderboardBenchmarkReporter(
        out_jsonl,
        out_md,
        title="#### FlashAttention-2 leaderboard ablation (batch=1, causal=True)",
    )

    row = LeaderboardBenchRow(
        variant=variant,
        dtype="bf16",
        seq_len=S,
        n_heads=H,
        d_head=Dh,
        fwd_ms=fwd_ms,
        bwd_ms=bwd_ms,
        e2e_ms=e2e_ms,
        status=status,
    )

    msg = f"[{variant:12s}] B={B:1d}, S={S:5d}, H={H:2d}, Dh={Dh:3d} -> {row.status}"
    if row.fwd_ms is not None:
        msg += f", fwd={row.fwd_ms:.3f}ms, bwd={row.bwd_ms:.3f}ms, e2e={row.e2e_ms:.3f}ms"
    print(msg)

    reporter.append(row)
    reporter.write_markdown()

if __name__ == "__main__":
    main()

Note：博主测试的主机是 RTX3030（12GB），在当前硬件资源下的测试的参数配置如下：

• batch_size = 1
• seq_len = 4096
• num_heads = 16
• d_head = 64

执行 bench_flash_leaderboard.py 后的输出如下图所示：

可以看到目前未做优化时的 forward 耗时、bacdword 耗时以及端到端耗时分别是：

impl	batch_size	seq_len	num_heads	d_head	fwd_ms	bwd_ms	e2e_ms
baseline	1	4096	16	64	7.090	88.848	95.497

这个结果可以作为我们的 baseline 与后续优化进行对比

6.1 优化 (autotune)

OK，接下来我们来看第一个优化：autotune 调整 kernel 的 tile size 及其他配置参数

目前我们的 tile size 给的是 Bq=32, Bk=32，num_warps 给的是 4，这个配置可能并非最优，我们需要利用 autotune 尽可能找一个最优的 kernel 配置参数，那具体该怎么做呢？

我们需要调整下之前 flash_triton 的代码，首先加入 configs，如下所示：

_fwd_configs = [
    triton.Config({"Q_TILE_SIZE": 32,  "K_TILE_SIZE": 32}, num_warps=4, num_stages=2),
    triton.Config({"Q_TILE_SIZE": 64,  "K_TILE_SIZE": 32}, num_warps=4, num_stages=2),
    triton.Config({"Q_TILE_SIZE": 64,  "K_TILE_SIZE": 64}, num_warps=4, num_stages=2),
    triton.Config({"Q_TILE_SIZE": 128, "K_TILE_SIZE": 32}, num_warps=4, num_stages=2),
    triton.Config({"Q_TILE_SIZE": 128, "K_TILE_SIZE": 64}, num_warps=4, num_stages=2),
    triton.Config({"Q_TILE_SIZE": 32,  "K_TILE_SIZE": 32}, num_warps=4, num_stages=3),
    triton.Config({"Q_TILE_SIZE": 64,  "K_TILE_SIZE": 32}, num_warps=4, num_stages=3),
    triton.Config({"Q_TILE_SIZE": 64,  "K_TILE_SIZE": 64}, num_warps=4, num_stages=3),
    triton.Config({"Q_TILE_SIZE": 128, "K_TILE_SIZE": 32}, num_warps=4, num_stages=3),
    triton.Config({"Q_TILE_SIZE": 128, "K_TILE_SIZE": 64}, num_warps=4, num_stages=3),
]

@triton.autotune(
    configs=_fwd_configs,
    key=["N_QUERIES", "N_KEYS", "D"],
)
@triton.jit
def flash_fwd_kernel(
    Q_ptr, K_ptr, V_ptr,
    O_ptr, L_ptr,
    stride_qb, stride_qq, stride_qd,
    stride_kb, stride_kk, stride_kd,
    stride_vb, stride_vk, stride_vd,
    stride_ob, stride_oq, stride_od,
    stride_lb, stride_lq,
    N_QUERIES: tl.constexpr,
    N_KEYS: tl.constexpr,
    scale,
    D: tl.constexpr,
    Q_TILE_SIZE: tl.constexpr,
    K_TILE_SIZE: tl.constexpr,
    IS_CAUSAL: tl.constexpr
):
    ...

然后调整 grid、Q_TILE_SIZE 等参数的传入：

class FlashAttention2Triton(torch.autograd.Function):
    @staticmethod
    def forward(ctx, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, is_causal: bool = False):
        ...
        # grid = (triton.cdiv(Q, Bq), B)
        grid = lambda META: (triton.cdiv(Q, META["Q_TILE_SIZE"]), B)

        flash_fwd_kernel[grid](
            q, k, v,
            o, L,
            q.stride(0), q.stride(1), q.stride(2),
            k.stride(0), k.stride(1), k.stride(2),
            v.stride(0), v.stride(1), v.stride(2),
            o.stride(0), o.stride(1), o.stride(2),
            L.stride(0), L.stride(1),
            N_QUERIES=Q,
            N_KEYS=K,
            scale=scale,
            D=D,
            # Q_TILE_SIZE=Bq,
            # K_TILE_SIZE=Bk,
            IS_CAUSAL=is_causal,
            # num_warps=4
        )                

Note：configs 可以自行调整，但需要注意 configs 越多，autotune 越慢，因为它会对每个 config 做小基准，因此建议先用小集合找个大概，再逐步加密搜索

接着我们可以写一个 autotune.py 测试脚本，调用一次 forward 过程，打印出 autotune 选出的最优的配置，代码实现如下：

import torch

import cs336_systems.flash_triton as ft
from cs336_systems.flash_triton import FlashAttention2Triton


def flash(q, k, v):
    return FlashAttention2Triton.apply(q, k, v, True)


def main():
    B = 1
    S = 4096
    H = 16
    Dh = 64
    dtype = torch.bfloat16
    device = "cuda"

    q = torch.randn(B * H, S, Dh, device=device, dtype=dtype, requires_grad=True)
    k = torch.randn(B * H, S, Dh, device=device, dtype=dtype, requires_grad=True)
    v = torch.randn(B * H, S, Dh, device=device, dtype=dtype, requires_grad=True)

    _ = flash(q, k, v)
    torch.cuda.synchronize()

    kern = ft.flash_fwd_kernel

    print("kernel type:", type(kern))
    print("keys:", getattr(kern, "keys", None))
    print("cache size:", len(getattr(kern, "cache", {})))

    bc = getattr(kern, "best_config", None)
    if bc is not None:
        print("best.kwargs:", bc.kwargs)
        print("best.num_warps:", bc.num_warps, "best.num_stages:", bc.num_stages)

    cache = getattr(kern, "cache", {})
    if cache:
        for key, best in cache.items():
            print("\nkey:", key)
            print("  best.kwargs:", best.kwargs)
            print("  best.num_warps:", best.num_warps, "best.num_stages:", best.num_stages)

    timings = getattr(kern, "configs_timings", None)
    if timings:
        print("\nconfigs_timings:")
        for cfg, ts in timings.items():
            print("  cfg.kwargs:", cfg.kwargs, "warps:", cfg.num_warps, "stages:", cfg.num_stages, "times:", ts)


if __name__ == "__main__":
    main()

执行 uv run cs336_systems/autotune.py 脚本后输出如下所示：

可以看到在博主当前硬件以及参数下的最优配置是：

• Q_TILE_SIZE = 64
• K_TILE_SIZE = 64
• num_warps = 4
• num_stages = 3

我们可以把最优的参数固定写入到 forward 中，如下所示：

class FlashAttention2Triton(torch.autograd.Function):
    @staticmethod
    def forward(ctx, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, is_causal: bool = False):
        ...

        # tile sizes
        Bq = 64
        Bk = 64

        flash_fwd_kernel[grid](
            q, k, v,
            o, L,
            q.stride(0), q.stride(1), q.stride(2),
            k.stride(0), k.stride(1), k.stride(2),
            v.stride(0), v.stride(1), v.stride(2),
            o.stride(0), o.stride(1), o.stride(2),
            L.stride(0), L.stride(1),
            N_QUERIES=Q,
            N_KEYS=K,
            scale=scale,
            D=D,
            Q_TILE_SIZE=Bq,
            K_TILE_SIZE=Bk,
            IS_CAUSAL=is_causal,
            num_warps=4,
            num_stages=3
        )

此时我们再来执行下 bench_flash_leaderboard.py 性能测试看下提升如何：

可以看到 forward 部分性能有一些小提升，目前的性能提升对比表如下：

impl	batch_size	seq_len	num_heads	d_head	fwd_ms	bwd_ms	e2e_ms
baseline	1	4096	16	64	7.090	88.848	95.497
+autotune	1	4096	16	64	6.087	89.004	94.900

6.2 优化 (triton-backward)

我们接着来看作业中提到的第二个优化点：Triton 实现反向传播且对输入执行两次遍历

代码实现如下：

@triton.jit
def flash_bwd_dq_kernel(
    Q_ptr, K_ptr, V_ptr, O_ptr, DO_ptr,
    L_ptr,
    DQ_ptr,
    stride_qb, stride_qq, stride_qd,
    stride_kb, stride_kk, stride_kd,
    stride_vb, stride_vk, stride_vd,
    stride_ob, stride_oq, stride_od,
    stride_dob, stride_doq, stride_dod,
    stride_lb, stride_lq,
    stride_dqb, stride_dqq, stride_dqd,
    N_QUERIES: tl.constexpr,
    N_KEYS: tl.constexpr,
    scale,
    D: tl.constexpr,
    Q_TILE_SIZE: tl.constexpr,
    K_TILE_SIZE: tl.constexpr,
    IS_CAUSAL: tl.constexpr,
):
    pid_q = tl.program_id(0)  # q tile
    pid_b = tl.program_id(1)  # batch

    # offsets for causal and reductions
    q_idx = pid_q * Q_TILE_SIZE + tl.arange(0, Q_TILE_SIZE)   # [Bq]

    # base pointers
    Qb  = Q_ptr + pid_b * stride_qb
    Kb  = K_ptr  + pid_b * stride_kb
    Vb  = V_ptr  + pid_b * stride_vb
    Ob  = O_ptr  + pid_b * stride_ob
    DOb = DO_ptr + pid_b * stride_dob
    Lb  = L_ptr  + pid_b * stride_lb
    DQb = DQ_ptr + pid_b * stride_dqb

    # (Q, D) block pointers for q/o/do/dq/l
    Q_bp = tl.make_block_ptr(
        Qb,
        shape=(N_QUERIES, D),
        strides=(stride_qq, stride_qd),
        offsets=(pid_q * Q_TILE_SIZE, 0),
        block_shape=(Q_TILE_SIZE, D),
        order=(1, 0),
    )
    O_bp = tl.make_block_ptr(
        Ob,
        shape=(N_QUERIES, D),
        strides=(stride_oq, stride_od),
        offsets=(pid_q * Q_TILE_SIZE, 0), 
        block_shape=(Q_TILE_SIZE, D), 
        order=(1, 0),
    )
    DO_bp = tl.make_block_ptr(
        DOb, 
        shape=(N_QUERIES, D), 
        strides=(stride_doq, stride_dod),
        offsets=(pid_q * Q_TILE_SIZE, 0),
        block_shape=(Q_TILE_SIZE, D),
        order=(1, 0),
    )
    DQ_bp = tl.make_block_ptr(
        DQb, 
        shape=(N_QUERIES, D), 
        strides=(stride_dqq, stride_dqd),
        offsets=(pid_q * Q_TILE_SIZE, 0),
        block_shape=(Q_TILE_SIZE, D),
        order=(1, 0),
    )
    L_bp = tl.make_block_ptr(
        base=Lb,
        shape=(N_QUERIES,),
        strides=(stride_lq,),
        offsets=(pid_q * Q_TILE_SIZE,),
        block_shape=(Q_TILE_SIZE,),
        order=(0,),
    )

    # load Q, dO, O, L
    # keep the original dtype for the final store cast
    q_raw = tl.load(Q_bp,  boundary_check=(0, 1), padding_option="zero")
    q  = q_raw.to(tl.float32)
    do = tl.load(DO_bp, boundary_check=(0, 1), padding_option="zero").to(tl.float32)
    o  = tl.load(O_bp,  boundary_check=(0, 1), padding_option="zero").to(tl.float32)
    L  = tl.load(L_bp,  boundary_check=(0,),   padding_option="zero").to(tl.float32)

    D_row = tl.sum(do * o, axis=1)  # [Bq]
    dq_acc = tl.zeros((Q_TILE_SIZE, D), tl.float32)

    # (K, D) block pointers for sweeping K/V
    K_bp = tl.make_block_ptr(
        Kb, 
        shape=(N_KEYS, D),
        strides=(stride_kk, stride_kd),
        offsets=(0, 0),
        block_shape=(K_TILE_SIZE, D),
        order=(1, 0),
    )
    V_bp = tl.make_block_ptr(
        Vb,
        shape=(N_KEYS, D),
        strides=(stride_vk,stride_vd),
        offsets=(0, 0), 
        block_shape=(K_TILE_SIZE, D), 
        order=(1, 0),
    )

    K_it = K_bp
    V_it = V_bp

    # sweep K/V tiles
    for kb in range(0, N_KEYS, K_TILE_SIZE):
        k_idx = kb + tl.arange(0, K_TILE_SIZE)  # [Bk]

        k = tl.load(K_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bk, D]
        v = tl.load(V_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bk, D]

        # S: [Bq, Bk]
        S = tl.dot(q, tl.trans(k)) * scale
        if IS_CAUSAL:
            S = tl.where(q_idx[:, None] >= k_idx[None, :], S, -1.0e6)

        # P = exp(S - L)
        P = tl.exp(S - L[:, None])  # [Bq, Bk]

        # dP = dO @ V^T
        dP = tl.dot(do, tl.trans(v))  # [Bq, Bk]

        # dS = P * (dP - D_row)
        dS = P * (dP - D_row[:, None])  # [Bq, Bk]

        # dQ += dS @ K * scale
        dq_acc += tl.dot(dS, k) * scale

        # advance K/V block pointers to the next tile along the sequence dimension
        K_it = K_it.advance((K_TILE_SIZE, 0))
        V_it = V_it.advance((K_TILE_SIZE, 0))

    tl.store(DQ_bp, dq_acc.to(q_raw.dtype), boundary_check=(0, 1))


@triton.jit
def flash_bwd_dkdv_kernel(
    Q_ptr, K_ptr, V_ptr, O_ptr, DO_ptr,
    L_ptr,
    DK_ptr, DV_ptr,
    stride_qb, stride_qq, stride_qd,
    stride_kb, stride_kk, stride_kd,
    stride_vb, stride_vk, stride_vd,
    stride_ob, stride_oq, stride_od,
    stride_dob, stride_doq, stride_dod,
    stride_lb, stride_lq,
    stride_dkb, stride_dkk, stride_dkd,
    stride_dvb, stride_dvk, stride_dvd,
    N_QUERIES: tl.constexpr,
    N_KEYS: tl.constexpr,
    scale,
    D: tl.constexpr,
    Q_TILE_SIZE: tl.constexpr,
    K_TILE_SIZE: tl.constexpr,
    IS_CAUSAL: tl.constexpr,
):
    pid_k = tl.program_id(0)  # k tile
    pid_b = tl.program_id(1)

    # offsets for causal and reductions
    k_idx = pid_k * K_TILE_SIZE + tl.arange(0, K_TILE_SIZE)  # [Bk]

    # base pointers (batch)
    Qb  = Q_ptr  + pid_b * stride_qb
    Kb  = K_ptr  + pid_b * stride_kb
    Vb  = V_ptr  + pid_b * stride_vb
    Ob  = O_ptr  + pid_b * stride_ob
    DOb = DO_ptr + pid_b * stride_dob
    Lb  = L_ptr  + pid_b * stride_lb
    DKb = DK_ptr + pid_b * stride_dkb
    DVb = DV_ptr + pid_b * stride_dvb

    # (K, D) block pointers for this tile
    K_bp = tl.make_block_ptr(
        Kb,
        shape=(N_KEYS, D),
        strides=(stride_kk, stride_kd),
        offsets=(pid_k * K_TILE_SIZE, 0),
        block_shape=(K_TILE_SIZE, D),
        order=(1, 0),
    )
    V_bp = tl.make_block_ptr(
        Vb,
        shape=(N_KEYS, D),
        strides=(stride_vk, stride_vd),
        offsets=(pid_k * K_TILE_SIZE, 0),
        block_shape=(K_TILE_SIZE, D),
        order=(1, 0),
    )
    DK_bp = tl.make_block_ptr(
        DKb,
        shape=(N_KEYS, D),
        strides=(stride_dkk, stride_dkd),
        offsets=(pid_k * K_TILE_SIZE, 0),
        block_shape=(K_TILE_SIZE, D),
        order=(1, 0),
    )
    DV_bp = tl.make_block_ptr(
        DVb,
        shape=(N_KEYS, D),
        strides=(stride_dvk, stride_dvd),
        offsets=(pid_k * K_TILE_SIZE, 0),
        block_shape=(K_TILE_SIZE, D),
        order=(1, 0),
    )    

    # keep the original dtype for the final store cast
    k_raw = tl.load(K_bp, boundary_check=(0, 1), padding_option="zero")
    v_raw = tl.load(V_bp, boundary_check=(0, 1), padding_option="zero")
    k = k_raw.to(tl.float32)
    v = v_raw.to(tl.float32)

    dk_acc = tl.zeros((K_TILE_SIZE, D), tl.float32)
    dv_acc = tl.zeros((K_TILE_SIZE, D), tl.float32)

    # iter (Q, D) pointers for sweeping Q/O/DO/L
    Q_bp = tl.make_block_ptr(
        Qb,
        shape=(N_QUERIES, D),
        strides=(stride_qq, stride_qd),
        offsets=(0, 0),
        block_shape=(Q_TILE_SIZE, D),
        order=(1, 0),
    )
    O_bp = tl.make_block_ptr(
        Ob, 
        shape=(N_QUERIES, D), 
        strides=(stride_oq, stride_od),
        offsets=(0, 0), 
        block_shape=(Q_TILE_SIZE, D), 
        order=(1, 0),
    )
    DO_bp = tl.make_block_ptr(
        DOb,
        shape=(N_QUERIES, D),
        strides=(stride_doq, stride_dod),
        offsets=(0, 0),
        block_shape=(Q_TILE_SIZE, D),
        order=(1, 0),
    )
    L_bp = tl.make_block_ptr(
        base=Lb,
        shape=(N_QUERIES,),
        strides=(stride_lq,),
        offsets=(0,),
        block_shape=(Q_TILE_SIZE,),
        order=(0,),
    )

    Q_it  = Q_bp
    O_it  = O_bp
    DO_it = DO_bp
    L_it  = L_bp

    # sweep Q tiles
    for qb in range(0, N_QUERIES, Q_TILE_SIZE):
        q_idx = qb + tl.arange(0, Q_TILE_SIZE)  # [Bq]

        q  = tl.load(Q_it,  boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bq, D]
        o  = tl.load(O_it,  boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bq, D]
        do = tl.load(DO_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bq, D]
        L  = tl.load(L_it,  boundary_check=(0,),   padding_option="zero").to(tl.float32)  # [Bq]

        D_row = tl.sum(do * o, axis=1)   # [Bq]

        # S: [Bq, Bk]
        S = tl.dot(q, tl.trans(k)) * scale
        if IS_CAUSAL:
            S = tl.where(q_idx[:, None] >= k_idx[None, :], S, -1.0e6)

        # P = exp(S - L)
        P = tl.exp(S - L[:, None])  # [Bq, Bk]

        # dP = dO @ V^T
        dP = tl.dot(do, tl.trans(v))  # [Bq, Bk]

        # dS = P * (dP - D_row)
        dS = P * (dP - D_row[:, None])  # [Bq, Bk]

        # dV += P^T @ dO
        dv_acc += tl.dot(tl.trans(P), do)

        # dK += dS^T @ Q * scale
        dk_acc += tl.dot(tl.trans(dS), q) * scale

        Q_it  = Q_it.advance((Q_TILE_SIZE, 0))
        O_it  = O_it.advance((Q_TILE_SIZE, 0))
        DO_it = DO_it.advance((Q_TILE_SIZE, 0))
        L_it  = L_it.advance((Q_TILE_SIZE,))

    tl.store(DK_bp, dk_acc.to(k_raw.dtype), boundary_check=(0, 1))
    tl.store(DV_bp, dv_acc.to(v_raw.dtype), boundary_check=(0, 1))


class FlashAttention2Triton(torch.autograd.Function):
    @staticmethod
    def backward(ctx, do):
        (L, q, k, v, o) = ctx.saved_tensors
        B, Q, D = q.shape
        K = k.shape[1]
        scale = 1.0 / math.sqrt(D)
        is_causal = ctx.is_causal

        dq = torch.empty_like(q)
        dk = torch.empty_like(k)
        dv = torch.empty_like(v)

        Bq = 32
        Bk = 32

        grid_dq = (triton.cdiv(Q, Bq), B)
        flash_bwd_dq_kernel[grid_dq](
            q, k, v, o, do,
            L,
            dq,
            q.stride(0), q.stride(1), q.stride(2),
            k.stride(0), k.stride(1), k.stride(2),
            v.stride(0), v.stride(1), v.stride(2),
            o.stride(0), o.stride(1), o.stride(2),
            do.stride(0), do.stride(1), do.stride(2),
            L.stride(0), L.stride(1),
            dq.stride(0), dq.stride(1), dq.stride(2),
            N_QUERIES=Q,
            N_KEYS=K,
            scale=scale,
            D=D,
            Q_TILE_SIZE=Bq,
            K_TILE_SIZE=Bk,
            IS_CAUSAL=is_causal,
            num_warps=4,
        )

        grid_dkdv = (triton.cdiv(K, Bk), B)
        flash_bwd_dkdv_kernel[grid_dkdv](
            q, k, v, o, do,
            L,
            dk, dv,
            q.stride(0), q.stride(1), q.stride(2),
            k.stride(0), k.stride(1), k.stride(2),
            v.stride(0), v.stride(1), v.stride(2),
            o.stride(0), o.stride(1), o.stride(2),
            do.stride(0), do.stride(1), do.stride(2),
            L.stride(0), L.stride(1),
            dk.stride(0), dk.stride(1), dk.stride(2),
            dv.stride(0), dv.stride(1), dv.stride(2),
            N_QUERIES=Q,
            N_KEYS=K,
            scale=scale,
            D=D,
            Q_TILE_SIZE=Bq,
            K_TILE_SIZE=Bk,
            IS_CAUSAL=is_causal,
            num_warps=4,
        )

        return dq, dk, dv, None    

还记得我们之前提到的在 backward 阶段采用的是 recomput 路线吗？也就是不在 forward 保存完整的注意力矩阵 $S$ 或概率矩阵 $P$，而是在 backward 里按 tile 重新计算，并结合上游梯度 $dO$ 与 forward 输出 $O$、logsumexp $L$ 推导出 $dQ,dK,dV$

Triton kernel 实现和我们之前的 PyTorch 实现 flash_bwd_recompute_impl 的 数学路径完全一致，但 Triton 实现把这些操作改成了 按块加载 + kernel 内融合 + GPU 并行 来做，从而避免了巨大的中间张量与多次 kernel launch 的开销

我们的 backward 实现为 两个 Triton kernel：

• flash_bwd_dq_kernel：计算 $dQ$，grid 按 (query, batch) 并行，每个 program 负责一个 (Bq, D) 的 dQ tile，并在 kernel 内沿着 K 维 sweep 全部 key tiles，逐步累加得到该 tile 的 dQ
• flash_bwd_dkdv_kernel：计算 $dK,dV$，grid 按 (key_tile, batch) 并行，每个 program 固定一个 (Bk, D) 的 dK/dV tile，然后在 kernel 内沿着 Q 维 sweep 全部 query tiles，逐步累加得到该 tile 的 dK/dV

这也是作业中提到的优化点：在反向传播中对输入执行两次遍历：一次用于计算 $\mathbf{dQ}$，另一次用于计算 $\mathbf{dK}$ 和 $\mathbf{dV}$，从而避免 block 之间的原子操作或同步

backward 拆成两个 kernel 这件事其实非常自然，因为 $dQ$ 的自然并行维度是 Q（每个 query 行/块独立累加），而 $dK,dV$ 的自然并行维度是 K（每个 key 行/块独立累加），如果试图在同一个 kernel 里同时并行计算三者，会出现上提到的 写冲突与原子累加问题，因为它们的复用模式相反，难以同时高效。把它们拆开之后，每个 kernel 都能选择对自己最友好的数据驻留方式，从而避免原子操作，同时最大化 tile 复用与吞吐

代码实现比较简单，我们来主要来看下它和 Algorithm 2 中的公式是如何对应上的：

1) flash_bwd_dq_kernel：固定 Q tile，扫 K tiles 累加 dQ

# load Q, dO, O, L
# keep the original dtype for the final store cast
q_raw = tl.load(Q_bp,  boundary_check=(0, 1), padding_option="zero")
q  = q_raw.to(tl.float32)
do = tl.load(DO_bp, boundary_check=(0, 1), padding_option="zero").to(tl.float32)
o  = tl.load(O_bp,  boundary_check=(0, 1), padding_option="zero").to(tl.float32)
L  = tl.load(L_bp,  boundary_check=(0,),   padding_option="zero").to(tl.float32)

D_row = tl.sum(do * o, axis=1)  # [Bq]
dq_acc = tl.zeros((Q_TILE_SIZE, D), tl.float32)

这个 kernel 会加载当前 query tile 对应的 Q / O / dO / L，并计算每一行的 $D_{\text{row}}={\sum }_{d}dO\odot O$ 作为 softmax backward 的常用中间量

# (K, D) block pointers for sweeping K/V
K_bp = tl.make_block_ptr(
    Kb, 
    shape=(N_KEYS, D),
    strides=(stride_kk, stride_kd),
    offsets=(0, 0),
    block_shape=(K_TILE_SIZE, D),
    order=(1, 0),
)
V_bp = tl.make_block_ptr(
    Vb,
    shape=(N_KEYS, D),
    strides=(stride_vk,stride_vd),
    offsets=(0, 0), 
    block_shape=(K_TILE_SIZE, D), 
    order=(1, 0),
)

K_it = K_bp
V_it = V_bp

接着用 block pointer 的方式构造 (K, D) 的 K_bp / V_bp，并通过 advance((K_TILE_SIZE, 0)) 在循环中推进到下一个 key tile

# sweep K/V tiles
for kb in range(0, N_KEYS, K_TILE_SIZE):
    k_idx = kb + tl.arange(0, K_TILE_SIZE)  # [Bk]

    k = tl.load(K_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bk, D]
    v = tl.load(V_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bk, D]

    # S: [Bq, Bk]
    S = tl.dot(q, tl.trans(k)) * scale
    if IS_CAUSAL:
        S = tl.where(q_idx[:, None] >= k_idx[None, :], S, -1.0e6)

    # P = exp(S - L)
    P = tl.exp(S - L[:, None])  # [Bq, Bk]

    # dP = dO @ V^T
    dP = tl.dot(do, tl.trans(v))  # [Bq, Bk]

    # dS = P * (dP - D_row)
    dS = P * (dP - D_row[:, None])  # [Bq, Bk]

    # dQ += dS @ K * scale
    dq_acc += tl.dot(dS, k) * scale

    # advance K/V block pointers to the next tile along the sequence dimension
    K_it = K_it.advance((K_TILE_SIZE, 0))
    V_it = V_it.advance((K_TILE_SIZE, 0))

每次迭代加载一个 (Bk, D) 的 K, V tile，重新计算：

$$\begin{array}{rl}S& =Q{K}^{\top }\cdot \text{scale}\\ P& =\exp (S-L)\\ dP& =dO\cdot V^{\top }\\ dS& =P\odot (dP-D_{\text{row}})\\ dQ& +=dS\cdot K\cdot \text{scale}\end{array}$$

对应的核心更新逻辑在代码里就是上面这一段，并且循环末尾用 block pointer 前进

tl.store(DQ_bp, dq_acc.to(q_raw.dtype), boundary_check=(0, 1))

最后把 dq_acc 写回 DQ

值得注意的是，最后写回的梯度张量 dq 我们进行了强制转换，让它的 dtype 与输入张量保持一致。这是因为在 bfloat16 精度下如果省略这一步显式转换，会导致 Triton 编译失败（store dtype 不匹配），因此，本实现统一采用 float32 累加 + 按输入 dtype 写回 的策略，以同时保证数值稳定性与混合精度训练的正确性。

2) flash_bwd_dkdv_kernel：固定 K tile，扫 Q tiles 累加 dK 和 dV

# keep the original dtype for the final store cast
k_raw = tl.load(K_bp, boundary_check=(0, 1), padding_option="zero")
v_raw = tl.load(V_bp, boundary_check=(0, 1), padding_option="zero")
k = k_raw.to(tl.float32)
v = v_raw.to(tl.float32)

dk_acc = tl.zeros((K_TILE_SIZE, D), tl.float32)
dv_acc = tl.zeros((K_TILE_SIZE, D), tl.float32)

另一个 kernel 的策略反过来：program grid 固定一个 k_idx tile，先把该 tile 的 K, V load 到寄存器/片上，并初始化 dk_acc/dv_acc

# iter (Q, D) pointers for sweeping Q/O/DO/L
Q_bp = tl.make_block_ptr(
    Qb,
    shape=(N_QUERIES, D),
    strides=(stride_qq, stride_qd),
    offsets=(0, 0),
    block_shape=(Q_TILE_SIZE, D),
    order=(1, 0),
)
O_bp = tl.make_block_ptr(
    Ob, 
    shape=(N_QUERIES, D), 
    strides=(stride_oq, stride_od),
    offsets=(0, 0), 
    block_shape=(Q_TILE_SIZE, D), 
    order=(1, 0),
)
DO_bp = tl.make_block_ptr(
    DOb,
    shape=(N_QUERIES, D),
    strides=(stride_doq, stride_dod),
    offsets=(0, 0),
    block_shape=(Q_TILE_SIZE, D),
    order=(1, 0),
)
L_bp = tl.make_block_ptr(
    base=Lb,
    shape=(N_QUERIES,),
    strides=(stride_lq,),
    offsets=(0,),
    block_shape=(Q_TILE_SIZE,),
    order=(0,),
)

Q_it  = Q_bp
O_it  = O_bp
DO_it = DO_bp
L_it  = L_bp

然后用 (Q, D) 的 block points（Q_it/O_it/DO_it/L_it）沿着 Q 维 sweep

# sweep Q tiles
for qb in range(0, N_QUERIES, Q_TILE_SIZE):
    q_idx = qb + tl.arange(0, Q_TILE_SIZE)  # [Bq]

    q  = tl.load(Q_it,  boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bq, D]
    o  = tl.load(O_it,  boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bq, D]
    do = tl.load(DO_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bq, D]
    L  = tl.load(L_it,  boundary_check=(0,),   padding_option="zero").to(tl.float32)  # [Bq]

    D_row = tl.sum(do * o, axis=1)   # [Bq]

    # S: [Bq, Bk]
    S = tl.dot(q, tl.trans(k)) * scale
    if IS_CAUSAL:
        S = tl.where(q_idx[:, None] >= k_idx[None, :], S, -1.0e6)

    # P = exp(S - L)
    P = tl.exp(S - L[:, None])  # [Bq, Bk]

    # dP = dO @ V^T
    dP = tl.dot(do, tl.trans(v))  # [Bq, Bk]

    # dS = P * (dP - D_row)
    dS = P * (dP - D_row[:, None])  # [Bq, Bk]

    # dV += P^T @ dO
    dv_acc += tl.dot(tl.trans(P), do)

    # dK += dS^T @ Q * scale
    dk_acc += tl.dot(tl.trans(dS), q) * scale

    Q_it  = Q_it.advance((Q_TILE_SIZE, 0))
    O_it  = O_it.advance((Q_TILE_SIZE, 0))
    DO_it = DO_it.advance((Q_TILE_SIZE, 0))
    L_it  = L_it.advance((Q_TILE_SIZE,))

每次加载一个 query tile 的 Q/O/dO/L，算出 D_row，在 recompute 本 tile 的 S,P,dP,dS，并做两条累加：

$$\begin{array}{rl}dV& +=P^{\top }\cdot dO\\ dK& +=d{S}^{\top }\cdot Q\cdot \text{scale}\end{array}$$

这部分实现对应上面这段代码

tl.store(DK_bp, dk_acc.to(k_raw.dtype), boundary_check=(0, 1))
tl.store(DV_bp, dv_acc.to(v_raw.dtype), boundary_check=(0, 1))

最后把 dk_acc、dv_acc 分别写回 DQ 和 DV

OK，我们来看下将 backward 改成 Triton 实现的版本后性能如何：

可以看到提升相当可观，backward 耗时相比之前降低了约 2/3，端到端耗时也跟着下降，目前的性能提升对比表如下：

impl	batch_size	seq_len	num_heads	d_head	fwd_ms	bwd_ms	e2e_ms
baseline	1	4096	16	64	7.090	88.848	95.497
+autotune	1	4096	16	64	6.087	89.004	94.900
+autotune+triton-bwd	1	4096	16	64	6.081	31.706	37.329

6.3 优化 (causal mask)

我们接着来看作业中提到的第三个优化点：causal mask 的特殊处理

我们先来理解下作业中提到的 causal mask 这部分为什么能被优化：

1) 先回顾下 attention 中 causal mask 的概念

标准 attention 的 logits 是：

$$S_{ij}=(q_i\cdot k_j)\cdot \text{scale}$$

softmax 得到权重 $P_{ij}$，再加权求和得到输出：

$$P_{ij}=\text{softmax}(S_{i,:}{)}_j,o_i=\sum _j{P}_{ij}{v}_j$$

causal masking 的限制是：第 $i$ 个 query 只能看见它 “之前（含自己）” 的位置：

$$j\le i\Rightarrow \text{可见},j>i\Rightarrow \text{不可见}$$

通常实现上是把不可见位置的 logits 设为 $-\infty$，这样 softmax 后的概率约为 0

如果把 $S$ 看出一个 $Q\times K$ 的矩阵（行是 query index $i$，列是 key index $j$），causal mask 会把 上三角 全部屏蔽：

j (keys) →
i (queries) ↓

      0 1 2 3 4 5 6 ...
0     ✔ x x x x x x
1     ✔ ✔ x x x x x
2     ✔ ✔ ✔ x x x x
3     ✔ ✔ ✔ ✔ x x x
4     ✔ ✔ ✔ ✔ ✔ x x
...
(✔: 需要算; x: 结果必为 0)

因此，如果我们还像 non-causal 那样把整个矩阵算一遍，然后再用 mask 把上三角置零，那相当于做了接近一半的无用功

2) 进一步：FlashAttention 的 tile 视角下， 上三角 tile 是整块无效的

我们知道 FlashAttention 不会显式构造整个 $Q\times K$ 矩阵，而是按 tile 分块计算，假设：

• 一个 Q tile 覆盖连续的 $B_q$ 个 query（比如 32）
• 一个 K tile 覆盖连续的 $B_k$ 个 key（比如 32）

那么注意力矩阵就被切分成很多小方块（tiles）：

       K tiles →
     [0] [1] [2] [3] [4] ...
Q  [0]  D   X   X   X   X
t  [1]  N   D   X   X   X
i  [2]  N   N   D   X   X
l  [3]  N   N   N   D   X
e  ...

这里的含义是：

• N（non-mask tile，严格在对角线下方）：这个 tile 里的每个 $(i,j)$ 都满足 $j<i$，因此 完全不需要做 mask 比较，直接算就行
• D（diagonal tile，对角线 tile）：tile 与对角线相交，tile 内一部分可见、一部分不可见，因此 需要做 mask 比较（例如 q_idx >= k_idx）
• X（fully-masked tile，严格在对角线上方）：这个 tile 内所有 $(i,j)$ 都满足 $j>i$，因此整块都不可见，softmax 权重为 0，这整个 tile 的 dot/exp/归约都是浪费，应该直接跳过

这就是作业中提到的两条优化点：

1. 提前终止 / 跳过必然全为 0 的 tiles

2. 把 non-mask tiles 和 diagonal tile 分开处理：前者不做索引比较，后者才做一次比较

通过这种方式，我们在 “算力 + 访存” 两个方面都能得到显著优化：

• 少算：跳过上三角 tile，减少大量 QK^T / PV 的矩阵乘
• 少访问：不再加载无效的 K/V（或 Q/O/DO）

OK，接下来我们来看下这些优化点在代码中该如何实现，我们以 flash_fwd_kernel 为例，backward 的两个 kernel 类似：

首先我们来看下没有做 causal mask 优化时的代码：

for kb in range(0, N_KEYS, K_TILE_SIZE):
    # load one (K_TILE_SIZE, D) tile of K and V
    k = tl.load(K_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bk, D]
    v = tl.load(V_it, boundary_check=(0, 1), padding_option="zero")                 # [Bk, D]

    # S = q @ k^T * scale -> [Bq, Bk]
    S = tl.dot(q, tl.trans(k)) * scale  # float32

    # causal mask: keep if q_idx >= k_idx else -1e-6
    if IS_CAUSAL:
        k_abs = kb + tl.arange(0, K_TILE_SIZE)
        S = tl.where(q_abs[:, None] >= k_abs[None, :], S, -1.0e6)

    # online softmax update
    m_new = tl.maximum(m, tl.max(S, axis=1))  # [Bq]
    p = tl.exp(S - m_new[:, None])            # [Bq, Bk]

    alpha = tl.exp(m - m_new)                 # [Bq]
    l_new = alpha * l + tl.sum(p, axis=1)     # [Bq]

    # acc = alpha * acc + p @ v
    # p needs to match v dtype before dot
    p = p.to(v.dtype)
    acc = alpha[:, None] * acc
    acc = tl.dot(p, v, acc=acc)

    m = m_new
    l = l_new

    # advance K/V block pointers to the next tile along the sequence dimension
    K_it = K_it.advance((K_TILE_SIZE, 0))
    V_it = V_it.advance((K_TILE_SIZE, 0))

在不做 causal 专门优化时，一个 pid_q（一个 Q tile）通常会：

• 依次扫过所有 K tiles：for kb in range(0, N_KEYS, K_TILE_SIZE)
• 每个 tile 都做：S = Q @ K^T → mask（对每个元素 q_idx >= k_idx 比较）→ online softmax 更新 → acc += P @ V

这就对应我们在前面提到的浪费：对角线之后（上三角）的 tile 其实完全不可见，但仍然做了 load/dot/softmax，只是最后被 mask 掉

我们接着看下 causal mask 优化后的代码：

if IS_CAUSAL:
    # Phase 1: Non-diagonal tiles (fully inside the causal mask)
    # Iterate over K tiles [0, pid_q * Q_TILE_SIZE]
    # These tiles are completely visible to the current Q tile, so no mask is needed.
    limit_nonding = min(N_KEYS, pid_q * Q_TILE_SIZE)
    for _ in range(0, limit_nonding, K_TILE_SIZE):
        k = tl.load(K_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)
        v = tl.load(V_it, boundary_check=(0, 1), padding_option="zero")

        # S = q @ k^T * scale -> [Bq, Bk]
        S = tl.dot(q, tl.trans(k)) * scale  # float32

        # online softmax update (no causal mask check here)
        m_new = tl.maximum(m, tl.max(S, axis=1))  # [Bq]
        p = tl.exp(S - m_new[:, None])            # [Bq, Bk]
        alpha = tl.exp(m - m_new)                 # [Bq]
        l_new = alpha * l + tl.sum(p, axis=1)     # [Bq]

        # acc = alpha * acc + p @ v
        p = p.to(v.dtype)
        acc = alpha[:, None] * acc
        acc = tl.dot(p, v, acc=acc)
        m = m_new
        l = l_new            

        # advance K/V block pointers
        K_it = K_it.advance((K_TILE_SIZE, 0))
        V_it = V_it.advance((K_TILE_SIZE, 0))

    # Phase 2: Diagonal tile (partial causal mask)
    # Only process if the diagonal tile exists (i.e., within N_KEYS)
    if pid_q * Q_TILE_SIZE < N_KEYS:
        k = tl.load(K_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)
        v = tl.load(V_it, boundary_check=(0, 1), padding_option="zero")
        S = tl.dot(q, tl.trans(k)) * scale

        # causal mask: keep if q_idx >= k_idx else -1e-6
        k_abs = pid_q * Q_TILE_SIZE + tl.arange(0, K_TILE_SIZE)
        S = tl.where(q_abs[:, None] >= k_abs[None, :], S, -1.0e6)

        # online softmax update
        m_new = tl.maximum(m, tl.max(S, axis=1))
        p = tl.exp(S - m_new[:, None])
        alpha = tl.exp(m - m_new)
        l_new = alpha * l + tl.sum(p, axis=1)

        # acc = alpha * acc + p @ v
        p = p.to(v.dtype)
        acc = alpha[:, None] * acc
        acc = tl.dot(p, v, acc=acc)
        m = m_new
        l = l_new
        # No need to advace further, as tiles > pid_q are fully masked out
else:
    # Non-causal path: process all tiles
    for _ in range(0, N_KEYS, K_TILE_SIZE):
        # load one (K_TILE_SIZE, D) tile of K and V
        k = tl.load(K_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bk, D]
        v = tl.load(V_it, boundary_check=(0, 1), padding_option="zero")                 # [Bk, D]

        # S = q @ k^T * scale -> [Bq, Bk]
        S = tl.dot(q, tl.trans(k)) * scale  # float32

        # online softmax update
        m_new = tl.maximum(m, tl.max(S, axis=1))  # [Bq]
        p = tl.exp(S - m_new[:, None])            # [Bq, Bk]

        alpha = tl.exp(m - m_new)                 # [Bq]
        l_new = alpha * l + tl.sum(p, axis=1)     # [Bq]

        # acc = alpha * acc + p @ v
        # p needs to match v dtype before dot
        p = p.to(v.dtype)
        acc = alpha[:, None] * acc
        acc = tl.dot(p, v, acc=acc)

        m = m_new
        l = l_new

        # advance K/V block pointers to the next tile along the sequence dimension
        K_it = K_it.advance((K_TILE_SIZE, 0))
        V_it = V_it.advance((K_TILE_SIZE, 0))

Note：causal mask 优化的代码实现来自 GLM-4.7，完整的代码请参考：https://github.com/Melody-Zhou/stanford-cs336-spring2025-assignments

我们现在的实现中把 causal 的处理清晰拆成了两个 phase：

Phase 1：Non-diagonal tiles（完全可见）— 不做 mask 比较

limit_nonding = min(N_KEYS, pid_q * Q_TILE_SIZE)
for _ in range(0, limit_nonding, K_TILE_SIZE):
    ...
    S = tl.dot(q, tl.trans(k)) * scale
    # online softmax update (no causal mask check here)

这部分代码正对应作业建议里：“将非 mask tile 与对角线 tile 分开处理：前者完全不需要索引比较”

为什么 limit_nonding = min(N_KEYS, pid_q * Q_TILE_SIZE) 合理？这是因为当 K tile 的起始位置 < pid_q * Bq 时，意味着这个 tile 对应的 key index 范围整体都在当前 Q tile 的 “过去”，严格满足 k < q（属于下三角），所以 整个 tile 可见，根本不需要 q_idx >= k_idx 的逐元素比较。

Phase 2：Diagonal tile（与对角线相交）— 只对这一块做一次 mask

if pid_q * Q_TILE_SIZE < N_KEYS:
    ...
    S = tl.dot(q, tl.trans(k)) * scale
    k_abs = pid_q * Q_TILE_SIZE + tl.arange(0, K_TILE_SIZE)
    S = tl.where(q_abs[:, None] >= k_abs[None, :], S, -1.0e6)
    # online softmax update

这部分代码对应：“对角线 tile 只需一次索引比较”

注意这里的比较只发生在 一个 tile 上（对每个 pid_q 最多一个对角线 tile），而不是发生在所有 tiles 上。此外，对角线处理之后我们不再继续 advance，而是直接结束 causal 分支：

# No need to advace further, as tiles > pid_q are fully masked out

这就是作业里提到的提前终止 program instance。

OK，我们接着来看反向传播中的两个 kernel 的改动

flash_bwd_dq_kernel 的改动如下：

if IS_CAUSAL:
    # Phase 1: Non-diagonal K tiles [0, pid_q * K_TILE_SIZE)
    limit_nondiag = min(N_KEYS, pid_q * Q_TILE_SIZE)
    for _ in range(0, limit_nondiag, K_TILE_SIZE):
        k = tl.load(K_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)
        v = tl.load(V_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)

        # S: [Bq, Bk]
        S = tl.dot(q, tl.trans(k)) * scale
        
        # No mask needed for k < q
        # P = exp(S - L)
        P = tl.exp(S - L[:, None])  # [Bq, Bk]
        
        # dP = dO @ V^T
        dP = tl.dot(do, tl.trans(v))  # [Bq, Bk]
        
        # dS = P * (dP - D_row)
        dS = P * (dP - D_row[:, None])  # [Bq, Bk]
        
        # dQ += dS @ K * scale
        dq_acc += tl.dot(dS, k) * scale
        K_it = K_it.advance((K_TILE_SIZE, 0))
        V_it = V_it.advance((K_TILE_SIZE, 0))

    # Phase 2: Diagonal K tile (kb = pid_q * K_TILE_SIZE)
    if pid_q * Q_TILE_SIZE < N_KEYS:
        k = tl.load(K_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)
        v = tl.load(V_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)
        S = tl.dot(q, tl.trans(k)) * scale

        # Apply mask for the diagonal
        k_idx = pid_q * Q_TILE_SIZE + tl.arange(0, K_TILE_SIZE)
        S = tl.where(q_idx[:, None] >= k_idx[None, :], S, -1.0e6)
        
        # P = exp(S - L)
        P = tl.exp(S - L[:, None])

        # dP = dO @ V^T
        dP = tl.dot(do, tl.trans(v))

        # dS = P * (dP - D_row)
        dS = P * (dP - D_row[:, None])

        # dQ += dS @ K * scale
        dq_acc += tl.dot(dS, k) * scale
else:
    # Non-causal path
    for _ in range(0, N_KEYS, K_TILE_SIZE):
        k = tl.load(K_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bk, D]
        v = tl.load(V_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bk, D]

        # S: [Bq, Bk]
        S = tl.dot(q, tl.trans(k)) * scale

        # P = exp(S - L)
        P = tl.exp(S - L[:, None])  # [Bq, Bk]

        # dP = dO @ V^T
        dP = tl.dot(do, tl.trans(v))  # [Bq, Bk]

        # dS = P * (dP - D_row)
        dS = P * (dP - D_row[:, None])  # [Bq, Bk]

        # dQ += dS @ K * scale
        dq_acc += tl.dot(dS, k) * scale

        # advance K/V block pointers to the next tile along the sequence dimension
        K_it = K_it.advance((K_TILE_SIZE, 0))
        V_it = V_it.advance((K_TILE_SIZE, 0))

flash_bwd_dq_kernel 的 causal 优化与 forward 完全一致：同样把 K tile 分成 non-diagonal 和 diagonal 两段遍历，并在对角线之后停止遍历，从而避免在被 mask 的上三角区域做无效的 QK^T / softmax 相关计算

flash_bwd_dkdv_kernel 的改动如下：

if IS_CAUSAL:
    # Optimization: Start from the diagonal Q tile (pid_k * Q_TILE_SIZE)
    # Previous Q tiles (i < k) are masked out and contribute 0 gradient to K[k], V[k]
    start_q_offset = pid_k * Q_TILE_SIZE

    # Advance pointers to the start Q tile
    Q_it  = Q_it.advance((start_q_offset, 0))
    O_it  = O_it.advance((start_q_offset, 0))
    DO_it = DO_it.advance((start_q_offset, 0))
    L_it  = L_it.advance((start_q_offset,))
    
    # Phase 1: Diagonal Q tile (needs mask)
    if start_q_offset < N_QUERIES:
        q  = tl.load(Q_it,  boundary_check=(0, 1), padding_option="zero").to(tl.float32)
        o  = tl.load(O_it,  boundary_check=(0, 1), padding_option="zero").to(tl.float32)
        do = tl.load(DO_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)
        L  = tl.load(L_it,  boundary_check=(0,),   padding_option="zero").to(tl.float32)
        D_row = tl.sum(do * o, axis=1)   # [Bq]

        # S: [Bq, Bk]
        S = tl.dot(q, tl.trans(k)) * scale
        
        # Mask: q_idx >= k_idx
        q_idx = start_q_offset + tl.arange(0, Q_TILE_SIZE)
        S = tl.where(q_idx[:, None] >= k_idx[None, :], S, -1.0e6)
        
        # P = exp(S - L)
        P = tl.exp(S - L[:, None])  # [Bq, Bk]
        
        # dP = dO @ V^T
        dP = tl.dot(do, tl.trans(v))  # [Bq, Bk]
        
        # dS = P * (dP - D_row)
        dS = P * (dP - D_row[:, None])  # [Bq, Bk]
        
        # dV += P^T @ dO
        dv_acc += tl.dot(tl.trans(P), do)
        
        # dK += dS^T @ Q * scale
        dk_acc += tl.dot(tl.trans(dS), q) * scale
        
        # Advance to next tile
        Q_it  = Q_it.advance((Q_TILE_SIZE, 0))
        O_it  = O_it.advance((Q_TILE_SIZE, 0))
        DO_it = DO_it.advance((Q_TILE_SIZE, 0))
        L_it  = L_it.advance((Q_TILE_SIZE,))

    # Phase 2: Non-diagonal Q tiles (qb > start_q_offset, no mask)
    for _ in range(start_q_offset + Q_TILE_SIZE, N_QUERIES, Q_TILE_SIZE):
        q  = tl.load(Q_it,  boundary_check=(0, 1), padding_option="zero").to(tl.float32)
        o  = tl.load(O_it,  boundary_check=(0, 1), padding_option="zero").to(tl.float32)
        do = tl.load(DO_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)
        L  = tl.load(L_it,  boundary_check=(0,),   padding_option="zero").to(tl.float32)
        D_row = tl.sum(do * o, axis=1)

        # S: [Bq, Bk]
        S = tl.dot(q, tl.trans(k)) * scale
        
        # No mask needed for q > k
        # P = exp(S - L)
        P = tl.exp(S - L[:, None])

        # dP = dO @ V^T
        dP = tl.dot(do, tl.trans(v))

        # dS = P * (dP - D_row)
        dS = P * (dP - D_row[:, None])

        # dV += P^T @ dO
        dv_acc += tl.dot(tl.trans(P), do)

        # dK += dS^T @ Q * scale
        dk_acc += tl.dot(tl.trans(dS), q) * scale

        # Advance to next tile
        Q_it  = Q_it.advance((Q_TILE_SIZE, 0))
        O_it  = O_it.advance((Q_TILE_SIZE, 0))
        DO_it = DO_it.advance((Q_TILE_SIZE, 0))
        L_it  = L_it.advance((Q_TILE_SIZE,))
else:
    # Non-causal path: standard sweep
    for _ in range(0, N_QUERIES, Q_TILE_SIZE):
        q  = tl.load(Q_it,  boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bq, D]
        o  = tl.load(O_it,  boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bq, D]
        do = tl.load(DO_it, boundary_check=(0, 1), padding_option="zero").to(tl.float32)  # [Bq, D]
        L  = tl.load(L_it,  boundary_check=(0,),   padding_option="zero").to(tl.float32)  # [Bq]

        D_row = tl.sum(do * o, axis=1)   # [Bq]

        # S: [Bq, Bk]
        S = tl.dot(q, tl.trans(k)) * scale

        # P = exp(S - L)
        P = tl.exp(S - L[:, None])  # [Bq, Bk]

        # dP = dO @ V^T
        dP = tl.dot(do, tl.trans(v))  # [Bq, Bk]

        # dS = P * (dP - D_row)
        dS = P * (dP - D_row[:, None])  # [Bq, Bk]

        # dV += P^T @ dO
        dv_acc += tl.dot(tl.trans(P), do)

        # dK += dS^T @ Q * scale
        dk_acc += tl.dot(tl.trans(dS), q) * scale

        Q_it  = Q_it.advance((Q_TILE_SIZE, 0))
        O_it  = O_it.advance((Q_TILE_SIZE, 0))
        DO_it = DO_it.advance((Q_TILE_SIZE, 0))
        L_it  = L_it.advance((Q_TILE_SIZE,))

flash_bwd_dkdv_kernel 的思路和 forward/dq 方向相反：它是 “固定一个 K tile（pid_k），去扫 Q tiles”，在 causal 下，梯度贡献只来自满足 q >= k 的那部分，因此它的主要变化是：

变化 1：把 Q 的遍历起点从 0 改为对角线位置 start_q_offset = pid_k * Q_TILE_SIZE

start_q_offset = pid_k * Q_TILE_SIZE
Q_it  = Q_it.advance((start_q_offset, 0))
O_it  = O_it.advance((start_q_offset, 0))
DO_it = DO_it.advance((start_q_offset, 0))
L_it  = L_it.advance((start_q_offset,))

对固定的 K tile（索引区间为 [k0, k1]），当 Q tile 完全位于其之前（q1 <= k0）时，causal mask 使得这些位置的注意力权重恒为 0，因此它们对 dK/dV 的贡献也为 0.于是可以将 Q 的扫描起点直接移动到对角线处（q0 = k0），从而跳过整段无贡献的 Q tiles

变化 2：仍然分成“对角线 tile（带 mask）+ 后续 tiles（无 mask）”

我们先在对角线 Q tile 做一次 mask：

q_idx = start_q_offset + tl.arange(0, Q_TILE_SIZE)
S = tl.where(q_idx[:, None] >= k_idx[None, :], S, -1.0e6)

然后进入 Phase 2：后续 Q tiles 都满足 q > k，不再需要 mask：

# Phase 2: Non-diagonal Q tiles ... no mask
for _ in range(start_q_offset + Q_TILE_SIZE, N_QUERIES, Q_TILE_SIZE):
    ...
    # No mask needed for q > k

OK，我们来看下将 causal mask 优化后性能如何：

可以看到提升还是不错的，forward 和 backward 的耗时都大幅降低了，最终的性能提升对比表如下：

impl	batch_size	seq_len	num_heads	d_head	fwd_ms	bwd_ms	e2e_ms
baseline	1	4096	16	64	7.090	88.848	95.497
+autotune	1	4096	16	64	6.087	89.004	94.900
+autotune+triton-bwd	1	4096	16	64	6.081	31.706	37.329
+autotune+triton-bwd+causal_mask	1	4096	16	64	3.052	21.887	24.681

从上面的结果图表中我们可以看到，通过实现作业建议的三个优化点博主成功在当前硬件资源和参数下将 forward 耗时降低了 2.32x，backward 耗时降低了 4.06x，端到端耗时降低了 3.87x，效果还是相当不错的

OK，以上就是本次 FlashAttention-2 作业实现的全部内容了

结语

本篇文章我们完整实现并验证了 CS336 Assignment2 中 FlashAttention-2 的前向与反向传播流程，从最初的纯 PyTorch tiled 实现出发，逐步过渡到 Triton kernel 融合实现。

整个实现过程清晰地体现了 FlashAttention-2 的核心工程思想：通过重新组织计算顺序与内存访问模式，在不改变数学定义的前提下，避免显式物化 O(S^2) 级别的中间结果。online softmax 所维护的 (m, l, acc) 状态，使得注意力计算能够以线性显存复杂度完成，这一优势在长序列场景下尤为关键。而反向传播部分通过 recomputation 策略避免了 softmax 的复杂反传，成功将显存压力从存储激活转移为了额外的计算开销。

通过对前向与反向计算公式的重新组织，可以看到 FlashAttention-2 并未改变注意力机制本身的数学定义，而是通过更贴合 GPU 硬件特性的执行策略，将原本 “被内存限制” 的计算重新拉回到 “受计算吞吐限制” 的状态，这正是现代高性能深度学习系统设计中最具代表性的优化范式之一。

在接下来的作业中，我们将探索多 GPU 训练方法，并重点关注分布式数据并行（DDP）训练，敬请期待🤗

源码下载链接

• https://github.com/Melody-Zhou/stanford-cs336-spring2025-assignments

参考

• https://github.com/stanford-cs336/assignment2-systems