ZML: a Zig and MLIR inference stack that compiles one model to CUDA, ROCm, OneAPI, TPU and Trainium

The problem ZML targets: one model definition, many accelerators

Most inference stacks make you choose early. You write against CUDA, or you write against a framework that abstracts the accelerator and pays for it in kernel coverage. ZML's stated goal is to decouple AI workloads from proprietary hardware: the README describes it as a production inference stack with "Any model, many hardwares, one codebase, peak performance", compiling directly to NVIDIA, AMD, INTEL, TPU and Trainium. The audience is therefore narrow and identifiable. It is for engineers who already operate Bazel builds and who have more than one accelerator type in their fleet, or who expect to. It is also for people who want to read the generated program rather than trust a runtime: the repository ships mlir/, kernels/, pjrt/ and ffi/ as top-level directories, so the compiler internals are in the tree rather than behind a wheel. If you only ever run on one NVIDIA GPU and you are happy in Python, the abstraction ZML provides costs you more than it returns. The value appears when the second accelerator shows up and the model code does not have to change.

How the Zig, MLIR and Bazel layers fit together

The stack has three visible layers. Model code is written in Zig against a small tensor API. The README's MNIST sample shows the shape of it: a struct with two Layer fields, each holding a weight and a bias as zml.Tensor, an init that pulls tensors out of a zml.io.TensorStore.View, and a forward that calls .dot, .add and .relu on tensors carrying named dimensions. Those named dimensions are the project's own idea; the docs call them tagged dimensions and give them a dedicated tutorial, working_with_tensors.md. Loading is a separate step from construction. The sample's load function calls zml.io.load with explicit options: shardings, parallelism, dma_chunks and dma_chunk_size (16 * 1024 * 1024 in the example). That is where weights move from a store onto a platform, and where multi-device placement is decided. Below that sits MLIR: the tensor graph is lowered through MLIR and then to the accelerator backends, which is why the same forward function can be built for CUDA, ROCm, OneAPI, TPU or Neuron. Bazel sits outside all of it as the build system that fetches dependencies, selects the platform, and compiles the MLIR toolchain. The practical consequence is that platform selection happens at build time, not at runtime. You do not pass a device string to a binary; you rebuild with a flag.

Installing ZML and running the 30-second MNIST smoke test

The README is explicit that Bazel is the only prerequisite, and it recommends installing it through Bazelisk. On macOS that is a single Homebrew command; on Linux the README downloads the Bazelisk binary directly to /usr/local/bin/bazel and marks it executable.

brew install bazelisk

The Linux path fetches a pinned Bazelisk release and makes it runnable:

curl -L -o /usr/local/bin/bazel 'https://github.com/bazelbuild/bazelisk/releases/download/v1.28.0/bazelisk-linux-amd64'
chmod +x /usr/local/bin/bazel

With Bazel present, the smoke test is one target. According to the README it downloads a small pretrained MNIST model, compiles it, loads the weights, and classifies a random handwritten digit:

bazel run //examples/mnist

If that completes, the toolchain, the MLIR lowering and the CPU backend all work on your machine. That is the point of the example: it is the smallest end-to-end model run in the repository, and it exercises the same load path the LLM example uses.

Running an LLM from Hugging Face, S3 or a local directory

The LLM example is the one most people will actually run. The README lists its supported architectures as Llama 3.1 / 3.2, Qwen 3.5 and LFM 2.5. Gated repositories such as Meta Llama need authentication, either through the project's own tool or an environment variable.

bazel run //tools/hf -- auth login

Alternatively the README says you can set the HF_TOKEN environment variable. A single prompt run looks like this, with the model referenced through the hf:// scheme:

bazel run //examples/llm -- --model=hf://meta-llama/Llama-3.2-1B-Instruct --prompt="What is the capital of France?"

Omitting --prompt opens an interactive chat loop instead. The model argument accepts three forms: hf:// for Hugging Face, a plain local directory such as /var/models/meta-llama/Llama-3.2-1B-Instruct, and s3://bucket/path/to/model. The examples/io target exists to inspect those same sources, since the README describes it as loading local, hf://, https:// and s3:// repositories through the VFS layer. Note the double dash before the example's own flags: everything after it belongs to the example binary, not to Bazel.

Selecting an accelerator: build-time platform flags

There is no runtime device switch. You append one or more platform flags to the Bazel invocation, and the README lists them as --@zml//platforms:cuda=true, --@zml//platforms:rocm=true, --@zml//platforms:oneapi=true, --@zml//platforms:tpu=true and --@zml//platforms:neuron=true for AWS Trainium and Inferentia 2. CPU compilation can be turned off with --@zml//platforms:cpu=false, which the Intel example uses alongside the OneAPI flag.

bazel run //examples/llm --@zml//platforms:cuda=true -- --model=hf://meta-llama/Llama-3.2-1B-Instruct --prompt="Write a haiku about Zig"

This design has a clear cost. Because the flag is a build setting, switching accelerators means a different compiled binary, and the flag has to be repeated on every invocation or captured in .bazelrc. The repository does ship a .bazelrc at the top level, which is the natural place to pin a default platform for a given machine. It also means a CI matrix has to build per platform rather than test one artifact everywhere. In exchange, the compiler can specialize the generated MLIR for the target instead of dispatching at runtime. The kernels/ and pjrt/ directories in the tree are where that specialization lives. Whether the trade is worth it depends on how often your target changes: if it never does, the build-time coupling is pure overhead.

Where ZML is the wrong tool

The largest limitation is the language boundary. Model code is Zig. The README's own walkthrough is a Zig struct with init, load, unloadBuffers and forward methods, and the tutorials it points to (write_first_model.md, working_with_tensors.md) teach that API. There is no Python entry point described anywhere in the README. A team whose researchers publish PyTorch checkpoints and expect to tweak a forward pass in a notebook will not be able to use this without a translation step, and the README does not describe one. Second, the supported model list for the LLM example is short and explicit: Llama 3.1 / 3.2, Qwen 3.5, LFM 2.5. An architecture outside that list is not covered by the example, and the README does not claim general architecture coverage. Third, the prerequisite is Bazel, which is a real organizational cost: a repository that builds with Bazel wants the rest of your build to be Bazel-aware too, and the top-level MODULE.bazel, MODULE.bazel.lock and .bazelversion files show this is a fully Bazel-native project rather than one with a Bazel wrapper. If your team has no Bazel experience, the first day is spent on Bazel, not on inference. Finally, the README does not document rollback, version pinning across releases, or a compatibility matrix between ZML versions and model checkpoints, so upgrade planning has to come from the release history and the repository itself.

How ZML differs from ONNX Runtime and vendor runtimes

The closest familiar alternative is ONNX Runtime. The difference is where the model boundary sits. ONNX Runtime takes a serialized graph, produced by an exporter, and executes it with per-vendor execution providers selected at runtime. ZML takes source-level model code in Zig and compiles it, with the target chosen at build time. That means ZML can specialize aggressively for the accelerator because the target is known when the compiler runs, but it also means you cannot ship one binary and let it discover the hardware. ONNX Runtime also has a Python API and a large set of pre-exported models; ZML's README points at a handful of example directories instead. A second comparison is the vendor runtimes themselves, TensorRT or the ROCm stack. Those are deeper on their own hardware and give you no portability at all. ZML's bet is that portability across NVIDIA, AMD, INTEL, TPU and Trainium is worth more than depth on any one of them. That bet only pays off if you genuinely run on more than one accelerator, which is a smaller group than the marketing phrasing suggests. For a single-vendor deployment, a vendor runtime is the more direct route and skips Bazel entirely.

Licence, releases and what an upgrade actually costs

ZML is licensed under Apache-2.0, per the LICENSE file at the repository root and the README's licence section. That is a permissive licence with an explicit patent grant, which matters for a compiler that lowers to vendor backends; it does not carry the source-disclosure obligations of a copyleft licence. This is a factual note about the identifier, not legal advice, and anyone embedding ZML in a shipped product should read the LICENSE text and their own counsel's view rather than this paragraph. On releases, the recent release list shows zml-smi-v0.3.0 published on 2026-04-22. That name is worth noticing: it is a tool-level release, not a ZML core version, and the README does not document a versioning scheme for the core library or a compatibility policy between core versions and model checkpoints. The last push to the default branch was on 2026-09-15, so the repository is being worked on, but commit activity is not the same as a stable API. The upgrade cost is dominated by the Bazel dependency graph: MODULE.bazel.lock is checked in, which pins dependency resolution, so a ZML upgrade can pull in a wider set of changes than the ZML diff alone suggests. Budget for a full rebuild per platform rather than an incremental swap.

Editorial conclusion

ZML fits teams that already build with Bazel and want one model definition to reach CUDA, ROCm, OneAPI, TPU or Trainium without rewriting it per vendor, and who are willing to write model code in Zig. It does not fit anyone who needs a Python API today, or who wants a runtime that installs from a package manager: the README lists Bazel as the only prerequisite, so Bazelisk comes first. Before committing, check that your target model is one of the architectures the LLM example lists (Llama 3.1 / 3.2, Qwen 3.5, LFM 2.5), confirm your accelerator flag actually builds, and read docs/howtos/deploy_on_server.md before assuming a deployment story exists.

Frequently asked questions

Official sources

1. License: Apache-2.0
2. Project website
3. README
4. Releases
5. zml/zml on GitHub