flash_attention_test.cpp

#include "flash_attention.h"

#include <algorithm>
#include <cmath>
#include <limits>
#include <ostream>
#include <random>
#include <string>
#include <utility>
#include <vector>

#include <gtest/gtest.h>

namespace {

// Fill a buffer with uniform random values in [lo, hi] from a fixed seed so
// failures are reproducible.
void fill_random(std::vector<float> &buf, float lo, float hi, unsigned seed) {
  std::mt19937 gen(seed);
  std::uniform_real_distribution<float> dist(lo, hi);
  for (float &x : buf)
    x = dist(gen);
}

// Owns all buffers for one call of the naive forward pass.
struct Buffers {
  int N, d;
  std::vector<float> Q, K, V, S, P, O;

  Buffers(int N_, int d_)
      : N(N_), d(d_), Q(N * d), K(N * d), V(N * d), S(N * N), P(N * N),
        O(N * d) {}

  void run() {
    attention_forward_naive(Q.data(), K.data(), V.data(), S.data(), P.data(),
                            O.data(), N, d);
  }
};

} // namespace

// N = 2, d = 2, worked out by hand.
//   Q = K = I, V = [[1, 2], [3, 4]]
//   S = Q K^T = I
//   P = [[e/(e+1), 1/(e+1)], [1/(e+1), e/(e+1)]]
//   O = P V
TEST(AttentionForwardNaive, HandComputed2x2) {
  Buffers b(2, 2);
  b.Q = {1, 0, 0, 1};
  b.K = {1, 0, 0, 1};
  b.V = {1, 2, 3, 4};
  b.run();

  const float tol = 1e-5f;

  EXPECT_NEAR(b.S[0], 1.0f, tol);
  EXPECT_NEAR(b.S[1], 0.0f, tol);
  EXPECT_NEAR(b.S[2], 0.0f, tol);
  EXPECT_NEAR(b.S[3], 1.0f, tol);

  const float a = 0.731058578f; // e / (e + 1)
  const float c = 0.268941421f; // 1 / (e + 1)
  EXPECT_NEAR(b.P[0], a, tol);
  EXPECT_NEAR(b.P[1], c, tol);
  EXPECT_NEAR(b.P[2], c, tol);
  EXPECT_NEAR(b.P[3], a, tol);

  EXPECT_NEAR(b.O[0], 1.537882842f, tol);
  EXPECT_NEAR(b.O[1], 2.537882842f, tol);
  EXPECT_NEAR(b.O[2], 2.462117157f, tol);
  EXPECT_NEAR(b.O[3], 3.462117157f, tol);
}

// Softmax rows must sum to 1 and contain only values in [0, 1].
TEST(AttentionForwardNaive, RowsOfPSumToOne) {
  Buffers b(64, 16);
  fill_random(b.Q, -1.0f, 1.0f, 1);
  fill_random(b.K, -1.0f, 1.0f, 2);
  fill_random(b.V, -1.0f, 1.0f, 3);
  b.run();

  for (int i = 0; i < b.N; ++i) {
    float sum = 0;
    for (int j = 0; j < b.N; ++j) {
      float p = b.P[i * b.N + j];
      EXPECT_GE(p, 0.0f);
      EXPECT_LE(p, 1.0f);
      sum += p;
    }
    EXPECT_NEAR(sum, 1.0f, 1e-5f) << "row " << i;
  }
}

// If every row of K is identical, every score in a row of S is equal, so P is
// uniform and each row of O is the column mean of V.
TEST(AttentionForwardNaive, IdenticalKeysGiveMeanOfV) {
  Buffers b(32, 8);
  fill_random(b.Q, -1.0f, 1.0f, 4);
  fill_random(b.V, -1.0f, 1.0f, 5);

  std::vector<float> key(b.d);
  fill_random(key, -1.0f, 1.0f, 6);
  for (int j = 0; j < b.N; ++j)
    for (int k = 0; k < b.d; ++k)
      b.K[j * b.d + k] = key[k];

  b.run();

  std::vector<float> mean(b.d, 0.0f);
  for (int j = 0; j < b.N; ++j)
    for (int k = 0; k < b.d; ++k)
      mean[k] += b.V[j * b.d + k] / b.N;

  for (int i = 0; i < b.N; ++i)
    for (int k = 0; k < b.d; ++k)
      EXPECT_NEAR(b.O[i * b.d + k], mean[k], 1e-5f) << "i=" << i << " k=" << k;
}

// With Q = s * I and K = I for a large s, row i of P is nearly one-hot at
// position i, so O should reproduce V.
TEST(AttentionForwardNaive, SharpAttentionSelectsRowOfV) {
  const int N = 8;
  Buffers b(N, N);
  const float s = 50.0f;
  for (int i = 0; i < N; ++i) {
    b.Q[i * N + i] = s;
    b.K[i * N + i] = 1.0f;
  }
  fill_random(b.V, -1.0f, 1.0f, 7);
  b.run();

  // exp(-50) * (N - 1) leakage is far below float resolution.
  for (int i = 0; i < N * N; ++i)
    EXPECT_NEAR(b.O[i], b.V[i], 1e-5f) << "index " << i;
}

// Scores far outside the range of exp() must not overflow to inf/NaN thanks
// to the max subtraction.
TEST(AttentionForwardNaive, LargeScoresStayFinite) {
  Buffers b(16, 4);
  fill_random(b.Q, -1000.0f, 1000.0f, 8);
  fill_random(b.K, -1000.0f, 1000.0f, 9);
  fill_random(b.V, -1.0f, 1.0f, 10);
  b.run();

  for (float p : b.P)
    EXPECT_TRUE(std::isfinite(p));
  for (float o : b.O)
    EXPECT_TRUE(std::isfinite(o));

  for (int i = 0; i < b.N; ++i) {
    float sum = 0;
    for (int j = 0; j < b.N; ++j)
      sum += b.P[i * b.N + j];
    EXPECT_NEAR(sum, 1.0f, 1e-5f) << "row " << i;
  }
}

// ---------------------------------------------------------------------------
// attention_forward_tiled
// ---------------------------------------------------------------------------

namespace {

// Runs the naive reference and returns only O.
std::vector<float> run_naive(const std::vector<float> &Q,
                             const std::vector<float> &K,
                             const std::vector<float> &V, int N, int d) {
  std::vector<float> S(N * N), P(N * N), O(N * d);
  attention_forward_naive(Q.data(), K.data(), V.data(), S.data(), P.data(),
                          O.data(), N, d);
  return O;
}

// Runs the tiled version. O is pre-filled with NaN so the test fails loudly if
// the implementation relies on the caller to zero it.
std::vector<float> run_tiled(const std::vector<float> &Q,
                             const std::vector<float> &K,
                             const std::vector<float> &V, int N, int d, int Br,
                             int Bc) {
  std::vector<float> O(N * d, std::numeric_limits<float>::quiet_NaN());
  attention_forward_tiled(Q.data(), K.data(), V.data(), O.data(), N, d, Br, Bc);
  return O;
}

// Elementwise comparison with a combined absolute/relative tolerance.
void expect_close(const std::vector<float> &got, const std::vector<float> &want,
                  float atol, float rtol) {
  ASSERT_EQ(got.size(), want.size());
  for (size_t i = 0; i < got.size(); ++i) {
    float tol = atol + rtol * std::fabs(want[i]);
    EXPECT_NEAR(got[i], want[i], tol) << "index " << i;
  }
}

struct TileCase {
  int N, d, Br, Bc;
};

std::ostream &operator<<(std::ostream &os, const TileCase &c) {
  return os << "N=" << c.N << " d=" << c.d << " Br=" << c.Br << " Bc=" << c.Bc;
}

// Block shapes shared by every parametrized tiled suite.
std::vector<TileCase> block_shapes() {
  return {
      // Blocks divide N exactly.
      {64, 16, 16, 16},
      {64, 16, 8, 32},
      {64, 16, 32, 8},
      // Single block covers everything: tiling degenerates to the reference.
      {32, 8, 32, 32},
      // Blocks larger than N: the ragged-edge path on the very first block.
      {20, 8, 64, 64},
      // One row / one column per block: maximum number of merges.
      {16, 4, 1, 1},
      {16, 4, 1, 16},
      {16, 4, 16, 1},
      // N not divisible by Br, Bc, or both.
      {100, 16, 32, 32},
      {100, 16, 7, 32},
      {100, 16, 32, 7},
      {37, 8, 5, 6},
      // Larger head dimension.
      {48, 64, 16, 16},
  };
}

} // namespace

class AttentionForwardTiled : public ::testing::TestWithParam<TileCase> {};

// The tiled forward must reproduce the naive reference for every block shape.
TEST_P(AttentionForwardTiled, MatchesNaive) {
  const TileCase c = GetParam();
  std::vector<float> Q(c.N * c.d), K(c.N * c.d), V(c.N * c.d);
  fill_random(Q, -1.0f, 1.0f, 11);
  fill_random(K, -1.0f, 1.0f, 12);
  fill_random(V, -1.0f, 1.0f, 13);

  auto want = run_naive(Q, K, V, c.N, c.d);
  auto got = run_tiled(Q, K, V, c.N, c.d, c.Br, c.Bc);
  expect_close(got, want, 1e-5f, 1e-4f);
}

INSTANTIATE_TEST_SUITE_P(BlockShapes, AttentionForwardTiled,
                         ::testing::ValuesIn(block_shapes()));

// Same hand-computed case as the naive test, but with 1x1 blocks so every
// output element is assembled from several merges.
TEST(AttentionForwardTiledFixed, HandComputed2x2) {
  std::vector<float> Q = {1, 0, 0, 1};
  std::vector<float> K = {1, 0, 0, 1};
  std::vector<float> V = {1, 2, 3, 4};
  auto O = run_tiled(Q, K, V, 2, 2, 1, 1);

  const float tol = 1e-5f;
  EXPECT_NEAR(O[0], 1.537882842f, tol);
  EXPECT_NEAR(O[1], 2.537882842f, tol);
  EXPECT_NEAR(O[2], 2.462117157f, tol);
  EXPECT_NEAR(O[3], 3.462117157f, tol);
}

// Identical keys: every row of O is the column mean of V, regardless of how
// the columns are split into blocks.
TEST(AttentionForwardTiledFixed, IdenticalKeysGiveMeanOfV) {
  const int N = 40, d = 8;
  std::vector<float> Q(N * d), K(N * d), V(N * d);
  fill_random(Q, -1.0f, 1.0f, 14);
  fill_random(V, -1.0f, 1.0f, 15);

  std::vector<float> key(d);
  fill_random(key, -1.0f, 1.0f, 16);
  for (int j = 0; j < N; ++j)
    for (int k = 0; k < d; ++k)
      K[j * d + k] = key[k];

  std::vector<float> mean(d, 0.0f);
  for (int j = 0; j < N; ++j)
    for (int k = 0; k < d; ++k)
      mean[k] += V[j * d + k] / N;

  auto O = run_tiled(Q, K, V, N, d, 8, 12);
  for (int i = 0; i < N; ++i)
    for (int k = 0; k < d; ++k)
      EXPECT_NEAR(O[i * d + k], mean[k], 1e-5f) << "i=" << i << " k=" << k;
}

// Sharp attention with Q = 50 * I, K = I. The dominant key for row i lands in
// a different K block for different i, so the running max must move as blocks
// are visited, exercising the rescaling of the partial O.
TEST(AttentionForwardTiledFixed, SharpAttentionSelectsRowOfV) {
  const int N = 12;
  std::vector<float> Q(N * N, 0.0f), K(N * N, 0.0f), V(N * N);
  for (int i = 0; i < N; ++i) {
    Q[i * N + i] = 50.0f;
    K[i * N + i] = 1.0f;
  }
  fill_random(V, -1.0f, 1.0f, 17);

  auto O = run_tiled(Q, K, V, N, N, 4, 5);
  for (int i = 0; i < N * N; ++i)
    EXPECT_NEAR(O[i], V[i], 1e-5f) << "index " << i;
}

// Scores in the thousands spread across blocks: the running max changes by
// hundreds between blocks, so exp(m_old - m_new) underflows to exactly zero
// for early blocks. The result must stay finite and match the reference.
TEST(AttentionForwardTiledFixed, LargeScoresStayFiniteAndMatchNaive) {
  const int N = 48, d = 4;
  std::vector<float> Q(N * d), K(N * d), V(N * d);
  fill_random(Q, -1000.0f, 1000.0f, 18);
  fill_random(K, -1000.0f, 1000.0f, 19);
  fill_random(V, -1.0f, 1.0f, 20);

  auto want = run_naive(Q, K, V, N, d);
  auto got = run_tiled(Q, K, V, N, d, 8, 8);

  for (float o : got)
    EXPECT_TRUE(std::isfinite(o));
  expect_close(got, want, 1e-5f, 1e-4f);
}

// Different block shapes must agree with each other, not just with the
// reference. This catches errors that scale with the number of merges.
TEST(AttentionForwardTiledFixed, BlockShapeInvariance) {
  const int N = 50, d = 16;
  std::vector<float> Q(N * d), K(N * d), V(N * d);
  fill_random(Q, -2.0f, 2.0f, 21);
  fill_random(K, -2.0f, 2.0f, 22);
  fill_random(V, -1.0f, 1.0f, 23);

  auto base = run_tiled(Q, K, V, N, d, N, N);
  expect_close(run_tiled(Q, K, V, N, d, 1, 1), base, 1e-5f, 1e-4f);
  expect_close(run_tiled(Q, K, V, N, d, 16, 16), base, 1e-5f, 1e-4f);
  expect_close(run_tiled(Q, K, V, N, d, 7, 13), base, 1e-5f, 1e-4f);
}

// ---------------------------------------------------------------------------
// attention_forward_tiled<Mask>: Algorithm 2 with scale, mask, and l, m outputs
// ---------------------------------------------------------------------------

struct ForwardResult {
  std::vector<float> O, l, m;
};

// Runs the Algorithm 2 template. Outputs are pre-filled with NaN so a value
// the implementation forgets to write fails loudly.
template <class Mask>
ForwardResult run_tiled_masked(const std::vector<float> &Q,
                               const std::vector<float> &K,
                               const std::vector<float> &V, int N, int d,
                               int Br, int Bc, float scale, Mask mask) {
  const float nan = std::numeric_limits<float>::quiet_NaN();
  ForwardResult r{std::vector<float>(N * d, nan), std::vector<float>(N, nan),
                  std::vector<float>(N, nan)};
  attention_forward_tiled(Q.data(), K.data(), V.data(), r.O.data(), r.l.data(),
                          r.m.data(), N, d, Br, Bc, scale, mask);
  return r;
}

// Double-precision reference written straight from the definitions, applying
// the mask at global (i, j):
//   S = scale * Q K^T,  m_i = max over kept j of S_ij,
//   l_i = sum over kept j of exp(S_ij - m_i),
//   O_i = sum over kept j of exp(S_ij - m_i) V_j / l_i.
template <class Mask>
ForwardResult reference_masked(const std::vector<float> &Q,
                               const std::vector<float> &K,
                               const std::vector<float> &V, int N, int d,
                               float scale, Mask mask) {
  ForwardResult r{std::vector<float>(N * d), std::vector<float>(N),
                  std::vector<float>(N)};
  std::vector<double> s(N), o(d);
  std::vector<bool> kept(N);
  for (int i = 0; i < N; ++i) {
    double mi = -std::numeric_limits<double>::infinity();
    for (int j = 0; j < N; ++j) {
      kept[j] = !mask(i, j);
      if (!kept[j])
        continue;
      double dot = 0;
      for (int k = 0; k < d; ++k)
        dot += static_cast<double>(Q[i * d + k]) * K[j * d + k];
      s[j] = scale * dot;
      mi = std::max(mi, s[j]);
    }

    double li = 0;
    std::fill(o.begin(), o.end(), 0.0);
    for (int j = 0; j < N; ++j) {
      if (!kept[j])
        continue;
      const double p = std::exp(s[j] - mi);
      li += p;
      for (int c = 0; c < d; ++c)
        o[c] += p * V[j * d + c];
    }

    for (int c = 0; c < d; ++c)
      r.O[i * d + c] = static_cast<float>(o[c] / li);
    r.l[i] = static_cast<float>(li);
    r.m[i] = static_cast<float>(mi);
  }
  return r;
}

void expect_result_close(const ForwardResult &got, const ForwardResult &want,
                         float atol, float rtol) {
  {
    SCOPED_TRACE("O");
    expect_close(got.O, want.O, atol, rtol);
  }
  {
    SCOPED_TRACE("l");
    expect_close(got.l, want.l, atol, rtol);
  }
  {
    SCOPED_TRACE("m");
    expect_close(got.m, want.m, atol, rtol);
  }
}

class AttentionForwardTiledMasked : public ::testing::TestWithParam<TileCase> {
};

// Every mask, with the standard 1/sqrt(d) scale, against the double-precision
// reference, for every block shape. Ragged shapes matter here: a mask applied
// with tile-local instead of global indices passes only when Br = Bc = 1.
TEST_P(AttentionForwardTiledMasked, MatchesReference) {
  const TileCase c = GetParam();
  std::vector<float> Q(c.N * c.d), K(c.N * c.d), V(c.N * c.d);
  fill_random(Q, -2.0f, 2.0f, 24);
  fill_random(K, -2.0f, 2.0f, 25);
  fill_random(V, -1.0f, 1.0f, 26);
  const float scale = 1.0f / std::sqrt(static_cast<float>(c.d));

  auto check = [&](const char *label, auto mask) {
    SCOPED_TRACE(label);
    auto want = reference_masked(Q, K, V, c.N, c.d, scale, mask);
    auto got = run_tiled_masked(Q, K, V, c.N, c.d, c.Br, c.Bc, scale, mask);
    expect_result_close(got, want, 1e-5f, 1e-4f);
  };
  check("no mask", masks::None{});
  check("causal", masks::Causal{});
  check("key padding", masks::KeyPadding{(2 * c.N + 2) / 3});
  check("window", masks::Window{std::max(1, c.N / 4)});
}

// With no mask and scale 1, Algorithm 2 reduces to Algorithm 1, so the two
// implementations should agree to rounding.
TEST_P(AttentionForwardTiledMasked, NoMaskUnitScaleReducesToAlgorithm1) {
  const TileCase c = GetParam();
  std::vector<float> Q(c.N * c.d), K(c.N * c.d), V(c.N * c.d);
  fill_random(Q, -1.0f, 1.0f, 11);
  fill_random(K, -1.0f, 1.0f, 12);
  fill_random(V, -1.0f, 1.0f, 13);

  auto want = run_tiled(Q, K, V, c.N, c.d, c.Br, c.Bc);
  auto got =
      run_tiled_masked(Q, K, V, c.N, c.d, c.Br, c.Bc, 1.0f, masks::None{});
  expect_close(got.O, want, 1e-7f, 1e-6f);
}

INSTANTIATE_TEST_SUITE_P(BlockShapes, AttentionForwardTiledMasked,
                         ::testing::ValuesIn(block_shapes()));

// Hand-computed causal case with Q = K = I, V = [[1, 2], [3, 4]]:
//   row 0 sees only key 0:  O_0 = V_0, l_0 = 1, m_0 = S_00 = 1
//   row 1 sees S_1 = [0, 1]: m_1 = 1, l_1 = 1 + 1/e,
//                            P_1 = [1/(e+1), e/(e+1)], O_1 = P_1 V
// Run with 1x1 blocks (row 1 needs a merge) and with one 2x2 block.
TEST(AttentionForwardTiledMaskedFixed, HandComputed2x2Causal) {
  std::vector<float> Q = {1, 0, 0, 1};
  std::vector<float> K = {1, 0, 0, 1};
  std::vector<float> V = {1, 2, 3, 4};
  const float tol = 1e-5f;

  for (int B : {1, 2}) {
    SCOPED_TRACE("block size " + std::to_string(B));
    auto r = run_tiled_masked(Q, K, V, 2, 2, B, B, 1.0f, masks::Causal{});
    EXPECT_NEAR(r.O[0], 1.0f, tol);
    EXPECT_NEAR(r.O[1], 2.0f, tol);
    EXPECT_NEAR(r.O[2], 2.462117157f, tol);
    EXPECT_NEAR(r.O[3], 3.462117157f, tol);
    EXPECT_NEAR(r.l[0], 1.0f, tol);
    EXPECT_NEAR(r.l[1], 1.367879441f, tol);
    EXPECT_NEAR(r.m[0], 1.0f, tol);
    EXPECT_NEAR(r.m[1], 1.0f, tol);
  }
}

// Same inputs with scale 2 and no mask: S = 2I, so
//   P = [[e^2, 1], [1, e^2]] / (e^2 + 1),  l = 1 + e^-2,  m = 2.
TEST(AttentionForwardTiledMaskedFixed, HandComputed2x2Scale) {
  std::vector<float> Q = {1, 0, 0, 1};
  std::vector<float> K = {1, 0, 0, 1};
  std::vector<float> V = {1, 2, 3, 4};
  auto r = run_tiled_masked(Q, K, V, 2, 2, 1, 1, 2.0f, masks::None{});

  const float tol = 1e-5f;
  const float c = 0.119202922f; // 1 / (e^2 + 1)
  EXPECT_NEAR(r.O[0], 1.0f + 2 * c, tol);
  EXPECT_NEAR(r.O[1], 2.0f + 2 * c, tol);
  EXPECT_NEAR(r.O[2], 3.0f - 2 * c, tol);
  EXPECT_NEAR(r.O[3], 4.0f - 2 * c, tol);
  EXPECT_NEAR(r.l[0], 1.135335283f, tol);
  EXPECT_NEAR(r.l[1], 1.135335283f, tol);
  EXPECT_NEAR(r.m[0], 2.0f, tol);
  EXPECT_NEAR(r.m[1], 2.0f, tol);
}

// l and m must be the row statistics of the softmax the naive version
// materializes: m_i is the row max of S, and exp(S_ij - m_i) / l_i is P_ij.
TEST(AttentionForwardTiledMaskedFixed, LAndMReproduceNaiveP) {
  const int N = 40, d = 8;
  Buffers b(N, d);
  fill_random(b.Q, -1.0f, 1.0f, 30);
  fill_random(b.K, -1.0f, 1.0f, 31);
  fill_random(b.V, -1.0f, 1.0f, 32);
  b.run();

  auto r = run_tiled_masked(b.Q, b.K, b.V, N, d, 8, 12, 1.0f, masks::None{});
  for (int i = 0; i < N; ++i) {
    float row_max = -std::numeric_limits<float>::max();
    for (int j = 0; j < N; ++j)
      row_max = std::max(row_max, b.S[i * N + j]);
    EXPECT_NEAR(r.m[i], row_max, 1e-6f) << "row " << i;

    for (int j = 0; j < N; ++j)
      EXPECT_NEAR(std::exp(b.S[i * N + j] - r.m[i]) / r.l[i], b.P[i * N + j],
                  1e-6f)
          << "i=" << i << " j=" << j;
  }
}

// Scaling the scores by tau is the same as scaling Q by tau, so the two ways
// of asking for it must agree on O, l, and m.
TEST(AttentionForwardTiledMaskedFixed, ScaleEqualsPrescalingQ) {
  const int N = 50, d = 16;
  const float tau = 0.3f;
  std::vector<float> Q(N * d), K(N * d), V(N * d);
  fill_random(Q, -2.0f, 2.0f, 33);
  fill_random(K, -2.0f, 2.0f, 34);
  fill_random(V, -1.0f, 1.0f, 35);

  std::vector<float> Qs = Q;
  for (float &x : Qs)
    x *= tau;

  auto scaled = run_tiled_masked(Q, K, V, N, d, 16, 7, tau, masks::None{});
  auto prescaled = run_tiled_masked(Qs, K, V, N, d, 16, 7, 1.0f, masks::None{});
  expect_result_close(scaled, prescaled, 1e-5f, 1e-4f);
}

// Under the causal mask, query i attends to keys 0..i only, so its output must
// equal the last row of unmasked attention over the first i + 1 positions. The
// naive reference has no scale parameter, so Q is pre-scaled instead.
TEST(AttentionForwardTiledMaskedFixed, CausalMatchesPrefixAttention) {
  const int N = 24, d = 8;
  const float scale = 1.0f / std::sqrt(static_cast<float>(d));
  std::vector<float> Q(N * d), K(N * d), V(N * d);
  fill_random(Q, -2.0f, 2.0f, 27);
  fill_random(K, -2.0f, 2.0f, 28);
  fill_random(V, -1.0f, 1.0f, 29);

  auto got = run_tiled_masked(Q, K, V, N, d, 5, 7, scale, masks::Causal{});

  for (int i = 0; i < N; ++i) {
    const int n = i + 1;
    Buffers b(n, d);
    for (int t = 0; t < n * d; ++t) {
      b.Q[t] = scale * Q[t];
      b.K[t] = K[t];
      b.V[t] = V[t];
    }
    b.run();

    for (int c = 0; c < d; ++c)
      EXPECT_NEAR(got.O[i * d + c], b.O[i * d + c], 1e-5f)
          << "i=" << i << " c=" << c;

    float row_max = -std::numeric_limits<float>::max();
    for (int j = 0; j < n; ++j)
      row_max = std::max(row_max, b.S[i * n + j]);
    double row_sum = 0;
    for (int j = 0; j < n; ++j)
      row_sum += std::exp(b.S[i * n + j] - row_max);
    EXPECT_NEAR(got.m[i], row_max, 1e-5f) << "i=" << i;
    EXPECT_NEAR(got.l[i], row_sum, 1e-5f + 1e-4f * row_sum) << "i=" << i;
  }
}

// Keys and values behind the padding mask must not influence any output, so
// replacing them with unrelated data leaves O, l, and m unchanged. With
// Bc = 12 and 29 valid keys, block 24..35 is partly padding and block 36..39
// is entirely padding.
TEST(AttentionForwardTiledMaskedFixed, PaddedKeysDoNotInfluenceOutput) {
  const int N = 40, d = 8, valid = 29;
  std::vector<float> Q(N * d), K(N * d), V(N * d);
  fill_random(Q, -1.0f, 1.0f, 36);
  fill_random(K, -1.0f, 1.0f, 37);
  fill_random(V, -1.0f, 1.0f, 38);
  auto base =
      run_tiled_masked(Q, K, V, N, d, 8, 12, 0.5f, masks::KeyPadding{valid});

  std::vector<float> garbage(N * d);
  fill_random(garbage, -100.0f, 100.0f, 39);
  std::vector<float> K2 = K, V2 = V;
  for (int t = valid * d; t < N * d; ++t) {
    K2[t] = garbage[t];
    V2[t] = -garbage[t];
  }
  auto got =
      run_tiled_masked(Q, K2, V2, N, d, 8, 12, 0.5f, masks::KeyPadding{valid});
  expect_result_close(got, base, 1e-6f, 1e-6f);
}

// Query 0 of a causal attention sees only key 0: its output is exactly the
// first row of V, l_0 is exactly 1, and m_0 is the scaled score of that one
// kept entry.
TEST(AttentionForwardTiledMaskedFixed, CausalRowZeroCopiesFirstValueRow) {
  const int N = 37, d = 8;
  const float scale = 0.35f;
  std::vector<float> Q(N * d), K(N * d), V(N * d);
  fill_random(Q, -1.0f, 1.0f, 41);
  fill_random(K, -1.0f, 1.0f, 42);
  fill_random(V, -1.0f, 1.0f, 43);

  auto r = run_tiled_masked(Q, K, V, N, d, 5, 6, scale, masks::Causal{});
  for (int c = 0; c < d; ++c)
    EXPECT_FLOAT_EQ(r.O[c], V[c]) << "c=" << c;
  EXPECT_FLOAT_EQ(r.l[0], 1.0f);

  float s00 = 0;
  for (int k = 0; k < d; ++k)
    s00 += Q[k] * K[k];
  EXPECT_NEAR(r.m[0], scale * s00, 1e-6f);
}

// Causal results must agree across block shapes, including shapes whose
// blocks straddle the diagonal differently.
TEST(AttentionForwardTiledMaskedFixed, BlockShapeInvarianceCausal) {
  const int N = 50, d = 16;
  const float scale = 1.0f / std::sqrt(static_cast<float>(d));
  std::vector<float> Q(N * d), K(N * d), V(N * d);
  fill_random(Q, -2.0f, 2.0f, 44);
  fill_random(K, -2.0f, 2.0f, 45);
  fill_random(V, -1.0f, 1.0f, 46);

  auto base = run_tiled_masked(Q, K, V, N, d, N, N, scale, masks::Causal{});
  for (auto [Br, Bc] : {std::pair{1, 1}, std::pair{16, 16}, std::pair{7, 13}}) {
    SCOPED_TRACE("Br=" + std::to_string(Br) + " Bc=" + std::to_string(Bc));
    auto got = run_tiled_masked(Q, K, V, N, d, Br, Bc, scale, masks::Causal{});
    expect_result_close(got, base, 1e-5f, 1e-4f);
  }
}

// ---------------------------------------------------------------------------
// attention_backward_naive: Algorithm 3, standard attention backward
//
// Faithful to the paper: Algorithm 3 is the backward of Algorithm 0, whose
// scores are S = Q K^T with no scale. Scaled attention is handled at the call
// site by passing tau * Q as the query and multiplying the returned dQ by tau
// (see ScaledAttentionViaPrescaledQuery); dK and dV need no correction.
// ---------------------------------------------------------------------------

namespace {

struct BackwardResult {
  std::vector<float> dQ, dK, dV;
};

// Runs the backward. Outputs are pre-filled with NaN so anything the
// implementation forgets to write, or accumulates into, fails loudly.
BackwardResult run_backward_naive(const std::vector<float> &Q,
                                  const std::vector<float> &K,
                                  const std::vector<float> &V,
                                  const std::vector<float> &dO,
                                  const std::vector<float> &P, int N, int d) {
  const float nan = std::numeric_limits<float>::quiet_NaN();
  BackwardResult r{std::vector<float>(N * d, nan),
                   std::vector<float>(N * d, nan),
                   std::vector<float>(N * d, nan)};
  attention_backward_naive(Q.data(), K.data(), V.data(), dO.data(), P.data(),
                           r.dQ.data(), r.dK.data(), r.dV.data(), N, d);
  return r;
}

// P = softmax(Q K^T) through the real pipeline, the naive forward.
std::vector<float> naive_P(const std::vector<float> &Q,
                           const std::vector<float> &K,
                           const std::vector<float> &V, int N, int d) {
  std::vector<float> S(N * N), P(N * N), O(N * d);
  attention_forward_naive(Q.data(), K.data(), V.data(), S.data(), P.data(),
                          O.data(), N, d);
  return P;
}

// P = softmax(mask(Q K^T)) computed in double from the definition, because
// the naive forward has no mask yet. Masked entries are exactly zero.
template <class Mask>
std::vector<float> masked_P(const std::vector<float> &Q,
                            const std::vector<float> &K, int N, int d,
                            Mask mask) {
  std::vector<float> P(N * N, 0.0f);
  std::vector<double> s(N);
  for (int i = 0; i < N; ++i) {
    double mi = -std::numeric_limits<double>::infinity();
    for (int j = 0; j < N; ++j) {
      if (mask(i, j))
        continue;
      double dot = 0;
      for (int k = 0; k < d; ++k)
        dot += static_cast<double>(Q[i * d + k]) * K[j * d + k];
      s[j] = dot;
      mi = std::max(mi, s[j]);
    }
    double li = 0;
    for (int j = 0; j < N; ++j)
      if (!mask(i, j))
        li += std::exp(s[j] - mi);
    for (int j = 0; j < N; ++j)
      if (!mask(i, j))
        P[i * N + j] = static_cast<float>(std::exp(s[j] - mi) / li);
  }
  return P;
}

std::vector<double> to_double(const std::vector<float> &v) {
  return std::vector<double>(v.begin(), v.end());
}

// Double-precision forward returning only O, for the finite-difference check.
// Takes a scale so the same helper can later check the tiled backward.
template <class Mask>
std::vector<double> forward_double(const std::vector<double> &Q,
                                   const std::vector<double> &K,
                                   const std::vector<double> &V, int N, int d,
                                   double scale, Mask mask) {
  std::vector<double> O(N * d, 0.0), s(N);
  for (int i = 0; i < N; ++i) {
    double mi = -std::numeric_limits<double>::infinity();
    for (int j = 0; j < N; ++j) {
      if (mask(i, j))
        continue;
      double dot = 0;
      for (int k = 0; k < d; ++k)
        dot += Q[i * d + k] * K[j * d + k];
      s[j] = scale * dot;
      mi = std::max(mi, s[j]);
    }
    double li = 0;
    for (int j = 0; j < N; ++j) {
      if (mask(i, j))
        continue;
      const double p = std::exp(s[j] - mi);
      li += p;
      for (int c = 0; c < d; ++c)
        O[i * d + c] += p * V[j * d + c];
    }
    for (int c = 0; c < d; ++c)
      O[i * d + c] /= li;
  }
  return O;
}

// The scalar loss L = <dO, O(Q, K, V)>. Its gradients with respect to Q, K
// and V are exactly what a backward must return for this dO.
template <class Mask>
double loss_double(const std::vector<double> &Q, const std::vector<double> &K,
                   const std::vector<double> &V, const std::vector<double> &dO,
                   int N, int d, double scale, Mask mask) {
  auto O = forward_double(Q, K, V, N, d, scale, mask);
  double L = 0;
  for (size_t t = 0; t < O.size(); ++t)
    L += dO[t] * O[t];
  return L;
}

enum class WhichInput { Q, K, V };

// Central finite differences of the loss with respect to one input, in
// double. With eps = 1e-4 the truncation error is around 1e-8 and rounding
// around 1e-12, far below the float rounding of the analytic gradient.
template <class Mask>
std::vector<float> fd_gradient(WhichInput which, std::vector<double> Q,
                               std::vector<double> K, std::vector<double> V,
                               const std::vector<double> &dO, int N, int d,
                               double scale, Mask mask) {
  std::vector<double> &X =
      which == WhichInput::Q ? Q : which == WhichInput::K ? K : V;
  const double eps = 1e-4;
  std::vector<float> g(X.size());
  for (size_t t = 0; t < X.size(); ++t) {
    const double x0 = X[t];
    X[t] = x0 + eps;
    const double lp = loss_double(Q, K, V, dO, N, d, scale, mask);
    X[t] = x0 - eps;
    const double lm = loss_double(Q, K, V, dO, N, d, scale, mask);
    X[t] = x0;
    g[t] = static_cast<float>((lp - lm) / (2 * eps));
  }
  return g;
}

// Compares gradients dQ, dK, dV with finite differences of the loss for
// attention with the given scale and mask.
template <class Mask>
void expect_matches_finite_differences(const std::vector<float> &dQ,
                                       const std::vector<float> &dK,
                                       const std::vector<float> &dV,
                                       const std::vector<float> &Q,
                                       const std::vector<float> &K,
                                       const std::vector<float> &V,
                                       const std::vector<float> &dO, int N,
                                       int d, float scale, Mask mask) {
  const auto Qd = to_double(Q), Kd = to_double(K), Vd = to_double(V),
             dOd = to_double(dO);
  {
    SCOPED_TRACE("dQ");
    expect_close(dQ,
                 fd_gradient(WhichInput::Q, Qd, Kd, Vd, dOd, N, d, scale, mask),
                 1e-4f, 1e-3f);
  }
  {
    SCOPED_TRACE("dK");
    expect_close(dK,
                 fd_gradient(WhichInput::K, Qd, Kd, Vd, dOd, N, d, scale, mask),
                 1e-4f, 1e-3f);
  }
  {
    SCOPED_TRACE("dV");
    expect_close(dV,
                 fd_gradient(WhichInput::V, Qd, Kd, Vd, dOd, N, d, scale, mask),
                 1e-4f, 1e-3f);
  }
}

struct BackwardCase {
  int N, d;
};

std::ostream &operator<<(std::ostream &os, const BackwardCase &c) {
  return os << "N=" << c.N << " d=" << c.d;
}

} // namespace

class AttentionBackwardNaive : public ::testing::TestWithParam<BackwardCase> {
};

// The analytic gradients must agree with central finite differences of the
// loss <dO, O>, computed in double, for unmasked attention at several sizes.
// P comes from the naive forward, as it will in practice.
TEST_P(AttentionBackwardNaive, MatchesFiniteDifferences) {
  const BackwardCase c = GetParam();
  std::vector<float> Q(c.N * c.d), K(c.N * c.d), V(c.N * c.d), dO(c.N * c.d);
  fill_random(Q, -1.0f, 1.0f, 47);
  fill_random(K, -1.0f, 1.0f, 48);
  fill_random(V, -1.0f, 1.0f, 49);
  fill_random(dO, -1.0f, 1.0f, 50);

  auto P = naive_P(Q, K, V, c.N, c.d);
  auto r = run_backward_naive(Q, K, V, dO, P, c.N, c.d);
  expect_matches_finite_differences(r.dQ, r.dK, r.dV, Q, K, V, dO, c.N, c.d,
                                    1.0f, masks::None{});
}

INSTANTIATE_TEST_SUITE_P(Sizes, AttentionBackwardNaive,
                         ::testing::Values(BackwardCase{6, 4}, BackwardCase{5, 3},
                                           BackwardCase{8, 2}, BackwardCase{7, 5},
                                           BackwardCase{12, 8}));

// Hand-computed: Q = K = I, V = [[1, 2], [3, 4]], dO = I.
//   P  = [[a, c], [c, a]] with a = e/(e+1), c = 1/(e+1)
//   dV = P^T dO = P
//   dP = dO V^T = [[1, 3], [2, 4]],  D_i = sum_j P_ij dP_ij
//   dS = P o (dP - D) = 2ac * [[-1, 1], [-1, 1]]      (rows sum to zero)
//   dQ = dS K = dS,  dK = dS^T Q = dS^T
TEST(AttentionBackwardNaiveFixed, HandComputed2x2) {
  std::vector<float> Q = {1, 0, 0, 1};
  std::vector<float> K = {1, 0, 0, 1};
  std::vector<float> V = {1, 2, 3, 4};
  std::vector<float> dO = {1, 0, 0, 1};
  const float a = 0.731058579f, c = 0.268941421f;
  std::vector<float> P = {a, c, c, a};

  auto r = run_backward_naive(Q, K, V, dO, P, 2, 2);

  const float tol = 1e-5f;
  const float g = 0.393223866f; // 2ac = 2e / (e+1)^2
  EXPECT_NEAR(r.dQ[0], -g, tol);
  EXPECT_NEAR(r.dQ[1], g, tol);
  EXPECT_NEAR(r.dQ[2], -g, tol);
  EXPECT_NEAR(r.dQ[3], g, tol);
  EXPECT_NEAR(r.dK[0], -g, tol);
  EXPECT_NEAR(r.dK[1], -g, tol);
  EXPECT_NEAR(r.dK[2], g, tol);
  EXPECT_NEAR(r.dK[3], g, tol);
  EXPECT_NEAR(r.dV[0], a, tol);
  EXPECT_NEAR(r.dV[1], c, tol);
  EXPECT_NEAR(r.dV[2], c, tol);
  EXPECT_NEAR(r.dV[3], a, tol);
}

// Same case with the query pre-scaled, Q = 2I, which is how scaled attention
// reaches this backward. S = 2I, so P uses a2 = e^2/(e^2+1), c2 = 1/(e^2+1),
// and dS = 2 a2 c2 * [[-1, 1], [-1, 1]]. Then dQ = dS K = dS, while
// dK = dS^T Q = 2 dS^T picks up the factor 2 from Q itself.
TEST(AttentionBackwardNaiveFixed, HandComputed2x2PrescaledQuery) {
  std::vector<float> Q = {2, 0, 0, 2};
  std::vector<float> K = {1, 0, 0, 1};
  std::vector<float> V = {1, 2, 3, 4};
  std::vector<float> dO = {1, 0, 0, 1};
  const float a2 = 0.880797078f, c2 = 0.119202922f;
  std::vector<float> P = {a2, c2, c2, a2};

  auto r = run_backward_naive(Q, K, V, dO, P, 2, 2);

  const float tol = 1e-5f;
  const float h = 0.209987171f; // 2 a2 c2 = 2e^2 / (e^2+1)^2
  EXPECT_NEAR(r.dQ[0], -h, tol);
  EXPECT_NEAR(r.dQ[1], h, tol);
  EXPECT_NEAR(r.dQ[2], -h, tol);
  EXPECT_NEAR(r.dQ[3], h, tol);
  EXPECT_NEAR(r.dK[0], -2 * h, tol);
  EXPECT_NEAR(r.dK[1], -2 * h, tol);
  EXPECT_NEAR(r.dK[2], 2 * h, tol);
  EXPECT_NEAR(r.dK[3], 2 * h, tol);
  EXPECT_NEAR(r.dV[0], a2, tol);
  EXPECT_NEAR(r.dV[1], c2, tol);
  EXPECT_NEAR(r.dV[2], c2, tol);
  EXPECT_NEAR(r.dV[3], a2, tol);
}

// The backward takes no mask: P already encodes it. With a causal P it must
// therefore match finite differences of the causal loss.
TEST(AttentionBackwardNaiveFixed, CausalMatchesFiniteDifferences) {
  const int N = 7, d = 4;
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO(N * d);
  fill_random(Q, -1.0f, 1.0f, 51);
  fill_random(K, -1.0f, 1.0f, 52);
  fill_random(V, -1.0f, 1.0f, 53);
  fill_random(dO, -1.0f, 1.0f, 54);

  auto P = masked_P(Q, K, N, d, masks::Causal{});
  auto r = run_backward_naive(Q, K, V, dO, P, N, d);
  expect_matches_finite_differences(r.dQ, r.dK, r.dV, Q, K, V, dO, N, d, 1.0f,
                                    masks::Causal{});
}

// Padded keys have P_ij = 0 in every row, so their dK and dV rows must be
// exactly zero, and the remaining gradients must match finite differences.
TEST(AttentionBackwardNaiveFixed, KeyPaddingZeroesPaddedRowsAndMatchesFD) {
  const int N = 9, d = 4, valid = 6;
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO(N * d);
  fill_random(Q, -1.0f, 1.0f, 55);
  fill_random(K, -1.0f, 1.0f, 56);
  fill_random(V, -1.0f, 1.0f, 57);
  fill_random(dO, -1.0f, 1.0f, 58);

  const masks::KeyPadding mask{valid};
  auto P = masked_P(Q, K, N, d, mask);
  auto r = run_backward_naive(Q, K, V, dO, P, N, d);

  for (int j = valid; j < N; ++j)
    for (int c = 0; c < d; ++c) {
      EXPECT_EQ(r.dK[j * d + c], 0.0f) << "j=" << j << " c=" << c;
      EXPECT_EQ(r.dV[j * d + c], 0.0f) << "j=" << j << " c=" << c;
    }
  expect_matches_finite_differences(r.dQ, r.dK, r.dV, Q, K, V, dO, N, d, 1.0f,
                                    mask);
}

// The calling convention for scaled attention, checked against finite
// differences of the scaled loss: run the unscaled backward on tau * Q with
// P = softmax(tau Q K^T), then dQ_scaled = tau * dQ, while dK and dV are
// already correct because dS^T (tau Q) = tau dS^T Q. This is exactly what
// the Algorithm 4 tests will use to compare against this reference.
TEST(AttentionBackwardNaiveFixed, ScaledAttentionViaPrescaledQuery) {
  const int N = 10, d = 6;
  const float tau = 0.3f;
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO(N * d);
  fill_random(Q, -1.0f, 1.0f, 59);
  fill_random(K, -1.0f, 1.0f, 60);
  fill_random(V, -1.0f, 1.0f, 61);
  fill_random(dO, -1.0f, 1.0f, 62);

  std::vector<float> Qs = Q;
  for (float &x : Qs)
    x *= tau;

  auto P = naive_P(Qs, K, V, N, d);
  auto r = run_backward_naive(Qs, K, V, dO, P, N, d);
  std::vector<float> dQ = r.dQ;
  for (float &x : dQ)
    x *= tau;

  expect_matches_finite_differences(dQ, r.dK, r.dV, Q, K, V, dO, N, d, tau,
                                    masks::None{});
}

// If every key is the same vector k, S is constant along each row, P is
// uniform, and O = mean(V) does not depend on Q at all, so dQ must vanish.
// Uniform P also makes every row of dV the column mean of dO.
TEST(AttentionBackwardNaiveFixed, IdenticalKeysGiveZeroQueryGradient) {
  const int N = 12, d = 5;
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO(N * d), key(d);
  fill_random(Q, -1.0f, 1.0f, 63);
  fill_random(V, -1.0f, 1.0f, 64);
  fill_random(dO, -1.0f, 1.0f, 65);
  fill_random(key, -1.0f, 1.0f, 66);
  for (int j = 0; j < N; ++j)
    for (int k = 0; k < d; ++k)
      K[j * d + k] = key[k];

  auto P = naive_P(Q, K, V, N, d);
  auto r = run_backward_naive(Q, K, V, dO, P, N, d);

  for (int t = 0; t < N * d; ++t)
    EXPECT_NEAR(r.dQ[t], 0.0f, 1e-5f) << "index " << t;

  std::vector<float> mean(d, 0.0f);
  for (int i = 0; i < N; ++i)
    for (int c = 0; c < d; ++c)
      mean[c] += dO[i * d + c] / N;
  for (int j = 0; j < N; ++j)
    for (int c = 0; c < d; ++c)
      EXPECT_NEAR(r.dV[j * d + c], mean[c], 1e-5f) << "j=" << j << " c=" << c;
}

// Adding the same vector to every key shifts each row of S by a constant,
// which softmax ignores, so the gradient with respect to K summed over all
// keys is zero. Equivalently, every row of dS sums to zero. Holds with masks.
TEST(AttentionBackwardNaiveFixed, KeyGradientsSumToZero) {
  const int N = 9, d = 4;
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO(N * d);
  fill_random(Q, -1.0f, 1.0f, 67);
  fill_random(K, -1.0f, 1.0f, 68);
  fill_random(V, -1.0f, 1.0f, 69);
  fill_random(dO, -1.0f, 1.0f, 70);

  auto check = [&](const char *label, const std::vector<float> &P) {
    SCOPED_TRACE(label);
    auto r = run_backward_naive(Q, K, V, dO, P, N, d);
    for (int c = 0; c < d; ++c) {
      float sum = 0;
      for (int j = 0; j < N; ++j)
        sum += r.dK[j * d + c];
      EXPECT_NEAR(sum, 0.0f, 1e-5f) << "c=" << c;
    }
  };
  check("no mask", naive_P(Q, K, V, N, d));
  check("causal", masked_P(Q, K, N, d, masks::Causal{}));
}

// With Q = 50 I and K = I, P is one-hot to float precision. The softmax is
// saturated, so dS vanishes: dQ and dK are ~0 while dV = P^T dO = dO.
TEST(AttentionBackwardNaiveFixed, SaturatedSoftmaxHasVanishingScoreGradients) {
  const int N = 6;
  std::vector<float> Q(N * N, 0.0f), K(N * N, 0.0f), V(N * N), dO(N * N);
  for (int i = 0; i < N; ++i) {
    Q[i * N + i] = 50.0f;
    K[i * N + i] = 1.0f;
  }
  fill_random(V, -1.0f, 1.0f, 71);
  fill_random(dO, -1.0f, 1.0f, 72);

  auto P = naive_P(Q, K, V, N, N);
  auto r = run_backward_naive(Q, K, V, dO, P, N, N);

  for (int t = 0; t < N * N; ++t) {
    EXPECT_NEAR(r.dV[t], dO[t], 1e-5f) << "index " << t;
    EXPECT_NEAR(r.dQ[t], 0.0f, 1e-5f) << "index " << t;
    EXPECT_NEAR(r.dK[t], 0.0f, 1e-5f) << "index " << t;
  }
}

// The backward is linear in the upstream gradient dO.
TEST(AttentionBackwardNaiveFixed, LinearInUpstreamGradient) {
  const int N = 8, d = 4;
  const float alpha = 0.7f, beta = -1.3f;
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO1(N * d), dO2(N * d);
  fill_random(Q, -1.0f, 1.0f, 73);
  fill_random(K, -1.0f, 1.0f, 74);
  fill_random(V, -1.0f, 1.0f, 75);
  fill_random(dO1, -1.0f, 1.0f, 76);
  fill_random(dO2, -1.0f, 1.0f, 77);

  std::vector<float> dO(N * d);
  for (int t = 0; t < N * d; ++t)
    dO[t] = alpha * dO1[t] + beta * dO2[t];

  auto P = naive_P(Q, K, V, N, d);
  auto r1 = run_backward_naive(Q, K, V, dO1, P, N, d);
  auto r2 = run_backward_naive(Q, K, V, dO2, P, N, d);
  auto r = run_backward_naive(Q, K, V, dO, P, N, d);

  auto combine = [&](const std::vector<float> &a, const std::vector<float> &b) {
    std::vector<float> out(a.size());
    for (size_t t = 0; t < a.size(); ++t)
      out[t] = alpha * a[t] + beta * b[t];
    return out;
  };
  {
    SCOPED_TRACE("dQ");
    expect_close(r.dQ, combine(r1.dQ, r2.dQ), 1e-5f, 1e-4f);
  }
  {
    SCOPED_TRACE("dK");
    expect_close(r.dK, combine(r1.dK, r2.dK), 1e-5f, 1e-4f);
  }
  {
    SCOPED_TRACE("dV");
    expect_close(r.dV, combine(r1.dV, r2.dV), 1e-5f, 1e-4f);
  }
}

// Under the causal mask, query 0 sees a single key, so its softmax is the
// constant 1 and has no dependence on Q: row 0 of dQ must be zero.
TEST(AttentionBackwardNaiveFixed, CausalRowZeroHasZeroQueryGradient) {
  const int N = 9, d = 4;
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO(N * d);
  fill_random(Q, -1.0f, 1.0f, 78);
  fill_random(K, -1.0f, 1.0f, 79);
  fill_random(V, -1.0f, 1.0f, 80);
  fill_random(dO, -1.0f, 1.0f, 81);

  auto P = masked_P(Q, K, N, d, masks::Causal{});
  auto r = run_backward_naive(Q, K, V, dO, P, N, d);
  for (int c = 0; c < d; ++c)
    EXPECT_NEAR(r.dQ[c], 0.0f, 1e-6f) << "c=" << c;
}

// ---------------------------------------------------------------------------
// attention_backward_tiled<Mask>: Algorithm 4, FlashAttention backward
// ---------------------------------------------------------------------------

namespace {

// Runs the tiled forward to obtain O, l, m, then the tiled backward on the
// same inputs. Gradients are pre-filled with NaN so anything the
// implementation forgets to write fails loudly.
template <class Mask>
BackwardResult run_backward_tiled(const std::vector<float> &Q,
                                  const std::vector<float> &K,
                                  const std::vector<float> &V,
                                  const std::vector<float> &dO, int N, int d,
                                  int Br, int Bc, float scale, Mask mask) {
  auto fwd = run_tiled_masked(Q, K, V, N, d, Br, Bc, scale, mask);
  const float nan = std::numeric_limits<float>::quiet_NaN();
  BackwardResult r{std::vector<float>(N * d, nan),
                   std::vector<float>(N * d, nan),
                   std::vector<float>(N * d, nan)};
  attention_backward_tiled(Q.data(), K.data(), V.data(), fwd.O.data(),
                           dO.data(), fwd.l.data(), fwd.m.data(), r.dQ.data(),
                           r.dK.data(), r.dV.data(), N, d, Br, Bc, scale, mask);
  return r;
}

// Algorithm 3 as the reference for scaled, masked attention, through the
// pre-scaling recipe: P = softmax(mask(tau Q K^T)) built in double, the
// unscaled backward run on tau * Q, and dQ multiplied by tau afterwards.
template <class Mask>
BackwardResult reference_backward(const std::vector<float> &Q,
                                  const std::vector<float> &K,
                                  const std::vector<float> &V,
                                  const std::vector<float> &dO, int N, int d,
                                  float scale, Mask mask) {
  std::vector<float> Qs = Q;
  for (float &x : Qs)
    x *= scale;
  auto P = masked_P(Qs, K, N, d, mask);
  auto r = run_backward_naive(Qs, K, V, dO, P, N, d);
  for (float &x : r.dQ)
    x *= scale;
  return r;
}

void expect_gradients_close(const BackwardResult &got,
                            const BackwardResult &want, float atol,
                            float rtol) {
  {
    SCOPED_TRACE("dQ");
    expect_close(got.dQ, want.dQ, atol, rtol);
  }
  {
    SCOPED_TRACE("dK");
    expect_close(got.dK, want.dK, atol, rtol);
  }
  {
    SCOPED_TRACE("dV");
    expect_close(got.dV, want.dV, atol, rtol);
  }
}

} // namespace

class AttentionBackwardTiled : public ::testing::TestWithParam<TileCase> {};

// Every mask, scale 1/sqrt(d), every block shape: the tiled backward must
// reproduce Algorithm 3. dQ is accumulated block by block here, so expect
// rounding-level differences only.
TEST_P(AttentionBackwardTiled, MatchesAlgorithm3) {
  const TileCase c = GetParam();
  std::vector<float> Q(c.N * c.d), K(c.N * c.d), V(c.N * c.d), dO(c.N * c.d);
  fill_random(Q, -1.0f, 1.0f, 82);
  fill_random(K, -1.0f, 1.0f, 83);
  fill_random(V, -1.0f, 1.0f, 84);
  fill_random(dO, -1.0f, 1.0f, 85);
  const float scale = 1.0f / std::sqrt(static_cast<float>(c.d));

  auto check = [&](const char *label, auto mask) {
    SCOPED_TRACE(label);
    auto want = reference_backward(Q, K, V, dO, c.N, c.d, scale, mask);
    auto got = run_backward_tiled(Q, K, V, dO, c.N, c.d, c.Br, c.Bc, scale, mask);
    expect_gradients_close(got, want, 1e-4f, 1e-3f);
  };
  check("no mask", masks::None{});
  check("causal", masks::Causal{});
  check("key padding", masks::KeyPadding{(2 * c.N + 2) / 3});
  check("window", masks::Window{std::max(1, c.N / 4)});
}

INSTANTIATE_TEST_SUITE_P(BlockShapes, AttentionBackwardTiled,
                         ::testing::ValuesIn(block_shapes()));

// Same hand-computed case as the naive backward (Q = K = I, V = [[1, 2],
// [3, 4]], dO = I, scale 1), with 1x1 blocks so every gradient is assembled
// from several block updates, and with one 2x2 block.
TEST(AttentionBackwardTiledFixed, HandComputed2x2) {
  std::vector<float> Q = {1, 0, 0, 1};
  std::vector<float> K = {1, 0, 0, 1};
  std::vector<float> V = {1, 2, 3, 4};
  std::vector<float> dO = {1, 0, 0, 1};
  const float a = 0.731058579f, c = 0.268941421f;
  const float g = 0.393223866f; // 2ac = 2e / (e+1)^2
  const float tol = 1e-5f;

  for (int B : {1, 2}) {
    SCOPED_TRACE("block size " + std::to_string(B));
    auto r = run_backward_tiled(Q, K, V, dO, 2, 2, B, B, 1.0f, masks::None{});
    EXPECT_NEAR(r.dQ[0], -g, tol);
    EXPECT_NEAR(r.dQ[1], g, tol);
    EXPECT_NEAR(r.dQ[2], -g, tol);
    EXPECT_NEAR(r.dQ[3], g, tol);
    EXPECT_NEAR(r.dK[0], -g, tol);
    EXPECT_NEAR(r.dK[1], -g, tol);
    EXPECT_NEAR(r.dK[2], g, tol);
    EXPECT_NEAR(r.dK[3], g, tol);
    EXPECT_NEAR(r.dV[0], a, tol);
    EXPECT_NEAR(r.dV[1], c, tol);
    EXPECT_NEAR(r.dV[2], c, tol);
    EXPECT_NEAR(r.dV[3], a, tol);
  }
}

// Finite differences of the scaled loss, applied directly to the tiled
// backward, on a shape where neither Br nor Bc divides N.
TEST(AttentionBackwardTiledFixed, MatchesFiniteDifferences) {
  const int N = 7, d = 4, Br = 3, Bc = 2;
  const float scale = 0.5f;
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO(N * d);
  fill_random(Q, -1.0f, 1.0f, 86);
  fill_random(K, -1.0f, 1.0f, 87);
  fill_random(V, -1.0f, 1.0f, 88);
  fill_random(dO, -1.0f, 1.0f, 89);

  auto r = run_backward_tiled(Q, K, V, dO, N, d, Br, Bc, scale, masks::None{});
  expect_matches_finite_differences(r.dQ, r.dK, r.dV, Q, K, V, dO, N, d, scale,
                                    masks::None{});
}

// Causal, with blocks that straddle the diagonal.
TEST(AttentionBackwardTiledFixed, CausalMatchesFiniteDifferences) {
  const int N = 8, d = 3, Br = 3, Bc = 3;
  const float scale = 0.6f;
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO(N * d);
  fill_random(Q, -1.0f, 1.0f, 90);
  fill_random(K, -1.0f, 1.0f, 91);
  fill_random(V, -1.0f, 1.0f, 92);
  fill_random(dO, -1.0f, 1.0f, 93);

  auto r = run_backward_tiled(Q, K, V, dO, N, d, Br, Bc, scale, masks::Causal{});
  expect_matches_finite_differences(r.dQ, r.dK, r.dV, Q, K, V, dO, N, d, scale,
                                    masks::Causal{});
}

// Sliding window: for rows far from the start, whole leading K blocks are
// masked, so the backward meets tile rows with nothing kept.
TEST(AttentionBackwardTiledFixed, WindowMatchesFiniteDifferences) {
  const int N = 12, d = 3, Br = 4, Bc = 2;
  const float scale = 0.7f;
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO(N * d);
  fill_random(Q, -1.0f, 1.0f, 94);
  fill_random(K, -1.0f, 1.0f, 95);
  fill_random(V, -1.0f, 1.0f, 96);
  fill_random(dO, -1.0f, 1.0f, 97);

  const masks::Window mask{3};
  auto r = run_backward_tiled(Q, K, V, dO, N, d, Br, Bc, scale, mask);
  expect_matches_finite_differences(r.dQ, r.dK, r.dV, Q, K, V, dO, N, d, scale,
                                    mask);
}

// Padded keys must receive exactly zero dK and dV, including the key block
// that is entirely padding, and the rest must match finite differences.
TEST(AttentionBackwardTiledFixed, KeyPaddingZeroesPaddedRowsAndMatchesFD) {
  const int N = 9, d = 4, Br = 4, Bc = 3, valid = 5;
  const float scale = 0.5f;
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO(N * d);
  fill_random(Q, -1.0f, 1.0f, 98);
  fill_random(K, -1.0f, 1.0f, 99);
  fill_random(V, -1.0f, 1.0f, 100);
  fill_random(dO, -1.0f, 1.0f, 101);

  const masks::KeyPadding mask{valid};
  auto r = run_backward_tiled(Q, K, V, dO, N, d, Br, Bc, scale, mask);
  for (int j = valid; j < N; ++j)
    for (int c = 0; c < d; ++c) {
      EXPECT_EQ(r.dK[j * d + c], 0.0f) << "j=" << j << " c=" << c;
      EXPECT_EQ(r.dV[j * d + c], 0.0f) << "j=" << j << " c=" << c;
    }
  expect_matches_finite_differences(r.dQ, r.dK, r.dV, Q, K, V, dO, N, d, scale,
                                    mask);
}

// Different block shapes must agree with each other, including shapes whose
// blocks straddle the causal diagonal differently.
TEST(AttentionBackwardTiledFixed, BlockShapeInvariance) {
  const int N = 30, d = 8;
  const float scale = 1.0f / std::sqrt(static_cast<float>(d));
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO(N * d);
  fill_random(Q, -1.0f, 1.0f, 102);
  fill_random(K, -1.0f, 1.0f, 103);
  fill_random(V, -1.0f, 1.0f, 104);
  fill_random(dO, -1.0f, 1.0f, 105);

  auto base = run_backward_tiled(Q, K, V, dO, N, d, N, N, scale, masks::Causal{});
  for (auto [Br, Bc] : {std::pair{1, 1}, std::pair{8, 8}, std::pair{7, 13}}) {
    SCOPED_TRACE("Br=" + std::to_string(Br) + " Bc=" + std::to_string(Bc));
    auto got = run_backward_tiled(Q, K, V, dO, N, d, Br, Bc, scale, masks::Causal{});
    expect_gradients_close(got, base, 1e-5f, 1e-4f);
  }
}

// With dO = 0 every gradient is exactly zero. Nothing else is allowed to leak
// in, in particular nothing from padded rows or columns of a ragged block.
TEST(AttentionBackwardTiledFixed, ZeroUpstreamGradientGivesZeroGradients) {
  const int N = 10, d = 3, Br = 4, Bc = 3;
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO(N * d, 0.0f);
  fill_random(Q, -1.0f, 1.0f, 106);
  fill_random(K, -1.0f, 1.0f, 107);
  fill_random(V, -1.0f, 1.0f, 108);

  auto r = run_backward_tiled(Q, K, V, dO, N, d, Br, Bc, 0.5f, masks::None{});
  for (int t = 0; t < N * d; ++t) {
    EXPECT_EQ(r.dQ[t], 0.0f) << "index " << t;
    EXPECT_EQ(r.dK[t], 0.0f) << "index " << t;
    EXPECT_EQ(r.dV[t], 0.0f) << "index " << t;
  }
}

// The backward is linear in the upstream gradient dO.
TEST(AttentionBackwardTiledFixed, LinearInUpstreamGradient) {
  const int N = 8, d = 4, Br = 3, Bc = 5;
  const float scale = 0.5f, alpha = 0.7f, beta = -1.3f;
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO1(N * d), dO2(N * d);
  fill_random(Q, -1.0f, 1.0f, 109);
  fill_random(K, -1.0f, 1.0f, 110);
  fill_random(V, -1.0f, 1.0f, 111);
  fill_random(dO1, -1.0f, 1.0f, 112);
  fill_random(dO2, -1.0f, 1.0f, 113);

  std::vector<float> dO(N * d);
  for (int t = 0; t < N * d; ++t)
    dO[t] = alpha * dO1[t] + beta * dO2[t];

  auto r1 = run_backward_tiled(Q, K, V, dO1, N, d, Br, Bc, scale, masks::None{});
  auto r2 = run_backward_tiled(Q, K, V, dO2, N, d, Br, Bc, scale, masks::None{});
  auto r = run_backward_tiled(Q, K, V, dO, N, d, Br, Bc, scale, masks::None{});

  BackwardResult want{r1.dQ, r1.dK, r1.dV};
  for (std::vector<float> *out : {&want.dQ, &want.dK, &want.dV}) {
    const std::vector<float> &b =
        out == &want.dQ ? r2.dQ : out == &want.dK ? r2.dK : r2.dV;
    for (size_t t = 0; t < out->size(); ++t)
      (*out)[t] = alpha * (*out)[t] + beta * b[t];
  }
  expect_gradients_close(r, want, 1e-5f, 1e-4f);
}

// The backward must use the forward's statistics rather than recompute the
// softmax: P = exp(S - m) / l, so doubling l halves P and therefore halves
// dV = P^T dO exactly.
TEST(AttentionBackwardTiledFixed, DVScalesInverselyWithL) {
  const int N = 9, d = 4, Br = 4, Bc = 3;
  const float scale = 0.5f;
  std::vector<float> Q(N * d), K(N * d), V(N * d), dO(N * d);
  fill_random(Q, -1.0f, 1.0f, 114);
  fill_random(K, -1.0f, 1.0f, 115);
  fill_random(V, -1.0f, 1.0f, 116);
  fill_random(dO, -1.0f, 1.0f, 117);

  auto fwd = run_tiled_masked(Q, K, V, N, d, Br, Bc, scale, masks::None{});
  std::vector<float> l2 = fwd.l;
  for (float &x : l2)
    x *= 2.0f;

  const float nan = std::numeric_limits<float>::quiet_NaN();
  std::vector<float> dQ(N * d, nan), dK(N * d, nan), dV(N * d, nan);
  std::vector<float> dQ2(N * d, nan), dK2(N * d, nan), dV2(N * d, nan);
  attention_backward_tiled(Q.data(), K.data(), V.data(), fwd.O.data(),
                           dO.data(), fwd.l.data(), fwd.m.data(), dQ.data(),
                           dK.data(), dV.data(), N, d, Br, Bc, scale,
                           masks::None{});
  attention_backward_tiled(Q.data(), K.data(), V.data(), fwd.O.data(),
                           dO.data(), l2.data(), fwd.m.data(), dQ2.data(),
                           dK2.data(), dV2.data(), N, d, Br, Bc, scale,
                           masks::None{});
  for (int t = 0; t < N * d; ++t)
    EXPECT_NEAR(dV2[t], 0.5f * dV[t], 1e-6f) << "index " << t;
}