Expand description
Kopitiam Runtime: the Qwen-family transformer forward pass, running entirely in Rust on CPU, offline.
This crate is the payoff of the layers beneath it. kopitiam-core
defines the shared vocabulary (kopitiam_core::DType,
kopitiam_core::Shape); kopitiam-tensor implements the CPU tensor
and its ops (matmul, softmax, RMSNorm, SiLU, embedding gather, GGUF
block-quantization decoding); kopitiam-loader parses GGUF/SafeTensors
files into raw bytes plus dtype/shape metadata; kopitiam-tokenizer
turns text into token ids and back. None of those crates constructs a
model or runs a forward pass — this one does.
§Pipeline
kopitiam_loader::load_model(path)
-> LoadedModel
-> QwenModel::from_loaded_model (crate::model)
- QwenConfig::from_metadata (crate::config)
- ModelWeights::load (crate::weights, via crate::bridge)
- RotaryEmbedding::new (crate::rope)
-> generate(&model, &tokenizer, prompt, ...) (crate::generate)
per new token:
embedding lookup (Tensor::gather_rows)
per layer: block_forward (crate::block)
RMSNorm -> attention_forward (crate::attention)
RoPE (split-half) (crate::rope)
grouped-query KV repeat (crate::attention::repeat_kv_heads)
causal mask + softmax
KV cache read/append (crate::kv_cache)
RMSNorm -> swiglu_mlp (crate::mlp)
final RMSNorm -> output projection
[optional] constraint mask -> -inf (crate::constraint)
greedy sampling (crate::sampling)§Module map
- [
bridge] — the loader/tensor glue: bytes + dtype + shape ->Tensor.kopitiam-loaderandkopitiam-tensorwere built concurrently and deliberately do not depend on each other; this crate is the first one downstream of both, so this is where that gap is bridged. - [
config] —config::QwenConfig, the architecture hyperparameters resolved (with documented fallbacks and validation) fromkopitiam_loader::ModelMetadata. - [
rope] — rotary position embeddings, split-half (“NEOX”) pairing. Read this module’s docs before touching anything position-related; a swapped pairing convention is this crate’s single easiest place to introduce silent, undetectable-by-type-system wrongness. - [
kv_cache] — the per-layer, growable key/value cache that makes autoregressive decoding linear instead of quadratic in sequence length. - [
attention] — grouped-query causal self-attention: repeating shared KV heads across their query-head group, causal masking, and wiring RoPE and the KV cache into one attention forward pass. - [
mlp] — the SwiGLU feed-forward block. - [
linear] — the singlex @ W^T + bhelper every projection in this crate goes through. - [
block] — one pre-norm transformer block (attention sub-layer, MLP sub-layer, both with a residual connection). - [
weights] — loads every named GGUF weight tensor a block/model needs. - [
model] —model::QwenModel: wires embedding, every block, the final norm, and the (possibly tied) output projection into atraits::Modelimplementation. traits—traits::Modelandtraits::Backend, the Model Runtime boundary every layer above this crate should code against.- [
sampling] — turning a row of logits into a token id:sampling::GreedySampler(argmax) andsampling::StochasticSampler(temperature/top-k/top-p/min-p/ repetition penalty, composed as a pipeline and driven by a seeded PRNG — see that module’s docs for the pipeline shape and why seedability is mandatory, not optional). - [
constraint] — grammar-constrained decoding (the keystone): mask the tokens aconstraint::TokenConstraintforbids to-infat the front of the sampling path, before temperature/top-k/top-p, so the model physically cannot emit invalid structure. Ships a fixed allowed-token-set (constraint::AllowedTokens, e.g. a tool-name enum) and a structural JSON constraint (constraint::JsonStructure);constraint::mask_logitsis the masking step andconstraint::ConstrainedSamplerthe drop-in sampler wrapper, driven end to end bygenerate::generate_constrained. See that module’s docs for why mask-before-sample and-inf-not-0.0(AID-0045). - [
gguf_tokenizer] — builds akopitiam_tokenizer::BpeTokenizerdirectly from a GGUF file’s embeddedtokenizer.ggml.*vocabulary (no companiontokenizer.jsonneeded). generate— the end-to-end entry point:prompt -> tokens -> forward passes -> sampled ids -> text, with streaming per-token output.
§What is here as of Phase 2, and what is still deliberately not
As of Phase 2: stochastic sampling (sampling::StochasticSampler —
temperature/top-k/top-p/min-p/repetition penalty) alongside greedy, and
a fused quantized matmul for Q4_0/Q8_0 matmul-operand weights (see
kopitiam_tensor::Tensor::quantized_matmul and
bridge::load_matmul_weight — weights whose on-disk dtype is
quantized now stay quantized in memory instead of being eagerly
dequantized to f32, which is what makes a multi-gigabyte model’s
resident memory footprint match its on-disk size rather than balloon by
4-8x).
Still not here: no batching across multiple concurrent prompts; no
scheduler or execution graph (see the parent epic, kopitiam-082,
Phase 3, and this crate’s benchmark module for why a general graph
engine was judged not to earn its keep yet); no SIMD. “Correct before
fast” per this workspace’s engineering principles remains the ordering
principle — every fast path added so far (quantized matmul) ships
alongside, and is tested against, the plain reference it replaces.
Re-exports§
pub use traits::Backend;pub use traits::CpuBackend;pub use traits::Model;
Modules§
Structs§
- Allowed
Tokens - A fixed allowed-token-set constraint: the same ids every step, regardless of decode state.
- Constrained
Sampler - Wraps any
Samplerso the constraint’s mask is applied at the front of the sampling path — the drop-in integration seam for real decoding. - Decode
State - The decode state a
TokenConstraintreasons about when it decides which next tokens are valid. - Generation
Config - Generation limits and stop conditions for
generate. - Greedy
Sampler - Always picks the highest-scoring token:
argmax(logits). - Json
Structure - A
TokenConstraintthat keeps the model’s output structurally valid JSON. - KvCache
- Per-layer, growable key/value cache for one generation session.
- Qwen
Config - The resolved shape of a Qwen-family transformer: everything the forward
pass needs to know about the model it is about to run, with no further
Options or metadata lookups once construction succeeds. - Qwen
Model - A loaded, ready-to-run Qwen-family transformer.
- Rotary
Embedding - Precomputed
cos/sintables for every position up tomax_position, soRotaryEmbedding::applyis a table lookup per (position, pair) rather than acos/sincall — cheap enough to matter once this runs once per token per layer per head during decode. - Sampling
Config - Configuration for
StochasticSampler’s logit-transform pipeline. - Shape
- The dimensions of a tensor, outermost first.
- Slice
Vocab - The obvious
TokenVocabover an ownedid -> bytestable. Handy for tests and for any caller that already has the vocabulary as aVec. - Stochastic
Sampler - Temperature / top-k / top-p / min-p / repetition-penalty sampling,
composed as a pipeline of logit transforms (see this module’s docs) and
driven by a seeded [
Rng].
Enums§
- Allowed
Set - Which next-token ids a constraint permits at the current step.
- Constrained
Generate Error - Either a normal runtime error, or the constraint leaving no valid token —
the two failure modes
generate_constrainedcan hit. - Constraint
Error - What went wrong when a constraint could not be satisfied.
- DType
- The element type of a tensor.
- Device
- Where a tensor’s storage lives and where kernels run on it.
- Error
- Every way a Kopitiam Runtime operation can fail.
Traits§
- Sampler
- Chooses the next token id from one row of logits (length
vocab_size). - Token
Constraint - A rule that says which next tokens are valid, given what’s been decoded so far.
- Token
Vocab - A token id -> its byte string, for constraints that must reason about the bytes a candidate token would append (not just its id).
Functions§
- generate
- Greedily generates a completion for
prompt. - generate_
constrained - Like
generate_with_sampler, but every step is grammar-constrained: theConstrainedSampler’sTokenConstraintmasks disallowed tokens to-infbefore its inner sampler runs, so the model physically cannot emit a token the constraint forbids. This is the keystone that makes a small model reliably produce valid JSON / valid tool names / valid structure — see [crate::constraint] for the full rationale and provenance (AID-0045). - generate_
with_ sampler - Identical to
generateexcept the token-selection strategy is pluggable: anySamplerimpl (greedy, or acrate::sampling::StochasticSamplerconfigured for temperature/ top-k/top-p/min-p/repetition-penalty sampling) drives which token is picked at every step, instead of alwaysargmax. Every other detail — prefill, one-token-at-a-time KV-cache decoding, EOS handling, streaming viaon_token, the returned completion text — is exactlygenerate’s behaviour, becausegenerateis defined in terms of this function; see this module’s docs for why that direction of composition (not the reverse) is what keeps the two entry points from drifting apart. - greedy_
argmax argmax(logits), ties broken toward the lowest id (the first maximum encountered) — aPartialOrdtie-break rule that is total and deterministic even in the presence ofNaN(f32::partial_cmp’sNonecase), unlike anassert-free.max_by(...)over rawf32, which would panic or silently misbehave on aNaNlogit.- load_
matmul_ weight - Loads
namethe way every matmul-operand weight (wq/wk/wv/wo, the SwiGLU MLP’s three matrices, andoutput.weight— see [crate::weights::LayerWeights] / [crate::weights::ModelWeights]) should: keeping a block-quantized on-disk dtype only when there is a fused kernel to compute on it, and otherwise dequantizing it tof32at load. - load_
matmul_ weight_ opt - Like
load_matmul_weight, but returnsOk(None)instead ofError::MissingTensorwhennameis absent — the matmul-operand counterpart toload_tensor_f32_opt, for the same optional weights (per-projection attention biases,output.weightwhen tied to the token embedding). - load_
tensor_ f32 - Looks up
nameinmodel, builds aTensorfrom its bytes, and dequantizes/upcasts it tof32— the dtype every op inkopitiam-tensoractually computes in. This is the call every weight-loading site in [crate::weights] goes through: model files may shipf16,bf16, or any of the five GGUF block-quantized formats, and this one function makes the rest of the forward pass indifferent to which. - load_
tensor_ f32_ opt - Like
load_tensor_f32, but returnsOk(None)instead ofError::MissingTensorwhennameis absent. - mask_
logits - The masking step itself: set every logit whose id is not in
allowedtof32::NEG_INFINITY, in place. - tensor_
from_ entry - Builds a
Tensorfrom oneTensorEntry’s bytes, preserving its on-diskDTypeexactly (still block-quantized ifentry.dtypeis quantized, stillf16/bf16if the file stored it that way). - tokenizer_
from_ gguf - Builds a
BpeTokenizerfrommodel’s embedded GGUF vocabulary.
Type Aliases§
- Result
- The runtime’s result alias.