Poem ops
Declarations and operations for .poem model files.
A Poem package is a Starlark file that declares a model's weights, state spaces, and forward computation graph. Write a model in Poem walks through a complete package.
Declarations
| Name | Purpose |
|---|---|
MODELS | model(id, template=, tokenizer=, arch=, layers=, vocab=) for each model variant in the file. |
DEPLOYMENTS | deployment(id, weights=, kv=, ...) in default import preference order. |
PUBLISHED | Speculative draft heads exported by the package. |
layout(id, deploy) | Returns the weight struct; field paths name the tensors in .zt files. |
formats(m) | format(name, recognizes=, read=, states=): maps Hugging Face checkpoint names onto layout. |
caches(m, c) | Declares KV cache spaces and recurrent state slots. |
forward(m, inputs) | Defines the forward pass and registers outputs with output(name, value). |
media(id, deploy) | Declares how images and audio are patchified for multimodal models. |
template(format) | Selects a built-in chat template: chatml, harmony, gemma, deepseek, dots, glm, kimi, kimi3, inkling, atem, mimo, minimax, hy, lines, or raw. |
tokenizer(markers=, pinned=, parts=) | Declares required special tokens. |
Load shared helpers from models/lib with load("//lib/blocks.poem", "attention", "swiglu").
Data types (dtype)
Supported weight and KV cache dtypes include dtype.bf16, dtype.e4m3 (FP8), dtype.e4m3tile128a8 (block-scaled FP8), dtype.nvfp4, and dtype.u4g64 (4-bit integer, group size 64).
Operations (ops)
| Namespace | Operations |
|---|---|
ops.linear | matmul, matmul_grouped, glu, gelu_tanh, lm_head, moe_route, moe_matmul_select, moe_weighted_sum, group_routes |
ops.elemwise | add, mul, mul_scalar, scale, clamp, copy, silu, gelu, rmsnorm and variants, layernorm, standardize, residual_add, gated_residual_add, rope, rope_axes, sinusoid, modulate, add_bias, gate_sigmoid_mul |
ops.attn | attend, prefill, plan_prefill, kv_append, logit_softcap, MLA (mla_latents, mla_kv_append, mla_absorb_q, mla_absorb_out), pooled and indexed attention (pool_kv_append, pool_gather, index_topk, index_kv_append, index_rope) |
ops.ssm | gdn_gates, causal_conv1d, causal_conv1d_chunked |
ops.gather | embed, embed_weighted, gather_rows, split_rows, split_qkv, split_q_gate, concat_cols, pack_rows, unpack_rows, group_rows, scatter_live_rows, pixels |
ops.conv | conv3d, cells, patchify, unpatchify, pixel_shuffle, upsample_nearest |
ops.sampling | argmax |
inputs
Inside forward(m, inputs), inputs provides access to the current pass: tokens(), positions(), readout_rows(), kv(space), adapter_routes(), and fold_layers(layers, y, block), which applies block across every layer.