Triton-Puzzles-Lite: Triton Kernel Exercises That Run on CPU

What Triton-Puzzles-Lite Solves, and for Whom

Triton is a Python-embedded language for writing GPU kernels, and the fastest way to learn its block-level model is to write small kernels and watch them fail. The original Triton-Puzzles by Sasha Rush and others is a Jupyter notebook that carries visualization code and a heavier dependency set. Triton-Puzzles-Lite is a fork that strips that away. The README states it is "decoupling it from many unnecessary dependencies and making it more accessible for beginner users." The stated dependency is torch alone, because Triton and NumPy arrive with PyTorch. The audience is narrow and clear: someone who has read about tl.load and tl.store and wants to write them, not someone looking for a production kernel library. The exercises are aimed at the point where Triton's programming model stops being obvious, which is the mapping between program IDs, block shapes and memory offsets.

One File, One Interpreter, No GPU Required

The architecture is deliberately flat. All exercises live in puzzles.py, and reference answers live in puzzles_ans.py. There is no package to import, no plugin system, no test harness beyond the script's own argument parser. Execution happens through the Triton interpreter when the environment variable TRITON_INTERPRET is set to 1, which the README describes as running "on CPU through Triton interpreter." That single environment variable is the mechanism that matters. Under it, a Triton kernel written for a GPU executes step by step on the host, which is why a machine with torch-cpu installed is sufficient. The README notes you can also run them on GPU by omitting the variable. The data flow is therefore: you edit a kernel body inside puzzles.py, the script runs it, and a modified test function compares your output against the expected result. The README says the test function was enhanced to print both outputs and "the positions of the different values" when they disagree, which is more useful than a bare assertion failure.

Getting It Running: Commands and Version Pins

The README gives one installation command inside a virtual environment or conda environment: pip install torch==2.5.0. It then notes the matching Triton version as triton==3.1.0, and states that NumPy comes along with PyTorch. That pin is not decoration. The known issues section documents compatibility problems between the Triton interpreter and NumPy 2.0. To run every puzzle and stop at the first failure: TRITON_INTERPRET=1 python3 puzzles.py -a. To run a single puzzle: TRITON_INTERPRET=1 python3 puzzles.py -p 1. To run the reference answers: TRITON_INTERPRET=1 python3 puzzles_ans.py -a. The script exposes further flags through python3 puzzles.py -h, and there is a demo mode, TRITON_INTERPRET=1 python3 puzzles.py -i, which the README recommends as a first check. Puzzle descriptions with pictures are in puzzles.md, so the intended workflow is to read that file alongside the code rather than treat puzzles.py as self-documenting. Debugging inside a kernel uses plain print in interpreter mode; on GPU the README directs you to tl.static_print or tl.device_print instead.

The Debugging Hooks Are the Real Feature

Most exercise repositories stop at a pass or fail message. This one adds two inspection paths that are worth more than the puzzles themselves for a first-time Triton user. The first is a differential output print: when your result differs from the expected one, the harness prints both and the positions where values diverge. The second is memory access reporting. The README states that when invalid memory access is detected, the tool prints access offsets and a valid/invalid mask, and that this is "implemented by hooking the Triton interpreter, so only in CPU mode." That sentence is the most important constraint in the repository. The feature that makes the project attractive for beginners is exactly the feature that disappears the moment you run on real hardware. If your goal is to learn how out-of-bounds loads behave under the interpreter, you are learning the interpreter's model, not the GPU's. The README does not claim otherwise, but it is easy to over-read the screenshots.

Known Failures and the NumPy 2.0 Trap

Two problems are documented, and both should shape how you use the project. First, puzzles 11 and 12 fail in GPU mode. The README lists this without a workaround. So the GPU path is not a superset of the CPU path; it is a partially broken path. If you install a GPU build of torch expecting to do everything the CPU route does, you will hit a wall near the end of the set. Second, the Triton interpreter has compatibility issues with NumPy 2.0. The README gives a concrete diagnostic: run the demo mode and inspect Demo 1. If the output shows the sequence 0 through 7 on the first line and eight zeros on the second, the README says you should fix your version problem first and points to issue 1 for solutions. That is a real failure mode with a real symptom, and it is the reason the torch and triton pins matter more here than in a typical tutorial. A modern environment that resolves to NumPy 2.x can leave you debugging your own setup instead of the puzzles.

How It Differs From the Original Triton-Puzzles

The honest alternative is the upstream project this was forked from, Sasha Rush's Triton-Puzzles. The difference is not in the exercises but in the delivery. Upstream is a Jupyter notebook with visualization code and a larger dependency surface; this fork removes the visualization, converts the notebook into a script, and drops everything except torch. It also carries corrections to the original text. The README lists them: puzzle 6 swaps (i, j) to (j, i) and makes x two-dimensional; puzzle 9's description and notation were rewritten because the original was confusing; puzzle 10 renames the index variable k to l to avoid colliding with the kernel named k; puzzle 12 adds notes about the difference between shift in the formula and in the tests. Those are small fixes, but they are the kind that cost a beginner an hour of confusion. The cost of the fork is the loss of visualizations and the notebook format. If you learn better from an executed cell with rendered output, upstream is the better fit. If you want to run exercises from a terminal, in a script, or in CI, this version is the one that fits.

Maintenance, Licence and What You Are Actually Adopting

There are no releases retrieved for this repository, so there is no versioned artifact to pin beyond the git revision itself. The last push recorded is 2026-03-17. The licence is Apache-2.0, which permits commercial and private use and modification, and requires that you preserve the licence and notice files and state changes you make. That matters here because the project is explicitly a modified derivative of Triton-Puzzles, so if you copy puzzles into your own teaching material, you are carrying both the upstream provenance and the Apache-2.0 obligations. This is not legal advice; check the LICENSE file and the upstream project's terms yourself. The practical maintenance cost is low but not zero: the exercises are tied to triton==3.1.0 and torch==2.5.0, and the interpreter path is the fragile part. Upgrading either dependency without re-running the demo check is how you turn a working exercise set into an afternoon of environment debugging.

Editorial conclusion

Adopt Triton-Puzzles-Lite if you want to learn Triton block semantics on a laptop or CI machine without a GPU, and you can pin torch==2.5.0 with triton==3.1.0. Do not adopt it if you need GPU-only puzzle coverage (puzzles 11 and 12 fail in GPU mode) or a maintained exercise set with releases; there are none retrieved. Before writing any kernel, run TRITON_INTERPRET=1 python3 puzzles.py -i and confirm Demo 1 prints the expected values, then read issue 1 if it does not.

Official sources

1. Issues
2. License: Apache-2.0
3. README
4. SiriusNEO/Triton-Puzzles-Lite on GitHub