Skip to content

Repository files navigation

maxuna

A from-scratch Rust inference engine for poolside's Laguna S 2.1 (118B-total / 8B-active MoE coding model), built on candle, for running that one model at maximum speed with minimal ceremony.

maxuna is opinionated. One model, one machine class (Apple Silicon, Metal, 128GB unified memory), batch=1, the official GGUF quants — no CPU fallback, no multi-GPU. Every default is already the fastest correct configuration: official checkpoint, speculative decoding on, the full 256k trained context, a prefix cache in RAM and on disk. There are no llama.cpp-style command lines to assemble; flags exist to opt out of things or point at your own files, not to make it work. The reasoning behind these choices — and the measured dead ends behind several of them — is written down in docs/decisions.md.

Quick start

cargo build --release
./target/release/maxuna

That's the whole thing. A bare maxuna starts the server with a live dashboard on 127.0.0.1:5241, speaking both the Anthropic Messages API (POST /v1/messages) and the OpenAI Chat Completions API (POST /v1/chat/completions), streaming and non-streaming, tool calls included. The first run downloads the official Q4_K_M (68GB) and the DFlash drafter (2.2GB) into the standard Hugging Face cache (~/.cache/huggingface/hub, honoring HF_HUB_CACHE/HF_HOME); the ~70GB model load itself waits for the first request (~15s cold, ~7s warm — the weights are mmap'd straight out of the page cache into Metal buffer views, never copied). The model field of a request is echoed back verbatim, so clients pinned to another model name need no special-casing.

The other verbs:

maxuna --no-tui                                 # serve with plain stderr logs
maxuna --init                                   # write ~/.config/maxuna/serve.toml, annotated
maxuna chat                                     # interactive chat REPL
maxuna generate -p "Write fizzbuzz in Rust." --stats
maxuna generate --raw -p "def fibonacci(n):" --temp 0   # no template, BOS prepended
maxuna fetch                                    # prefetch both checkpoints, print paths
maxuna inspect                                  # dump GGUF metadata + parsed config

(maxuna here means ./target/release/maxuna, or cargo run --release --bin maxuna --.) The binary is self-contained — the checkpoint tokenizer is compiled into it, and everything else resolves through the Hugging Face cache — so cargo install --path . gives a maxuna that runs from any directory. --tokenizer <tokenizer.json> overrides the embedded vocabulary on the subcommands that take prompts. maxuna serve still exists and takes the same flags as the bare invocation. Serve flags at the top level and subcommands don't mix — flags for generate/chat go after the subcommand (maxuna generate --no-draft -p …, not maxuna --no-draft generate …; the parse error says as much). DFlash speculative decoding is on everywhere by default — --no-draft decodes plain, --draft <gguf> swaps in a custom drafter — and --model <gguf> overrides the checkpoint, each on the bare invocation and every subcommand that loads the model.

Dev setup: nix develop (or direnv via the checked-in .envrc) enters a devshell that installs the lefthook git hooks — pre-commit runs rustfmt on staged .rs files. Rust itself comes from the system toolchain, not the flake.

Serving

One model process, one inference at a time by construction (a dedicated engine thread owns the model; the HTTP layer never touches Metal), with an inspectable queue in front (--queue-capacity, default 16; overflow answers 429 + Retry-After: 1). Both API dialects are on by default and individually disableable; GET /v1/models and GET /health are always served. Setting api_key requires it on every request except /health, worth doing whenever host leaves loopback.

--init writes a config template carrying every setting at its built-in default with the reasoning behind each one. Settings resolve CLI flag > config file > default, and a flag that contradicts an explicit config value prints a warning naming both — a forgotten flag never silently wins.

The dashboard is the default sink: the running job (prefill progress bar with ETA, then a decode line separating thinking from answer tokens), the queue, what each cache slot holds, the last 50 finished requests, and the scrollable log. It exists because the server is at its quietest exactly when it is busiest — a multi-minute prefill runs with nothing to print at ~0% CPU, indistinguishable from a hang without it. It draws on stderr and steps aside by itself when stderr is not a terminal (a 2> redirect, a pipe, CI — no flag needed), and --no-tui forces plain logging, whose lines are byte-identical to what they have always been. q or Ctrl-C shuts down gracefully.

Thinking is forced on every assistant turn by default, with prior turns' thinking stripped from the context — except across the trailing run of assistant and tool-result messages since the last real user message, where reasoning is preserved verbatim so a tool loop keeps its train of thought. Requests steer it via Anthropic's thinking.budget_tokens or OpenAI's reasoning_effort; requests that carry neither get the server's default_budget (4096 — the checkpoint is prone to open-ended reasoning loops that can outlast an agent client's timeout). A budget that cannot fit alongside the reply is quietly lowered or dropped rather than failing the request. Budgets drive the same graduated exit schedule the CLI's --max-think uses (described under Generation), and thinking comes back as Anthropic thinking blocks / OpenAI reasoning_content.

The prefix cache keeps several conversations warm at once (--cache-slots, default 4): one lives in the GPU cache, the rest as host-RAM images that page back in at 8–13 ms where re-prefilling costs seconds to minutes. Matching is longest-common-prefix over token histories — no session ids, which is what lets a hit cross dialects. A follow-up turn reaches its first token in 0.09s against 1.68s cold. Host RAM is the budget to watch (~54 GB at the defaults if every slot fills; the generated config carries the formula), and cache reuse is reported per request in both dialects' usage fields.

Warm conversations also survive the process: slots are written under --cache-dir as a tree of prefix segments (shared prefixes stored once, a fork splits the stored segment at the divergence point), and a restarted server resumes a 13k-token conversation in 2.0s where it cost 41.5s the restart before. A fresh conversation sharing only a long system prompt with a stored one resumes at the end of that prompt — agent clients get this for free. --disk-max-gib (default 64) bounds the store; everything invalid, stale or damaged degrades to an ordinary cache miss; --no-disk-cache turns it off.

Scheduling is shortest-estimated-prefill-first (prompt tokens minus what the cache already holds), with a 20s age guard against starvation and --schedule fifo to restore arrival order. What the scheduler cannot fix: a job already running blocks everything behind it for its whole duration — batch=1 is the design. Every request gets a wall-clock ceiling derived from its own size; a client that disconnects cancels within a token, and cancellation keeps the cache, so an identical retry resumes warm.

Tool calls work in both dialects, streaming and non-streaming, with definitions rendered byte-for-byte as the checkpoint was trained on and argument deltas that always concatenate to one valid JSON object. Client text is inert: a pasted file discussing <think> or </tool_call> encodes as literal characters, never control tokens. tool_choice supports auto and none; anything that would force a call is a 400 (no constrained tool grammars — a considered decision, see decisions.md). Note the checkpoint writes at most one call per turn. Absent: images, logprobs.

Structured output constrains the reply to a JSON schema by masking every answer-section draw to schema-legal tokens (llguidance; see decisions.md "Constrained decoding"). OpenAI: response_format with json_schema (the schema in json_schema.schema) or json_object; Anthropic: output_config.format with {"type": "json_schema", "schema": ...}. The reply is the JSON value — generation ends when the schema completes, with a plain stop finish. Thinking is untouched (the model reasons freely, the constraint arms at </think>), speculative decoding keeps working, and an unsupported schema (recursion via remote $ref, lookarounds, not) is a 400 naming the offending construct at request time. Combining a schema with tools is refused — the schema masks the whole reply, which would leave tools silently uncallable.

The native APIPOST /maxuna/v1/generate — exposes what the compat dialects cannot spell: steering the current turn itself. A continue block writes the start of the assistant turn into the prompt — thinking injects reasoning (left open for the model to extend, or sealed with "close_thinking": true to skip free-form reasoning entirely), and prefix is answer text the model must continue (requires the reasoning block closed). All of it composes with format (json_schema/json_object, same compilation as the compat dialects): a prefix is fed through the grammar before decoding, so the mask continues that document — a prefix the schema rejects is a 400 naming the offending token. The prefix is never echoed back; it counts as input tokens. Requests are strict (deny_unknown_fields, no field is ever silently dropped); absent thinking/max_thinking_tokens follow the server policy; max_thinking_tokens bounds newly generated reasoning only. Multi-turn note: the API is stateless, so a client replaying a continued turn as history must itself concatenate injected + generated thinking and prefix + generated text. scripts/classify.ts is a worked example — a tagging task where an injected one-sentence scaffold replaces free reasoning (close_thinking) under a schema, cutting a classification to a few dozen output tokens.

The full knob list is in maxuna serve --help and the annotated --init template.

Generation

generate and chat carry a few controls of their own (they go after the subcommand):

--max-think N bounds the reasoning block without truncating it: a graduated schedule (damping the "wait"/"but" reasoning-restart cues, then lifting </think> toward the leading logit, then one injected wrap-up sentence at a sentence boundary, then a bounded grace before forcing the exit) so the model concludes instead of being cut mid-sentence. Injected tokens are drawn through the sampler with every other logit at −inf, so they advance the RNG exactly as sampled tokens would — the invariant the speculative path's output identity rests on. --min-think N is the opposite floor.

--ban-string <s> (repeatable) drops every token whose decoded text contains s — banning an em dash masks 34 ids. Structural tokens are protected: a pattern colliding with the added vocabulary (--ban-string assistant matches <assistant>) skips it with a warning.

DFlash speculation proposes a block of tokens from a diffusion drafter reading the target's residual taps, and the target verifies them in one batched forward with KV rollback on partial accept. Code-shaped output measures ~1.3x plain decode; a wall-clock auto-pause controller keeps low-acceptance prose at parity instead of a loss, and drafting hands over to plain decode past --draft-ctx (default 8192) tokens of context, where it stops paying. Greedy output matches spec-off token-for-token except where batched-verify numerics flip a near-tie.

How it works

Laguna S 2.1: 48 layers; interleaved attention (every 4th layer global with YaRN rope, the rest sliding-window 512 with plain rope); per-layer head counts (48/72); QK-norm; per-head softplus attention output gating; 256 routed experts (sigmoid router, top-10, DeepSeek-style selection bias) + 1 shared expert.

Nearly every hot kernel is vendored — our own Metal sources under src/ops/, runtime-compiled and driven by our own host dispatch, mostly following ggml's current geometry. candle supplies the tensor/allocator/quantization plumbing and the sdpa vector kernels; it does not supply the matmuls:

kernel vendored source replaces
MoE prefill gather-matmul mm_id.metal (+ mm_id_t_hp.metal) candle's baked kernel_mul_mm_id_*, unusable at 256 experts
MoE decode gather-matvec mv.metal candle's baked kernel_mul_mv_id_*
attention projections f16.metal, f16_t.metal, q8.metal candle matmul, which cannot express mixed-dtype operands
prefill attention flash.metal candle sdpa + a materialized mask tensor
drafter matmuls (BF16 weights) bf16.metal, bf16_t.metal candle matmul over widened-to-f32 copies
glue: gate, rope, permute/cast, SwiGLU, MoE combine attn_glue.metal, rope.metal, silu_mul.metal, combine.metal multi-dispatch candle op chains

The glue kernels are bit-identical by construction to the chains they collapse (proven by bitwise tests); flash is value-identical to a pinned f32 sdpa configuration; the matmuls reorder accumulation and are held to the parity tier that matches their path. Nearly all carry a MAXUNA_*_CLASSIC kill-switch and record a provenance field the parity gate enforces. Expert routing and the expert matmuls run entirely on-GPU with one GPU→CPU sync per decoded token. candle is a pinned git dependency; the residency patch the mmap load relies on landed upstream (candle PR #3771).

Module map: config.rs (GGUF metadata → config), gguf.rs (mmap/no-copy loader, quant-agnostic weights, zero-copy ExpertStack), rope.rs (dual YaRN/plain), kv_cache.rs (prealloc full / 512-ring SWA, plus verify rollback), attention.rs, moe.rs (routing + fused/reference expert runners), ops/ (vendored Metal kernels + dispatch), model.rs (assembly + parity/spec taps), dflash.rs (the self-contained drafter), tokenizer.rs/chat.rs/sampler.rs (I/O), generate.rs (prefill/decode + speculative loop), parity_schema.rs (dump provenance), serve/ (HTTP server: config merge, engine thread, one module per API dialect), bin/maxuna/ (CLI + chat REPL), bin/logits-dump.rs (parity harness), bin/spec-verify-bench.rs.

Checkpoint

The supported file is the official laguna-s-2.1-Q4_K_M.gguf from poolside/Laguna-S-2.1-GGUF (2026-07-24 re-release, 68.2GB): all attention weights Q8_0 — which the dual-storage attention path targets, a dequantized f16 plane for the prefill gemm and a no-copy alias of the raw Q8_0 blocks for the decode gemv at half the streamed bytes — with shared experts, embeddings, lm_head and the dense layer-0 FFN Q8_0, and all routed expert stacks Q4_K.

Read the GGUF metadata, not the HF config.json: the official conversion is 256k-context with YaRN factor 32, where the HF config says 1M / factor 128.

Verification

Correctness is proven by parity against poolside's llama.cpp fork (vendored at reference/llama.cpp-laguna-branch, built via scripts/build-llamacpp.sh), which runs the identical GGUF. bun scripts/parity-gate.ts runs the whole cycle; docs/parity.md is the runbook. The gate is three-tier by change kind (strict oracle cosine for the classic fallback, fork-equivalence for the tiled prefill, greedy replay plus a perplexity bound for the decode kernels — why in docs/decisions.md), with floors calibrated per checkpoint quant mix. Unit tests: cargo test --release (the ops tests require the Metal device).

Status

Correct and parity-gated end-to-end. At the fork's decode throughput and just under it on prefill, measured in the same power mode — see CLAUDE.md for the current anchors and the power-mode caveat (an absolute number is meaningless unless you know which mode produced it). Perf work ongoing; TODO.md is the deferred-work ledger.

Docs: docs/decisions.md (why things are the way they are, including the refuted directions), docs/parity.md (verification runbook), docs/log.md (the dated engineering narrative behind both).

About

Laguna S 2.1, tuned to the wall for 128GB M5 Max

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages