teenygrad kernels / Kernels in a Real Model
Building for Another Target
The machine you develop on and the machine that runs your model are often not the same. A Jetson on a robot has an Arm CPU, a different GPU generation, and no appetite for compiling anything.
Two separate problems, and it helps to keep them apart:
- The CPU code — your program — must be compiled for the board’s architecture.
- The GPU code — your kernels — must be compiled for the board’s compute capability.
cargo-teeny handles both, with a different subcommand each.
Capability, not architecture
Chapter 4 listed the compute capabilities. The one that matters here is sm_87,
Jetson Orin — because it is the case where a developer’s desktop and the target
differ in a way that silently produces a binary that will not run.
PTX gives you some slack. It is an intermediate form, and the driver compiles it
to machine code at load time, so PTX built for an older capability generally
runs on a newer card. What it will not do is use instructions the older
capability did not have. Build for sm_75 and run on an sm_90 and you get a
working kernel that leaves the newer Tensor Cores idle.
So: build for the capability you will run on.
let target = Target::new(Capability::Sm87);
let ptx_path = compile_kernel(&kernel, &target, false)?;
Nothing about this needs an sm_87 device present. Compiling for a capability
and having one are unrelated — which is what makes cross-building possible at
all.
Cross-compiling the program
cargo teeny build --target jetson-orin-nano
cargo teeny check --target jetson-orin-nano # faster feedback
cargo teeny clippy --target jetson-orin-nano
cargo teeny build --target jetson-orin-nano --example yolo26
This delegates to cross, which builds
inside a container holding the target’s toolchain, and handles two things
cross does not do on its own:
- It resolves the teenygrad workspace root from your
[patch.crates-io]entries and mounts it, becausecrossmounts individual crate directories and that is not enough for workspace inheritance. - It mounts the host’s CUDA aarch64 target directory where the board’s image expects it.
cargo teeny check --target jetson-orin-nano is a genuinely useful thing to run
on a laptop. It type-checks everything, including your kernels, for a board you
do not own.
If the board needs libraries you do not have locally, cargo teeny sysroot lays
out an FHS-style tree and can rsync the real thing off the device:
cargo teeny sysroot --host aarch64-unknown-linux-gnu --path ./sysroot \
--type jetson-orin-nano --rsync-from ubuntu@jetson
Compiling kernels ahead of time
Chapter 3 said kernel compilation happens at run time, on the first call, cached on disk. On a development machine that is a one-off cost you never notice.
On a deployed board it is a problem. The first inference pays for compiling every
kernel in the model, teenyc has to be installed on the board, and a read-only
or space-constrained filesystem may not have anywhere to put a cache.
Ahead-of-time compilation moves that to build time:
cargo teeny aot --example yolo26 --device cuda --options "capability=sm_87,ptx-version=82"
The mechanism is worth understanding, because it is unusual. aot builds your
binary for the host — not cross-compiled — and runs it, with flags telling
it to compile its kernels for the named capability and stop. So the program
itself does the compiling; cargo-teeny forwards --device, --options,
--cache-dir and --force verbatim without parsing them.
What lands in the cache directory is PTX for sm_87, produced on your desktop.
Packaging and deploying
cargo teeny package # cross-compiles the binary + AOT-compiles its kernels
cargo teeny deploy # ships the result over SSH
package runs both of the previous steps in one go, forcing the AOT cache
directory to <dest>/cache so the layout is right — cache/ beside bin/.
That layout is not arbitrary. default_cache_dir looks for a cache/ directory
next to the executable and uses it if present, falling back to
TEENYC_CACHE_DIR or /tmp/teenyc_cache. So a packaged binary finds its
precompiled kernels with no environment variables set — and a normal cargo run
during development, whose executable lives under target/debug/ with no cache/
sibling, is unaffected.
What to check before you ship
The capability matches. Building sm_89 PTX for an Orin gets you a runtime
failure, not a compile error.
The PTX version is one the board’s driver accepts. This is the failure that
catches people out. On some Blackwell cards teenyc’s default PTX version is
rejected outright:
PTX .version 8.6 does not support .target sm_120a
The fix is TEENYC_PTX_VERSION=87, or ptx-version= in the AOT options. It is
a teenyc-side default, not something the SDK can work around, and this tree’s
benches carry a note about it.
The cache shipped. A packaged binary with an empty cache/ will try to
compile at run time and fail, because teenyc is not installed on the board.
The dtypes match. A model quantised on the host must be loaded as the dtype it was written as. Chapter 19 covers what silently goes wrong when it is not.
Next: how much of any of this transfers.