// Author: Simon-Pierre Boucher — contact@spboucher.ai
#include "ops/cpu/cpu_ops.h"

#include <dispatch/dispatch.h>

#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <limits>

namespace forge::cpu {

namespace {

void check(bool cond, const char* msg) {
    if (!cond) {
        std::fprintf(stderr, "forge/cpu: %s\n", msg);
        std::abort();
    }
}

// Row-parallel helper (GCD). Serial below the threshold so tiny test
// tensors don't pay dispatch overhead.
template <typename F>
void parallel_rows(int64_t n, F&& body) {
    if (n < 32) {
        for (int64_t i = 0; i < n; ++i) body(i);
    } else {
        dispatch_apply(size_t(n), DISPATCH_APPLY_AUTO,
                       ^(size_t i) { body(int64_t(i)); });
    }
}

int32_t token_at(const Tensor& ids, int64_t i) {
    return ids.dtype() == DType::I32 ? ids.data<int32_t>()[i]
                                     : int32_t(ids.data<uint16_t>()[i]);
}

} // namespace

// ---- init -------------------------------------------------------------------

void fill_normal(Tensor& t, float mean, float stddev, std::mt19937_64& rng) {
    std::normal_distribution<float> dist(mean, stddev);
    float* p = t.data<float>();
    for (int64_t i = 0; i < t.numel(); ++i) p[i] = dist(rng);
}

void fill_uniform_int(Tensor& t, int64_t low, int64_t high, std::mt19937_64& rng) {
    std::uniform_int_distribution<int64_t> dist(low, high - 1);
    for (int64_t i = 0; i < t.numel(); ++i) t.set_item(i, float(dist(rng)));
}

// ---- matmul -----------------------------------------------------------------

void matmul(const Tensor& a, const Tensor& b, Tensor& c,
            bool transpose_a, bool transpose_b, bool accumulate) {
    check(a.dtype() == DType::F32 && b.dtype() == DType::F32 && c.dtype() == DType::F32,
          "matmul: f32 only");
    check(a.ndim() == 2 && b.ndim() == 2 && c.ndim() == 2, "matmul: 2-D only");
    check(a.is_contiguous() && b.is_contiguous() && c.is_contiguous(),
          "matmul: contiguous only");

    const int64_t M = transpose_a ? a.size(1) : a.size(0);
    const int64_t K = transpose_a ? a.size(0) : a.size(1);
    const int64_t Kb = transpose_b ? b.size(1) : b.size(0);
    const int64_t N = transpose_b ? b.size(0) : b.size(1);
    check(K == Kb, "matmul: inner dims mismatch");
    check(c.size(0) == M && c.size(1) == N, "matmul: output shape mismatch");

    const float* A = a.data<float>();
    const float* B = b.data<float>();
    float* C = c.data<float>();
    const int64_t lda = a.size(1);
    const int64_t ldb = b.size(1);

    // i-k-j order streams contiguous rows of B and C in the common case.
    parallel_rows(M, [&](int64_t i) {
        float* crow = C + i * N;
        if (!accumulate)
            for (int64_t j = 0; j < N; ++j) crow[j] = 0.0f;
        for (int64_t k = 0; k < K; ++k) {
            const float aik = transpose_a ? A[k * lda + i] : A[i * lda + k];
            if (!transpose_b) {
                const float* brow = B + k * ldb;
                for (int64_t j = 0; j < N; ++j) crow[j] += aik * brow[j];
            } else {
                for (int64_t j = 0; j < N; ++j) crow[j] += aik * B[j * ldb + k];
            }
        }
    });
}

// ---- elementwise ------------------------------------------------------------

void add(const Tensor& a, const Tensor& b, Tensor& out) {
    check(a.numel() == b.numel() && a.numel() == out.numel(), "add: numel mismatch");
    const float* pa = a.data<float>();
    const float* pb = b.data<float>();
    float* po = out.data<float>();
    for (int64_t i = 0; i < a.numel(); ++i) po[i] = pa[i] + pb[i];
}

void mul(const Tensor& a, const Tensor& b, Tensor& out) {
    check(a.numel() == b.numel() && a.numel() == out.numel(), "mul: numel mismatch");
    const float* pa = a.data<float>();
    const float* pb = b.data<float>();
    float* po = out.data<float>();
    for (int64_t i = 0; i < a.numel(); ++i) po[i] = pa[i] * pb[i];
}

void scale(const Tensor& a, float s, Tensor& out) {
    check(a.numel() == out.numel(), "scale: numel mismatch");
    const float* pa = a.data<float>();
    float* po = out.data<float>();
    for (int64_t i = 0; i < a.numel(); ++i) po[i] = pa[i] * s;
}

void add_bias(const Tensor& x, const Tensor& bias, Tensor& out) {
    const int64_t C = bias.numel();
    const int64_t N = x.numel() / C;
    const float* px = x.data<float>();
    const float* pb = bias.data<float>();
    float* po = out.data<float>();
    for (int64_t i = 0; i < N; ++i)
        for (int64_t j = 0; j < C; ++j) po[i * C + j] = px[i * C + j] + pb[j];
}

void add_bias_backward(const Tensor& dout, Tensor& dbias) {
    const int64_t C = dbias.numel();
    const int64_t N = dout.numel() / C;
    const float* pd = dout.data<float>();
    float* pb = dbias.data<float>();
    for (int64_t i = 0; i < N; ++i)
        for (int64_t j = 0; j < C; ++j) pb[j] += pd[i * C + j];
}

void silu(const Tensor& x, Tensor& out) {
    const float* px = x.data<float>();
    float* po = out.data<float>();
    for (int64_t i = 0; i < x.numel(); ++i) {
        const float v = px[i];
        po[i] = v / (1.0f + std::exp(-v));
    }
}

void silu_backward(const Tensor& x, const Tensor& dout, Tensor& dx) {
    const float* px = x.data<float>();
    const float* pd = dout.data<float>();
    float* pdx = dx.data<float>();
    for (int64_t i = 0; i < x.numel(); ++i) {
        const float v = px[i];
        const float sig = 1.0f / (1.0f + std::exp(-v));
        pdx[i] += pd[i] * sig * (1.0f + v * (1.0f - sig));
    }
}

void gelu(const Tensor& x, Tensor& out) {
    constexpr float k = 0.7978845608028654f; // sqrt(2/pi)
    const float* px = x.data<float>();
    float* po = out.data<float>();
    for (int64_t i = 0; i < x.numel(); ++i) {
        const float v = px[i];
        po[i] = 0.5f * v * (1.0f + std::tanh(k * (v + 0.044715f * v * v * v)));
    }
}

void gelu_backward(const Tensor& x, const Tensor& dout, Tensor& dx) {
    constexpr float k = 0.7978845608028654f;
    const float* px = x.data<float>();
    const float* pd = dout.data<float>();
    float* pdx = dx.data<float>();
    for (int64_t i = 0; i < x.numel(); ++i) {
        const float v = px[i];
        const float u = k * (v + 0.044715f * v * v * v);
        const float t = std::tanh(u);
        const float du = k * (1.0f + 3.0f * 0.044715f * v * v);
        pdx[i] += pd[i] * (0.5f * (1.0f + t) + 0.5f * v * (1.0f - t * t) * du);
    }
}

void relu2(const Tensor& x, Tensor& out) {
    const float* px = x.data<float>();
    float* po = out.data<float>();
    for (int64_t i = 0; i < x.numel(); ++i) {
        const float v = std::max(px[i], 0.0f);
        po[i] = v * v;
    }
}

void relu2_backward(const Tensor& x, const Tensor& dout, Tensor& dx) {
    const float* px = x.data<float>();
    const float* pd = dout.data<float>();
    float* pdx = dx.data<float>();
    for (int64_t i = 0; i < x.numel(); ++i)
        pdx[i] += pd[i] * 2.0f * std::max(px[i], 0.0f);
}

void softcap(const Tensor& x, float cap, Tensor& out) {
    const float* px = x.data<float>();
    float* po = out.data<float>();
    for (int64_t i = 0; i < x.numel(); ++i) po[i] = cap * std::tanh(px[i] / cap);
}

void softcap_backward(const Tensor& y, const Tensor& dout, float cap, Tensor& dx) {
    const float* py = y.data<float>();
    const float* pd = dout.data<float>();
    float* pdx = dx.data<float>();
    for (int64_t i = 0; i < y.numel(); ++i) {
        const float t = py[i] / cap;
        pdx[i] += pd[i] * (1.0f - t * t);
    }
}

// ---- norms --------------------------------------------------------------------

void rmsnorm(const Tensor& x, const Tensor& w, float eps, Tensor& out) {
    const int64_t C = w.numel();
    const int64_t N = x.numel() / C;
    const float* pw = w.data<float>();
    const float* px = x.data<float>();
    float* po = out.data<float>();
    parallel_rows(N, [&](int64_t i) {
        const float* row = px + i * C;
        float ss = 0.0f;
        for (int64_t j = 0; j < C; ++j) ss += row[j] * row[j];
        const float inv_rms = 1.0f / std::sqrt(ss / float(C) + eps);
        for (int64_t j = 0; j < C; ++j) po[i * C + j] = pw[j] * row[j] * inv_rms;
    });
}

void rmsnorm_backward(const Tensor& x, const Tensor& w, float eps,
                      const Tensor& dout, Tensor& dx, Tensor& dw) {
    const int64_t C = w.numel();
    const int64_t N = x.numel() / C;
    const float* px = x.data<float>();
    const float* pw = w.data<float>();
    const float* pd = dout.data<float>();
    float* pdx = dx.data<float>();
    float* pdw = dw.data<float>();
    // dw is a cross-row reduction: keep it serial for determinism.
    for (int64_t i = 0; i < N; ++i) {
        const float* row = px + i * C;
        const float* drow = pd + i * C;
        float ss = 0.0f;
        for (int64_t j = 0; j < C; ++j) ss += row[j] * row[j];
        const float inv_rms = 1.0f / std::sqrt(ss / float(C) + eps);
        float dot = 0.0f; // sum_j g_j w_j x_j
        for (int64_t j = 0; j < C; ++j) dot += drow[j] * pw[j] * row[j];
        const float coef = dot * inv_rms * inv_rms * inv_rms / float(C);
        for (int64_t j = 0; j < C; ++j) {
            pdx[i * C + j] += drow[j] * pw[j] * inv_rms - row[j] * coef;
            pdw[j] += drow[j] * row[j] * inv_rms;
        }
    }
}

void layernorm(const Tensor& x, const Tensor& w, const Tensor& b, float eps, Tensor& out) {
    const int64_t C = w.numel();
    const int64_t N = x.numel() / C;
    const float* px = x.data<float>();
    const float* pw = w.data<float>();
    const float* pb = b.data<float>();
    float* po = out.data<float>();
    parallel_rows(N, [&](int64_t i) {
        const float* row = px + i * C;
        float mean = 0.0f;
        for (int64_t j = 0; j < C; ++j) mean += row[j];
        mean /= float(C);
        float var = 0.0f;
        for (int64_t j = 0; j < C; ++j) var += (row[j] - mean) * (row[j] - mean);
        var /= float(C);
        const float inv_std = 1.0f / std::sqrt(var + eps);
        for (int64_t j = 0; j < C; ++j)
            po[i * C + j] = pw[j] * (row[j] - mean) * inv_std + pb[j];
    });
}

void layernorm_backward(const Tensor& x, const Tensor& w, float eps,
                        const Tensor& dout, Tensor& dx, Tensor& dw, Tensor& db) {
    const int64_t C = w.numel();
    const int64_t N = x.numel() / C;
    const float* px = x.data<float>();
    const float* pw = w.data<float>();
    const float* pd = dout.data<float>();
    float* pdx = dx.data<float>();
    float* pdw = dw.data<float>();
    float* pdb = db.data<float>();
    for (int64_t i = 0; i < N; ++i) {
        const float* row = px + i * C;
        const float* drow = pd + i * C;
        float mean = 0.0f;
        for (int64_t j = 0; j < C; ++j) mean += row[j];
        mean /= float(C);
        float var = 0.0f;
        for (int64_t j = 0; j < C; ++j) var += (row[j] - mean) * (row[j] - mean);
        var /= float(C);
        const float inv_std = 1.0f / std::sqrt(var + eps);
        // dxhat = g*w ; dx = inv_std * (dxhat − mean(dxhat) − xhat*mean(dxhat∘xhat))
        float mean_dxhat = 0.0f, mean_dxhat_xhat = 0.0f;
        for (int64_t j = 0; j < C; ++j) {
            const float xhat = (row[j] - mean) * inv_std;
            const float dxhat = drow[j] * pw[j];
            mean_dxhat += dxhat;
            mean_dxhat_xhat += dxhat * xhat;
        }
        mean_dxhat /= float(C);
        mean_dxhat_xhat /= float(C);
        for (int64_t j = 0; j < C; ++j) {
            const float xhat = (row[j] - mean) * inv_std;
            pdx[i * C + j] += inv_std * (drow[j] * pw[j] - mean_dxhat - xhat * mean_dxhat_xhat);
            pdw[j] += drow[j] * xhat;
            pdb[j] += drow[j];
        }
    }
}

// ---- softmax ------------------------------------------------------------------

void softmax(const Tensor& x, Tensor& out) {
    const int64_t C = x.shape().back();
    const int64_t N = x.numel() / C;
    const float* px = x.data<float>();
    float* po = out.data<float>();
    parallel_rows(N, [&](int64_t i) {
        const float* row = px + i * C;
        float m = row[0];
        for (int64_t j = 1; j < C; ++j) m = std::max(m, row[j]);
        float sum = 0.0f;
        for (int64_t j = 0; j < C; ++j) {
            const float e = std::exp(row[j] - m);
            po[i * C + j] = e;
            sum += e;
        }
        const float inv = 1.0f / sum;
        for (int64_t j = 0; j < C; ++j) po[i * C + j] *= inv;
    });
}

void softmax_backward(const Tensor& p, const Tensor& dout, Tensor& dx) {
    const int64_t C = p.shape().back();
    const int64_t N = p.numel() / C;
    const float* pp = p.data<float>();
    const float* pd = dout.data<float>();
    float* pdx = dx.data<float>();
    for (int64_t i = 0; i < N; ++i) {
        const float* prow = pp + i * C;
        const float* drow = pd + i * C;
        float dot = 0.0f;
        for (int64_t j = 0; j < C; ++j) dot += drow[j] * prow[j];
        for (int64_t j = 0; j < C; ++j) pdx[i * C + j] += prow[j] * (drow[j] - dot);
    }
}

// ---- QAT / MoE ------------------------------------------------------------------

void fake_quant(const Tensor& w, int mode, Tensor& out) {
    const int64_t R = w.size(0), C = w.size(1);
    const float* pw = w.data<float>();
    float* po = out.data<float>();
    for (int64_t i = 0; i < R; ++i) {
        const float* row = pw + i * C;
        float* orow = po + i * C;
        if (mode == 0) { // int8: symmetric absmax
            float amax = 0.0f;
            for (int64_t j = 0; j < C; ++j) amax = std::max(amax, std::fabs(row[j]));
            const float s = std::max(amax / 127.0f, 1e-12f);
            for (int64_t j = 0; j < C; ++j) orow[j] = std::rint(row[j] / s) * s;
        } else {         // ternary: BitNet b1.58 absmean
            float asum = 0.0f;
            for (int64_t j = 0; j < C; ++j) asum += std::fabs(row[j]);
            const float s = std::max(asum / float(C), 1e-12f);
            for (int64_t j = 0; j < C; ++j)
                orow[j] = std::min(1.0f, std::max(-1.0f, std::rint(row[j] / s))) * s;
        }
    }
}

void sigmoid(const Tensor& x, Tensor& out) {
    const float* px = x.data<float>();
    float* po = out.data<float>();
    for (int64_t i = 0; i < x.numel(); ++i) po[i] = 1.0f / (1.0f + std::exp(-px[i]));
}

void sigmoid_backward(const Tensor& y, const Tensor& dout, Tensor& dx) {
    const float* py = y.data<float>();
    const float* pd = dout.data<float>();
    float* pdx = dx.data<float>();
    for (int64_t i = 0; i < y.numel(); ++i) pdx[i] += pd[i] * py[i] * (1.0f - py[i]);
}

namespace {
// Kept set of the k largest (score + bias) entries, ties to the lower index
// (matches Metal). The returned sum is over the BIASLESS scores — the bias
// only steers selection (DeepSeek-V3 noaux routing).
void topk_select(const float* row, const float* bias, int64_t E, int64_t k,
                 bool* kept, float* sum) {
    for (int64_t j = 0; j < E; ++j) kept[j] = false;
    float S = 0.0f;
    for (int64_t sel = 0; sel < k; ++sel) {
        float best = -std::numeric_limits<float>::max();
        int64_t arg = 0;
        for (int64_t j = 0; j < E; ++j) {
            const float v = row[j] + bias[j];
            if (!kept[j] && v > best) { best = v; arg = j; }
        }
        kept[arg] = true;
        S += row[arg];
    }
    *sum = S;
}
} // namespace

void topk_renorm(const Tensor& p, const Tensor& bias, int64_t k, bool norm,
                 Tensor& out) {
    const int64_t E = p.shape().back();
    const int64_t N = p.numel() / E;
    const float* pp = p.data<float>();
    const float* pb = bias.data<float>();
    float* po = out.data<float>();
    bool kbuf[64];
    for (int64_t i = 0; i < N; ++i) {
        const float* row = pp + i * E;
        float S = 0.0f;
        topk_select(row, pb, E, k, kbuf, &S);
        const float inv = norm ? 1.0f / std::max(S, 1e-12f) : 1.0f;
        for (int64_t j = 0; j < E; ++j) po[i * E + j] = kbuf[j] ? row[j] * inv : 0.0f;
    }
}

void topk_renorm_backward(const Tensor& p, const Tensor& bias, const Tensor& dout,
                          int64_t k, bool norm, Tensor& dp) {
    const int64_t E = p.shape().back();
    const int64_t N = p.numel() / E;
    const float* pp = p.data<float>();
    const float* pb = bias.data<float>();
    const float* pd = dout.data<float>();
    float* pdp = dp.data<float>();
    bool kbuf[64];
    for (int64_t i = 0; i < N; ++i) {
        const float* row = pp + i * E;
        const float* drow = pd + i * E;
        float S = 0.0f;
        topk_select(row, pb, E, k, kbuf, &S);
        if (!norm) {
            for (int64_t j = 0; j < E; ++j)
                if (kbuf[j]) pdp[i * E + j] += drow[j];
            continue;
        }
        const float inv = 1.0f / std::max(S, 1e-12f);
        float dot = 0.0f;
        for (int64_t j = 0; j < E; ++j)
            if (kbuf[j]) dot += drow[j] * row[j] * inv;
        for (int64_t j = 0; j < E; ++j)
            if (kbuf[j]) pdp[i * E + j] += (drow[j] - dot) * inv;
    }
}

void expert_counts(const Tensor& gates, Tensor& counts) {
    const int64_t E = gates.shape().back();
    const int64_t N = gates.numel() / E;
    const float* pg = gates.data<float>();
    float* pc = counts.data<float>();
    for (int64_t i = 0; i < N; ++i)
        for (int64_t e = 0; e < E; ++e)
            if (pg[i * E + e] != 0.0f) pc[e] += 1.0f;
}

void row_scale(const Tensor& x, const Tensor& gates, int64_t e, Tensor& out) {
    const int64_t C = x.shape().back();
    const int64_t N = x.numel() / C;
    const int64_t E = gates.shape().back();
    const float* px = x.data<float>();
    const float* pg = gates.data<float>();
    float* po = out.data<float>();
    for (int64_t i = 0; i < N; ++i) {
        const float s = pg[i * E + e];
        for (int64_t j = 0; j < C; ++j) po[i * C + j] = px[i * C + j] * s;
    }
}

void row_scale_accumulate(const Tensor& x, const Tensor& gates, int64_t e, Tensor& dst) {
    const int64_t C = x.shape().back();
    const int64_t N = x.numel() / C;
    const int64_t E = gates.shape().back();
    const float* px = x.data<float>();
    const float* pg = gates.data<float>();
    float* pd = dst.data<float>();
    for (int64_t i = 0; i < N; ++i) {
        const float s = pg[i * E + e];
        for (int64_t j = 0; j < C; ++j) pd[i * C + j] += px[i * C + j] * s;
    }
}

void row_scale_gate_backward(const Tensor& dout, const Tensor& x, int64_t e,
                             Tensor& dgates) {
    const int64_t C = x.shape().back();
    const int64_t N = x.numel() / C;
    const int64_t E = dgates.shape().back();
    const float* pd = dout.data<float>();
    const float* px = x.data<float>();
    float* pg = dgates.data<float>();
    for (int64_t i = 0; i < N; ++i) {
        float acc = 0.0f;
        for (int64_t j = 0; j < C; ++j) acc += pd[i * C + j] * px[i * C + j];
        pg[i * E + e] += acc;
    }
}

// ---- embedding ------------------------------------------------------------------

void embedding(const Tensor& weight, const Tensor& ids, Tensor& out) {
    const int64_t C = weight.size(1);
    const int64_t N = ids.numel();
    const float* pw = weight.data<float>();
    float* po = out.data<float>();
    for (int64_t i = 0; i < N; ++i) {
        const int32_t tok = token_at(ids, i);
        const float* src = pw + int64_t(tok) * C;
        float* dst = po + i * C;
        for (int64_t j = 0; j < C; ++j) dst[j] = src[j];
    }
}

void embedding_backward(const Tensor& ids, const Tensor& dout, Tensor& dweight) {
    const int64_t C = dweight.size(1);
    const int64_t N = ids.numel();
    const float* pd = dout.data<float>();
    float* pw = dweight.data<float>();
    for (int64_t i = 0; i < N; ++i) {
        const int32_t tok = token_at(ids, i);
        float* dst = pw + int64_t(tok) * C;
        const float* src = pd + i * C;
        for (int64_t j = 0; j < C; ++j) dst[j] += src[j];
    }
}

// ---- RoPE ----------------------------------------------------------------------

namespace {
void rope_impl(const float* in, float* out, int64_t B, int64_t T, int64_t H, int64_t hd,
               const float* freqs, int64_t pos_offset, bool inverse, bool accumulate) {
    const int64_t C = H * hd;
    parallel_rows(B * T, [&](int64_t bt) {
        const int64_t t = bt % T;
        const float pos = float(t + pos_offset);
        const float* src = in + bt * C;
        float* dst = out + bt * C;
        for (int64_t h = 0; h < H; ++h) {
            for (int64_t k = 0; k < hd / 2; ++k) {
                const float angle = pos * freqs[k];
                const float c = std::cos(angle);
                const float s = inverse ? -std::sin(angle) : std::sin(angle);
                const int64_t i0 = h * hd + 2 * k;
                const float x0 = src[i0], x1 = src[i0 + 1];
                const float y0 = x0 * c - x1 * s;
                const float y1 = x0 * s + x1 * c;
                if (accumulate) {
                    dst[i0] += y0;
                    dst[i0 + 1] += y1;
                } else {
                    dst[i0] = y0;
                    dst[i0 + 1] = y1;
                }
            }
        }
    });
}
} // namespace

void rope(const Tensor& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset,
          Tensor& out) {
    check(x.ndim() == 3, "rope: expected [B,T,C]");
    const int64_t hd = x.size(2) / n_heads;
    check(hd % 2 == 0, "rope: head_dim must be even");
    check(freqs.numel() == hd / 2, "rope: freqs must have head_dim/2 entries");
    rope_impl(x.data<float>(), out.data<float>(), x.size(0), x.size(1), n_heads, hd,
              freqs.data<float>(), pos_offset, /*inverse=*/false, /*accumulate=*/false);
}

void rope_backward(const Tensor& dout, int64_t n_heads, const Tensor& freqs,
                   int64_t pos_offset, Tensor& dx) {
    const int64_t hd = dout.size(2) / n_heads;
    rope_impl(dout.data<float>(), dx.data<float>(), dout.size(0), dout.size(1), n_heads,
              hd, freqs.data<float>(), pos_offset, /*inverse=*/true, /*accumulate=*/true);
}

// ---- attention -------------------------------------------------------------------

void attention(const Tensor& q, const Tensor& k, const Tensor& v,
               int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,
               Tensor& out, Tensor* probs_out, int64_t window, float attn_softcap) {
    check(q.ndim() == 3 && k.ndim() == 3 && v.ndim() == 3, "attention: expected [B,T,C]");
    const int64_t B = q.size(0), T = q.size(1);
    const int64_t hd = q.size(2) / n_heads;
    const int64_t rep = n_heads / n_kv_heads;
    const int64_t Cq = n_heads * hd, Ckv = n_kv_heads * hd;
    check(probs_out != nullptr, "attention: probs_out required (reference impl)");
    check(probs_out->numel() == B * n_heads * T * T, "attention: probs shape");

    const float* pq = q.data<float>();
    const float* pk = k.data<float>();
    const float* pv = v.data<float>();
    float* po = out.data<float>();
    float* pp = probs_out->data<float>();

    parallel_rows(B * n_heads, [&](int64_t bh) {
        const int64_t b = bh / n_heads;
        const int64_t h = bh % n_heads;
        const int64_t hkv = h / rep;
        float* P = pp + bh * T * T;

        for (int64_t i = 0; i < T; ++i) {
            const float* qi = pq + (b * T + i) * Cq + h * hd;
            const int64_t jmax = causal ? i : T - 1;
            // sliding window: attend only to the last `window` keys (incl. self)
            const int64_t jmin = window > 0 ? std::max<int64_t>(0, i - window + 1) : 0;
            // scores (masked positions never written; treated as prob 0)
            float m = -INFINITY;
            for (int64_t j = jmin; j <= jmax; ++j) {
                const float* kj = pk + (b * T + j) * Ckv + hkv * hd;
                float s = 0.0f;
                for (int64_t d = 0; d < hd; ++d) s += qi[d] * kj[d];
                s *= scale;
                if (attn_softcap > 0.0f)
                    s = attn_softcap * std::tanh(s / attn_softcap);
                P[i * T + j] = s;
                m = std::max(m, s);
            }
            float sum = 0.0f;
            for (int64_t j = jmin; j <= jmax; ++j) {
                const float e = std::exp(P[i * T + j] - m);
                P[i * T + j] = e;
                sum += e;
            }
            const float inv = 1.0f / sum;
            for (int64_t j = 0; j < jmin; ++j) P[i * T + j] = 0.0f;
            for (int64_t j = jmin; j <= jmax; ++j) P[i * T + j] *= inv;
            for (int64_t j = jmax + 1; j < T; ++j) P[i * T + j] = 0.0f;

            float* oi = po + (b * T + i) * Cq + h * hd;
            for (int64_t d = 0; d < hd; ++d) oi[d] = 0.0f;
            for (int64_t j = jmin; j <= jmax; ++j) {
                const float p = P[i * T + j];
                const float* vj = pv + (b * T + j) * Ckv + hkv * hd;
                for (int64_t d = 0; d < hd; ++d) oi[d] += p * vj[d];
            }
        }
    });
}

void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
                        const Tensor& probs, const Tensor& out, const Tensor& dout,
                        int64_t n_heads, int64_t n_kv_heads, float scale,
                        Tensor& dq, Tensor& dk, Tensor& dv, float attn_softcap) {
    const int64_t B = q.size(0), T = q.size(1);
    const int64_t hd = q.size(2) / n_heads;
    const int64_t rep = n_heads / n_kv_heads;
    const int64_t Cq = n_heads * hd, Ckv = n_kv_heads * hd;

    const float* pq = q.data<float>();
    const float* pk = k.data<float>();
    const float* pv = v.data<float>();
    const float* pp = probs.data<float>();
    const float* po = out.data<float>();
    const float* pd = dout.data<float>();
    float* pdq = dq.data<float>();
    float* pdk = dk.data<float>();
    float* pdv = dv.data<float>();

    // Serial over heads: dk/dv rows are shared across q-heads under GQA.
    for (int64_t bh = 0; bh < B * n_heads; ++bh) {
        const int64_t b = bh / n_heads;
        const int64_t h = bh % n_heads;
        const int64_t hkv = h / rep;
        const float* P = pp + bh * T * T;

        for (int64_t i = 0; i < T; ++i) {
            const float* qi = pq + (b * T + i) * Cq + h * hd;
            const float* doi = pd + (b * T + i) * Cq + h * hd;
            float* dqi = pdq + (b * T + i) * Cq + h * hd;

            // dP_ij = dO_i · V_j ; dS = P ∘ (dP − D_i), and the FA2 identity
            // gives D_i = Σ_j dP_ij P_ij = dO_i · O_i in one pass.
            const float* oi = po + (b * T + i) * Cq + h * hd;
            float row_dot = 0.0f;
            for (int64_t d = 0; d < hd; ++d) row_dot += doi[d] * oi[d];
            for (int64_t j = 0; j < T; ++j) {
                const float p = P[i * T + j];
                if (p == 0.0f) continue;
                const float* vj = pv + (b * T + j) * Ckv + hkv * hd;
                const float* kj = pk + (b * T + j) * Ckv + hkv * hd;
                float* dvj = pdv + (b * T + j) * Ckv + hkv * hd;
                float* dkj = pdk + (b * T + j) * Ckv + hkv * hd;

                float dp = 0.0f;
                for (int64_t d = 0; d < hd; ++d) dp += doi[d] * vj[d];
                float ds = p * (dp - row_dot) * scale;
                if (attn_softcap > 0.0f) {
                    // chain through s' = cap·tanh(s/cap): recompute the raw
                    // score, factor is 1 − tanh²
                    float s = 0.0f;
                    for (int64_t d = 0; d < hd; ++d) s += qi[d] * kj[d];
                    const float t = std::tanh(s * scale / attn_softcap);
                    ds *= 1.0f - t * t;
                }

                for (int64_t d = 0; d < hd; ++d) {
                    dvj[d] += p * doi[d];
                    dqi[d] += ds * kj[d];
                    dkj[d] += ds * qi[d];
                }
            }
        }
    }
}

// ---- cross entropy ------------------------------------------------------------

float cross_entropy(const Tensor& logits, const Tensor& targets, Tensor* dlogits) {
    const int64_t V = logits.size(1);
    const int64_t N = logits.size(0);
    const float* pl = logits.data<float>();

    int64_t n_valid = 0;
    for (int64_t i = 0; i < N; ++i)
        if (token_at(targets, i) >= 0) ++n_valid;
    if (n_valid == 0) return 0.0f;

    double total = 0.0;
    const float inv_n = 1.0f / float(n_valid);
    float* pd = dlogits ? dlogits->data<float>() : nullptr;

    for (int64_t i = 0; i < N; ++i) {
        const int32_t tgt = token_at(targets, i);
        if (tgt < 0) continue;
        const float* row = pl + i * V;
        float m = row[0];
        for (int64_t j = 1; j < V; ++j) m = std::max(m, row[j]);
        double sum = 0.0;
        for (int64_t j = 0; j < V; ++j) sum += std::exp(double(row[j] - m));
        const double lse = double(m) + std::log(sum);
        total += lse - double(row[tgt]);

        if (pd) {
            const float inv_sum = float(1.0 / sum);
            float* drow = pd + i * V;
            for (int64_t j = 0; j < V; ++j) {
                const float p = std::exp(row[j] - m) * inv_sum;
                drow[j] += (p - (j == tgt ? 1.0f : 0.0f)) * inv_n;
            }
        }
    }
    return float(total / double(n_valid));
}

} // namespace forge::cpu