绕过 FlashAttention-2 限制：在 Turing 架构上使用 PyTorch 实现 FlashAttention

背景：FlashAttention 的架构限制与现实困境

FlashAttention【github】系列，由斯坦福大学 AI 实验室的 Dao 等人提出，通过巧妙的 I/O 感知算法（I/O-aware algorithm）和自定义的 CUDA 内核（CUDA kernel）显著提升了 Transformer 模型中注意力机制的训练和推理速度，尤其是在长序列处理方面。

然而，在使用过程中，我们可能会遇到一个常见的限制：FlashAttention-2 官方声明不支持较旧的 Turing 架构 （如 NVIDIA RTX 20 系列显卡）。尽管项目作者曾表示会很快支持，但在实际等待中，我们发现这一支持迟迟未能到来。

与此同时，许多最新的模型和代码库都是基于 FlashAttention-2 版本进行开发的，如果仅仅为了适配旧架构而回退到 FlashAttention-1.x，不仅需要大量的代码修改，还可能因此失去部分性能优化，这无疑是一个费时费力的选择。

那么，有没有一种方法，既能利用 FlashAttention 的核心思想，又能避免版本兼容性问题，让我们在 Turing 架构上也能快速、高效地运行最新的模型呢？

笔者想到的答案是：通过 PyTorch 自行实现 FlashAttention 的核心部分。

源码文件见：在Turing 架构上使用 PyTorch实现FlashAttention

实现原理：从 CUDA Kernel 到 PyTorch Tensor 操作

FlashAttention 的核心思想是分块计算（Tiled computation）。它将长序列的注意力计算拆分为多个小块，通过在显存（SRAM）上多次加载和计算，减少了对带宽瓶颈的 DRAM 的访问，从而显著提升了计算效率。

虽然我们无法直接复现其底层的 CUDA kernel，但我们可以利用 PyTorch 强大的张量操作能力，模拟这种分块计算的逻辑。具体来说，我们可以：

将 Q、K、V 矩阵进行分块（tiling）：将输入矩阵在序列维度上切分成多个小块。
迭代计算分块注意力：在循环中，依次取出 Q 的一个块和 K、V 的一个块进行计算。
累积结果：在每次迭代中，计算出局部分块的注意力得分，并将其与之前的累积结果进行合并。

这种方法虽然无法完全达到原生 CUDA kernel 的极限速度，但它成功地将 FlashAttention 的分块思想带入了 PyTorch，使得在没有原生支持的情况下，我们也能获得接近的性能和效果，尤其是在 PyTorch 2.x 引入的 torch.compile 等优化后，性能差距进一步缩小。

具体实现：核心代码详解与实现样例

我们将重点实现flash_attn.flash_attn_varlen_qkvpacked_func函数。

1. 核心实现文件（可直接使用）：flash_attn_torch.py

这个文件包含了我们用 PyTorch 张量操作实现的 FlashAttention 核心逻辑。以下是关键部分的伪代码：

python 复制代码

# 伪代码

def pytorch_flash_attention_varlen_qkvpacked(
    qkv,
    cu_seqlens,
    max_seqlen,
    dropout_p=0.0,
    softmax_scale=None,
    causal=False,
    window_size=(-1, -1),
    softcap=0.0,
    alibi_slopes=None,
    deterministic=False, # Note: PyTorch SDP deterministic behavior might depend on backend/version
    return_attn_probs=False,
):
    """
    PyTorch implementation mimicking flash_attn.flash_attn_varlen_qkvpacked_func.
    Handles variable-length sequences specified by cu_seqlens.
    Uses optimized vectorized padding/unpadding preprocessing.

    Args:
        qkv (torch.Tensor): Packed QKV tensor [total_tokens, 3, num_heads, head_dim].
        cu_seqlens (torch.Tensor): Cumulative sequence lengths [batch_size + 1].
        max_seqlen (int): Maximum sequence length in the batch.
        dropout_p (float): Dropout probability. Default is 0.0.
        softmax_scale (float, optional): Softmax scaling factor. Default is 1/sqrt(head_dim).
        causal (bool): Apply causal masking. Default is False.
        window_size (tuple): **Unsupported**. Must be (-1, -1).
        softcap (float): **Unsupported**. Must be 0.0.
        alibi_slopes (torch.Tensor, optional): **Unsupported**. Must be None.
        deterministic (bool): Attempt deterministic execution.
        return_attn_probs (bool): **Unsupported**. Must be False.

    Returns:
        torch.Tensor: Output context tensor [total_tokens, num_heads, head_dim].
    """

    # --- QKV Splitting ---
    q, k, v = qkv.unbind(dim=1)

    # Use vectorized assignment to place data into padded tensors
    padded_q[batch_indices, within_seq_indices] = q
    padded_k[batch_indices, within_seq_indices] = k
    padded_v[batch_indices, within_seq_indices] = v


    output = F.scaled_dot_product_attention(
                padded_q, padded_k, padded_v,
                attn_mask=None,
                dropout_p=dropout_p,
                is_causal=causal,
                scale=softmax_scale
            )

    # --- Optimized Unpadding ---
    # Transpose back to [bs, max_seqlen, nheads, headdim]
    output = output.transpose(1, 2)

    # Create the boolean mask efficiently using broadcasting
    mask = torch.arange(max_seqlen, device=q.device)[None, :] < seqlens[:, None] # Shape: [batch_size, max_seqlen]

    # Use the boolean mask to select only the valid tokens
    # Result shape: [total_tokens, nheads, headdim]
    unpadded_output = output[mask]


    return unpadded_output

2. 验证文件：`flash_attention_verify.py`

这个文件用于验证我们自实现的 PyTorch 版本与官方 FlashAttention 的效果和速度差异。

效果一致性验证 ：

我们生成随机的 Q、K、V 张量，分别使用官方的 FlashAttention和我们自实现的 PyTorch 版本进行计算，然后比较两者的输出张量是否在数值上接近（使用 torch.allclose）。
里面包含测试过程中实现的若干个版本