Tutorial 18: End-to-End — Putting It All Together

This is the final tutorial. We assemble every component from the previous 17 tutorials into a complete, working inference engine with a CLI and a server. By the end, you should be able to load a real model, generate text, and serve it over HTTP.

---

1. The Complete Architecture

┌──────────────────────────────────────────────────────────┐
│  CLI (cli.py)                                            │
│    run     → single-request streaming generation          │
│    serve   → OpenAI-compatible HTTP server                │
│    bench   → performance benchmarking                     │
└──────────────────────────────────────────────────────────┘
        │
        ▼
┌──────────────────────────────────────────────────────────┐
│  Engine (engine/)                                        │
│    generator.py  → stream_generate() for single requests  │
│    scheduler.py  → BatchedEngine for concurrent requests  │
│    sampler.py    → temperature / top-p / top-k sampling   │
└──────────────────────────────────────────────────────────┘
        │
        ▼
┌──────────────────────────────────────────────────────────┐
│  Model (core/)                                           │
│    model.py      → Qwen3ForCausalLM (dense + MoE)        │
│    attention.py  → SDPA, GQA, paged attention            │
│    rope.py       → RoPE                                   │
│    norms.py      → RMSNorm                                │
│    embedding.py  → Embedding + tied weights               │
│    moe.py        → Mixture of Experts                     │
│    quantize.py   → 4-bit quantization                     │
│    layers.py     → linear, silu, softmax                  │
└──────────────────────────────────────────────────────────┘
        │
        ▼
┌──────────────────────────────────────────────────────────┐
│  State (cache/)     Kernels (ext + kernels.py + Metal)   │
│    dense.py          rms_norm, rope, swiglu,             │
│    paged.py          quantized_matmul, paged_attention,  │
│                      decode_attention                    │
└──────────────────────────────────────────────────────────┘
        │
        ▼
┌──────────────────────────────────────────────────────────┐
│  Loading (config.py + weights.py)                        │
│    config.json  → config dataclass                       │
│    safetensors  → WeightStore                            │
└──────────────────────────────────────────────────────────┘

---

2. Loading a Real Model

from anllm.core.model import load_model, ModelFeatures

# Load a real Qwen3 model (downloaded from HuggingFace, or local)
# Smallest options: Qwen3-0.6B / Qwen3-1.7B / Qwen3-4B (4-bit)
model = load_model(
    "mlx-community/Qwen3-0.6B-4bit",
    features=ModelFeatures(kernels=False),  # use pure Python for now
)
print(f"Loaded model: {model.config.hidden_size} hidden, "
      f"{model.num_hidden_layers} layers, "
      f"{model.vocab_size} vocab")

What Happens Under the Hood

def load_model(model_path, features=None):
    from pathlib import Path

    # 1. Locate config.json
    model_path = Path(model_path)
    config_file = (
        model_path / "config.json" if model_path.is_dir()
        else model_path.parent / "config.json"
    )

    # 2. Parse config into a dataclass (handles Qwen3-specific fields)
    config = Qwen3Config.from_json_file(config_file)

    # 3. Load safetensors weights + detect quantization
    store = WeightStore.from_path(model_path)

    # 4. Construct the model with those weights
    return Qwen3ForCausalLM(config, store, features)

---

3. Running Single-Request Generation (CLI)

from anllm.cli import run

# The CLI has three commands:
#   anllm run   "Tell me a joke"   → generate text
#   anllm serve                    → start HTTP server
#   anllm bench                    → benchmark performance

def run_generation():
    # Tokenize (in practice, use a tokenizer library like transformers)
    prompt = "Once upon a time, "
    prompt_tokens = tokenize(prompt)  # → [100, 200, ...]

    # Configure the sampler
    from anllm.engine.sampler import make_sampler
    sampler = make_sampler(temp=0.7, top_p=0.9, top_k=50)

    # Generate
    from anllm.engine.generator import stream_generate
    tokens = []
    for token in stream_generate(
        model,
        prompt_tokens,
        max_new_tokens=100,
        sampler=sampler,
        eos_token_ids=[151645, 151643],  # <eos>, <endoftext>
    ):
        tokens.append(token)
        text = detokenize(tokens)
        print(text, end="", flush=True)  # print as we generate

What Happens Per Token

Step 0 (prefill):
  Input: [100, 200, 300, 400, 500]   ← full prompt
  Output: token 700                  ← predicted next token
  KV cache now holds 5 tokens

Step 1 (decode):
  Input: [700]                       ← just the new token
  Output: token 250
  KV cache now holds 6 tokens

Step 2 (decode):
  Input: [250]
  Output: token 900
  KV cache now holds 7 tokens
  ...

---

4. Running the Server

# From the terminal:
anllm serve --model mlx-community/Qwen3-0.6B-4bit --max-batch 4

# Or in Python:
from aioserver import App
import uvicorn

app = create_server(model, batch_size=4)

# Start the server
uvicorn.run(app, host="0.0.0.0", port=8000)

Test It

# Chat completion (non-streaming)
curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "qwen3", "messages": [{"role": "user", "content": "Hello!"}], "max_tokens": 64}'

# Chat completion (streaming)
curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "qwen3", "messages": [{"role": "user", "content": "Tell me a story"}], "stream": true}'

# List models
curl http://localhost:8000/v1/models

---

5. The Full Generation Pipeline

Let's trace a complete request through the entire stack:

Client sends: {"messages": [{"role": "user", "content": "What is 2+2?"}]}
    │
    ▼
Validator: ChatCompletionRequest
    │
    ▼
Tokenizer: "What is 2+2?" → [1001, 125, 355, 3, 1050, 525, 30]
    │
    ▼
Scheduler.submit(prompt_tokens, max_new_tokens=64, temp=0.0)
    │
    ▼
Queue → Prefilling (chunked): process 7 tokens
    │
    ▼
Awaiting Slot → Decode Slot (when free)
    │
    ▼
Each step:
    Embed → TransformerBlocks(32) → Norm → LM Head
    → logits (1, 151936)
    → argmax (temp=0) → next token
    → append to KV cache
    → yield token
    │
    ▼
Client receives: "It is 4."

Finish reason: "stop" (hit EOS) or "length" (hit max_tokens)

---

6. Performance Tuning

Set 1: Correctness First

features = ModelFeatures(
    kernels=False,             # pure Python — always works
    decode_attention_kernel=False,
    force_dequantized=False,
)

Set 2: Enable Metal Kernels (if extension is compiled)

features = ModelFeatures(
    kernels=True,              # use native RMSNorm, RoPE, SwiGLU
    decode_attention_kernel=True,  # specialized decode attention
    force_dequantized=False,   # keep quantized matmul (fastest)
)

Set 3: Force Dequantization (debugging quantization issues)

features = ModelFeatures(
    kernels=True,
    decode_attention_kernel=False,
    force_dequantized=True,    # use full-precision weights (memory heavy)
)

Scheduler Tuning

from anllm.engine.scheduler import BatchedEngine

engine = BatchedEngine(
    model,
    batch_size=4,        # concurrent decode slots — bigger = more throughput
    max_seq_len=512,     # max tokens per request (incl. prompt + generated)
    prefill_step=128,    # chunk size for prefill — smaller = lower latency
)

Balance:

• Larger batch_size → more throughput, more memory
• Larger prefill_step → faster prefill, longer per-step latency
• Smaller prefill_step → more responsive, more steps

---

7. Debugging Checklist

Wrong output / NaN logits

1. Check the config: num_attention_heads % num_key_value_heads == 0
2. Check RoPE: dims must be even, base must match the checkpoint
3. Check QK-norm: only Qwen3 has it (q_norm, k_norm tensors)
4. Check tie_word_embeddings: if weight is missing, lm_head may be separate

Out of memory

1. Lower batch_size or max_seq_len
2. Enable quantization (4-bit) at load time
3. Use paged cache instead of dense
4. Reduce prefill_step

Slow generation

1. Enable Metal kernels
2. Use logits_to_keep=1 — don't compute logits for all positions
3. Verify batch size is reasonable (each additional slot adds memory)

Wrong results with quantization

1. Force dequantization: force_dequantized=True
2. Verify group_size matches the checkpoint config
3. Check scales/biases are loaded (not just weight)

---

8. What You Have Built

Across these 18 tutorials, you have built:

1. Math foundations — linear algebra, softmax, matrix multiplication
2. Core layers — embeddings, RMSNorm, RoPE, SiLU, linear
3. Attention — SDPA, GQA, causal masking
4. Model architecture — transformer blocks, MoE, quantization
5. State management — KV cache, paged attention
6. Generation — autoregressive loops, sampling
7. Concurrency — continuous batching, scheduling
8. Performance — custom Metal kernels
9. Serving — OpenAI-compatible HTTP API

Each component is a complete, working piece that connects to the others through well-defined interfaces. You can extend it with new models (see ADD.md), new kernels, or new serving features.

---

9. Next Steps

• Add a new model family: Follow ADD.md to add Llama, Mistral, Qwen2.5.
• Write custom Metal kernels: Optimize the hot paths (quantized matmul, paged attention, decode kernel).
• Add speculative decoding: Use a small draft model to predict 4-8 tokens, verify with the large model in one pass.
• Add multi-GPU support: Split layers across MLX's multi-device API.
• Add the tokenizer: Integrate a local tokenizer (e.g., sentencepiece) so the server accepts raw text instead of token IDs.
• Add logprobs to the API: Expose per-token probabilities for downstream applications.
• Add stop sequences: Currently only EOS is checked; multi-string stop support requires more logic.

Happy learning to build systems!