A Zig Inference Stack That Refuses to Pick a GPU Vendor

What it does

ZML is an inference framework written in Zig that compiles AI models to multiple hardware backends—NVIDIA CUDA, AMD ROCm, Google TPU, and AWS Trainium—through MLIR. You define models as ordinary Zig structs with tagged tensors, and the stack handles lowering, weight loading from local or remote stores, and sharded execution. It is pitched as a way to keep a single codebase while avoiding vendor-specific lock-in.

The interesting bit

Most ML infrastructure is either Python-centric or welded to a single vendor’s toolchain. ZML flips that by using Zig’s compile-time metaprogramming and MLIR’s multi-target machinery to treat hardware backends as interchangeable compilation targets. The sharding and VFS layers suggest it is aiming for production deployment, not just toy examples.

Key highlights

• Targets NVIDIA, AMD, TPU, and AWS Trainium from one Zig codebase.
• Model definitions are plain Zig structs with tensor operations; no Python indirection required.
• Includes a VFS layer for loading weights from local disk, Hugging Face, HTTP, or S3.
• Supports model parallelism via logical mesh sharding and partitioners.
• Apache 2.0 licensed.

Verdict

Worth a look if you are building inference infrastructure and want to hedge against single-vendor hardware lock-in—especially if you already tolerate Bazel and Zig’s ecosystem. If you need the maturity of PyTorch or a Python-native workflow, this is not your stack yet.