CUDA Tape Operations: Attention

Multi-head self-attention

Forward structure matches Spec.MultiHeadAttention.forward:

1. Q = x @ Wq, K = x @ Wk, V = x @ Wv
2. reshape to heads (numHeads, n, headDim)
3. attention per head (batched): softmax(Q Kᵀ / sqrt(headDim)) @ V
4. combine heads, then output projection @ Wo

Masking:

• Every provider uses hard-mask semantics: blocked entries contribute zero softmax numerator.
• The selected capsule determines the forward provider and VJP owner. Native CUDA uses its fused VJP; the TorchLean and LibTorch-forward routes record the same TorchLean local VJP.
• This incurs a host-to-device copy for the mask (since the mask is a host Tensor Bool).

def Runtime.Autograd.Cuda.Tape.Internal.multiHeadAttention {n numHeads dModel headDim : ℕ} (_hSeq : n ≠ 0) (batch : ℕ) (_hBatch : batch ≠ 0) (inputShape outputShape : Spec.Shape) (nodeName : String) (t : Tape) (wqId wkId wvId woId xId : ℕ) (mask : Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) := none) (attentionCapsule : NN.Backend.KernelCapsule := NN.Backend.Attention.torchLeanComposed) : IO (Result (Tape × ℕ))

Shared implementation behind the single-sample and batched attention entrypoints.

def Runtime.Autograd.Cuda.Tape.multiHeadAttention {n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) (t : Tape) (wqId wkId wvId woId xId : ℕ) (mask : Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) := none) (attentionCapsule : NN.Backend.KernelCapsule := NN.Backend.Attention.torchLeanComposed) : IO (Result (Tape × ℕ))

Single-sample multi-head self-attention.

This is a thin shape wrapper around the batch-aware implementation so both paths use the same TorchLean forward and VJP code.

def Runtime.Autograd.Cuda.Tape.batchedMultiHeadAttention {batch n numHeads dModel headDim : ℕ} (hBatch : batch ≠ 0) (h1 : n ≠ 0) (t : Tape) (wqId wkId wvId woId xId : ℕ) (mask : Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) := none) (attentionCapsule : NN.Backend.KernelCapsule := NN.Backend.Attention.torchLeanComposed) : IO (Result (Tape × ℕ))

Batch-aware multi-head self-attention.

The leading batch is folded into the projection row axis and the (batch, head) pair becomes the BMM batch axis. TorchLean still records one typed node and computes the complete local VJP, including accumulation of the shared projection-weight gradients over every sample.