Generation

In short: Sampler applies Hugging Face's steps in order (repetition penalty, temperature, top-k, top-p, min-p) and draws from what is left, with ties broken by token id so a seeded run is reproducible; temperature 0 is greedy. Session::Generate takes the whole context so far, keeps the part of its KV cache that the new prompt shares, runs the rest as one prefill, then samples, reports and feeds back one token at a time until a stop id, the token limit, a full cache or the caller says stop. SpeculativeSession lets a smaller model guess up to k tokens, checks them all in one batched target pass and keeps each with probability min(1, p / q), so the output distribution is exactly the target's. rml-chat loads a GGUF file or Hugging Face directory with its tokenizer and chat format and streams replies, optionally with a draft model.

The concepts are on /ml-models: decoding and autoregressive models. The code is sampler, generation and speculative in rml/llm, and examples/rml_chat. The forward pass that produces the logits is on The decoder, the chat formats and token bytes on Tokenizers.

The sampler

The header lists the steps in the order they run (sampler.h L4–21):

Filter does steps 1–5 on a copy of the logits (L75–146). Probabilities are kept unnormalized, relative to the best candidate, which gets exactly 1: exp((logit − best) / T) (L105–116). That makes min-p a plain comparison against min_p, and the draw needs only the running total. Ranking compares logits and breaks ties by the lower id (L69–73). All scratch (a logit copy, the candidate list, the probabilities, the penalty flags) is sized to the vocabulary when the sampler is made, so Sample never allocates (L41–47). Create rejects a negative or non-finite temperature, a negative top-k, top-p outside (0, 1], min-p outside [0, 1), a non-positive penalty and a negative window (L16–39).

Speculative decoding needs the distribution itself, not just a draw, so the sampler exposes two more calls: Distribution writes the filtered, normalized probabilities densely over the vocabulary (all mass on the arg max when greedy), and Draw samples from any non-negative weights (sampler.h L54–63; sampler.cc L165–196).

Worked example

Four tokens with logits [2, 1, 0, −1], temperature 1, top-p 0.9, min-p 0.05. Relative to the best, the probabilities are e⁰, e⁻¹, e⁻², e⁻³ = 1, 0.368, 0.135, 0.050, total 1.553. Top-p sorts them and keeps adding until the running sum reaches 0.9 × 1.553 = 1.398: after two tokens it is 1.368, after three 1.503, so three are kept. Min-p then keeps those at least 0.05 of the best, which all three are (the fourth, at 0.0498, would have fallen to min-p anyway). The draw is over 1, 0.368 and 0.135, so 66.5%, 24.5% and 9.0%. With repetition_penalty 1.1 and token 0 in the window, its logit would first become 2 / 1.1 = 1.82; a logit of −1 would become −1.1.

The generation loop

A Session keeps one conversation's KV cache between calls. Each call passes the whole context so far, the session finds how much of it the cache already holds and runs only the rest, so a chat turn costs its new tokens, not the history (generation.h L4–11). Session::Create(model, sampler, context_size, stop_ids, kv_dtype = kF16) makes the sampler and a cache through DecoderLM::NewCache (generation.cc L27–33).

Details worth knowing:

• Prefix reuse. std::mismatch finds how many cached tokens the new prompt starts with; at most prompt − 1 are reused, so there is always one token to run and logits to sample from. The cache is truncated to that length (L63–70).
• The context the sampler sees is the session's cached tokens, so the repetition penalty covers the prompt and the reply so far.
• Stop order. A stop id ends the call before it is counted or reported; then the callback; then the token limit; then a full cache. The token that hits the limit is reported but never run, so the next call's prefix match stops just before it (L85–102).
• Errors. An empty prompt or non-positive limit is InvalidArgument; a prompt longer than the cache is OutOfRange; a failed Forward is returned as is (L54–61).
• Stats. GenerationStats records prompt tokens run, tokens reused, tokens generated and the prefill and decode wall-clock seconds (generation.h L34–40).

Against the documented decode loop

docs/architecture.md describes the decode loop in three steps: load (parse, mmap, tokenizer, size the KV cache and a step arena), prefill in chunks of up to 512 tokens with logits only for the last position, and a per-token decode whose every buffer is preallocated (architecture.md L92–101). The code matches on the decode step (the order of operations, logits for the last position only, zero allocations, tested by decoder_alloc_test), with two differences at the pin:

• Scratch is DecoderLM's workspace of vectors that grow to the largest call, not an Arena reset each step.
• Generate runs the whole unmatched prompt as one Forward call; chunked prefill is planned. So is the backend dispatch table the same document has the decode loop call through (L83–90); today DecoderLM calls the kernels in rml/llm/kernels.h directly.

Speculative decoding

Each round the draft samples k tokens from its distributions q_1..k; the target runs the pending token and the k guesses as one batch, giving p_1..k+1. Guess i is kept with probability min(1, p_i(x_i) / q_i(x_i)); the first rejected one is replaced by a draw from max(0, p_i − q_i), renormalized, and the round ends; if every guess is kept a bonus token comes from p_k+1. Every kept token is then distributed exactly as the target alone would sample it (Leviathan et al. 2023; Chen et al. 2023), so the draft changes the speed and never the output; with greedy sampling the output is the target's token for token (speculative.h L4–20).

Measured. rml's own run, on the 4-core x86 development container: SmolLM2-360M-Instruct with the 135M model as draft gives greedy replies byte-identical to the 360M model alone, but decodes slower, 8.3 tok/s with 3 drafted tokens at 50% acceptance against 12.7 tok/s alone, because a batched pass over 4 tokens costs about two single steps there and the draft is only 2.7× smaller. A real speedup needs a memory-bound target (8B at Q4 on the Mac) and is unmeasured (progress.md L77); that measurement is planned (L291–292).

rml-chat

examples/rml_chat keeps its logic in a library, chat_lib (chat.h / chat.cc), with a thin main.cc that parses flags and prints, as the project's rule for entry points asks (BUILD.bazel).

Because Send re-encodes the whole conversation and the session reuses the shared prefix, a second turn runs only the new user message and the tokens around it. The command line (main.cc L33–55):

rml-chat MODEL [-p TEXT] [--system TEXT] [--format NAME]
         [--temp 0.7] [--top-k 40] [--top-p 0.95] [--min-p 0.05] [--repeat-penalty 1.1]
         [--seed 0] [--ctx 4096] [--kv f16|f32|i8] [--max-tokens 512] [--threads N]
         [--draft MODEL] [--draft-tokens 4]

The defaults above are rml-chat's, set in ParseFlags (L83–89), not SamplerConfig's. With -p it answers once and exits; otherwise each line is a user turn, /reset starts over and an empty line quits. After each reply it prints prompt tokens (and how many were cached), prefill and decode tok/s, and with a draft the acceptance rate and tokens per round to stderr (L154–175). A draft is loaded with the main model's chat format, and --threads gives both models one pool (L191–243).

Tests

These run on the toy models from tools/make_llm_fixtures.py (see The decoder). The program itself is listed with the others on Examples.