CUDA Float32 Buffers

Low-level buffer operations for the native CUDA autograd runtime. CUDA builds use csrc/cuda/tensor/torchlean_cuda_tensor.cu; ordinary CPU builds link the parity implementation in csrc/cuda/tensor/torchlean_cuda_tensor_stub.c so that the same runtime interfaces remain testable.

Runtime Availability

inductive Runtime.Autograd.Cuda.Buffer.RuntimeStatus : Type

What implementation sits behind the CUDA FFI symbols in the current process.

• cpuStub : RuntimeStatus

  Default non-CUDA builds provide host-memory parity stubs for low-level tests.

• nativeAvailable : RuntimeStatus

  The project was built with CUDA and at least one CUDA device is visible.

• nativeUnavailable : RuntimeStatus

  The project was built with CUDA, but no usable CUDA device is visible.

@[implicit_reducible]

instance Runtime.Autograd.Cuda.Buffer.instDecidableEqRuntimeStatus : DecidableEq RuntimeStatus

@[implicit_reducible]

instance Runtime.Autograd.Cuda.Buffer.instReprRuntimeStatus : Repr RuntimeStatus

def Runtime.Autograd.Cuda.Buffer.instReprRuntimeStatus.repr : RuntimeStatus → Nat → Std.Format

@[extern torchlean_cuda_runtime_status]

opaque Runtime.Autograd.Cuda.Buffer.runtimeStatusRaw (token : UInt32) : UInt32

def Runtime.Autograd.Cuda.Buffer.runtimeStatus (token : UInt32 := 0) : RuntimeStatus

Query whether the linked CUDA symbols are native or the CPU parity stubs.

def Runtime.Autograd.Cuda.Buffer.requireNativeRuntime : IO Unit

Require real CUDA execution for a user-selected CUDA session.

Deterministic Reductions Mode

TorchLean's CUDA runtime uses atomicAdd in a few kernels to accumulate float32 results. This is fast, but floating-point addition is non-associative, and CUDA does not fix a global order for the interleaving of atomic updates. As a result, some kernels can be bit-nondeterministic across runs.

TorchLean therefore exposes an opt-in deterministic mode that replaces those atomic accumulation paths with fixed-order reductions. This trades performance for reproducibility.

This flag is a runtime setting affecting only the CUDA/stub backends; it has no effect on the pure Lean Spec.

@[extern torchlean_cuda_set_deterministic_reductions]

opaque Runtime.Autograd.Cuda.Buffer.setDeterministicReductionsRaw (on : UInt32) : Unit

@[extern torchlean_cuda_get_deterministic_reductions_u]

opaque Runtime.Autograd.Cuda.Buffer.getDeterministicReductionsRaw (u : UInt32) : UInt32

@[extern torchlean_cuda_set_deterministic_reductions_checked]

opaque Runtime.Autograd.Cuda.Buffer.setDeterministicReductionsCheckedRaw (on : UInt32) : UInt32

def Runtime.Autograd.Cuda.Buffer.setDeterministicReductionsChecked (on : Bool) : Bool

Enable/disable deterministic reductions mode and return the observed flag value.

Why this helper exists: the raw setter returns Unit, so if you write let _ := set... in Lean, the compiler is free (under pure semantics) to reorder or eliminate that call. The runtime therefore provides a *_checked wrapper that both sets the flag and returns the observed value, giving us a single call with an explicit return value dependency.

def Runtime.Autograd.Cuda.Buffer.setDeterministicReductions (on : Bool) : Unit

Enable/disable deterministic reductions mode (see module docstring).

def Runtime.Autograd.Cuda.Buffer.getDeterministicReductions : Bool

Query whether deterministic reductions mode is enabled.

Allocator Telemetry

@[extern torchlean_cuda_allocator_live_bytes]

opaque Runtime.Autograd.Cuda.Buffer.allocatorLiveBytesRaw (u : UInt32) : UInt64

@[extern torchlean_cuda_allocator_peak_bytes]

opaque Runtime.Autograd.Cuda.Buffer.allocatorPeakBytesRaw (u : UInt32) : UInt64

@[extern torchlean_cuda_allocator_alloc_count]

opaque Runtime.Autograd.Cuda.Buffer.allocatorAllocCountRaw (u : UInt32) : UInt64

@[extern torchlean_cuda_allocator_free_count]

opaque Runtime.Autograd.Cuda.Buffer.allocatorFreeCountRaw (u : UInt32) : UInt64

@[extern torchlean_cuda_wrapper_live_count]

opaque Runtime.Autograd.Cuda.Buffer.wrapperLiveCountRaw (u : UInt32) : UInt64

@[extern torchlean_cuda_wrapper_peak_count]

opaque Runtime.Autograd.Cuda.Buffer.wrapperPeakCountRaw (u : UInt32) : UInt64

@[extern torchlean_cuda_wrapper_alloc_count]

opaque Runtime.Autograd.Cuda.Buffer.wrapperAllocCountRaw (u : UInt32) : UInt64

@[extern torchlean_cuda_wrapper_finalize_count]

opaque Runtime.Autograd.Cuda.Buffer.wrapperFinalizeCountRaw (u : UInt32) : UInt64

@[extern torchlean_cuda_allocator_device_free_bytes]

opaque Runtime.Autograd.Cuda.Buffer.allocatorDeviceFreeBytesRaw (u : UInt32) : UInt64

@[extern torchlean_cuda_allocator_device_total_bytes]

opaque Runtime.Autograd.Cuda.Buffer.allocatorDeviceTotalBytesRaw (u : UInt32) : UInt64

@[extern torchlean_cuda_allocator_cache_bytes]

opaque Runtime.Autograd.Cuda.Buffer.allocatorCacheBytesRaw (u : UInt32) : UInt64

@[extern torchlean_cuda_allocator_cache_cap_bytes]

opaque Runtime.Autograd.Cuda.Buffer.allocatorCacheCapBytesRaw (u : UInt32) : UInt64

structure Runtime.Autograd.Cuda.Buffer.AllocatorStats : Type

Snapshot of the CUDA buffer allocator.

liveBytes/peakBytes count device or stub payloads allocated by this runtime layer. The wrapper counters track the Lean external objects that own those payloads; in a steady workload, wrapperAllocCount - wrapperFinalizeCount should remain bounded. deviceFreeBytes and deviceTotalBytes come from cudaMemGetInfo in the CUDA build and are 0 in the CPU stub. Together these fields distinguish payload leaks, wrapper-lifetime leaks, and broader CUDA memory pressure or fragmentation.

cacheBytes is the device memory held in the buffer reuse cache: dropped buffers awaiting reuse, which are not counted in liveBytes. cacheCapBytes is the limit selected by TORCHLEAN_CUDA_CACHE_CAP_BYTES; 0 means unbounded. Both fields are 0 in the CPU stub, which keeps no cache.

• liveBytes : UInt64
• peakBytes : UInt64
• allocCount : UInt64
• freeCount : UInt64
• wrapperLiveCount : UInt64
• wrapperPeakCount : UInt64
• wrapperAllocCount : UInt64
• wrapperFinalizeCount : UInt64
• deviceFreeBytes : UInt64
• deviceTotalBytes : UInt64
• cacheBytes : UInt64
• cacheCapBytes : UInt64

def Runtime.Autograd.Cuda.Buffer.instReprAllocatorStats.repr : AllocatorStats → Nat → Std.Format

@[implicit_reducible]

instance Runtime.Autograd.Cuda.Buffer.instReprAllocatorStats : Repr AllocatorStats

def Runtime.Autograd.Cuda.Buffer.allocatorStatsWithToken (token : UInt32) : IO AllocatorStats

Read the current CUDA allocator counters.

token is ignored by the native implementation. It exists so call sites that sample repeatedly can pass a changing value (for example, the training step), preventing Lean from treating repeated FFI reads as identical pure expressions.

def Runtime.Autograd.Cuda.Buffer.allocatorStats : IO AllocatorStats

Read the current CUDA allocator counters. Prefer allocatorStatsWithToken in repeated loops.

def Runtime.Autograd.Cuda.Buffer.mibString (bytes : UInt64) : String

Format a byte count as MiB for allocator progress messages.

def Runtime.Autograd.Cuda.Buffer.AllocatorStats.format (s : AllocatorStats) : String

One-line allocator report for progress logs.

@[extern torchlean_cuda_buffer_of_float_array]

opaque Runtime.Autograd.Cuda.Buffer.ofFloatArray (a : FloatArray) : Buffer

Create a device buffer by copying from a host FloatArray (casts each element to float32).

This primitive has a pure Lean type, but the native implementation allocates a fresh device buffer. Runtime code that repeatedly uploads the same host value should prefer ofFloatArrayIO, which adds an IO token so two uploads cannot be collapsed into the same external object after one is released.

@[extern torchlean_cuda_buffer_of_float_array_with_token]

opaque Runtime.Autograd.Cuda.Buffer.ofFloatArrayWithToken (a : FloatArray) (token : UInt32) : Buffer

def Runtime.Autograd.Cuda.Buffer.ofFloatArrayIO (a : FloatArray) : IO Buffer

Effectful host-to-device upload.

The token is ignored by C/CUDA. Its purpose is semantic: repeated uploads of the same FloatArray must still allocate distinct device buffers. Without a changing token, Lean can treat the extern as a pure expression, which is not the ownership model we want for long eager CUDA training loops.

@[extern torchlean_cuda_buffer_to_float_array]

opaque Runtime.Autograd.Cuda.Buffer.toFloatArray (b : Buffer) : FloatArray

Copy a buffer back to a host FloatArray (casts float32 elements to Float).

@[extern torchlean_cuda_buffer_to_float_array_io]

opaque Runtime.Autograd.Cuda.Buffer.toFloatArrayIO (b : Buffer) : IO FloatArray

@[extern torchlean_cuda_buffer_to_float32_bytes_io]

opaque Runtime.Autograd.Cuda.Buffer.toFloat32BytesIO (b : Buffer) : IO ByteArray

Copy a buffer to its raw float32 byte representation.

This is primarily used by streaming checkpoints. Unlike toFloatArrayIO, it does not widen every element to Lean Float, so a large CUDA parameter can be written without constructing a second double-precision host array.

@[extern torchlean_cuda_buffer_of_float32_bytes_io]

opaque Runtime.Autograd.Cuda.Buffer.ofFloat32BytesIO (bytes : ByteArray) : IO Buffer

Upload a raw float32 byte payload to a fresh buffer.

@[extern torchlean_float_array_to_float32_bytes]

opaque Runtime.Autograd.Cuda.Buffer.floatArrayToFloat32Bytes (values : FloatArray) : ByteArray

Encode a host FloatArray as raw float32 bytes.

@[extern torchlean_float32_bytes_to_float_array]

opaque Runtime.Autograd.Cuda.Buffer.float32BytesToFloatArray (bytes : ByteArray) : FloatArray

Decode raw float32 bytes into a host FloatArray.

@[extern torchlean_cuda_buffer_size]

opaque Runtime.Autograd.Cuda.Buffer.size (b : Buffer) : UInt32

Number of float32 elements in the buffer.

@[extern torchlean_cuda_buffer_size_with_token]

opaque Runtime.Autograd.Cuda.Buffer.sizeWithToken (b : Buffer) (token : UInt32) : UInt32

def Runtime.Autograd.Cuda.Buffer.sizeIO (b : Buffer) : IO UInt32

Read a buffer size at a specific point in an IO ownership sequence.

@[extern torchlean_cuda_buffer_release]

opaque Runtime.Autograd.Cuda.Buffer.release (b : Buffer) : UInt32

Release the device allocation held by a buffer, returning 1 when a live allocation was released.

This is a runtime pressure valve for eager training loops that create many short-lived CUDA buffers. The C finalizer is still safe after an explicit release because the pointer is nulled out.

@[extern torchlean_cuda_buffer_release_with_token]

opaque Runtime.Autograd.Cuda.Buffer.releaseWithToken (b : Buffer) (token : UInt32) : UInt32

def Runtime.Autograd.Cuda.Buffer.releaseIO (b : Buffer) : IO UInt32

Effectfully release a device allocation owned by a completed runtime scope.

The changing token makes the release depend on the surrounding IO sequence. Use this form at ownership boundaries; the pure release primitive is reserved for expressions that thread its result into another native buffer operation. Every alias becomes invalid, so callers must own the complete tape or workspace containing the buffer.

@[extern torchlean_cuda_buffer_release_then]

opaque Runtime.Autograd.Cuda.Buffer.releaseThen (workspace keep : Buffer) : Buffer

Release workspace and return keep.

This exists for pure CUDA tape code: because the returned buffer is used downstream, Lean cannot erase the native release call as dead code.

def Runtime.Autograd.Cuda.Buffer.releaseManyThen (workspace : List Buffer) (keep : Buffer) : Buffer

Release a collection of workspace buffers and return keep.

Many CUDA tape formulas create a group of intermediate buffers, then continue with one final result buffer. Threading cleanup through the result keeps ownership local to the formula and avoids waiting for external-object finalizers in long training loops.

structure Runtime.Autograd.Cuda.Buffer.WithWorkspace : Type

A CUDA result together with workspace buffers that were needed to compute it.

This is the common ownership shape for eager CUDA formulas. Some forward computations need intermediate buffers again during the backward pass, so the tape keeps those buffers on the node and releases them when the node is retired. Backward formulas use the same shape when they recompute a value only to differentiate through it.

• value : Buffer
• workspace : List Buffer

def Runtime.Autograd.Cuda.Buffer.WithWorkspace.releaseWorkspaceThen (r : WithWorkspace) (keep : Buffer) : Buffer

Return keep after releasing all workspace buffers owned by this result.

def Runtime.Autograd.Cuda.Buffer.WithWorkspace.releaseAllThen (r : WithWorkspace) (keep : Buffer) : Buffer

Return keep after releasing both the result buffer and its workspace buffers.

@[extern torchlean_runtime_collect_allocator]

opaque Runtime.Autograd.Cuda.Buffer.collectAllocatorRaw (force : UInt32) : UInt32

Ask the Lean runtime allocator (mimalloc) to collect abandoned/free pages.

This allocator-only operation releases abandoned or free pages without changing TorchLean values. It is intended for long native eager loops that create many short-lived tape closures and external-buffer wrappers at every step.

In the CUDA build, force = true also releases cached device blocks held by the native buffer pool. That gives training code a way to trade reuse for returning memory to the CUDA driver at clear phase boundaries.

def Runtime.Autograd.Cuda.Buffer.collectAllocator (force : Bool := true) : UInt32

Collect the native allocator's free pages.

@[extern torchlean_cuda_buffer_zeros]

opaque Runtime.Autograd.Cuda.Buffer.zeros (n : UInt32) : Buffer

Allocate a length-n buffer filled with zeros.

@[extern torchlean_cuda_buffer_full]

opaque Runtime.Autograd.Cuda.Buffer.full (n : UInt32) (v : Float) : Buffer

Allocate a length-n buffer filled with v (host Float, cast to float32).

@[extern torchlean_cuda_buffer_full_with_token]

opaque Runtime.Autograd.Cuda.Buffer.fullWithToken (n : UInt32) (v : Float) (token : UInt32) : Buffer

def Runtime.Autograd.Cuda.Buffer.fullIO (n : UInt32) (v : Float) : IO Buffer

Allocate a fresh filled buffer inside a repeated runtime loop.

Deterministic RNG (device-side)

These are low-level building blocks used by TorchLean's seeded RNG ops (rand_uniform, bernoulli_mask) when running on the eager CUDA backend.

They use the same SplitMix64-style mixing as TorchLean.Random so results are deterministic given (seed, counter) and a row-major linear index.

@[extern torchlean_cuda_buffer_rand_uniform]

opaque Runtime.Autograd.Cuda.Buffer.randUniform (n : UInt32) (key : UInt64) : Buffer

Deterministic U[0,1) generator: returns a length-n buffer (float32) keyed by key.

@[extern torchlean_cuda_buffer_rand_normal]

opaque Runtime.Autograd.Cuda.Buffer.randNormal (n : UInt32) (mean std : Float) (key : UInt64) : Buffer

Deterministic normal generator using Box-Muller on the device.

@[extern torchlean_cuda_buffer_bernoulli_mask]

opaque Runtime.Autograd.Cuda.Buffer.bernoulliMask (n : UInt32) (keepProb : Float) (key : UInt64) : Buffer

Deterministic {0,1} mask generator: returns a length-n buffer keyed by key.

@[extern torchlean_cuda_buffer_abs]

opaque Runtime.Autograd.Cuda.Buffer.abs (b : Buffer) : Buffer

Absolute value applied pointwise to a CUDA buffer.

@[extern torchlean_cuda_buffer_abs_bwd]

opaque Runtime.Autograd.Cuda.Buffer.absBwd (x dLdy : Buffer) : Buffer

Backward for abs: dx = sign(x) * dLdy (with sign(0)=0).

@[extern torchlean_cuda_buffer_sqrt]

opaque Runtime.Autograd.Cuda.Buffer.sqrt (b : Buffer) : Buffer

@[extern torchlean_cuda_buffer_sqrt_bwd]

opaque Runtime.Autograd.Cuda.Buffer.sqrtBwd (x dLdy : Buffer) : Buffer

Backward for sqrt.

Uses the TorchLean convention: dx = dLdy * (1 / (2*sqrt(x))) for x > 0, else 0.

@[extern torchlean_cuda_buffer_exp]

opaque Runtime.Autograd.Cuda.Buffer.exp (b : Buffer) : Buffer

@[extern torchlean_cuda_buffer_log]

opaque Runtime.Autograd.Cuda.Buffer.log (b : Buffer) : Buffer

@[extern torchlean_cuda_buffer_inv]

opaque Runtime.Autograd.Cuda.Buffer.inv (b : Buffer) : Buffer

Reciprocal: 1/x.

@[extern torchlean_cuda_buffer_clamp]

opaque Runtime.Autograd.Cuda.Buffer.clamp (b : Buffer) (lo hi : Float) : Buffer

Clamp each element to [lo, hi] (bounds are host Floats).

@[extern torchlean_cuda_buffer_clamp_bwd]

opaque Runtime.Autograd.Cuda.Buffer.clampBwd (x dLdy : Buffer) (lo hi : Float) : Buffer

Backward for clamp.

Uses the TorchLean convention: derivative is 1 strictly inside (lo, hi), else 0.

@[extern torchlean_cuda_buffer_max]

opaque Runtime.Autograd.Cuda.Buffer.max (a b : Buffer) : Buffer

Pointwise maximum of two equal-length CUDA buffers.

@[extern torchlean_cuda_buffer_max_bwd]

opaque Runtime.Autograd.Cuda.Buffer.maxBwd (a b dLdy : Buffer) : Buffer × Buffer

Backward for max, returning (dA, dB).

Tie-breaking follows the spec: when a = b, split upstream gradient evenly (0.5) between both.

@[extern torchlean_cuda_buffer_min]

opaque Runtime.Autograd.Cuda.Buffer.min (a b : Buffer) : Buffer

@[extern torchlean_cuda_buffer_min_bwd]

opaque Runtime.Autograd.Cuda.Buffer.minBwd (a b dLdy : Buffer) : Buffer × Buffer

Backward for min, returning (dA, dB).

Tie-breaking follows the spec: when a = b, split upstream gradient evenly (0.5) between both.

@[extern torchlean_cuda_buffer_div]

opaque Runtime.Autograd.Cuda.Buffer.div (a b : Buffer) : Buffer

Pointwise division of two equal-length CUDA buffers.

@[extern torchlean_cuda_buffer_relu]

opaque Runtime.Autograd.Cuda.Buffer.relu (b : Buffer) : Buffer

Pointwise ReLU activation on a CUDA buffer.

@[extern torchlean_cuda_buffer_relu_bwd]

opaque Runtime.Autograd.Cuda.Buffer.reluBwd (x dLdy : Buffer) : Buffer

Backward for relu: dx = dLdy where x > 0, else 0.

@[extern torchlean_cuda_buffer_gelu]

opaque Runtime.Autograd.Cuda.Buffer.gelu (x : Buffer) : Buffer

Tanh-approximate GELU evaluated by one pointwise CUDA kernel.

@[extern torchlean_cuda_buffer_gelu_bwd]

opaque Runtime.Autograd.Cuda.Buffer.geluBwd (x dLdy : Buffer) : Buffer

Backward for tanh-approximate GELU using Activation.geluDerivSpec.

@[extern torchlean_cuda_buffer_add]

opaque Runtime.Autograd.Cuda.Buffer.add (a b : Buffer) : Buffer

Elementwise addition (sizes must match).

@[extern torchlean_cuda_buffer_sub]

opaque Runtime.Autograd.Cuda.Buffer.sub (a b : Buffer) : Buffer

Elementwise subtraction (sizes must match).

@[extern torchlean_cuda_buffer_mul]

opaque Runtime.Autograd.Cuda.Buffer.mul (a b : Buffer) : Buffer

Elementwise multiplication (sizes must match).

@[extern torchlean_cuda_buffer_scale]

opaque Runtime.Autograd.Cuda.Buffer.scale (b : Buffer) (c : Float) : Buffer

Multiply each element by a scalar c (host Float, cast to float32).

This is a primitive building block for many ops (e.g. scaling gradients).

def Runtime.Autograd.Cuda.Buffer.copy (b : Buffer) : Buffer

Device-to-device copy, implemented as a scale-by-one kernel.

@[extern torchlean_cuda_buffer_copy_and_release]

opaque Runtime.Autograd.Cuda.Buffer.copyAndRelease (b : Buffer) : Buffer

Copy a buffer and release the source after the copy has been produced.

The native operation creates the destination before it retires the source, so the compiler cannot reorder the two lifetime events. Use this at ownership-transfer boundaries in the sparse CUDA tape.

@[extern torchlean_cuda_buffer_axpy]

opaque Runtime.Autograd.Cuda.Buffer.axpy (a b : Buffer) (c : Float) : Buffer

Fused multiply-add: a + c * b (sizes must match; c is a host Float, cast to float32).

This is the classic BLAS-style axpy primitive and is useful for optimizers and bias-like updates.

@[extern torchlean_cuda_buffer_adam_step]

opaque Runtime.Autograd.Cuda.Buffer.adamStep (parameters gradient firstMoment secondMoment : Buffer) (beta1 oneMinusBeta1 beta2 oneMinusBeta2 firstMomentCorrection secondMomentCorrection epsilon decay updateScale : Float) : Buffer × Buffer × Buffer

Perform one Adam-family update in a single CUDA pass.

The result is (parameters, firstMoment, secondMoment). Passing decay = 0 gives Adam; passing decay = -(learningRate * weightDecay) gives AdamW's decoupled parameter decay. The caller computes the two bias-correction scales from the step counter, exactly as in Optim.Adam.update and Optim.AdamW.update.

This primitive changes only the execution plan. TorchLean's optimizer definitions remain the semantic reference, while this native boundary avoids materializing every intermediate tensor in the pointwise update.

@[extern torchlean_cuda_buffer_reduce_sum]

opaque Runtime.Autograd.Cuda.Buffer.reduceSum (b : Buffer) : Buffer

Reductions (return a length-1 buffer).

@[extern torchlean_cuda_buffer_reduce_mean]

opaque Runtime.Autograd.Cuda.Buffer.reduceMean (b : Buffer) : Buffer