flash_attention.cpp

#include "flash_attention.h"

#include <algorithm>
#include <cmath>
#include <limits>
#include <vector>

constexpr float someRandVal = 42.0f;

/* Algorithm 0 Standard Attention Implementation
  Q, K, V - N x d
  S - N x N
  P - N x N
  O - N x d
  N - sequence length
  d - head dimension
*/
void attention_forward_naive(const float *Q, const float *K, const float *V,
                             float *S, float *P, float *O, int N, int d) {
  for (int i = 0; i < N; ++i) {
    for (int j = 0; j < N; ++j) {
      float Sij = 0;
      for (int k = 0; k < d; ++k)
        Sij += Q[i * d + k] * K[j * d + k];
      S[i * N + j] = Sij;
    }
  }

  for (int i = 0; i < N; ++i) {
    float Si_max = -std::numeric_limits<float>::max();
    for (int j = 0; j < N; ++j)
      Si_max = std::max(Si_max, S[i * N + j]);

    float Si_norm = 0;
    for (int j = 0; j < N; ++j) {
      float Pij = std::exp(S[i * N + j] - Si_max);
      P[i * N + j] = Pij;
      Si_norm += Pij;
    }

    for (int j = 0; j < N; ++j)
      P[i * N + j] /= Si_norm;
  }

  for (int i = 0; i < N; ++i) {
    for (int j = 0; j < d; ++j) {
      float Oij = 0;
      for (int k = 0; k < N; ++k)
        Oij += P[i * N + k] * V[k * d + j];
      O[i * d + j] = Oij;
    }
  }
}

// Algorithm 1 FlashAttention
void attention_forward_tiled(const float *Q, const float *K, const float *V,
                             float *O, int N, int d, int Br, int Bc) {
  for (int i = 0; i < N; ++i)
    for (int j = 0; j < d; ++j)
      O[i * d + j] = 0;

  std::vector<float> l(N, 0);
  std::vector<float> m(N, -std::numeric_limits<float>::max());

  const int Tr = (N + Br - 1) / Br;
  const int Tc = (N + Bc - 1) / Bc;

  for (int j = 0; j < Tc; ++j) {
    std::vector<float> Kj(Bc * d, someRandVal); // SRAM
    std::vector<float> Vj(Bc * d, someRandVal); // SRAM
    for (int row = 0; row < Bc && Bc * j + row < N; ++row) {
      for (int col = 0; col < d; ++col) {
        Kj[row * d + col] = K[(Bc * j + row) * d + col];
        Vj[row * d + col] = V[(Bc * j + row) * d + col];
      }
    }
    for (int i = 0; i < Tr; ++i) {
      std::vector<float> Qi(Br * d, someRandVal); // SRAM
      std::vector<float> Oi(Br * d, someRandVal); // SRAM
      std::vector<float> li(Br, someRandVal);
      std::vector<float> mi(Br, someRandVal);

      for (int row = 0; row < Br && Br * i + row < N; ++row) {
        for (int col = 0; col < d; ++col) {
          Qi[row * d + col] = Q[(Br * i + row) * d + col];
          Oi[row * d + col] = O[(Br * i + row) * d + col];
        }
        li[row] = l[Br * i + row];
        mi[row] = m[Br * i + row];
      }

      std::vector<float> Sij(Br * Bc, someRandVal); // SRAM
      for (int row = 0; row < Br && Br * i + row < N; ++row) {
        for (int col = 0; col < Bc && Bc * j + col < N; ++col) {
          float tmp = 0;
          for (int k = 0; k < d; ++k)
            tmp += Qi[row * d + k] * Kj[col * d + k];
          Sij[row * Bc + col] = tmp;
        }
      }
      for (int row = 0; row < Br && Br * i + row < N; ++row) {
        float mij = -std::numeric_limits<float>::max();
        for (int col = 0; col < Bc && Bc * j + col < N; ++col)
          mij = std::max(mij, Sij[row * Bc + col]);

        for (int col = 0; col < Bc && Bc * j + col < N; ++col)
          Sij[row * Bc + col] = std::exp(Sij[row * Bc + col] - mij);

        float lij = 0;
        for (int col = 0; col < Bc && Bc * j + col < N; ++col)
          lij += Sij[row * Bc + col];

        const float mi_new = std::max(mi[row], mij);
        const float li_new =
            std::exp(mi[row] - mi_new) * li[row] + std::exp(mij - mi_new) * lij;

        const float inv_li_new = 1 / li_new;
        const float o_factor = li[row] * std::exp(mi[row] - mi_new);
        const float pv_factor = std::exp(mij - mi_new);
        for (int col = 0; col < d; ++col) {
          float pv = 0;
          for (int k = 0; k < Bc && Bc * j + k < N; ++k)
            pv += Sij[row * Bc + k] * Vj[k * d + col];

          O[(Br * i + row) * d + col] =
              inv_li_new * (o_factor * Oi[row * d + col] + pv_factor * pv);
        }

        l[Br * i + row] = li_new;
        m[Br * i + row] = mi_new;
      }
    }
  }
}

// Algorithm 2 FlashAttention Forward Pass
// TODO: Add `dropout`
template <typename Mask>
void attention_forward_tiled(const float *Q, const float *K, const float *V,
                             float *O, float *l, float *m, int N, int d, int Br,
                             int Bc, float scale, Mask mask) {
  for (int i = 0; i < N; ++i)
    for (int j = 0; j < d; ++j)
      O[i * d + j] = 0;

  for (int i = 0; i < N; ++i)
    l[i] = 0;
  for (int i = 0; i < N; ++i)
    m[i] = -std::numeric_limits<float>::max();

  const int Tr = (N + Br - 1) / Br;
  const int Tc = (N + Bc - 1) / Bc;

  for (int j = 0; j < Tc; ++j) {
    constexpr float someRandVal = 42.0f;
    std::vector<float> Kj(Bc * d, someRandVal); // SRAM
    std::vector<float> Vj(Bc * d, someRandVal); // SRAM
    for (int row = 0; row < Bc && Bc * j + row < N; ++row) {
      for (int col = 0; col < d; ++col) {
        Kj[row * d + col] = K[(Bc * j + row) * d + col];
        Vj[row * d + col] = V[(Bc * j + row) * d + col];
      }
    }
    for (int i = 0; i < Tr; ++i) {
      std::vector<float> Qi(Br * d, someRandVal); // SRAM
      std::vector<float> Oi(Br * d, someRandVal); // SRAM
      std::vector<float> li(Br, someRandVal);
      std::vector<float> mi(Br, someRandVal);

      for (int row = 0; row < Br && Br * i + row < N; ++row) {
        for (int col = 0; col < d; ++col) {
          Qi[row * d + col] = Q[(Br * i + row) * d + col];
          Oi[row * d + col] = O[(Br * i + row) * d + col];
        }
        li[row] = l[Br * i + row];
        mi[row] = m[Br * i + row];
      }

      std::vector<float> Sij(Br * Bc, someRandVal); // SRAM
      for (int row = 0; row < Br && Br * i + row < N; ++row) {
        for (int col = 0; col < Bc && Bc * j + col < N; ++col) {
          if (mask(Br * i + row, Bc * j + col)) {
            Sij[row * Bc + col] = -std::numeric_limits<float>::max();
            continue;
          }

          float tmp = 0;
          for (int k = 0; k < d; ++k)
            tmp += Qi[row * d + k] * Kj[col * d + k];
          Sij[row * Bc + col] = scale * tmp;
        }
      }
      for (int row = 0; row < Br && Br * i + row < N; ++row) {
        float mij = -std::numeric_limits<float>::max();
        for (int col = 0; col < Bc && Bc * j + col < N; ++col)
          mij = std::max(mij, Sij[row * Bc + col]);

        for (int col = 0; col < Bc && Bc * j + col < N; ++col)
          Sij[row * Bc + col] = std::exp(Sij[row * Bc + col] - mij);

        float lij = 0;
        for (int col = 0; col < Bc && Bc * j + col < N; ++col)
          lij += Sij[row * Bc + col];

        const float mi_new = std::max(mi[row], mij);
        const float li_new =
            std::exp(mi[row] - mi_new) * li[row] + std::exp(mij - mi_new) * lij;

        const float inv_li_new = 1 / li_new;
        const float o_factor = li[row] * std::exp(mi[row] - mi_new);
        const float pv_factor = std::exp(mij - mi_new);
        for (int col = 0; col < d; ++col) {
          float pv = 0;
          for (int k = 0; k < Bc && Bc * j + k < N; ++k)
            pv += Sij[row * Bc + k] * Vj[k * d + col];

          O[(Br * i + row) * d + col] =
              inv_li_new * (o_factor * Oi[row * d + col] + pv_factor * pv);
        }

        l[Br * i + row] = li_new;
        m[Br * i + row] = mi_new;
      }
    }
  }
}

#define INSTANTIATE_ATTENTION_FORWARD_TILED(MASK)                              \
  template void attention_forward_tiled<MASK>(                                 \
      const float *, const float *, const float *, float *, float *, float *,  \
      int, int, int, int, float, MASK);

INSTANTIATE_ATTENTION_FORWARD_TILED(masks::None)
INSTANTIATE_ATTENTION_FORWARD_TILED(masks::Causal)
INSTANTIATE_ATTENTION_FORWARD_TILED(masks::KeyPadding)
INSTANTIATE_ATTENTION_FORWARD_TILED(masks::Window)

/* Algorithm 3 Standard Attention Backward Pass
  Q, K, V, dO - N x d
  P - N x N
  dQ, dK, dV - N x d
*/
void attention_backward_naive(const float *Q, const float *K, const float *V,
                              const float *dO, const float *P, float *dQ,
                              float *dK, float *dV, int N, int d) {
  for (int i = 0; i < N; ++i) {
    for (int j = 0; j < d; ++j) {
      float dVij = 0;
      for (int k = 0; k < N; ++k)
        dVij += P[k * N + i] * dO[k * d + j];
      dV[i * d + j] = dVij;
    }
  }

  std::vector<float> dP(N * N, someRandVal);
  for (int i = 0; i < N; ++i) {
    for (int j = 0; j < N; ++j) {
      float dPij = 0;
      for (int k = 0; k < d; ++k)
        dPij += dO[i * d + k] * V[j * d + k];
      dP[i * N + j] = dPij;
    }
  }

  std::vector<float> dS(N * N, someRandVal);
  for (int i = 0; i < N; ++i) {
    float PdP = 0;
    for (int j = 0; j < N; ++j)
      PdP += P[i * N + j] * dP[i * N + j];
    for (int j = 0; j < N; ++j)
      dS[i * N + j] = P[i * N + j] * (dP[i * N + j] - PdP);
  }

  for (int i = 0; i < N; ++i) {
    for (int j = 0; j < d; ++j) {
      float dQij = 0;
      for (int k = 0; k < N; ++k)
        dQij += dS[i * N + k] * K[k * d + j];
      dQ[i * d + j] = dQij;
    }
  }

  for (int i = 0; i < N; ++i) {
    for (int j = 0; j < d; ++j) {
      float dKij = 0;
      for (int k = 0; k < N; ++k)
        dKij += dS[k * N + i] * Q[k * d + j];
      dK[i * d + j] = dKij;
    }
  }
}

template <typename Mask>
void attention_backward_tiled(const float *Q, const float *K, const float *V,
                              const float *O, const float *dO, const float *l,
                              const float *m, float *dQ, float *dK, float *dV,
                              int N, int d, int Br, int Bc, float scale,
                              Mask mask) {
  const int Tr = (N + Br - 1) / Br;
  const int Tc = (N + Bc - 1) / Bc;

  for (int i = 0; i < N; ++i) {
    for (int j = 0; j < d; ++j) {
      dQ[i * d + j] = 0;
      dK[i * d + j] = 0;
      dV[i * d + j] = 0;
    }
  }

  for (int j = 0; j < Tc; ++j) {
    std::vector<float> Kj(Bc * d, someRandVal);
    std::vector<float> Vj(Bc * d, someRandVal);
    for (int row = 0, grow = Bc * j; row < Bc && grow < N; ++row, ++grow) {
      for (int col = 0; col < d; ++col) {
        Kj[row * d + col] = K[grow * d + col];
        Vj[row * d + col] = V[grow * d + col];
      }
    }

    std::vector<float> dKj(Bc * d, someRandVal);
    std::vector<float> dVj(Bc * d, someRandVal);
    for (int row = 0; row < Bc; ++row) {
      for (int col = 0; col < d; ++col) {
        dKj[row * d + col] = 0;
        dVj[row * d + col] = 0;
      }
    }

    for (int i = 0; i < Tr; ++i) {
      std::vector<float> Qi(Br * d, someRandVal);
      std::vector<float> Oi(Br * d, someRandVal);
      std::vector<float> dQi(Br * d, someRandVal);
      std::vector<float> dOi(Br * d, someRandVal);
      std::vector<float> li(Br, someRandVal);
      std::vector<float> mi(Br, someRandVal);
      for (int row = 0, grow = Br * i; row < Br && grow < N; ++row, ++grow) {
        for (int col = 0; col < d; ++col) {
          Qi[row * d + col] = Q[grow * d + col];
          Oi[row * d + col] = O[grow * d + col];
          dQi[row * d + col] = dQ[grow * d + col];
          dOi[row * d + col] = dO[grow * d + col];
        }
        li[row] = l[grow];
        mi[row] = m[grow];
      }

      std::vector<float> Sij(Br * Bc, someRandVal);
      for (int row = 0, grow = Br * i; row < Br && grow < N; ++row, ++grow) {
        const float inv_li = 1 / li[row];
        const float mi_val = mi[row];

        for (int col = 0, gcol = Bc * j; col < Bc && gcol < N; ++col, ++gcol) {
          float Sij_val;
          if (mask(grow, gcol))
            Sij_val = -std::numeric_limits<float>::max();
          else {
            Sij_val = 0;
            for (int k = 0; k < d; ++k)
              Sij_val += Qi[row * d + k] * Kj[col * d + k];
            Sij_val *= scale;
          }

          const float Pij_val = inv_li * std::exp(Sij_val - mi_val);
          Sij[row * Bc + col] = Pij_val;
        }
      }

      for (int row = 0, grow = Bc * j; row < Bc && grow < N; ++row, ++grow) {
        for (int col = 0; col < d; ++col) {
          float tmp = 0;
          for (int k = 0; k < Br && Br * i + k < N; ++k)
            tmp += Sij[k * Bc + row] * dOi[k * d + col];
          dVj[row * d + col] += tmp;
        }
      }

      std::vector<float> dPij(Br * Bc, someRandVal);
      for (int row = 0, grow = Br * i; row < Br && grow < N; ++row, ++grow) {
        for (int col = 0, gcol = Bc * j; col < Bc && gcol < N; ++col, ++gcol) {
          float tmp = 0;
          for (int k = 0; k < d; ++k)
            tmp += dOi[row * d + k] * Vj[col * d + k];
          dPij[row * Bc + col] = tmp;
        }
      }

      std::vector<float> Di(Br, someRandVal);
      for (int row = 0, grow = Br * i; row < Br && grow < N; ++row, ++grow) {
        float tmp = 0;
        for (int col = 0; col < d; ++col)
          tmp += dOi[row * d + col] * Oi[row * d + col];
        Di[row] = tmp;
      }

      std::vector<float> dSij(Br * Bc, someRandVal);
      for (int row = 0, grow = Br * i; row < Br && grow < N; ++row, ++grow) {
        const float di_val = Di[row];
        for (int col = 0, gcol = Bc * j; col < Bc && gcol < N; ++col, ++gcol) {
          dSij[row * Bc + col] =
              Sij[row * Bc + col] * (dPij[row * Bc + col] - di_val);
        }
      }

      for (int row = 0, grow = Br * i; row < Br && grow < N; ++row, ++grow) {
        for (int col = 0; col < d; ++col) {
          float tmp = 0;
          for (int k = 0; k < Bc && Bc * j + k < N; ++k)
            tmp += dSij[row * Bc + k] * Kj[k * d + col];
          dQ[grow * d + col] = dQi[row * d + col] + tmp * scale;
        }
      }

      for (int row = 0, grow = Bc * j; row < Bc && grow < N; ++row, ++grow) {
        for (int col = 0; col < d; ++col) {
          float tmp = 0;
          for (int k = 0; k < Br && Br * i + k < N; ++k)
            tmp += dSij[k * Bc + row] * Qi[k * d + col];
          dKj[row * d + col] += tmp * scale;
        }
      }
    }

    for (int row = 0, grow = Bc * j; row < Bc && grow < N; ++row, ++grow) {
      for (int col = 0; col < d; ++col) {
        dK[grow * d + col] = dKj[row * d + col];
        dV[grow * d + col] = dVj[row * d + col];
      }
    }
  }
}

#define INSTANTIATE_ATTENTION_BACKWARD_TILED(MASK)                             \
  template void attention_backward_tiled<MASK>(                                \
      const float *, const float *, const float *, const float *,              \
      const float *, const float *, const float *, float *, float *, float *,  \
      int, int, int, int, float, MASK);

INSTANTIATE_ATTENTION_BACKWARD_TILED(masks::None)
INSTANTIATE_ATTENTION_BACKWARD_TILED(masks::Causal)
INSTANTIATE_ATTENTION_BACKWARD_TILED(masks::KeyPadding)
INSTANTIATE_ATTENTION_BACKWARD_TILED(masks::Window)