Public API

TileInfer

A single-GPU LLM inference engine for Llama-style models (Llama 3, Mistral, Qwen3, Apertus), with its hot kernels written in cuTile.jl. Built to learn how inference engines work, not for production use.

Start with pull, load_model, load_tokenizer and generate or ChatSession; the tileinfer command line (cli) wraps the same calls.

pull(repo_id; revision = "main") -> snapshot directory

Download the inference files of repo_id into the hub cache and return the directory of the snapshot. Files already in the cache are not downloaded again. Gated repos need an access token in HF_TOKEN or ~/.cache/huggingface/token.

model_dir(model; revision = "main") -> directory

Resolve model (a local directory or a pulled repo id) to the directory holding its files, without network access. A repo counts as pulled once its config.json, which pull downloads last, is in the cache.

List the models in the hub cache that have safetensors weights.

Remove a pulled model (all revisions, or only revision) from the hub cache.

CachedModel

A model in the hub cache: repo id, revision (a ref such as "main", or a commit hash), snapshot files and their total size in bytes.

ModelConfig

Hyperparameters of a Llama-style decoder. Build from a model directory with load_config.

Read config.json from a model directory (see model_dir).

DecoderLM{T}

A model on the GPU: configuration, weights, RoPE tables and the op backend. Load with load_model; run with forward!.

load_model(model; T = BFloat16, backend = RefBackend(), max_len = 4096, quantize = nothing) -> DecoderLM

Load a pulled model (repo id) or a model directory. max_len bounds the positions the RoPE tables cover, i.e. the longest sequence the model can process. quantize = :int8 stores the linear weights as Int8 with per-channel scales (see Int8Weight).

max_length(model) -> Int

The longest sequence model can process: the positions its RoPE tables cover, set by max_len in load_model.

forward!(model, cache, seq, ids; all_logits = false) -> CuMatrix{Float32}

Run the 0-based token ids through the model at the positions following those of seq, storing their keys and values in cache and appending them to seq. Returns the logits (vocab, n) of every token if all_logits, otherwise (vocab, 1) of the last token. Runs eagerly with buffers allocated for this step; ModelRunner reuses them.

forward!(runner, seq, ids) -> CuMatrix{Float32}

forward! through runner: the logits (vocab, 1) of the last token. For short steps they live in the step's static buffers and are overwritten by the next step of the same length, so read them before running another step.

BPETokenizer

Byte-level BPE tokenizer. Load with load_tokenizer; use encode and decode.

Load tokenizer.json from a model directory.

encode(tok, text; add_special_tokens = true) -> Vector{Int}

Token ids (0-based) of text. Added tokens written in the text, such as <|im_end|>, become their ids. add_special_tokens controls the post-processor prefix (BOS).

decode(tok, ids; skip_special_tokens = false) -> String

Inverse of encode. Bytes that do not form valid UTF-8 on their own (a token boundary inside a multi-byte character) are kept as-is.

One turn of a conversation. role is "system", "user" or "assistant".

chat_prompt(model_type, messages; add_generation_prompt = true, enable_thinking = true)

Render messages (a vector of ChatMessage) into the prompt string the model was trained on. add_generation_prompt opens an assistant turn for the model to complete. The prompt includes any BOS token the template writes, so encode it with add_special_tokens = false, as transformers' apply_chat_template does.

Port of Qwen3's template without tool calling. Two details matter:

• assistant turns before the last user message drop their <think>…</think> reasoning, so earlier reasoning does not use up context;
• enable_thinking = false pre-fills an empty think block, which makes the model answer directly.

Port of Apertus 1.5's template (chat_template.jinja) for text messages without tools. A system message may only come first; without one, the template's default is used. The developer turn says whether the model may deliberate (think between <|inner_prefix|> and <|inner_suffix|>). Consecutive assistant messages form one turn. Earlier reasoning is kept.

PagedKVCache{T}(config, num_pages; page_size = 64, quantize = nothing, key_smoothing = nothing)

Key/value page pools for every layer and the allocator for their pages, for a model whose activations are T. The pools hold T, or with quantize = :int8 Int8 values with one Float32 scale per cached vector (Int8Pages): about half the memory, and half the bytes decode attention reads. key_smoothing holds each layer's smoothing of the keys of an INT8 cache (see key_smoothing), or nothing for none.

Sequence

A sequence's place in the paged KV cache: its block table and the token ids whose keys and values are cached, in position order.

A new, empty sequence in cache.

truncate!(blocks_or_cache, seq, n)

Drop the cached positions of seq after the first n and return its unused pages to the pool. Speculative decoding uses it to roll back rejected drafts, and chat to roll back to the prefix a new turn shares with the cache.

release!(blocks_or_cache, seq)

Return all pages of seq to the pool.

OutOfPagesError(needed, free)

Thrown by reserve! when the KV cache has fewer free pages than a step needs. A ChatSession can recover with reset!.

SamplingParams(; temperature = 0, top_k = 0, top_p = 1)

temperature = 0 is greedy decoding. top_k = 0 and top_p = 1 disable those filters.

generate(model, prompt_ids; kwargs...) -> Vector{Int}

One-shot generate! in a fresh cache sized for the prompt and max_new_tokens, quantized as kv_quantize (see PagedKVCache).

generate!(runner, seq, ids; max_new_tokens = 256, sampling = SamplingParams(),
          stop_ids = model.config.eos_token_ids, rng = Random.default_rng(),
          on_token = nothing, pipeline = true, speculative = nothing,
          spec_stats = SpecStats()) -> Vector{Int}
generate!(model, cache, seq, ids; kwargs...)

Continue seq with the 0-based token ids and generate up to max_new_tokens tokens. Returns the new token ids, without the stop token. on_token(id) is called after each new token, for streaming. Afterwards seq holds ids and every generated token except the last one, which was sampled but never run through the model.

The tokens are a function of the logits and one seed drawn from rng (see sample_host), however they are sampled: on the device when the sampling is on_device, pipelined (generate_pipelined!) on a runner with CUDA graphs unless pipeline = false; otherwise on the host, step by step.

With a proposer as speculative (such as NgramProposer), decoding checks its drafts several tokens per step (generate_speculative!) and counts them in spec_stats; the sampling must be on_device.

generation_defaults(model_dir, config) -> (SamplingParams, stop_ids)

The model authors' recommended sampling settings and stop tokens, from generation_config.json when the model ships one; greedy and the config's EOS otherwise.

TextStream(tokenizer, io)

Prints generated tokens to io as they arrive. Bytes of a UTF-8 character split across tokens are held back until the character is complete.

ModelRunner(model, cache; graphs = supports_graphs(model.backend))

Forward steps of model on cache, with CUDA graphs for steps of up to 8 tokens when graphs. Run steps with forward!(runner, seq, ids).

ChatSession(model, tokenizer; sampling, stop_ids, enable_thinking = true,
            max_tokens = max_length(model), rng = Random.default_rng(), speculative = nothing,
            kv_quantize = nothing)

A conversation with model, with its own paged KV cache for up to max_tokens positions, quantized as kv_quantize (see PagedKVCache). Send messages with reply!. speculative is a proposer for speculative decoding (see generate!).

reply!(session, text; max_new_tokens = 1024, on_token = nothing) -> (reply, TurnStats)

Add the user message text, generate the assistant's reply and add it to the conversation. On error (e.g. the cache is full) the conversation is left unchanged.

Clear the conversation and free its cache pages.

Statistics of one chat turn: prompt tokens reused from the cache and newly prefilled.

NgramProposer(; k = 4, max_ngram = 3, min_ngram = 1)

Proposes up to k tokens by prompt lookup: find the latest earlier occurrence of the last n tokens (n from max_ngram down to min_ngram) and propose the tokens that followed it.

DraftProposer(draft_model; k = 4, sample = false, max_tokens = max_length(draft_model),
              graphs = supports_graphs(draft_model.backend))

Proposes the next k tokens that draft_model decodes greedily, on a KV cache of its own for up to max_tokens positions; or, with sample and a target that samples, that it samples with the target's sampling parameters, verified by the p/q rule. The draft model must share the target's vocabulary.

DFlashProposer(drafter; k = drafter.block_size - 1, sample = false, max_tokens = max_length(drafter),
               graphs = supports_graphs(drafter.backend))

Proposes the first k of the block_size - 1 tokens that the DFlash drafter (see load_dflash) drafts in one pass, on a KV cache of its own for up to max_tokens positions: the argmax at each block position, or, when sample and the target samples, a sample with the target's parameters (each position independently: the block's positions do not see each other's drafts, so q of draft j is the drafter's distribution at position j). Reads the hidden states of the target's latest step, which generation enables on the target's runner.

DFlashDrafter

A DFlash drafter on the GPU: its Qwen3 layers, fc and norms, the target's embedding table and lm_head (shared), its RoPE table, the target layers it reads (aux_layers, 1-based), the block size and the mask token. Load with load_dflash.

load_dflash(drafter, target::DecoderLM; max_len = max_length(target)) -> DFlashDrafter

Load a DFlash drafter for target (a pulled repo id or a directory), in the format of z-lab's checkpoints (layers.*, fc, hidden_norm, dflash_config).

load_dflash(drafter, target_config, embed, lm_head, backend; max_len) -> DFlashDrafter

The same, with the target's configuration, embedding table and lm_head given separately (the tests run the drafter in Float32 without the whole target).

propose(proposer, tokens, k) -> Vector{Int}

Up to k draft tokens to follow tokens (possibly none).

SpecStats()

Counters of a speculative generation: verify steps, drafted and accepted tokens.

Fraction of drafted tokens that were accepted (NaN before any draft).

Backend

Selects the implementation of every op by dispatch: RefBackend (CUDA.jl broadcasts and cuBLAS, the ground truth) or TileBackend (cuTile kernels). A model holds its backend, so load_model(…; backend = TileBackend()) runs it on the cuTile kernels.

Reference ops: CUDA.jl broadcasts and cuBLAS. Ground truth for correctness tests.

cuTile.jl kernels.

embed(backend, table, ids) -> (hidden, tokens)

Columns of table (hidden, vocab) for the 0-based token ids (a device vector).

rmsnorm(backend, x, w, eps)

RMS-normalize x along its first dimension and scale by w; x is (hidden, tokens) or (head_dim, heads, tokens).

linear(backend, W, x) -> (out, tokens)

Wᵀx for a weight stored (in, out).

rope(backend, x, cos, sin, positions) -> same shape

Rotary embedding of x (head_dim, heads, n) for tokens at the 0-based positions (a device vector of n Int32). cos/sin are the model's full tables (head_dim, max_len), whose column p + 1 is position p. Uses the "rotate half" pairing of Llama/Qwen: element i is paired with element i + head_dim/2.

silu_mul(backend, gate_up) -> (intermediate, tokens)

SwiGLU activation silu(gate) .* up, where gate_up stacks gate rows over up rows.

xielu(backend, x, act::XIELU) -> same shape

The xIELU activation of an ungated MLP (Apertus) with a layer's scalars act, rounding to T after each operation as transformers' XIELUActivation does.

XIELU{T}

The learned scalars of a layer's xIELU activation (Apertus, arXiv:2411.13010), in the form transformers applies them, each rounded to T as there: x > 0 ? αp x² + β x : (expm1(min(x, ε)) - x) αn + β x.

store_kv!(backend, k_pages, v_pages, k, v, slots)

Write the keys and values (head_dim, kv_heads, n) of n tokens into the page pools (head_dim, page_size, kv_heads, num_pages) at the 1-based cache slots (a device vector, see slots): slot s is offset (s - 1) % page_size of page (s - 1) ÷ page_size + 1.

attention(backend, q, k_pages, v_pages, block_table, positions, scale; causal = true) -> (head_dim, heads, n)

Causal grouped-query attention of the queries q (head_dim, heads, n) of one sequence, at the consecutive 0-based positions (a device vector of n Int32), over its first positions[end] + 1 cached positions. The sequence's keys and values are in the pages listed by block_table (a device vector, which may be padded past the sequence's pages). With causal = false, every query sees all positions[end] + 1 positions, including those of the later queries of the step (a DFlash drafter's block).

attention_workspace(backend, q, k_pages)

Scratch memory attention! needs for queries shaped like q, or nothing.

embed!(backend, y, table, ids) -> y

In-place embed, writing into y.

Every op has an in-place form like this one, which writes into a preallocated output and returns it. TileBackend writes directly and allocates nothing, so a step made of these calls can be captured in a CUDA graph. The generic fallbacks run the allocating op and copy its result.

rmsnorm!(backend, y, x, w, eps) -> y

In-place rmsnorm, writing into y (see embed!).

linear!(backend, y, W, x) -> y

In-place linear, writing into y (see embed!).

rope!(backend, y, x, cos, sin, positions) -> y

In-place rope, writing into y (see embed!).

silu_mul!(backend, y, gate_up) -> y

In-place silu_mul, writing into y (see embed!).

xielu!(backend, y, x, act) -> y

In-place xielu, writing into y (see embed!).

attention!(backend, out, workspace, q, k_pages, v_pages, block_table, positions, scale; causal = true) -> out

In-place attention, writing into out (see embed!). workspace is the backend's scratch memory, from attention_workspace.

attention!(::TileBackend, out, workspace, q, k_pages, v_pages, block_table, positions, scale; causal = true)

The decode kernels when there is a split-KV workspace, the prefill kernel otherwise.

attention_inputs!(backend, q, k_pages, v_pages, qkv, q_norm, k_norm, table, positions, slots, eps)

Everything between the fused q/k/v projection qkv ((heads + 2 kv_heads) head_dim, n) and attention, as one op: the optional Qwen3 q/k RMSNorm (q_norm/k_norm or nothing), rotary embedding with the RotaryTable table at positions, and writing the keys and values into the page pools at slots. Writes the queries into q (head_dim, heads, n) and returns it. The fallback chains the primitive ops; TileBackend fuses them.

add_rmsnorm!(backend, x, h, delta, w, eps)

The residual update x .+= delta (rounded to T), then h = rmsnorm(x, w, eps), in one op. Returns h.

linear_swiglu!(backend, y, gate_up, W, x)

y = silu_mul(linear(W, x)) for the stacked gate/up weight W, in one op where the backend can fuse it; gate_up is scratch for backends that cannot. Returns y.

linear_swiglu!(::TileBackend, y, gate_up, W, x)

SwiGLU of the gate_up projection: fused into the GEMV for one token, otherwise a GEMM into gate_up followed by silu_mul.

linear_xielu!(backend, y, W, x, act)

y = xielu(linear(W, x), act) for the up weight W of an ungated MLP, in one op where the backend can fuse it. Returns y.

linear_xielu!(::TileBackend, y, W, x, act)

xIELU of the up projection: fused into the GEMV for one token, otherwise a GEMM followed by xielu! in place.

sample_token!(backend, ids, mirror, logits, sampler)

Sample a token from the logits (vocab, 1) with the parameters and seed of the TokenSampler sampler and its next counter, which it advances. The 0-based token goes to ids[1] and mirror[1], as sample_host picks it. On TileBackend this runs on the device and can be captured in a graph.

sample_token!(backend, ids, mirror, logits, sampler, (q_ids, q_probs))

sample_token! that also writes the distribution the token is sampled from, for speculative sampling (accept_token!): the 0-based ids of the filtered_dist survivors, best first, into q_ids (-1 past them) and their probabilities into q_probs (0 past them), vectors of MAX_TOP_K entries.

sample_tokens!(backend, ids, mirror, logits, sampler, scratch, dists = nothing)

sample_token! on each column j of logits (vocab, n) in turn, into ids[j] and mirror[j] (and with dists, a pair of (MAX_TOP_K, n) matrices, column j's distribution into their column j), in a few launches on TileBackend. scratch is a zeroed sample_cols_scratch.

accept_token!(backend, ids, mirror, logits, sampler, draft, (q_ids, q_probs))

The token a verify step emits at a position drafted as draft[1] from the distribution (q_ids, q_probs) (see accept_host): draft[1] with probability min(1, p/q) under the sampler's parameters, else a token from max(p - q, 0) that differs from it. Uses and advances the sampler's counter like sample_token!.

Int8Weight

A linear weight (in, out) as Int8 values q with one Float32 scale per output column: W[:, j] ≈ scale[j] * q[:, j]. Made by quantize_int8.

quantize_int8(W) -> Int8Weight

Round-to-nearest symmetric quantization of each column of W (in, out) to [-127, 127], with the scale max |W[:, j]| / 127 (1 for an all-zero column).

Int8Pages

A key or value page pool stored as Int8 values q (head_dim, page_size, kv_heads, num_pages) with one Float32 scale (page_size, kv_heads, num_pages) per cached vector: pool[:, o, h, p] ≈ scale[o, h, p] * q[:, o, h, p]. A layer's keys and values are the two halves of one Int8 array and one scale array (layer_q, layer_scale, keys first), so a kernel takes both pools as a single argument and stays under 32 parameters. Vectors are stored divided by smooth (head_dim), the keys' smoothing factors (ones for values): pool[:, o, h, p] ≈ smooth .* scale[o, h, p] .* q[:, o, h, p]. Made in pairs by int8_kv_pools.

supports_graphs(backend) -> Bool

Whether a forward step on backend can be captured in a CUDA graph: its in-place ops neither allocate nor synchronize with the host.

Run the CLI on args; returns the process exit code.