reference the flash attention GitHub

bert_padding.py

@@ -1,5 +1,10 @@
 # Adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py
 
+""""
+The implementation was further adapted from
+https://github.com/Dao-AILab/flash-attention/blob/43950dda456e095969d842fca7a73c5bfe3cecd0
+"""
+
 import torch
 import torch.nn.functional as F
 from einops import rearrange, repeat

block.py

@@ -1,5 +1,10 @@
 # Copyright (c) 2024, Tri Dao.
 
+""""
+The implementation was adopted from
+https://github.com/Dao-AILab/flash-attention/blob/43950dda456e095969d842fca7a73c5bfe3cecd0
+"""
+
 from functools import partial
 from typing import Optional
 

embedding.py

@@ -1,5 +1,10 @@
 # Copyright (c) 2022, Tri Dao.
 
+""""
+The implementation was adopted from
+https://github.com/Dao-AILab/flash-attention/blob/43950dda456e095969d842fca7a73c5bfe3cecd0/flash_attn/models/bert.py
+"""
+
 import torch
 import torch.nn as nn
 from torch import Tensor

mha.py

@@ -1,5 +1,14 @@
 # Copyright (c) 2023, Tri Dao.
 
+""""
+The implementation was adopted from
+https://github.com/Dao-AILab/flash-attention/blob/43950dda456e095969d842fca7a73c5bfe3cecd0
+and made modifications to
+    - support QK normalization
+    - make ALiBi run with MHA (needed to cast alibi slopes to fp32)
+    - make ALiBi run on CPU
+"""
+
 import math
 from functools import partial
 

mlp.py

@@ -1,5 +1,10 @@
 # Copyright (c) 2023, Tri Dao.
 
+""""
+The implementation was adopted from
+https://github.com/Dao-AILab/flash-attention/blob/43950dda456e095969d842fca7a73c5bfe3cecd0
+"""
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F