The inference core

One crate, three targets

curta-core is one Rust crate with no dependencies. It holds the weight format, the tokenizer, the resumable forward pass, the sampler, the turn planner and the readers for the program’s accounts. The same source builds for three targets:

# from backend/
cargo build --manifest-path crates/curta-core/Cargo.toml --features std --release
cargo rustc --manifest-path crates/curta-core/Cargo.toml --features std --target wasm32-unknown-unknown --release --lib -- --crate-type cdylib

The crate’s modules:

formatmodule

The artifact header, Config, Metadata and the table lookups. The weight format is its layout.

tokenizermodule

Lookup, encoding and decoding over the tokenizer stored in section zero.

inferencemodule

The state layout, init_state, turn and the arithmetic kernels.

plannermodule

Phase codes, the cursor, the compute charges and plan_turn.

protocolmodule

Sizes, tags and readers for the Config, Model, Run and weight accounts.

convertmodulestd only

The checkpoint converter, the only code in the crate that uses floating point.

Fixed-point arithmetic

int8 weights, fixed-point activations, lookup tables. Bit-exact across targets. Every value in the forward pass is an integer with a fixed count of fractional bits:

A matrix row

The step the forward pass repeats most is a row of int8 weights times a vector of activations. The weights are stored as unsigned bytes, each the weight plus 128, so the kernel multiplies unsigned bytes by 64-bit words with no sign extension, eight at a time, and corrects once per vector:

sum  = x[0] + x[1] + … + x[columns − 1]                 // once per input vector
acc  = byte[0] × x[0] + … + byte[columns − 1] × x[columns − 1]
y    = ((acc − 128 × sum) × scale) >> 24

The products wrap in 64 bits, which gives the same result as the signed dot product. The activation buffer is aligned to eight bytes, and the core checks that alignment before a turn runs.

RMS normalization

rms   = isqrt(sum(x[i] × x[i]) / dimension + 42,950)      // integer square root, at least 1
y[i]  = ((x[i] × weight[i]) / rms) >> 8

The constant 42,950 is the epsilon 0.00001 with 32 fractional bits. Division rounds toward zero.

The attention cache

Keys and values are kept for every position a run has passed, at the same 16 fractional bits as the activations, in 32-bit words. Every store into the cache, including a key after rotary encoding, checks that the value fits 32 bits. Attention reads the cached words back sign-extended to 64 bits.

Lookup tables

The two nonlinear functions of the forward pass are tables in section zero of the weights, each of 4,097 entries, computed at conversion.

exponentiale^(−d)softmax

The input d is a difference of two scores with 16 fractional bits, zero or more. The index is d >> 8, capped at 4,096, so the table steps by 1/256 and covers differences from 0 to 16. The result has 30 fractional bits.

SiLUx × sigmoid(x)feed-forward

For an input from -16 to 16 the index is (x + 16) >> 9 in fixed point, a step of 1/128. An input of 16 or more returns itself, and an input of -16 or less returns zero.

The forward pass, in order

A run moves through numbered phases. The run’s state names the phase it is in, and a turn continues from there. Phases 0 to 2 tokenize the prompt and run once. Phases 3 to 16 run once per position, with 4 to 16 repeated for each layer. Phases 17 to 19 run for positions that generate a token.

A prompt position before the last goes from phase 16 straight to the next position’s embedding: its keys and values are in the cache, and the classifier and the sampler run only where a token is generated.

Attention, stage by stage

Phase 9 handles one head at a time, and each head goes through four stages. With p positions so far, the current one included:

1. Scores. For each position, the dot product of the head’s query with that position’s cached key, shifted right 16 bits, times one over the square root of the head size, shifted right 24 bits. The largest score is kept.
2. Exponentials. Each score becomes the exponential table’s value for the largest score minus that score. Their sum is kept.
3. Normalization. Each exponential is shifted left 24 bits and divided by the sum, which makes it a probability with 24 fractional bits.
4. Weighted values. For each element of the head’s output, the sum over positions of probability times cached value, shifted right 24 bits.

With fewer key and value heads than heads, head h reads the keys and values of key and value head h / (heads / kvHeads).

The resumable state

Everything a run has computed lives in one buffer, the core state. On chain it is the run account after its 256-byte header. A turn reads its position from the state, works, and writes its position back, so the pass can stop between any two chunks of work and continue in a later transaction on any target.

The state opens with a 128-byte header tagged CURTAST2. All offsets count from the start of the state.

After the header come the run’s buffers, each starting at a multiple of eight bytes:

Cache positions is the count of positions the run can reach: the prompt’s bytes plus max tokens plus one, up to the context length. state_size computes the whole size from the configuration, the prompt’s length in bytes and max tokens, and output_tokens and output_text read the output out of a state.

On the 16-byte prompt “Once upon a time”, a 12-token run on stories260K has 29 cache positions and a run account of 48,304 bytes, header included. An 8-token run on stories15M has 25 positions and 624,832 bytes. With an empty prompt and max tokens equal to the context, the accounts are 736,328 and 3,852,680 bytes.

The turn planner

The planner decides how far one turn goes. It divides the forward pass by compute units, from the configuration and the state alone, so a turn takes no input from its sender and every target cuts the pass at the same places.

The unit of work is a chunk: a run of rows of one matrix, a run of positions or elements of one attention stage, a run of vocabulary entries of one sampling pass, or one whole vector operation. The planner charges each chunk a number of compute units and adds chunks until one of three things happens:

• the next chunk does not fit what is left of the turn’s 1,350,000 compute units
• the next chunk needs a third matrix section
• the run is complete

A turn starts with 18,000 units charged for the instruction itself: accounts, payment and the event. Every chunk adds 160 units for reading the cursor, dispatching and writing the cursor back. A transaction may use 1,400,000 compute units, and the planner’s 1,350,000 sits inside that.

A turn runs across phase, layer and position boundaries: it finishes a token, samples it, and starts the next one when units remain. Tokenizing the prompt reads section zero alone, and the turn that finishes it stops at the first embedding.

sections[u32; 2]

The matrix sections the turn reads, in the order it reaches them. Section zero is always read and is left out of this list.

section_countusize0 to 2

How many entries of sections are used. Zero for a turn that tokenizes.

estimated_cuu32

The compute units the planner charged for the turn.

max_generatedu32

The most tokens the turn can complete. The program reserves the royalty for this many before the turn runs.

endCursor

Where the state’s cursor stands after the turn.

plan_turn returns this plan for a state. The program plans each turn to check the weight accounts the transaction names and to price the turn, and the SDK plans the same turn to choose those accounts. quote_turns tokenizes a prompt and then steps the planner to max tokens, which gives the count of turns a run takes. A run that produces token 1 or 2 early ends in fewer.

Compute costs

These are the planner’s charges, in compute units, before the 160 units per chunk. d is the dimension, h the hidden dimension and n the count of columns of a matrix.

Matrices

The row kernel runs 37 instructions for every eight weights. A chunk takes as many rows as fit, so a large matrix spans several turns and a small one shares a turn with its neighbours.

Vector operations

Attention

Each stage of a head is charged per unit, with q the head size and p the positions so far. A stage can stop after any unit and continue in the next turn.

Sampling

Tokenizing

A lookup is a binary search over the sorted token ids. A character’s lookup uses its own byte length and a merge’s lookup uses the joined length of its two tokens.

One position, added up

The planner’s charges for one generating position at the end of the context, at temperature zero: position 511 of stories260K, which takes 5 turns, and position 255 of stories15M, which takes 62. Each row sums every layer and every chunk.

The classifier of stories15M is 32,000 rows of 288 weights, over nine million multiply-accumulates for each generated token.

Sampling

The core samples each token itself, from the logits in the state, the run’s temperature and the run’s seed. Sampling is phase 19, and it makes one pass over the vocabulary at temperature zero and three passes above zero. Each pass is resumable at any vocabulary entry.

1. Maximum. Finds the largest logit and its index. At temperature zero that index is the token.
2. Sum. For each logit, takes the largest logit minus it, divides by the temperature, reads the exponential table, and adds the result to a sum.
3. Draw. Advances the generator once and reduces its output modulo the sum. Walks the vocabulary again, subtracting each entry’s exponential, and takes the first entry the remainder falls inside.

The generator is xorshift64*, with the run’s seed as its first state:

x ^= x >> 12
x ^= x << 25
x ^= x >> 27                                  // x is stored back as the generator's state
draw = (x × 0x2545F4914F6CDD1D mod 2^64) mod sum

A seed of zero starts the generator at 0x9E3779B97F4A7C15. The draw runs over the whole vocabulary, weighted by the exponentials.

The chosen token’s id is appended to the output tokens. Its text is the token’s bytes from the tokenizer, with two rules from llama2.c: a leading space is dropped right after the opening token, and a token written <0xXX> becomes that one byte. Tokens 1 and 2 add no text and complete the run. The run is also complete when the count of generated tokens reaches max tokens.

The Rust API

state_size(config, prompt_len, max_tokens)Result<usize, Error>

The size of the state for a prompt of prompt_len bytes, up to 1,024, and max_tokens from 1 to the context length.

init_state(config, state, prompt, max_tokens, temperature_q16, seed)Result<(), Error>

Writes a fresh header and the prompt into a buffer of exactly state_size bytes. The prompt is UTF-8 with no zero byte. The temperature is at most 16, which is 1,048,576 with 16 fractional bits.

planner::plan_turn(config, state)Result<TurnPlan, Error>

The plan of the next turn.

turn(config, state, metadata, sections)Result<TurnResult, Error>

Runs one turn. metadata is section zero, and sections are the matrix sections the plan names, each a WeightSection with its index and bytes.

output_tokens(config, state)&[u8]

The generated token ids, four little-endian bytes each.

output_text(config, state)&[u8]

The generated text, UTF-8.

StateInfo::read(state)Result<StateInfo, Error>

The header’s cursor: phase, layer, row, position, prompt_tokens, generated, text_len, turns and complete.

inference::quote_turns(config, metadata, prompt, max, temperature_q16, scratch)Result<u64, Error>

The count of turns a run takes, computed in a scratch buffer of state_size bytes.

turn returns a TurnResult: complete, generated (tokens so far), generated_this_turn, token (the last token this turn sampled, if any) and work (the compute units the planner charged).

use curta_core::{
    format::SECTION_TABLE, init_state, output_text, planner, read_u64, state_size, turn,
    Metadata, StateInfo, WeightSection,
};

fn main() {
    let artifact = std::fs::read("stories260K.curta").unwrap();
    // Section zero's length is the second field of the first section table entry.
    let zero = read_u64(&artifact, SECTION_TABLE + 8) as usize;
    let m = Metadata::new(&artifact[..zero]).unwrap();

    let prompt = b"Once upon a time";
    let mut state = vec![0u8; state_size(m.config, prompt.len(), 32).unwrap()];
    init_state(m.config, &mut state, prompt, 32, 0, 42).unwrap();

    while !StateInfo::read(&state).unwrap().complete {
        let plan = planner::plan_turn(m.config, &state).unwrap();
        let mut sections = [WeightSection { index: 0, bytes: &[] }; 2];
        for i in 0..plan.section_count {
            let index = plan.sections[i] as usize;
            let offset = m.section_offset(index);
            sections[i] = WeightSection {
                index: index as u32,
                bytes: &artifact[offset..offset + m.section_len(index)],
            };
        }
        turn(m.config, &mut state, m.bytes, &sections[..plan.section_count]).unwrap();
    }
    println!("{}", String::from_utf8_lossy(output_text(m.config, &state)));
}

WebAssembly exports

The WebAssembly build exports plain functions over the module’s linear memory. Pointers and lengths are u32. The configuration is passed as a pointer to 28 bytes, the seven u32 values in order. @curta/sdk wraps these as CoreWasm and CoreSession, loaded with loadCore().

alloc(length)→ pointer

Reserves length zeroed bytes in the module’s memory.

dealloc(pointer, length)

Releases bytes from alloc or convert.

state_size(config, promptLength, maxTokens)→ u32

The state’s size. 0xFFFFFFFF for parameters out of range.

init_state(config, state, stateLength, prompt, promptLength, maxTokens, temperatureQ16, seedLow, seedHigh)→ error

Starts a state. The seed is passed as its low and high 32 bits.

plan_turn(config, state, stateLength, output)→ error

Writes 20 bytes at output: section count, first section, second section, estimated compute units and the most tokens the turn can complete, each a u32.

turn_with_sections(config, state, stateLength, metadata, metadataLength, first, firstLength, second, secondLength)→ error

Runs one turn with the sections its plan names. A slot the plan leaves out is passed empty.

quote(config, metadata, metadataLength, prompt, promptLength, maxTokens, temperatureQ16)→ u64

The count of turns a run takes.

convert(checkpoint, checkpointLength, tokenizer, tokenizerLength)→ pointer

Converts a checkpoint and returns the artifact’s address, or zero. result_len() returns its length.

required_section(config, state, stateLength)→ u32

The first section the next turn reads, zero for a turn that reads section zero alone.

turn(config, state, stateLength, metadata, metadataLength, weights, weightsLength)→ error

Runs a turn whose plan names one section.

An error result of zero is success, and any other value is a core error’s number.

Core errors

The core’s functions return one of eight errors. The WebAssembly exports return its number, and the SDK reports it as Curta core error N. The program reports a core error inside a turn as CoreError, listed with the program’s other errors in Errors.

Measured on a validator

The program’s turns on the reference runs, on the prompt “Once upon a time” at temperature zero, from the receipts of their transactions:

Compute units per multiply-accumulate is the slope between two simulated turns that differ only in the count of matrix rows. Both runs produce the same tokens and the same text bytes on the program, the native build and the WebAssembly build.