rmsNormMatrixGPU function

GPUTensor<Matrix> rmsNormMatrixGPU(GPUTensor<Matrix> input, GPUTensor<Vector>? weight, double epsilon, bool scalePlusOne, CommandBuffer tape,)

///////////////////////////////// Transformer & LLM Ops (2000+) /// ///////////////////////////////// Applies Root Mean Square (RMS) Normalization across the feature dimension of an input matrix. Optionally scales the normalized features using a learnable weight vector. If scalePlusOne is true (as required by Gemma models), the scaling factor applied is (1.0 + weight). The reciprocal standard deviation is cached in rstdCache to accelerate backpropagation.

Implementation

/// Applies Root Mean Square (RMS) Normalization across the feature dimension of an input matrix.
/// Optionally scales the normalized features using a learnable [weight] vector. If [scalePlusOne] is true (as required by Gemma models), the scaling factor applied is `(1.0 + weight)`. The reciprocal standard deviation is cached in [rstdCache] to accelerate backpropagation.
GPUTensor<Matrix> rmsNormMatrixGPU(
    GPUTensor<Matrix> input,
    GPUTensor<Vector>? weight,
    double epsilon,
    bool scalePlusOne,
    CommandBuffer tape) {
  int numRows = input.shape[0];
  int numCols = input.shape[1];

  GPUTensor<Matrix> out = GPUTensor<Matrix>.empty(<int>[numRows, numCols]);
  GPUTensor<Vector> rstdCache = GPUTensor<Vector>.empty(<int>[numRows]);
  out.subMap["rstdCache"] = rstdCache;

  String weightId = '';
  String weightGradId = '';

  if (weight != null) {
    weightId = weight.id;
    weightGradId = weight.gradId;
  }

  tape.putInt(OP_RMS_NORM_FORWARD);
  tape.putString(input.id);
  tape.putString(weightId);
  tape.putString(out.id);
  tape.putString(rstdCache.id);
  tape.putFloat(epsilon);
  tape.putBool(scalePlusOne);

  List<GPUTensor> inputs = <GPUTensor>[input];
  if (weight != null) {
    inputs.add(weight);
  }

  out.creator = GPUNode(
    inputs,
    (CommandBuffer bTape) {
      bTape.putInt(OP_RMS_NORM_BACKWARD);
      bTape.putString(out.gradId);
      bTape.putString(input.id);
      bTape.putString(weightId);
      bTape.putString(rstdCache.id);
      bTape.putString(input.gradId);
      bTape.putString(weightGradId);
      bTape.putBool(scalePlusOne);
    },
    opName: 'rmsNormMatrixGPU',
    cost: numRows * numCols * 4,
  );

  return out;
}