FlashAttention2 with Custom Masks

FlashAttention2 with Custom Masks 简介

FlashAttention2 with Custom Masks 是一个针对 FlashAttention2 的非官方实现，旨在提供对任意自定义掩码的支持。在标准 FlashAttention 实现中，通常只支持特定的掩码，例如用于语言建模的因果掩码，而无法处理更加灵活的自定义掩码。该实现通过修改前向和反向传播的过程，允许用户为每个头和每个批次定义不同的掩码。

主要特性

• 支持任意自定义掩码，用户可以根据需求定义不同的掩码形式。
• 兼容标准 FlashAttention 以外的应用场景，实现了更高的灵活性。

安装与使用

用户需要创建一个 Python 环境（推荐版本 >= 3.8），并通过 pip 安装该库。可以参考以下命令进行安装：

pip install flashattention2-custom-mask

使用自定义掩码的 FlashAttention2 内核需要安装 triton 和 torch 等相关库。

示例代码

用户可以在标准注意力管道中插入该模块，使用如下示例代码进行操作：

from fa2_custom_mask import flash_attention_custom_mask

B, H, L, D = 4, 16, 4096, 64
sm_scale = 1 / (D ** 0.5)

fp32_q = torch.randn(B, H, L, D).float().cuda()
fp32_k = torch.randn(B, H, L, D).float().cuda()
fp32_v = torch.randn(B, H, L, D).float().cuda()
mask = torch.randint(0, 2, (B, 1, L, L)).int().cuda()
mask = torch.broadcast_to(mask, (B, H, L, L))

out = flash_attention_custom_mask(fp32_q, fp32_k, fp32_v, mask=mask, sm_scale=sm_scale)
out.backward(loss)

应用场景

FlashAttention2 with Custom Masks 的应用场景非常广泛，包括但不限于：

1. 自然语言处理：在语言模型中使用自定义掩码，支持更复杂的上下文注意力机制。
2. 图像处理：在视觉任务中，可根据特定需求调整注意力机制，使其对图像区域更具选择性。
3. 序列处理任务：例如在时间序列分析中，用户可以定义不同的掩码策略以优化处理效果。

性能评估

可以通过运行 test_benchmark.py 测试实施的正确性，并评估与官方 FlashAttention 实现及其他内存高效注意力实现（如 xformers）之间的性能比较。

注意事项

1. 该实现仅适用于 Ampere 及更高版本的设备。
2. 确保使用 triton>=3.0.0 以避免一些兼容性问题。

总之，FlashAttention2 with Custom Masks 提供了强大的灵活性，支持多种自定义掩码的需求，适用于各种深度学习任务，尤其是在需要复杂注意力机制的应用场景中。