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} [TorchLean.Storage α] (s : EagerSession α) [Context α] [TensorTransfer α] {batch n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) (wq wk wv : TensorRef α [dModel, numHeads * headDim]) (wo : TensorRef α [numHeads * headDim, dModel]) (x : TensorRef α [batch, n, dModel]) (mask : Option (TorchLean.Tensor Bool [n, n]) := none) : IO (TensorRef α [batch, n, dModel])

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} [TorchLean.Storage α] (s : EagerSession α) [Context α] [TensorTransfer α] {batch n numHeads dModel headDim : ℕ} (hBatch : batch ≠ 0) (h1 : n ≠ 0) (wq wk wv : TensorRef α [dModel, numHeads * headDim]) (wo : TensorRef α [numHeads * headDim, dModel]) (x : TensorRef α [batch, n, dModel]) (mask : Option (TorchLean.Tensor Bool [n, n]) := none) : IO (TensorRef α [batch, n, dModel])

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.