Kernel Capsules🔗

Suppose a graph reaches scaled dot-product attention. Its CUDA capsule describes a Lean composition of matrix products and softmax, with an explicit local VJP over saved probabilities. LibTorch supplies the numerical primitives while TorchLean owns the tape, saved buffers, and parameter gradients. The capsule's contract covers this composition; it does not select a fused attention implementation or record a LibTorch autograd graph.

A KernelCapsule records those differences:

-- These fields describe an implementation contract before
-- any handler is called.
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

Capsules are declared before the run and registered with the backend. The planner may select one only when its device, provider, gradient mode, and trust level fit the requested profile. If no capsule fits, planning stops with an error.

Capsules are collected in named Registry.CapsuleModules. Attention, LibTorch primitives, and portable reference operations contribute modules to the same registry. A downstream provider can prepend another module with BackendProfile.withCapsuleModules; it does not add a new model class or a branch to the graph walker. The model still lowers to ordinary BackendOps, and the planner either finds an admissible capsule for each operation or reports the missing operation.

The registry rejects duplicate module names:

Except.error "duplicate backend capsule module `libtorch`"

Except.ok ()

BackendOp names semantic operation families such as matrix multiplication, reduction, pooling, or convolution. Rank, axes, padding, strides, and index tensors remain in the graph payload. This keeps capability discovery general without erasing the information needed to state the operation correctly. The catalog is one flat list; both maintained profiles share it, and the device filter happens when a plan is made, not when a module is registered.

The sources are Capsule.lean for the record and Registry.lean for the modules. The installation guide has an installation and kernel overview.

Distinct modules can still offer the same operation; the profile decides which one to use.

Capsule Contracts And Evidence🔗

Each capsule has four obligations: shape, layout, forward value, and VJP. Each is a ContractDescriptor: a structured ContractClaim, a one-line summary for readers, and the evidence behind the claim.

The claim is one of shape safety, compatibility with a named tensor layout, refinement of a forward specification, refinement of a VJP specification, or vjpUnavailable for a forward-only capsule. Before planning, TorchLean checks that each descriptor has the right obligation kind and operation, and that a VJP descriptor names the capsule's declared VJP mode.

For example, moving a value test into the portable attention capsule's shape field makes the descriptor disagree with its assigned obligation:

(true, false)

A misaligned capsule is not merely reported; KernelCapsule.admissible requires contractsAligned, so the planner will not select it at all:

Except.error "no admissible kernel capsule for op attention on device cpu"

A value test placed in the shape field is therefore rejected rather than counted as shape evidence, and a forward-only capsule cannot describe .none as though it were a VJP refinement.

ContractEvidence classifies the evidence in four ways:

• runtimeGuard name records validation performed at an execution boundary;

• testSuite name records a regression or differential test suite;

• trustedBoundary reason names code whose correctness is assumed for the claim;

• notApplicable records an obligation that the capsule intentionally does not provide.

There is no "proved" evidence variant: this descriptor does not carry a Lean proof of native implementation correctness. What a capsule can say is which guards and tests stand behind it, or that it delegates to code TorchLean does not check. The related idea of proof-carrying code (Necula, 1997) requires a producer to ship a checkable certificate with its binary. A capsule carries a weaker record: a classified claim and its evidence source. The checker verifies that classification and the caller's policy, without verifying the kernel.

NumericalPolicy currently has one field, reduction : ReductionPolicy, with values fixedLeft, implementationDefined, and notApplicable. The portable matrix-product capsule records the fixed left fold used by the tensor semantics; the LibTorch CUDA capsules record an implementation-defined reduction. A fixed-left range trace therefore cannot be reused for an ATen schedule merely because both capsules implement matmul. The runtime has precision and deterministic-algorithm controls, but the capsule does not audit their readbacks or specify every rounding, subnormal, and multiply-add choice.

Trust has two levels, checked and trustedExternal. The maintained profiles accept checked capsules whose evidence is guards and tests. A capsule whose value contract rests on a trusted boundary is admitted only under the explicit external policy, even if the capsule itself is labelled checked: the contract check reads the evidence, not the label.

Alignment checks that each claim belongs in its assigned field. Whether a named test adequately covers that claim remains a separate question about the test and implementation.

For a matrix product, layout and numerical policy answer different questions. Layout determines which stored elements are interpreted as each row and column. Reduction policy determines the order in which products are added. Two implementations may agree on every address and still produce different last bits through their summation orders. A shape check can catch incompatible matrix dimensions but cannot establish either of those numerical relationships. This is why a capsule carries several obligations instead of one broad correctness label.

Kernel Selection Plans🔗

Planning is a pure function from a profile and a list of operations to a plan or an error message, so a plan can be printed. This helper prints one block per selected capsule:

matmul -> reference.matmul
  provider reference, trust checked
  reduction fixed-left
  value: covered by test suite NN.Tests.Runtime.Floats.Suite
relu -> reference.relu
  provider reference, trust checked
  reduction n/a
  value: covered by test suite NN.Tests.Runtime.Floats.Suite

The portable matrix product commits to a fixed left fold, and relu reports n/a because a pointwise operation performs no reduction at all. Those are the labels a range trace or a benchmark record should quote, and Runtime Approximation explains why a summation order is not a detail.

A missing implementation is an ordinary value, available before any kernel is launched:

no admissible kernel capsule for op selective_scan on device cpu
no admissible kernel capsule for op fft_fno on device cpu

The registered selective_scan and fft_fno capsules are CUDA-only, so the CPU profile cannot plan those operation-family requests. Generic differentiable scan and real-FFT reference paths can instead express their computation through supported ordinary operations. The error identifies a missing capsule for the requested family, not the absence of every CPU implementation of the underlying algorithm.

The reference.matmul and reference.relu rows also explain how to read a successful plan: the name identifies the implementation family, the provider identifies its owner, and the value-evidence line names the retained test source. The planner does not run that suite while printing the report. Its input is the registered evidence record. A benchmark that stores only checked_cpu loses these per-operation choices; storing the selected capsules makes later changes in preference or registry order inspectable.

The Attention Specification Theorem🔗

TorchLean's FlashAttention specification gives a precise target for refinement:

@Spec.flashAttention_eq_scaledDotProductAttention : ∀ {α : Type} [inst : Storage α] [inst_1 : Context α]
  [inst_2 : DecidableRel fun x1 x2 => x1 > x2] (config : Spec.FlashAttentionConfig) {nQ nK dModel : ℕ} {h1 : nQ ≠ 0}
  {h2 : nK ≠ 0} (ctx : Spec.AttentionContext α nQ nK dModel h1 h2),
  Spec.flashAttention config ctx = Spec.scaledDotProductAttention ctx

Both sides are Spec functions. The theorem says that the fused Lean specification has the same denotation as TorchLean's standard scaled-dot-product-attention specification, for every element type with the required operation instances and for every block configuration (Dao et al., 2022). That is what licenses a semantic graph rewrite, and it is what a fused implementation can be asked to refine. The Lean definition reuses the standard attention calculation and ignores tile sizes; this theorem does not verify an online tiled algorithm or its memory traffic. Its axiom audit is:

'Spec.flashAttention_eq_scaledDotProductAttention' depends on axioms: [propext, Classical.choice, Quot.sound]

The three standard Lean axioms, and no sorryAx. The companion theorem flashAttentionBackward_eq_scaledDotProductAttentionBackward has the same axiom set for the backward pass; both live in FlashAttention.lean.

Those theorems compare two Lean specifications. The LibTorch attention capsule records the runtime guards, regression tests, and checked trust level of the actual kernel. Connecting PTX or a library call all the way to the specification would require another refinement argument over Float32, layout, compiler, and hardware behavior, and no capsule field can stand in for it.

Both expressions use the same Context α; the theorem does not compare CPU and CUDA arithmetic. Its attention context also requires nonzero query and key dimensions.

Forward And Backward Ownership🔗

Inference asks for a forward value. Training asks for more: the value must remain connected to the derivative rule used by the optimizer.

TorchLean distinguishes three VJP modes:

• none: a capsule supplies no reverse rule; a profile requests forward execution only;

• torchLeanTape: a capsule composes its local VJP from TorchLean runtime operations;

• backendVJP: a capsule supplies its local VJP through a native bridge routine.

A profile requesting torchLeanTape admits either local VJP implementation while retaining TorchLean's tape and backward traversal. A profile requesting backendVJP specifically requires the native local rule. Attention's capsule uses torchLeanTape: Lean computes its matrix products and explicit softmax VJP over tape-owned probabilities, then adds the input cotangents to the surrounding graph's contributions. There is no LibTorch autograd graph to traverse. Reverse-mode accumulation remains the classical construction (Baydin et al., 2018).

GPU attention uses LibTorch automatically. There is no separate composed-attention provider to select. The registered LibTorch.attention contract records the shape checks, hard-mask convention, local VJP, and implementation-defined reduction order.

No-grad sessions request none automatically. During training, a differentiable operation cannot select a forward-only capsule. Seeded random sources are the deliberate exception: they create non-differentiable values, so they do not need a local VJP of their own.

Keep the selected capsule names with an experiment. A preference alone does not identify what ran, and libtorch.direct_attention still leaves the numerical primitives' implementation choices to the SDK.

Boolean Attention Masks🔗

TorchLean gives boolean attention masks one semantics across specifications and runtimes. A blocked entry contributes exactly zero to the softmax numerator, as if its score were negative infinity. The CUDA composition forms zero numerators for blocked entries and handles fully blocked rows explicitly. Additive score biases remain a separate operation. Modern Models runs that mask against PyTorch and shows the theorem which makes “exactly zero” exact over the reals rather than approximate.

Replacing blocked scores by a large negative sentinel and calling an ordinary softmax changes the function. The original transformer paper writes the mask as -\infty before the softmax (Vaswani et al., 2017), and \exp(-\infty)=0 exactly; a sentinel is a finite number and \exp of it is not zero. TorchLean therefore builds the softmax numerators directly: an allowed entry contributes \exp(\text{score}-\text{rowMax}) and a blocked entry contributes literally 0. In the three-token causal example below, each row is a query and each column is a key. A query may attend to its own position and earlier positions, so entries above the diagonal are blocked:

[[1.000000, 0.000000, 0.000000], [0.622459, 0.377541, 0.000000], [0.333333, 0.333333, 0.333333]]

Every entry above the diagonal prints 0.000000. In the real specification each nonempty row sums to one; the displayed Float entries approximate those normalized weights. The first query has only one allowed key, which receives all its mass. The second row normalizes its first two scores. The third row is uniform because all three keys are allowed and their scores are equal.

In row two, the ratio of the allowed weights is \exp(0.5)/\exp(0), so the first key receives about 0.622459 and the second 0.377541. The third key receives no mass even though its score is a finite number. This is the contract a provider must preserve when it accepts a boolean mask. Multiplying the resulting weights by value vectors then prevents the blocked value from contributing through that attention edge. Causality of a whole model also requires checking its other paths, such as convolutions and any unmasked attention layers.

(0.034425, 0.000000, 0.000000)

[[0.000000, 0.000000, 0.000000], [0.000000, 0.000000, 0.000000], [0.000000, 0.000000, 0.000000]]

[[0.333333, 0.333333, 0.333333], [0.333333, 0.333333, 0.333333], [0.333333, 0.333333, 0.333333]]

Contract Checks🔗

Planning and acceptance are separate steps. Planning finds capsules for graph operations. Auditing turns their contract fields into obligation reports. The evidence check classifies those reports under an AssurancePolicy; execution acceptance also rechecks the full KernelPolicy.

The implementation follows one explicit path:

1. Availability states the devices and providers declared for planning. Eager execution performs the separate linked-library and hardware probes before launching a kernel.

2. Registry supplies the capsule modules, and Planner chooses one PlannedKernel for each operation. Together these choices form a KernelPlan; for a graph, NN.Backend.IR plans one capsule per runtime-relevant node and Grouping coalesces adjacent nodes served by the same capsule.

3. Audit turns the plan into a KernelPlanAudit: one KernelAudit per selected capsule with its four contract descriptors and reduction policy.

4. KernelPlanAudit.checkContracts policy filters the obligations whose evidence the policy does not accept. Its result is a ContractCheck: accepted, or rejected with the failing ObligationReports. The complete acceptance gate also checks policy compatibility. Eager execution receives an AcceptedKernel for each operation; graph planning produces an AcceptedGraphKernelPlan that a later graph executor could consume. The current runtime does not execute from this graph metadata.

5. The eager session binds the selected capsule to a KernelHandler with the same operation, provider, and device through KernelCapsule.bind. If this build has no such handler, execution fails before entering a different provider's code.

6. Report renders providers, trust levels, VJP modes, reduction policies, and per-obligation evidence for logs and benchmark records.

To isolate the evidence check, make a teaching capsule whose forward-value evidence is an assumption. Its other descriptors come from the registered matmul capsule. This is a deliberately constructed diagnostic plan, not a new provider or an executable selection:

checked: rejected demo.assumed_matmul value
  trusted boundary: demonstration assumption
external: accepted

One obligation out of four fails, and the report names which one: the forward value. Shape and layout still have runtime guards, and the VJP still names its regression suite. Accepting external permits the demonstration assumption while retaining those other obligations. The registered LibTorch.matmul uses test-suite evidence for its forward value; the example changed that field specifically to expose the policy boundary.

The eager runtime caches and reports an accepted selection before looking up its handler. It then binds a matching handler and runs it, or fails if no handler is available. The cached selection is therefore not a record of completed calls. Inspection tools can also retain rejected graph plans and explain why they failed.

AcceptedKernel carries a proof that PlannedKernel.acceptable policy = true. This gate checks operation identity, forward support, contract alignment, trust, provider, device, VJP mode, and evidence. AcceptedGraphKernelPlan carries the corresponding result for the groups derived from its stored graph plan. A diagnostic plan alone does not authorize execution: these accepted values require the complete gate to succeed.

AssurancePolicy has one field, allowTrustedExternal. The checked policy leaves it false and admits runtime guards, test suites, and notApplicable; the external policy sets it true and also admits trusted boundaries. The same record decides which capsules the planner may select (by trust level) and which evidence the selected capsules may rely on (by evidence kind). A profile's acceptGraph uses its configured policy for both steps. Diagnostic evidence checks, such as bsCheck, can deliberately compare policies without producing an executable accepted value. Provider preference chooses between admissible implementations. The evidence check is a separate filter: asking for a provider cannot override an obligation's rejected evidence.

The evidence check accepts exactly when the filtered list of rejected obligations is empty. What it establishes is that every selected capsule rests on the evidence classes the caller agreed to. It does not establish the numerical correctness of any kernel, and no registered capsule claims a proof. The sources are Audit.lean and ContractCheck.lean.

The maintained CUDA wrappers also perform concrete checks at the FFI boundary. Convolution and pooling validate rank and dimension conversion, nonzero strides, representable element counts, buffer lengths, and operation-specific domains such as finite nonzero smooth-max \beta after conversion to Float32. The C++ adapter repeats critical size and overflow checks. These guards prevent malformed calls; they complement rather than replace a mathematical value-refinement argument.

Capsules record reduction order and layout claims but have no scalar-type field. The runtime configuration and native conversion checks determine which arithmetic and buffer format actually execute; a capsule label alone does not certify their relationship.

Neither result reports a numerical comparison: this example only checks which evidence classes the policy permits.

Runtime Configuration🔗

The model API stays independent of these implementation details. The same builder is used on every device, and the device is an ordinary argument:

bsTrainer gpu : Trainer [4] [1]

The trainer's type is Trainer [4] [1] on every device: the device is not part of the model's interface, so switching it cannot change which inputs the model accepts. The same holds for a bare runtime configuration:

"cuda"

This block only projects a field from the configuration; it does not validate or execute it. A build without LibTorch rejects this request when the session opens.

Graph Kernel Planning🔗

Everything above planned a list of operations. Graph-level planning additionally records source node identities: some nodes need no kernel, and adjacent selections with the same capsule can share one audit row. Lower the model to the canonical verification IR and inspect its plan. This is a diagnostic graph plan; eager execution selects handlers per operation.

nodes 18, planned 13
  1 NN.IR.OpKind.reshape [4] [1, 4] -> reference.reshape
  3 NN.IR.OpKind.transpose 0 1 -> reference.permute
  4 NN.IR.OpKind.matmul -> reference.matmul
  6 NN.IR.OpKind.broadcastTo [8] [1, 8] -> reference.broadcast
  7 NN.IR.OpKind.add -> reference.add
  8 NN.IR.OpKind.reshape [1, 8] [8] -> reference.reshape
  9 NN.IR.OpKind.relu -> reference.relu
  10 NN.IR.OpKind.reshape [8] [1, 8] -> reference.reshape
  12 NN.IR.OpKind.transpose 0 1 -> reference.permute
  13 NN.IR.OpKind.matmul -> reference.matmul
  15 NN.IR.OpKind.broadcastTo [1] [1, 1] -> reference.broadcast
  16 NN.IR.OpKind.add -> reference.add
  17 NN.IR.OpKind.reshape [1, 1] [1] -> reference.reshape

The plan has thirteen kernels for eighteen nodes. The five missing ones are the input and the four weight and bias constants, because NN.Backend.IR.op? maps input, const, and detach to none. A constant weight tensor is data the executor already holds and needs no operation kernel.

Every kernel here came from the reference.* family, because checkedCpu declares CPU availability and the portable module is the only one offering CPU capsules. The same graph under a CUDA profile selects LibTorch capsules for the same operations, and the audit rows change accordingly while the graph does not.

The node list exposes how a vector linear layer becomes ordinary tensor operations. A reshape introduces a batch axis, a transpose presents the stored weights in matrix-product orientation, and a broadcast aligns the bias with the result. The next reshape removes that batch axis. Those operations are part of the plan even though the user wrote one nn.linear. The four constants retain the model's two weight matrices and two biases. Reading the plan at this level helps locate where a provider gap arises without changing the model's public shape interface.

kernels 13, groups 13
  #[1] reshape
  #[3] permute
  #[4] matmul
  #[6] broadcast
  #[7] add
  #[8] reshape
  #[9] relu
  #[10] reshape
  #[12] permute
  #[13] matmul
  #[15] broadcast
  #[16] add
  #[17] reshape

kernels 3, groups 1
  #[1, 2, 3] relu reference.relu

accepted under checked
13 rows, distinct capsules:
  reference.reshape
  reference.permute
  reference.matmul
  reference.broadcast
  reference.add
  reference.relu

Backend Evidence And Guarantees🔗

These statements have different strengths:

• "the example ran on CUDA" reports an execution path;

• "CUDA matched the CPU reference on this test suite" reports finite parity evidence;

• "the fused attention spec equals standard attention" cites a Lean semantic theorem;

• "the Lean composition over ATen primitives implements the attention spec" requires an implementation refinement argument;

• "this LibTorch build passed the attention tests" reports the tested cases and configuration, while the foreign implementation remains outside the Lean proof.

A backend report records the selected provider and its declared evidence. Keep it beside a benchmark to identify the intended implementations, together with execution results, SDK version and runtime settings. The report alone does not establish that their handlers completed.

For example, a shape guard cannot justify a claim about attention's numerical result. Follow the selected capsule's value obligation to its cited test or trusted boundary, then check what that evidence covers.

Related Chapters🔗

Read Choosing How A Model Runs for the runtime API, and GPU and CUDA for the native implementation details. Runtime Approximation is where reduction order and Float32 rounding are related back to real-valued specifications. The Installation page lists platform commands and the profiles currently wired into the repository.

References🔗

The framework whose tensor library supplies CUDA execution is Paszke et al. (2019); the fused attention algorithm is Dao et al. (2022); reverse-mode accumulation is surveyed by Baydin et al. (2018); and the proof-carrying-code approach used for comparison is Necula (1997).

• PyTorch, C++ and LibTorch API, and the torch.nn.attention backend selector for the SDK permissions exposed by the controls above.

• NVIDIA, CUDA C++ Programming Guide.

• Lean, validating proofs, for what #print axioms is checking.