Annotate Triton kernels

An annotation answers one question: which part of which model does this kernel claim to compute, and under which preconditions? It is a decorator placed above @triton.jit (and above @triton.autotune, if any).

import triton
import triton.language as tl

import theformalera as tfe


@tfe.implements("Qwen/Qwen3.8-27B@1.0:layers[*].mlp.act", grid=("T", "cdiv(17408, BLOCK)"),
                bind={"GU": "gate_up", "Y": "y", "T": "T"})
@triton.jit
def silu_mul_kernel(GU, Y, T, BLOCK: tl.constexpr, I: tl.constexpr):
    row = tl.program_id(0).to(tl.int64)   # row * (2 * I) exceeds int32 for long prefills
    blk = tl.program_id(1)
    cols = blk * BLOCK + tl.arange(0, BLOCK)
    m = cols < I
    g = tl.load(GU + row * (2 * I) + cols, mask=m, other=0.0).to(tl.float32)
    u = tl.load(GU + row * (2 * I) + I + cols, mask=m, other=0.0).to(tl.float32)
    tl.store(Y + row * I + cols, g * tl.sigmoid(g) * u, mask=m)

Annotations are no-ops at runtime: implements returns the kernel unchanged. The CLI never imports your modules. It reads annotations from the source, so every argument must be a literal (string, number, list, tuple, dict) or a tfe.layout constructor call.

Parameters

string | string[]

required

A Spec Key, or a list of keys when the same kernel is used in several places. Each key is verified separately.

tuple[str, ...]

required

The launch grid, one expression per grid dimension (at most three). Expressions use shape-domain symbols (T, B, S), constexpr names, integer literals, + - * //, and cdiv(a, b). Examples: ("T",), ("T", "cdiv(17408, BLOCK)"), ("B", "heads").

dict[str, str]

Kernel parameter → spec tensor or scalar name. A value can also be an integer expression for an integer parameter: {"n": "T * 5120"}, or {"stride_bt": "cdiv(S, 16)"}. Entries you leave out are inferred from conventional names (x/X, w/W, y/out, eps, …).

dict[str, Layout]

Pointer parameter → layout mapping. The default is row-major, with the outer strides taken from the kernel’s stride parameters in signature order.

dict[str, tuple[int, int]]

Narrows the shape domain, for example {"T": (1, 8192)}. It can only narrow: domain: T ∈ [1, 300000] widens the default domain [1, 262144] is an error.

list[str]

Extra preconditions, for example ["T % 16 == 0", "contiguous(X)"]. They become hypotheses of the proof, and theformalera guards turns each one into a runtime check.

list[dict]

Constexpr combinations to verify. By default every @triton.autotune config is verified, plus module-level integer constants (BLOCK = 1024) as defaults.

str

The minimum guarantee level for this kernel (proof or bounded). It is recorded in the request; this release’s policy uses [policy] require_level and allow_bounded_on instead.

Binding

The engine has to know which kernel argument is which tensor of the spec node. Run theformalera spec show <key> to see the node’s tensor names:

mlp.act   TheFormalEra.Ops.siluMul  ·  64 layer(s)
├── y[t,i] = silu(gate_up[t,i]) · gate_up[t,17408+i]
├── constants: I=17408, act=silu
├── tensors
│   ├── gate_up  in      [T, 34816]  bf16
│   └── y        out     [T, 17408]  bf16
└── shape domain: T ∈ [1, 262144]

So bind={"GU": "gate_up", "Y": "y"} maps the pointers. "T": "T" binds the integer parameter T to the symbolic dimension. Model constants such as I = 17408 are constexprs, and the engine checks that the values you pass equal the spec’s. If the engine cannot match a spec tensor, it tells you what to add:

cannot bind spec tensor `gate_up` to a kernel parameter (add `bind={'<param>': 'gate_up'}`)

Autotune configs

Every config of @triton.autotune is a different program, and any of them may be chosen at runtime, so each is verified:

@tfe.implements("Qwen/Qwen3.8-27B@1.0:layers[*].mlp.act", grid=("T", "cdiv(17408, BLOCK)"),
                bind={"GU": "gate_up", "Y": "y", "T": "T"})
@triton.autotune(configs=[triton.Config({"BLOCK": 512}), triton.Config({"BLOCK": 1024}),
                          triton.Config({"BLOCK": 2048}), triton.Config({"BLOCK": 4096})], key=["T"])
@triton.jit
def silu_mul_tanh_kernel(GU, Y, T, I: tl.constexpr, BLOCK: tl.constexpr):
    ...

 silu_mul_tanh_kernel     layers[*].mlp.act              4/4   PASS      PROVED        12.7s

This covers exactly the case differential testing misses: the config picked on the test machine works, and the one picked on production hardware does not. To verify a fixed set instead, pass configs=[{"BLOCK": 1024}].

Kernels with several keys

@tfe.implements(
    ["Qwen/Qwen3.8-27B@1.0:layers[*].input_norm",
     "Qwen/Qwen3.8-27B@1.0:layers[*].post_attn_norm"],
    grid=("T",),
    bind={"X": "x", "W": "w", "Y": "y", "eps": "eps"},
)

Each key produces its own row in the results table and its own proof.

Launcher-level annotations

You can annotate a Python wrapper that launches one or more kernels instead of a kernel itself. This is useful for split-KV attention, where two kernels together compute attn.core:

@tfe.implements(
    "Qwen/Qwen3.8-27B@1.0:layers[*].attn.core",
    layout={"k_cache": tfe.layout.paged_kv(block_table="block_table", page=16),
            "v_cache": tfe.layout.paged_kv(block_table="block_table", page=16)},
    domain={"B": (1, 256), "S": (1, 262144)},
)
def paged_decode_attention(q, k_cache, v_cache, block_table, seq_lens, out):
    _decode_partial[(B, KV_HEADS, NUM_SPLITS)](q, k_cache, v_cache, block_table, seq_lens, tmp, ...)
    _decode_reduce[(B, Q_HEADS)](tmp, out, ...)

scan lists launchers, and trace records their launches for L3 sign-off. Static verification of launchers is on the roadmap. Today a launcher reports UNKNOWN, so for now put the annotation on each kernel.

Draft annotations automatically

theformalera annotate kernels/act_new.py

kernels/act_new.py:6 silu_and_mul  (sigmoid)
╭────────────────────────────────────────────────────────────────────────╮
│ @tfe.implements("Qwen/Qwen3.8-27B@1.0:layers[*].mlp.act", grid=("T",), │
│                 bind={'gate_up': 'gate_up', 'out': 'y'})               │
╰────────────────────────────────────────────────────────────────────────╯
A wrong draft can only make verification fail, never pass.

The draft is printed, never written. Check the grid: here the kernel’s second grid dimension (cdiv(17408, BLOCK)) has to be added by hand.

Require annotations

A kernel without an annotation is never verified, so it can never fail. Close that gap with:

theformalera.toml

[policy]
require_annotation = ["kernels/**/*.py"]

Any @triton.jit function under these paths that has no annotation, and is not called by an annotated launcher, fails verify:

 silu_and_mul   -                    0/1   FAILED    F-0001          0ms

Get started

Concepts

Guides

Model conformance

Case studies

Help

Parameters

Binding

Autotune configs

Kernels with several keys

Launcher-level annotations

Draft annotations automatically

Require annotations

​Parameters

​Binding

​Autotune configs

​Kernels with several keys

​Launcher-level annotations

​Draft annotations automatically

​Require annotations

Parameters

Binding

Autotune configs

Kernels with several keys

Launcher-level annotations

Draft annotations automatically

Require annotations