import { a as m } from "../../index-DSGwv2Yx.js";
class h {
  variableNames = ["q", "k"];
  outputShape;
  userCode;
  customUniforms = [
    { name: "divisor", type: "float" },
    { name: "pastLen", type: "int" },
    { name: "inf", type: "float" }
  ];
  constructor(t, e, s, n, a) {
    this.outputShape = [t, e, s, n], this.userCode = `
        void main() {
            ivec4 coords = getOutputCoords(); // [batch, nh, t1, t2]
            int b = coords.x;
            int h = coords.y;
            int t1 = coords.z;
            int t2 = coords.w;

            float sum = 0.0;
            for (int i = 0; i < ${a}; ++i) {
                float qv = getQ(b, h, t1, i);
                float kv = getK(b, h, t2, i);
                sum += qv * kv;
            }

            // Scale by divisor
            float scaled = sum * divisor;

            // Mask out future positions
            setOutput((t2 > t1 + pastLen) ? inf : scaled);
        }
        `;
  }
}
function l(o) {
  const { q: t, k: e } = o.inputs, { divisor: s, pastLen: n } = o.attrs, a = o.backend, i = t.shape[0], r = t.shape[2], c = e.shape[2], u = t.shape[1], p = t.shape[3], d = new h(i, u, r, c, p);
  return a.runWebGLProgram(d, [t, e], "float32", [[s], [n], [Number.NEGATIVE_INFINITY]]);
}
const f = {
  kernelName: "AttentionMask",
  backendName: "webgl",
  kernelFunc: l
};
m(f);
