zml/zml

ZML – A Production‑grade Inference Stack

What it is – ZML is an open‑source framework for running AI models (LLMs, vision models, etc.) at high performance on many kinds of accelerators (NVIDIA CUDA, AMD ROCm, Intel OneAPI, Google TPU, AWS Trainium/Inferentia). It lets you write a model once in Zig and have the same code compiled to the native instruction set of each hardware target, without hand‑written device‑specific code.

Key components

• Zig‑based library – the core API (zml.Tensor, zml.Platform, zml.Mesh, …) is written in the systems language Zig, giving low‑level control and safety.
• MLIR & Bazel – the build pipeline uses MLIR for intermediate representation and Bazel (via Bazelisk) to compile and link for each backend.
• Unified VFS – models can be loaded from local paths, Hugging‑Face (hf://), HTTP(S), or S3 URLs through a virtual file‑system layer.
• Platform flags – at build or run time you enable a target with flags like --@zml//platforms:cuda=true.

How to try it

1. Install Bazelisk (the only prerequisite). On macOS: brew install bazelisk; on Linux download the binary.
2. Run the tiny MNIST demo (downloads a pre‑trained model, compiles, and classifies a random digit):

   bazel run //examples/mnist
   

3. Run an LLM example (Llama‑3.2, Qwen‑3.5, or LFM‑2.5). Authenticate to Hugging Face if needed:

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

   Add --@zml//platforms:cuda=true (or rocm, oneapi, tpu, neuron) to run on a specific accelerator.

Example code snippet – a minimal MNIST model written in Zig shows the API style:

const Mnist = struct {
    fc1: Layer,
    fc2: Layer,
    // … layer definition omitted …
    pub fn forward(self: Mnist, input: zml.Tensor) zml.Tensor {
        var x = input.flatten().convert(.f32);
        for (self.layers) |layer| { x = layer.forward(x); }
        return x.argMax(0).indices.convert(.u8);
    }
};

The model is loaded with zml.io.load, compiled for the chosen platform, and executed with the same call site regardless of hardware.

Where to go next

• Browse the other examples (sharding, io, benchmark).
• Follow the tutorials in docs/tutorials/ to write your own model.
• Read the deployment guide (docs/howtos/deploy_on_server.md) for serving models in production.

Community & licensing – ZML is Apache‑2.0 licensed, has a Discord community, and accepts contributions via the usual GitHub workflow.

---

TL;DR – ZML lets you write AI models once in Zig and compile them to run efficiently on any major accelerator, using Bazel/MLIR under the hood. It provides ready‑to‑run examples for MNIST and several LLM families, a unified model‑loading VFS, and a clear path from prototype to production deployment.