causalMaskGPU function

GPUTensor<Tensor3D> causalMaskGPU(GPUTensor<Tensor3D> input, int batchSize, int numHeads, int seqLen, CommandBuffer tape,)

Applies a causal attention mask to prevent sequences from attending to future tokens. Modifies the [seqLen, seqLen] attention matrices within the 3D tensor by overriding elements above the main diagonal with -10000.0, effectively zeroing out their probability post-softmax. Gradients are blocked from flowing back through the masked positions.

Implementation

GPUTensor<Tensor3D> causalMaskGPU(GPUTensor<Tensor3D> input, int batchSize,
    int numHeads, int seqLen, 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_CAUSAL_MASK_FORWARD);
  tape.putString(input.id);
  tape.putString(out.id);
  tape.putInt(batchSize);
  tape.putInt(numHeads);
  tape.putInt(seqLen);

  out.creator = GPUNode(
    <GPUTensor>[input],
    (CommandBuffer bTape) {
      bTape.putInt(OP_CAUSAL_MASK_BACKWARD);
      bTape.putString(out.gradId);
      bTape.putString(input.gradId);
      bTape.putInt(batchSize);
      bTape.putInt(numHeads);
      bTape.putInt(seqLen);
    },
    opName: 'causalMaskGPU',
    cost: batchSize * numHeads * seqLen * seqLen,
  );

  return out;
}