Invent

Write a model in Poem

Every model in Pie is a short Poem file. Add or edit an architecture without recompiling.

Pie has no hardcoded C++ or Rust model classes. Every model is a Poem package: a short file that declares the model's weights, how they map from a Hugging Face checkpoint, what KV or recurrent state the model keeps, and what its forward pass computes. As the figure below shows, Pie uses that one file both to convert checkpoints into .zt artifacts and to compile fused kernels for each backend.

llama_3.poemMODELSids, template, tokenizerDEPLOYMENTSweight and KV formatslayoutevery weight, by nameformatscheckpoint names → layoutcachesKV and recurrent stateforwardthe computationpie model importwrites model.ztcompilerfused kernels per backend
One file says what the weights are, where they come from, and what to compute with them.

Poem is written in Starlark (a small Python-like language) over a typed set of tensor operations. The Llama 3 definition in llama_3.poem is about 200 lines, most of which map checkpoint tensor names.

Anatomy of a Poem file

A .poem file has five main parts, shown below from llama_3.poem.

1. Model metadata (MODELS)

MODELS names each model variant and sets its chat template, tokenizer markers, layer count, and vocabulary size:

LLAMA = template("llama")
TOKENIZER = tokenizer(markers = [["<|eot_id|>", "<|eom_id|>", "<|end_of_text|>"], ...])

MODELS = [
    model("llama31-8b", template = LLAMA, tokenizer = TOKENIZER, arch = "llama", layers = 32, vocab = 128256),
    model("llama33-70b", template = LLAMA, tokenizer = TOKENIZER, arch = "llama", layers = 80, vocab = 128256),
]

2. Deployments (DEPLOYMENTS)

DEPLOYMENTS lists the supported weight and KV cache dtypes, ordered by import preference:

DEPLOYMENTS = [
    deployment("llama31-8b", weights = dtype.bf16, kv = dtype.bf16),
    deployment("llama31-8b", weights = dtype.e4m3tile128a8, kv = dtype.e4m3),   # FP8
    deployment("llama31-8b", weights = dtype.nvfp4, kv = dtype.e4m3),           # NVFP4
    deployment("llama31-8b", weights = dtype.u4g64, kv = dtype.bf16),           # 4-bit groups
]

3. Weights (layout and formats)

layout returns a struct of every weight tensor with its shape and dtype. Field paths such as layers.3.q become the tensor names in the .zt file. formats maps a Hugging Face checkpoint's tensor names onto those paths:

("model.layers.{l}.self_attn.q_proj.weight", "layers.{l}.q"),

4. State (caches)

caches declares the KV cache spaces, or recurrent state slots for hybrid and SSM models:

def caches(m, c):
    kv = c.kv_space(m.kv)
    for w in m.layers:
        c.kv(kv, w.kv, [m.kv_width, m.kv_width], m.head_dim, heads = True)

5. Forward pass (forward)

forward defines the computation graph, which the compiler traces once into backend kernels:

def forward(m, inputs):
    y = ops.gather.embed(inputs.tokens(), m.embed, m.vocab)

    def block(l, w, y):
        x = ops.elemwise.rmsnorm(y, w.attn_norm, m.eps)
        y = ops.elemwise.residual_add(attention(x, inputs, m, w, inputs.positions()), y)
        x = ops.elemwise.rmsnorm(y, w.mlp_norm, m.eps)
        return ops.elemwise.residual_add(swiglu(x, w), y)

    y = inputs.fold_layers(m.layers, y, block)
    x = ops.elemwise.rmsnorm(y, m.final_norm, m.eps)
    x = ops.gather.gather_rows(x, inputs.readout_rows())
    output("logits", ops.linear.lm_head(x, m.head))

Any tensor passed to output(name, value) can be read by an epilogue via intrinsics.output(lane, name). Models can declare additional outputs beyond "logits".

Shared building blocks

Common layer blocks live in models/lib and can be loaded as modules: grouped-query attention, MLA, RoPE variants, gated delta-net (kda), diffusion transformer blocks (dit), VAEs, vision towers, and adapters.

load("//lib/blocks.poem", "attention", "swiglu")

Add or edit a model without recompiling

Pie unpacks its built-in .poem files into $PIE_HOME/models/ on startup. Any file you edit or add in $PIE_HOME/models/ is used immediately without rebuilding Pie:

  1. Copy the closest existing model (such as llama_3.poem) in $PIE_HOME/models/.
  2. Update the dimensions, weight mappings, and forward function.
  3. Run pie model import <org>/<name>. The importer checks every checkpoint tensor against your layout before writing the .zt file.

When porting a new architecture, add a small "mini" model variant with only 2 layers first so you can test the import and forward pass end to end in seconds.

Beyond text

The same language defines vision encoders (media), image and video diffusion models (FLUX.2, Wan 2.2, LTX 2.5), and hybrid autoregressive-diffusion models (HunyuanImage 3). See text-to-media for generating images and video from an inferlet.

On this page

Anatomy of a Poem file1. Model metadata (MODELS)2. Deployments (DEPLOYMENTS)3. Weights (layout and formats)4. State (caches)5. Forward pass (forward)Shared building blocksAdd or edit a model without recompilingBeyond text