Kernel Capsules🔗

Suppose a graph reaches scaled dot-product attention. TorchLean currently knows three maintained ways to compute it: a composed TorchLean expression, a native fused CUDA implementation, and a LibTorch forward bridge with a TorchLean-owned backward pass. The operation is the same; the implementation contract is not.

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 CapsuleModules. Built-in attention, native CUDA, portable reference, and optional LibTorch code 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.

The registry rejects the duplicate libtorch name. Distinct modules can still contribute implementations of 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, the evidence behind the claim, and optional provenance (source paths and native symbols). Provenance identifies the implementation being discussed; it is not evidence of correctness.

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 scaled_dot_product_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. A capsule cannot claim that a Lean theorem covers its native implementation, because no such theorem exists for any registered kernel. 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 CUDA and LibTorch capsules record an implementation-defined reduction. A fixed-left range trace therefore cannot be reused for a cuBLAS schedule merely because both capsules implement matmul. Rounding mode, subnormal handling, and multiply-add contraction are not recorded in the capsule, and nothing audits them.

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 returned shape does not reveal which attention
# implementation was selected.
q = torch.zeros(1, 1, 4, 8)
with sdpa_kernel([SDPBackend.FLASH_ATTENTION]):
    out = F.scaled_dot_product_attention(q, q, q)
print(tuple(out.shape))

(1, 1, 4, 8)

# This request changes both the device and dtype and may
# have no eligible kernel.
q64 = torch.zeros(1, 1, 4, 8, device="cuda", dtype=torch.float64)
with sdpa_kernel([SDPBackend.FLASH_ATTENTION]):
    F.scaled_dot_product_attention(q64, q64, q64)

UserWarning: Expected query, key and value to all be of dtype:
  {Half, BFloat16}. Got Query dtype: double, ...
RuntimeError: No available kernel. Aborting execution.

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 computes the standard attention stages and ignores the 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 native CUDA capsule records the runtime guards, regression tests, source provenance, 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.

Read the theorem's Context α as the operations used to give both sides their meaning. The statement quantifies one context shared by the two expressions; it does not compare CPU arithmetic with CUDA arithmetic. The nonzero query and key dimensions are part of the attention context's contract. The axiom output reports logical dependencies of this Lean theorem and says nothing about a compiled CUDA binary. It is therefore useful evidence about the specification rewrite, while the capsule audit remains the place to inspect the implementation boundary.

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: no gradient is requested;

• torchLeanTape: TorchLean owns the tape and backward traversal; each capsule declares whether its local VJP is expressed through TorchLean operations or a named backend kernel;

• backendVJP: require capsules whose local VJP is computed by a backend kernel.

The preferred external-forward design is therefore precise: a provider may compute a fast forward value, TorchLean records the same operation on its tape, and TorchLean applies the backward rule. This requires enough forward information to be retained for that rule. If the bridge cannot provide it, the implementation must fall back or expose a larger trust boundary. Reverse-mode accumulation itself is the classical construction (Baydin et al., 2018); what is being negotiated here is only who owns each local rule.

The maintained LibTorch-forward profile implements this design for scaled-dot-product attention. Selection is per operation. The following profiles differ in provider preference and assurance:

checked_cuda: #[torchlean.composed_attention]
libtorch_forward_cuda: #[libtorch.sdpa_forward]
libtorch_forward_strict: #[native_cuda.direct_attention]

The third profile asks for LibTorch by preference but keeps the checked assurance policy, so the LibTorch capsule is not admissible at all, and chooseCapsuleFor? falls back to the first admissible capsule in catalog order. Preference is a request, and the trust level is a filter that a request cannot override. Fallback follows catalog order. A profile that requires a particular provider must use only rather than a preference.

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.

The three selected capsule names make the interaction between preference and assurance visible. checked_cuda prefers the composed TorchLean route for this request. The external profile admits the LibTorch forward boundary and selects it. Tightening assurance while retaining that preference forces selection to another admissible implementation. A preference therefore cannot be used as a measurement label without checking the result: a run requesting LibTorch may legally select a different provider. If the experiment requires LibTorch specifically, provider restriction and failure are more informative than fallback.

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. Native CUDA skips blocked entries, while the LibTorch bridge passes a boolean mask directly to scaled dot-product attention. 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 is 0.000000 and every row still sums to one. 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.

The LibTorch-forward plan illustrates the evidence check. Holding the selected capsule fixed, check its obligations under both assurance policies:

checked: rejected libtorch.sdpa_forward value
  trusted boundary: LibTorch/CUDA SDPA implementation
external: accepted

One obligation out of four fails, and the report names which one: the forward value. The capsule's shape and layout obligations are runtime guards, its VJP obligation is TorchLean's own rule, and only the forward value relies on the named LibTorch boundary. Accepting external permits that assumption while retaining the other recorded obligations.

These are Lean data structures rather than an informal convention between command-line flags. The eager runtime consumes the accepted per-operation value, binds it to the implementation it will call, and records the capsule it actually used. Inspection tools can 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. The bsPrefer block above shows both halves of that sentence at once: under checked the LibTorch capsule was never selected, and under external it was selected and then accepted.

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/CUDA boundary repeats critical size and overflow checks. These guards prevent malformed launches; 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.

The rejected value obligation does not say that the LibTorch answer was numerically wrong. It says that the chosen policy does not permit that obligation to rest on an external assumption. Likewise, external: accepted does not report a numerical comparison; it records agreement between the evidence class and the caller's policy. This gives an application a precise way to allow an external forward implementation while continuing to require shape guards and a TorchLean-owned reverse path. A single unchecked/checked switch for the whole application would hide which part of the computation actually relies on that assumption.

Runtime Configuration🔗

The model API stays independent of these implementation details. The same builder is used for both execution modes, and the mode is an ordinary argument:

bsTrainer Runtime.Autograd.Torch.ExecutionMode.eager : Trainer [4] [1]

The trainer's type is Trainer [4] [1] either way: the execution mode is not part of the model's interface, so switching it cannot change which inputs the model accepts. Device selection lives in the same configuration record:

"cuda"

This block only projects a field from the configuration; it does not validate or execute it. The current typed graph trainer rejects this CUDA combination when the session opens.

Typed graph execution records and reuses a typed SSA graph. The execution chapter describes this lowering and its proof boundary; backend selection here determines which registered implementation may execute each operation. Graph construction alone is not an optimizing compiler and does not prove the stored derivative rules correct, which is why the two chapters stay separate.

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 native 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 native attention kernel implements the fused spec" requires a native refinement argument;

• "the LibTorch result is correct" depends on the explicitly named LibTorch boundary unless a stronger checker or theorem covers it.

A backend report records the selected provider and the evidence attached to it. Keeping that report beside a benchmark makes “CUDA” concrete: readers can see which operations were native or external and which guards and tests stand behind each one.

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 Execution Modes 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 kernel-selection behavior we compare against 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, torch.compile reference.

• PyTorch, C++ and LibTorch API, and the torch.nn.attention backend selector used in the transcripts above.

• NVIDIA, CUDA C++ Programming Guide.

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