Struct Int8Model 

pub struct Int8Model {
    pub d: usize,
    pub heads: usize,
    pub d_ff: usize,
    pub vocab: usize,
    pub ctx: usize,
    pub requant: BTreeMap<String, (i64, u32)>,
    pub weights: BTreeMap<String, Vec<i64>>,
    pub embed: Vec<i64>,
    pub positions: Vec<i64>,
    pub exp_lut: Vec<i64>,
}

A one-layer int8 transformer: its shapes, its requant pairs, its tables.

Fields

d: usize

heads: usize

d_ff: usize

vocab: usize

ctx: usize

Positions the cache holds — the positional table’s rows.

requant: BTreeMap<String, (i64, u32)>

Each requant’s (M, S), by its name in the model (r_qkv, M_q, M_score, …).

weights: BTreeMap<String, Vec<i64>>

The int8 weights, row-major, by name: Wq Wk Wv Wo (d × d), W1 (d × d_ff), W2 (d_ff × d), Eq (vocab × d, the tied embedding the head multiplies by its transpose).

embed: Vec<i64>

The token embedding as the stream reads it (vocab × d).

positions: Vec<i64>

The positional embedding (ctx × d).

exp_lut: Vec<i64>

The softmax’s exp table, indexed by the requanted distance from the row’s maximum.

Implementations

impl Int8Model

pub fn from_json(j: &Value) -> Result<Int8Model, String>

The model scripts/train_tiny_transformer.py exports (benchmarks/silicon/tiny_transformer.json).

pub fn step_reads() -> BTreeSet<String>

What a decode step’s host reads, besides the caches the step carries: next, the token it feeds back as the next step’s tok. A step lowered returning these (crate::spatial::pipeline::compile_program_to_spatial_returning) forwards every other value only its next statement reads.

pub fn reference_cuda( &self, prompt: &[i64], generated: usize, ) -> Result<String, String>

The hand-written CUDA reference of the decode step: one warp generating every token of the trace for prompt on the device in one launch, the generated tokens past the prompt printed (GEN i token), then the whole trace repeated (its first argument’s count, 2,000 by default) between two events and the time per token printed (US_PER_TOKEN). The GPU’s best for this model: the kernel written for the part by hand.

pub fn prompt_next_tokens( &self, prompt: &[usize], target: &SpatialTarget, ) -> Result<Vec<usize>, String>

The next token, 0-based, each step over prompt predicts on target: the decode step at each position lowered and run on the spatial evaluator, the cache each step writes carried into the next. What a board stepped through the prompt replies, and the token after the prompt’s last is the trainer’s first generated one.

pub fn decode_step_lg( &self, tok: usize, pos: usize, kc: &[i64], vc: &[i64], ) -> String

The decode step for token tok at position pos (both 0-based), the cache holding the key and value rows of positions 0..pos (pos · d values each, row-major). The step binds logits (1 × vocab) and next, the 1-based index of the first maximal logit, and writes this token’s rows into Kc and Vc as row pos + 1.

pub fn decode_step_tp_lg( &self, tok: usize, pos: usize, kc: &[i64], vc: &[i64], tp: usize, ) -> Result<(String, TpManifest), String>

The decode step of Int8Model::decode_step_lg split over tp chips, Megatron’s way, as one program and the manifest of where each of its tensors lives.

Chip c owns heads c·H/tp .. whole: its columns of Wq, Wk, Wv and of the KV cache, its heads’ rows of Wo, an exp table of its own, and FFN columns c·d_ff/tp .. of W1 with their rows of W2. Every chip-local table is named <table>_c<chip> and is a TpSlice of the model’s. A chip’s output projection is a partial accumulator (o_a_c<chip>, f_a_c<chip>), and an all-reduce, spelled as the fold of plus it means, sums the partials before the requant, so the step is the unsplit step bit for bit. The embedding, the residuals and the head are the replicated spine. On one chip the accumulators are whole and nothing is all-reduced.

Trait Implementations

impl Clone for Int8Model

fn clone(&self) -> Int8Model

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for Int8Model

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.