CUDA Tape Operations: Attention

def Runtime.Autograd.LibTorch.Buffer.scaledProduct (left right : Buffer) (batch m n p : UInt32) (scale : Float) (transposeA transposeB : Bool := false) : Buffer

Keep the original matrices available when scaling a native contraction loses range.

One device flag gates the host transfer. A flagged call downloads its original operands and reuses the reference contraction's finite-coefficient replay, in the buffer's actual precision. The guard is conservative: zero contractions, including exact orthogonality, can also take this CPU path. Nonfinite inputs retain the native result. This is range recovery, not a proof of ATen.

def Runtime.Autograd.LibTorch.Buffer.softmaxProductBackward (weights gradient value : Buffer) (batch rows keys width : UInt32) : Buffer

Apply the centered softmax VJP without discarding the original matrix cotangent.

Contraction range loss or a nonfinite centered result triggers recovery; otherwise the ordinary GPU operations are retained. A flagged finite contraction is reconstructed through the pullback using the saved probabilities, rather than first rounding gradient @ valueᵀ to infinity. Nonfinite original operands or saved probabilities retain the ordinary native buffer.

def Runtime.Autograd.LibTorch.Buffer.softmaxMatmulBackward (weights value multiplier gradient : Buffer) (batch rows keys width : UInt32) : Buffer × Buffer × Buffer

The normalized product's local pullback, including its differentiable multiplier. Recovery uses the actual saved probabilities and retains the original contraction through the multiplier. The existing scalar-product guard also checks the pointwise factors, requiring a device-flag synchronization in addition to the output-finiteness checks. Fully blocked rows suppress their upstream values. Elsewhere nonfinite originals keep native results rather than round-trip NaN payloads.

def Runtime.Autograd.LibTorch.Buffer.attentionForward (Q K V : Buffer) (mask : Option Buffer) (batch n d : UInt32) (scale : Float) : Except String (Buffer × Buffer)

Compose scaled dot-product attention from numerical buffer operations.

Inputs have shape (batch, n, d); an optional (batch, n, n) mask marks allowed entries. Returns the output and saved probabilities. The caller owns both buffers and keeps Q/K/V and probabilities alive until backward. Blocked probabilities and fully blocked rows are zero.

def Runtime.Autograd.LibTorch.Buffer.attentionBackward (Q K V probabilities dOut : Buffer) (batch n d : UInt32) (scale : Float) : Except String (Buffer × Buffer × Buffer)

Attention's local VJP, composed in Lean from the saved probabilities and numerical primitives.

dP = dOut Vᵀ. For each row, choose a maximal-probability coordinate and subtract its dP value: centered = dP - anchor, dS = P * (centered - sum(P * centered)). For normalized exact-real probabilities this is the usual softmax VJP; rounded probabilities need not sum exactly to one, so the centered and uncentered computations can differ. dQ = scale * dS K, dK = scale * dSᵀ Q, and dV = Pᵀ dOut. All arguments are borrowed. Fully blocked rows contribute zero even for nonfinite cotangents.

The contraction check can replay finite dOut and V through the complete centered pullback using the actual saved probabilities. This can recover a finite score cotangent even when dP cannot be represented on its own. Native operations, transfers, and codecs remain trust boundaries.

def Runtime.Autograd.LibTorch.Tape.softmaxMatmul {batch : Spec.Shape} {rows keys width : ℕ} (t : Tape SomeBuffer) (scoresId valueId multiplierId : ℕ) (mask : Option (TorchLean.Tensor Bool [rows, keys]) := none) : Result (Tape SomeBuffer × ℕ)

Record a row-normalized matrix product on the existing native tape. Only saved probabilities are workspace; input buffers remain borrowed tape parents.

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:

• Blocked entries contribute zero softmax numerator, and fully blocked rows return zero.
• Lean composes forward and backward and retains probabilities on its tape.
• LibTorch executes the numerical primitives without recording an autograd graph.
• The host Tensor Bool mask is copied to the device.

def Runtime.Autograd.LibTorch.Tape.Internal.attention {n numHeads dModel headDim : ℕ} (_hSeq : n ≠ 0) (batch : ℕ) (_hBatch : batch ≠ 0) (inputShape : Spec.Shape) (t : Tape SomeBuffer) (wqId wkId wvId woId xId : ℕ) (mask : Option (TorchLean.Tensor Bool [n, n]) := none) : IO (Result (Tape SomeBuffer × ℕ))

Attention on a fixed input shape, selected by the optional-batch public adapter.

def Runtime.Autograd.LibTorch.Tape.attention {n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) (t : Tape SomeBuffer) (wqId wkId wvId woId xId : ℕ) (mask : Option (TorchLean.Tensor Bool [n, n]) := none) (batch : Option ℕ := none) (hBatch : batch.getD 1 ≠ 0 := by decide) : IO (Result (Tape SomeBuffer × ℕ))

Self-attention with shared projection weights and an optional leading batch dimension.

Without batch, the input is [n, dModel]; batch := some b selects [b, n, dModel]. The head count is numHeads, including one for single-head attention.