CUDA Buffer Kernels FFI #
Foreign-function declarations for ATen operations on TorchLean's dtype-carrying
LibTorch.Buffer: reductions, indexing, matmul/BMM, normalization, Fourier transforms, scans,
and broadcast/view helpers. Lean composes attention from these primitives in Ops.Attention.
The declarations here are the Lean side of the LibTorch CUDA trust boundary documented in
docs/TRUST_BOUNDARIES.md.
Stable softmax along a selected axis of a contiguous tensor, including empty tensors.
The native boundary checks the supplied dimensions against the buffer and delegates to ATen. It does not record an autograd graph. Invalid dimensions, axes and device failures are IO errors.
Mean weighted logits cross-entropy. Original logits remain available for a weighted term whose unweighted log-probability overflows. The foreign routine does not record autograd.
Fixed-teacher mean cross-entropy, or its student gradient when gradient is true.
ATen evaluates the shared logit-difference formula without recording an autograd graph. The
result is scalar for the loss and has the input shape for the gradient.
Grouped convolution over a spatial suffix, with dilation and asymmetric zero padding.
Axes before the channel axis are folded into a batch. kernelDims describes the dense
(outChannels, inChannels, spatial...) layout used by the canonical IR; the bridge selects its
group-diagonal blocks before calling ATen. LibTorch supports one to three spatial dimensions.
Input lengths and all geometry are checked again at the foreign boundary. No gradient is recorded.
Sum across the columns of a 2D row-major buffer.
Input b has shape (rows, cols) and is stored as length rows*cols.
Output is length rows (sum across the columns for each row).
Maximum across the columns of a 2D row-major buffer.
Input b has shape (rows, cols) and is stored as length rows*cols.
Output is length rows (max across the columns for each row).
When cols = 0, the result contains rows zeros by TorchLean's native compatibility convention.
One at the first maximal coordinate in each nonempty row and zero elsewhere. Empty rows return an empty buffer. The mask retains the input dtype.
One for the k entries of each row of a flat (rows, cols) buffer that fewer than k entries of
the row beat, and zero for the others (Spec.MoE.kept). Entry j beats entry i when it is
larger, or when neither is larger and j < i. The kernel compares and counts, so it is exact.
Stable row-wise hard-masked softmax for flat (rows, cols) buffers.
The row maximum and denominator are computed only from entries whose mask value is nonzero. Blocked entries are exactly zero in the output. A row with no allowed entries is defined to be all zeros.
Concatenate two 1D buffers a (length n) and b (length m).
Slice a 1D buffer b (length n) starting at start for len elements.
Requires start + len ≤ n.
Broadcast a column-vector (length rows) to a (rows, cols) matrix.
Output is row-major of length rows*cols, with out[i, j] = vec[i].
Layer normalization over the columns of a row-major (rows, cols) buffer.
gamma and beta each have length cols. The result is (output, normalized, invStd), where
normalized has shape (rows, cols) and invStd has length rows. Keeping these two values is
enough for TorchLean's layer-normalization VJP; the native kernel does not create or own an
autograd graph.
TorchLean's layer-normalization VJP evaluated by LibTorch tensor operations.
Given the upstream derivative, cached normalized values and inverse standard deviations, and
gamma, returns (dX, dGamma, dBeta). The formula and parent association remain part of the
TorchLean tape; this primitive only evaluates that formula.
Batched matrix multiply over row-major buffers, optionally transposing either logical operand.
The logical multiplication always has shape (batch, m, n) × (batch, n, p). When
transposeA = 1, the stored shape of A is (batch, n, m); when transposeB = 1, the stored
shape of B is (batch, p, n). Other flag values are rejected by the native boundary.
LibTorch accepts these views. In particular, backward rules can request Aᵀ B or
A Bᵀ without first allocating a transposed buffer.
Scaled batched product and a range-loss flag for its original operands.
The native path splits a shrinking scale between the operands. Conservative product-range checks, tiny or zero results, and nonfinite intermediates request replay from the original buffers. Exact zero contractions can also request replay; this is not a certified error bound. Reading the combined flag synchronizes the device once. The caller owns the returned buffer.
Batched matrix multiplication A B for ordinary row-major operands.
Instances For
Batched multiplication A Bᵀ.
A is stored as (batch, m, n) and B as (batch, p, n); the result has shape
(batch, m, p).
Instances For
Batched multiplication Aᵀ B.
A is stored as (batch, n, m) and B as (batch, n, p); the result has shape
(batch, m, p).
Instances For
Real-valued 1D FFT over row-major batches, returning a packed half-spectrum.
Input:
x: lengthbatch*n, interpreted as shape(batch, n).
Output:
- length
batch*(n/2+1)*2, interpreted as shape(batch, n/2+1, 2); - the last channel stores
[real, imag]for each nonredundant frequency bin.
CUDA calls LibTorch's real FFT. This is a low-level runtime primitive; differentiable tensor/autograd wrappers should spell out their backward convention separately because half-spectrum packing has normalization and conjugate-symmetry edge cases.
Inverse of rfft1dPacked for packed half-spectra.
Input:
spec: lengthbatch*(n/2+1)*2, interpreted as(batch, n/2+1, 2).
Output:
- length
batch*n, interpreted as(batch, n).
The CUDA implementation calls LibTorch's inverse real FFT with 1/n normalization, matching
the CPU reference.
Unnormalized inverse real FFT with the same packed layout and endpoint handling as irfft1dPacked.
This numerical primitive omits the 1/n factor inside the transform. Lean uses it for the real
FFT adjoint, avoiding both overflow from scaling the spectrum by n beforehand and underflow
from normalizing a tiny result before scaling it back up.
Diagonal selective scan for state-space models.
Inputs:
A,B,h0: lengthstate, representing per-channel recurrence parameters and initial state,X: lengthseqLen*state, row-major token/state inputs.
Output:
- length
seqLen*state, row-major hidden states, withh[t,j] = A[j] * h[t-1,j] + B[j] * X[t,j], starting fromh0[j].
This is the runtime primitive corresponding to the proof layer affine scan contract in
NN.Spec.Layers.SelectiveScan and NN.MLTheory.Proofs.StateSpace.Scan.
The adapter composes ATen operations in a parallel affine scan for finite inputs whose
coefficients have magnitude at most one. Other inputs, or a nonfinite parallel result, use the
sequential recurrence. The parallel schedule can round differently from a left-to-right evaluation.
Backward kernel for selectiveScanDiagFwd.
Given out = selectiveScanDiagFwd A B X h0 and an upstream gradient dY with the same
seqLen*state layout as out, returns (dA, dB, dX, dH0).
Diagonal selective scan with token-dependent coefficients.
Inputs:
A,B,X: lengthseqLen*state, row-major by(time, flattened_state_channel),h0: lengthstate.
Output:
- length
seqLen*state, withh[t,j] = A[t,j] * h[t-1,j] + B[t,j] * X[t,j].
This is the runtime primitive corresponding to full Mamba-style selective scans where the token
controls the affine transition coefficients. It uses the same parallel or sequential scheduling
as selectiveScanDiagFwd.
Reverse accumulation for token-dependent diagonal coefficients.
With g[t] = dY[t] + A[t+1] * g[t+1], the returned arrays are
dA[t] = g[t] * h[t-1], dB[t] = g[t] * X[t], dX[t] = g[t] * B[t], and
dH0 = A[0] * g[0]. Empty sequences return an empty coefficient/input gradient and zero dH0.
The adapter applies the affine scan to the reversed sequence, then forms the pointwise gradients.
Broadcast a buffer to a new shape (TorchLean Shape.CanBroadcastTo semantics).
Arguments:
x: input bufferinDims: input dimension list (outermost-first)outDims: output dimension list (outermost-first)axisMap: lengthoutDims.size;axisMap[j] = 0means the output axisjis an inserted/broadcast axis (input coordinate is0), otherwiseaxisMap[j] = inAxis+1tells which input axis to read.
This shape-driven mapping is generated in Lean from a Shape.CanBroadcastTo proof so the kernel
does not need to interpret the proof object.
Adjoint of broadcastTo for sum-accumulation: reduce a broadcasted gradient back to the input
shape by summing over broadcasted axes.
This uses the same (inDims,outDims,axisMap) convention as broadcastTo.
Swap adjacent axes at depth for a contiguous buffer described by dims.
depth = 0 swaps the first two axes; depth = 1 swaps axes 1 and 2; etc.
Reduce-sum along axis for an N-D contiguous buffer described by dims (outermost-first).
The returned buffer is laid out row-major with shape dims with the axis dimension removed.
Gather k rows from a row-major matrix.
Input:
Output:
- shape
(k, cols)stored row-major as lengthk*cols
Indices that fit in UInt32 but are out of bounds are totalized to 0 rows.
Large Nat values outside the FFI index range are rejected by the runtime.
Scatter-add k rows given host indices.
Semantics: out = mat with out[indices[r], j] += values[r, j] for each r < k, j < cols.
Indices that fit in UInt32 but are out of bounds are ignored; repeated indices accumulate
(scatter-add). Large Nat values outside the FFI index range are rejected by the runtime.