tinyllm

tinyllm is an early-stage Rust-native framework for causal language model inference and serving. Its intended differentiation is the serving runtime: request scheduling, continuous batching, KV-cache management, model execution orchestration, streaming, observability, and performance measurement.

The project will initially reuse Candle and mature GPU primitives. It is not a new tensor framework and will not reimplement GEMM.

Status

Milestones 0 through 8 are complete; Milestone 5's continuous-admission lifecycle was delivered with Milestone 4. The scheduler, paged-cache ownership model, prefix-cache ownership, correctness suite, and paired CUDA benchmarks are recorded. Milestone 1 is complete for CPU/F32 correctness; the CUDA/BF16 path runs and is benchmarked, but exact parity with the PyTorch CUDA/BF16 oracle remains under investigation. Model execution runs behind a public non-HTTP engine on a dedicated worker, and an Axum frontend exposes a deliberately limited OpenAI-compatible HTTP subset.

The repository currently provides:

• one Cargo package with a reusable tinyllm library and thin CLI binary;
• local directory and Hugging Face Hub artifact resolution;
• config.json, tokenizer.json, and single or sharded SafeTensors discovery;
• explicit architecture detection for LlamaForCausalLM and Qwen3ForCausalLM;
• a model inspection command that validates the loading boundary;
• end-to-end generation for Qwen3-0.6B: typed config, tokenizer, SafeTensors loading (full-buffer Vec<u8> read), 28-layer model with GQA attention, BF16/F32 KV caching, cached greedy decode, timing metrics, resolved device/dtype reporting, and BF16 preflight probe;
• Qwen3 architecture components: RMSNorm, RoPE, GQA, SwiGLU MLP, causal attention, and contiguous plus paged-reference cache paths;
• a public non-HTTP Engine API backed by a named dedicated model-runner thread;
• one-time worker startup for the tokenizer, Qwen3 model, device, dtype, and EOS configuration;
• bounded Crossbeam command and per-request event channels, stable request IDs, incremental token events, terminal summaries or failures, cooperative cancellation, and explicit shutdown;
• bounded FIFO scheduling, continuous admission, token-budgeted chunked prefill, and variable-length batched decode for multiple active sequences;
• one runner-owned physical KV page pool, worker-owned allocator, logical block table and incremental decoder per active sequence, with deterministic reserve/commit/release and a retained contiguous comparison path;
• opt-in, bounded process-global prefix caching for immutable committed full-page prompt K/V, with exact namespace/token verification, active sequence leases, LRU eviction, and exact/partial hit metrics;
• GET /health, GET /ready, and GET /v1/models;
• streaming and non-streaming POST /v1/completions and POST /v1/chat/completions;
• bounded HTTP admission, bounded engine-to-Tokio event forwarding, disconnect cancellation, public request IDs, structured error envelopes, and graceful SIGINT/SIGTERM shutdown;
• Qwen3 chat-template rendering from tokenizer_config.json, with text-only system, user, and assistant messages and thinking disabled.
• an embedded, dependency-free browser chat client served from / that uses the existing streaming chat endpoint.

Generation for non-Qwen3 architectures is not implemented yet.

Known limitations:

• Chunked prefill defaults to 256 prompt tokens within a 512-token scheduler budget. The default accepted prompt limit remains 2,048 tokens and may be raised to the model's 40,960-position limit, but the Candle reference path still reconstructs the growing prefix and is not optimized for very long contexts.
• Checkpoint weights are loaded into an owned Vec<u8> buffer, approximately 1.5 GB for Qwen3-0.6B. Unsafe mmap is not used.
• One partially prefilling request advances at a time. Decode rows run before each prompt chunk, so established generations keep making progress, but a newly queued short prompt waits behind the active long prefill.
• Paged-reference decode gathers logical pages and right-pads rows on every layer. This readable correctness path copies temporary K/V data; direct paged-attention kernels remain deferred.
• The default KV budget is 1 GiB with 16-token pages. Override it with --kv-cache-budget-mib and compare with --kv-cache-mode contiguous.
• Prefix caching is disabled by default because this server has no authenticated tenant boundary. process-global mode retains prompt-derived K/V and token page identities in memory for reuse by other requests in the same process. Only complete pages are shared; a non-aligned suffix is recomputed.
• Model startup accepts a local directory or Hugging Face repository ID. Hub resolution is synchronous during worker startup, never in an HTTP handler.
• Cancellation is cooperative between model invocations; it does not interrupt an in-flight CPU or CUDA kernel.
• The HTTP API is not a complete OpenAI implementation. It supports one prompt or one text-only chat history, one greedy completion, and max_tokens plus stream. Sampling controls, stop, tools, multimodal content, response formats, prompt arrays, multiple choices, logprobs, and unknown fields are rejected rather than ignored.
• Authentication, TLS termination, and multi-model routing are not provided.
• CUDA/BF16 generation passes its device preflight and runs on an RTX 5090, but the tested greedy sequence diverges from PyTorch CUDA/BF16 after roughly 12 tokens. CPU/F32 remains the correctness baseline.

Supported architectures

Architecture	Inspection	Generation
LlamaForCausalLM	Yes	No
Qwen3ForCausalLM	Yes	Yes (Qwen3-0.6B)

Quick Start

Inspect a model

# Local directory
cargo run -- inspect --model /path/to/model

# Hugging Face repository
cargo run -- inspect --model meta-llama/Llama-3.2-1B --revision main

Private or gated repositories use the standard Hugging Face token and cache environment configured by hf-hub.

Generate text (Qwen3ForCausalLM only)

# CPU (use --dtype float32 for BF16-unsupported ops)
cargo run -- generate --model models/Qwen3-0.6B --prompt "Hello" --dtype float32

# CUDA/BF16 runtime path; exact reference parity is still pending
cargo run --release --features cuda -- generate \
  --model models/Qwen3-0.6B \
  --prompt "Hello" \
  --device cuda

# Adjust max new tokens
cargo run -- generate --model models/Qwen3-0.6B \
  --prompt "Hello world" \
  --max-new-tokens 64 \
  --dtype float32

# Compare the retained one-shot path or tune chunk scheduling
cargo run --release -- generate \
  --model models/Qwen3-0.6B --prompt "Hello world" \
  --prefill-mode chunked --max-prefill-chunk-tokens 128 \
  --max-batch-tokens 256

CUDA support is opt-in so the default build works on development and CI hosts without a CUDA toolkit:

cargo build --release --features cuda

That command enables the CUDA backend used by generate; inspect remains an artifact-only command.

Engine API

The CLI and HTTP server use the same synchronous, non-HTTP engine API:

use tinyllm::engine::{
    DTypeConfig, DeviceConfig, Engine, EngineConfig, GenerationEvent,
    GenerationRequest, KvCacheConfig, ModelSourceConfig, SamplingParams,
    SchedulerConfig,
};

let engine = Engine::start(EngineConfig {
    model: ModelSourceConfig::Local("models/Qwen3-0.6B".into()),
    device: DeviceConfig::Cpu,
    dtype: DTypeConfig::Float32,
    scheduler: SchedulerConfig::default(),
    kv_cache: KvCacheConfig::default(),
})?;

let handle = engine.submit(GenerationRequest {
    prompt: "Hello".into(),
    max_new_tokens: 32,
    sampling: SamplingParams::Greedy,
})?;

while let Some(event) = handle.recv()? {
    match event {
        GenerationEvent::Token(token) => print!("{}", token.text),
        GenerationEvent::Finished(summary) => {
            print!("{}", summary.trailing_text);
            break;
        }
        GenerationEvent::Failed(failure) => {
            return Err(std::io::Error::other(failure.message).into());
        }
    }
}

engine.shutdown()?;
# Ok::<(), Box<dyn std::error::Error>>(())

GenerationHandle::cancel requests cancellation at the next safe generation boundary. Dropping the handle also cancels the request. Token fragments use the tokenizer's incremental decoder; GenerationSummary::trailing_text contains any final buffered suffix needed to reconstruct decoded_text.

HTTP Serving

Start one configured model:

cargo run --release -- serve \
  --model models/Qwen3-0.6B \
  --served-model-name qwen3-0.6b \
  --bind 127.0.0.1:8000 \
  --device cpu \
  --dtype float32

The server waits for model startup before binding. For Hub startup, pass --model Qwen/Qwen3-0.6B --revision <revision>. Inspect all server limits with cargo run -- serve --help.

Enable the trusted-process prefix cache explicitly:

cargo run --release --features cuda -- serve \
  --model models/Qwen3-0.6B \
  --served-model-name qwen3-0.6b \
  --device cuda --dtype bfloat16 \
  --prefix-cache process-global \
  --prefix-cache-max-mib 256 \
  --prefix-cache-max-entries 1024

This mode requires paged-reference KV storage and chunked prefill. Cache identity includes the resolved model source/revision, model/config/weight digests, tokenizer artifacts, execution dtype/backend, KV layout/page size, and exact token pages. --prefix-cache-max-mib covers retained KV pages plus analytically accounted host index metadata; the entry limit is an additional bound. The HTTP API does not expose cache entries or hit state.

Scheduler defaults are a command queue of 8, pending queue of 12, 4 active sequences, and decode batches of 2. Tune them with --engine-command-queue-capacity, --max-pending-requests, --max-active-sequences, --max-decode-batch-size, and --max-commands-per-tick. Chunked prefill is controlled by --max-batch-tokens, --max-prefill-chunk-tokens, --max-prompt-tokens, and --max-total-sequence-tokens; use --prefill-mode one-shot only for differential comparison. Command-queue and scheduler-pending saturation map to structured HTTP 429 responses. Larger active/batch limits increase physical page use and temporary gather/padding memory. The default cache mode is paged-reference; contiguous cache comparison also requires --prefill-mode one-shot.

Open http://127.0.0.1:8000/ for the minimal web chat. The page discovers the configured model through GET /v1/models and sends the complete in-memory conversation to the existing streaming POST /v1/chat/completions endpoint. Conversation state exists only in the current page and is cleared by refresh or the Clear button. Stop aborts the active HTTP stream, which triggers the server's existing disconnect cancellation path.

The webpage intentionally provides only text chat, max_tokens, streaming, Stop, and Clear. It has no persistence, Markdown rendering, sampling controls, sessions, or authentication. Keep the server bound to a trusted interface. It requires a modern browser with ES modules, Fetch streaming, TextDecoder, ReadableStream, and AbortController support.

Non-streaming completion:

curl --fail-with-body http://127.0.0.1:8000/v1/completions \
  -H 'content-type: application/json' \
  -d '{
    "model": "qwen3-0.6b",
    "prompt": "Hello",
    "max_tokens": 16
  }'

Streaming chat completion:

curl --no-buffer --fail-with-body \
  http://127.0.0.1:8000/v1/chat/completions \
  -H 'content-type: application/json' \
  -d '{
    "model": "qwen3-0.6b",
    "messages": [{"role": "user", "content": "Say hello in one word."}],
    "max_tokens": 16,
    "stream": true
  }'

Successful streams use text/event-stream, emit ordered OpenAI-shaped chunks, and terminate with exactly data: [DONE]. Every response carries x-request-id. Liveness is GET /health; readiness is GET /ready and does not execute the model.

Only these request fields are accepted:

Endpoint	Accepted fields
/v1/completions	model, string prompt, max_tokens, stream
/v1/chat/completions	model, text-only messages, max_tokens, stream

The model name must exactly match --served-model-name. Chat histories may have one leading system message, then alternating user/assistant messages, and must end with user. See docs/milestone-3-report.md for protocol, lifecycle, test, and benchmark details.

The dependency-free frontend parser fixture can be run manually with a recent Node.js runtime:

node tools/test_web_frontend.mjs

Architecture

CLI or Axum HTTP frontend
        |
        v
public Engine API
        |
        v
bounded command channel
        |
        v
dedicated model-runner worker
   |        |        |
tokenizer  Qwen3   request executions + block tables
   |        |        |
   +--------+--------+
            v
   runner-owned physical KV page pool
        |
        v
bounded per-request GenerationEvent channel

The worker exclusively owns Candle model execution state; the model is not wrapped in Arc<Mutex<_>>. HTTP handlers translate DTOs and never access the tokenizer, Candle tensors, KV cache, device, or model runner. A finite spawn_blocking consumer forwards each accepted engine event stream into a bounded Tokio channel; response drop cancels the associated generation. See docs/architecture.md for detailed ownership and failure semantics.

Roadmap

1. Execute Qwen3-0.6B end to end with Candle and KV-cached greedy decoding: CPU/F32 correctness complete; CUDA/BF16 parity investigation remains open.
2. Introduce the non-HTTP engine, dedicated worker, streaming events, and request/sequence lifecycle: complete.
3. Add streaming OpenAI-compatible HTTP endpoints without moving model execution into async request tasks: complete.
4. Add correct decode batching and continuous request scheduling: complete.
5. Continuous admission lifecycle: completed as part of Milestone 4.
6. Replace fixed KV allocation with paged KV blocks and a typed attention metadata boundary: complete, with a correctness-first gather path.
7. Add chunked prefill and shared token-budget scheduling: complete, with a correctness-first Candle path.
8. Add bounded immutable full-page prefix caching: complete.
9. Specialize kernels only after benchmarks identify a concrete bottleneck.

Each milestone has dependencies, acceptance criteria, and required benchmark evidence in docs/roadmap.md. The Milestone 2 validation and benchmark record is in docs/milestone-2-report.md. The Milestone 3 HTTP compatibility and validation record is in docs/milestone-3-report.md. The Milestone 4 implementation and validation record is in docs/milestone-4-report.md. The Milestone 6 cache design, validation, and benchmark record is in docs/milestone-6-report.md. The Milestone 7 scheduler, correctness, and benchmark record is in docs/milestone-7-report.md. The Milestone 8 consolidation, ownership, correctness, privacy, and benchmark record is in docs/milestone-8-report.md.

Non-Goals

The initial project does not target training, autograd, a general tensor library, distributed inference, tensor or pipeline parallelism, multimodal or MoE models, speculative decoding, broad quantization, every Hugging Face architecture, custom GEMM, or CPU-first high-performance inference.

SafeTensors format support and model architecture support are separate claims. An artifact being readable does not mean its architectures value is supported.

Development

The local quality gate is:

cargo fmt --all --check
cargo build --locked
cargo test --locked --all-targets
cargo clippy --locked --all-targets -- -D warnings

See CONTRIBUTING.md for conventions and docs/dependencies.md for dependency policy and the initial dependency evaluation.