forward method

Tensor2D forward(Tensor2D x, { KVCache? cache, })

Implementation

Tensor2D forward(Tensor2D x, {KVCache? cache}) {
  final T = x.rows;
  final q = wq.forward(x); // (T, dModel)
  final k = wk.forward(x);
  final v = wv.forward(x);

  // split per head
  final qs = _splitHeads(q);
  final ks = _splitHeads(k);
  final vs = _splitHeads(v);

  // apply RoPE to Q/K per head if configured
  if (rope != null) {
    for (var h = 0; h < nHeads; h++) {
      rope!.applyInPlace(qs[h]);
      rope!.applyInPlace(ks[h]);
    }
  }

  final List<Tensor2D> headsOut = [];

  if (cache != null && T == 1) {
    // decode step: attend to cached + current
    for (var h = 0; h < nHeads; h++) {
      final khCat = _concatRows(cache.getK(h), ks[h]);
      final attn = _scaledDotSoftmax(qs[h], khCat, causal: false);
      final vhCat = _concatRows(cache.getV(h), vs[h]);
      final outH = Ops.matmul(attn, vhCat); // (1, headDim)
      headsOut.add(outH);
    }
    // append current K/V into cache
    cache.append(ks, vs);
  } else {
    // prefill/full attention with causal mask
    for (var h = 0; h < nHeads; h++) {
      final attn = _scaledDotSoftmax(qs[h], ks[h], causal: true);
      final outH = Ops.matmul(attn, vs[h]); // (T, headDim)
      headsOut.add(outH);
    }
    // prime cache if provided
    if (cache != null) {
      cache.append(ks, vs);
    }
  }

  final concat = _catHeads(headsOut); // (T, dModel)
  return wo.forward(concat);
}