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.

import * as Dit from "gizai:dit";

Types

Arch

type Arch = { family: string; hidden: i64; layers: i64; refiner: i64; heads: i64; head: i64; ffn: i64; latents: i64; audio: i64; patch: Vec<i64>; 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<i64>
• textDim: i64
• freqDim: i64
• timeHidden: i64
• timeDim: i64
• eps: f64
• qkEps: f64
• finalEps: f64

Block

type Block = { w: Vec<Tensor>; bias: Vec<Tensor>; codes: Vec<W4>; lora: Vec<Tensor>; }

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<M45.Tensor>
• bias: Vec<M45.Tensor>
• codes: Vec<M41.W4>
• lora: Vec<M45.Tensor>

Sequence

type Sequence = { rows: i64; map: Tensor; cos: Tensor; sin: Tensor; exact: bool; group: i64; runs: Vec<i64>; queries: i64; boost: Vec<N.Boost>; }

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<i64>
• queries: i64
• boost: Vec<M36.Boost>

Modulation

type Modulation = { attn: Vec<Tensor>; mlp: Vec<Tensor>; }

A block's modulation: shift, scale and gate of the attention, then of the MLP, each [modulation rows, hidden] f32.

• attn: Vec<M45.Tensor>
• mlp: Vec<M45.Tensor>

Staged

type Staged = { host: Pinned; bytes: i64; parts: Vec<Part>; counts: Vec<i64>; }

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<M23.Part>
• counts: Vec<i64>

Part

type Part = { shape: Vec<i64>; 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<i64>
• dtype: i64
• at: i64
• size: i64
• host: M45.Tensor

Work

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

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

function blockTensors(): Vec<string>

The names of the tensors of a block, in the order of the indexes, and whether each is a linear.

Effects: fail.

gluRows

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

function weightAt(sys: System, named: Map<string, Stored>, lora: Map<string, Stored>, name: string, gpu: Option<Gpu>, 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

function loadBlock(sys: System, named: Map<string, Stored>, lora: Map<string, Stored>, prefix: string, gpu: Option<Gpu>, 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

function rawAt(sys: System, s: Stored, gpu: Option<Gpu>, 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

function runsOf(map: Vec<i64>): Vec<i64>

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

function sequence(positions: Vec<f64>, map: Vec<i64>, inverse: Vec<f32>, gpu: Option<Gpu>, 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

function sequenceOf(positions: Vec<f64>, map: Vec<i64>, inverses: Vec<Vec<f32>>, gpu: Option<Gpu>, 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

function stagedBytes(b: Block): i64

The bytes of b's staged copy (Staged.bytes) without making it.

Effects: fail.

stage

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

function silu(x: Tensor): Tensor

Effects: fail, read, write, console, io.

sinusoids

function sinusoids(times: Vec<f64>, 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

function workOf(gpu: Gpu, b: Block, rows: i64, a: Arch): Work

Effects: fail, read, write, io.

linearInto

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

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

function quantizedOf(work: Work, x: Tensor, b: Block, ready: bool): Vec<Tensor>

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

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

function loraOf(work: Work, x: Tensor, b: Block, i: i64, qs: Vec<Tensor>): 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

function gatedInto(work: Work, x: Tensor, y: Tensor, b: Block, i: i64, w: Tensor, gate: Tensor, runs: Vec<i64>, 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

function int8Cores(gpu: Option<Gpu>): bool

Whether gpu is a GPU with INT8 matrix cores (K.hasI8Mma): there W4 and INT8 blocks run as stored.

Effects: fail.

normedInto

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

function runsIn(runs: Vec<i64>, r0: i64, r1: i64): Vec<i64>

The runs (Sequence.runs) of rows [r0, r1), from r0.

Effects: fail.

prepOf

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

function comfyInt8Rows(sys: System, named: Map<string, Stored>, names: Vec<string>, gpu: Gpu): Vec<Tensor> 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

function canonicalInto(w: Tensor, from: Vec<i64>, 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

function blockOf(w: Vec<Tensor>): Block

A block of weights w (Block.w): no biases, codes or LoRA.

Methods

Arch.width

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

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

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

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

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

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

function Block.parallel(self: Block, a: Arch, x: Tensor, m: Modulation, s: Sequence, slot: Option<RefSlot> = 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

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

const AXIS_FREQUENCIES = 16

The rotary frequencies of each of a head's (t, h, w) axes: HEAD / 8.

QKV

const QKV = 0

The index in Block.w of the fused query, key and value projection.

OUT

const OUT = 1

The index in Block.w of the attention's output projection.

UP

const UP = 2

The index in Block.w of the MLP's up projection (gate and up rows in pairs: gluRows).

DOWN

const DOWN = 3

The index in Block.w of the MLP's down projection.

NORM1

const NORM1 = 4

The index in Block.w of the norm weight before the attention.

NORM2

const NORM2 = 5

The index in Block.w of the norm weight before the MLP.

Q_NORM

const Q_NORM = 6

The index in Block.w of the queries' head norm weight.

K_NORM

const K_NORM = 7

The index in Block.w of the keys' head norm weight.

POST_ATTN

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

const POST_MLP = 9

The index in Block.w of a sandwich block's norm weight after the MLP (see POST_ATTN).

CHUNK

const CHUNK = 4096

Rows a chunk of a block's linears and MLP takes: their temporaries (Work) stay within a few hundred MB.