Introduction

On an NVIDIA GB10 Grace-Blackwell Superchip, Atlas reaches 131 tok/s on Qwen3.5-35B-A3B — 3.6× faster than NVIDIA's vLLM on the same model, same hardware. It reaches 104 tok/s on Qwen3-Next-80B-A3B and 46 tok/s on Qwen3.5-122B-A10B (EP=2, two nodes). On a head-to-head suite of 32 micro-benchmarks against PyTorch — attention, GEMM, SSM, RoPE, RMSNorm — Atlas wins 32 out of 32, with speedups from 1.04× up to 18.2×.

This book is the canonical long-form documentation for Atlas. It complements — rather than replaces — the source-of-truth material already in the repository:

• README.md — headline performance claims, model matrix, porting guides.
• QUICKSTART.md — Docker recipes for every supported model.
• AGENTS.md / CONTRIBUTING.md — contributor workflow.
• docs/ — design notes, release history, benchmark journeys.

Who this book is for

Three audiences, one narrative arc:

• Operators who want to serve one of the supported models on a GB10 today (22 (model, quant) targets ship in the image; the compatibility matrix is docs/GB10_DEPLOYMENT_GUIDE.md §2). Start with Installation and Quickstart, then jump to Operating Atlas for CLI flags, KV-cache dtypes, and multi-GPU bring-up.
• Model authors extending Atlas with a new architecture. Read Architecture → spark-model → Engineering Deep Dives in order, then follow the "Adding a new model" section of the repo README alongside crates/spark-model/src/weight_loader/minimax.rs as a template.
• Kernel engineers porting Atlas to a new hardware target or hyperoptimizing an existing kernel. Read Philosophy → Kernel Dispatch Pipeline → the CUDA Kernel Engineering deep dive, then use the repo's "Adding a new hardware target" walkthrough with kernels/gb10/ as a reference implementation.

What Atlas is not

• Not a training framework. Atlas serves; it does not fine-tune. Use vLLM-style baselines, trl, or axolotl upstream.
• Not a generic kernel. Atlas does not cover the matrix with one templated CUDA kernel. It covers it by specializing per (hardware, model, quantization) target and wrapping those kernels in abstractions designed for broad support — ComputeTarget (vendor-agnostic build), GpuBackend (vendor-agnostic runtime), CommBackend (vendor-agnostic collectives). The first hardware we shipped is GB10; the design is explicitly multi-vendor, and porting to H100, B200, Apple Silicon, AMD, or Intel is a well-scoped piece of work, not an architectural change.
• Not a Python wrapper. Atlas is pure Rust + GPU source. There is no Python in the serving path — no PyTorch, no Triton JIT, no runtime compilation. Every kernel is compiled to its hardware's native binary (PTX today) at build time and embedded in the Rust binary.

What you get

A single multi-model Docker image (avarok/atlas-gb10:latest), an OpenAI-compatible HTTP server, and the machinery to port Atlas to fresh (H, M_q) targets. The model matrix spans Qwen3 through Qwen3.6, Qwen3-VL, Gemma-4, Mistral-Small-4, MiniMax-M2.7, Nemotron-3 Nano and Super, and DeepSeek-V4-Flash — covering dense, hybrid SSM/attention, MoE, vision, MLA, and 256-expert routing. The engine ships with MTP speculative decoding, RadixAttention prefix caching with SSM snapshots, FP8 and NVFP4 KV caches, per-batch CUDA graphs, chunked prefill, tool calling in three formats, and RoCEv2-backed expert parallelism for models beyond a single GB10.

How to read this book

The TOC is linear but the parts are independent. If you came here to run a model, skip straight to Quickstart. If you came here to understand why Atlas is fast, read Philosophy first — the rest of the book is a working demonstration of that claim.

Philosophy

Atlas is built for broad hardware and model support — the kind of matrix that today includes GB10, tomorrow includes H100, B200, Apple Silicon, AMD MI300X and Intel GPUs, and a long tail of model architectures on top of each. The question the project starts from is the one every inference engine eventually has to answer:

How do you cover a large matrix of (Hardware, Model_q) targets and run each one at the hardware's theoretical peak?

Existing general-purpose engines answer by absorbing genericity into the kernel: templated CUDA, JIT compilation, runtime shape branching. Atlas answers the opposite way — specialize per target, but design the abstractions so that specialization can scale.

This is our version of AI Kernel HyperCompiling. Every chapter of this book is a consequence of taking that answer seriously.

The trade general frameworks make

vLLM, TensorRT-LLM, SGLang, and the other mainstream engines support thousands of models across dozens of GPU generations with a single binary. That is a real, useful thing. The cost of doing it well — a cost they pay every release — is a layer of abstraction between the kernel and the hardware: a templating engine, a just-in-time compiler, or a broadly-parameterized CUDA kernel that branches on shape, dtype, and arch.

Every one of those layers trims a few percent off peak throughput. Added up, on a specific hybrid-SSM/attention/MoE model running on a specific GPU, those trims are how a 3.6× gap opens up.

Atlas does not try to close that gap inside a general framework. We reject the framing — and we refuse to let specialization shrink the scope of what we support. Both at once.

Abstractions designed for many targets

The way Atlas gets broad support is by putting the genericity above the kernel layer, not inside it. Three traits do the load-bearing work:

• ComputeTarget (atlas-core/src/compute.rs) — the build-time trait. Given a hardware vendor and an architecture flag, it knows how to invoke the right compiler (nvcc, hipcc, xcrun metal, …) to turn a source file into a binary module. Adding a new hardware vendor is one impl ComputeTarget.
• GpuBackend (spark-runtime/src/gpu.rs) — the runtime trait. 31 methods cover memory, kernel launch, streams, events, and graphs. The model code, the scheduler, the HTTP server — none of them know whether they're running on CUDA, Metal, or HIP. They hold a &dyn GpuBackend. A new backend is one impl GpuBackend.
• CommBackend (spark-comm/src/lib.rs) — the multi-GPU trait. all_reduce, all_gather, reduce_scatter, broadcast. NCCL ships today; HIP's RCCL and Metal's MPS collective ops would drop into the same shape.

The Vendor enum in atlas-core already enumerates Nvidia, Amd, Apple, Intel. Nothing about the engine above the trait layer is NVIDIA-specific. The multi-vendor design is in the code today; the first (hw, model, quant) set we happened to ship was for GB10 because that's the hardware on our desks. Porting to H100, B200, MI300X, M4 Ultra, or Arc A770 means implementing two traits and writing the kernel source — the full walkthrough is in the Adding a new hardware target guide.

How specialization scales

The question this framing forces — and the question the rest of the codebase answers — is: how do you keep adding targets without turning into the general framework you rejected?

The Atlas answer has two parts:

1. Hyperoptimize in isolation. Every (H, M_q) target lives in its own directory: kernels/<hw>/<model>/<quant>/. The kernels there can use any tiling strategy, any shared-memory layout, any MMA instruction mix — they cannot accidentally slow down another target because they share no code with one. This is the opposite of the templating approach: instead of one kernel that branches, we have N kernels that each do exactly one thing. Adding a new target is a new directory; it is physically impossible for it to regress an existing one.
2. Share abstractions, not kernels. What is shared — the GpuBackend trait, the ComputeTarget build-time trait, the layer factory in spark-model — is abstraction above the kernel level. The shared code knows about "launch a kernel" and "allocate GPU memory"; it does not know, and does not want to know, what the kernel inside is doing. New hardware plugs in at the trait layer; new models plug in at the ModelWeightLoader trait — neither disturbs the other axis.

Adding a new model works the same way. One new ModelWeightLoader impl, one match arm in spark-model/src/factory.rs. The KV cache, buffer arena, scheduler, and HTTP server are model-agnostic — they do not need to change to support a fundamentally different architecture. Qwen3.5 hybrid SSM+attention+MoE, Nemotron-H Mamba-2, MiniMax 256-expert sigmoid MoE, Gemma-4 sliding+full alternating attention, Qwen3-VL vision all coexist in one binary today, sharing zero kernel code with each other.

Why specialization is finally practical

Writing specialized kernels was prohibitively expensive. One-off work by a human CUDA engineer, non-transferable, bit-rots as the hardware changes. The reason general frameworks won for a decade is that the specialist approach had no path to scale — you could not afford to write a new kernel for every GPU generation and every quantization scheme.

That has changed. AI-assisted kernel engineering — profiling a kernel, proposing tiling experiments, verifying correctness against a reference — is now good enough that we can dedicate real effort to every target. One reason Atlas exists is to prove this at scale: every kernel in the repo is AI-written, human-reviewed, benchmark-verified. We explicitly want new PRs to be AI-generated. Human-only contributions get reviewed by AI.

The specialization thesis does not require a superhuman human sitting behind every kernel. It requires a pipeline where specializing is cheap enough to do twelve times, then fifty, then a hundred — across hardware and across models.

What this means for you

• If you're an operator — expect Atlas to be fast on every target we've shipped. The matrix grows with each release; if your model or GPU is not in the matrix yet, it is a well-scoped piece of work to add it, not an architectural impossibility.
• If you're a model author — the trait you need to implement is ModelWeightLoader. Everything downstream is model-agnostic and will not change when you add a new architecture. See spark-model.
• If you're a kernel engineer — the directory you'll live in is kernels/<hw>/<model>/<quant>/. The abstractions above you exist to stay out of your way. Your job is to make that one (H, M_q) tuple run as close to silicon peak as you can. See CUDA Kernel Engineering.
• If you're a hardware vendor — the abstractions that let a new GPU family plug in are two traits (ComputeTarget, GpuBackend) plus kernel source. See Adding a new hardware target in the README.

The next chapter gets you running on the hardware we've shipped. The rest of the book earns the claim that specialization, done with the right abstractions, scales — across GPUs, across models, and across quantization schemes — without giving up a single percent of peak throughput.

A formal lens

If you want the same argument in the language of category theory — the target matrix as a product, the kernel registry as a coproduct, the trait layer as an algebraic theory, general frameworks as a factoring Atlas refuses — see the appendix A Category-Theoretic Perspective. It is optional reading; nothing else in the book depends on it.

Installation

Atlas ships as a single Docker image that contains the release binary plus every compiled (GB10, model, quant) PTX module — 22 target sets today, one per kernels/gb10/<model>/<quant>/ directory. There is no "install Atlas + download kernels" step — the kernels are baked in.

Hardware prerequisites

Atlas is designed for broad hardware support — the engine is vendor-agnostic above the kernel layer (ComputeTarget at build time, GpuBackend at runtime, CommBackend for collectives) and new hardware plugs in at the trait layer. The first shipped target is NVIDIA GB10 (SM121) — the Grace-Blackwell Superchip in the NVIDIA DGX Spark workstation. To run the shipped image you need:

• A DGX Spark (or any GB10-based system) with 119.7 GB of unified GPU memory
• NVIDIA driver supporting CUDA 13.0 or later
• docker with --gpus all support (recent nvidia-container-toolkit)
• Internet access for the first model download; models are cached under ~/.cache/huggingface after that

Other NVIDIA GPUs (H100, B200) and other vendors (AMD, Apple, Intel) are on the roadmap rather than in the shipped image. The PTX that ships today is compiled with -arch=sm_121 using SM121-specific tile shapes and a software E2M1 conversion — none of that is architectural, it's just the first target we hyperoptimized. Adding a new hardware target is two trait impls plus kernel source; the Adding a new hardware target guide in the README walks through an Apple Metal example end to end.

Pull the image

docker pull avarok/atlas-gb10:latest

The image contains the Rust release binary, all 22 PTX module sets, tokenizer dependencies, and the nvidia-container-runtime library surfaces. No Python, no CUDA toolkit.

Bring your own weights

Atlas loads HuggingFace safetensors directly. The image does not ship model weights. On first run, the binary resolves a HuggingFace model ID (e.g. Sehyo/Qwen3.5-35B-A3B-NVFP4) against ~/.cache/huggingface/hub — download the weights once with the hf CLI or let the server download-on-miss:

pip install -U huggingface_hub
hf download Sehyo/Qwen3.5-35B-A3B-NVFP4

The command is hf, not huggingface-cli: huggingface_hub 1.0 renamed the binary, and on 1.16+ the old name is gone entirely. Older docs and scripts still say huggingface-cli download …, which now fails with "command not found".

Mount the cache directory into the container:

-v ~/.cache/huggingface:/root/.cache/huggingface

Build from source (optional)

You only need to build from source if you are modifying Atlas. The rust-toolchain.toml pins stable; CUDA 13.0+ with nvcc on PATH (or CUDA_HOME set) is required for a real build. Clippy and fmt can run without CUDA via ATLAS_SKIP_BUILD=1.

git clone https://github.com/Avarok-Cybersecurity/atlas.git
cd atlas

# Full build — compiles every (gb10, model, quant) target (~6 min)
docker build -f docker/gb10/Dockerfile -t atlas-gb10 .

# Rust-only check (no CUDA). CUDARC_CUDA_VERSION is needed alongside
# ATLAS_SKIP_BUILD: without it cudarc's build script shells out to
# `nvcc --version` and panics on a host that has no CUDA toolkit.
# This pair is exactly what ci.yml exports. Deny-warnings comes from
# [workspace.lints], so `-- -Dwarnings` is not needed and CI does not pass it.
ATLAS_SKIP_BUILD=1 CUDARC_CUDA_VERSION=13000 cargo clippy --workspace --tests
cargo fmt --all -- --check

# Unit tests (uses MockGpuBackend; no GPU required)
cargo test --release

# Integration tests (require GPU + weights)
cargo test -p spark-server --release -- --ignored

The build system reads kernels/gb10/HARDWARE.toml for architecture flags, enumerates every (model, quant) subdirectory that matches the ATLAS_TARGET_* wildcards, compiles each .cu source file through nvcc, and emits a single target_ptx.rs that the atlas-kernels crate embeds in the final binary. Zero runtime compilation.

Verify the install

docker run --rm --gpus all avarok/atlas-gb10:latest --version
# → spark 1.0.0-beta-preview   (the workspace version in Cargo.toml)
docker run --rm --gpus all avarok/atlas-gb10:latest --help | head -20

If --version errors with "no compatible GPU", the nvidia-container-toolkit is not picking up the device. Check docker info | grep -i runtime and nvidia-smi on the host.

You are now ready for the Quickstart.

Quickstart

Goal: first successful chat completion in under five minutes, against the flagship Qwen3.5-35B-A3B model running at 131 tok/s on a single GB10.

1. Start the server

sudo docker run -d \
  --name atlas-35b \
  --network host --gpus all --ipc=host \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  avarok/atlas-gb10:latest \
  serve Sehyo/Qwen3.5-35B-A3B-NVFP4 \
    --port 8888 \
    --max-seq-len 8192 \
    --kv-cache-dtype nvfp4 \
    --gpu-memory-utilization 0.88 \
    --scheduling-policy slai \
    --speculative \
    --mtp-quantization nvfp4

What's happening:

• The container binds the host's network (--network host) so port 8888 is reachable on the host directly. --gpus all grants the GB10, --ipc=host enables shared memory for larger KV buffers.
• serve <model-id> selects the model; the binary auto-detects model_type from config.json and picks the matching kernel target.
• --kv-cache-dtype nvfp4 keeps the KV cache in 4-bit E2M1 — halves memory vs FP8, with no measurable coherence loss for Qwen3.5.
• --speculative --mtp-quantization nvfp4 turns on Multi-Token Prediction speculative decoding with the NVFP4 MTP head that ships in the checkpoint. This is the change that takes throughput from ~70 tok/s to ~131 tok/s.
• --scheduling-policy slai enables SLO-aware scheduling — prioritises decode steps approaching their TBT deadline.

First start-up takes 2–5 minutes: the loader reads 15–40 GB of safetensors through the O_DIRECT fast path, CUDA graphs are captured for every batch size, and the HTTP server binds. Watch the log:

sudo docker logs -f atlas-35b

You'll see loaded 125000 tensors in 34s, then captured graph for batch=1, then finally Listening on 127.0.0.1:8888. The server does not accept requests before that last line appears.

The address in that line is whatever --bind resolved to, and --bind defaults to 127.0.0.1 — so the default run is loopback-only and logs a second line saying so. You get Listening on 0.0.0.0:8888 only if you pass --bind 0.0.0.0, which also emits a LAN-exposure warning. Note the capital L: grepping for listening on 0.0.0.0 matches nothing.

2. Send a request

Once the server logs listening, a standard OpenAI request works:

curl -s http://localhost:8888/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "atlas",
    "messages": [{"role": "user", "content": "Explain the key idea behind speculative decoding in one paragraph."}],
    "max_tokens": 256
  }'

Atlas accepts any model string — it serves exactly one model per container, so the field is ignored. Use the real HF id if you want round-tripping through OpenAI clients to feel natural.

3. Stream tokens

curl -sN http://localhost:8888/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "atlas",
    "messages": [{"role": "user", "content": "Write a short poem about kernels."}],
    "max_tokens": 200,
    "stream": true
  }'

Each chunk is a standard data: {...} SSE frame with choices[0].delta.content. Tool calls stream as choices[0].delta.tool_calls chunks in the same format OpenAI emits.

4. From Python (OpenAI SDK)

from openai import OpenAI

client = OpenAI(base_url="http://localhost:8888/v1", api_key="unused")

stream = client.chat.completions.create(
    model="atlas",
    messages=[{"role": "user", "content": "Hello!"}],
    max_tokens=200,
    stream=True,
)
for chunk in stream:
    if chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end="", flush=True)

The same client works with Open WebUI — set Base URL: http://<host>:8888/v1, API key sk-dummy.

5. Stop

sudo docker stop atlas-35b && sudo docker rm atlas-35b

Troubleshooting

• error: out of memory during start-up — drop --gpu-memory-utilization to 0.85, or --max-seq-len to 4096. 35B has the headroom; it's usually leaked GPU state from a previous container. nvidia-smi should show ~0 MB used before starting.
• Server logs loaded 0 tensors — your HF cache is empty or the path is wrong. Verify with ls ~/.cache/huggingface/hub/models--Sehyo--Qwen3.5-35B-A3B-NVFP4.
• Connection refused on port 8888 — the server hasn't finished initialising. Watch the log; Listening on <bind>:<port> is the readiness marker. If it has printed and you're still refused from another machine, that's the --bind 127.0.0.1 default, not a start-up problem — see Quickstart §Network exposure.
• Tokens are gibberish — almost always a model/loader mismatch. Check that the HF model id in the command line matches the cached directory. If the kernel target the binary picked is wrong (unlikely — Atlas logs it on startup), open an issue; Atlas's house rule is never blame the model, always find the Atlas bug.

Next: pick a different model from Supported Models, or dive into the Architecture.

Supported Models

Twelve (GB10, model, quant) targets ship in the default image today. One multi-model binary, one Docker image, one serve <hf-id> command per model. The binary reads the model's config.json, computes the canonical model_type, and dispatches to the matching kernel set at startup.

The matrix

Throughput figures are p50 single-request decode on a short prompt (max_tokens ≤ 128, temperature ≤ 0.1). The "Best tok/s" column reflects the flag set that wins for that model (e.g. MTP enabled where supported, NVFP4 KV cache for Qwen models). Full numbers for alternative flag combinations are in Benchmarking.

Recently added (not yet tabulated): DeepSeek-V4-Flash — MLA + MoE + CSA/HCA hybrid attention + mHC, with native MXFP4 (E8M0) routed-expert loading and Phase-K E8M0 GEMM kernels — landed end-to-end on GB10 in #293. Its model_type dispatches through factory.rs (deepseek_v4).

How to pick

• Fastest — Qwen3.5-35B-A3B with MTP. The flagship. 131 tok/s.
• Largest on one node — Qwen3-Next-80B-A3B or Nemotron-3-Super-120B. Both fit in 119.7 GB with FP8 KV.
• Vision — Qwen3-VL-30B (pure attention) or Qwen3.6-35B (hybrid SSM + vision).
• Largest overall — MiniMax-M2.7 at 229B / 256-expert MoE, or Qwen3.5-122B-A10B. Both require EP=2 (two GB10 nodes over RoCEv2).
• Long reasoning traces — any Qwen3.5 model; thinking budget is configurable via --max-thinking-budget.
• Function calling — all Qwen-family and Nemotron models support OpenAI-style tools. See Tool Calling.

Per-model serve commands

Every command below uses avarok/atlas-gb10:latest, --network host --gpus all --ipc=host, and the -v ~/.cache/huggingface:/root/.cache/huggingface volume mount — omitted here for readability. Full copy-pasteable commands are in QUICKSTART.md.

Qwen3.5-35B-A3B (flagship)

serve Sehyo/Qwen3.5-35B-A3B-NVFP4 \
  --max-seq-len 8192 --kv-cache-dtype nvfp4 \
  --scheduling-policy slai --speculative --mtp-quantization nvfp4

Qwen3-Next-80B-A3B (largest single-node MTP)

serve nvidia/Qwen3-Next-80B-A3B-Instruct-NVFP4 \
  --max-seq-len 8192 --kv-cache-dtype nvfp4 \
  --speculative --mtp-quantization nvfp4

Qwen3-VL-30B (vision)

serve ig1/Qwen3-VL-30B-A3B-Instruct-NVFP4 \
  --max-seq-len 32768 --kv-cache-dtype nvfp4

Send images in OpenAI content-parts format (see Tool Calling & Streaming).

Nemotron-3-Nano-30B (Mamba-2 hybrid)

serve nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4 \
  --max-seq-len 8192 --kv-cache-dtype nvfp4

Qwen3.5-122B-A10B — EP=2

Two nodes connected via RoCEv2 (head on <head-ip>, worker on <worker-ip>). The canonical launcher is scripts/start-ep2.sh. See Multi-GPU & EP=2 for the full flow and the critical MTP-flag symmetry rule between head and worker.

Single-node 122B (tight budget)

serve Sehyo/Qwen3.5-122B-A10B-NVFP4 \
  --kv-cache-dtype fp8 --kv-high-precision-layers 2 \
  --max-batch-size 1 --max-prefill-tokens 2048 --oom-guard-mb 512

~32 tok/s. The --kv-high-precision-layers 2 keeps the first and last two attention layers at BF16 — costs a few hundred MB, buys coherence at very long context.

Adding a new model

The entire model-specific surface is one new ModelWeightLoader impl and one match arm in spark-model/src/factory.rs. The KV cache, buffer arena, scheduler, and HTTP server are all model-agnostic. The full walkthrough with a live example (Mistral-Small-4) is in the repo's Adding a new model guide. The chapter on spark-model covers the trait shape.

Troubleshooting

Problems that stop atlasctl run before a model ever loads, and what to do about each. If your problem is a model that starts and then misbehaves, the Quickstart has a section for that.

permission denied talking to Docker

permission denied while trying to connect to the Docker daemon socket
at unix:///var/run/docker.sock

What it means. Docker is installed and running. The daemon answered your request and refused it, because your user is not in the docker group. This is the single most common reason a fresh DGX Spark cannot launch a model, and it has nothing to do with the hardware.

The fix, once:

sudo usermod -aG docker $USER
newgrp docker          # or log out and back in

usermod changes your groups; it does not change the groups of a shell that is already running. newgrp docker starts a shell that has the new membership, and logging out and back in achieves the same thing for every shell.

Confirm it took:

docker info --format '{{.ServerVersion}}'

If that prints a version, atlasctl run will work. You do not need to restart the Atlas agent — it re-checks its own capability, so the control plane stops saying "this machine cannot run models" within a few seconds.

Do not use sudo atlasctl

It appears to work, and it is the wrong move:

• the model runs as root, and so does everything the container does;
• ~/.atlas collects root-owned files that your normal user then cannot read, so the next unprivileged atlasctl run fails in a way that looks unrelated;
• sudo uses root's PATH, so atlasctl is frequently "not found" even though which atlasctl finds it for you.

Fix the group membership instead. It is a one-time change.

Rootless Docker and Podman

If you run rootless Docker, the socket is under $XDG_RUNTIME_DIR rather than /var/run, and group membership is not the issue — check that DOCKER_HOST=unix://$XDG_RUNTIME_DIR/docker.sock is exported. Docker's own post-installation guide is the canonical reference for both the group and the rootless setups.

Cannot connect to the Docker daemon

Cannot connect to the Docker daemon at unix:///var/run/docker.sock.
Is the docker daemon running?

Nothing is listening. Start it:

sudo systemctl start docker
sudo systemctl enable docker    # so it survives a reboot

docker: command not found

Docker is not installed, or not on this shell's PATH. atlasctl list, atlasctl show and atlasctl run --print all work without a container engine — only atlasctl run needs one. See Installation for the prerequisites, including the NVIDIA container runtime.

atlasctl is not on PATH

The installer puts the binary in ~/.local/bin and tells you if that directory is not on your PATH. Add it:

echo 'export PATH="$PATH:$HOME/.local/bin"' >> ~/.bashrc

Note the absence of a trailing slash. A PATH entry that ends in / is legal and works, but it makes which print a doubled slash — /home/you/.local/bin//atlasctl — which looks like a bug and is not one: which joins the PATH entry to the program name without checking whether the entry already ends in a separator.

The control plane says this machine cannot run models

The browser is repeating what the local agent told it. The agent decides by running docker info, so the cause is almost always one of the Docker problems above — most often the permission one. Fix that and the page corrects itself within a few seconds; there is no need to restart the agent.

To see the agent's own view:

atlasctl doctor

doctor reports each check and exits non-zero if any of them found a problem, so it is safe to use in a script.

Philosophy: AI Kernel HyperCompiling

Part I's Philosophy chapter answered why Atlas specializes. This chapter answers how the specialization thesis forces specific design choices in the code, and what you should see — or, if you're writing a PR, what you should preserve — when you read the codebase.

The single design rule that every choice below derives from is:

Specialization is a directory, not a template. Everything that varies across (Hardware, Model_q) targets lives in its own directory. Everything shared across them is abstraction above the kernel layer, not parameters inside one.

Consequence 1: the kernel tree is a coordinate system

kernels/
  gb10/                               # Hardware
    HARDWARE.toml                     # vendor, arch, memory specs
    common/
      KERNEL.toml
      *.cu                            # the shared baseline — 160 files today
    qwen3-next-80b-a3b/               # Model
      MODEL.toml                      # layer counts, sampling defaults
      nvfp4/                          # Quantization
        KERNEL.toml                   # compiler flags
        *.cu                          # only the files this target overrides

Three levels of directory, over a common/ baseline. The shape is deliberate: a leaf directory owns exactly one (H, M_q) target, and a leaf file shadows its same-stem namesake in common/ (atlas-kernels/build.rs::collect_cu_files). Shadowing is whole-file, not per-symbol.

So a leaf holds only its divergences, not a full kernel set — qwen3.6-27b/nvfp4 carries 11 .cu, qwen3.6-35b-a3b/nvfp4 carries 5, qwen3-next-80b-a3b/nvfp4 carries 3, over the 160 in common/. Two leaves therefore do share source for everything neither of them overrides; what is guaranteed is that where a target does diverge, it diverges in a file nothing else compiles. When we say "Atlas ships N targets," we mean N independent leaves — 22 of them on GB10 today.

The corollary is a real failure class: because shadowing is whole-file, a shadow that has quietly become identical to common/ overrides nothing while masking every later common/ improvement, and two models keeping private byte-identical copies of the same shadow will drift apart. CI's kernel-structure job (scripts/check_kernel_shadows.py) rejects both; the sanctioned way to share one file between leaves is a relative symlink.

The three .toml files are the only metadata the build system consumes. HARDWARE.toml tells atlas-kernels/build.rs which ComputeTarget impl to use (nvidia, amd, apple, intel) and what arch flag to pass the compiler. MODEL.toml is the per-model behavior SSOT — sampling presets, thinking budgets, tool-call parser defaults. KERNEL.toml overrides compiler flags and module names.

Adding a model or a hardware target is, at the file-system level, creating a new directory. No code elsewhere in the repository needs to move.

Consequence 2: the runtime crate structure mirrors the axis split

Read the workspace Cargo.toml and you'll see nineteen workspace members. Group them by what axis of variation they insulate:

Each crate has exactly one reason to change. A new GPU vendor never touches spark-model. A new model family never touches spark-runtime. A new quantization scheme touches spark-model's format modules and atlas-kernels, but not the layer code. This orthogonality is not a happy accident of the crate layout — it is the architectural consequence of the specialization thesis.

Consequence 3: SBIO — business logic never touches I/O

Business logic — the layer code in spark-model, the scheduler in spark-server — never calls CUDA APIs, never opens a socket, never reads a file. Every such operation goes through a trait:

• GPU memory, launches, graphs → GpuBackend
• Collective comms → CommBackend
• Weight loading I/O → WeightStore (wraps safetensors + O_DIRECT)
• HTTP responses → axum handlers, tested against a mock channel

This is what the user instructions call SBIO (Separation of Business logic from I/O). The payoff is that 80%+ of the codebase is unit-testable without a GPU. MockGpuBackend records launches but does not execute them. SingleGpuBackend is a no-op CommBackend for single-GPU runs. The SBIO chapter shows the pattern in detail.

Consequence 4: zero runtime compilation

Every general-purpose framework has, somewhere, a codepath that compiles kernels at runtime. PyTorch has torch.compile. vLLM has Triton JIT. TensorRT-LLM has TRT engine builds. Each of those is a slow path the first time you hit a new shape, and an ongoing operational surface the ops team has to manage (cache directories, warm-up scripts, cold-start budgets).

Atlas has none of it. atlas-kernels/build.rs enumerates every (H, M_q) target matching the ATLAS_TARGET_* env vars, compiles every .cu file for every matching target, and emits one auto-generated target_ptx.rs that is include!'d into the crate. The release binary contains every PTX module we ship. Startup is "mmap the binary, upload PTX to the GPU, capture CUDA graphs for a handful of batch sizes, done".

This is what "embedded in the binary" means throughout the book. It is the concrete mechanism by which specialization does not cost operator pain.

Consequence 5: one binary per installation, N kernel sets

You deploy one Docker image. It contains one spark-server binary. It contains 22 (today) (gb10, model, quant) PTX sets embedded in that binary. At startup, the binary reads the model's config.json, computes the canonical model_type, looks up the matching KernelTarget, and uses that set.

The knobs that let this scale:

• ATLAS_TARGET_*=* at build time — compiles every matching target. The default image sets everything to * and ships the lot.
• ATLAS_TARGET_HW=gb10 ATLAS_TARGET_MODEL=qwen3.5-35b-a3b ATLAS_TARGET_QUANT=nvfp4 — compiles exactly one target. Used for per-model slim images in docker/gb10/<model>/.
• ATLAS_SKIP_BUILD=1 — emits a stub target_ptx.rs so that clippy, fmt, check, and any non-GPU test can run on a vanilla Linux host.

The same image works across all supported targets. The startup dispatcher picks the right kernels. Operators don't manage a kernel cache; they don't warm a JIT; they don't think about it.

What the rest of this book is about

Every subsequent chapter is an elaboration of one of the consequences above. The workspace layout chapter walks the directory tree. The dispatch chapter traces a single request from HTTP to kernel launch. The SBIO chapter shows how the testability claim actually holds.

The deep-dive chapters in Part IV show what the kernels look like — what a hand-tuned kernel set per target buys you, and how you'd write new ones when you're porting Atlas to your own (H, M_q) target.

Reading the architecture categorically

The design choices above have precise names in category theory. The target set 𝒯 = Hw × Mod × Quant is a categorical product; the crate split is that product made syntactically real, which is why orthogonality of axes is a structural fact and not a convention. The kernel registry is a coproduct (disjoint union of per-target PTX sets), which is why adding a summand cannot regress existing summands. The GpuBackend trait defines an algebraic theory with two ship-worthy models — AtlasCudaBackend and MockGpuBackend — and that is what makes the test suite runnable without a GPU. A general framework is, in this vocabulary, an engine that factors Kernels : 𝒯 → 𝐒𝐞𝐭 through a smaller "essence" category; Atlas refuses the factoring, and the 3.6× gap against vLLM is the cost of the factoring that Atlas does not pay.

The appendix A Category-Theoretic Perspective works through each of these structures at appendix length. It is a design reference, not a prerequisite.

Workspace Layout

Atlas is a nineteen-member Cargo workspace plus a build-time kernel tree (count them in the root Cargo.toml members list). This chapter maps every top-level directory to its role, and the crates to the axes of variation they each insulate.

Repository tree (top level)

atlas/
├── README.md                     headline, benchmarks, porting guides
├── QUICKSTART.md                 per-model Docker recipes
├── CONTRIBUTING.md, AGENTS.md    contributor workflow
├── SECURITY.md                   disclosure
├── CLA.md                        contributor license agreement
├── LICENSE                       AGPL-3.0-only
├── Cargo.toml                    workspace root (12 members)
├── Cargo.lock
├── rust-toolchain.toml           pins stable
├── deny.toml                     cargo-deny allow/deny lists
├── crates/                       Rust source for every crate
├── kernels/                      CUDA source, organized as (hw, model, quant)
├── docker/                       per-hardware Dockerfiles
├── scripts/                      bench, model-sweep, release helpers
├── tests/                        cross-crate integration tests (run_all_models.py lives here)
├── docs/                         design notes, history, release notes
├── paper/                        LaTeX paper (ArXiv)
├── jinja-templates/              chat templates for models that need custom ones
├── bench/                        stable benchmark harness outputs (tracked)
├── book/                         this book (mdBook source)
└── vendor/                       vendored deps (e.g. xgrammar-rs)

The workspace members

Cargo.toml lists:

members = [
    "crates/atlas-core",
    "crates/atlas-kernels",
    "crates/atlas-plugin",
    "crates/atlas-tier",
    "crates/atlas-rdma",
    "crates/spark-runtime",
    "crates/spark-comm",
    "crates/spark-model",
    "crates/spark-nllb",
    "crates/spark-server",
    "crates/spark-storage",
    "crates/atlas-spark-bench",
    "crates/cufile-sys",
    "crates/xgrammar",
]

Each is its own crate with its own Cargo.toml, its own unit tests, and its own responsibility:

The dependency graph runs strictly downward in the table above — atlas-core has no internal deps, every crate above it builds on crates below. There are no cycles.

The kernel tree

kernels/
└── gb10/                                        # One directory per hardware target
    ├── HARDWARE.toml                            # vendor, arch, memory specs
    ├── qwen3-next-80b-a3b/                      # One directory per model target
    │   ├── MODEL.toml                           # layer counts, sampling presets, behavior
    │   └── nvfp4/                               # One directory per quantization target
    │       ├── KERNEL.toml                      # compile flags, module name overrides
    │       └── *.cu                             # ~35 hand-written CUDA kernels
    ├── qwen3.5-35b-a3b/
    │   └── nvfp4/
    │       └── *.cu
    ├── qwen3.6-35b-a3b/
    │   └── fp8/
    ├── nemotron-3-nano-30b-a3b/
    │   └── nvfp4/
    ├── mistral-small-4-119b/
    │   └── nvfp4/
    ├── minimax-m2-229b/
    │   └── nvfp4/
    └── ... (one leaf per (model, quant) target — 22 under kernels/gb10/ today)

Every leaf directory is a fully self-contained (gb10, model, quant) target. The kernels inside a leaf can use any tile shape, any register budget, any shared-memory layout — they are physically incapable of regressing a different target.

This is the mechanism that makes kernels/ a scalable structure. Adding a new GPU is kernels/<new-hw>/. Adding a new model is kernels/<hw>/<new-model>/. Adding a new quantization is kernels/<hw>/<model>/<new-quant>/. Nothing else moves.

Docker layout

docker/
├── gb10/
│   ├── Dockerfile                         multi-model image — compiles every target
│   ├── qwen3-next-80b-a3b/nvfp4/          per-model slim image
│   ├── qwen3.5-35b-a3b/nvfp4/
│   └── ... (one slim Dockerfile per supported model)
└── docker-guide.md                        build + run instructions

The multi-model Dockerfile at docker/gb10/Dockerfile is what ships as avarok/atlas-gb10:latest. Per-model Dockerfiles exist for operators who want a smaller image containing only one target — the kernel registry still uses KernelTarget at runtime, but only one target set is baked in.

Docs, design records, history, releases

Inside docs/:

• adr/ — architecture decision records (licensing, pure-Rust, hybrid SSM/attention, NVFP4/FP8 quantization, TP/EP composition, EP batched decode, etc.). Treat these as the long-form rationale behind code changes; commit messages are deliberately terse and point here. Top-level notes like ARCHITECTURE.md, ATLAS_KERNELS.md, and HARDWARE.md sit alongside them.
• ATLAS_SPARK_JOURNEY.md — benchmark journey and retrospective across the Spark line. Useful context, but not a contract.
• releases/ — human-readable release notes keyed by release (README.md plus per-release files).

The book you're reading in book/ synthesises all of this into a single narrative — it is not a canonical rewrite of those documents. The design records in docs/adr/ remain the authoritative reference and the book links to them directly from the deep-dive chapters.

What changes when you add a…

Each row touches a small, bounded set of files. That bounded-ness is the architectural payoff of the workspace being split along axes of variation. Read Kernel Dispatch Pipeline next to see the runtime side, or SBIO to see how the trait layering makes the whole thing testable without a GPU.

Kernel Dispatch Pipeline

One Atlas binary contains kernels for every (Hardware, Model, Quantization) target it was built for. This chapter traces a single chat completion from the moment the HTTP request arrives to the moment a kernel launches on the GPU, so you know exactly where each piece of dispatch lives.

The high-level flow

             1. HTTP request                                    7. Kernel launch
┌──────────────┐   ┌──────────┐   ┌───────────┐   ┌───────────┐   ┌──────────┐
│ OpenAI       │──►│ axum     │──►│ scheduler │──►│ engine    │──►│ PTX on   │
│ client       │   │ (server) │   │ (server)  │   │ (model)   │   │ GPU      │
└──────────────┘   └──────────┘   └───────────┘   └───────────┘   └──────────┘
                                                        │               ▲
                                                        ▼               │
                                                 ┌──────────────────────┴───┐
                                                 │ KernelTarget → PtxModule │
                                                 │ (atlas-kernels)          │
                                                 └──────────────────────────┘

The dispatch decisions happen in two distinct phases:

• Build time — atlas-kernels/build.rs decides which PTX to embed.
• Startup — spark-server::main decides which embedded PTX to upload to the GPU based on the model being served.

After startup, the fast path is deterministic: a layer's forward(ctx) always calls the same KernelHandles, always on the same GPU stream, always in the same order. There is no per-request dispatch decision. This is the payoff of the specialization thesis — no branching, no polymorphism across kernel variants, no cache miss.

Phase 1 — build time: which PTX gets embedded

atlas-kernels/build.rs runs during cargo build. Its job:

1. Read the three wildcards:
   • ATLAS_TARGET_HW (default gb10; accepts *)
   • ATLAS_TARGET_MODEL (default * — all)
   • ATLAS_TARGET_QUANT (default * — all)
2. Enumerate kernels/<hw>/<model>/<quant>/ leaves matching the wildcards.
3. For each leaf, read HARDWARE.toml to learn the vendor, and call resolve_compute_target(vendor) to get a Box<dyn ComputeTarget>:
   • Vendor::Nvidia → NvidiaTarget { nvcc }
   • Vendor::Apple → AppleTarget { xcrun } (planned)
   • Vendor::Amd → AmdTarget { hipcc } (planned)
4. Call compute_target.compile(source, out, arch, extra_flags) on every .cu / .metal / .hip file in the leaf.
5. Parse KERNEL.toml for module-name overrides (some kernels are compiled from e2m1_branchless.cu but exposed at runtime as the e2m1 module).
6. Emit an auto-generated Rust file, $OUT_DIR/target_ptx.rs, that include!()'s back into atlas-kernels/src/lib.rs. The generated file contains one pub static PTX_<TARGET>: &[PtxModule] per target plus an all_ptx_sets() function that returns the whole set.

The output is one single PTX set per target, embedded in the final spark-server binary as a byte slice. This is why "one Docker image, one binary, zero runtime compilation" is true.

ATLAS_SKIP_BUILD=1 short-circuits the whole phase: build.rs emits a stub target_ptx.rs with empty constants so that clippy, fmt, and non-GPU tests can compile on a Linux host with no nvcc. The CI in .github/workflows/ci.yml uses this.

Phase 2 — startup: which embedded PTX gets uploaded

When the user runs spark serve <model-id>, spark-server/src/main.rs does the following, roughly in order:

1. Parse the model config. atlas_core::config::ModelConfig::from_hf(&model_path) reads config.json and its nested text/vision configs.
2. Canonicalize model_type. Lowercase, replace - and . with _. "Qwen3.5_NextForCausalLM" becomes "qwen3_5_next_for_causal_lm". This is the key we dispatch on.
3. Resolve the KernelTarget. Given model_type and the selected quantization (from config or --kv-cache-dtype when overriding), atlas-kernels::select_target(hw, model, quant) looks up the matching KernelTarget. Fail fast with a clear error if there's no match.
4. Instantiate the GpuBackend. AtlasCudaBackend::new(gpu_ordinal, &ptx_set.modules) uploads every embedded PTX module for the chosen target to the GPU, via cuModuleLoadData. Kernel handles are cached per (module_name, function_name) pair.
5. Instantiate the ModelWeightLoader. spark_model::factory::loader_for_config(&config) matches on the canonical model_type and returns Box<dyn ModelWeightLoader>.
6. Load weights. The loader translates HF weight names (model.layers.0.self_attn.q_proj.weight) into Atlas layer types (Qwen3AttentionLayer), going through WeightStore (the O_DIRECT fast path) and the quantization helpers in spark_model::weight_map.
7. Build layer trait objects. Each loaded layer becomes a Box<dyn TransformerLayer> stored in the InferenceEngine.
8. Capture CUDA graphs. For each supported batch size, engine.capture_graph(bs) runs a single decode step inside a graph region. Subsequent decodes replay the graph — one GPU launch for the whole forward pass.
9. Bind the HTTP endpoint. axum::Router::new()...serve(&addr) starts listening.

At this point dispatch is frozen. Every request goes through the same kernels, the same graph, the same streams.

Phase 3 — per-request path

POST /v1/chat/completions
 │
 ▼
spark_server::api::chat_completions    (axum handler)
 │
 ▼  1. Apply jinja chat template
 │  2. Tokenize
 │  3. Enqueue Request {prompt_ids, sampling, stream?, tools?}
 │
 ▼
spark_server::scheduler                (SLAI or FIFO)
 │
 ▼  1. Allocate KV pages for prefix
 │  2. Chunked prefill through InferenceEngine
 │  3. Enter the decode loop
 │
 ▼
spark_model::engine::InferenceEngine::decode_step
 │
 ▼  for layer in layers:
 │      layer.forward(&ctx)      ← dyn dispatch, one per layer
 │          └─ calls into ops.rs kernel launches
 │              └─ GpuBackend::launch(KernelHandle, grid, block, args, stream)
 │                  └─ CUDA cuLaunchKernel    (PTX on GPU)
 │
 ▼
Sampler                              (argmax / top-p / top-n-sigma / min-p)
 │
 ▼
Detokenize → stream chunk → HTTP response

Two dynamic-dispatch points:

• dyn TransformerLayer — one virtual call per layer per step. Layer types (Qwen3AttentionLayer, MoeLayer, Qwen3SsmLayer, NemotronMamba2Layer, VisionEncoder) hold their own pre-resolved KernelHandles for the ops they need. The virtual call is cheap — typically ~ns — against a forward pass that takes ~0.1–1 ms per token.
• &dyn GpuBackend — one virtual call per kernel launch. Same argument; the overhead is negligible compared to the kernel itself.

Both virtual calls are unavoidable consequences of the specialization thesis: we want spark-server to not know what GpuBackend it's talking to, and we want InferenceEngine to not know what layer shape it's running. That's how new hardware and new models plug in.

With CUDA graphs enabled (the default in production), steps 5–6 collapse to a single cuGraphLaunch — the dynamic dispatch cost disappears into the graph capture phase.

Where to look in the code

The spark-runtime chapter expands on GpuBackend; spark-model on the layer/factory split; atlas-kernels on the build-time codegen. The SBIO chapter explains why every arrow in the diagram above goes through a trait.

SBIO: Business Logic vs I/O

SBIO — Separation of Business logic and I/O — is the user-level naming of a specific pattern the Atlas codebase applies aggressively: business logic never performs I/O directly. It calls a trait. Real I/O is implemented behind that trait. Tests swap in mock implementations.

The payoff is concrete. The Atlas test suite runs ~80% of the code without a GPU. You can verify the scheduler's fairness properties, the OpenAI/Anthropic protocol parsers, the sampler's numeric behavior, the tokenizer's template expansion, and every weight loader's shape checks on a vanilla Linux laptop. The only tests that need a GB10 are the ones that exercise a real CUDA kernel end-to-end.

The I/O surfaces in Atlas

Four kinds of I/O happen at runtime. Each goes through a dedicated trait:

Everything above these traits — the scheduler, the engine, the layer assembly, the sampler, the tokenizer, the tool parser, the rate limiter — is pure Rust. It contains no cudaXxx, no socket call, no file-open, no mpi_*.

What the trait boundary looks like

GpuBackend is the most load-bearing of the four. Simplified shape (full trait in crates/spark-runtime/src/gpu.rs):

#![allow(unused)]
fn main() {
pub trait GpuBackend: Send + Sync {
    // Memory
    fn alloc(&self, bytes: usize) -> Result<DevicePtr>;
    fn free(&self, ptr: DevicePtr) -> Result<()>;
    fn copy_h2d(&self, src: &[u8], dst: DevicePtr) -> Result<()>;
    fn copy_d2h(&self, src: DevicePtr, dst: &mut [u8]) -> Result<()>;
    fn copy_d2d(&self, src: DevicePtr, dst: DevicePtr, bytes: usize) -> Result<()>;
    fn memset(&self, ptr: DevicePtr, value: u8, bytes: usize) -> Result<()>;
    fn total_memory(&self) -> Result<u64>;
    fn free_memory(&self) -> Result<u64>;

    // Kernel launch
    fn kernel(&self, module: &str, func_name: &str) -> Result<KernelHandle>;
    fn launch(
        &self, func: KernelHandle, grid: [u32; 3], block: [u32; 3],
        shared_mem: u32, stream: u64, params: &mut [*mut c_void],
    ) -> Result<()>;

    // Streams
    fn default_stream(&self) -> u64;
    fn synchronize(&self, stream: u64) -> Result<()>;

    // Optional (default no-op impls)
    fn begin_capture(&self, stream: u64) -> Result<()> { Ok(()) }
    fn end_capture(&self, stream: u64) -> Result<GraphHandle> { unimplemented!() }
    fn launch_graph(&self, graph: GraphHandle, stream: u64) -> Result<()> { unimplemented!() }
    // ... events, host-pinned alloc, thread binding, etc.
}
}

A layer's forward pass calls gpu.launch(...). It does not know, and cannot know, whether gpu is AtlasCudaBackend or MockGpuBackend. That opacity is the whole point.

The mock backend

MockGpuBackend is in crates/spark-runtime/src/gpu.rs alongside the trait. It does not talk to a GPU — it keeps a bump-allocator of fake DevicePtr values, records every launch in a Vec<LaunchRecord>, and returns Ok(()) from every op. Typical test:

#![allow(unused)]
fn main() {
#[test]
fn engine_runs_correct_layer_sequence() {
    let gpu = MockGpuBackend::new();
    let cfg = ModelConfig::fixture_qwen3_5_small();
    let engine = InferenceEngine::build_for_test(&cfg, &gpu).unwrap();

    engine.decode_step(&mut ctx).unwrap();

    let launches = gpu.drain_launches();
    assert_eq!(launches.len(), cfg.num_hidden_layers * KERNELS_PER_LAYER);
    assert_eq!(launches[0].module, "attention");
    assert_eq!(launches[0].function, "prefill_attn_v47");
}
}

No GPU, no nvcc, no cudarc ever opens the driver. The test verifies a behavioral property of the business logic — the number and order of kernel launches — without depending on the kernel actually executing correctly.

The single-GPU CommBackend

SingleGpuBackend in crates/spark-comm/src/lib.rs is the distributed-comms analogue:

#![allow(unused)]
fn main() {
pub struct SingleGpuBackend;

impl CommBackend for SingleGpuBackend {
    fn all_reduce(&self, _ptr: u64, _bytes: usize) -> Result<()> { Ok(()) }
    fn all_gather(&self, _send: u64, _recv: u64, _bytes: usize) -> Result<()> { Ok(()) }
    fn reduce_scatter(&self, _s: u64, _r: u64, _b: usize) -> Result<()> { Ok(()) }
    fn broadcast(&self, _ptr: u64, _bytes: usize, _root: usize) -> Result<()> { Ok(()) }
    fn rank(&self) -> usize { 0 }
    fn world_size(&self) -> usize { 1 }
    // ...
}
}

Every collective op is a no-op. The single-GPU serving path holds a Box<dyn CommBackend> that happens to be SingleGpuBackend; the multi-GPU path holds an NcclBackend. The layer code never knows which one is live.

Business logic that benefits

The SBIO pattern makes the following blocks fully testable on CI without any GPU:

• Scheduler (spark-server/src/scheduler/): SLAI deadline logic, chunked-prefill budget enforcement, KV page allocation and eviction. Tested with a MockGpuBackend standing in for the KV cache.
• Engine (spark-model/src/engine.rs): layer ordering, speculative-decode verify + accept logic, sampler integration. The layer trait objects hold mock kernel handles.
• Tool parsers (spark-server/src/tool_parser.rs): Hermes, Qwen3-coder, Mistral formats. Input is a plain string of model output; tested with fixtures.
• Rate limiter (spark-server/src/rate_limiter.rs): token-bucket arithmetic. Pure CPU.
• Refusal / citation extraction (spark-server/src/refusal.rs): post-processing regex over plain strings. Pure CPU.
• Weight loader per family: shape/name checks, quantization-scheme dispatch. Tested with fixture safetensor files and MockGpuBackend.

The things that still require a GPU:

• Kernel correctness vs a PyTorch/reference implementation (covered by atlas-spark-bench and the integration tests in tests/).
• End-to-end model coherence (covered by tests/run_all_models.py).
• Multi-node collective ops (covered by scripts/test-minimax-ep2.sh and the EP=2 test harness).

The ATLAS_SKIP_BUILD gate

The matching idea at build time: ATLAS_SKIP_BUILD=1 makes atlas-kernels/build.rs emit a stub target_ptx.rs with empty constants. The workspace compiles, cargo clippy and cargo fmt both work, unit tests run. nvcc is not on the PATH of the GHA runner that runs the ci.yml workflow, and that is on purpose — CI catches type and lint regressions without needing a GPU CI pool.

The only test category that requires CUDA is the ones marked #[ignore] in the cargo test run, which the integration CI (not currently in this repo) would run on a real GB10 host.

Anti-patterns we don't use

• No global cuda::init(). Every GPU op takes &dyn GpuBackend. There is no lazily-initialised global driver context that tests would need to mock.
• No cfg(test) swaps. The same code path runs in production and in tests — test doubles are first-class impls, not conditional-compilation ghosts.
• No Result<T, Cuda*> leaking upward. Every GpuBackend method returns anyhow::Result<T>. Callers can't depend on which driver error surfaced.

SBIO is a discipline, not a mechanism. But because the trait surface is small, the discipline is easy to enforce in review: if you see cudaXxx or socket or fs::open inside spark-model or spark-server, that is a bug.

When you're adding code

The rule is: I/O goes through a trait; business logic goes nowhere near a syscall. Three concrete checks when you're about to merge:

1. The file you edited — does it import cudarc, nccl_sys, std::net, or std::fs? If it's inside spark-model, spark-server's non-handler code, or any atlas-* primitive crate, the answer should be no. Route through the matching trait.
2. Is the function unit-tested without #[ignore]? If yes, it's on the SBIO side of the line. If no, either move it over or explain in the PR why not.
3. If you're adding a new trait method, ask: is there a no-op mock impl that makes sense? If not, the method is probably the wrong shape.

Read Kernel Dispatch next to see how SBIO composes with the runtime dispatch path — or spark-comm for the collective-ops version of the same pattern.

atlas-core

Role: the type + trait vocabulary every other Atlas crate builds on. Key traits: ComputeTarget (build-time compiler abstraction), Vendor, DType, KernelTarget, TensorRef, ModelConfig. Dependencies: none from the workspace — this is the bottom of the stack.

Module map

crates/atlas-core/src/
├── lib.rs            re-exports every public module
├── compute.rs        ComputeTarget trait + Vendor enum (Nvidia/Amd/Apple/Intel)
├── target.rs         KernelTarget — the (arch, model, quant) dispatch key
├── dtype.rs          DType enum: E2M1, FP8E4M3, FP8E5M2, BF16, FP16, FP32
├── tensor.rs         TensorRef — zero-copy handle (ptr, shape, strides, dtype)
├── config.rs         ModelConfig — deserialized HF config.json (SSOT for model dims)
├── device.rs         Device {ordinal, total/free memory}
├── stream.rs         Stream abstraction (u64 wrapping CUstream or Metal queue)
├── kernel.rs         Kernel launch descriptor types
├── registry.rs       Kernel module registry shared by atlas-kernels codegen
├── capabilities.rs   Hardware capability flags (tensor cores, async copy, graphs)
└── error.rs          anyhow re-export + Result alias

ComputeTarget: the hardware-vendor abstraction

#![allow(unused)]
fn main() {
pub trait ComputeTarget {
    fn source_extension(&self) -> &str;       // "cu", "metal", "hip", "cl"
    fn output_extension(&self) -> &str;       // "ptx", "metallib", "hsaco"
    fn output_is_text(&self) -> bool;         // PTX is text; metallib is binary
    fn find_compiler(&self) -> Option<PathBuf>; // nvcc, xcrun, hipcc, icpx
    fn compile(&self, src: &Path, out: &Path, arch: &str, flags: &[String]) -> Result<(), String>;
    fn vendor(&self) -> Vendor;
}
}

The trait is consumed at build time by atlas-kernels/build.rs, which reads kernels/<hw>/HARDWARE.toml, instantiates the matching Box<dyn ComputeTarget>, and calls compile() on every source file in every matching (model, quant) leaf. Today only NvidiaTarget { nvcc } is implemented; AppleTarget, AmdTarget, IntelTarget are the planned impls.

Vendor and the matrix

#![allow(unused)]
fn main() {
pub enum Vendor { Nvidia, Amd, Apple, Intel }
}

Parsed from HARDWARE.toml's vendor = "nvidia" / "amd" / "apple" / "intel" field. Vendor::from_str lives here, used by both build.rs (to pick ComputeTarget) and spark-server::main (to pick the runtime GpuBackend).

KernelTarget: the runtime dispatch key

#![allow(unused)]
fn main() {
pub struct KernelTarget {
    pub arch: &'static str,   // "sm_121", "sm_100a", ...
    pub model: &'static str,  // "qwen3-next-80b-a3b"
    pub quant: &'static str,  // "nvfp4", "fp8", "bf16"
}
}

One const per supported target (KernelTarget::GB10_QWEN3_NVFP4, GB10_QWEN35_NVFP4, GB10_QWEN35_122B_NVFP4, …). Equality + hashing work from the three string components; the struct is Copy.

The atlas-kernels crate's auto-generated target_ptx.rs emits one pub static PTX_<TARGET>: &[PtxModule] per known KernelTarget. Startup dispatch in spark-server::main calls atlas_kernels::select_target(&ptx_sets, &target) to pick the right module set.

DType: the supported numeric formats

#![allow(unused)]
fn main() {
pub enum DType {
    E2M1,        // 4-bit float: {0, 0.5, 1, 1.5, 2, 3, 4, 6} — NVFP4 weights
    FP8E4M3,     // 8-bit float, 4e/3m — block scales, FP8 weights
    FP8E5M2,     // 8-bit float, 5e/2m — alternative FP8
    BF16,        // brain float — activations, residual stream
    FP16,        // IEEE 754 half — rarely used
    FP32,        // full precision — accumulation
}
}

DType::element_size_bits() returns the per-element bit count (not bytes — E2M1 is sub-byte). Buffer sizing in spark-runtime::buffers uses this.

TensorRef: the kernel-argument view

#![allow(unused)]
fn main() {
pub struct TensorRef {
    pub ptr: u64,              // CUdeviceptr / MTLBuffer / ...
    pub shape: Vec<usize>,
    pub strides: Vec<usize>,   // in elements, not bytes
    pub dtype: DType,
}
}

Every kernel argument that points at GPU memory is a TensorRef. The layer code in spark-model builds TensorRefs from DevicePtrs owned by BufferArena (in spark-runtime), passes them to primitive op traits (Normalize::rms_norm, Activation::silu_mul, Reduce::topk), and those impls extract the raw ptr: u64 to pass to GpuBackend::launch().

TensorRef itself does not own memory and is Clone — it's always a view.

ModelConfig: the single source of truth for model shape

ModelConfig::from_hf(&path) reads config.json (plus nested text/vision configs for VLMs) and returns a struct that every downstream buffer/KV-cache size derives from. Fields include hidden_size, num_hidden_layers, num_attention_heads, num_key_value_heads, intermediate_size, num_experts, num_experts_per_token, vocab_size, MoE-specific moe_intermediate_size, SSM-specific ssm_state_size, vision tower sub-config, and layer-type arrays for hybrid models.

Per the PCND principle, there are no implicit defaults for fields that materially affect correctness — missing critical fields produce a bail! at load time.

Capabilities: hardware feature flags

HardwareCapabilities exposes booleans like tensor_cores, async_copy, cuda_graphs, native_fp4_mma. HARDWARE.toml populates them; kernel-selection code in some layers reads them to pick between code paths (e.g., GB10 has no native FP4 MMA so the NVFP4 kernels fall back to software E2M1 conversion — see the NVFP4 deep dive).

What's explicitly not here

• No actual compilation. ComputeTarget::compile() is a trait method; the impls live in the callers (or in atlas-kernels/build.rs for the bootstrap).
• No actual GPU allocation. That's spark-runtime.
• No weight loader types. Those are spark-model.
• No HTTP types. Those are spark-server.

atlas-core is small on purpose. It contains only the vocabulary that every crate downstream needs.

Adding a vendor

Implementing a new hardware vendor starts here:

1. Extend Vendor if needed (all four major ones are already enumerated).
2. Write your struct XxxTarget in your own crate or inline, implement ComputeTarget.
3. Register it in atlas-kernels/build.rs::resolve_compute_target().
4. Continue to the spark-runtime chapter for the runtime trait (GpuBackend).

atlas-kernels

Role: the bridge between the CUDA source tree and the Rust workspace. Every PTX module every other crate launches is defined here. Key file: src/lib.rs (hand-written glue, ~50 lines) + build.rs (auto-generates the heavy file).

The trick: auto-generated PTX embedding

atlas-kernels/src/lib.rs ends with a single line:

#![allow(unused)]
fn main() {
include!(concat!(env!("OUT_DIR"), "/target_ptx.rs"));
}

Everything inside target_ptx.rs — per-target PTX byte constants, the ptx_modules() lookup function, the all_ptx_sets() multi-target registry — is produced by build.rs at every Cargo build. You will not find target_ptx.rs in the repository; it is generated fresh into OUT_DIR each time.

The generated file looks roughly like:

#![allow(unused)]
fn main() {
pub static PTX_GB10_QWEN3_NVFP4: &[PtxModule] = &[
    PtxModule { name: "prefill_attn_v47", ptx: include_bytes!("...sm121/prefill_v47.ptx") },
    PtxModule { name: "decode_attn",       ptx: include_bytes!("...sm121/decode_attn.ptx") },
    PtxModule { name: "moe_w4a16",         ptx: include_bytes!("...sm121/moe_w4a16.ptx") },
    // ~35 modules per target
];
pub static PTX_GB10_QWEN35_NVFP4: &[PtxModule] = &[ /* ... */ ];

pub fn ptx_modules(target: &KernelTarget) -> Option<&'static [PtxModule]> {
    match (target.arch, target.model, target.quant) {
        ("sm_121", "qwen3-next-80b-a3b", "nvfp4") => Some(PTX_GB10_QWEN3_NVFP4),
        ("sm_121", "qwen3.5-35b-a3b",   "nvfp4") => Some(PTX_GB10_QWEN35_NVFP4),
        // ...
        _ => None,
    }
}
}

At runtime, spark-server::main resolves the current model's KernelTarget, calls ptx_modules(&target), and passes the resulting slice to AtlasCudaBackend::new which uploads each PTX module to the GPU via cuModuleLoadData.

What build.rs actually does

1. Read ATLAS_TARGET_* — three env vars (HW, MODEL, QUANT). Wildcards (*) expand to "every matching directory".
2. Walk kernels/<hw>/<model>/<quant>/ — for each leaf that matches the wildcards, read HARDWARE.toml, MODEL.toml, KERNEL.toml.
3. Resolve the compiler — resolve_compute_target(vendor) returns a Box<dyn ComputeTarget>. Today always NvidiaTarget { nvcc }.
4. Compile every source file — for each *.cu in the leaf, call compute_target.compile(src, out, arch, flags). Flags come from KERNEL.toml's extra_nvcc_flags = [...] plus the arch-specific ones from HARDWARE.toml.
5. Apply module-name overrides — KERNEL.toml's [modules] section lets kernels with different file stems (e2m1_branchless.cu) expose themselves under shorter module names (e2m1). This is cosmetic but keeps GpuBackend::kernel("e2m1", "convert_f32_to_e2m1") readable at the call site.
6. Parse MODEL.toml → SamplingPresets + ModelBehavior — non-kernel metadata that the server consumes directly (default temperature, thinking_budget, etc.). Emitted as Rust constants alongside the PTX.
7. Write target_ptx.rs — one PtxModule array per target, the dispatch match, and a const ALL_TARGETS: &[KernelTarget] listing everything that got compiled.

The whole phase is idempotent — rerun-if-changed directives on the kernel tree mean cargo only recompiles what changed.

ATLAS_SKIP_BUILD=1 — the escape hatch

On a Linux laptop with no nvcc, the crate would fail to build without this. The escape hatch is a single env var. When set, build.rs:

• Does not invoke any compiler.
• Emits a stub target_ptx.rs with an empty ALL_TARGETS and ptx_modules returning None for everything.
• The crate compiles cleanly.

ci.yml uses this. So does local cargo clippy. The Kernel Dispatch chapter covers the broader flow.

Per-target PTX bytes — sizes and counts

Rough numbers for the default multi-model build at avarok/atlas-gb10:latest:

Total embedded PTX in the default multi-model binary: ~65 MB. The binary itself lands at ~200 MB in release builds. This is why the Docker image is ~8 GB once you add the CUDA userspace (libcudart, libnvrtc), tokenizer deps, and Ubuntu base — the actual Atlas footprint is small.

What gets added to this crate when you…

• …add a new (hw, model, quant) leaf? Nothing in the atlas-kernels/src/ directory. build.rs picks it up automatically on the next cargo build. You do need to add a matching KernelTarget const in atlas-core::target::KernelTarget so downstream code can refer to it by name.
• …add a new kernel to an existing leaf? Drop the .cu in the leaf directory; build.rs picks it up. If you want a non-stem module name, add an entry to the leaf's KERNEL.toml.
• …add a new hardware vendor? Extend resolve_compute_target(vendor) in build.rs to return your new ComputeTarget impl. Everything else flows from there.

The rest of the kernel-engineering story is in the CUDA Kernel Engineering deep dive.

spark-runtime

Role: everything that touches the GPU directly. GpuBackend trait + CUDA implementation, KV cache, prefix cache, buffer arena, sampler, fast weight loader. Key files: gpu.rs, cuda_backend.rs, kv_cache.rs, prefix_cache.rs, radix_tree.rs, buffers.rs, sampler.rs, fast_weights/mod.rs, weights.rs.

The load-bearing trait: GpuBackend

27 methods across five concerns:

• Memory — alloc, free, copy_h2d/d2h/d2d, memset, total_memory, free_memory, alloc_host_pinned.
• Kernel launch — kernel(module, func) returns a KernelHandle; launch(handle, grid, block, shared, stream, args) fires the launch.
• Streams — default_stream, create_stream, synchronize, bind_to_thread.
• CUDA graphs — begin_capture, end_capture → GraphHandle, launch_graph.
• Events — create_event, record_event, stream_wait_event.

Required methods have no default; optional methods have default panics/no-ops so a partial backend (e.g. a Metal backend without CUDA-graph support yet) still compiles. The trait's SBIO role is discussed in Part II.

The production impl: AtlasCudaBackend

In cuda_backend.rs. Built on cudarc (Rust bindings over the CUDA driver API). On construction:

1. cudarc::driver::CudaDevice::new(ordinal) acquires a driver context.
2. For each PtxModule in the provided target set: cuModuleLoadData uploads the PTX, then cuModuleGetFunction caches one CUfunction per exported kernel.
3. Kernel handles are stored in a HashMap<(module: &str, func: &str), KernelHandle>.

Per-launch the backend unpacks the kernel args into void*[], sets the stream, and calls cuLaunchKernel. The launch-side ceremony — converting DevicePtr/scalar args into a void* pointer array — is abstracted one level up by the KernelLaunch builder pattern in spark-model::layers::ops.

The paged KV cache (kv_cache.rs)

Atlas uses paged attention (à la vLLM) with block-level allocation. Core types:

#![allow(unused)]
fn main() {
pub enum KvCacheDtype {
    Bf16,    // 2 bytes/element — unquantized baseline
    Fp8,     // 1 byte/element — E4M3 + per-tensor scale
    Nvfp4,   // 0.5 bytes + per-group FP8 scale — maximum compression
    Turbo4,  // 4-bit WHT + Lloyd-Max (TurboQuant) — lower MSE than NVFP4
    Turbo3,  // 3-bit WHT + Lloyd-Max — smallest
    Turbo8,  // WHT + FP8 — outlier-resistant FP8
}
}

PagedKvCache holds a pool of fixed-size blocks (configurable, typically 16 tokens per block). KvCacheConfig derives pool sizing from ModelConfig + --max-seq-len + --max-batch-size. Allocation is O(1) from a free list; eviction is handled by the scheduler.

The TurboQuant family (turbo3, turbo4, turbo8) is specific to Atlas: Walsh-Hadamard rotation followed by Lloyd-Max quantization to optimal Gaussian codebook levels. For the same bit rate, turbo4 has ~2× lower MSE than NVFP4 on the kinds of activations transformers produce, because WHT flattens outliers before quantization. See docs/turboquant-plus.md and FP8 / NVFP4 chapters.

Prefix caching (prefix_cache.rs, radix_tree.rs)

RadixAttention: the system prompt shared by every request can be KV-cached once, reused forever. Implementation:

• radix_tree.rs — in-memory radix tree keyed on token sequences. Each node owns the KV pages for its token prefix.
• prefix_cache.rs — the orchestration layer. When a request arrives, the scheduler calls prefix_cache.lookup(tokens) which walks the tree to the deepest matching node. The KV pages for that prefix are already resident on the GPU.

Hit rates are high in practice — system prompts and few-shot examples dominate, and chat agents reuse most of their tool schemas across turns. TTFT drops ~10× on warm-cache hits. This is the feature enabled by --enable-prefix-caching.

Marconi (SSM snapshots) extends the idea to SSM layers: a full SSM state is ~GB on a 35B model, so prefix cache hits for hybrid models also need a snapshotted SSM state to be genuinely equivalent. That machinery lives partly here and partly in spark-model. See docs/adr/0003-hybrid-ssm-attention.md for the SSM-snapshot-cache design.

Buffer arena (buffers.rs)

One BufferArena per serve. Allocates every scratch buffer once at startup, sized for the worst-case batch × seq_len combination allowed by CLI flags. Includes:

• hidden_states, residual, norm_output — residual stream and its post-norm staging
• qkv_output, attn_output — attention projection outputs
• gate_logits, moe_output — MoE intermediates
• logits — the final [M, vocab_size] output
• ssm_qkvz_scratch — Mamba/GDN projections (sized for 3× positions for MRoPE on models that need it)
• expert_outputs — sized for max(k_max, max_batch_tokens) to cover both speculative decode (K=3) and batched MoE prefill

The arena never reallocates during serving. This is one of the invariants that makes CUDA graph capture viable — buffer addresses are graph-stable.

Sampler (sampler.rs)

SamplingParams — temperature, top_p, top_k, top_n_sigma, min_p, repetition_penalty, presence_penalty. The sampler:

1. Applies penalties (presence, repetition) in-place on the logits buffer.
2. Applies top_n_sigma (entropy-based filter).
3. Applies top_p + top_k + min_p.
4. Softmax.
5. Multinomial sampling or argmax (if temperature == 0).

A known bug with temperature=0 && repetition_penalty=0 was fixed in wave-8 of the bug sweeps; the sampler now has explicit div-by-zero guards. --adaptive-sampling toggles an entropy-gated greedy path that avoids the full softmax+sample when the logits are effectively one-hot.

Fast weight loader (fast_weights/)

Atlas's production weight loader. Modeled on scitix/InstantTensor:

• Each safetensors shard is opened with O_DIRECT (bypasses the page cache — critical on GB10 where the page cache shares physical memory with the GPU).
• One reader thread pre-fetches the next tensor's bytes into a page-aligned buffer.
• The main thread copy_h2d's the current tensor while the next one is being read.
• A per-shard heuristic auto-picks between O_DIRECT and buffered reads: shards with > 5000 tensors pay too much per-tensor syscall overhead for O_DIRECT, and kernel readahead wins there.
• If the filesystem rejects O_DIRECT (tmpfs, overlayfs), falls back to mmap automatically.

Cold-load speedups measured vs mmap (with posix_fadvise(DONTNEED) between runs):

On by default. --no-fast-load reverts to mmap.

MockGpuBackend and testing

The test double lives in gpu.rs next to the trait. Records every launch; returns success for every op. Enables the ~80% of the test suite that doesn't need a real GPU — see SBIO.

What's explicitly not here

• No model layers. That's spark-model.
• No HTTP. That's spark-server.
• No collective ops. That's spark-comm.

spark-runtime is the bottom of the "things that move bits on a GPU" stack and the top of the "things a layer is allowed to call directly" stack.

spark-comm

Role: the multi-GPU collective-ops abstraction. One trait, two impls: NCCL for real distributed runs and a no-op backend for single-GPU. Key file: src/lib.rs (CommBackend trait), nccl.rs (raw NCCL FFI), nccl_backend.rs (the NcclBackend impl).

Why this is its own crate

Multi-GPU in Atlas is Expert Parallelism (EP) — the MoE experts of models beyond one GB10's weight budget (122B, 119B, 229B) are split across two nodes connected via RoCEv2. Token dispatch between ranks goes through NCCL all-reduces and send/recv.

spark-comm isolates the NCCL surface so that:

1. Single-GPU deployments never link against NCCL (the binary loads a SingleGpuBackend).
2. Tests for the scheduler and the layer code can run against the no-op impl on CI.
3. Porting to a different collective-ops library (RCCL for AMD, Metal MPS, oneCCL) is a new CommBackend impl, nothing else changes.

The trait

#![allow(unused)]
fn main() {
pub trait CommBackend: Send + Sync {
    fn all_reduce(&self, ptr: u64, bytes: usize) -> Result<()>;
    fn all_gather(&self, send: u64, recv: u64, bytes: usize) -> Result<()>;
    fn reduce_scatter(&self, send: u64, recv: u64, bytes: usize) -> Result<()>;
    fn broadcast(&self, ptr: u64, bytes: usize, root: usize) -> Result<()>;
    fn send(&self, ptr: u64, bytes: usize, peer: usize) -> Result<()>;
    fn recv(&self, ptr: u64, bytes: usize, peer: usize) -> Result<()>;
    fn barrier(&self) -> Result<()>;
    fn rank(&self) -> usize;
    fn world_size(&self) -> usize;
    fn stream(&self) -> u64;
    fn set_stream(&mut self, stream: u64);
}
}

All pointer arguments are u64 (matching CUDA's CUdeviceptr) to avoid coupling the crate to spark-runtime's DevicePtr. Every op is stream-associated — collectives and kernel launches can be pipelined through CUDA graph capture.

SingleGpuBackend — the no-op

#![allow(unused)]
fn main() {
pub struct SingleGpuBackend;

impl CommBackend for SingleGpuBackend {
    fn all_reduce(&self, _ptr: u64, _bytes: usize) -> Result<()> { Ok(()) }
    fn all_gather(&self, _s: u64, _r: u64, _b: usize) -> Result<()> { Ok(()) }
    // ...
    fn rank(&self) -> usize { 0 }
    fn world_size(&self) -> usize { 1 }
}
}

This is what runs in single-GPU serving. The expert-parallel code paths in spark-model::layers::moe still call comm.all_reduce(...) — the op is a no-op under SingleGpuBackend and an actual NCCL call under NcclBackend. The caller never branches on world_size.

NcclBackend — the real impl

In nccl_backend.rs. Uses the unsafe NCCL FFI in nccl.rs. Construction flow:

1. The master rank (0) calls ncclGetUniqueId and publishes the id to the scheduler's rendezvous port (--master-addr, --master-port, default 29500).
2. Every rank (including master) dials the rendezvous, receives the id.
3. All ranks call ncclCommInitRank(world_size, id, rank) in parallel; the call is collective and blocks until every rank has joined.

The NCCL env layer is fussy on GB10 — the scripts in scripts/start-ep2.sh + scripts/start-minimax-ep2.sh pin the critical vars:

These are worth the paragraph — a mis-set NCCL_SOCKET_IFNAME on GB10 will silently fall back to the 1 GbE management interface and drop EP=2 throughput by an order of magnitude.

The EP=2 throughput path

For Qwen3.5-122B-A10B NVFP4 at EP=2:

• 128 experts per rank (256 total).
• Token dispatch: the gate runs on every rank, top-k expert IDs are selected, tokens destined for remote experts are reduce_scatter'd to the owning rank.
• Expert compute happens locally.
• Expert outputs are all_gather'd back.
• Result: ~46 tok/s sustained on 600-token decodes (see Multi-GPU).

The bandwidth pressure is all in the dispatch + gather, which is why RoCEv2 with GDR matters. A plain TCP NCCL falls off by 3×.

The critical MTP-flag symmetry rule

A subtle footgun: when the head (rank 0) runs with --speculative --mtp-quantization nvfp4 --num-drafts N, the worker must be started with the same flags. If not, the MTP verify command from the head lands in the worker's SSM layer without intermediate buffers allocated and you get an SSM intermediate-buffer error. scripts/start-ep2.sh handles this; a manual two-command launch does not, and it has bit multiple contributors. See the Multi-GPU chapter.

NCCL safety in tests

The unit tests for the expert-parallel layer code do not instantiate NcclBackend. They hold a Box<dyn CommBackend> = Box::new(SingleGpuBackend) and verify the code path by checking that the layer calls all_reduce at the right moment — the launch recorder in MockGpuBackend plus a trace in SingleGpuBackend is enough. The real NCCL path is validated by scripts/test-minimax-ep2.sh against a live two-node cluster.

What's explicitly not here

• No kernel code. The EP=2 token-dispatch logic lives in Rust at crates/spark-model/src/layers/moe/forward_ep.rs, and the routed grouped-GEMM kernel in kernels/gb10/<model>/<quant>/moe_w4a16_grouped_gemm.cu.
• No scheduler logic. That's spark-server::scheduler.
• No RDMA-specific code. Atlas talks through NCCL; NCCL talks through libibverbs/librdmacm. We do not bypass.

Adding a new collective-ops library is a single impl CommBackend in a new module here plus a selection arm in spark-server::main that picks the right backend given the vendor.

spark-model

Role: the model assembly crate. Translates loaded weights and config into Box<dyn TransformerLayer> objects, drives the inference engine loop, implements speculative decoding and vision preprocessing. Key files: engine.rs, model.rs, factory.rs, layer.rs + layers/*.rs, weight_loader/*.rs, weight_map.rs, speculative.rs, vision_preprocess.rs, traits.rs, quant_format.rs, mistral_loader.rs, preflight.rs.

This is the largest crate in the workspace — ~18k lines — because every model architecture Atlas supports has its own loader here. The design centers on two small, heavily-used traits.

The central traits

Model

#![allow(unused)]
fn main() {
pub trait Model {
    fn alloc_sequence(&self) -> Result<Sequence>;
    fn free_sequence(&self, seq: &mut Sequence) -> Result<()>;
    fn prefill(&self, seq: &mut Sequence, tokens: &[u32], stream: u64) -> Result<()>;
    fn decode(&self, seq: &mut Sequence, token: u32, stream: u64) -> Result<u32>;
    // + vision, tool, MTP extension methods
}
}

One-model-per-server. TransformerModel (in model.rs) is the concrete type — owns the loaded Vec<Box<dyn TransformerLayer>>, the shared embedding + LM head, the KV cache, the prefix cache, the buffer arena. Threaded via Arc into the scheduler.

TransformerLayer

#![allow(unused)]
fn main() {
pub trait TransformerLayer: Send + Sync {
    fn forward(&self, ctx: &mut LayerContext) -> Result<()>;
    fn kind(&self) -> LayerKind; // Attention / SsmAttention / Moe / DenseFfn / Vision
    // + MTP-aware variants, KV cache allocation hooks
}
}

Every layer type is a trait object. The decode-step loop in engine.rs iterates self.layers.iter().map(|l| l.forward(ctx)). No match layer.kind() branches in the hot loop — the virtual call is the entire dispatch.

The layer menagerie (layers/*.rs)

New models reuse these where possible. Writing a new layer type is rare — the MiniMax 256-expert sigmoid-routed MoE is the most recent example, and it was a new file because the routing semantics genuinely differ from softmax-topk. Gemma-4's sliding+full alternation got a new file because the attention window masks alternate per layer.

Weight loaders (weight_loader/*.rs)

One file per model family. Each implements:

#![allow(unused)]
fn main() {
pub trait ModelWeightLoader {
    fn load_layers(&self, store: &WeightStore, config: &ModelConfig,
                   gpu: &dyn GpuBackend, layer_kv_dtypes: &[KvCacheDtype])
        -> Result<Vec<Box<dyn TransformerLayer>>>;
    fn load_embedding(&self, store: &WeightStore, config: &ModelConfig) -> Result<DenseWeight>;
    fn load_final_norm(&self, store: &WeightStore, config: &ModelConfig,
                       gpu: &dyn GpuBackend) -> Result<DenseWeight>;
    fn load_lm_head(&self, store: &WeightStore, config: &ModelConfig) -> Result<DenseWeight>;
    fn load_mtp_weights(&self, store: &WeightStore, config: &ModelConfig,
                        gpu: &dyn GpuBackend) -> Result<Option<MtpWeights>>;
}
}

Current files:

• qwen3.rs — Qwen3-Next (NVFP4, hybrid SSM+Attention+MoE with MTP).
• qwen35.rs — Qwen3.5 MoE (35B, 122B) with GDN + MTP.
• qwen35_dense.rs — Qwen3.5 Dense (27B), hybrid without MoE.
• qwen3_vl.rs — Qwen3-VL (30B, vision + attention + MoE).
• gemma4.rs — Gemma-4 (26B MoE + 31B dense; GeGLU; sliding/full alternation).
• nemotron.rs — Nemotron-H Nano + Super (Mamba-2 + MoE + attention).
• minimax.rs — MiniMax-M2 / M2.7 (256-expert sigmoid MoE).
• mistral_loader.rs — the one outlier (lives one level up in spark-model/src/ because the Mistral-Small-4 loader predates the weight_loader/ submodule reorganisation).

Each loader knows the HF weight-name patterns for its family and translates them into Box<dyn TransformerLayer> via the helpers in weight_map.rs:

• load_attention(store, layer_idx, prefix) — reads q_proj, k_proj, v_proj, o_proj + optional RoPE scales.
• load_moe(store, layer_idx, num_experts, ...) — reads the expert weights, gate, optional shared experts.
• load_ssm(store, layer_idx, config) — reads the Mamba/GDN A, B, C, D, dt projections.
• load_dense_ffn(...) — gate + up + down.
• dequant_nvfp4_to_bf16, dequant_fp8_to_bf16 — on-the-fly quant conversion for layers that run BF16 even if the checkpoint is quantized (e.g. first/last attention layers under --kv-high-precision-layers).

The factory (factory.rs)

The single point where a model type becomes a loader:

#![allow(unused)]
fn main() {
fn loader_for_config(config: &ModelConfig) -> Result<Box<dyn ModelWeightLoader>> {
    let normalized = config.model_type
        .to_lowercase()
        .replace('-', "_")
        .replace('.', "_");
    match normalized.as_str() {
        "qwen3_next_for_causal_lm"     => Ok(Box::new(Qwen3WeightLoader)),
        "qwen3_5_next_for_causal_lm"   => Ok(Box::new(Qwen35WeightLoader)),
        "qwen3_5_for_causal_lm"        => Ok(Box::new(Qwen35DenseWeightLoader)),
        "qwen3_vl_for_causal_lm"       => Ok(Box::new(Qwen3VLWeightLoader)),
        "gemma_4_for_causal_lm"        => Ok(Box::new(Gemma4WeightLoader)),
        "nemotron_h_for_causal_lm"     => Ok(Box::new(NemotronHWeightLoader)),
        "minimax_m2_for_causal_lm"     => Ok(Box::new(MinimaxM2WeightLoader)),
        "mistral_small_4_for_causal_lm" => Ok(Box::new(MistralWeightLoader)),
        other => bail!("Unsupported model type: '{}'", other),
    }
}
}

This is the single code site where model_type strings are matched. Everything downstream of factory::build holds Box<dyn Model> and is model-agnostic. See the top-level repo Adding a new model guide.

The engine (engine.rs)

#![allow(unused)]
fn main() {
pub fn generate(
    model: &dyn Model,
    prompt_tokens: &[u32],
    params: &SamplingParams,
) -> Result<GenerateResult>;
}

Prefill → decode loop, sampler integration, finish-reason detection ("stop" / "length"), EOS + stop-token handling. The scheduler in spark-server drives this — engine.rs itself is stateless per call; the per-sequence state lives on Sequence (allocated/freed around the generate).

Speculative decoding (speculative.rs)

Wraps the MTP draft-then-verify loop. Draft tokens are produced by the MTP head, verified by the main model in one forward pass, and accepted-to-longest-match. Sibling stride bugs (qwen3_attention + qwen3_ssm with K≠2) were fixed in the Pass-16/Pass-22 bug sweeps; the current code handles K=1, 2, and 3 for the families that support it. See the MTP chapter.

Vision preprocessing (vision_preprocess.rs)

For Qwen3-VL and Qwen3.6: accept image input (JPEG/PNG/base64), resize/normalise to the model's patch grid, produce pixel-values tensor + MRoPE position IDs (H/W/T triples). Handles the 3× positions scratch the MRoPE path needs.

Quant format runtime dispatch (quant_format/)

Sniffs the checkpoint shape on load and picks the right Dequantize implementation. Introduced in the Pass-25 sweep to replace a load-time heuristic that had produced EP=2 CUDA illegal-address errors on the ModelOpt-NVFP4 variant of M2.7.

Preflight (preflight.rs)

Runs a small synthetic decode step before the HTTP server binds. Catches:

• Weight-loading shape mismatches.
• KV-cache budget overruns (pre-OOM).
• Missing kernel modules for the selected KernelTarget.

This is what "OOM pre-flight" in the feature matrix is. Failing preflight produces a clear, early error; passing it means the hot path is safe.

What's explicitly not here

• No HTTP. That's spark-server.
• No GPU ops. Every GPU touch is via spark-runtime::GpuBackend.
• No collective ops. Every multi-GPU touch is via spark-comm::CommBackend.

Adding a new model is almost always: one new weight_loader/<family>.rs, one match arm in factory.rs, optional reuse of existing layers/*.rs, optional new layers/<family>_attention.rs if the attention shape genuinely differs.

spark-server

Role: the binary. OpenAI- and Anthropic-compatible HTTP server, request scheduler, tokenizer, tool parsing, streaming, rate limiter, CLI. Key files: main.rs, cli.rs, api.rs, openai.rs, anthropic.rs, scheduler.rs, scheduling_policy.rs, tool_parser.rs, reasoning_parser.rs, tokenizer.rs, rate_limiter.rs, refusal.rs, metrics.rs, conversation_store.rs, response_store.rs, session_manager.rs, model_resolver.rs, grammar.rs, hint_injector.rs, citation.rs, adaptive_sampler.rs, ngram.rs.

This crate is the only bin in the workspace — building spark-server produces the spark executable that the Docker image ships. Everything above is lib — spark-server ties it all together.

Startup sequence (in main.rs)

1. Parse CLI (cli::Cli::parse()).
2. Resolve model path — HF id via HF_HUB_CACHE / ~/.cache/huggingface/hub or explicit --model-from-path.
3. Load ModelConfig from config.json.
4. Resolve KernelTarget from the config's model_type + --kv-cache-dtype.
5. Instantiate GpuBackend — AtlasCudaBackend::new(ordinal, &ptx_modules).
6. Instantiate CommBackend — NcclBackend if --world-size > 1, else SingleGpuBackend.
7. factory::build(config, gpu, comm) → Arc<dyn Model> — the model weights land on the GPU.
8. Load the tokenizer (tokenizers crate, optional chat template in jinja-templates/<family>.j2).
9. Run preflight::check(&model) — one synthetic decode, fails fast on shape / budget / kernel issues.
10. Spawn the scheduler — a dedicated tokio task with a bounded mpsc channel of Requests.
11. Start axum — serve(&addr) with the route table from api.rs.
12. Capture CUDA graphs in the background while the first request is in flight (warm-start).

The scheduler is driven by an inbound mpsc::Receiver<Request>; each HTTP handler enqueues a request and awaits a mpsc::Sender<ResponseChunk> handed back to it.

HTTP routes (api.rs + openai.rs + anthropic.rs)

Streaming is the default for chat; non-streaming aggregates and returns ChatCompletionResponse. Tool-call chunks are emitted as delta.tool_calls in the SSE stream. Anthropic streaming populates stop_sequence on message_delta events (fixed in wave-12).

The server also implements the Responses API with a stateful backend (response_store, conversation_store) for multi-turn conversations, citation extraction, and a streaming refusal filter.

The scheduler (scheduler.rs + scheduling_policy.rs)

Two policies:

• FIFO — first come, first served. Decode step picks up to max_batch_size active sequences and runs a batched forward.
• SLAI — SLO-aware. Each sequence has a time-between-tokens (TBT) deadline; the scheduler prioritizes the sequence closest to its deadline. Under mixed load, this is the difference between smooth streaming and bursty output.

Allocation discipline: KV pages are claimed on prefill start, released on completion. The scheduler tracks the KV budget and chunks prefills when memory is tight (--max-prefill-tokens caps the per-iteration tokens so scratch sizes stay bounded).

Active context compaction (the compact_messages function in api.rs) applies at the HTTP level before tokenization: if the tokenized prompt approaches --max-seq-len, the server progressively truncates middle tool responses (stage 2), replaces middle responses with pointers (stage 3), drops oldest middle pairs (stage 4), and finally trims the system prompt + keeps only the last 4 messages (stage 5). References arXiv:2603.05344 (OpenDev).

Tool-call parsing (tool_parser.rs)

The server supports three tool-call formats, auto-detected from the model's MODEL.toml with a --tool-call-parser override:

The Qwen3.5 coder parser had several robustness improvements in the bug sweeps (literal </tool_call> recovery, missing </parameter> recovery, empty {} tool-calls) — the parser is now tolerant of slightly-malformed model output.

Reasoning / thinking (reasoning_parser.rs)

Models that emit <think>...</think> blocks (Qwen3.5, Nemotron-H, MiniMax) stream the thinking content to the client as a separate SSE channel keyed on "reasoning" (per OpenAI's o1-family convention). --max-thinking-budget caps the total thinking tokens. --disable-thinking is a kill-switch.

Several subtle fixes in this area were important:

• Template-forced thinking — some models emit <think> seeded by the chat template; Atlas's detector had to distinguish that from the model's own <think>. The pass-16 fix required the opening <think> to be unclosed to count as the model's own.
• Closed empty thinking — <think>\n\n</think>\n\n is a template no-op, not a reasoning block. Wave-4 fixed the false-positive.
• Multi-block reasoning — models occasionally emit multiple <think> blocks; the extractor concatenates them.

Tokenizer (tokenizer.rs)

Wraps the HF tokenizers crate. Adds jinja chat-template expansion (via minijinja). Resolves special tokens (<|im_start|>, <think>, <minimax:tool_call>, etc.) from tokenizer_config.json so the raw token ids match what the model was trained on.

Rate limiter (rate_limiter.rs)

Per-key token bucket. Wave-9 added a MAX_KEYS guard to prevent DoS via cardinality explosion. Body-size limits are env-configurable.

What's explicitly not here

• No GPU kernels. Every GPU call delegates through spark-runtime.
• No CUDA. The crate's Cargo.toml does not depend on cudarc.
• No model-specific weight code. That's spark-model.

Adding a new HTTP shape

• A new API endpoint (e.g. an Atlas-native /v1/sessions/create) — one handler in api.rs, one route in the router bindings in main.rs.
• A new tool-call format — one new parser module, one enum variant, one --tool-call-parser option.
• A new reasoning tag (<scratchpad>, etc.) — extend reasoning_parser.rs.
• A new chat template — a file in jinja-templates/<model>.j2, auto-picked up by the tokenizer layer if named after the HF repo.

atlas-spark-bench

Role: the benchmark harness that produces every throughput number in this book. HTTP client that targets a running Atlas Spark server and measures token rate, TTFT, and concurrency behaviour. Key file: src/lib.rs.

Design

This is a client-side harness. It does not link against spark-runtime or atlas-kernels — it just speaks OpenAI-compatible HTTP to a server on localhost:8888 (or wherever ATLAS_BENCH_URL points).

That shape is deliberate:

1. No GPU pollution. Running benches from the same process that serves the model would compete for GPU memory and distort results.
2. Same surface as real clients. If the HTTP stack has a latency bug, atlas-spark-bench sees it. If the tokenizer is slow on some prompt shape, it shows up.
3. Same harness tests correctness and perf. The integration tests in tests/ reuse the client to drive coherence checks against every model.

What it measures

The harness uses ureq for blocking requests and std::thread::Barrier for synchronising concurrent launches. No async — simpler, more reproducible.

The bench shapes

The canonical shapes live in bench/ at the repo root (the harness loads JSON fixtures) and include:

• Short prompt, short output — "What is the capital of France?", max_tokens ≤ 30. This is the number the READMEs quote: it emphasizes the decode hot loop, not prefill.
• Long output, single request — "Explain the theory of relativity", max_tokens = 200. Shows CUDA-graph sustained throughput.
• Concurrency sweep — the same prompt, 1 / 2 / 4 / 8 / 16 parallel streams. Reveals scheduler + KV-allocation behaviour.
• Prefix warmup — preflights the system prompt, then measures cold vs warm TTFT.
• Tool-calling — a single-tool request with a well-known function signature; measures tool-emit latency and token-streaming behaviour during tool blocks.

Running a bench

# Start a server somewhere
docker run -d ... avarok/atlas-gb10:latest serve <model>

# In another terminal
export ATLAS_BENCH_URL=http://localhost:8888
cargo bench -p atlas-spark-bench

Criterion stores results under target/criterion/. The repo's bench/ directory retains the stable JSON snapshots that the README tables are derived from; ephemeral bench runs are gitignored.

Scripts that drive it

• scripts/sweep_all_models.sh — boots each model in turn, runs the short-prompt bench, and writes the README.md throughput table.
• scripts/run_conc_benchmark.sh — runs the N-stream sweep (bench/bench_concurrency.py is the underlying driver).
• scripts/test-minimax-ep2.sh — doubles as perf + correctness for EP=2 MiniMax.

All three live in the top-level scripts/ directory, not in the crate.

require_server() — the safety rail

#![allow(unused)]
fn main() {
pub fn require_server() -> String {
    let url = server_url();
    match ureq::get(&format!("{url}/health")).call() {
        Ok(resp) if resp.status() == 200 => url,
        _ => panic!("Server not reachable at {url}. Start Atlas Spark first."),
    }
}
}

Every bench starts with this. Failing fast if the server isn't up beats a confusing timeout minutes into a run.

Where results live

• Hand-vetted snapshots: bench/*.json (gitignored at the file level but some stable ones are tracked). These feed the Benchmarks chapter.
• Criterion outputs: target/criterion/ (gitignored).
• Concurrency-sweep logs: pinned result files under bench/ (tracked, updated manually when a significant run completes).

What this crate is not

• Not a generic LLM benchmark tool. It knows about Atlas's server and its SSE streaming format.
• Not a load tester. For that you want vegeta or locust pointed at the same server.
• Not a kernel micro-benchmark. Those live in atlas-spark-bench's sister tests under each primitive crate's benches/, Criterion-driven, no HTTP.

The job of atlas-spark-bench is to produce end-to-end, apples-to-apples numbers that survive the tokenizer, the scheduler, the HTTP layer, and the kernel set. When the README says "131 tok/s on Qwen3.5-35B", that number came from here.

CUDA Kernel Engineering

This is the chapter you read when you're about to write a kernel. It covers the conventions every Atlas kernel follows, the tools that matter on GB10 SM121, and the workflow that takes an idea from a profile to a merged PR.

The kernel inventory

A default (GB10, <model>, <quant>) leaf ships ~30–40 kernels. Canonical roles:

Kernels that differ between models (e.g., Nemotron's Mamba-2 vs Qwen3.5's GDN) live in different (model, quant) leaves but share file names when the shapes match.

SM121 hardware budget (quick reference)

Grace-Blackwell GB10 / SM121 numbers you will care about when writing kernels:

The native FP4 MMA caveat is load-bearing: SM120/SM121 does not expose the cvt.rn.satfinite.e2m1x2.f32 instruction or native FP4 tensor-core paths. Every NVFP4 kernel on GB10 uses software E2M1 conversion (the "branchless" kernel) and dequantises-to-BF16 for the MMA. This is not a performance bug — it is the silicon. The community benchmarks that cite "native FP4 throughput" on newer Blackwell parts do not transfer. The NVFP4 deep dive walks the workaround in detail.

Conventions every Atlas kernel follows

• SPDX header line 1. // SPDX-License-Identifier: AGPL-3.0-only. Enforced by the license-headers job in CI.
• extern "C" __global__ entry points with a stable name. The name is what GpuBackend::kernel(module, func) looks up.
• All pointer args are typed at the right level — const __nv_bfloat16*, const int8_t*, not const void*. The BF16 + E2M1 types come from <cuda_bf16.h> and <cuda_fp8.h>; module-local aliases are fine but don't hide the precision.
• Grid/block dimensions are passed from the Rust side. Never compute block dims from runtime GPU properties inside the kernel — let KernelLaunch::new().grid(...).block(...) own it.
• One kernel per file where possible. Fused variants belong in their own files (e.g. silu_mul_quant.cu vs silu_mul.cu) so the KERNEL.toml [modules] override maps cleanly.
• No <iostream>, no printf in hot paths. Use #if 0 stubs during development, strip before merging. nvcc warns on printf inside __device__ code when it bloats the PTX.
• Shared-memory layouts are always annotated. A comment near each __shared__ alias documents the row/column order and any padding added to avoid bank conflicts.

The profiling workflow

1. Start with nsys profile against the live server running a benchmark. The first question is always which kernel is the bottleneck — do not tune in the abstract.

   nsys profile --trace=cuda,cudnn,cublas,osrt -o atlas.qdrep \
     /path/to/spark serve <model>
   # in another terminal: drive bench load
   nsys stats --report cuda_gpu_kern_sum atlas.qdrep | head -20
   

2. For the top 2–3 kernels, drill into ncu (Nsight Compute). The metrics that matter on GB10:
   • smsp__cycles_active.avg.pct_of_peak_sustained_elapsed — SM utilisation.
   • l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_ld.sum — shared-memory bank conflicts.
   • dram__bytes.sum vs dram__bytes_read.sum.peak_sustained — how close to the 273 GB/s ceiling.
   • sm__warps_active.avg.pct_of_peak_sustained_elapsed — warp occupancy.
3. Know which side of the roofline you're on. GB10's compute-vs-BW roofline is steeper than a desktop GPU's (273 GB/s vs several thousand tensor-core TFLOPs). Most decode kernels are memory-bound; most prefill kernels are tensor-core-bound once cp.async is hiding the K/V load.

The three performance levers (in order of payoff)

1. Tiling + cp.async. The biggest single win on SM121 is pipelining the next tile's global-memory load behind the current tile's compute. __pipeline_commit()/__pipeline_wait_prior(N) with two or three in-flight stages typically doubles attention throughput vs a naive loop. Prefill v47 uses 2 stages and two CTAs per SM.
2. Shared-memory layout. Bank conflicts destroy kernels silently. The usual fix is an xor-swizzle or an extra padding column. The ncu metric above tells you how far you are from zero conflicts.
3. Register budget. Hit the 255-register-per-thread cliff and you spill to local memory, which on GB10 means LPDDR. __launch_bounds__(256, 2) (max 256 threads, 2 blocks per SM) is the decoration that constrains nvcc's register allocator. Use it. Check with --ptxas-options=-v.

Things that matter less on GB10 than on an H100:

• Shared-memory capacity (100 KB is generous for these shapes).
• Warp specialisation. SM121's scheduler is good enough that explicit producer/consumer warp roles rarely pay back on kernels of the shapes Atlas runs.
• Distributed shared memory. No NVLS / multi-CTA clusters to lean on — NVLS_ENABLE=0 is forced in the NCCL env.

CUDA graphs

Every supported batch size gets a captured graph at startup. The engine replays graphs with cuGraphLaunch, eliminating per-launch cuLaunchKernel overhead (~microseconds each, compounded over ~300 kernels per forward pass). Graph stability requires:

• Buffer addresses do not change. BufferArena pre-allocates at startup; the engine reuses the same DevicePtrs for every step. No alloc/free in the hot loop.
• Kernel launch parameters are data-dependent in a bounded way. Per-layer kernel launches pass a handful of scalar args that vary per step (current token count, current block index). Those are captured as CUgraphNodeParams.

Turning graphs off (--profile) disables capture — useful when you are profiling under nsys because graphs collapse every kernel into a single graph-launch event and defeat per-kernel timing.

Writing a new kernel — the minimum

1. Find the (hw, model, quant) leaf. Typically kernels/gb10/<model>/<quant>/.
2. Drop your_kernel.cu with the SPDX header and a extern "C" __global__ entry.
3. Decide the module name. Default is the file stem (your_kernel). Override in KERNEL.toml if you want a short name.
4. Call it from the layer. In spark-model/src/layers/<your_layer>.rs, gpu.kernel("your_kernel_module", "your_kernel_function") returns a KernelHandle. Store it in the layer struct at load time, not per-step.
5. Wire the launch via the KernelLaunch builder in spark-model/src/layers/ops.rs:

   #![allow(unused)]
   fn main() {
   KernelLaunch::new(gpu, self.kernel_handle)
       .grid([num_tokens, 1, 1])
       .block([256, 1, 1])
       .shared_mem(shared_bytes)
       .arg_ptr(input)
       .arg_ptr(output)
       .arg_u32(hidden_size)
       .arg_f32(eps)
       .launch(stream)
   }

6. Benchmark. Add a shape in atlas-spark-bench or a micro-benchmark in the relevant primitive crate. A kernel without a benchmark is not allowed to claim "faster".
7. Verify correctness against a PyTorch reference on a fixture tensor. Numerical diff tolerance: for BF16 outputs, abs-tol 1e-3 / rel-tol 1e-2 is a typical starting point.

Anti-patterns

• Don't branch on runtime flags inside the hot loop. If a kernel needs two variants (NVFP4 vs BF16 KV), make them two kernels.
• Don't try to be generic. The whole point is specialisation. A kernel that works for three batch sizes is usually slower than three kernels that work for one each.
• Don't call cudaDeviceSynchronize anywhere inside a kernel launch path. GpuBackend::synchronize(stream) exists for explicit syncs. Random device-sync calls break CUDA graph capture and defeat pipelining.
• Don't mutate __constant__ memory at runtime. Upload it once at load time; use it forever.

What to read next

• Per-format specifics: NVFP4, FP8
• Per-op specifics: Attention & Paged KV Cache, MoE, SSM
• Per-feature specifics: Speculative Decoding, XGrammar
• Measuring results: Benchmarking
• Authoritative designs: docs/adr/ in the repo

NVFP4 Quantization

NVFP4 is Atlas's flagship format on GB10 — 4-bit weights with FP8 block scales. Most Qwen and Nemotron checkpoints ship in it, and it's the KV-cache dtype that hits the best compression/quality balance for the Qwen3.5 family.

The numeric format

Each NVFP4 tensor is stored as two pieces:

• weight — [N, K/2] bytes, two E2M1 nibbles packed per byte.
• weight_scale — [N, K/16] bytes, one FP8 E4M3 scale per 16-element block along K.

E2M1 encoding (per nibble):

Eight values after sign. Per 16-element block, a scalar FP8 scale is stored — so the reconstructed value is scale[block] * e2m1_lut[nibble].

Why blocks of 16

Two reasons that both land on 16:

1. Tensor-core fragment sizes. SM121's MMA instructions work on 16 × k fragments along the K dim. Block-aligning the scale to 16 lets a single tile load reach all the scales it needs without a second indexed load.
2. Accuracy. 16 is small enough that outliers inside a block are rare; dequantizing against a scale computed over 16 elements is close enough to per-element calibration for modern transformer activation statistics. Larger blocks (64, 128) lose perplexity; smaller blocks (4, 8) waste scale bytes.

16 is the community standard for NVFP4 and what every HF checkpoint we load uses.

SM121's E2M1 conversion problem — and the software fix

Later Blackwell parts have a single-instruction conversion: cvt.rn.satfinite.e2m1x2.f32 takes two floats and emits two E2M1 nibbles. SM121 does not have this instruction. It is the headline hardware limitation of GB10 NVFP4.

Atlas's software fix — e2m1_branchless.cu — does the conversion in 7 ALU ops using the IEEE-754 bit pattern:

// Simplified sketch: convert a positive f32 to a 3-bit magnitude nibble.
// (sign bit is handled separately)
//
// The E2M1 value set for positive magnitude is:
//   {0, 0.5, 1, 1.5, 2, 3, 4, 6}
// corresponding to (exponent_field, mantissa) pairs:
//   (0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1), (3, 0), (3, 1)
//
// 7 compares against the mid-points of adjacent E2M1 levels map each
// f32 to one of the 8 values without a single branch.

uint32_t bits = __float_as_uint(x);
uint32_t abs  = bits & 0x7fffffff;
uint32_t sign = (bits >> 31) << 3;

// Thresholds (bit patterns of mid-points in f32)
bool ge_025 = abs >= 0x3e800000;  // 0.25
bool ge_075 = abs >= 0x3f400000;  // 0.75
bool ge_125 = abs >= 0x3fa00000;  // 1.25
bool ge_175 = abs >= 0x3fe00000;  // 1.75
bool ge_250 = abs >= 0x40200000;  // 2.5
bool ge_350 = abs >= 0x40600000;  // 3.5
bool ge_500 = abs >= 0x40a00000;  // 5.0
// Accumulate the magnitude index in 0..=7
uint32_t idx = ge_025 + ge_075 + ge_125 + ge_175 + ge_250 + ge_350 + ge_500;

uint32_t nibble = sign | idx;  // final 4-bit E2M1

Seven compares, seven adds, one bit shift, one OR. Fully branchless. Two nibbles per byte are produced by running the same sequence on lo and hi halves of a 64-bit pair and packing.

The payoff: NVFP4 round-trip (dequant → compute → requant) runs at full pipeline speed on SM121, despite the missing instruction. Exhaustive testing on 19 experiments (logged in trtllm-ngram-experiments.csv) established 29.6 tok/s as the TRT-LLM ceiling on the same model; Atlas's vLLM-path approach running NVFP4 through software-E2M1 CUTLASS hits 36.4 tok/s (CUTLASS MoE) and 59.9 tok/s (Marlin + MTP) on the same hardware. Software E2M1 is a 32× speedup vs the first "enable E2M1" build.

Dequantization in the GEMM kernel

Atlas does not pre-materialize BF16 from NVFP4. Instead, the GEMM kernel loads NVFP4 directly into shared memory and dequantizes on the fragment boundary just before the MMA:

1. cp.async a tile of weight (packed nibbles) + its weight_scale into shared memory.
2. On the consumer warp, unpack a 16×8 fragment of weights: convert each pair of nibbles to two BF16s using the E2M1 LUT in shared/constant memory, multiply by the block scale (broadcast from the scale tile).
3. Feed the BF16 fragment to mma.sync.aligned.m16n8k16.row.col.bf16.bf16.f32.
4. Accumulate in FP32, downcast to BF16 on output.

Pre-materializing BF16 would require ~8× more shared memory and throw away most of the compression benefit. The fragment-time dequant is the key to NVFP4 actually being fast.

NVFP4 KV cache

The same E2M1 format is available as a KV-cache dtype (--kv-cache-dtype nvfp4). K and V tensors are stored as NVFP4 with per-block FP8 scales; the paged cache allocator budgets 0.5 bytes + scale_bytes per element.

When a request hits prefix caching, the cached K/V is in NVFP4; the attention kernel reads directly and dequantizes at the MMA boundary, same pattern as the weight GEMM.

For coherence at long context, --kv-high-precision-layers N keeps the first and last N attention layers at BF16 — those layers' KVs are the most sensitive to precision loss. The default is 0; production deployments of the 122B model use 2.

When not to use NVFP4

• When a model's weights don't ship in it. A few checkpoints (Qwen3.6-35B-FP8) are FP8-native. Don't re-quantize — load them FP8 and use the FP8 KV cache. See the FP8 chapter.
• When you have only one GB10 and the model fits in FP8. FP8 weights require no software-E2M1 path, so the prefill/decode hot loops are slightly simpler and slightly faster per kernel-launch. The trade is the 2× memory increase vs NVFP4.
• When you're debugging a coherence regression. Fall back to BF16 first, then FP8, then NVFP4 — narrows the bug source quickly.

Why Atlas is not pursuing native FP4 MMA on SM121

Discovering this was the point of a multi-week research dive (logged as the "FP4 MMA GB10" project in the repo). The short version:

• SM121 silicon does not expose the relevant MMA or conversion instructions.
• CUTLASS 4.3's SM120 builders enforce cooperative-pipeline scheduling via static_assert — you cannot swap in pingpong mode where it would help.
• Every alternate MoE backend (TRTLLM, CuteDSL, DeepGEMM, Triton) fails with NotImplementedError: SM120 and above or crashes outright.
• The community benchmarks that claim native FP4 throughput on "Blackwell" are on SM100a / SM101a parts and do not apply.

Atlas's 131 tok/s on Qwen3.5-35B and 104 tok/s on Qwen3-Next-80B are, in this sense, the real GB10 ceiling — achieved through software E2M1, Marlin-style dequant-to-BF16, and MTP speculative decoding. New NVIDIA silicon would unlock another axis of improvement; on today's GB10, the numbers in the README are the answer.

Files to read

• kernels/gb10/<model>/nvfp4/e2m1_branchless.cu — the conversion.
• kernels/gb10/<model>/nvfp4/moe_prefill.cu, dense_gemm_nvfp4.cu — the GEMM tiles + fragment-time dequant.
• kernels/gb10/<model>/nvfp4/paged_decode_attn_nvfp4.cu — the NVFP4 KV attention.
• docs/adr/0004-nvfp4-fp8-quantization.md — the quantization decision record covering --kv-high-precision-layers.

FP8 Native Serving

FP8 is the second-most-common quantization format in Atlas after NVFP4. It's also the format where the most recent engineering work has landed — Qwen3.6 ships FP8-native, Nemotron's checkpoints are FP8, and Atlas now runs them end-to-end without a BF16 upcast on the critical paths.

The two FP8 checkpoint shapes

Atlas sees two layouts on disk, both handled by the format modules under spark-model/src/quant_format/:

1. Per-tensor scaled — weight (FP8 E4M3 bytes) + weight_scale (one f32 scalar per tensor). Common in vLLM-exported checkpoints.
2. Block-scaled — weight (FP8 E4M3) + weight_scale_inv (BF16, one scale per block_size × block_size tile, typically 128 × 128). Used by compressed-tensors FP8 checkpoints from Qwen and Nemotron.

Per-tensor scaled checkpoints can be read as a degenerate block case (block_size = ∞, one scale covers everything). The kernel code handles both with the same fragment-time dequant.

FP8 E4M3: range and quirks

E4M3 is sign(1) | exp(4) | mantissa(3), bias 7. Finite range is [-448, +448]. There is no infinity encoding; 0xFF / 0x7F are NaNs. The per-tensor scale maps the activation's dynamic range into E4M3's representable window.

Atlas ships a 256-entry FP8_E4M3_LUT in atlas-core/src/numeric.rs for CPU sanity checks and scale-inversion arithmetic at weight-load time. That module is also where the f32_to_bf16 round-to-nearest-even cast lives, byte-exact against PyTorch's float32 → bfloat16. The GPU hot path does not use the LUT — it uses the PTX instruction cvt.rn.bf16.e4m3 (FP8 → BF16 on the fragment boundary), which is available on SM121 unlike the NVFP4 instruction.

Native FP8 vs dequant-to-BF16

Two paths exist in the code today, selected per layer in spark-model::quant_format:

• Dequant-to-BF16 — read FP8 from memory, convert to BF16 in the GEMM fragment, do MMA in BF16. Simple, always correct, works for every FP8 checkpoint.
• Native FP8 — read FP8, feed directly to the FP8 MMA instructions (mma.sync.aligned.m16n8k32.row.col.e4m3.e4m3.f32). Keeps FP8 all the way through tensor cores, only converts to BF16 when writing activations back.

Native FP8 is the faster path and the one "FP8 native serving" refers to. The two gotchas:

1. KV-cache interaction. If the KV cache is FP8, the attention kernel needs to either keep it FP8 through the MMA (native) or convert it per-fragment (dequant). The paged_decode_attn_fp8.cu kernel does the former; the older paged_decode_attn_fp8kv.cu does the latter.
2. Calibration scales. Native FP8 demands well-conditioned scales. A calibration pass (--fp8-kv-calibration-tokens N) runs for the first N tokens, collecting online max-‖K‖ / max-‖V‖ stats and updating the KV scales before CUDA-graph capture. Without calibration, long-context FP8 KV drifts; with it, it matches BF16 to within measurement noise.

Typical deployment: --kv-cache-dtype fp8 --fp8-kv-calibration-tokens 256. 256 is enough warm-up; larger values cost prefill time without measurable quality lift.

The Qwen3.6 FP8 story

Qwen3.6-35B-A3B is FP8-native: weights, KV, MTP head, vision tower all FP8. Atlas's support here was a sequence of fixes logged across several bug sweeps:

• FP8 weight loading for native MTP (wave-6) — the NVFP4 MTP loader was force-BF16 when ignore_modules listed mtp.*; fixed to fall through to FP8 dequant when the scales were BF16-block rather than NVFP4-group.
• FP8 prefill shared-experts allreduce reorder (wave-6) — the shared-expert path was all-reducing FP8 activations across EP=2 before the final BF16 downcast, which silently dropped precision. Reordered so the allreduce sees BF16.
• FP8 KV calibration during CUDA-graph capture — graph capture froze the scales at their t=0 values; now calibration runs before capture, and capture picks up the converged scales.
• Spontaneous <think> fix — Qwen3.6's FP8 path exposed a reasoning-parser bug where the model emitted <think> outside the template's expected position. Fixed across four codepaths.

End-to-end result: Qwen3.6-35B-A3B serves at ~90 tok/s on GB10 with full FP8 coherence including Claude Code-style tool use. See project_coder_next_fp8.md (in the history notes) for the before/after.

When FP8 beats NVFP4

On GB10:

• Qwen3.6 (FP8-native) — obviously. Don't re-quantize.
• Nemotron-3 — FP8 is the only quant that preserves the Mamba-2 A/B/C/D projections' numeric character; NVFP4 drifts on long-context Mamba.
• Anything where you have the memory budget — FP8 is 2× NVFP4 bytes but doesn't pay the software-E2M1 cost per fragment, which can matter at small batch × large K shapes.

When NVFP4 beats FP8:

• 122B-class models that only fit in NVFP4. 76 GB NVFP4 weights leave room on a single GB10; 152 GB FP8 do not.
• Short-context deployments where the compression just saves money.
• Qwen3.5 — the family has native NVFP4 checkpoints and calibrates well.

KV cache dtypes (--kv-cache-dtype)

The full list:

The Turbo family is Atlas-specific: Walsh-Hadamard rotates out the outlier structure typical of transformer K/V activations before quantizing with an optimally-placed codebook. For the same bit count, turbo4 gives measurably lower per-token error than NVFP4 on models with large RMSNorm weights. It is purely additive — you opt in via --kv-cache-dtype turbo4; the NVFP4 path is unchanged. See docs/turboquant-plus.md.

Files to read

• kernels/gb10/<model>/fp8/ — per-model FP8 kernel sets (Qwen3.6 has its own leaf).
• kernels/gb10/<model>/<quant>/paged_decode_attn_fp8.cu — native FP8 KV attention.
• crates/atlas-core/src/numeric.rs — the FP8 E4M3 LUT and the f32 → BF16 RNE cast, with the PyTorch-parity vectors.
• crates/spark-model/src/quant_format/ — per-format descriptors and runtime dispatch.
• crates/spark-runtime/src/kv_cache.rs — KvCacheDtype::Fp8 sizing + calibration plumbing.
• docs/adr/0004-nvfp4-fp8-quantization.md — the authoritative quantization decision record.

Attention & Paged KV Cache

The attention path is where Atlas's biggest speedups land — up to 6.02× vs PyTorch decode, up to 4.95× vs PyTorch prefill. This chapter walks the kernels, the KV cache allocation model, and the pieces that make them fast on GB10.

Two kernels, two shapes

Prefill and decode look like different problems:

• Prefill — seq_len is large (prompt length), batch is typically 1 per forward. The GEMM has one long axis; memory pressure is on the Q·K^T tile.
• Decode — seq_len for the new query is 1; the KV cache accumulated so far is long (thousands of tokens). The GEMM is matrix-vector; memory pressure is on the full-history K/V load.

Different kernel shapes:

Prefill: Flash Attention v2 on SM121

The v47 prefill kernel is the "production" prefill path. It follows Flash Attention v2's tiled online-softmax pattern:

1. Load a tile of Q into shared memory once; it stays resident for the whole kernel.
2. Stream tiles of K and V through shared memory using cp.async (2 stages in flight).
3. For each K/V tile:
   • Compute partial Q·K^T into a register fragment.
   • Apply causal masking + softmax rescaling against the running max + running sum.
   • Multiply by V, accumulate.
4. After all tiles, normalize by the final sum.

Key design choices that matter on SM121:

• 2 CTAs/SM, not 1. The shared-memory budget is generous enough that two blocks per SM fit comfortably, and the second block absorbs scheduler bubbles from the first. 1 CTA/SM was the baseline; going to 2 was a ~15% win.
• mma.sync.aligned.m16n8k16 for Q·K^T and for A·V. These are the bread-and-butter BF16 MMA fragments on Blackwell.
• Shared-memory xor-swizzle on the Q tile to avoid bank conflicts during the mma load.
• Causal masking as a predicate inside the mma loop, not a separate kernel. A predicate add into the softmax scratch is almost free; a separate pre-mask would be a full kernel.

The _fp8kv variant adds one extra step: the K/V tiles are E4M3 in memory, not BF16. Loads are half the bytes; the first thing the consumer warp does is cvt.rn.bf16.e4m3 on each fragment. Net effect: ~1.8× BW saved on the K/V load, at the cost of one extra instruction per fragment.

Prefill numbers on Qwen3-Next-80B shapes (hidden=2048, 16Q / 2KV heads, head_dim=256):

The dramatic win at seq=256 is the kernel hitting its sweet spot where the cp.async pipeline is fully saturated and shared memory is the bottleneck instead of DRAM. At small seq the fixed launch cost dominates; at large seq PyTorch's Flash-Attention-2 backend catches up.

Decode: split-K online softmax

Decode attention is a different kernel because the shapes are different. For each new token, we need Q (one row) × K (full history). The hot axis is K — thousands of elements of history per head, dozens of heads.

Atlas's decode kernel parallelises across two axes:

1. Head — one warp per (Q-head, KV-head) pair.
2. Split-K — the full K/V history is chopped into N chunks; each chunk gets a CTA that produces a partial softmax + partial attention output. A second pass reduces across chunks.

The split count N is adaptive — chosen per call based on history length and current batch. Long history → more splits. This is the "adaptive split count" in the kernel table.

The online softmax pattern survives from prefill but with a different shape: each CTA maintains a running max and a running exp-sum for its K-chunk. The reduction across chunks at the end is a single warp's work.

NVFP4 K/V in decode is where the headline throughput comes from. The decode kernel reads packed E2M1 nibbles, unpacks two per byte in registers, multiplies by the block scale, and feeds BF16 to the Q·K^T and A·V MMAs. K and V in memory are half the bytes of BF16 → 2× BW saved → 2× decode throughput on BW-bound steps.

Decode numbers (same Qwen3-Next-80B shapes):

At 4k history we're at ~63% of GB10's 273 GB/s peak — tight but well below the roof. Further improvement here is the hot lane of kernel work.

The paged KV cache

Atlas follows vLLM's paged-attention model: KV is allocated in fixed-size blocks (default 16 tokens per block) from a pool, not per-sequence. Key advantages:

• No fragmentation. A completed request returns its blocks to the pool; a new request claims fresh ones. No moving, no compaction.
• Copy-on-write prefix sharing. When prefix-caching hits, the prefix's blocks are shared between the cached and new sequence until divergence.
• Scheduler-friendly. Memory budget is expressible in "free blocks", a scalar — easy to reason about under load.

spark-runtime::kv_cache::PagedKvCache owns the pool. KvCacheConfig derives the pool size from:

• --max-seq-len × --max-batch-size → total token capacity
• Model's num_hidden_layers × num_key_value_heads × head_dim × 2 (K and V) × bytes-per-elt (from KvCacheDtype) → per-token storage
• Block size (16) → number of blocks

Block allocation is O(1) from a free list. Eviction is policy-driven — the scheduler picks victims (LRU by default) when the pool is full.

Paged attention in the kernel

The decode kernel takes three extra arguments beyond a non-paged version:

• block_tables — [batch, max_blocks_per_seq] of block indices.
• context_lens — [batch] of valid token counts.
• block_size — compile-time constant for indexing arithmetic.

Inside the kernel, each CTA computes its K/V pointer for a given history position by indexing into the block table: block_ptr = k_cache + block_tables[seq][pos/16] * block_size * head_dim * bytes_per_elt. Gather-SMEM-MMA: gather the block pointers into shared memory, then run MMA against the resulting shared tile. Pattern from the FlashInfer paper (MLSys 2025 Best Paper, cited in the README).

RadixAttention prefix caching

Prefix caching is a radix-tree lookup on the token prefix. The tree's nodes own KV blocks; a lookup that matches N tokens of prefix returns those N tokens' blocks already resident. --enable-prefix-caching turns it on.

Typical hit rates in production:

• Agent workloads (Claude Code, OpenCode, Cline): 90%+ on the system prompt + tool schemas.
• Multi-turn chat: 60–80% on the conversation context.
• Cold workload (single one-shot): 0%.

TTFT goes from ~400ms cold to ~40ms warm on Qwen3.5-35B. The prefix-cache chapter of the engine test suite validates the hit rate and the byte-identical output under warm vs cold.

For hybrid SSM+attention models, a matching "Marconi" SSM snapshot cache lives in spark-runtime::prefix_cache — the SSM state at the end of the prefix is checkpointed alongside the attention KV, so a warm hit reconstructs the full model state, not just the attention cache. Without this, prefix caching on an SSM model would produce silently incorrect output.

Files to read

• kernels/gb10/<model>/<quant>/inferspark_prefill_v47.cu — the prefill kernel.
• kernels/gb10/<model>/<quant>/paged_decode_attn_*.cu — decode kernel variants.
• kernels/gb10/<model>/<quant>/kv_cache_append.cu — per-token KV write.
• crates/spark-runtime/src/kv_cache.rs, prefix_cache.rs, radix_tree.rs — Rust side.
• README "Citations" — links to the Flash Attention 2, Flash Attention 4, FlashInfer, SageAttention 3, and LeanAttention papers.

MoE Routing & Experts

Mixture-of-Experts is where the Atlas supported-model matrix gets most of its diversity. Qwen3.5 routes 128 experts top-10. MiniMax-M2.7 routes 256 experts top-8 with sigmoid gating (not softmax). Gemma-4 has shared experts. Nemotron-H has an MoE-only layer variant (no mixer, just routing + FFN). The engineering problem is making all these shapes share one kernel pipeline without losing per-shape efficiency.

The generic MoE block

residual ─► gate (linear) ─► softmax/sigmoid ─► topk ──► dispatch ──► experts ──► gather ──► weighted sum ──► residual
                                                    │         │
                                                    │         └─ (N different FFN weights, one per expert)
                                                    │
                                                    └─ (tokens routed to ≤ k experts each)

Five kernels contribute to one MoE block:

1. Gate — linear projection. Uses the same GEMM kernels as attention projections.
2. Top-k + softmax/sigmoid — per-token expert selection.
3. Dispatch — token scatter to expert-local arrangement. Conceptually a permutation.
4. Expert FFN — k copies of silu_mul_quant(up(x)) → down(...). Grouped GEMM with one batch per expert.
5. Gather + weighted sum — reduce expert outputs back to the residual shape.

Prefill: grouped GEMM

For prefill (many tokens per step), the expert FFN is a grouped GEMM: M tokens split across N expert buckets, each bucket does its own (B_e, K) × (K, N_out) matmul.

Atlas's moe_prefill.cu kernel takes the "sort tokens by expert, then one GEMM per expert group" approach. The sort is parallel — dispatch tables are computed in a single pre-kernel pass. The GEMM itself uses the same 16×8×16 BF16 MMAs as attention, one CTA per expert-bucket.

This is the kernel that hits the headline MoE W4A16 256-expert 80-token: 3.87× PyTorch number in the benchmark table.

Decode: token-level MoE

For decode (one token per sequence per step), grouped GEMM is the wrong shape — each "group" has one or two tokens, and the launch overhead dominates. Atlas's decode MoE kernels (moe_expert_relu2_down_shared.cu and moe_shared_expert_fused_fp8.cu) take the opposite approach: one warp per (token, expert-in-topk) pair, compute the expert's up-proj and down-proj in a single kernel.

The shared-expert path for models like Gemma-4-26B fuses the shared expert's compute with the per-token dispatch, avoiding a round-trip through memory.

The 256-expert case: MiniMax-M2.7

MiniMax-M2.7 is the extreme end of the MoE support matrix: 256 experts, top-8, sigmoid routing (not softmax). The sigmoid variant lets multiple experts contribute independently without the softmax normalisation; the tradeoff is that topk has to pick from a longer tail of meaningful scores.

Atlas's minimax_moe layer uses:

• norm_topk_prob = true semantics (topk weights normalised by their sum, not softmax).
• A dedicated topk_sigmoid kernel path that reads sigmoid scores rather than post-softmax probabilities.
• An EP=2 token-dispatch pipeline — with 256 experts, a single node is impractical; MiniMax ships only as EP=2.

The wave-17 bug fix that landed the "M2.7-NVFP4 EP=2 full PASS" milestone addressed four subtle issues:

1. rms_norm placement — the rms_norm output was being reused after the MoE dispatch modified it, causing subtle coherence drift.
2. norm_topk_prob — initial implementation used softmax-normalised weights; corrected to sum-normalised sigmoid.
3. FP8-free path — MiniMax weights are NVFP4 end-to-end; early code had an accidental FP8 upcast in the shared-expert path.
4. Template-forced thinking detection — MiniMax's chat template seeds <think> differently from Qwen; the detector needed to distinguish.

Shared experts

Several models (Gemma-4, Qwen3-VL) have a "shared expert" FFN that runs for every token in addition to the k routed experts:

moe_out = topk_weighted_sum(expert_out) + shared_expert(x)

The shared expert is a standard dense FFN run alongside the routed experts. Atlas fuses the shared FFN with the gather step in the decode kernel (moe_shared_expert_fused_fp8.cu) to save a memory round-trip.

Expert parallelism (EP=2)

For 122B, 119B, 229B models, the experts don't fit on one GB10. Split them across two ranks:

• Gate runs on every rank (replicated).
• Token top-k assigns each token to k experts; tokens destined for remote experts are sent via reduce_scatter.
• Expert FFN runs locally on the owning rank.
• Results come back via all_gather.

The collective ops go through spark-comm::CommBackend; the EP dispatch logic in crates/spark-model/src/layers/moe/forward_ep.rs handles the local-vs-remote bucketing. See spark-comm and Multi-GPU & EP=2.

Routing edge cases

A handful of bug-sweep wave findings landed in the MoE code:

• Sibling stride bugs — the attention and SSM layer code assumed K=2 MTP strides in the expert outputs; corrected to support K=1/2/3 uniformly.
• Slot-keyed verify*_graph caches — CUDA graph instances for MTP-verify were keyed by batch size alone; needed (batch, k) to avoid replaying a K=1 graph on a K=2 step.
• MoE topk bounds — a weight-loader edge case where num_experts_per_token exceeded num_experts for a few experimental checkpoints; added an assertion at load.
• MoE topk weights == 0 guard — a numerical edge case where all-zero gate logits produced NaN after normalisation; now clamped.

Why MoE is fast on GB10

Three design choices stack:

1. Grouped-GEMM prefill dispatches fewer kernels than per-token — the fixed overhead per expert group is tiny compared to the math.
2. Fragment-time dequant — NVFP4 expert weights are never materialised to BF16; they unpack at the MMA boundary, same as attention (see NVFP4 deep dive).
3. Fused SiLU×Mul + quant — the up-proj output is never written to memory in BF16; it's SwiGLU'd and re-quantized in one kernel before the down-proj reads it back. Saves the entire BF16 intermediate (a major BW win on MoE, where the intermediate is batch × topk × moe_intermediate_size).

The end result on a 256-expert MoE at batch=80: Atlas at 8.43 ms vs PyTorch at 32.65 ms — 3.87×. That's the biggest single-kernel win in the benchmark table.

Files to read

• kernels/gb10/<model>/<quant>/moe_prefill.cu — grouped-GEMM prefill.
• kernels/gb10/<model>/<quant>/moe_expert_relu2_down_shared.cu — token-level decode MoE.
• kernels/gb10/<model>/<quant>/moe_shared_expert_fused_fp8.cu — fused shared-expert path.
• kernels/gb10/minimax-m2-229b/nvfp4/moe_w4a16_grouped_gemm.cu — routed grouped-GEMM kernel.
• crates/spark-model/src/layers/moe/ (forward.rs, forward_prefill.rs, forward_ep.rs, …) — Rust side; forward_ep.rs holds the EP=2 token dispatch.
• docs/adr/0007-tp-ep-composition.md — TP/EP composition design record.
• docs/adr/0011-ep-batched-decode-optimization.md — EP batched-decode optimization.

SSM / Mamba / GDN Layers

State-Space Models (SSMs) are the non-attention half of every hybrid model in the support matrix: Qwen3.5 (GDN), Qwen3-Next (SSM), Nemotron-H (Mamba-2), Qwen3.6 (GDN + vision). Their speedups vs PyTorch are the largest in the whole kernel suite — up to 9.95× on conv1d prefill, 7.89× on GDR prefill.

The three variants in Atlas

All three share structure: a QKVZ projection (expanded 4-way linear), a causal conv1d, a selective linear recurrence, and a gated output normalisation. The differences are in the recurrence and the gating.

Why SSMs are fast on Atlas

Three things:

1. Fused QKVZ preprocess. The ssm_preprocess.cu kernel deinterleaves the 4-way projection output (Q, K, V, Z) and computes the GDN gate (softplus + sigmoid) in one pass. PyTorch does this in three separate kernels plus a reshape; Atlas does it in one.
2. Fused Gated Delta Rule (GDR). The linear recurrence itself is hand-written — one warp per (batch, head) walks the token sequence, maintaining the hidden state in registers. Causal conv1d is fused into the same pass.
3. Hand-rolled causal conv1d. The conv1d kernel takes advantage of the fixed small filter size (kernel width 4) to hold the entire filter in registers and stream inputs through.

Benchmark numbers (dim=8192):

The 9.95× at conv1d seq=512 is the largest compute speedup in the whole repo. PyTorch's causal_conv1d_fn on this shape walks the full sequence per batch element; Atlas's kernel fuses the whole thing into one launch with shared-memory tiling.

The SSM state

Unlike attention, SSMs carry a compressed hidden state across the sequence. For Mamba-2 with d_inner = 8192 and d_state = 128, the state is [8192, 128] FP32 per layer — about 4 MB per layer per sequence. For a 36-layer model, that's 150 MB per sequence — comparable in size to a full attention KV cache.

Chunked SSM prefill

A chronic issue: prefilling a long prompt through an SSM layer requires computing the full linear recurrence from scratch. The intermediate state can be gigabytes if the prompt is 16k tokens and the batch is moderate.

The chunked prefill path in kernels/gb10/<model>/<quant>/ breaks the prefill into chunks of (typically) 1024 tokens. Each chunk:

1. Starts from the state at the end of the previous chunk.
2. Processes its tokens through the recurrence.
3. Writes the ending state for the next chunk.

Savings: ~7–9 GB of scratch memory for long-context prefill, at negligible perf cost. Before chunked prefill, 122B at 8k context wouldn't fit in 119.7 GB. After it, it does.

See docs/adr/0003-hybrid-ssm-attention.md for the chunked-SSM-prefill design (there was a BF16 paged-dispatch bug where the FP8 kernel was called on a BF16 cache, producing NaN; fixed in wave-4).

Marconi: SSM state snapshots for prefix caching

A naive prefix cache on an SSM model reads the attention KV for the prefix and — because it doesn't have the SSM state — recomputes the SSM layers from scratch. That defeats most of the win.

Atlas's Marconi (inside-joke name for the SSM snapshot cache) stores the SSM state at the end of each cached prefix alongside the KV. A warm prefix-cache hit restores both the attention KV and the SSM state. Output is byte-identical to the cold run.

Costs: an extra ~GB of snapshot storage per top-level cache entry. Worth it for repeat agent workloads. Controlled by --ssm-cache-slots (default 16) and --ssm-checkpoint-interval (default 256). See the marconi.md note and crates/spark-runtime/src/prefix_cache.rs.

GDN: softplus + sigmoid fusion

Gated Delta Rule does a per-token gate computation:

g = softplus(dt) * sigmoid(beta)

where dt is the delta-time projection. The gate is used both to attenuate the state update and to form the output. Atlas fuses this into ssm_preprocess.cu — the gate appears as a byproduct of the QKVZ deinterleave.

Numerically, softplus is the chronic problem: softplus(large x) = x, but softplus(very large x) can overflow in FP32 if you compute naively. Atlas uses the standard max(0, x) + log1p(exp(-|x|)) stable form.

Known gotchas (lessons from the bug sweeps)

• SSM catastrophic forgetting — an older version had a bug where the snapshot state was sometimes restored with a stale conv1d buffer, causing a slow drift of coherence over long agentic sessions. Fixed by also snapshotting the conv1d tail-buffer.
• Upstream Mamba bug on chunked prefill — the vLLM fix approach (different from Mamba-2 reference) is what Atlas uses; the reference implementation has a subtle boundary issue at chunk seams.
• GDN register-tile experiments — the bench gdn_regtile_results.md tracks a long tail of tile-shape experiments. The current production choice is a middle ground; alternate shapes win on specific seq_len ranges but fail the full regression suite.

What the code looks like

crates/spark-model/src/layers/qwen3_ssm.rs and nemotron_mamba2.rs contain the layer-level state machines. They:

1. Run the fused preprocess kernel to produce Q, K, V, Z and gate.
2. Run the causal conv1d on the Q/K/V branches.
3. Run the selective linear recurrence (GDR or Mamba-2 variant) to produce the output.
4. Run gated_rms_norm(output, Z, weight) for the final gated normalisation.

Each step calls into spark-runtime::GpuBackend via the layer's cached KernelHandles — the per-step Rust code is short because the real work is in the 3–4 kernel calls it issues.

Files to read

• kernels/gb10/<model>/<quant>/ssm_preprocess.cu, gdr.cu, causal_conv1d.cu
• crates/spark-model/src/layers/qwen3_ssm.rs, nemotron_mamba2.rs
• crates/spark-runtime/src/prefix_cache.rs (Marconi SSM snapshot)
• README "Atlas Spark" section — the SSM/GDN story in narrative form

Speculative Decoding (MTP)

MTP = Multi-Token Prediction. The flagship feature that takes Qwen3.5-35B from ~70 tok/s to 131 tok/s on a single GB10, and Qwen3-Next-80B from ~70 to 104. Speculative decoding with a model-native draft head.

The idea

Plain decode generates one token per forward pass. If the model is certain about the next few tokens, we could generate more than one per pass — the bottleneck is almost always memory bandwidth, not compute, so generating 2 or 3 tokens costs roughly the same as 1.

MTP does exactly that:

1. The main model forward produces logits for the next token (call it $t_{+1}$).
2. A small MTP head — a few transformer layers that hang off the main-model hidden state — predicts $t_{+2}$, $t_{+3}$, …, $t_{+k}$ as drafts.
3. The main model runs one extra forward pass on the drafted positions to verify them, producing the "true" logits at each drafted position.
4. We accept the longest prefix of draft tokens where the verify logits' argmax agrees with the draft. After the first disagreement (or the end), we take the main model's next token.

For K=2 (one draft), the best case is 2× throughput: one draft + one verify yields two tokens per verify pass. The expected speedup depends on draft acceptance rate — Qwen3.5 + NVFP4 MTP achieves ~85% acceptance on short-prompt benchmarks, which maps to ~1.8× throughput in practice.

Atlas's numbers:

The MTP head

The MTP head lives in the checkpoint under the mtp.* prefix. For Qwen3.5, it's 1–2 transformer blocks (depending on K) fed from the penultimate layer's hidden state. It is trained jointly with the main model, so draft distributions match what the main model would actually emit.

--num-drafts controls K (the number of draft tokens). The per-model default comes from MODEL.toml. Most models cap at K=2; some (MiniMax-M2.7) support K=1/K=2 natively. The MultiModuleMtpHead caps naturally — more layers means more compute per verify, at some point not worth it.

--mtp-quantization must match the main-model checkpoint quantization. For an NVFP4 main checkpoint with an NVFP4 MTP head, --mtp-quantization nvfp4. Mixing is an error and will produce gibberish.

The verify loop

Pseudocode for a K=2 step:

# One step produces up to 3 accepted tokens.
main_logits, main_hidden = main_model.forward(cur_tokens, kv_cache=shared)
t_1 = sample(main_logits[-1])

# Draft the next two tokens from the MTP head
t_2, t_3 = mtp_head.draft(main_hidden, prev=t_1)

# Verify both in one main-model forward
verify_logits = main_model.forward([t_1, t_2, t_3],
                                   kv_cache=shared,
                                   append_kv=True)
a_2 = argmax(verify_logits[1])   # true next token after t_1
a_3 = argmax(verify_logits[2])   # true next token after (t_1, t_2)

# Accept the longest matching prefix
if a_2 == t_2:
    accept(t_1, t_2)
    if a_3 == t_3:
        accept(t_3)
        next_seed = t_3
    else:
        accept(a_3)
        next_seed = a_3
else:
    accept(t_1, a_2)   # drop drafts
    next_seed = a_2

Three subtle correctness requirements that must hold throughout:

1. KV cache must track accepted tokens, not drafts. If a draft is rejected, its KV contribution must be unwound.
2. The SSM state must track accepted tokens, not drafts. For hybrid models, the SSM recurrence is stateful — verify passes update the state; rejects must roll back.
3. Sampler state (running RNG, penalty counters) must track accepted tokens, not drafts.

All three are implemented in crates/spark-model/src/speculative.rs and the paired rewind_kv_cache + rewind_mamba_state hooks in spark-runtime.

The bug-sweep history

MTP is the single subsystem with the most documented bug-sweep history. Lessons that shaped the current code:

• seq_len += k - 1 off-by-one (Pass-16). The MTP scheduler bootstrap violated the "tokens[0] already in seq.tokens" precondition; caused Fibonacci drift on 80B-MTP. Fixed to seq_len += k.
• WY (whisperer/verifier) state desynchronisation (Pass-16). State clamping across the draft/verify boundary was off by one step.
• v_contiguous + MTP (Pass-6). The v_contiguous optimisation broke MTP's KV append because it assumed one token per call. Fixed with an explicit K-aware KV-append path.
• Sibling stride bugs (Pass-22). qwen3_attention and qwen3_ssm layers assumed K=2 for stride arithmetic; corrected to handle K=1/2/3.
• Slot-keyed verify*_graph caches (Pass-22). CUDA graphs for verify were keyed by batch only; needed (batch, k) to avoid replaying a K=1 graph on a K=2 step.
• NVFP4 MTP loader force-BF16 (Wave-6). When the checkpoint's ignore_modules listed mtp.*, the loader accidentally forced BF16 — disabled MTP on models where it should have worked. Fixed to fall through to the proper dequant.
• MTP logit masking (most recent). Masking MTP draft logits at propose time (disallowing tokens the main model can't produce from the current position) improved tool-call throughput by +37%.

MTP with tools

A hidden value of MTP is tool-call throughput. Tool calls are structured (JSON or XML); most tokens in a well-formed call are predictable conditional on the opening delimiter. The logit mask at propose time filters draft tokens that would break the grammar; acceptance jumps from ~70% to ~95% inside a tool call.

On agentic workloads (Claude Code, OpenCode, Cline), this compounds because a large fraction of generated tokens are tool calls. The +37% throughput is measured end-to-end against a real agent.

Self-speculative (no MTP weights)

--self-speculative is the fallback for models without an MTP head: layer-skipping drafts. The "drafter" runs the main model with some attention + FFN layers skipped, producing a fast-but-approximate draft; the full model verifies.

Acceptance rate is lower (~60%) than MTP (~85%), but it works on any model. Atlas ships this for coverage; operators typically use MTP when the checkpoint supports it.

N-gram speculative (CPU-side)

--ngram-speculative is the other fallback: an n-gram pattern matcher on recent output. If the model is repeating a token pattern (e.g. verbatim quoting a document), the matcher predicts the continuation directly. Acceptance is binary (0 or 100%), and the average rate on open-ended generation is low, but on certain workloads (summarisation, re-ranking) it's free throughput.

N-gram speculative was experimented with heavily on TRT-LLM (see the project_ngram_* notes); Atlas's Rust implementation lives in spark-server/src/ngram.rs and is much simpler.

Files to read

• crates/spark-model/src/speculative.rs — the verify + accept loop.
• crates/spark-runtime/src/kv_cache.rs — rewind_kv_cache.
• kernels/gb10/<model>/<quant>/ — there isn't a "MTP kernel"; MTP reuses the main model's attention/MoE kernels with different shapes.
• docs/SPEC-DECODING-TODO.md — authoritative design + outstanding items.
• docs/ATLAS_SPARK_JOURNEY.md — release journey and bug-sweep history.

Constrained Decoding (XGrammar)

Constrained decoding lets Atlas force the model to produce output that conforms to a grammar — the subset of tokens that could continue the current partial output while keeping the output valid is computed at every step; invalid tokens get their logits set to -inf before sampling.

Atlas uses XGrammar for this. It is the machinery that makes tool calls reliable (no invented field names, no broken JSON/XML), and it's the substrate for any "structured output" feature (response_format, JSON schema conforming, etc.).

The problem

Without constrained decoding, an LLM producing a tool call can:

• Invent field names that don't exist in the schema.
• Forget commas.
• Close JSON objects with the wrong bracket type.
• Break inside a string literal because the tokenizer merged characters across a boundary.
• Switch formats mid-call (XML close tag when the template expects JSON).

Every one of these has been seen in the wild on large models, even Qwen3.5-class models. A clean tool-call parser can't fix them — the broken token sequence comes out of the sampler before the parser sees it.

Constrained decoding intervenes upstream: at the sampler, we know which tokens are syntactically legal next given the current state of the output. Everything else gets masked.

XGrammar's trick

The naive version of constrained decoding is expensive: at every sampling step, compile the grammar against the current output prefix and enumerate the allowed next tokens. That's a parser pass over the vocabulary (tens of thousands of tokens) per step — prohibitive.

XGrammar (paper, Xiamen University / CMU 2024) does two things that make it tractable:

1. Pre-compiles a token-bitmap automaton. For each grammar state, precompute a bitmask over the whole vocabulary indicating which tokens are legal. At runtime, transition the automaton by the sampled token, look up the new bitmap — O(1) per step.
2. Handles tokeniser boundary cases. Real-world tokenizers merge characters across grammar-legal boundaries (e.g. the byte-pair-encoded token ",\n" crosses a JSON key/value boundary). XGrammar's compiler handles these at compile time by enumerating all byte-level prefixes each token can legally complete.

The cost is a grammar compilation step (~ms for typical JSON schemas), amortised across all requests using that schema.

How Atlas uses it

Atlas ships XGrammar as a pure-Rust in-tree crate, crates/xgrammar — a from-scratch port of mlc-ai/xgrammar v0.1.32, with no C++ core, no cxx FFI bridge and no build script. (It was previously a vendored vendor/xgrammar-rs/ binding wrapping the C++ engine; that directory is gone.) The call sites:

• Tool calls — when the request includes tools: [...], Atlas derives an XGrammar grammar from the function schemas + the model's tool-call format (Hermes JSON, Qwen3-coder XML, Mistral JSON). The grammar enforces: opening delimiter → valid function name → opening args bracket → schema-conforming JSON/XML → closing delimiter. --tool-max-tokens caps the total argument-generation length.
• Response-format structured output — OpenAI-compatible response_format: {type: json_schema, json_schema: {...}}. Atlas compiles the schema into an XGrammar grammar and constrains the entire response.
• Reasoning boundaries — the reasoning parser uses a lightweight grammar to enforce that <think>...</think> blocks close cleanly when --max-thinking-budget kicks in, preventing the unclosed-think bug that blocked Claude Code compatibility on Qwen3.6.

At the sampling step:

#![allow(unused)]
fn main() {
let logits = /* model logits [vocab_size] */;
if let Some(grammar) = active_grammar_for_request {
    let mask = grammar.current_token_mask();    // &[u32] bitmap
    apply_mask_in_place(&mut logits, mask);     // -inf for disallowed tokens
}
let token = sampler.sample(&logits);
grammar.advance(token);                         // transition automaton
}

The mask-apply and advance calls are both O(1) — a single bitmap test per token, a single state transition per step.

The integration history

XGrammar integration shipped across two substantial work items:

• xgrammar-integration-plan (initial) — wiring the FFI, compiling grammars on request, applying masks in the sampler, streaming boundary handling.
• xgrammar2-upgrade-plan (current) — moved to the 2.0 API with better handling of long schemas, on-the-fly grammar recompilation, and support for Anthropic-style nested-XML tool-call formats.

The payoff compounds with MTP: constrained decoding inside an MTP draft mask blocks draft tokens that would break the grammar, raising the draft acceptance rate from ~70% to ~95% during tool calls. The +37% tool-call throughput win (referenced in the MTP chapter) is a direct result.

Opencode & markdown fences

A specific bug worth noting: when a model emits a tool call inside a markdown code fence, Atlas's tool-call parser originally ate the surrounding fence characters — the closing backticks came through as "extra content" and broke downstream code that expected clean JSON. Fixed in wave-1 of the bug sweeps by making the parser markdown-fence aware; XGrammar then enforces the fence is balanced.

A related hallucination class: the Qwen3-coder XML format allows the model to emit the literal string </tool_call> inside a JSON string value. The parser now disambiguates, and XGrammar's grammar masks it at the source.

When to turn it off

XGrammar is lightweight but not free. For vanilla free-form generation (no tools, no response_format), the sampler skips the mask path entirely — there's no active_grammar. For workloads that explicitly want the model to deviate from a schema (creative tool exploration), setting tool_choice: "none" disables the grammar.

The one place where constrained decoding can interact badly with sampling: very low-entropy grammars combined with temperature=0 greedy sampling can produce repetitive output if the grammar masks the "natural" next token. --default-top-n-sigma and --default-min-p help; dropping temperature below 0.1 is rarely worth it on constrained paths.

Files to read

• crates/xgrammar/ — the engine itself; DESIGN.md there covers the Tier-3 synthesis and PORT_PLAN.md the port roadmap.
• crates/spark-server/src/grammar/ — per-request grammar compilation and mask application.
• crates/spark-server/src/tool_parser/ — tool-call format → grammar translation.
• docs/adr/0010-vendor-xgrammar.md — the decision record behind the original vendoring (superseded by the pure-Rust port).

OpenAI-Compatible Server

Atlas serves via spark-server — an OpenAI and Anthropic compatible HTTP API over axum. This chapter is the operator's reference for CLI flags, protocols supported, and the knobs that matter in production. The authoritative flag list is always spark serve --help; the headings below match the groupings in the CLI so cross-referencing is easy.

CLI structure

spark serve <MODEL> [--flags...]
spark serve --model-from-path <PATH> [--flags...]
spark --version
spark --help

Every runtime configuration flag has a long-form name. Most are documented inline with #[arg] doc-strings in crates/spark-server/src/cli/serve_args.rs.

Model selection and I/O

Memory / budget

Production rule of thumb for tight single-GPU deployments of 100B+ models: drop --max-prefill-tokens to 2048 and --max-batch-size to 1. The default 8192 sizes the scratch arena, not the KV pool; tuning down frees hundreds of MB.

KV cache precision

The flag takes a number or one of three words: auto (a fixed alias for 2, not a heuristic) and max/all (every attention layer). Anything else that fails to parse warns and falls back to 0.

--kv-high-precision-layers 0 is not "no promotion" for the turbo* family. 0 means defer to the per-dtype automatic value (main_modules/kv_dtypes.rs::auto_high_precision_layers, applied at serve_phases/kv_cache.rs). For bf16 / fp8 / nvfp4 the automatic value is None, so 0 really is zero. For every turbo* and asymmetric variant it is not: turbo2 and bf16k_turbo3v promote max(4, ⌈4·L/5⌉) layers, and all the others promote max(2, ⌈L/3⌉) — roughly a third of attention layers forced to BF16, which also shrinks the KV pool. Pass an explicit non-zero value if you want to control it; there is no spelling of this flag that promotes nothing under a turbo dtype.

See FP8 and NVFP4 for the trade-offs. Atlas's recommendation per model family:

• Qwen3.5 family → nvfp4 KV.
• Qwen3.6 / Nemotron-H → fp8 with calibration.
• 122B-class (memory-constrained) → nvfp4 + --kv-high-precision-layers 2.
• Everything else → fp8 (safe default).

Speculative decoding

See the MTP deep dive. Use only one of --speculative, --self-speculative, --ngram-speculative — but note this is guidance, not an enforced constraint: none of the three carries a clap conflicts_with and cli/validate.rs has no rule for the combination, so passing several parses and serves. The scheduler then resolves them by silent precedence (ngram → self-spec → MTP) rather than rejecting the config. --dflash is enforced — it declares conflicts_with = "speculative".

--num-drafts is also not a plain constant: when it is still 1, the model's MODEL.toml default_num_drafts replaces it (serve_phases/config.rs). On qwen3.6-27b that is 3, i.e. K=4.

Scheduling / caching

Agent workloads (Claude Code, OpenCode): always enable --enable-prefix-caching and --scheduling-policy slai. The prefix cache dominates wall-clock for system prompts + tool schemas; SLAI keeps streaming smooth under concurrent load.

Multi-GPU (Expert Parallelism)

See Multi-GPU & EP=2 for the full setup, including the NCCL env vars that matter on GB10.

Reasoning / tools

Observability / experimental

Endpoints

Rate limiting and auth

• --require-auth (with --auth-token <key> or --auth-tokens-file <path>) — requires an Authorization: Bearer <key> header on write endpoints. The presented token must match one of the loaded tokens (constant-time compare); there is no "accept any key" mode.
• Token-bucket rate limiter per key (crates/spark-server/src/rate_limiter.rs). Off by default; enable by setting ATLAS_RATE_LIMIT_RPM (requests/min) and/or ATLAS_RATE_LIMIT_TPM (tokens/min) > 0 (bursts via ATLAS_RATE_LIMIT_BURST_RPM / ATLAS_RATE_LIMIT_BURST_TPM, default = the cap). A MAX_KEYS DoS guard bounds the key table.
• Body-size limit env-configurable via ATLAS_MAX_BODY_BYTES (default 32 MiB — main_modules/serve_router.rs).

Changing the model without restarting

A running server can replace its model in place. Three routes reach it, and they share one code path:

• The dashboard. spark serve with no MODEL starts the listener and opens the Library, where a model and one of its recipes can be picked. Selecting one loads it and returns to the Main view.
• The Library, on a server that is already serving. Same flow; the running model is released first.
• A client request naming a different model, Ollama-style — off unless --auto-swap is passed.

--no-auto-swap forbids request-triggered loading outright and wins over --auto-swap regardless of order. For deployments where the served model is part of the contract, pass it: no client can then change what the endpoint is running, whatever else is on the command line.

Even with --auto-swap, a swap needs a request whose model resolves to a DIFFERENT model with a known recipe. A name that is absent, unrecognised, or already live is served by the current model, exactly as before the flag existed.

What a swap preserves. Stored responses and conversations, rate-limiter buckets, the API-key policy and --dump all belong to the process, not the model, and survive unchanged — including while no model is loaded, when GET /v1/conversations/{id} still answers rather than reporting the model missing.

What it cannot change. The listening socket is bound once for the process lifetime, so a recipe's host/port are ignored with a warning and the model serves on the address already bound. Hot-swap is single-node: a recipe needing more than one rank is refused rather than half-applied.

While it runs. In-flight requests finish on the model they started on. New requests that need a model get 503 model_not_loaded — retriable, and the same shape /health already reports during startup. A load that fails restores the previous model, and one this build has no kernels for is refused before anything is released.

Chat templating

Tokenization uses the HF tokenizers crate plus minijinja for chat templates. Atlas ships its own template overrides for a handful of models in jinja-templates/<family>.j2 when the upstream template has known issues (e.g. template-forced <think> seeding). Naming convention: filename matches the HF repo.

Observability

Prometheus-style metrics are exposed on /metrics (optional — gated behind a build feature):

• spark_requests_total{status, model}
• spark_tokens_generated_total
• spark_ttft_seconds{model}
• spark_decode_throughput_tok_per_sec{model}
• spark_kv_pool_utilization_ratio
• spark_active_sequences

Structured logs go to stderr; they're the primary operational signal. Atlas logs a brief line per completed request (model, prompt tokens, generated tokens, TTFT, elapsed, tools used). Per-token DECODE spam is deliberately not logged — it's useless.

A safe production config (Qwen3.5-35B, agents)

serve Sehyo/Qwen3.5-35B-A3B-NVFP4 \
  --port 8888 \
  --max-seq-len 16384 \
  --max-batch-size 4 \
  --kv-cache-dtype nvfp4 \
  --gpu-memory-utilization 0.88 \
  --scheduling-policy slai \
  --tbt-deadline-ms 100 \
  --enable-prefix-caching \
  --speculative --mtp-quantization nvfp4 \
  --auto-compact 0.85 \
  --adaptive-sampling

Claude Code uses 16k+ context for tool use; running with --max-seq-len 4096 will make agents fail mid-session. Always size up when running an agent workload.

Tool Calling & Streaming

Atlas supports OpenAI-compatible function calling across three wire formats and full SSE streaming (OpenAI + Anthropic conventions). This chapter is the operator reference for running agents against Atlas — how to enable tools, stream responses, handle multi-turn tool results, and recognise the failure modes that used to bite real agents.

Enable tools

Tool calling is on by default — just include tools: [...] in your request. Atlas auto-selects the wire format from the model's MODEL.toml. Overriding: --tool-call-parser <FORMAT>.

All three formats are parsed on the server and emitted to the client as standard OpenAI tool_calls blocks — you do not need to handle the wire format in your client.

Minimal tool-call request

curl -s http://localhost:8888/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "atlas",
    "messages": [{"role": "user", "content": "What is the weather in Paris?"}],
    "tools": [{
      "type": "function",
      "function": {
        "name": "get_weather",
        "description": "Get current weather for a location",
        "parameters": {
          "type": "object",
          "properties": {
            "location": {"type": "string", "description": "City name"}
          },
          "required": ["location"]
        }
      }
    }],
    "max_tokens": 512
  }'

Response (abridged):

{
  "choices": [{
    "message": {
      "role": "assistant",
      "content": null,
      "tool_calls": [{
        "id": "call_00000000",
        "type": "function",
        "function": {
          "name": "get_weather",
          "arguments": "{\"location\":\"Paris\"}"
        }
      }]
    },
    "finish_reason": "tool_calls"
  }]
}

Multi-turn: sending tool results back

Standard OpenAI pattern — append the assistant's tool_calls message and then a role: "tool" message with the tool's output:

{
  "messages": [
    {"role": "user", "content": "What is the weather in Paris?"},
    {"role": "assistant", "content": null, "tool_calls": [{
      "id": "call_00000000", "type": "function",
      "function": {"name": "get_weather", "arguments": "{\"location\":\"Paris\"}"}
    }]},
    {"role": "tool", "tool_call_id": "call_00000000", "name": "get_weather",
     "content": "{\"temperature\": 15, \"condition\": \"cloudy\"}"}
  ],
  "tools": [...]
}

Atlas expands the multi-turn conversation through the chat template and runs a fresh forward.

tool_choice

"required" is implemented via the XGrammar grammar (see XGrammar) — the grammar masks the "no-tool-call" path, so the sampler can only produce a valid tool-call opening.

Streaming

Streaming is enabled with "stream": true:

curl -sN http://localhost:8888/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"atlas","messages":[...],"stream":true}'

Output is standard OpenAI SSE:

data: {"choices":[{"delta":{"role":"assistant","content":"Once "}}]}

data: {"choices":[{"delta":{"content":"upon "}}]}

...

data: {"choices":[{"delta":{},"finish_reason":"stop"}]}

data: [DONE]

Tool calls stream as delta.tool_calls chunks:

data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_00000000","type":"function","function":{"name":"get_weather","arguments":""}}]}}]}

data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\""}}]}}]}

data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"location"}}]}}]}

...

data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}]}

data: [DONE]

Anthropic Messages

For clients built against Anthropic's API:

curl -s http://localhost:8888/v1/messages \
  -H "Content-Type: application/json" \
  -d '{
    "model": "atlas",
    "max_tokens": 500,
    "messages": [{"role": "user", "content": "Hello!"}]
  }'

Streaming uses Anthropic's event conventions — message_start, content_block_start, content_block_delta, content_block_stop, message_delta, message_stop. Atlas populates stop_sequence on message_delta when a stop token was hit (fixed in wave-12 — earlier builds left it null).

Tool use on /v1/messages uses Anthropic's nested content-block format:

{
  "content": [
    {"type": "text", "text": "Let me check that."},
    {"type": "tool_use", "id": "toolu_...", "name": "get_weather",
     "input": {"location": "Paris"}}
  ],
  "stop_reason": "tool_use"
}

Reasoning / <think> blocks

Models that emit <think> (Qwen3.5, Nemotron-H, MiniMax) stream reasoning content as a separate channel:

data: {"choices":[{"delta":{"reasoning":"Let me think step by step. First, ..."}}]}

data: {"choices":[{"delta":{"reasoning":" the user is asking about..."}}]}

data: {"choices":[{"delta":{"content":"The answer is 42."}}]}

data: {"choices":[{"delta":{},"finish_reason":"stop"}]}

This matches OpenAI's o1 family convention. Clients that don't parse reasoning chunks will ignore them cleanly.

--max-thinking-budget caps the total reasoning tokens; --disable-thinking strips them entirely. For agent workloads that want reasoning but don't want unbounded think time, a budget of 2048–4096 is typical.

Vision requests

Qwen3-VL and Qwen3.6 accept images in OpenAI content-parts format:

{
  "model": "atlas",
  "messages": [{
    "role": "user",
    "content": [
      {"type": "text", "text": "Describe this image."},
      {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,<DATA>"}}
    ]
  }],
  "max_tokens": 512
}

image_url accepts data: URLs (base64-encoded) or http(s): URLs (the server fetches them). Multiple images per message are supported.

Known pitfalls and how Atlas addresses them

• Tool-call hallucination inside markdown fences. Older builds' parsers ate the surrounding fence characters. Fixed in wave-1 (markdown-fence aware parser) + XGrammar grammar enforcement.
• Broken tool-call XML (Qwen3-coder format). The parser now tolerates literal </tool_call> inside JSON string values, missing </parameter> tags, and empty {} tool-call bodies (wave-7).
• Streaming Responses store tool_calls. The server now persists tool calls to the Responses-API session store mid-stream so multi-turn conversations across the Responses API see them on the next turn (wave-11).
• Balanced markdown URL parens. The citation extractor used to choke on Wikipedia URLs containing parentheses; now uses a balanced parser (wave-11).
• Template-forced thinking false-positive. Qwen3.6's <think>\n\n</think>\n\n template prologue was triggering the reasoning parser; now requires the opening <think> to be unclosed (wave-4).
• Spontaneous <think> outside the template position. Qwen3.6 occasionally emits <think> in response mid-stream; Atlas now detects this in all four affected code paths (wave-3).

All of these have regression tests under crates/spark-server/src/tool_parser.rs and reasoning_parser.rs.

Running against real agents

Minimum Atlas config for running Claude Code, OpenCode, Cline, or nanobot:

• --max-seq-len 16384 or higher (agents regularly exceed 4k).
• --enable-prefix-caching (massive TTFT win on tool schemas).
• --scheduling-policy slai (keeps streaming smooth).
• --speculative --mtp-quantization nvfp4 if the model supports it (agents are 50%+ tool calls; MTP + constrained decoding = +37% throughput).
• --auto-compact 0.85 so long agent sessions don't crash into the seq-len wall.

Files to read

• crates/spark-server/src/tool_parser.rs — the three parser impls.
• crates/spark-server/src/reasoning_parser/ — <think> detection + extraction.
• crates/spark-server/src/openai/, anthropic/ — request/response structs.
• crates/spark-server/src/api/ — the HTTP handlers.
• docs/ARCHITECTURE.md — system overview covering the tool-call path.
• XGrammar deep dive for the constrained-decoding side.

Multi-GPU & EP=2

Expert Parallelism across two GB10 nodes is the only way to run the largest MoE models (Qwen3.5-122B-A10B, Mistral-Small-4-119B, MiniMax-M2.7) — their experts don't fit on one node. Atlas's multi-GPU support is specifically EP=2 over RoCEv2; the scheduler and HTTP API run on rank 0.

What "EP=2" means here

Two GB10 nodes connected by InfiniBand or RoCE. The model's MoE experts are split 50/50 (128 experts per node for a 256-expert model). Every other layer (attention, SSM, dense FFN, embeddings, LM head) is replicated on both ranks.

Per decode step:

1. Both ranks run the attention / SSM / dense layers on their local portion of the batch.
2. At each MoE layer, the gate runs on both ranks; top-k expert IDs are computed.
3. Tokens destined for remote experts are sent via reduce_scatter.
4. Local experts compute.
5. Results all_gather back.
6. Continue.

Only rank 0 runs the HTTP server and the scheduler. Rank 1 is a silent compute worker that joins at startup via NCCL rendezvous.

Network layer

Atlas's production two-node setup uses InfiniBand RoCE over a Mellanox ConnectX HCA (mlx5_0). The two-node network is dedicated — the public/management interface is separate. Canonical IPs:

• Head: <head-ip>
• Worker: <worker-ip>

If you don't have InfiniBand, EP=2 works over plain Ethernet with a 3–4× throughput penalty. GB10's EP=2 numbers in this book assume RoCE.

Launching — the canonical scripts

scripts/start-ep2.sh handles 99% of deployments. Defaults to Qwen3.5-122B-A10B-NVFP4. Usage:

# Default model
bash scripts/start-ep2.sh

# Explicit model
bash scripts/start-ep2.sh Sehyo/Qwen3.5-122B-A10B-NVFP4

# MiniMax (script auto-strips --speculative since MiniMax doesn't have MTP)
bash scripts/start-ep2.sh lukealonso/MiniMax-M2.7-NVFP4

On each node, the script does:

1. Cleans any stale containers.
2. Sets the NCCL + GLOO env vars (see below) for RoCE.
3. Runs docker run ... atlas-gb10:latest serve <model> --rank {0|1} --world-size 2 --master-addr <head-ip> --master-port 29500 ....

Manual launch

If you need custom flags, the manual flow is:

Head (rank 0, node 0, ):

sudo docker run -d --name atlas-122b-r0 \
  --network host --gpus all --ipc=host \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  -e NCCL_SOCKET_IFNAME=enp1s0f0np0 \
  -e NCCL_IB_HCA=mlx5_0 -e NCCL_IB_DISABLE=0 \
  -e NCCL_NET_GDR_LEVEL=5 -e NCCL_NVLS_ENABLE=0 \
  -e GLOO_SOCKET_IFNAME=enp1s0f0np0 \
  avarok/atlas-gb10:latest \
  serve Sehyo/Qwen3.5-122B-A10B-NVFP4 \
    --port 8888 \
    --rank 0 --world-size 2 \
    --master-addr <head-ip> --master-port 29500 \
    --max-seq-len 4096 \
    --max-batch-size 1 \
    --kv-cache-dtype nvfp4 \
    --gpu-memory-utilization 0.70 \
    --scheduling-policy slai \
    --speculative --mtp-quantization nvfp4

Worker (rank 1, node 1, ):

sudo docker run -d --name atlas-122b-r1 \
  --network host --gpus all --ipc=host \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  -e NCCL_SOCKET_IFNAME=enp1s0f0np0 \
  -e NCCL_IB_HCA=mlx5_0 -e NCCL_IB_DISABLE=0 \
  -e NCCL_NET_GDR_LEVEL=5 -e NCCL_NVLS_ENABLE=0 \
  -e GLOO_SOCKET_IFNAME=enp1s0f0np0 \
  avarok/atlas-gb10:latest \
  serve Sehyo/Qwen3.5-122B-A10B-NVFP4 \
    --port 8889 \
    --rank 1 --world-size 2 \
    --master-addr <head-ip> --master-port 29500 \
    --max-seq-len 4096 \
    --max-batch-size 1 \
    --kv-cache-dtype nvfp4 \
    --gpu-memory-utilization 0.70 \
    --scheduling-policy slai \
    --speculative --mtp-quantization nvfp4

Start both; they rendezvous on <head-ip>:29500. Only the head serves HTTP.

The critical MTP-flag symmetry rule

The worker must be started with the same --speculative --mtp-quantization <value> --num-drafts <N> flags as the head. Otherwise the head's MTP verify command arrives at the worker's SSM layer with no intermediate buffers allocated, and the step crashes with an SSM intermediate-buffer error.

start-ep2.sh enforces this. Manual launches where only the head has MTP enabled have bit multiple contributors; the in-code EP=2 MTP guard (wave-5) now produces a clearer error, but the rule is: head and worker flags must match for speculative, num-drafts, and mtp-quantization.

NCCL env vars — what matters on GB10

These are baked into the scripts/start-*ep2.sh scripts. If you bring up EP=2 from scratch, carry them forward — a missing NCCL_NVLS_ENABLE=0 will crash silently mid-warmup.

Testing

scripts/test-minimax-ep2.sh is the canonical end-to-end harness. It runs:

• Coherence check (3 prompts, 3 tokens each, verifies byte-identical output across 3 repeats)
• Fibonacci generation (tests long-output coherence)
• Tool calls (tests the EP=2 + tool-call + MTP interaction)
• TPS benchmark (tests decode throughput)

Last known-green run on alpha-2.35: MiniMax-M2.7 EP=2 scored 8/10 on the suite; the Qwen3.5-122B EP=2 equivalent hits ~46 tok/s on 600-token decodes with the flags shown above.

Performance reference

Troubleshooting

• Rendezvous timeout — master-addr unreachable from worker, or master-port blocked by firewall.
• NCCL stuck at "all ranks ready" — NCCL_SOCKET_IFNAME wrong on one side. Both ranks need the same ifname (both enp1s0f0np0).
• Silent NaN after first request — NCCL_NVLS_ENABLE=1 leaked through. Force off.
• SSM intermediate buffer error on worker — MTP flags mismatched between head and worker. See the symmetry rule above.
• Health endpoint on head returns 200 but requests hang — worker failed to come up and the head's initial barrier is blocking. Check the worker's logs.

Why not tensor parallel?

Tensor parallelism (TP) was the traditional approach for splitting dense models across GPUs. Atlas on GB10 does not use TP because:

1. GB10 is unified memory — there's no NVLink island to exploit for intra-node TP.
2. The models Atlas targets are MoE-dominant beyond a single GB10's memory. EP is the natural split — one expert per rank is cleaner than splitting a weight matrix.
3. TP requires per-layer all-reduces. EP requires per-MoE-layer token dispatch. At the model shapes we care about, EP is lower collective traffic than TP.

TP for dense models on future multi-GPU hardware is on the roadmap; today there's no supported use case where it would win.

Files to read

• scripts/start-ep2.sh, scripts/start-minimax-ep2.sh, scripts/test-minimax-ep2.sh
• crates/spark-comm/src/nccl_backend.rs — NCCL impl.
• crates/spark-model/src/layers/moe/forward_ep.rs (EP=2 dispatch); MiniMax loader in crates/spark-model/src/weight_loader/minimax.rs.
• kernels/gb10/minimax-m2-229b/nvfp4/moe_w4a16_grouped_gemm.cu — routed grouped-GEMM kernel.
• docs/adr/0007-tp-ep-composition.md — TP/EP composition design record.
• docs/GB10_DEPLOYMENT_GUIDE.md §7 — field-tested EP=2 troubleshooting.

Benchmarking

Atlas's performance claims are measurable. This chapter shows what the benchmarks measure, how to reproduce them, and how to read the numbers.

The headline numbers

From the repo README, distilled:

And the kernel micro-benchmark summary: Atlas wins 32/32 against PyTorch on attention, GEMM, SSM, RMSNorm, RoPE, SiLU×Mul, and conv1d, with speedups from 1.04× up to 18.2×.

Two kinds of benchmark

Atlas has two benchmark surfaces:

1. End-to-end HTTP throughput — atlas-spark-bench (client-side Criterion harness targeting a running server). This is what "131 tok/s" means.
2. Per-kernel micro-benchmarks — Criterion benches in each primitive crate, run with cargo bench. This is where "4.95× prefill attention" comes from.

Different things; both are meaningful. The E2E number is what an operator sees. The per-kernel number is what tells the kernel engineer where effort is paying back.

Running end-to-end benchmarks

Start a server:

sudo docker run -d --name atlas-35b \
  --network host --gpus all --ipc=host \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  avarok/atlas-gb10:latest \
  serve Sehyo/Qwen3.5-35B-A3B-NVFP4 \
    --max-seq-len 8192 --kv-cache-dtype nvfp4 \
    --scheduling-policy slai \
    --speculative --mtp-quantization nvfp4

Wait for listening. Then:

export ATLAS_BENCH_URL=http://localhost:8888
cd /path/to/atlas
cargo bench -p atlas-spark-bench

Criterion saves results to target/criterion/. The stable JSON snapshots that the README quotes are pinned under bench/.

The scripts/sweep_all_models.sh helper boots each model in turn, runs the canonical short-prompt bench, and writes the README.md throughput table. That's how the table in the README gets regenerated.

Running per-kernel benchmarks

cargo bench -p spark-runtime        # KV cache ops, sampler micro
cargo bench -p atlas-spark-bench    # end-to-end client benchmarks

Criterion-driven, from each crate's benches/*.rs. Reference shapes come from Qwen3-Next-80B (hidden=2048, 16 Q-heads, 2 KV-heads, head_dim=256, intermediate=512, num_experts=256, topk=10).

The full kernel numbers table:

32/32 wins. Peak achieved memory bandwidth in this set: 599 GB/s (2.2× the 273 GB/s LPDDR5X spec — that's the L2 cache effect for small SiLU×Mul inputs).

Concurrency sweep

scripts/run_conc_benchmark.sh drives N parallel streams against one server. Reveals the scheduler + KV allocator under load. Typical pattern on Qwen3.5-35B:

Aggregate throughput grows super-linearly up to the batch size cap (where graph amortisation kicks in) and then super-linearly until KV pool pressure; after that latency degrades more than throughput improves. Sweet spot on 35B: concurrency 4–8.

TTFT and prefix-cache behaviour

Run a request, note the TTFT. Run the same request again — with --enable-prefix-caching, TTFT drops to ~40ms (prefix cache hit). Agent workloads observe this as the difference between "the first response was slow" and "everything after is fast". The bench harness's prefix-warmup suite measures cold vs warm TTFT.

Where raw results live

• Pinned snapshots (tracked): result files under bench/. These feed the book and the README.
• Ephemeral Criterion runs (gitignored): target/criterion/.
• Historical benchmark journeys: docs/ATLAS_SPARK_JOURNEY.md — the benchmark retrospective across the Spark line.

Apples-to-apples notes

When comparing Atlas to vLLM or TensorRT-LLM:

• Same hardware. GB10 SM121 numbers do not transfer to H100 / B200 / MI300X.
• Same model. "Qwen3.5-35B-A3B at 36 tok/s" is vLLM's NVIDIA GB10 benchmark on the NVFP4 CUTLASS MoE path, same HF checkpoint.
• Same prompt shape. The 131 tok/s number is on a short prompt ("What is the capital of France?", max_tokens ≤ 30). Longer prompts show slightly different numbers because prefill cost amortizes differently.
• Same precision. Atlas NVFP4 vs vLLM NVFP4; Atlas FP8 vs vLLM FP8. Never compare across quant schemes.

The headline "3.6× faster than NVIDIA's 36 tok/s" is apples-to-apples against NVIDIA's own vLLM numbers on the same (GB10, Qwen3.5-35B-A3B, NVFP4) target.

Files to read

From the CLI

The benchmark suite the dashboard runs is also a subcommand, so a benchmark can be scripted, run in CI, or driven over SSH with no terminal attached.

spark benchmark list                      # the suite
spark benchmark list concurrency-sweep    # one benchmark's parameters
spark benchmark run  concurrency-sweep --model <served-model>
spark benchmark history

run drives an endpoint that is already serving — it neither loads a model nor touches the GPU. The one exception is --pull-request-gate, which does start a server: it serves the benchmark's own recipe on a free port (900 s boot timeout, for a cold NVFP4 load) and tears it down on drop (cli/bench_selfstart.rs).

A benchmark can be defined on more than one model variant — one BENCH.toml entry per checkpoint, each carrying its own serve recipe and its own thresholds (spark benchmark list <id> prints them, and the TUI shows them as a step after selecting the benchmark). A gate run serves the variant marked default = true unless you name another:

spark benchmark run agentic-webserver --yes --pull-request-gate   --checkpoint unsloth/Qwen3.8-27B-NVFP4

Records are keyed by variant too — a non-default variant's record gets the checkpoint slug in its filename, and only the default variant's records can discharge the required gate. Numbers never compare across variants: the dense 27B's wall band is roughly 2× the 35B MoE's, which is exactly why the thresholds live per checkpoint.

Point run somewhere else with --url:

spark benchmark run concurrency-sweep \
  --url http://10.10.10.3:8888 --model Qwen/Qwen3.6-35B-A3B-FP8 \
  --param concurrencies=1,2,4 --param isls=128 --param osl=64

--param takes any key from spark benchmark list <id>; anything you leave out takes the schema default. An unknown key is an error listing the valid ones, because a silently-ignored override produces a run measuring something other than what you asked for.

Exit codes

1 and 2 are distinct so a script can tell "the harness broke" from "the model missed the bar". --no-fail-on-verdict collapses 2 into 0 when you are collecting numbers rather than gating on them.

Run history

Every run — from the CLI or the dashboard — is recorded under ~/.atlas/runs/<benchmark-id>/, carrying the result, every parameter used (not just the ones you overrode), the target, the source, and the Atlas version. So a stored run says what it measured and can be reproduced.

spark benchmark history --id concurrency-sweep --limit 5
spark benchmark history --run run-1785000000123456789 --format json | jq .params

One store, both directions: a CLI run appears in the dashboard's Benchmarks → History pane, and a dashboard run appears in spark benchmark history marked tui.

Machine-readable output goes to stdout, progress to stderr, so --format json > run.json is a clean file. ATLAS_HOME relocates the store.

• crates/atlas-spark-bench/src/lib.rs — E2E harness.
• Each primitive crate's benches/*.rs — per-kernel micro.
• bench/*.json — pinned result snapshots.
• scripts/sweep_all_models.sh, scripts/run_conc_benchmark.sh — automation.
• docs/ATLAS_SPARK_JOURNEY.md — benchmark journey and retrospective.
• README "Benchmark Results" section — the authoritative long-form table.

Contributing

The canonical references are CONTRIBUTING.md and AGENTS.md. This chapter gives a working overview for anyone reading the book first.

The AI-first policy

Atlas is explicitly an AI-first codebase. From CONTRIBUTING.md:

• All PRs are expected to be AI-generated. Use the best AI tools available to write your kernels, Rust code, and benchmarks.
• Human-written code must be justified. Indicate which parts are human-authored and explain why.
• Human-only contributions will be reviewed by AI.

This is not branding — it's the operational consequence of the specialization thesis. If AI can hyperoptimize CUDA kernels for specific hardware targets, it can write the infrastructure too. Ports to new (H, M_q) targets are the clearest example: each is a bounded, well-scoped piece of work, and that's the unit AI-assisted engineering handles best.

What kinds of PRs are welcome

The README's Contributing section lists four categories:

• New (H, M_q) targets. Porting Atlas kernels to new hardware (H100, B200, MI300X, Apple M4, Intel) or new models. Each target is a self-contained body of work. See the Adding a new hardware target and Adding a new model guides.
• Kernel optimization. Profile existing kernels, experiment with tiling strategies, register pressure, shared-memory layouts. If you can beat the numbers in the Benchmarks chapter, send the PR.
• Benchmark coverage. Add shapes and configurations not yet tested. More data points sharpen the hypercompiler.
• Bug reports. Include hardware details, repro steps, and kernel timings.

Local checks before a PR

These are what CI runs (.github/workflows/ci.yml). Run them locally first:

# 1. Formatting
cargo fmt --all -- --check

# 2. Lints. BOTH env vars are needed: ATLAS_SKIP_BUILD stubs the PTX build,
#    CUDARC_CUDA_VERSION stops cudarc shelling out to `nvcc --version`.
#    Deny-warnings comes from [workspace.lints], so CI passes no -D flag
#    and no --all-features. This is verbatim what ci.yml runs.
ATLAS_SKIP_BUILD=1 CUDARC_CUDA_VERSION=13000 cargo clippy --workspace --tests

# 3. License headers
bash scripts/check-license-headers.sh

# 4. Typos
typos     # install once: cargo install typos-cli

All four are required to pass. Real CUDA build + test cycles require a GB10 host — not the laptop, the DGX Spark itself.

Ground rules (from AGENTS.md)

• SPDX header on every source file. // SPDX-License-Identifier: AGPL-3.0-only on line 1 of every .rs, .cu, .cuh, .h, .hpp, .cpp. Enforced by the license-headers CI job.
• License is AGPL-3.0-only. Don't mix in permissive-only code without confirming compatibility. deny.toml controls allowed dependency licenses.
• Don't regress supported models. The matrix in Supported Models is the contract; docs/GB10_DEPLOYMENT_GUIDE.md §2 is its SSOT, and kernels/gb10/ carries 22 (model, quant) leaves. If your PR might touch a hot path, validate against tests/run_all_models.py on a GB10 before opening.
• One logical change per commit. Don't bundle cleanup with a bug fix.
• Commit message format. <area>: <imperative summary> — e.g. spark-server: preserve template-forced thinking through EP=2.

Failure modes that cost the project time

These are the classes of bug that have burned days. Know them; avoid introducing them.

• Protocol drift between OpenAI (api/) and Anthropic (anthropic/) surfaces. A fix on one side often needs a matching change on the other.
• Template mismatches subtly breaking tool-calling — different <tool_call> vs <minimax:tool_call> tokens, <think> seeded by the template vs emitted by the model, thinking budget enforcement.
• FP8 / KV / quantization edge cases — BF16 paged cache routed into an FP8 kernel → silent NaN. If your change touches numeric paths, verify with a real model before claiming success.
• Docs drift — CLI flags, release commands, quick-start snippets. Verify against the current binary, not memory.

The cardinal rule

Never assume the model is at fault. Always look for the Atlas bug first.

The test matrix has caught many issues that would have looked like "model hallucination" in a lesser codebase. The heuristic is: if the model used to produce coherent output on this input and now doesn't, there's an Atlas bug, not a model bug.

The CLA

By contributing, you agree to the Contributor License Agreement. Your work goes out under AGPL-3.0 in the Community Edition, and you grant Avarok the right to relicense for the Enterprise Edition.

The CLA Assistant bot automatically comments on every PR. You must explicitly acknowledge and sign before merge.

Adding a new hardware target

High-level (full walkthrough in the repo README):

1. kernels/<hw>/HARDWARE.toml with vendor = "...".
2. impl ComputeTarget in atlas-core/src/compute.rs (or inline in your crate).
3. Arm in atlas-kernels/build.rs — resolve_targets() reads ATLAS_TARGET_HW (default gb10) and the leaf HARDWARE.toml's vendor picks the ComputeTarget.
4. impl GpuBackend in spark-runtime/src/<vendor>_backend.rs — 27 methods, some optional.
5. Kernel sources under kernels/<hw>/common/ (the GB10 baseline is 160 .cu files / 318 __global__ entry points), plus per-model shadows only where a target diverges.
6. MODEL.toml + KERNEL.toml for at least one model.
7. Backend selection branch in spark-server/src/main.rs.
8. Dockerfile for the new hardware.

Adding a new model

The model-specific surface is tiny:

1. crates/spark-model/src/weight_loader/<your_model>.rs implementing ModelWeightLoader (~200–500 lines depending on architecture complexity).
2. Module declaration + pub use in crates/spark-model/src/weight_loader/mod.rs.
3. One match arm in crates/spark-model/src/factory.rs::loader_for_config.
4. Optional: kernels/<hw>/<your-model>/MODEL.toml for sampling / behavior defaults.
5. Optional: tool-call parser under crates/spark-server/src/tool_parser/.
6. Entry in tests/run_all_models.py for regression coverage.
7. Entry in Supported Models.

Existing loaders for patterns: qwen35.rs, minimax.rs, nemotron.rs cover dense, SSM+MoE hybrid, and attention+MoE shapes respectively.

PR process

1. Fork and create a feature branch.
2. Atomic commits. Enforced by reviewers; squash only at the reviewer's request.
3. CI must pass: ci.yml runs fmt, clippy, license-headers, typos, kernel-structure, cargo test --workspace, test-macos-metal and release-matrix; security.yml runs cargo-deny; file-size-cap.yml the 500-LoC cap; docs.yml mdBook + cargo doc. The pr-benchmark-gate job is advisory (continue-on-error).
4. PR template asks for:
   • What — summary of the change.
   • Why — motivation and context.
   • Benchmarks — before/after numbers for perf-related changes.
   • Authorship — AI / human / mixed; justify human-written sections.
5. Sign the CLA when the bot asks.
6. A maintainer (and/or AI reviewer) merges.

Scope escalation

If a task is ambiguous, ask in the issue/PR before implementing. If scope grows past "one PR", split it. If you're modifying a shared trait, a build script, or CI config, flag it in the PR description so reviewers catch it.

References

• CONTRIBUTING.md — canonical.
• AGENTS.md — practical contributor guide.
• CLA.md — the CLA text.
• SECURITY.md — disclosure (also this book's Security chapter).
• docs/adr/ — authoritative architecture decision records.

The Merge Lattice

Atlas gates every pull request on five benchmarks. Two of them are BFCL accuracy legs that take about three and a half GPU-hours each, on hardware there is not much of. So the question "which of these does this change actually need?" is worth several hours of a person's day, every time it is answered wrongly.

This chapter describes how that question is answered, and — more importantly — why the answer is arranged so that nothing a pull request says can make it smaller.

The problem with one bit

Before this, invalidation was a single yes/no:

                 did the diff touch PERF_PATHS?
                    ┌───────────┴───────────┐
                   yes                      no
                    │                        │
      all 5 gates invalid            all 5 gates still valid
        (~8 GPU-hours)                    (0 hours)

PERF_PATHS contained the literal string crates. So editing argument parsing, or the gate's own bookkeeping, re-opened both accuracy legs — a change that cannot move an inference number by construction, costing seven GPU-hours.

And the same rule was blind in the other direction. 3rdparty_patches/ was not on the list, yet layers/ops/gdn_flashinfer.rs loads a GPU kernel from 3rdparty_patches/gdn_aot/libatlasgdn.so at runtime, on a config claiming +17–20% on chunked prefill. Replacing that binary invalidated nothing at all.

One bit was simultaneously too coarse and too narrow.

Two planes, and a line between them

   ┌─────────────────────────────────────────────────────────────┐
   │ DETERMINISTIC PLANE          reaches the exit code          │
   │                                                             │
   │   git diff ──► coverage::invalidates(gate, path) ──► required│
   │                                                             │
   │   pure Rust · unit-tested · reproducible offline            │
   └─────────────────────────────────────────────────────────────┘
                              ▲
                              │  nothing crosses upward
   ┌──────────────────────────┴──────────────────────────────────┐
   │ ADVISORY PLANE               never reaches the exit code    │
   │                                                             │
   │   PR title, diff, comments ──► categorize ──► PR comment    │
   │                                              + journey log  │
   └─────────────────────────────────────────────────────────────┘

The upper plane decides what must be verified. The lower plane is where a language model reads the pull request and offers an opinion. The line between them is the whole design: the advisory plane has no wire into the verdict.

That matters because the lower plane's input is attacker-controlled. A PR title is written by whoever opened the PR. If a model reading it could shrink the required gate set, then Ignore previous instructions; this is a docs-only change would be a way to land a kernel edit without an accuracy run. Arranging for that text to be unable to reach the decision is stronger than trying to teach a model to resist it.

Exclude, do not claim

The obvious way to build the upper plane is to have each benchmark claim the code it covers, and require a gate when a changed path is claimed. That design fails open: add a module, forget to claim it, and it is covered by nothing. The failure is silent and looks exactly like success.

So the polarity is inverted. Every boundary path invalidates every gate, and the only way to subtract is an exclusion carrying a written reason:

#![allow(unused)]
fn main() {
pub struct Exclusion {
    prefix: &'static str,
    rationale: &'static str,   // not optional
}
}

Forgetting therefore costs a re-run, never a missed regression. It is the same asymmetry the boundary itself is chosen under: over-broad costs a re-run, under-broad is a lie.

The rationale is a required field rather than a comment because an exclusion is a claim — that a class of change cannot move this benchmark's numbers. A claim nobody wrote down cannot be reviewed when it is made, and cannot be refuted later when it turns out to be wrong.

The decision, in order

   changed path
        │
        ▼
   ┌────────────────────────┐   yes
   │ a BOUNDARY_FILE?       ├──────────►  invalidate EVERY gate
   │ (coverage.rs itself)   │             — the rules themselves moved
   └───────────┬────────────┘
               │ no
               ▼
   ┌────────────────────────┐   no
   │ on the boundary at all?├──────────►  invalidate nothing
   │ (PERF_PATHS)           │             — docs, scripts, harness
   └───────────┬────────────┘
               │ yes
               ▼
   ┌────────────────────────┐   yes
   │ matches an Exclusion   ├──────────►  this gate stays valid
   │ for THIS gate?         │             — and the file says why
   └───────────┬────────────┘
               │ no
               ▼
        invalidate this gate      ◄── the default, and the safety property

Step three's default is what makes the whole thing safe. A path nobody has classified invalidates, so an unclassified new subsystem over-tests rather than escaping.

The map guards itself

An exclusion table that could exempt the file it lives in would be a lock whose key is kept inside it. A pull request could add "exclude everything", and that very edit would trigger no gate to catch it.

Hence the first question in the diagram above. Any change to coverage.rs invalidates all five gates, and it is checked before exclusions are consulted, so a blanket exclusion cannot reach it. A test writes the attack out explicitly — a gate excluding all of crates — and asserts the boundary file still invalidates.

Component-wise matching

   "crates"  vs  "crates2/src/lib.rs"        →  NOT under        (starts_with says yes)
   "Cargo.toml" vs "Cargo.toml.orig"         →  NOT under        (starts_with says yes)
   "crates"  vs  "crates/spark-model/x.rs"   →  under
   "crates"  vs  "crates"                    →  under

A naive prefix test matches the first two. That would invalidate gates for unrelated files, which teaches people the gate is noise, which ends with someone turning it off. So matching is p == entry || p.starts_with(entry + "/"), and a test runs a battery of lookalike names through it.

Why it is a lattice

The required set is ordered by inclusion, and the only operation that builds it is union:

                    {all five gates}          ⊤  — unclassified paths land here
                     /      |      \
              {bfcl×2}   {ttft×2}  {agentic}
                     \      |      /
                        {  }                  ⊥  — docs-only changes

Gates join upward and never meet downward. invalidated_by contains no branch that removes an element from its result, and a test asserts the consequence directly: adding a changed file never removes a required gate, over both benign and adversarial inputs.

This is the same shape as a security lattice in the information-flow sense, and it buys the same thing: monotonicity means you can reason about the worst case without enumerating the cases. Whatever a pull request contains, the answer is at least the floor.

What it costs and what it buys

The last two rows are the ones that matter most. One is the saving; the other is a hole that was open the entire time the gate has existed.

It cannot excuse itself

The pull request that introduced this machinery touches kernels/gb10/common/paged_decode_attn_fp8.cu and layers/ops/fp8_moe.rs. The floor therefore demands all five gates of it, and a test pins exactly that file list so the property cannot quietly lapse.

A governance system whose first act is to exempt itself is not a governance system. This one owed — and paid — the full bill.

When a gate is open, the message says why

NONE  bfcl-subset — latest record is for fe99349724 (2026-08-08-fe99349724.json)
      — invalidated by crates/atlas-kernels/tests/kernel_arity.rs,
        crates/spark-model/src/layers/mtp_head.rs, … and 16 more

Reporting only that a gate is open turns a twenty-second fix into a bisect. The check knows which files re-opened it, so it says so.

Auditing the rules

The exclusions are claims, and claims rot. Tests check that every exclusion names a path that exists (a rule matching nothing is either a rename that was missed or a mistake), that every one lies on the boundary (a rule with no effect that a reader would assume has one), that every registered benchmark is either gated or explicitly excused with a reason, and that the benchmark drivers do not import each other — the precondition the per-driver exclusions rest on.

That last one is the interesting case: TTFT excludes the BFCL driver on the grounds that one cannot affect the other. If somebody later makes BFCL import from TTFT, that reasoning silently becomes false. The test turns it into a compile-visible event instead.

Below the path floor: what a target actually compiles

The floor above answers at the granularity of a path list, and for kernels/ that is very coarse. kernels/gb10/common/ holds 160 shared kernels; each model directory shadows only 5–18 of them, and nothing shadows paged_decode_attn_fp8.cu at all. Under the path rule, editing one shared kernel re-opens every gate for all 28 targets. At roughly three and a half GPU-hours per accuracy leg, that is a cost people route around, and a gate people route around is worse than a slower one.

So a second rung sits on top of the floor. It can only ever narrow, never widen, and only for paths inside kernels/:

 changed paths
      │
      ├─ any path outside kernels/ ─────────────► every gate re-opens (unchanged)
      │
      └─ all paths inside kernels/
             │
             └─ for each target those paths can reach:
                    closure hash now == closure hash when measured?
                       ├─ yes for every one ─────► the record still stands
                       └─ no for any one ───────► that gate re-opens, and the
                                                  message names which targets

Why a file set is not enough

The tempting version hashes each target's resolved .cu set after shadowing: if the set is unchanged, the record still covers. It is wrong twice, and both were found by reading the tree rather than reasoning about it.

A shadow file may #include the very file it shadows. kernels/gb10/qwen3.6-27b/nvfp4/inferspark_prefill_paged_indirect.cu contains #include "../../common/inferspark_prefill_paged_indirect.cu", and eight files do this. A set hash reports "this model shadows that stem, so the common copy cannot reach it" — while the edited bytes are compiled straight into the model's kernel. Silent, and it fails open, on exactly the change class the scheme exists to scope.

Headers are in no set at all. The resolver matches *.cu non-recursively, so the nine common/*.cuh files — including the one carrying BR64 — are invisible. Editing a header would invalidate nothing.

Following includes dissolves both, because an included file's bytes are inside the hash wherever it lives.

Two-sided, or it proves nothing

The hash is baked into the binary by build.rs, at the moment the kernels are compiled, and copied from there into the record. It is deliberately not recomputed from the working tree when the record is written: the tree and the binary differ precisely when it matters — a stale target/, a dirty tree, an image carried between boxes — and a tree-side attestation would paper over all three while looking correct.

Verification then recomputes from the tree using the record's own stored arch, compiler and flags, so the only thing that can move the hash is a source change. Substituting the checker's environment would let whichever machine ran CI invalidate every record.

Two implementations of "what are this target's sources" now exist — build.rs uses collect_cu_files, the gate uses taxon::sources — and if they ever drift, the hashes never match, every record stays invalidated, and it looks exactly like "the kernels changed". spark-server/tests/closure_attestation.rs is the only place they are compared; it recomputes every baked hash from the tree and prints the count it checked, because "3 passed" reads identically at 21 targets and at 22.

What it does not cover

Angle-bracket includes (covered coarsely by the recorded compiler version); headers reached through an -I search path; #if/#ifdef, which are not evaluated, so an include in an untaken branch is walked anyway — over-including, which costs re-runs rather than soundness. Host code stays outside entirely. Equal hash proves equal device code, not equal outcome under load, which is why bitwise output gating remains valid only at C=1.

Thresholds live beside the model

kernels/<hw>/<model>/BENCH.toml, sibling of MODEL.toml. One file per model, [[benchmarks]] entries keyed first by quant and naming their gate, so hardware and model are implied by the path and cannot disagree with the contents.

Thresholds are per checkpoint, not per model — two checkpoints of one model differ by several BFCL points and cannot share a bar.

Three rules the schema enforces, each a way a threshold file can lie:

• status = "unmeasured" entries carry no metrics table. Absence is the TODO. A guessed number a run can clear is worse than no number, because it reports PASS for something nobody measured.
• measured entries must carry metrics, so the status cannot overstate.
• Exactly one checkpoint per (gate, hardware) sets default = true. There is no "the only entry wins" — a second checkpoint added later would silently move which one the gate scores.

BENCH.toml is under kernels/, a boundary path, so it is exempted by exact filename. Without that, raising a bar would invalidate every record — including the run that proved the new bar reachable. The exemption is safe only because nothing compiles the file, and it is checked after the boundary-file rule, so it can never exempt the rules that grant it.

The telemetry plane

Everything above judges one PR. Nothing in it can answer are these green together: two PRs touching one kernel target are each measured against a baseline the other will move, so whichever lands second is gated on a number that no longer describes the tree. Both were genuinely green when measured, which is why a merge queue cannot see it.

A scheduled workflow renders one comment, rewritten in place, carrying the per-PR blast radius, the collisions, a suggested order, CODEOWNERS mentions, and every target in the tree — including untouched ones, because listing only the affected ones would convert ungated into unaffected by omission.

It is advisory and fails nothing. The blocking decisions stay with the committed records. The judgement lives in gate::telemetry as a pure function of the PR facts plus the tree, so which targets, which order and who to mention are all unit-testable with no network and no fixture repository; the workflow only fetches and posts.

Security Policy

Canonical: SECURITY.md. This chapter summarises the policy and the threat model.

Reporting a vulnerability

Do not open a public issue for security vulnerabilities.

Email security@avarok.net with:

1. Description — what the vulnerability is and its potential impact.
2. Reproduction steps — minimal.
3. Environment — OS, CUDA version, GPU model, Rust version.
4. Affected component — which crate or kernel.

Receipt acknowledged within 48 hours. Initial assessment within 7 days.

Supported versions

Atlas moves fast; there are no LTS branches. Run from main or a recently-tagged release.

Threat model

Atlas is an inference server that runs locally with GPU access. The primary surfaces:

1. CUDA kernel safety

• Out-of-bounds reads/writes in kernels.
• Integer overflow in kernel grid/block parameter computation.
• Buffer overflows in shared-memory layouts.

Automated: nothing — there is no static analyser on the CUDA sources. Human: kernel reviews require the PR author to document tile shapes and memory accesses.

2. HTTP API input

• Malformed JSON — axum + serde handles schema validation; unknown fields are rejected by default.
• Oversized request bodies — ATLAS_MAX_BODY_BYTES caps inbound body size. The default is 32 MiB, not 8 (main_modules/serve_router.rs); size your reverse proxy against 32.
• Prompt injection via the chat template — the model is the primary defense; Atlas does not attempt content-level filtering.
• Rate-limit exhaustion — per-key token bucket with a MAX_KEYS DoS guard against cardinality explosion.

3. Weight loading

• Malicious safetensor files — the safetensors crate handles format parsing; Atlas validates shapes against ModelConfig before any GPU upload.
• Path traversal during model load — paths are resolved through PathBuf::canonicalize and checked against the configured cache root.
• Disk exhaustion — model downloads from HF can be many GB; operators should size the HF_HUB_CACHE volume accordingly.

4. Unsafe Rust

Atlas uses unsafe blocks for:

• CUDA FFI via cudarc (driver calls, raw pointer arithmetic).
• NCCL FFI via the vendored nccl_sys bindings.
• MaybeUninit scratch buffers in a handful of hot paths.

Every unsafe block is annotated with the safety invariant it relies on. Reviewers treat unsafe introductions as high-priority and typically block the PR until the invariant is written down.

5. Dependency supply chain

• cargo deny audits dependencies for known advisories (RustSec), license compliance (AGPL-compatible only), and banned crates. Runs on every PR and weekly via cron.
• deny.toml controls allow/deny lists. Permissive licenses (MIT, Apache-2.0, BSD-3) are allowed; GPL variants incompatible with AGPL-3.0 are denied.

Automated security in CI

This table previously listed a cppcheck CUDA static-analysis row. No such job has ever existed; it was removed rather than left as an advertised control nobody runs.

The -D correctness -D suspicious gate is deliberate: stylistic clippy lints churn across toolchain releases and are not worth blocking PRs, but the correctness + suspicious categories map to real bugs and always block.

Disclosure policy

Coordinated disclosure. On a valid report:

1. Fix lands in main.
2. New tagged release.
3. Credit to the reporter unless anonymity is requested.

Out of scope

Some things are not a security concern under this policy — they're bugs, but not security bugs:

• Slow kernels. Performance regressions go through the normal PR/bench workflow.
• Model hallucinations. The model is not Atlas. SECURITY.md does not cover what the model chooses to say.
• Operator misconfiguration. --gpu-memory-utilization 1.0 will OOM; that's not a vulnerability.

If you found something

Email security@avarok.net. Include what you need, keep the repro minimal, and do not exploit the vulnerability against production deployments you do not own. The team has fixed every credibly-reported issue within the 7-day initial-assessment window; known-good practice gets a prompt response.

Release Notes

Atlas's release notes are per-version markdown files in docs/releases/. This chapter links to them and summarises the big themes across recent alphas. For the latest release, check the repo — this page is a stable pointer, not a ticker.

Release naming

alpha-<major>.<minor><letter> — e.g. alpha-2.43, alpha-2.44, alpha-2.14c. Minor versions bump on any meaningful feature or fix; letters (a, b, c) are patch-level iterations on the same minor.

Since Atlas is pre-1.0 and under aggressive development, semantic versioning does not apply. Any release can break API or CLI compatibility — the per-release notes document what.

Where to read them

The multi-model Docker image always tracks the latest alpha at avarok/atlas-gb10:latest. Specific versions are tagged as avarok/atlas-gb10:alpha-2.44 etc.

Recent themes

Rather than duplicate every release note, here's the shape of recent work. Each theme maps to architecture decision records under docs/adr/ and the benchmark history in docs/ATLAS_SPARK_JOURNEY.md.

alpha-2.0 → alpha-2.20 — coherence and the model matrix

The long effort to get the full 12-model matrix to pass an end-to-end coherence suite. Highlights:

• Fast safetensors loader (InstantTensor-style, O_DIRECT + pipelined).
• MTP speculative decoding landed for Qwen3.5-35B, Qwen3-Next-80B, Qwen3.5-122B.
• RadixAttention prefix caching + Marconi SSM snapshot caching.
• Qwen3-VL vision tower integration (ViT block + merger layer + MRoPE image position IDs).
• Nemotron-H Mamba-2 integration.
• Chunked SSM prefill (saves 7–9 GB for long-context prefill).

alpha-2.20 → alpha-2.35 — MiniMax and the 256-expert problem

Getting MiniMax-M2 and M2.7 to pass:

• 256-expert sigmoid MoE routing (distinct from softmax-topk).
• EP=2 token dispatch kernel for >256 experts.
• rms_norm placement fix in the MiniMax MoE path.
• norm_topk_prob semantics (sum-normalised, not softmax).
• FP8-free enforcement on NVFP4 shared-expert path.
• Template-forced thinking detection that distinguishes MiniMax-style <think> seeding from Qwen-style.

alpha-2.35 was the first release where M2.7-NVFP4 EP=2 passed the full coherence + tool-call + TPS suite.

alpha-2.35 → alpha-2.44 — OSS prep + bug sweeps

Thirteen waves of systematic audit-framework bug sweeps (project_bug_sweep_wave1_2026_04_22.md through wave13). Net effect:

• Wave 1 — attention/SSM sibling stride bugs (K≠2 paths), slot-keyed CUDA graph caches, compact_sequence pointer leak.
• Wave 5 — Responses flat-form tools, vision prefix-cache contamination skip, EP=2 MTP guard.
• Wave 6 — NVFP4 MTP loader force-BF16 when ignore_modules lists mtp.*, FP8 prefill shared-experts allreduce reorder.
• Wave 7 — SSM dummy slot defensive fix, qwen3 parser literal-</tool_call> + missing-</parameter> recovery.
• Wave 8 — Sampler temperature=0 && rep_penalty=0 div-by-zero guards, longest-first stop-sequence matching.
• Wave 9 — Rate-limiter MAX_KEYS DoS guard, body-size env-configurable.
• Wave 10 — Responses function_call(_output) items + instructions stacking, multi-block reasoning extractor, MoE topk bounds, weight loader scale=0 guard.
• Wave 11 — Streaming Responses store tool_calls, balanced markdown URL parens (Wikipedia URLs), self-spec rollback fail-fast on SSM.
• Wave 12 — Anthropic streaming stop_sequence populated, /tokenize+/detokenize gated behind auth (then ATLAS_REQUIRE_AUTH; the env var has since been replaced by the --require-auth flag).

Vision fixes (7 ViT + MRoPE image position IDs) landed all four vision models passing the Mona-Lisa test.

Pass-N passes

Alongside the bug sweeps, "Pass-N" work is the systematic model-matrix regression suite. Each Pass runs the full 171-test suite across all models and all flag combinations. Milestones:

• Pass-14 — 152/171 (88.9%), 10/19 perfect. 14 fixes delivered.
• Pass-16 — 233/247 (94.3%), 15/19 perfect. Fixed the 80B-MTP seq_len += k-1 off-by-one bootstrap.
• Pass-21 — 233/247, 14/19. Tool-parser +4; 122B-nvfp4 LC regressed from minimax-m2 branch commits.
• Pass-22 — 237/247 (96.0%), 15/19 perfect. HARDWARE.toml SSOT + workspace lints + Cluster B error propagation.

What's next

OSS release prep (alpha-2.43-share) was the major non-code milestone: archive tags, docs cleanup, atlas-internal/ separation for proprietary artefacts.

For the current roadmap, check the repo's pinned issues and the authoritative decision records at docs/adr/.

Paper Summary

The Atlas team maintains a technical paper in paper/atlas.tex — a two-column LaTeX document titled:

"Atlas: A Custom CUDA Inference Engine for Hybrid Mamba/Attention MoE Models on NVIDIA Blackwell GB10"

The ArXiv version is the academic-facing companion to this book. Where the book is a guide for operators and contributors, the paper is the reference you cite from another piece of research.

Abstract (paraphrased)

Atlas is a pure-Rust LLM inference engine targeting a single (Hardware, Model, Quantization) tuple at a time and hyperoptimizing each tuple independently. On NVIDIA's GB10 Grace-Blackwell Superchip (SM121), running Qwen3.5-35B-A3B in NVFP4 with MTP speculative decoding, Atlas reaches 131 tokens/second — 3.6× NVIDIA's vLLM on the same hardware and model. The paper describes the kernel registry mechanism, the SBIO-based Rust trait layer that enables testing without a GPU, the NVFP4 software E2M1 conversion that works around SM121's missing native FP4 MMA, and the Marconi SSM snapshot cache that makes prefix caching correct on hybrid SSM+attention models.

Key claims the paper makes

• Specialization scales. Per-(H, M_q) kernel sets, combined with vendor-agnostic runtime traits, scale to many targets without regressing existing ones.
• Software E2M1 on SM121 is viable. Branchless FP32 → E2M1 conversion in 7 ALU ops closes the gap left by the missing hardware instruction; Atlas's NVFP4 throughput on GB10 is the silicon ceiling.
• MTP + constrained decoding is a throughput multiplier on agent workloads. XGrammar-masked MTP drafts achieve ~95% acceptance inside tool calls, yielding +37% throughput on agentic traces.
• Hybrid SSM+attention prefix caching requires state snapshots. Marconi (the SSM snapshot cache) produces byte-identical warm-cache output; without it, prefix caching would silently diverge on hybrid models.

How the book and the paper relate

• The book covers operations + architecture + contribution workflow. If you want to run or extend Atlas, start here.
• The paper covers the research claims + benchmark methodology + comparisons to contemporaneous work (vLLM NVFP4, TRT-LLM NVFP4, SGLang, FlashInfer). If you're writing a related paper or a systems course, cite it.

Both share kernel benchmark numbers, the supported-model matrix, and the architectural rationale. The book is the more expansive document; the paper is the tighter academic framing.

References

• Paper source: paper/atlas.tex (build with pdflatex).
• Citations the paper relies on (also in the README's Citations section):
  • FlashAttention-2 (ICLR 2024) — tiled online softmax
  • FlashAttention-4 (2025) — software polynomial exp, conditional softmax rescaling
  • FlashInfer (MLSys 2025) — block-sparse paged KV, gather-SMEM-MMA
  • SageAttention 3 (NeurIPS 2025) — native FP4 attention on newer Blackwell
  • LeanAttention (2024) — stream-K tile scheduling for decode
  • XGrammar (2024) — token-bitmap automaton for constrained decoding

A Category-Theoretic Perspective

The Atlas book argues its case in prose. The prose carries the claim: for every (Hardware, Model, Quantization) target, there exists a kernel configuration that runs at the hardware's theoretical peak; general frameworks cannot reach that peak because they pay a genericity tax; Atlas refuses the tax by specializing per target while keeping abstractions above the kernel layer.

Category theory gives precise names for the structures that claim leans on. This appendix names them. It is not a proof of performance, not a tutorial in category theory, and not required reading for anyone wanting to run or extend Atlas. It is a lens. Read it if you want to see the same design with fewer words.

Standard references for the underlying mathematics: Saunders Mac Lane, Categories for the Working Mathematician (second edition); Emily Riehl, Category Theory in Context (freely available). Everything below uses only the first two chapters of either.

1. The target category 𝒯

A category is a collection of objects together with arrows (morphisms) between them, closed under composition and equipped with an identity arrow on every object. In symbols: ob(𝒯) is a class, and for every ordered pair A, B ∈ ob(𝒯) there is a set 𝒯(A, B) of arrows.

Atlas's target category 𝒯 has one object per supported (H, M, q) triple. In code, these objects are atlas_core::target::KernelTarget values — GB10_QWEN35_NVFP4, GB10_QWEN3_NVFP4, GB10_QWEN35_122B_NVFP4, and nine siblings. The const declarations in crates/atlas-core/src/target.rs are a literal list of ob(𝒯).

The non-obvious choice is the morphism set: for every distinct pair A ≠ B, 𝒯(A, B) = ∅. The only arrows are identities. 𝒯 is a discrete category.

This choice matters. A non-identity arrow f : A → B would mean "a canonical way to go from kernel set A to kernel set B" — a declared compatibility. Such compatibilities are temptations that collapse specialization: the moment you posit f : (GB10, Qwen3.5-35B, NVFP4) → (GB10, Qwen3-Next-80B, NVFP4), you have committed to a kernel set that serves both, or at least to a shared essence that both factor through. That is the shape of vLLM. Atlas refuses by making 𝒯 discrete.

The specialization thesis, in one sentence: 𝒯 is discrete, and all performance claims are local to an object.

2. 𝒯 as a product

The three axes decompose: 𝒯 ≅ Hw × Mod × Quant, where each factor is itself a discrete category (one object per supported value). The product comes with projection functors — π_Hw : 𝒯 → Hw, π_Mod : 𝒯 → Mod, π_Quant : 𝒯 → Quant — that read off one coordinate.

A functor is a structure-preserving map between categories: it sends objects to objects and arrows to arrows, respecting identities and composition. On discrete categories a functor is just an object-to-object function.

The product decomposition is visible in three places in the repo:

• The directory tree: kernels/<hw>/<model>/<quant>/ mirrors the three-factor product exactly. A leaf is an object of 𝒯.
• The build-time wildcards: ATLAS_TARGET_HW, ATLAS_TARGET_MODEL, ATLAS_TARGET_QUANT in atlas-kernels/build.rs select subsets of each factor independently.
• The workspace crate split: spark-runtime/spark-comm insulate the Hw axis, spark-model insulates the Mod axis, spark-model/src/quant_format/ + atlas-kernels insulate the Quant axis.

Orthogonality of axes is not a lucky accident — it is the defining property of a categorical product. Adding an object to Hw does not touch Mod × Quant; the projection π_{Mod×Quant} is unchanged. This is exactly the empirical fact that "adding a new hardware vendor is two trait impls and a directory".

3. Kernels as a functor

The primary structure over 𝒯 is the kernel assignment:

Kernels : 𝒯 → 𝐒𝐞𝐭

𝐒𝐞𝐭 is the category of sets. Kernels sends each target to its set of compiled PTX modules. The auto-generated file atlas-kernels/src/target_ptx.rs is this functor materialised in code. ptx_modules(target: &KernelTarget) -> Option<&'static [PtxModule]> is the functor applied to an object.

Because 𝒯 is discrete, there are no naturality squares to draw — Kernels has complete freedom per object, which is the whole point. The image Kernels(H, M, q) in the default multi-model image has ~30–40 elements; no two targets share an element by construction.

4. Build-to-runtime as a composition of functors

Three categories and two functors sit in a line:

Sources  ──[ComputeTarget.compile]──►  Binaries  ──[embed + load]──►  KernelHandles

Sources has one object per (H, M, q) leaf directory whose underlying data is the set of .cu / .metal / .hip files inside. Binaries has one object per leaf whose underlying data is the set of compiled PTX / metallib / HSACO byte blobs. KernelHandles holds the runtime-resident entries returned by GpuBackend::kernel(module, function).

The first arrow is the ComputeTarget trait in crates/atlas-core/src/compute.rs. It is a vendor-indexed family of functors — one concrete functor per Vendor (Nvidia, Amd, Apple, Intel). Adding a new hardware vendor means adding a new member to the family. The rest of the diagram commutes unchanged: Binaries → KernelHandles doesn't care how the binaries were produced.

This is the categorical reading of "the abstractions sit above the kernel layer, not inside it". The abstractions are the arrows in the diagram. The kernels are elements of the objects. Arrows and elements live at different levels; only arrows need to be generic.

5. The GpuBackend trait as an algebraic theory

An algebraic theory is a signature (operation symbols with arities) plus equations that any implementation must satisfy. A model of the theory is a set together with operations that satisfy the equations. Different sets can be different models of the same theory — this is the mathematical name for "multiple implementations of the same trait".

The GpuBackend trait in crates/spark-runtime/src/gpu.rs is such a theory. Its operations are alloc, free, kernel, launch, synchronize, copy_h2d, and so on (27 methods). The (unwritten, but real) equations include "free after alloc returns memory to the pool", "synchronize serialises previously-launched work on the given stream", and "launch of a kernel with pointer arguments passes the addresses unchanged to the kernel".

Two models ship:

• AtlasCudaBackend — implements the theory by delegating to the CUDA driver API.
• MockGpuBackend — implements the theory by recording launches and returning the opaque successes the equations demand.

The business logic — scheduler, engine, layer code — is polymorphic over the choice of model. In category-theoretic language, business logic is an arrow in the category of GpuBackend-algebras, and evaluating it requires picking a model. The cargo test suite evaluates in MockGpuBackend; production evaluates in AtlasCudaBackend. Both evaluations agree on all facts that depend only on the algebraic theory — sequence of launches, argument correctness, allocation hygiene. This is why ~80% of the test surface runs without a GPU.

This is the formal meaning of SBIO: business logic never directly performs I/O because it never commits to a model. Commitment happens at the top of main.

6. The kernel registry as a coproduct

A coproduct (or disjoint union) in 𝐒𝐞𝐭 is the set-theoretic union of pairwise-disjoint copies:

all_ptx  ≅  ∐_{(H,M,q) ∈ 𝒯}  Kernels(H, M, q)

In code, all_ptx_sets() in atlas-kernels/src/lib.rs returns this coproduct. Each (H, M, q) contributes a summand; the summands share no elements by construction, because different leaf directories produce different PTX blobs with different module names.

The coproduct has a universal property that is worth stating because it matches the design discipline: for any set S and family of functions f_{H,M,q} : Kernels(H, M, q) → S, there is a unique function f : all_ptx → S that restricts to each f_{H,M,q}. The registry dispatch at runtime — "given a target, return the right PTX set" — is the inverse construction: a function out of all_ptx that factors through the target index.

Adding a new target adds a new summand. The universal property says the existing f_{H,M,q} for other targets don't need to change. This is the formal meaning of "specialization is a directory, not a template".

7. Where general frameworks sit in this picture

A general framework — call one 𝒢 — offers a kernel assignment Kernels_𝒢 : 𝒯 → 𝐒𝐞𝐭 that factors through a smaller "essence" category ℰ:

Kernels_𝒢  :  𝒯  ──F──►  ℰ  ──G──►  𝐒𝐞𝐭

ℰ has richer morphisms than 𝒯. Examples of non-identity arrows in ℰ:

• Shape polymorphism (a single templated kernel covers seq_len = 128 and seq_len = 256 via a compile-time branch).
• Dtype dispatch (a single kernel handles BF16 and FP16 via a runtime tag).
• Just-in-time specialisation (a single source file JITs per shape on first call).

The factoring is attractive because the image of F can be small: you write one kernel in ℰ and cover many objects of 𝒯. The cost is paid by G: every time G realises a morphism from the ℰ-image down to a specific 𝒯-object, real work happens — a branch, a dispatch, a JIT compilation, a dequant-to-BF16 fallback. Those costs are the genericity tax.

Atlas refuses the factoring. There is no ℰ. Kernels : 𝒯 → 𝐒𝐞𝐭 is defined directly, object by object, with no intermediate. This is why atlas-kernels has no runtime compilation and no dispatch branching: there is nothing to branch over.

The 3.6× gap on Qwen3.5-35B against NVIDIA's vLLM is the cost of NVIDIA's G on that particular object. The benchmarks in Benchmarks report what the cost is, per kernel and end-to-end, across the whole matrix.

8. Reading the book through this lens

The rest of the book, re-read categorically:

• The Part I philosophy chapter argues the refusal of ℰ in operator terms: the abstractions that enable scaling to new targets live above the kernel layer, not inside kernels themselves.
• The Part II philosophy chapter shows how the refusal forces specific code structure: the kernel tree is the coordinate system, the crate split is the product decomposition.
• The workspace chapter is a walk through ob(𝒯) and the trait layer above it.
• The dispatch chapter traces a single request through the functor composition of Section 4.
• The SBIO chapter is the operational version of Section 5 — two models of one theory.
• The crate chapters describe one vertex of the diagram each.
• The deep-dive chapters describe Kernels(H, M, q) at a single object each — the inside of one summand of the coproduct in Section 6.

Nothing in the book changes when you put on the categorical lens. What changes is the vocabulary you have for arguing about proposals — "does this preserve the product structure of 𝒯?", "does this force a factoring through ℰ?", "is this an arrow between models of the algebraic theory, or an operation inside one model?" These are questions a code review benefits from asking aloud.

9. What this perspective does not prove

Category theory names structures. It does not measure throughput, does not verify kernel correctness, does not port Atlas to a new hardware vendor, and does not write tool-call parsers. Everything the formalism claims follows from the code already being organised along these lines; the formalism is a mirror, not an engine.

In particular:

• Performance is empirical. See Benchmarking.
• Correctness is tested. See Contributing.
• Porting a vendor is design work, not paperwork. The categorical answer ("one new ComputeTarget impl, one new GpuBackend impl, kernel source") names the files but not the effort.

The formalism earns its keep when it helps catch a design drift early. When a PR proposes a cross-cutting trait that couples two axes of the product — say, a method on GpuBackend that only makes sense for one model family — the categorical reading surfaces it immediately: this proposal introduces an arrow between factors of Hw × Mod × Quant, breaking the product. That reading has saved review time in the past and will again. It is why the appendix is worth writing down.

Further reading. For the mathematics used above: Mac Lane, Categories for the Working Mathematician, chapters I–III; Riehl, Category Theory in Context, chapters 1–4. For the engineering the formalism describes: Philosophy (Part I), Philosophy (Part II), Kernel Dispatch Pipeline, SBIO.

Glossary

Short definitions for the acronyms and names that recur in this book and the Atlas codebase.

Further Reading

Curated references that informed Atlas's design. Not exhaustive — just the papers, articles, and prior-art projects worth reading if you want to understand why Atlas is shaped the way it is.

Kernel engineering

• FlashAttention-2 — Tri Dao (ICLR 2024). arXiv:2307.08691. Foundation of Atlas's prefill kernels — tiled online softmax, Q/K/V tiling with shared memory, causal masking.
• FlashAttention-4 — Shah, Bikshandi, Zhang, Thakkar, Ramani, Dao (2025). arXiv:2603.05451. Conditional softmax rescaling (skip ~90% of rescale ops), software polynomial exponential (sw_exp, avoids SFU bottleneck).
• FlashInfer — Ye et al. (MLSys 2025 Best Paper). arXiv:2501.01005. Block-sparse paged KV, gather-SMEM-MMA pattern. Informed Atlas's paged decode kernel.
• SageAttention 3 — Zhang et al. (NeurIPS 2025). arXiv:2505.11594. Native FP4 attention on newer Blackwell. Planned direction when SM12x+ silicon lands.
• LeanAttention — Roy, Vassilieva, Willke, Mendis (2024). arXiv:2405.10480. Stream-K tile scheduling for decode attention. Planned for SM occupancy improvements.
• CUTLASS documentation and examples. The BF16 MMA fragment shapes and cp.async pipelining patterns in Atlas's kernels follow CUTLASS conventions.

State-space models

• Mamba: Linear-Time Sequence Modeling with Selective State Spaces — Gu, Dao (2023). arXiv:2312.00752. The original selective SSM.
• Mamba-2 — Dao, Gu (ICML 2024). arXiv:2405.21060. The variant Nemotron-H uses.
• Gated Delta Networks — Yang, Dao, et al. (2024). Closer to Qwen3.5's GDN formulation.
• GDN register-tile results — Atlas's internal experiments in gdn_regtile_results.md track tile-shape tradeoffs on GB10.

Quantization

• NVFP4 / FP4 microscaling — NVIDIA's blog posts on Blackwell FP4. The public docs for SM120 coverage are thin; much of Atlas's SM121 workaround has no upstream equivalent yet.
• SmoothQuant — Xiao, Lin, et al. (ICML 2023). Scale factor calibration ideas used indirectly in Atlas's FP8 KV calibration.
• Compressed-Tensors format — the HF compressed-tensors library's on-disk FP8 block-scaled layout.
• TurboQuant (Atlas internal) — docs/turboquant-plus.md. WHT + Lloyd-Max 4-bit KV cache with ~2× lower MSE than NVFP4 at the same bit rate.

Speculative decoding

• Fast Inference from Transformers via Speculative Decoding — Leviathan, Kalman, Matias (ICML 2023). The foundational paper.
• Medusa: Multi-Token Prediction Heads — Cai et al. (2024). MTP-style draft heads.
• Self-speculative Decoding — Layer-skipping drafter. Atlas implements a variant.

Constrained decoding

• XGrammar — Li, Chen, Chen, et al. (2024). arXiv:2411.15100. The token-bitmap automaton approach Atlas uses.
• SGLang — Zheng et al. (NeurIPS 2024). Broader exploration of structured-output techniques; Atlas shares design elements with SGLang's regex_fsm.

Inference systems

• vLLM — Kwon et al. (SOSP 2023). The PagedAttention paper. Atlas's paged KV cache follows the vLLM model with Atlas-specific kernel work below it.
• TensorRT-LLM documentation and source. Atlas's TRT-LLM benchmark comparisons in docs/ATLAS_SPARK_JOURNEY.md are informed by reading the TRT-LLM codebase; that journey records the 29.6 tok/s ceiling for NVFP4 on SM121 TRT-LLM.
• SGLang — structured-generation inference framework.
• Triton — inference server. Peripheral; Atlas does not use it but the operational patterns are informative.

MoE routing

• Mixture-of-Experts with Sigmoid Routing — various 2023–2024 papers. MiniMax-M2.7's 256-expert sigmoid-routed MoE is an unusual design Atlas had to support natively.
• Switch Transformer — Fedus, Zoph, Shazeer (2021). The top-1 routing baseline.

Adjacent projects worth studying

• scitix/InstantTensor — the fast safetensors loader Atlas's O_DIRECT + pipelined reader is modeled on.
• huggingface/tokenizers — the tokenizer library Atlas wraps.
• huggingface/safetensors — format spec + Rust impl.
• PyKeOps / symbolic autograd — not used, but the philosophy (compile once per shape, amortise forever) resonates with AI Kernel HyperCompiling.

Atlas-internal references

Inside the repo, the canonical long-form references are the architecture decision records in docs/adr/, plus the top-level notes alongside them. Notable:

• docs/adr/0004-nvfp4-fp8-quantization.md — NVFP4/FP8 quantization, including why --kv-high-precision-layers exists.
• docs/adr/0003-hybrid-ssm-attention.md — hybrid SSM/attention design and chunked SSM prefill.
• docs/adr/0007-tp-ep-composition.md, docs/adr/0011-ep-batched-decode-optimization.md — EP=2 MoE dispatch and batched decode.
• docs/adr/0010-vendor-xgrammar.md — constrained decoding via vendored XGrammar.
• docs/turboquant-plus.md — TurboQuant KV.
• docs/ARCHITECTURE.md, docs/ATLAS_KERNELS.md, docs/HARDWARE.md — the system, kernel, and hardware overviews.

For the broader research context that informed Atlas's direction, see docs/atlas-spark-research-articles.md in the repo — a rolling curated list that's longer and more current than this page.

API Reference

If you are not redirected automatically, follow this link to the Rust API reference.

The API reference is generated from the crate source with cargo doc --workspace --no-deps on every merge to main. Top-level crates:

• atlas_core — target abstractions, tensor, dtype, kernel registry, host-side FP8/BF16 numerics
• atlas_kernels — embedded PTX registry
• spark_runtime — GPU backend, KV cache, sampler
• spark_comm — collective-op trait + NCCL impl
• spark_model — layer assembly, weight loaders, engine
• spark_server — HTTP server, tool parsing
• atlas_spark_bench — benchmark client