import { attentionMask as R } from "../ops/attentionMask.js";
import J from "./BaseLayer.js";
import { r as v } from "../rope-CC5RjmKU.js";
import { appendCache as A } from "../ops/appendCache.js";
import { k as c, t as L } from "../index-DSGwv2Yx.js";
import { softmax16 as y } from "../ops/softmax16.js";
import { b as M } from "../matMul16-BIT70Vya.js";
import { p as K } from "../pack16-WlOSOuZA.js";
import { transpose16 as j } from "../ops/transpose16.js";
import { dot16 as E } from "../ops/dot16.js";
import { reshape16 as _ } from "../ops/reshape16.js";
import { isPackedTensor as l } from "../utilities/packed.js";
import { qkv as q } from "../ops/qkv.js";
import { normRMS as O } from "../ops/normRMS.js";
import { dropout16 as x } from "../ops/dropout16.js";
import { v as P } from "../variable-SSATClyt.js";
import { r as S } from "../random_normal-CIm8lk2-.js";
class it extends J {
  constructor(t, o, s, i) {
    super(o, i), this.attentionConfig = s, this.index = t, this.units = o.nEmbed * 3, this.projUnits = o.nEmbed, this.ATTN = `block_${this.index}_cAttn`, this.PROJ = `block_${this.index}_cProj`, this.addVariable(this.ATTN), this.addVariable(this.PROJ), this.divisor = 1 / Math.sqrt(o.nEmbed / o.nHead);
  }
  divisor;
  index;
  units;
  projUnits;
  ATTN;
  PROJ;
  build() {
    this.hasVariable(this.ATTN) === !1 && this.setVariable(
      this.ATTN,
      P(S([this.config.nEmbed, this.units], 0, 0.02), !0, this.ATTN)
    ), this.hasVariable(this.PROJ) === !1 && this.setVariable(
      this.PROJ,
      P(S([this.projUnits, this.config.nEmbed], 0, 0.02), !0, this.PROJ)
    );
  }
  getAttentionScores(t, o, s) {
    const i = R(t, o, this.divisor, s), e = y(i);
    return i.dispose(), e;
  }
  getQKV(t) {
    const o = l(t) ? K(this.getVariable(this.ATTN)) : this.getVariable(this.ATTN), s = q(t, o, this.config.nHead);
    return l(t) && o.dispose(), s;
  }
  getOutputProjection(t) {
    const o = t.shape[0], s = t.shape[2], i = this.config.nEmbed, e = l(t), r = j(t, [0, 2, 1, 3]), n = _(r, [o, s, e ? i / 2 : i]);
    r.dispose();
    const p = e ? K(this.getVariable(this.PROJ)) : this.getVariable(this.PROJ), a = E(n, p);
    return e && p.dispose(), n.dispose(), a;
  }
  updateCache(t, o, s) {
    const i = this.config.blockSize, e = t.shape[2], r = s.length || 0, n = A(t, i, r, s.k);
    t.dispose(), s.k && s.k.dispose();
    const p = A(o, i, r, s.v);
    o.dispose(), s.v && s.v.dispose();
    const a = Math.min(r + e, i), h = s.cumulativeLength + e;
    s.length = a, s.cumulativeLength = h, s.k = c(n), s.v = c(p);
  }
  forward(t, o) {
    return L(() => {
      this.startMemory();
      const [s, i, e] = this.getQKV(o), r = t.pastKV ? t.pastKV.cumulativeLength : 0, n = t.ropeCache, p = n ? v(s, n, r) : s, a = n ? v(i, n, r) : i, h = this.attentionConfig.useQKNorm ?? !1, m = h ? O(p) : p;
      h && p.dispose();
      const f = h ? O(a) : a;
      h && a.dispose(), n && (s.dispose(), i.dispose());
      const T = t.pastKV ? t.pastKV.length : 0;
      t.pastKV && !t.training && this.updateCache(f, e, t.pastKV);
      const u = t.pastKV?.k ? t.pastKV.k : f, V = t.pastKV?.v ? t.pastKV.v : e;
      let d;
      T > 0 ? d = this.getAttentionScores(m, u, T) : d = this.getAttentionScores(m, u), m.dispose(), t.pastKV || u.dispose();
      const g = M(d, V), b = t.attentionScores !== void 0 && t.attentionScores.attentionOut !== void 0;
      b || d.dispose(), t.pastKV || V.dispose();
      const k = this.getOutputProjection(g);
      if (g.dispose(), b && t.attentionScores && t.attentionScores.attentionOut !== void 0) {
        const N = d.shape[1], C = d.shape[2];
        t.attentionScores.attentionOut?.push(
          c(d.slice([0, 0, 0, 0], [1, -1, -1, -1]).reshape([N, C, -1]))
        );
      }
      return this.endMemory("CausalSelfAttention"), t.dropout && t.dropout > 0 ? x(k, t.dropout) : k;
    });
  }
}
export {
  it as default
};
