Batch-Aware Attention For The Eager Trainer

This module supplies the eager implementation used by the generic Ops instance for a leading batch axis. The CPU path fixes the semantics by mapping the existing single-sample operation. The CUDA path folds batch and head axes into one native launch while retaining TorchLean's local VJP.

def Runtime.Autograd.Torch.Internal.EagerSession.batchedMultiHeadAttentionCpuFallback {α : Type} (s : EagerSession α) [Context α] [CudaBridge.TensorConv α] [DecidableEq Spec.Shape] {batch n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) (wq wk wv : TensorRef α (Spec.Shape.dim dModel (Spec.Shape.dim (numHeads * headDim) Spec.Shape.scalar))) (wo : TensorRef α (Spec.Shape.dim (numHeads * headDim) (Spec.Shape.dim dModel Spec.Shape.scalar))) (x : TensorRef α (Spec.Shape.dim batch (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar)))) (mask : Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) := none) : IO (TensorRef α (Spec.Shape.dim batch (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar))))

CPU reference path for batch-aware attention.

It records the established single-sample attention node for each entry of the leading axis. The CUDA path below can execute those samples together, while this definition fixes the exact per-sample forward and backward meaning.

def Runtime.Autograd.Torch.Internal.EagerSession.batchedMultiHeadAttention {α : Type} (s : EagerSession α) [Context α] [CudaBridge.TensorConv α] [DecidableEq Spec.Shape] {batch n numHeads dModel headDim : ℕ} (hBatch : batch ≠ 0) (h1 : n ≠ 0) (wq wk wv : TensorRef α (Spec.Shape.dim dModel (Spec.Shape.dim (numHeads * headDim) Spec.Shape.scalar))) (wo : TensorRef α (Spec.Shape.dim (numHeads * headDim) (Spec.Shape.dim dModel Spec.Shape.scalar))) (x : TensorRef α (Spec.Shape.dim batch (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar)))) (mask : Option (Spec.Tensor Bool (Spec.Shape.dim n (Spec.Shape.dim n Spec.Shape.scalar))) := none) : IO (TensorRef α (Spec.Shape.dim batch (Spec.Shape.dim n (Spec.Shape.dim dModel Spec.Shape.scalar))))

Batch-aware eager attention with a TorchLean-owned local VJP.

The CUDA executor folds (batch, head) into one BMM batch axis. Provider selection remains explicit, and the checked default uses TorchLean's hard-masked softmax and backward rule.