applyRopeGPU function

GPUTensor<Tensor3D> applyRopeGPU(GPUTensor<Tensor3D> input, GPUTensor<Matrix> cosTable, GPUTensor<Matrix> sinTable, int batchSize, int seqLen, int numHeads, int headDim, CommandBuffer tape,)

Applies Rotary Position Embeddings (RoPE) to a multi-head attention state tensor. Encodes relative positional information by pairing and rotating feature dimensions using precomputed values supplied in cosTable and sinTable. The backward pass computes the inverse rotation gradients by intrinsically transposing the rotation matrix (negating the sine terms).

Implementation

GPUTensor<Tensor3D> applyRopeGPU(
    GPUTensor<Tensor3D> input,
    GPUTensor<Matrix> cosTable,
    GPUTensor<Matrix> sinTable,
    int batchSize,
    int seqLen,
    int numHeads,
    int headDim,
    CommandBuffer tape) {
  int depth = input.shape[0];
  int height = input.shape[1];
  int width = input.shape[2];

  GPUTensor<Tensor3D> out =
      GPUTensor<Tensor3D>.empty(<int>[depth, height, width]);

  tape.putInt(OP_ROPE_FORWARD);
  tape.putString(input.id);
  tape.putString(cosTable.id);
  tape.putString(sinTable.id);
  tape.putString(out.id);
  tape.putInt(batchSize);
  tape.putInt(seqLen);
  tape.putInt(numHeads);
  tape.putInt(headDim);

  out.creator = GPUNode(
    <GPUTensor>[input, cosTable, sinTable],
    (CommandBuffer bTape) {
      bTape.putInt(OP_ROPE_BACKWARD);
      bTape.putString(out.gradId);
      bTape.putString(cosTable.id);
      bTape.putString(sinTable.id);
      bTape.putString(input.gradId);
      bTape.putInt(batchSize);
      bTape.putInt(seqLen);
      bTape.putInt(numHeads);
      bTape.putInt(headDim);
    },
    opName: 'applyRopeGPU',
    cost: batchSize * seqLen * numHeads * headDim,
  );

  return out;
}