Internals

Llama 3.x frequency-dependent RoPE scaling (rope_type = "llama3").

Shards of a model.safetensors.index.json weight map holding only skipped tensors.

Top-level files needed for inference: weights, configs, tokenizer, chat template (which newer repos ship as chat_template.jinja). Everything else (pytorch_model.bin, onnx/, gguf, README, …) is skipped.

Checkpoint

Lazily opened, memory-mapped tensors of every *.safetensors file in a model directory, by name. Nothing is read from disk until a tensor is uploaded.

Tensor name in file layout: column-major with the PyTorch shape reversed.

The weights of the decoder layer whose tensors are named prefix * "self_attn.q_proj.weight", …

load_weights(T, config, dir; quantize = nothing) -> ModelWeights{T}

Read the checkpoint in dir onto the GPU with element type T, fusing q/k/v and gate/up. With quantize = :int8, the linear weights and the lm_head are stored as Int8Weight (the embedding table stays in T).

The activation's scalars as transformers' XIELUActivation computes them in dtype T.

Quantize a freshly uploaded weight and free its T copy at once, so that loading a model that only fits quantized never holds more than one T matrix at a time.

Copies x into consecutive elements of dst starting at linear index offset.

Upload tensor name as a device array of element type T, in file layout. With cols, only its first cols columns (rows of the PyTorch tensor), which are contiguous in the file.

Upload several 2-D tensors as one matrix whose columns are the tensors' columns side by side. In file layout a PyTorch (out, in) weight is (in, out), so this concatenates along the output dimension: the fused q|k|v and gate|up projections.

Tensor name uploaded in file layout and stored as an Int8Weight (or as a T matrix when quantize === nothing), with only its first cols columns if given. Quantization is per column, so it is done block by block: the T copy on the device never exceeds QUANTIZE_COLUMNS columns (a 131k-token lm_head is 1 GB in BFloat16).

LayerWeights{T,M,V}

The weights of one decoder layer: its two RMSNorm weights, the fused wqkv and w_gate_up projections, wo and w_down, the optional Qwen3 q/k norms and the optional xIELU scalars (Apertus). Linear weights are stored (in, out); see the field comments for shapes.

ModelWeights{T,E,M,V}

All weights of a decoder on the device: the embedding table (hidden, vocab), the LayerWeights of each layer, the final norm and the lm_head. Made by load_weights.

The distinct device arrays and Int8 weights of w (a tied lm_head is listed once).

:int8 if the model's linear weights are quantized, else nothing.

Total bytes of the distinct device arrays in w (a tied lm_head is counted once).

Number of weight elements (parameters) in w; an Int8 weight's scales are not counted.

Per-position cos/sin tables (head_dim, max_len) on the device; column p+1 is position p. cos_sin holds both, interleaved as (head_dim, 2 max_len): columns 2p+1 and 2p+2 are the cos and sin of position p, so a fused kernel needs one array argument for them.

transformers' _compute_llama3_parameters: long wavelengths are divided by factor, short ones kept, and the band in between interpolated smoothly.

StepBuffers(model, cache, n; all_logits = false, aux_layers = Int[])

Device memory for one forward step of n tokens: the step's inputs (token ids, positions, cache slots and a block table long enough for every page of the pool), every intermediate activation, and the logits of the last token, or of all tokens with all_logits. With aux_layers, the step also keeps the residual stream after each of those layers (1-based), stacked in aux, for a DFlash drafter (see DFlashProposer). A step that only writes into these buffers allocates nothing, so on a backend that supports_graphs it can be captured in a CUDA graph and replayed (see ModelRunner).

StepBuffers{T}(config, backend, cache, n; all_logits = false, aux_layers = Int[], vocab = config.vocab_size)

StepBuffers from a configuration, for drafters that are not a DecoderLM; vocab sizes the logits.

prepare_step!(buffers, model, cache, seq, ids)
prepare_step!(buffers, model, cache, seq, nothing)

Reserve cache pages for the 0-based token ids after the positions of seq, and upload the step's inputs into buffers with one asynchronous copy, ordered before the step on the stream. With nothing, the ids are already on the device (written by a previous step's sampling). Host side only; seq.tokens is not changed.

run_step!(model, cache, buffers) -> logits

The forward pass of a prepared step, entirely on the device and only writing into buffers and the KV cache. Returns buffers.logits.

Keep the residual stream after layer l if it is one of the step's aux layers.

BlockManager(num_pages, page_size)

Hands out the pages of a pool of num_pages pages of page_size positions each.

Number of pages not held by any sequence.

reserve!(blocks, seq, n)

Make sure seq has pages for n more positions; throws an OutOfPagesError when the pool runs out.

slots(b, seq, start, n) -> Vector{Int32}

Cache slots of positions start:start+n-1 of seq (pages must be reserved). Slot (page - 1) * page_size + offset with a 1-based offset identifies one position of the pool; it is what store_kv! consumes.

Positions per KV-cache page unless the cache is built with another page_size.

Bytes of cache memory per cached position (keys and values of every layer).

Total pages of cache (free or held by a sequence).

Positions per page of cache.

:int8 for an INT8 cache, nothing for one in the activations' type.

Added tokens to split out of a text, matched leftmost-longest as Hugging Face tokenizers does. Indexed by first character and length rather than compiled into one regex: Apertus 1.5 has 136k of them (image and audio codes).

Repeatedly merges the adjacent pair with the lowest rank, all of its occurrences at once.

Calls f(piece, nothing) for each run of text between added tokens and f(token, id) for each added token, in order.

Only a fixed prefix of special tokens is supported (e.g. Llama 3's BOS); Qwen's ByteLevel post-processor adds nothing.

Steps of up to this many tokens are captured in CUDA graphs.

The static buffers of a step of one token count, and its CUDA graph once captured.

State of generate_pipelined!: two step buffers, their graphs and the host-mapped tokens.

enable_aux!(runner, layers)

Make every later step keep the residual stream after layers (1-based) in its buffers' aux, for a DFlash drafter; the latest step's buffers are runner.last_step. Drops the captured steps when the layers change.

generate_pipelined!(runner, seq, ids; max_new_tokens, stop_ids, on_token) -> Vector{Int}

generate! with the runner's sampler (set it with reset_sampler!), pipelined so that the host runs one step behind the GPU. Needs a runner with CUDA graphs.

Number of captured CUDA graphs (forward and verify steps).

TokenSampler(vocab)

Device state for sample_token!: the sampling parameters, the seed and the counter of tokens sampled so far, and the backend's scratch. Set with reset_sampler! before a generation; each sampled token advances the counter on the device, so a captured step samples a fresh token on every replay.

accept_host(logits, p, draft, q_ids, q_probs, ua, u) -> Int

The 0-based token emitted at a position drafted as draft (0-based), drawn from a distribution q (0-based q_ids with probabilities q_probs; other ids have probability 0): draft when ua * q(draft) < p(draft), else a token sampled by u from max(p - q, 0), renormalized over p's survivors best first. The reference for the device verifier.

The uniform number that decides whether the counter-th token accepts its draft.

The seed a drafter samples its drafts with, for a generation sampled under seed.

filtered_dist(logits, p) -> (ids, probs)

The distribution sample_host samples from: the 1-based ids of the tokens that survive temperature, then top-k, then top-p (nucleus) over the renormalized survivors, as in transformers' default logits-processor order, best first, and their probabilities. Greedy decoding keeps only the argmax.

lowbias32 (Chris Wellons): a bijective mix of 32 bits with low avalanche bias. (cuTile shifts need a shift count of the operand's type.)

Whether tokens can be sampled with p on the device: greedy, or top-k with k ≤ 64.

Inverse transform sampling over weights w (best first) by u ∈ [0, 1): the first entry whose cumulative weight reaches u × the total, among those with positive weight (the last of them on round-off).

reset_sampler!(s, p, seed)

Sample with p (which must be on_device) from counter 0 under seed. Greedy decoding samples the top-1 token.

sample_host(logits, p, u) -> Int

The 0-based token sampled from logits with p by the uniform number u ∈ [0, 1): the reference for the device samplers. Inverse transform sampling over filtered_dist, best first.

sampling_uniform(seed, counter) -> Float32

The uniform number in [0, 1) that decides the counter-th token sampled under seed. A pure function of its arguments, evaluated identically on the host and in kernels.

set_counter!(s, counter)

Sample the next token with counter, as when rewinding to a token that was sampled ahead. Ordered on the stream before the steps that follow.

The ids of the k largest logits, best first, ties to the smaller id.

Length of the longest common prefix of a and b.

Verify steps take up to this many tokens (a DFlash block); all are captured in CUDA graphs.

DraftDists(k)

The distributions q that up to k sampled drafts were drawn from: column j holds the 0-based ids of draft j's candidates (-1 past them) and their probabilities, as sample_token!(…, q) writes them.

A step that returns the token the target model samples after each of its n inputs: the forward step with the logits of every token, then the runner's sampler on each column (counters c, c + 1, …), all captured together in a CUDA graph. With the distributions of sampled drafts, the columns of drafted positions apply the p/q rule instead. The tokens land in host-mapped memory.

attach!(proposer, runner)

Prepare proposer to draft for the target model of runner, before a generation.

Draft j's distribution, as the q argument of sample_token! and accept_token!.

draft_dists(proposer, runner) -> Union{DraftDists,Nothing}

The distributions of the drafts proposer samples for the target of runner, or nothing when its drafts are a function of the context (then verification compares them with the target's own samples).

drafted(proposer, k) -> Vector{Int}

The k drafts of the latest propose_on_device!, after the stream has run them.

generate_speculative!(runner, seq, ids, proposer; max_new_tokens, stop_ids, on_token,
                      stats = SpecStats()) -> Vector{Int}

generate! with the runner's sampler (set it with reset_sampler!) and drafts from proposer checked by verify!: the same tokens as plain decoding with that sampler up to near-ties (a verify step of k + 1 tokens runs other kernels than a decode step, so the logits differ in the last bits), and seq ends in the same state. A proposer that samples its drafts (draft_dists) yields tokens with the same distribution instead, not the same tokens. Counts steps and accepted drafts in stats.

Where the sample of chain step i (or of the forward step before the chain, i = 0) goes. A proposer that can sample always records q (one-hot when greedy), so that its captured chain serves greedy and sampled generations alike.

propose_on_device!(proposer, tokens, k) -> Union{Nothing,CuVector{Int32}}

Enqueue drafting k tokens to follow tokens without waiting for them, and return the device vector they land in (0-based ids), for verify! to read on the device; the host reads them with drafted once the stream has run that far. nothing when the proposer cannot (then propose runs instead).

Run a proposer's step: replay graph, or run it eagerly and capture it when the proposer uses graphs. Returns the graph.

Whether a proposer that can sample its drafts does so for runner: only when the target samples (with greedy decoding, both draft and target take the argmax).

verify!(runner, seq, ids, counter; dists = nothing, drafts = nothing) -> Vector{Int}

Run the 0-based token ids (the last sampled token, then the drafts) through the model, appending them to seq, and return the token the runner's sampler picks after each of them, with counters counter, counter + 1, … (the argmax when greedy). When the drafts were sampled from the distributions dists (column j for draft j), the token after each input but the last follows the p/q rule instead (accept_token!): the next draft if accepted, else a different token. Steps of each length, up to 16 tokens, are captured in a CUDA graph like forward! steps. With drafts, a device vector (see propose_on_device!), the drafts are read from it on the device rather than from ids[2:end], which then only hold places (and go into seq.tokens as they are).

Buffers of one drafter step of c context positions and a block: the step's inputs and per-column activations for n = c + B columns (b.h holds the layers' input: the context states, then the normed block), and the target's stacked states of the context.

The target layers whose outputs the drafter reads, counted from 1. DFlash's target_layer_ids index transformers' hidden_states minus one, i.e. number layer outputs from 0; without them, they are spread evenly from layer 1 to L - 3.

The drafter on a prepared step whose aux holds the target's context states; leaves the logits of block positions 2:B in p.logits.

Enqueue the drafting of the block after context into p.sink and p.mirror (and its logits into p.logits), where the target's latest step ran positions start:length(context)-2 and left their states in the first columns of aux. Does not wait for the drafts.

The drafts stay on the device, for the verify step to copy into its inputs: the host enqueues the drafter's step and the target's verify step back to back and waits once.

has_kernel(backend, op) -> Bool

Whether backend implements op itself, rather than delegating to another backend.

Keys (or values) of positions 1:len of one KV head, gathered from the sequence's pages into a dense (head_dim, len) matrix.

The weight quantize_int8 stands for, in element type T.

Dequantized and unsmoothed in Float32 (not rounded to T): the kernels scale Float32 scores and probabilities.

int8_kv_pools(head_dim, page_size, kv_heads, num_pages; key_smoothing = nothing) -> (keys, values)

The key and value Int8Pages of one layer, zeroed; the keys smoothed by key_smoothing (see key_smoothing), or not.

key_smoothing(k_norm) -> CuVector{Float32}

The factors that keys normalized by the weight k_norm are divided by in an INT8 cache: max(|w_i|, |w_j|) for each rotary pair (i, j = i + head_dim/2), so the factor is the same on both channels a rotation mixes (1 where both are zero).

The layer arrays behind a key/value pair, for the kernels: Int8 values as (head_dim, 2 page_size kv_heads num_pages), keys' columns first, the scales as one vector, and the keys' smoothing.

Quantization of the columns of x (Float32): q = floor(x 127 / max|x| + 1/2) and s = max|x| / 127, as the kernels compute them (0 for an all-zero column).

Residual add + RMSNorm, one block per column: pass 1 sums the squares of x + delta, pass 2 recomputes it, stores it into x and the normalized value into H. Each tile is read in pass 2 before it is overwritten, and pass 1 never sees stored values.

Embedding lookup, one block per token, walking the column in tiles.

expm1 without exp(x) - 1's cancellation (Kahan): with u = exp(x) rounded, (u - 1) x / log(u) is accurate to a few ulps, since u's rounding error cancels between the two factors. Below -20, expm1 is -1 in Float32 (and log(u) underflows further down).

Rotary embedding, one block per token over all heads. The two halves of head_dim are loaded as separate tiles (tile index 1 and 2 along the first axis), which is exactly the "rotate half" pairing: out1 = x1 cos1 - x2 sin1, out2 = x2 cos2 + x1 sin2. The cos/sin columns are looked up from the token's position, so positions can live on the device.

SwiGLU: out = silu(gate) * up, with gate and up as separate (strided) views of the stacked gate_up activations. Grid: (tiles along intermediate, tokens).

Quantize the vector whose two halves are the Float32 (HALF, 1) tiles x1 and x2 into column col of KV and its scale into S[col], as quantize_columns does.

KV-cache write, one block per token. Arrays are 2-D views, which keeps the launch under 32 kernel parameters: the pools as (head_dim, page_size × kv_heads × pages), whose column for (offset, head h, page), all 1-based, is offset + P (h - 1) + P H (page - 1), and k/v as (head_dim, kv_heads × n). Tile indices along the size-1 axis are element indices.

xIELU, grid (tiles along the rows, tokens). The scalars are per layer and fixed, so launch arguments are fine inside a CUDA graph.

xIELU of a Float32 tile, rounding to T after each operation as the reference does.

INT8 GEMM tiles. The kernel converts each weight tile to T in registers, which BF16's tiles leave no room for (255 registers and spills at 128 × 128), and its loads of the Int8 tile are not overlapped like the BF16 operands', so its blocks are latency-bound: it wants two blocks per SM on every grid, narrower tiles at 512 tokens, and at ≤ 64 tokens (bandwidth-bound) deeper K steps that read as many bytes per step as BF16's. From a sweep of 15 tiles × occupancy 1, 2 on the layers of Qwen3-0.6B, Qwen3-4B and Apertus 8B at n = 8, 64, 512: within 1.5 % of the best tile in total (worst 1.11×), 1.23× the BF16 rule (down at 64 tokens: 2.1× on 4B, 1.7× on 8B).

Two GEMM blocks per SM when the grid has enough of them: the second hides the first's latency (4B gate_up at n = 512: 1.11×), but tileiras trims registers and shared memory to fit both, which slows grids that leave one block per SM (0.6B o/down at n = 512: 0.86×).

Int8 GEMV tiles of half the BF16 width: TK × TM ≈ 4096 elements was fastest on Apertus 8B's layers (TK = 1024, TM = 4: 600–670 GB/s; TM = 8 was 7–13 % slower).

Scaled, rounded GEMV outputs of column tile col, as Float32 (1, TM).

Float32 sums of TM weight columns against x, from column tile col: the GEMV loop.

GEMV with a SwiGLU epilogue for the stacked gate_up weight (in, 2I): block b reduces TM gate columns and the TM up columns I apart, and writes silu(gate) * up, rounding as silu_mul_kernel does after a plain GEMV. Needs I divisible by TM.

GEMV with an xIELU epilogue (Apertus' up projection): the plain GEMV, then xielu_tile on the rounded sums, as xielu_kernel computes after it.

1-based (block row, block column) of block bid, walking groups of GROUP block rows column by column so that consecutive blocks share W and X tiles in L2.

The same into an INT8 cache: KV, S and C are the layer's Int8 pools, scales and the keys' smoothing (store_int8!).

Grid (n, kv_heads): key head h (normalized, rotated) and value head h into the pools, viewed as (head_dim, page_size × kv_heads × pages) as in store_kv_kernel.

Head col of qkv, RMS-normalized with W when NORM (as rmsnorm_kernel: round the normalized value to T, scale, round again) and rotated by the angles of position pos, as two (HALF, 1) halves (as rope_kernel).

Grid (n, heads): query heads into Q (head_dim, heads × n).

An INT8 cache: the keys and values are converted to T for the tensor cores, the key's scale multiplies its row of scores (s_k q_kᵀ x) and the value's scale its probability (V diag(s_v) p). Converting whole (64, 128) pages of keys and values takes 255 registers and spills, so the kernels take 32 keys at a time, and K, V, the queries and the accumulator in two halves of the head dimension: q = (q1, q2), acc = (acc1, acc2).

One key tile of the online softmax for the query tile q (D, M) at positions qpos (1, M): its KEYS keys, at positions first_key + 1 : first_key + KEYS, are tile column kcol of K (width KEYS) and their values tile column vcol of V. S holds the scales of an INT8 cache (Int8Pages), at the same tile index as their vectors; it is nothing otherwise.

Grid (GROUP n, kv_heads): merges the splits of one packed query column.

Grid (kv_heads, SPLITS). Column c of the packed query tile is head (h - 1) GROUP + g of token t, with c = g + GROUP (t - 1); NT ≥ n is the padded token count (decode_tokens). The partial of split s goes to column (h - 1) SPLITS + s of the workspace: acc in rows 1:D, m in row D + 1 and l in row D + 2.

The kernels take the pools as K and V with S = nothing, or an INT8 layer as KV (keys, then values from VALUES key tiles on) with its scales S and the keys' smoothing C, which multiplies the queries (attend_operands).

Padded token count of the packed decode query tile: at least DECODE_MIN_ROWS rows.

decode_workspace(q, k_pages) -> CuArray{Float32,3}

Split-KV partials for decode attention of the queries q (head_dim, heads, n): for each packed query row, split and KV head, the unnormalized output (rows 1:head_dim), running max and running sum (rows head_dim + 1 and + 2).

Keys per step of the KV loop: a page, or half of one for an INT8 cache.

The 2-D views the kernels take (no copies).

The online-softmax update for the scores qk (KEYS, M) of keys at positions first_key + 1 : first_key + KEYS, for queries at positions qpos (1, M): the probabilities p (Float32), the factor alpha that rescales the accumulator, and the new running max and sum.

The (value, position) pair that is larger, ties to the smaller position: argmax without its NaN handling (logits are never NaN, padding is -Inf), which shortens the reduction.

The token emitted at a position drafted as Draft[1] from the distribution (QIds, QProbs), as accept_host: the draft with probability min(1, p/q), else a sample of max(p - q, 0).

The id at the first slot whose cumulative weight reaches u × the total, among the slots with positive weight (the last of them on round-off), as pick_host.

Scratch for sample_tokens! over cols columns: each column's candidates, zero-padded.

Scratch for sample_token!: the top-k candidate ids of every vocabulary chunk.

The survivors among the chunks' candidates Ids: their 1-based ids (best first, 0 past top_k), their probabilities under temperature and top-k, zeroed outside the top-p prefix, and the mask of survivors. With several columns of logits (sample_tokens!), Ids and Logits are the flat arrays and col the column: its candidates are entries (col - 1) N + 1 : col N of Ids, and its logits start at (col - 1) V + 1.

The 1-based positions of the k (≤ K) largest entries of the N-tile x, best first, ties to the smaller position; 0 past k.

The candidates of every chunk of the vocabulary, into the sampler's scratch.

topk_sample_kernel that also writes the distribution it samples from: the survivors' 0-based ids (-1 past them) and their probabilities.

Streams tokens to io and records when the first two arrived.

chat_loop(session, input, output; max_new_tokens)

Read user messages line by line from input and stream replies to output until EOF or /exit. /reset starts a new conversation. Statistics go to stderr.

Model, tokenizer and generation settings shared by run and chat.

The proposer of –speculate / –draft / –dflash, or nothing.

Splits args into positionals and --name value options; names in flags take no value and are set to "true".

In a fresh process, time to first token includes compiling the prefill kernels, and the first decode step compiles the single-token ones. Decode speed is measured from the second token on so that it excludes both.