# gizai:dit the diffusion transformer's block that every generation model shares (MiniMax H3, LTX-2.5, FLUX.3, Z-Image, Krea 2, OmniVoice). A block is an AdaLN transformer layer: modulated norms, attention over a packed sequence with rotary coordinates, and a gated MLP. This module owns its weights as stored (BF16, ComfyUI's INT8 ConvRot or FP8, H3's W4, with LoRA adapters), the packed sequence, the time embedding, the host reference and the GPU kernels, and the staging of a block's weights in pinned memory for streamed models. ```gizlang import * as Dit from "gizai:dit"; ``` ## Types ### `Arch` ```gizlang type Arch = { family: string; hidden: i64; layers: i64; refiner: i64; heads: i64; head: i64; ffn: i64; latents: i64; audio: i64; patch: Vec; textDim: i64; freqDim: i64; timeHidden: i64; timeDim: i64; eps: f64; qkEps: f64; finalEps: f64; } ``` What a model computes, as its family's recipe reads config.json. - `family`: `string` - `hidden`: `i64` - `layers`: `i64` - `refiner`: `i64` - `heads`: `i64` - `head`: `i64` - `ffn`: `i64` - `latents`: `i64` - `audio`: `i64` - `patch`: `Vec` - `textDim`: `i64` - `freqDim`: `i64` - `timeHidden`: `i64` - `timeDim`: `i64` - `eps`: `f64` - `qkEps`: `f64` - `finalEps`: `f64` ### `Block` ```gizlang type Block = { w: Vec; bias: Vec; codes: Vec; lora: Vec; } ``` One block's weights, by the indexes below: its 4 linears [out, in] as stored, then its 4 norms (f32). On a GPU a linear is BF16, or W4 when `codes` holds the decoding of each (`w` then holds its codes); on the host it is f32 values (W4: the code bytes as numbers). `bias` holds the linears' biases (f32, one per output) or none; `lora` the A ([rank, in]) and B ([out, rank]) of each linear's LoRA, applied beside it (BF16 on a GPU; a W4 block's A there is INT8 as N.int8Rows makes it, its row scales in B's columns), or none. UP's rows are the gate's and the up projection's in pairs (gluRows; SwiGLU: silu(gate) times up). - `w`: `Vec` - `bias`: `Vec` - `codes`: `Vec` - `lora`: `Vec` ### `Sequence` ```gizlang type Sequence = { rows: i64; map: Tensor; cos: Tensor; sin: Tensor; exact: bool; group: i64; runs: Vec; queries: i64; boost: Vec; } ``` The rows of one packed sequence: per row its modulation row (`map`, as f32 values), and the cosines and sines of its rotary angles (per axis its frequencies: H3's t, h and w; FLUX.3's four). Tensors where the block runs. `exact`: its attention in BF16 where 8-bit matrix cores would take it (the reference of N.attentionRows quantized). `boost`: keys whose attention weights some queries multiply (N.Boost; Krea 2's edit). - `rows`: `i64` - `map`: `M45.Tensor` - `cos`: `M45.Tensor` - `sin`: `M45.Tensor` - `exact`: `bool` - `group`: `i64` - `runs`: `Vec` - `queries`: `i64` - `boost`: `Vec` ### `Modulation` ```gizlang type Modulation = { attn: Vec; mlp: Vec; } ``` A block's modulation: shift, scale and gate of the attention, then of the MLP, each [modulation rows, hidden] f32. - `attn`: `Vec` - `mlp`: `Vec` ### `Staged` ```gizlang type Staged = { host: Pinned; bytes: i64; parts: Vec; counts: Vec; } ``` A block in page-locked host memory: the bytes of its GPU tensors (blockParts) one after another, each from a multiple of 256, copied to a slot on the GPU ahead of each use (Gpu.copyAhead). `counts`: Block's lists' lengths. - `host`: `M10.Pinned` - `bytes`: `i64` - `parts`: `Vec` - `counts`: `Vec` ### `Part` ```gizlang type Part = { shape: Vec; dtype: i64; at: i64; size: i64; host: Tensor; } ``` A tensor of a staged block: its shape, type, and bytes from `at`; a host tensor (`at` -1) stays as `host`. - `shape`: `Vec` - `dtype`: `i64` - `at`: `i64` - `size`: `i64` - `host`: `M45.Tensor` ### `Work` ```gizlang type Work = { gpu: Gpu; x: Buffer; q: Buffer; s: Buffer; down: Buffer; mid: Buffer; w: Buffer; } ``` Device memory a block's forward reuses for each chunk of rows (gpuForward): the input rows of a linear (BF16), their INT8 values and scales, its LoRA's down product, SwiGLU, and the INT8 weights of the linears of a W4 block (the largest two). - `gpu`: `M10.Gpu` - `x`: `M10.Buffer` - `q`: `M10.Buffer` - `s`: `M10.Buffer` - `down`: `M10.Buffer` - `mid`: `M10.Buffer` - `w`: `M10.Buffer` ### `RefSlot` ```gizlang type RefSlot = { first: i64; count: i64; rows: Tensor; record: bool; } ``` A joint block's reference rows kept across steps (Flux.sample's `refresh`): their keys and values, [count, 2 width] BF16 in `rows`. A step that records copies them from rows `first` to `first` + `count` of the block's qkv; the others, which update only the rows before `first`, put them back there. - `first`: `i64` - `count`: `i64` - `rows`: `M45.Tensor` - `record`: `bool` ## Functions ### `blockTensors` ```gizlang function blockTensors(): Vec ``` The names of the tensors of a block, in the order of the indexes, and whether each is a linear. Effects: fail. ### `gluRows` ```gizlang function gluRows(t: Tensor): Tensor ``` The rows of t ([gate; up] along its first axis, 2 ffn of them) interleaved as UP holds them: gate j at 2 j, up j at 2 j + 1, the pairs of the products' SwiGLU (Lazy.gluInto). Its type unchanged. Effects: fail, read, write, console, io. ### `weightAt` ```gizlang function weightAt(sys: System, named: Map, lora: Map, name: string, gpu: Option, linear: bool, head: i64 = 0): Tensor throws string ``` Weight `name` where it runs (as storedOn), plus its LoRA's update B A when `lora` has one (W + B A, as the adapters of H3 apply). With `head` > 0 the weight is a fused QKV projection stored as H3's checkpoints hold it, a head's query, key and value rows together (heads, 3, head); its rows are reordered to all queries, all keys, all values first, the order of the LoRA's rows and of this module's attention. Effects: fail, read, write, console, io. ### `loadBlock` ```gizlang function loadBlock(sys: System, named: Map, lora: Map, prefix: string, gpu: Option, interleaved: bool): Block throws string ``` Block `prefix` (as "blocks.3.") of the checkpoint's tensors on `gpu` when there is one, with the LoRA `lora` (an empty map: none): W4 or ComfyUI's INT8 ConvRot as stored when the checkpoint holds it so (w4Block, int8Block), else BF16 with the LoRA merged, its QKV rows a head at a time when `interleaved` (see weightAt; production's files hold them grouped). ComfyUI's scaled FP8 is quantized to INT8 (Block.quantized) after the merge: as BF16 it would take twice the GPU memory. Effects: fail, read, write, console, io. ### `rawAt` ```gizlang function rawAt(sys: System, s: Stored, gpu: Option, dtype: i64): Tensor throws string ``` A stored tensor of bytes (U8 or I8 codes, or E4M3 scales) where it runs: its bytes as `dtype` on a GPU, their values on the host (codes as unsigned bytes). Effects: fail, read, write, io. ### `runsOf` ```gizlang function runsOf(map: Vec): Vec ``` The runs of rows of one modulation row of `map`: (first, end, row) each, for products that gate in their epilogue (gatedInto). Effects: fail. ### `sequence` ```gizlang function sequence(positions: Vec, map: Vec, inverse: Vec, gpu: Option, exact: bool = false, group: i64 = 0): Sequence ``` The sequence of `positions` (rows of a position per axis: H3's t, h, w) whose rows use modulation rows `map`, with the rotary inverse frequencies `inverse` of each axis (H3's rope.inv_freq: head / 8 of them), on `gpu` when there is one. Rows attend within groups of `group` rows (0: all of them), as tiles of a VAE. Effects: fail, read, write, console, io. ### `sequenceOf` ```gizlang function sequenceOf(positions: Vec, map: Vec, inverses: Vec>, gpu: Option, exact: bool = false, group: i64 = 0): Sequence ``` `sequence` with each axis its own inverse frequencies (Z-Image's 16, 24 and 24). Effects: fail, read, write, console, io. ### `stagedBytes` ```gizlang function stagedBytes(b: Block): i64 ``` The bytes of b's staged copy (Staged.bytes) without making it. Effects: fail. ### `stage` ```gizlang function stage(g: Gpu, b: Block): Staged ``` Block b in page-locked host memory (Staged): its GPU tensors copied, its host ones kept. Effects: fail, read, write, io. ### `silu` ```gizlang function silu(x: Tensor): Tensor ``` Effects: fail, read, write, console, io. ### `sinusoids` ```gizlang function sinusoids(times: Vec, dim: i64): Tensor ``` Rows [cos | sin] of each time times the dim / 2 frequencies 10000^(-i / (dim / 2)), in f32 as diffusers' Timesteps (flip_sin_to_cos) computes them: the timestep embedding of H3 and FLUX.3. Effects: fail. ### `workOf` ```gizlang function workOf(gpu: Gpu, b: Block, rows: i64, a: Arch): Work ``` Effects: fail, read, write, io. ### `linearInto` ```gizlang function linearInto(work: Work, x: Tensor, b: Block, i: i64, w: Tensor, out: Tensor, ready: bool = false): void ``` The product x Wᵀ of linear i (`w` its weights as weightsOf gives them, [out, in]) on the matrix cores, plus its bias and its LoRA's product (withLora), as BF16 into `out` (x BF16). A W4 or INT8 linear (`codes`) takes x through ConvRot to INT8 (N.int8Rows) against its INT8 weights, the product times x's scale and W's: WanGP's asym_w4a8_int8 product. Effects: fail, read, write, console, io. ### `withLora` ```gizlang function withLora(b: Block, i: i64, down: Tensor, y: T.Lazy): T.Lazy ``` Linear i's output y plus its LoRA's B product of `down` (loraOf's rows), each rounded to BF16 as production stores them; in y's product's epilogue (Product.low), so neither is stored. y alone without a LoRA. Effects: fail, read, write, console, io. ### `quantizedOf` ```gizlang function quantizedOf(work: Work, x: Tensor, b: Block, ready: bool): Vec ``` A W4 or INT8 block's rows of x as INT8 and their scales (work.q, work.s; N.int8RowsInto unless `ready`), which its linears and a LoRA's INT8 A take; x and x for a BF16 block. Effects: fail, read, write, console, io. ### `productOf` ```gizlang function productOf(b: Block, i: i64, w: Tensor, x: Tensor, q: Tensor, s: Tensor): T.Lazy ``` Linear i's product of x's rows (q and s: quantizedOf's) with its bias, for a chain to go on in its epilogue. Effects: fail, read, write, console, io. ### `loraOf` ```gizlang function loraOf(work: Work, x: Tensor, b: Block, i: i64, qs: Vec): Tensor ``` Linear i's LoRA A product of x (its INT8 rows when A is INT8) into work.down; x itself without a LoRA. Effects: fail, read, write, console, io. ### `gatedInto` ```gizlang function gatedInto(work: Work, x: Tensor, y: Tensor, b: Block, i: i64, w: Tensor, gate: Tensor, runs: Vec, ready: bool): void ``` x += gate times linear i of y, in a product's epilogue: a product for each run of rows of one modulation row (Sequence.runs), whose gate row the chain takes, so the linear's output is never stored; its LoRA's product too (withLora). Effects: fail, read, write, console, io. ### `int8Cores` ```gizlang function int8Cores(gpu: Option): bool ``` Whether `gpu` is a GPU with INT8 matrix cores (K.hasI8Mma): there W4 and INT8 blocks run as stored. Effects: fail. ### `normedInto` ```gizlang function normedInto(work: Work, b: Block, a: Arch, xs: Tensor, w: Tensor, shift: Tensor, scale: Tensor, map: Tensor, y: Tensor): void ``` The modulated norm of rows xs into y (BF16), as a block's linears take it: an INT8 or INT4 block's to its rows and scales (work.q, work.s) in the same pass (N.modulatedInt8RowsInto), y then written for a LoRA of BF16 A alone. Effects: fail, read, write, console, io. ### `runsIn` ```gizlang function runsIn(runs: Vec, r0: i64, r1: i64): Vec ``` The runs (Sequence.runs) of rows [r0, r1), from r0. Effects: fail. ### `prepOf` ```gizlang function prepOf(b: Block, s: Sequence, a: Arch): N.HeadPrep ``` The head norm and H3's split rotary of a block's packed qkv rows (N.HeadPrep): each head of q and k over its RMS times its norm weight, then its first 2 half values turned by the sequence's cos and sin rows, as attention reads them. Effects: fail. ### `comfyInt8Rows` ```gizlang function comfyInt8Rows(sys: System, named: Map, names: Vec, gpu: Gpu): Vec throws string ``` Linears `names` of a ComfyUI INT8 ConvRot file (each name's .weight and .weight_scale: INT8 rows rotated by the regular Hadamard of 256, a scale a row) stacked by rows on `gpu`, as N.int8Rows' products read them (canonicalInto): [the INT8 weights [rows, in], their scales [1, rows] F32]. Effects: fail, read, write, console, io. ### `canonicalInto` ```gizlang function canonicalInto(w: Tensor, from: Vec, out: Tensor, bits: i64 = 8): void ``` INT8 rows w ([out, in]) rotated in the regular Hadamard's own order, as ConvRot files store them, into `out`: its row r is w's row from[r], its inputs in the order N.int8Rows rotates inputs to (as W4.decodeInto's codes). With 4 `bits` w holds INT4 pairs ([out, in / 2] bytes, the even input in the low nibble: convrot_w4a4) and out is I4 [out, in]. Effects: fail, read, write, console, io. ### `blockOf` ```gizlang function blockOf(w: Vec): Block ``` A block of weights `w` (Block.w): no biases, codes or LoRA. ## Methods ### `Arch.width` ```gizlang function Arch.width(self: Arch): i64 ``` Values of the attention's queries (and of its keys, and of its values) per row. Effects: fail. ### `Sequence.active` ```gizlang function Sequence.active(self: Sequence, queries: i64): Sequence ``` Sequence `s` where only its first `queries` rows are updated, the others (cached keys and values: RefSlot) being read only: the same positions and modulation rows, the runs cut to the first rows. Effects: fail. ### `Block.forward` ```gizlang function Block.forward(self: Block, a: Arch, x: Tensor, m: Modulation, s: Sequence): Tensor ``` The block over x ([rows, hidden], f32): x + gate (attention of the modulated norm), then the same with the MLP. On a GPU x is updated in place and returned; on the host a new tensor. Effects: fail, read, write, console, io. ### `Staged.on` ```gizlang function Staged.on(self: Staged, g: Gpu, slot: Buffer): Block ``` The staged block on the GPU memory `slot` (its bytes there once copied). Effects: fail. ### `Block.quantized` ```gizlang function Block.quantized(self: Block, bits: i64 = 8): Block ``` The block (BF16) with each linear quantized once as N.int8Rows quantizes x: its rows through ConvRot to INT8 over a scale a row (INT4 with 4 `bits`, where K.hasI4Mma: about 1.8 times the products' speed, at a loss of quality). Products then run as a W4 block's (linear). On the host, or on a GPU without INT8 matrix cores (K.hasI8Mma: Apple GPUs), the block stays BF16 and its products run on the F16 and BF16 matrix cores. Effects: fail, read, write, console, io. ### `Block.dequantized` ```gizlang function Block.dequantized(self: Block): Block ``` The block (W4 as stored, or INT8 as Block.quantized and ComfyUI's files hold it) as BF16: on a GPU without INT8 matrix cores, and before a requantization. Each linear is its INT8 weights times each output's scale, turned back from ConvRot's rotation: N.int8Rows takes each 256 inputs x to y = R x, so the weights of x are those of y times R. A LoRA stays as it is: its A must be BF16, not INT8. Effects: fail, read, write, console, io. ### `Block.parallel` ```gizlang function Block.parallel(self: Block, a: Arch, x: Tensor, m: Modulation, s: Sequence, slot: Option = undefined): Tensor ``` The parallel block over x ([rows, hidden], f32): x + gate (attention + SwiGLU MLP), both of one modulated LayerNorm. Its MLP's input linear has the gate and up rows interleaved (fluxBlock): row 2j is gate j, 2j + 1 up j. On a GPU x is updated in place and returned; on the host a new tensor. Effects: fail, read, write, console, io. ## Constants ### `HEAD` ```gizlang const HEAD = 128 ``` Values of an attention head of H3, and the rotary frequencies of each of its (t, h, w) axes. A head of any multiple of 32 values (`Arch.head`) rotates its first 3/4: the t, h and w angles of head / 8 frequencies each, twice (split halves). ### `AXIS_FREQUENCIES` ```gizlang const AXIS_FREQUENCIES = 16 ``` The rotary frequencies of each of a head's (t, h, w) axes: HEAD / 8. ### `QKV` ```gizlang const QKV = 0 ``` The index in Block.w of the fused query, key and value projection. ### `OUT` ```gizlang const OUT = 1 ``` The index in Block.w of the attention's output projection. ### `UP` ```gizlang const UP = 2 ``` The index in Block.w of the MLP's up projection (gate and up rows in pairs: gluRows). ### `DOWN` ```gizlang const DOWN = 3 ``` The index in Block.w of the MLP's down projection. ### `NORM1` ```gizlang const NORM1 = 4 ``` The index in Block.w of the norm weight before the attention. ### `NORM2` ```gizlang const NORM2 = 5 ``` The index in Block.w of the norm weight before the MLP. ### `Q_NORM` ```gizlang const Q_NORM = 6 ``` The index in Block.w of the queries' head norm weight. ### `K_NORM` ```gizlang const K_NORM = 7 ``` The index in Block.w of the keys' head norm weight. ### `POST_ATTN` ```gizlang const POST_ATTN = 8 ``` The indexes in Block.w of a sandwich block's norm weights after the attention and after the MLP (Z-Image): its outputs are normed before their gates. A block without them adds its gated outputs as they are. ### `POST_MLP` ```gizlang const POST_MLP = 9 ``` The index in Block.w of a sandwich block's norm weight after the MLP (see POST_ATTN). ### `CHUNK` ```gizlang const CHUNK = 4096 ``` Rows a chunk of a block's linears and MLP takes: their temporaries (Work) stay within a few hundred MB.