5.3. From A Tensor Operation To A GPU Kernel
When the running MLP evaluates
y=W x+b,
the mathematical specification sees matrix-vector multiplication and addition. The CUDA runtime sees something quite different: contiguous buffers, dimensions, an execution stream, a matrix kernel, a broadcast, and several error checks. TorchLean keeps both views and records the contract at the boundary between them.
This chapter follows one training step all the way to the GPU. CUDA is the maintained accelerator today; the capsule and target types are intentionally not CUDA-specific, so a future Metal, ROCm, TPU, or custom-chip provider can state the same kinds of obligations without pretending that its implementation already exists.
5.3.1. Build A Native CUDA Runtime
An ordinary CPU build compiles stub archives for CUDA symbols. The stubs let the package link on a machine without the NVIDIA toolchain, but they reject CUDA session creation. To compile the native implementation:
lake build -R -K cuda=true
Run this command from the repository root. -R rebuilds targets affected by the Lake configuration,
and -K cuda=true selects the CUDA source and link configuration. The build compiles TorchLean's
CUDA code and links the CUDA runtime, cuBLAS, and cuFFT where those libraries are used.
Run two optimizer steps and print the selected kernel contracts:
lake -R -K cuda=true exe torchlean quickstart_mlp \ --device cuda --steps 2 --seed 2026 --show-backend
The model reports the same 25-example dataset as the CPU run. The backend report includes:
matmul: native_cuda.matmul
provider=native-cuda trust=checked vjp=backend-vjp
numeric=[round=nearest-even,
subnormal=implementation-defined,
contract=implementation-defined,
reduce=implementation-defined]
add: native_cuda.add
provider=native-cuda trust=checked vjp=backend-vjp
relu: native_cuda.relu
provider=native-cuda trust=checked vjp=backend-vjp
mse_loss: native_cuda.mse_loss
provider=native-cuda trust=checked vjp=backend-vjp
This output is worth reading closely. The model description contains two linear layers, but the runtime decomposes them into reshape, permutation, matrix multiplication, broadcast, and addition. Capsules describe the operations that actually crossed a backend boundary, not only the layer names in source code.
5.3.2. The Path Of One Linear Layer
For an unbatched input x : [2] and weight W : [8,2], the eager CUDA path proceeds roughly as:
typed Tensor [2]
↓ upload / existing CUDA handle
opaque contiguous float32 buffer
↓ reshape and matrix-layout preparation
native_cuda.matmul
↓ broadcast bias [8] to output shape
native_cuda.add
↓
typed runtime Tensor [8]
The Lean wrapper knows the logical shapes and element counts. The device buffer is opaque; Lean does not inspect its contents by reducing a theorem. Before each FFI call, wrappers check the conditions they can observe:
-
the session really targets CUDA;
-
the selected capsule implements the requested operation;
-
flat lengths match the logical shapes;
-
every shape axis, runtime index, and element count fits the
UInt32CUDA ABI; -
ranks, axes, strides, and operation-specific dimensions are supported.
Shape-erased values, output seeds, and supplied gradients pass through AnyBuffer.validate before
shape-derived indexing. Convolution and pooling also check spatial products independently, so a
zero channel dimension cannot hide an overflowing spatial shape. Their forward wrappers validate
the native output length immediately after the FFI returns, before storing the result on the tape.
Those checks prevent many ABI and memory errors. They do not prove that a CUDA thread computes the correct arithmetic expression.
5.3.3. What A Kernel Capsule Contains
A KernelCapsule is the audit record for one operation-provider pair:
structure KernelCapsule where name : String op : BackendOp provider : Provider device : Device trustLevel : TrustLevel supportsForward : Bool vjpMode : VJPMode shapeContract : ContractDescriptor layoutContract : ContractDescriptor valueContract : ContractDescriptor vjpContract : ContractDescriptor numericalPolicy : NumericalPolicy
The four descriptors answer different questions.
5.3.3.1. Shape contract
For matrix multiplication, are the dimensions compatible and is the declared result shape correct? This field is often supported by Lean-side runtime guards and shape-indexed source objects.
5.3.3.2. Layout contract
Does [m,n] mean contiguous row-major storage with the last axis varying fastest? A correct matrix
formula paired with a transposed native layout still computes the wrong model.
5.3.3.3. Value contract
Which mathematical function should the forward result refine? Evidence may be a Lean theorem, a sound checker, a test suite, a fuzz oracle, or an explicit trusted boundary. The evidence variant is visible in reports.
5.3.3.4. VJP contract
Who owns backward, and which derivative specification should it implement? backend-vjp means the
runtime calls a backend derivative kernel while retaining TorchLean's tape structure.
torchLeanTape means TorchLean owns the graph and reverse traversal even if a capsule uses a named
backend kernel for its local VJP. A capsule marked backend-vjp makes that local numerical boundary
explicit in the audit.
Capsules record contracts. Runtime code supplies a typed KernelHandler for the result type of the
operation. Binding produces an ExecutableKernel only when the handler and capsule have equal
operation, provider, and device fields:
structure KernelHandler (β : Type) where name : String op : BackendOp provider : Provider device : Device execute : KernelCapsule → IO β structure ExecutableKernel (β : Type) where capsule : KernelCapsule handler : KernelHandler β operation_matches : handler.op = capsule.op provider_matches : handler.provider = capsule.provider device_matches : handler.device = capsule.device
The identity check prevents native_cuda.matmul from running the reference CPU closure under that
name. Numerical correctness remains exactly as strong as the capsule's theorem, checker, test, or
trusted-boundary evidence.
5.3.4. Planning, Selection, And Execution
Three stages are intentionally separate:
registry
all known capsule descriptions
↓ profile + availability + assurance gate
accepted plan
capsules allowed for this target
↓ bind matching typed handler
executable kernel
operation/provider/device identities agree
↓ provider-aware runtime dispatch
executed operation
the native symbol actually called
BackendProfile.acceptGraph exposes an AcceptedGraphPlan only after planning, lowering, and the
assurance gate accept every obligation. This is still not a hardware probe: a profile target
declares build capabilities, while CUDA session creation separately calls
Cuda.Buffer.requireNativeRuntime to distinguish native CUDA, a native build with no visible GPU,
and the host-memory parity stubs.
This separation catches two bad failure modes:
-
a registry entry cannot silently claim that an unavailable library is executable;
-
a runtime wrapper cannot silently call native CUDA after the profile selected another provider.
Most fixed CUDA wrappers currently require provider = nativeCuda. If a profile selects an
unwired provider, execution fails with an error. There is no hidden CPU fallback for an unsupported
CUDA operation, because moving a tensor between devices behind the user's back would change both
performance and the execution claim.
5.3.5. Try A Deliberate Failure
Build without native CUDA and request it:
lake build lake exe torchlean quickstart_mlp --device cuda --steps 1
On a stub build, session initialization rejects the request. This is not an accidental limitation:
the CLI's Device.cuda value describes the requested target, while
Cuda.Buffer.requireNativeRuntime probes whether this build can execute it.
Likewise, names such as metal, rocm, tpu, and trainium parse as devices so profiles and
future integrations can describe them. The maintained profile lookup currently returns no runtime
for those devices. Selecting one produces a clear “no maintained runtime profile” error rather than
running on CPU.
5.3.6. Native CUDA Versus cuBLAS
“CUDA” names the NVIDIA programming platform, not one kernel implementation. TorchLean's CUDA path can use:
-
custom
.cukernels for elementwise, reduction, indexing, convolution, attention, and other operations; -
cuBLAS for tuned dense linear algebra;
-
cuFFT for Fourier transforms;
-
CUDA runtime calls for allocation, copies, streams, and launch management.
cuBLAS is a vendor library inside the CUDA ecosystem. Calling it can be much faster than a simple handwritten matrix kernel because it chooses algorithms specialized for dimensions, datatype, and GPU generation. It is also a larger external trust boundary. The capsule should identify the provider and numerical policy precisely enough that “native CUDA” does not hide whether the arithmetic came from custom code or a vendor library.
5.3.7. Maintained CUDA Operations
The eager CUDA tape currently covers elementwise arithmetic and activations, reductions and broadcasting, shape transforms, gather/scatter, dense and batched matrix multiplication, normalization and softmax, two-dimensional and N-dimensional convolution/transposed convolution, max/average/smooth-max pooling, composed attention, a direct attention reference kernel, and fused spectral convolution. Backward rules are recorded as tape nodes rather than delegated to an invisible global autograd engine. Lower-level kernels additionally expose packed real FFT and selective-scan entrypoints used by focused runtime paths and tests.
Layer normalization and tanh-approximate GELU are examples of the distinction between semantics and scheduling. Each is one operation on the TorchLean tape, with a local forward rule and VJP. The CUDA implementation evaluates that rule with one forward kernel and one backward kernel rather than constructing the formula from a chain of temporary tensors. The fusion changes launch and allocation behavior; it does not replace the TorchLean operation or hand backward ownership to an external autograd engine.
Matrix derivatives use the same principle without introducing a fused model operation. Products
such as A^\mathsf{T}B and AB^\mathsf{T} are passed to cuBLAS with logical transpose flags.
The old path first copied an operand into a transposed buffer and then multiplied it. The new path
reads the original row-major storage directly. This applies to ordinary and batched matrix
multiplication, linear layers, projection weights, and any model built from them; attention is only
one caller. TorchLean still records the original matrix operation and its local VJP. Only the
schedule used to evaluate that VJP has changed.
This list does not mean every TorchLean operation has a CUDA implementation. Provider-aware
wrappers reject unsupported capsules and shapes; they do not copy a tensor to CPU and continue
silently. The native source map in NN.Runtime.Autograd.Engine.Cuda.NativeSources identifies the
Lean declarations and corresponding files under csrc/cuda.
5.3.8. Batched Attention
A transformer block receives a tensor of shape (batch, tokens, modelDim). The mathematical
operation applies the same attention layer to every batch entry, with shared projection matrices.
TorchLean keeps that description in the specification and compiled graph. Each sample is expressed
through the existing attention node, so its forward map, JVP, and VJP are the same definitions used
for unbatched attention.
The eager CUDA runtime schedules the work differently. It flattens (batch, tokens) for the four
shared projections and folds (batch, head) into the batch axis of the matrix multiplications.
Hard-masked softmax and the local backward rule still belong to TorchLean. The backward pass sums
the four projection-matrix gradients across the full batch.
This removes the old host loop over samples without introducing a second mathematical operation. The verifier lowers batched attention to the per-sample graph, while CUDA executes one batch-aware tape node. Regression tests compare the batched forward value, input gradient, and shared weight gradients with repeated single-sample attention. The comparison is runtime evidence; the cuBLAS calls and float32 behavior remain covered by the capsule's stated boundary.
5.3.9. Forward Values And Backward Ownership
There are three useful configurations for an operation:
Forward | Backward | TorchLean owns |
|---|---|---|
TorchLean native | TorchLean native VJP | graph, tape, value/VJP rules, native boundary |
external fast kernel | TorchLean VJP | graph, tape, backward semantics; external forward boundary |
external autograd | external autograd | only the surrounding contract and imported gradients |
The middle row is the preferred scaling direction when practical. An external provider computes a fast forward value, but the wrapper still records a TorchLean tape node and applies TorchLean's selected VJP. This keeps parameter ownership and optimizer flow local.
It does not magically prove the external forward. The forward capsule still needs checked, tested, fuzzed, or trusted evidence connecting its value to the spec.
The checked CUDA profile does not use this hybrid LibTorch row by default. Its attention path is
the composed TorchLean operation described above; LibTorch forward is available only through the
explicit libTorchForwardCuda profile.
5.3.10. The LibTorch Attention Example
LibTorch is PyTorch's C++ distribution. ATen is the lower tensor/operator library used inside it. TorchLean's optional adapter calls LibTorch/ATen scaled-dot-product attention; it does not embed a Python interpreter.
Build the optional provider with:
lake -R -K cuda=true -K libtorch=true build lake -K cuda=true -K libtorch=true exe libtorch_sdpa_test
The maintained libTorchForwardCuda profile delegates scaled-dot-product-attention forward to
LibTorch while keeping a TorchLean tape VJP. A raw LibTorch backward test exists to compare
gradients, but the maintained profile does not hand default backward ownership to LibTorch
autograd.
Programmatic selection is explicit:
let run : Trainer.RunConfig :=
({} : Trainer.RunConfig)
|>.withBackendProfile NN.Backend.BackendProfile.libTorchForwardCuda
|>.withBackendReport true
This cannot be obtained by merely spelling --device cuda; ordinary CUDA selects
BackendProfile.checkedCuda.
5.3.11. Hard Attention Masks
TorchLean's boolean attention meaning is:
true = this key participates false = this key has exactly zero softmax numerator
A fully blocked row returns zero. Native fused attention and the LibTorch adapter must preserve this
convention. The adapter constructs a boolean CUDA mask rather than replacing false by -1000.
Why not use a finite sentinel? If an allowed logit is -5000 and a blocked logit receives
-1000,
the blocked entry becomes larger and can dominate softmax. Negative infinity expresses a hard
support restriction; a finite additive bias expresses a different operation.
The attention regression suite includes masked forward values, \mathrm dQ, \mathrm dK,
\mathrm dV, and fully blocked
rows. Those are tests of concrete cases. The pure FlashAttention theorem separately proves that the
spec-level tiled online-softmax definition equals the ordinary attention specification.
5.3.12. Numerical Policy Is Not Decoration
The backend report for matmul says:
round=nearest-even subnormal=implementation-defined contract=implementation-defined reduce=implementation-defined
This is more honest than claiming bitwise equality with a left-fold real specification.
-
Hardware operations usually round to nearest-even, subject to instruction and library choices.
-
Subnormal inputs may be preserved or flushed depending on hardware mode and kernel.
-
multiplication and addition may be contracted into FMA;
-
parallel reductions can choose a tree that differs from the pure interpreter's order.
For floating-point addition,
(a+b)+c
need not equal
a+(b+c).
A mathematically equivalent reduction tree can therefore produce different final bits. Numerical certificates must use an error model or an exact policy strong enough for the claim; shape safety alone is insufficient.
5.3.13. Deterministic Reduction Mode
For reproducibility experiments, the CUDA buffer runtime can replace supported atomic reductions with fixed-order implementations:
let enabled := Runtime.Autograd.Cuda.Buffer.setDeterministicReductionsChecked true
The environment variable TORCHLEAN_CUDA_DETERMINISTIC_REDUCTIONS=1 selects the same mode before
startup. Coverage includes scalar/axis/broadcast reductions, gather/scatter accumulation, and the
pooling backward kernels. It does not make seeded RNG unnecessary, force cuBLAS to a universal
bitwise contract, or prove equality across GPU models. getDeterministicReductions reports the
active setting so a training log can record it.
5.3.14. Bound The Device Reuse Cache
Dropped CUDA buffers are not always returned immediately to the driver. TorchLean keeps exact-size
blocks in a process-wide reuse cache so a later allocation can avoid another cudaMalloc. This is
useful in a steady training loop, but a workload that visits many distinct tensor sizes can retain
more device memory than it will reuse.
Set a byte limit before starting the process:
TORCHLEAN_CUDA_CACHE_CAP_BYTES=$((512 * 1024 * 1024)) \ lake -K cuda=true exe torchlean gpt2 --device cuda --steps 100
The limit is read once by the native allocator. A returned block that would exceed it waits for its
recorded CUDA event and is then freed instead of cached. 0, or an unset variable, keeps the cache
unbounded. The value must be a decimal byte count; malformed and overflowing values produce a
warning and are ignored.
This limit applies only to released blocks in the reuse cache. It does not cap live model parameters, activations, gradients, or workspaces.
Allocator telemetry distinguishes live tensors from reusable memory:
open Runtime.Autograd.Cuda def printCudaMemory : IO Unit := do let stats ← Buffer.allocatorStats IO.println stats.format
liveBytes counts payloads owned by live TorchLean buffers. cacheBytes counts released device
blocks waiting for reuse and is therefore not included in liveBytes. cacheCapBytes reports the
limit installed by the native parser, with 0 meaning unbounded. The formatted report includes all
three values together with wrapper counts and cudaMemGetInfo totals.
The CUDA stress test starts fresh subprocesses because the limit cannot change after the allocator has initialized. It checks a finite cap, an explicit unbounded control, malformed input, and integer overflow against the real allocator. The CPU parity stub reports zero cached bytes and zero cap because it frees host buffers directly.
5.3.15. What The Tests Establish
The maintained CUDA checks are:
lake build -R -K cuda=true scripts/checks/check.sh --cuda
They exercise allocation, uploads and downloads, shapes, operation values, gradients, error paths, and selected numerical behavior. Native sanitizer runs add memory diagnostics. The LibTorch SDPA test compares native and external forward/backward results within declared tolerances.
Passing these checks establishes that those cases worked on the tested build, driver, and GPU. It does not universally prove:
-
the C/CUDA source refines every spec operation;
-
every GPU architecture behaves identically;
-
every reduction is deterministic;
-
the compiler and driver preserve source semantics.
Those components remain named in TRUST_BOUNDARIES.md and in capsule evidence.