# gizml:npullm Gemma 4 E2B's forward pass on an AMD NPU (XDNA 2), with the AIE kernels of FastFlowLM (the layer and lm_head xclbins, MIT runtime, binaries free to use unmodified). gizml:llm loads such a folder as a model like any other (its chat, sampling and serving); this module gives it the logits of the tokens it reads (Model.read). A token runs the 35 decoder layers as 35 runs of the layer kernel and then one run of the lm_head kernel; this module writes each run's control code (an AIE transaction) for the current context length, and the CPU does the embeddings and the rotary tables. The model folder holds config.json, tokenizer.json, model.q4nx (a safetensors file of 4.125-bit Q41 blocks, bf16 and q8_0 tensors) and xclbins/. ```gizlang import * as Npullm from "gizml:npullm"; ``` ## Types ### `Shape` ```gizlang type Shape = { d: i64; layers: i64; ffn: i64; pli: i64; heads: i64; kvHeads: i64; dh: i64; swaDh: i64; window: i64; period: i64; shared: i64; vocab: i64; padded: i64; softcap: f64; wide: bool; maxL: i64; } ``` The sizes of the model (config.json). A layer is global (full attention, head size `dh`) when (i + 1) is a multiple of `period`, else sliding-window (head size `swaDh`, the last `window` positions); the last `shared` layers reuse the cache of the last layer of their kind before them. - `d`: `i64` - `layers`: `i64` - `ffn`: `i64` - `pli`: `i64` - `heads`: `i64` - `kvHeads`: `i64` - `dh`: `i64` - `swaDh`: `i64` - `window`: `i64` - `period`: `i64` - `shared`: `i64` - `vocab`: `i64` - `padded`: `i64` - `softcap`: `f64` - `wide`: `bool` - `maxL`: `i64` ### `Model` ```gizlang type Model = { shape: Shape; npu: Session; layer: Vec; head: Program; commands: Vec; x: Buffer; proj: Vec; rms: Vec; rope: Vec; kv: Vec; lmHead: Buffer; logits: Buffer; prompt: Pass; embed: Vec; tokenizer: Tokenizer.Tokenizer; stops: Vec; length: i64; } ``` A loaded model: the NPU session, the programs of the two kernels (`layer` holds one for each of the four layer kinds), the buffers, the prompt pass, the embedding tensors (read from the file a row at a time), the tokenizer, and the context length so far. - `shape`: `M28.Shape` - `npu`: `M14.Session` - `layer`: `Vec` - `head`: `M1.Program` - `commands`: `Vec` - `x`: `M1.Bo` - `proj`: `Vec` - `rms`: `Vec` - `rope`: `Vec` - `kv`: `Vec` - `lmHead`: `M1.Bo` - `logits`: `M1.Bo` - `prompt`: `M28.Pass` - `embed`: `Vec` - `tokenizer`: `M34.Tokenizer` - `stops`: `Vec` - `length`: `i64` ### `Mm` ```gizlang type Mm = { m: i64; k: i64; n: i64; args: Vec; clamp: bool; } ``` A bf16 matrix-multiply kernel: C (M by N, row-major) = A (M by K, row-major) times B, whose columns come in groups of `n`, each group n by K, row-major. `m`, `k` and `n` are a core's tile; `args` are the kernel arguments of A, B and C; `clamp` says that the run-time parameters end with an output clamp. - `m`: `i64` - `k`: `i64` - `n`: `i64` - `args`: `Vec` - `clamp`: `bool` ### `Pass` ```gizlang type Pass = { mm: Vec; dequant: Vec; attention: Vec; rows: i64; start: i64; stop: i64; hidden: Buffer; q: Buffer; k: Buffer; v: Buffer; y: Buffer; gate: Buffer; up: Buffer; hid: Buffer; pliDown: Buffer; pliGate: Buffer; pliHid: Buffer; w: Vec>; cache: Vec; pliDownW: Buffer; pliW: Vec; norms: Vec; scales: Vec; pliNorm: bytes; } ``` The programs and buffers of the prompt pass. `rows`, `start` and `stop` are the geometry the row and attention codes have now. The buffers: the rows of the pass (hidden state, q, k, v, attention output, gate, up, their product, PLI down, gate and product), the bf16 projections of the current layer (q, k and v; o; gate; up; down), the prompt caches of a sliding-window and a global layer, and the bf16 PLI projections. The host copies of each layer's norms: input, post attention, pre and post MLP, q, k, post PLI; and its output scale. - `mm`: `Vec` - `dequant`: `Vec` - `attention`: `Vec` - `rows`: `i64` - `start`: `i64` - `stop`: `i64` - `hidden`: `M1.Bo` - `q`: `M1.Bo` - `k`: `M1.Bo` - `v`: `M1.Bo` - `y`: `M1.Bo` - `gate`: `M1.Bo` - `up`: `M1.Bo` - `hid`: `M1.Bo` - `pliDown`: `M1.Bo` - `pliGate`: `M1.Bo` - `pliHid`: `M1.Bo` - `w`: `Vec>` - `cache`: `Vec` - `pliDownW`: `M1.Bo` - `pliW`: `Vec` - `norms`: `Vec` - `scales`: `Vec` - `pliNorm`: `bytes` ## Functions ### `shapeOf` ```gizlang function shapeOf(config: Json.Json): Shape ``` The shape of the model whose config.json (the "text_config" part, or the whole) is `config`. Effects: fail. ### `kindOf` ```gizlang function kindOf(s: Shape, i: i64): i64 ``` The kind of layer `i`: bit 0 global, bit 1 shared cache. Effects: fail. ### `projOffsets` ```gizlang function projOffsets(s: Shape, kind: i64): Vec ``` The byte offsets in the projection buffer of a layer of `kind`: q, k, v (none on a shared layer), o, up and gate, down, then the bf16 PLI down, gate and up projections; the last is the buffer's size. Effects: fail. ### `ropeOffsets` ```gizlang function ropeOffsets(s: Shape, kind: i64): Vec ``` The bf16 element offsets in the rope buffer of a layer of `kind`: cos and sin, q norm, k norm, PLI embedding, PLI norm, post-PLI norm, the layer scale; the last is the elements with the padding. Effects: fail. ### `layerCode` ```gizlang function layerCode(s: Shape, kind: i64, length: i64): bytes ``` The control code of a layer of `kind` that attends to `length` positions (the new one included). Effects: fail. ### `headCode` ```gizlang function headCode(s: Shape): bytes ``` The control code of the lm_head kernel: the final norm and the logit soft cap (`softcap`, as f32 bits) go to the 32 cores; each round the 8 shim tiles send the Q41 rows of 1024 tokens and take their logits back. Effects: fail. ### `reordered` ```gizlang function reordered(data: bytes, columns: i64, vertical: i64): bytes ``` The Q41 blocks of `data` (rows of blocks, `columns` / 256 blocks a row) in the order the kernels read them: each column of blocks of `vertical` rows of blocks together. Effects: fail. ### `load` ```gizlang function load(sys: System, folder: string): Model throws string ``` Loads the model in `folder` onto the first NPU of this host. Effects: fail, io. ### `lmMm` ```gizlang function lmMm(): Mm ``` The kernel of the language model's projections. Effects: fail. ### `mmCode` ```gizlang function mmCode(g: Mm, rows: i64, inputs: i64, outputs: i64, b: i64, act: i64): bytes ``` The control code of `g` over A of `rows` by `inputs` and an output of `outputs` columns, with B from element `b` of its buffer and activation `act`. The 8 shim tiles each take a group of B columns and their C; the even ones also take the rows of A, 4 by m at a time; each queue holds two descriptors (ping-pong). Effects: fail. ### `dequantCode` ```gizlang function dequantCode(offset: i64, inputs: i64, outputs: i64, band: i64): bytes ``` The control code of dequant.xclbin over the Q41 matrix at byte `offset` of the projection buffer (`inputs` columns, `outputs` rows): the bf16 matrix in groups of 128 rows, each 128 by `inputs`, row-major in pieces 512 wide (B of an Mm). `band` (0 none, 1 up, 2 gate) takes one matrix of the up and gate band, where they alternate every ROUND_ROWS rows. Each of the 8 shim tiles takes 128 rows a round. Effects: fail. ### `attentionCode` ```gizlang function attentionCode(s: Shape, global: bool, begin: i64, end: i64): bytes ``` The control code of an attention kernel (attn.xclbin for a global layer, swa.xclbin for a sliding-window one) for the query rows of positions `begin` to `end` (multiples of CHUNK_ROWS): y (argument 0) gets each head's output from q (argument 1, the rows of the pass), and the k and v of the positions before go in from the prompt cache (argument 2: k of `maxL` positions, then v). A global kernel works two heads at a time on 2 groups of 4 columns, a sliding one four heads on 4 groups of 2; each takes 128 query rows a round. Effects: fail. ## Methods ### `Model.step` ```gizlang function Model.step(self: mut Model, sys: System, id: i64): Vec throws string ``` Runs token `id` at the next position and returns its logits, the vocabulary's. Effects: fail, io. ### `Model.read` ```gizlang function Model.read(self: mut Model, sys: System, at: i64, ids: Vec): Vec throws string ``` Reads `ids` from position `at` (the cache holds the tokens before it) and returns the logits after the last, the vocabulary's: passes of up to PASS_TOKENS tokens, then a rest of fewer than PASS_MIN tokens one at a time. Effects: fail, io. ### `Model.pass` ```gizlang function Model.pass(self: mut Model, sys: System, ids: Vec): Vec throws string ``` Runs the tokens `ids` (at most PASS_TOKENS) at the next positions in one pass and returns the logits of the last. Effects: fail, io.