Skip to content

Repository files navigation

Litmus

ci Space Dataset License

Try it in your browser (no install, no GPU, nothing sent to a server): huggingface.co/spaces/NagaYu/litmus · Corpus: litmus-kernels dataset

Your kernel test is green. That is not the same as your kernel being correct.

On this repository's corpus of deliberately broken Triton kernels, 88% of the planted bugs pass the test everyone actually runs — fixed power-of-two shapes, torch.rand(), allclose(rtol=1e-2). Litmus catches 100% of them, with 0% false positives on the correct kernels, in 41 seconds for the whole corpus on a laptop with no GPU.

the gap

Litmus does not generate kernels — that space is crowded. It does the other half: compute a sound upper bound on the error a correct kernel may have, then go looking for an input that exceeds it.


Three answers, never two

ACCEPT    a certified bound exists, safety is clean, and the targeted search found nothing
REJECT    a proved safety defect, or a concrete input that breaks the tolerance
UNKNOWN   no bound could be certified — reported as its own outcome, not folded into ACCEPT

That third value is the point. A verifier that says pass when it merely failed to look is worse than no verifier, and inside an RL loop it is a reward-hacking surface: the cheapest way to raise a reward built on a boolean is to write kernels the checker cannot read. litmus check exits 0 / 1 / 2 for ACCEPT / REJECT / UNKNOWN so CI can decide for itself what "unverified" means.


Use it in three lines

pip install -e .
litmus run softmax_nomaxsub          # full pipeline on a corpus kernel
litmus check my_kernel.py --ref my_ref.py --domain "x in [-1e3,1e3]" --dims N=1024
from litmus import ToleranceBound, GPUSafetyCheck, TargetedFalsification

How it works

flowchart TD
    K["Triton kernel<br/>(source)"] --> IR["tritonir<br/>symbolic index model"]
    S["Spec<br/>(what it should compute)"] --> TB
    D["Domain R<br/>x in [-1e3, 1e3]"] --> TB

    IR --> SAFE["<b>GPUSafetyCheck</b><br/>OOB · races · barrier divergence"]
    IR --> SHP["<b>SymbolicShapeSpace</b><br/>boundary classes from the guards"]

    TB["<b>ToleranceBound</b><br/>interval + affine arithmetic<br/>+ explicit reduction-order term"]

    SAFE -->|proof for all shapes| PROVE{"proved<br/>safe?"}
    PROVE -->|no| SOLVE["constraint solver<br/>→ concrete witness"]
    PROVE -->|yes| OK1["clean"]

    TB -->|bounded| THR["threshold B"]
    TB -->|unsupported / divergent| HEUR["heuristic tolerance<br/><i>labelled, not sound</i>"]

    THR --> FALS
    HEUR --> FALS
    SHP --> FALS
    FALS["<b>TargetedFalsification</b><br/>13 aimed generators × shape classes<br/>+ SPSA ascent"]

    FALS -->|violation| MIN["delta-debug<br/>→ minimal counterexample"]
    FALS -->|nothing| OK2["no violation found"]

    SOLVE --> R["REJECT"]
    MIN --> R
    OK1 --> V{"certified<br/>bound?"}
    OK2 --> V
    V -->|yes| A["ACCEPT"]
    V -->|no| U["UNKNOWN"]
Loading

The inversion is the whole design: the bound supplies the threshold, the falsifier supplies the input. Bounding the error of an unknown wrong program is impossible; bounding what a correct one may do is classical numerical analysis, and any measured deviation above it is a proof of incorrectness.

1. ToleranceBound

Interval arithmetic with outward rounding, plus affine forms so that cancellation survives — x - max(x) and (x - mean) * rstd are exactly where plain intervals go vacuous. The budget is reported decomposed:

component what it is
rounding per-operation IEEE-754 rounding, propagated with affine forms
reduction_order the explicit order-nondeterminism term. A parallel reduction may associate its N-1 additions in any order, so deviation from exact is bounded by γ_{N-1}·Σ|x_i| (Higham, ASNA 2nd ed., Thm 4.1), and two orders can differ by twice that. Kept as its own term because it is the part a kernel author cannot remove — and the reason bitwise comparison is the wrong test.
libm ULP budget for transcendentals (Triton lowers to libdevice, which is not correctly rounded)
dtype_cast precision changes, e.g. an fp32 accumulator stored to fp16

Three relational refinements are applied, each a stated lemma rather than a heuristic: R1 the softmax max-shift (x - max(x) ∈ [lo-hi, 0]), R2 the softmax denominator (≥ 1, since the max term contributes exp(0)), and R3 the centred second moment for layernorm, which exploits Σ(x_i - mean) = 0 to get a relative variance bound with an e_m² floor. Without R3 no interval method can bound layernorm at all: it must consider "variance near zero" and "error at its domain maximum" simultaneously, though they cannot co-occur.

2. GPUSafetyCheck

Reads the Triton source with ast and interprets the index computation symbolically while ignoring the data. pid * BLOCK + tl.arange(0, BLOCK) becomes pid*32 + lane0 with 0 ≤ lane0 < 32; offs < N becomes a constraint. Shapes stay symbolic.

Each property is proved first, searched second. A masked tiled matmul is proved in-bounds for all M, N, K by a polynomial argument, and proved race-free by a mixed-radix injectivity argument on the address — the same reason a row-major layout does not alias. Only when a proof fails does the constraint solver look for a witness, and a failed proof is never on its own reported as a defect.

3. TargetedFalsification

Uniform sampling and the region kernels break in barely intersect:

  • torch.rand() is non-negative → a max reduction initialised to 0.0 is never wrong;
  • torch.rand() lives in [0,1) → a softmax missing its max-shift never overflows;
  • benchmark shapes are powers of two → a tail tile never exists;
  • benchmark tensors are contiguous → a row-stride mistake is a no-op;
  • random signs → an fp16 accumulator random-walks instead of saturating.

So the 13 generators are aimed (large magnitude, all-negative, subnormal, NaN/Inf, constant rows, cancellation, ill-conditioned, sorted, saturating, …), crossed with shape classes and dtype classes rather than a size sweep. When those miss, SPSA ascends the error functional — a gradient estimate from two extra evaluations per step; no autograd is claimed, because the kernel is not differentiable and does not live in a tape. Winners are delta-debugged down to a median of 2 elements.

4. SymbolicShapeSpace

A kernel's behaviour depends on a size only through its guards, and pid*BLOCK + lane < N partitions all sizes by N mod BLOCK. So Litmus enumerates the equivalence classes1, BLOCK-1, BLOCK, BLOCK+1, 2·BLOCK, ragged tails, coprime, and a deep-reduction class — and asks the solver for one witness each. O(#dims × 8) configurations, independent of how large shapes may get. dtypes are a second categorical axis, where the fp16-accumulator bug lives.


Benchmark

Three regimes on one corpus of 25 kernels (9 correct, 16 with planted bugs):

  • (A) standard — fixed power-of-two shapes, torch.rand(), allclose(1e-2). What kernel benchmarks and RL rewards run today.
  • (B) fuzzing+ — random (non-power-of-two) shapes, torch.randn(), float64 reference, per-operation tolerances derived from the reduction length. A good-faith strong baseline, not a straw man.
  • (C) Litmus — bound + symbolic safety + targeted falsification.
metric (A) standard (B) fuzzing+ (C) Litmus
bug detection rate 12% 62% 100%
detection, numerically observable bugs only 14% 71% 100%
false positives on correct kernels 0% 0% 0%
wall clock, whole corpus 0.5s 1.0s 40.9s
  • 88% of planted bugs pass condition (A) (14 of 16) — benchmark metric (2), the central claim.
  • Certified bound available for 92% of kernels; median tightness (bound ÷ worst deviation measured on the correct kernel, using the adversarial generators) 1.3×10³.
  • Median counterexample 2 elements.
  • 2 of 16 planted bugs have no numerical signature at all — an out-of-bounds write past the tensor and a divergent barrier. No output comparison detects those at any shape or distribution; they are what separates "tested" from "verified".
figure
fig1_headline_gap.png pass rate vs. truth
fig2_detection_by_bug_class.png per-family detection
fig3_bound_tightness.png every bound above its measured maximum
fig4_verification_time.png CI cost
fig5_counterexample_size.png minimality
python benchmarks/run.py          # regenerate the table and all five figures
python benchmarks/run.py --quick  # CI-sized budgets

The corpus

Six planted bug families, each with the specific reason a standard test lets it through:

family example why (A) misses it
mask_leak tail tile unmasked N=512 is a multiple of BLOCK=32, so no tail exists
reduction_init max seeded with 0.0 not -inf torch.rand() is non-negative, so the identity is never exercised
no_max_subtract softmax without the max shift [0,1) inputs never approach the fp32 range limit
race non-atomic accumulate into one cell on real hardware the lost update is intermittent
dtype_promotion fp16 accumulator stagnates only above ~2048; N=512 never gets there
stride row base row*N on a padded tensor benchmarks pass contiguous tensors, where the bug is a no-op

Each buggy kernel ships with the input that exposes it — the answer key, so anyone can measure a detector's recall without re-deriving ground truth.


It runs in a browser tab

The demo Space is static: Pyodide loads numpy, the Litmus sources are embedded in the page, and the bound engine, Triton parser, constraint solver and CPU kernel surrogate all execute client-side. Nothing is uploaded and no server runs. That is only possible because the whole verifier is pure Python over numpy — which is the same property that makes the GPU-free claim below real rather than a footnote.

python scripts/build_static_space.py     # regenerate space_static/index.html
python app.py                            # or run the Gradio app locally

Two notes for anyone reproducing the deployment: Hugging Face requires a paid tier for Gradio Spaces but hosts static ones for free, and gradio-lite — the obvious route — currently fails to bootstrap because the gradio wheel pins huggingface-hub<1.0 while micropip resolves 1.27. Hence a hand-written page over raw Pyodide. litmus/webdemo.py holds the report logic and is shared with the Gradio app, so the two front ends cannot drift apart.

Everything runs without a GPU

This was a hard requirement, and it shapes the repo. Every kernel exists twice: as real annotated Triton (analysed statically, runs on CUDA when present) and as a numpy transliteration built from block-structured primitives — flat buffers, explicit masks, explicit accumulator dtypes — so a bug written into the Triton is written into the numpy the same way.

Two modelling choices, stated openly:

  • Out-of-bounds reads return a deterministic poison region after the live data, so an unmasked tail reads the next row for interior rows (exactly what a row-major tensor does) and poison past the end.
  • Races are modelled as a deterministic lost update. On real hardware the outcome is intermittent, so this favours the baselines: a numerical test sees the race every run here, where a GPU would show it only sometimes.

triton and z3 are optional. Without them Litmus runs the numpy path and the built-in constraint solver, and says which one it used.


RL reward

from litmus.rewards import LitmusReward
trainer = GRPOTrainer(..., reward_funcs=[LitmusReward()])

Also available as a plain function (litmus_reward) and a verifiers-style rubric function (litmus_reward_func). Anti-hacking properties:

  • silence is not successUNKNOWN scores strictly below ACCEPT, and an unreadable kernel scores 0;
  • memory safety is a gate, not a term — a proved OOB or race zeroes the score regardless of numerical agreement, because a kernel can match its own reference while scribbling on someone else's tensor;
  • a demonstrated counterexample caps the reward at 0.10, so "plausible but wrong" never outscores "simple and right";
  • shape-specialisation gains nothing — reward comes from the whole symbolic shape space, not the benchmark's fixed shape;
  • static mode is capped at 0.65 — without a reference, numerical correctness is unverified and unpaid.

Every score comes with a RewardBreakdown you can log.


Support map and limits

Run litmus support for the authoritative version.

fragment status
elementwise ✅ certified — + - * /, max/min, exp/log/sqrt/rsqrt/tanh/sigmoid/erf over an interval domain
row-wise reduction ✅ certified — one level of sum/mean/max with elementwise pre/post, plus the Norm composite (softmax, layernorm, rms-norm, row sums)
tiled matmul ✅ certified — one contraction of symbolic length K with an elementwise epilogue, including a distinct accumulator dtype
everything else not certified — chained contractions (attention: QK^T → softmax → PV), nested reductions, data-dependent control flow, atomics into shared accumulators, scan/sort

Limitations, stated plainly — this field punishes over-claiming:

  1. Bounds are relative to a declared domain. Outside the stated region nothing is claimed, and the falsifier is clamped to it. Layernorm declares a narrower region ([-10, 10]) than the rest because the R3 lemma requires eps to dominate the squared error of the computed mean; outside that, Litmus falls back to a coarse-but-sound triangle bound rather than refusing.
  2. Without z3, the constraint backend is a bounded search. SAT (a witness) is real; UNSAT only means no counterexample inside the box, and SolveResult.proved_unsat stays False. Litmus never upgrades a bounded UNSAT into a theorem.
  3. Proofs are for all shapes; witness searches are not. Searches run at shrunk block sizes (small-scope), and every witness reports the block size it was found at.
  4. Where no bound exists, falsification runs against a heuristic tolerance. It is not sound, it is tuned to favour precision over recall, and every result carries tolerance_kind="heuristic".
  5. ACCEPT is not a proof of correctness. It means a certified bound exists and this search did not break it. The verdict text says so every time.
  6. The corpus is deliberately bug-heavy. The "36% actually correct" in the headline figure is a property of this corpus, not of kernels in general. The finding is the gap.
  7. γ_n diverges at n·u ≥ 1 and Litmus reports DIVERGENT instead of inventing a number — which is exactly what happens to an fp16 accumulator at n = 2048.
  8. Litmus does not generate kernels. Verification and refutation only.

Layout

litmus/
  affine.py     interval + affine arithmetic, outward rounded
  bounds.py     ToleranceBound — the certified fragments and the three lemmas
  spec.py       the expression DSL and fragment classification
  tritonir.py   Triton source → symbolic memory model
  prove.py      polynomial proofs: in-bounds and injectivity, shapes symbolic
  smt.py        integer constraints; z3 when present, bounded CP otherwise
  safety.py     GPUSafetyCheck
  shapes.py     SymbolicShapeSpace
  falsify.py    TargetedFalsification + delta-debugging minimiser
  simulate.py   the GPU-free execution model
  corpus.py     25 labelled kernels with witnesses
  check.py      the pipeline
  cli.py        litmus check / safety / shapes / run / bench / support
  webdemo.py    demo report logic, shared by the Gradio app and the browser build
  space.py      the Gradio UI
  rewards/      RL reward adapters
benchmarks/run.py    the A/B/C benchmark and all five figures
scripts/build_corpus.py   HF Dataset builder
app.py               Gradio app entry point
space_static/        the generated in-browser (Pyodide) build
tests/               221 tests: soundness, detection, false positives, minimality

Every public function's docstring names the claim it substantiates — detection rate (C1), soundness (C2), low false positives (C3), CI speed (C4), minimality (metric 5), or honesty about coverage (C6).

pytest -q        # 221 passed in ~60s, no GPU required

Links

License

Apache-2.0. See LICENSE.

About

Sound error bounds, symbolic GPU safety checks and targeted falsification for Triton kernels. 88% of planted bugs pass the standard fixed-shape allclose test; Litmus catches 100% with 0% false positives, on CPU.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages