Inside an LLM: Pure Python Inference for Qwen3-0.6B on MLX

Type "The capital of France is" and a model finishes it with "Paris", in roughly 900 milliseconds. What happens inside that sliver of time? Which chain of transformations turns five words into a prediction?

The best way to find out is to build the whole inference pipeline yourself instead of calling a library that hides the details. The walkthrough below does that with real Qwen3-0.6B weights, loaded into plain Python through MLX, Apple's array framework. The complete script is about 200 lines and runs on any Apple Silicon Mac. Setup instructions come at the end.

The pipeline at a glance

Qwen3, GPT-4, Llama and their relatives all run through one shared flow:

Text → Tokenize → Embed → [Transformer Block × 28] → LM Head → Logits → Sample → Next Token
                           ↑                                                              |
                           └-------------------------- KV Cache --------------------------┘

Text goes in, becomes numbers, passes through 28 transformer blocks, and leaves as a probability distribution over 151,936 candidate next tokens. One token gets picked and appended, and the loop starts over.

That loop has two phases:

• Prefill. All input tokens are processed in one go. The model reads the prompt and builds its internal state.
• Decode. Tokens are produced one at a time. A new token only has to attend to earlier ones, whose keys and values sit in the KV cache, so nothing is reprocessed.

Prefill runs in parallel and decode runs sequentially. The KV cache is what keeps decode quick. Without it, producing token N would mean recomputing attention across all N-1 tokens before it.

Now for each stage in turn.

Step 1: Tokenization turns text into integers

A model never sees words, only integers.

from tokenizers import Tokenizer
tok = Tokenizer.from_file("Qwen3-0.6B/tokenizer.json")
ids = tok.encode("The capital of France is").ids
# [785, 6722, 315, 9625, 374]

Each integer is an index into a vocabulary of 151,936 tokens. The tokenizer uses byte pair encoding, so token boundaries don't always line up with word boundaries. " capital" is one token and " France" is another, with the leading space belonging to the token. Writing BPE from scratch isn't the interesting part here. The tokenizer is a lookup table, and the script simply uses the published one.

Step 2: Embedding turns integers into vectors

The first real computation maps each token ID to a 1,024-dimensional vector.

x = W["model.embed_tokens.weight"][mx.array(token_ids)]  # (seq, 1024)

Nothing more than a table lookup: the embedding matrix is shaped (151936, 1024), one row per token and one column per hidden dimension, and the code indexes into it. Think of the matrix as the model's dictionary: every token has a learned 1,024-number fingerprint that encodes meaning. The fingerprints for "Paris" and "France" point in similar directions, while "Paris" and "computer" point in very different ones.

Step 3: RMSNorm

Before attention, the hidden states get normalised. Modern models use RMSNorm in place of the older LayerNorm.

In words, the formula divides each vector by its root-mean-square value, which is the square root of the mean of its squared entries plus a tiny constant, and then multiplies by a learned weight per feature.

def rms_norm(x, w):
    return x * mx.rsqrt(mx.mean(x * x, axis=-1, keepdims=True) + EPS) * w

So each vector is scaled by its own magnitude, then adjusted by learned per-feature weights. There's no mean subtraction and no bias, which makes it simpler than LayerNorm without costing quality. EPS (1e-6) prevents division by zero, and the weight w, learned in training, lets the model rescale each of the 1,024 dimensions independently.

Step 4: RoPE weaves position into the vectors

This is where many "LLM from scratch" articles lose their readers, but it's simpler than it looks.

The problem. Attention is permutation-invariant. Shuffle the input tokens and the output is the same shuffle, so the model can't tell which token came first.

The old fix. Add a position number to each token's embedding. That works, but it doesn't generalise, and the model struggles to reason about relative distances.

The modern fix. Turn the query and key vectors through an angle that is a function of where the token sits. That's Rotary Position Embedding, or RoPE. Each pair of dimensions is treated as a point in a plane and turned by an angle set by the token's position.

def apply_rope(q, k, positions):
    half = D_HEAD // 2  # 64
    freqs = 1.0 / (THETA ** (mx.arange(half, dtype=mx.float32) / half))
    angles = mx.array(positions, dtype=mx.float32)[:, None] * freqs[None, :]
    cos = mx.cos(angles)[None, :, :]
    sin = mx.sin(angles)[None, :, :]

    def rotate(x):
        x1 = x[..., :half]
        x2 = x[..., half:]
        return mx.concatenate([x1 * cos - x2 * sin, x1 * sin + x2 * cos], axis=-1)

    return rotate(q), rotate(k)

Picture each pair of dimensions as a 2D plane. Tokens at different positions turn by different amounts, so when attention takes the dot product of a query and a key, the answer depends on how far apart they are and not only on what they contain. Here are three tokens at positions 0, 1 and 2, each drawn as a 2D arrow rotated a bit further than the last:

Position 0 (angle = 0°)        Position 1 (angle = 24°)       Position 2 (angle = 48°)

      q                                      q                                q
      |                                      /                                /
      |                                     /                                /
      |                                    /                                /
      |                                   /                                /
      +------ x                          +------ x                        +------ x

  Same content, 0° rotation         Rotated 24° counterclockwise    Rotated 48° counterclockwise
  "I'm at the start"               "I'm one step in"               "I'm two steps in"

In the original drawing the third arrow is not rotated further than the second, so read the diagram as a rough picture and not an exact one.

Attention compares a query arrow with a key arrow through a dot product. Arrows at the same angle give a big product, and arrows at very different angles give a small one. So the product reflects the distance between two tokens, with no explicit position number stored anywhere. RoPE doesn't add position as a separate signal. Rotation builds position into every vector.

Qwen3 sets THETA = 1,000,000, a high base frequency that lets the model cope with very long contexts, as many as 262K tokens without extensions.

Step 5: Attention with grouped query heads

This is the centrepiece. Understand attention and you understand the transformer.

From MHA to MQA to GQA

Multi-Head Attention (MHA):  16 Q heads, 16 K heads, 16 V heads  — expensive
Multi-Query Attention (MQA): 16 Q heads,  1 K head,  1 V head    — too aggressive
Grouped Query Attention:     16 Q heads,  8 K heads, 8 V heads   — sweet spot

Qwen3-0.6B uses GQA, with 16 query heads and only 8 key-value heads, so every KV head serves two query heads. That halves the KV cache with little loss in quality.

The whole attention function

def attention(x, layer_idx, cache, positions):
    p = f"model.layers.{layer_idx}.self_attn"

    # Project to Q, K, V
    q = x @ W[f"{p}.q_proj.weight"].T   # (seq, 2048) = 16 heads × 128
    k = x @ W[f"{p}.k_proj.weight"].T   # (seq, 1024) =  8 heads × 128
    v = x @ W[f"{p}.v_proj.weight"].T   # (seq, 1024) =  8 heads × 128

    # Split into heads
    q = q.reshape(-1, N_HEADS, D_HEAD).transpose(1, 0, 2)    # (16, seq, 128)
    k = k.reshape(-1, N_KV, D_HEAD).transpose(1, 0, 2)       # (8, seq, 128)
    v = v.reshape(-1, N_KV, D_HEAD).transpose(1, 0, 2)       # (8, seq, 128)

    # QK-Norm (Qwen3-specific: normalize each head before RoPE)
    q = rms_norm(q, W[f"{p}.q_norm.weight"])
    k = rms_norm(k, W[f"{p}.k_norm.weight"])

    # Apply RoPE
    q, k = apply_rope(q, k, positions)

    # KV Cache: append new keys/values
    if cache[layer_idx] is not None:
        ck, cv = cache[layer_idx]
        k = mx.concatenate([ck, k], axis=1)
        v = mx.concatenate([cv, v], axis=1)
    cache[layer_idx] = (k, v)

    # GQA: expand 8 KV heads → 16 (each KV head shared by 2 Q heads)
    k = mx.repeat(k, N_HEADS // N_KV, axis=0)   # (16, total, 128)
    v = mx.repeat(v, N_HEADS // N_KV, axis=0)

    # Scaled dot-product attention
    scale = D_HEAD ** -0.5
    scores = (q @ k.transpose(0, 2, 1)) * scale   # (16, seq, total)

    # Causal mask (only during prefill)
    if seq_len > 1:
        mask = mx.triu(mx.full((seq_len, total), -1e9, dtype=scores.dtype), k=1)
        scores = scores + mask

    # Softmax + weighted sum of V
    attn = mx.softmax(scores, axis=-1)
    out  = attn @ v

    # Merge heads + output projection
    out = out.transpose(1, 0, 2).reshape(-1, N_HEADS * D_HEAD)
    return out @ W[f"{p}.o_proj.weight"].T

A few parts of it deserve explanation.

QK-Norm. Qwen3 runs RMSNorm over every head's Q and K vectors ahead of RoPE, a newer trick that steadies training. Its weights q_norm and k_norm are 128 numbers long, a scale per dimension within a head.

The KV cache. This is the most important optimisation in inference. During prefill, K and V are computed for every input token and stored. During decode, only the new token needs K and V computed. They're appended to the cache and attention runs over the full history. Skipping the cache would force K and V to be recomputed for every previous token, which turns linear work into quadratic.

The causal mask. During prefill, the token at position 3 mustn't see the token at position 4, which lies in its future. The code adds -1e9 to the upper triangle of the score matrix to enforce that. In decode, where the sequence length is 1, no mask is needed, because the single new token may look at everything before it.

The formula. Attention computes softmax of (Q times K transposed, divided by the square root of the head dimension), then multiplies the result by V. The dot product measures how well each query matches each key. Dividing by the square root of the head size keeps the scores from growing so large that softmax saturates. Softmax turns scores into probabilities, and multiplying by V yields a weighted average of the value vectors.

Step 6: The SwiGLU feed-forward layer

After attention, every token passes through a feed-forward network. Modern models use SwiGLU instead of the older ReLU-based version. In words: two projections are computed from the input, one is passed through SiLU and used as a gate, the two are multiplied element by element, and a final projection brings the result back down.

def ffn(x, layer_idx):
    p = f"model.layers.{layer_idx}.mlp"
    gate = x @ W[f"{p}.gate_proj.weight"].T   # (seq, 3072)
    up   = x @ W[f"{p}.up_proj.weight"].T     # (seq, 3072)
    silu = gate * mx.sigmoid(gate)             # SiLU(x) = x · σ(x)
    return (silu * up) @ W[f"{p}.down_proj.weight"].T  # (seq, 1024)

The feed-forward network is where the model keeps much of its knowledge, which is one reason large dense models feel so capable. It expands the 1,024-dimensional vector to 3,072 dimensions, applies the SiLU gate, and projects back to 1,024. The gate learns which information to let through and which to block, a bit like a bouncer for neurons. Expanding from 1024 to 3072 and back gives the model room to store and transform information, and gating is why modern FFNs beat the ReLU ones.

Step 7: The transformer block

A block combines all of the above:

def transformer_block(x, layer_idx, cache, positions):
    p = f"model.layers.{layer_idx}"
    # Pre-norm attention + residual
    h = rms_norm(x, W[f"{p}.input_layernorm.weight"])
    x = x + attention(h, layer_idx, cache, positions)
    # Pre-norm FFN + residual
    h = rms_norm(x, W[f"{p}.post_attention_layernorm.weight"])
    x = x + ffn(h, layer_idx)
    return x

The sequence is normalise, attend, add the residual, normalise, feed forward, add the residual again.

Those residual connections (x = x + ...) are vital. They let gradients flow straight through the network during training, and they let information skip individual blocks at inference time. A 28-layer model without them would be close to untrainable.

Qwen3-0.6B stacks 28 of these blocks, each refining the representation a little further. By the 28th, the hidden state holds enough understanding of the input to predict the next token.

Step 8: LM head and sampling

After the last block there's one more normalisation, then a projection to vocabulary size:

x = rms_norm(x, W["model.norm.weight"])
logits = x @ W["lm_head.weight"].T   # (seq, 151936)

The output is 151,936 scores, called logits, one for every possible token. A higher score means a likelier token. For "The capital of France is", the top five probabilities after softmax look like this:

| Token     | Probability |
|-----------|-------------|
| Paris     | 0.6582      |
| located   | 0.0274      |
| the       | 0.0233      |
| ...       | 0.0157      |
| ____      | 0.0136      |

"Paris" wins with 65.8%, a confident prediction. Sampling then converts the probabilities into a single token ID:

def sample(logits, temperature=0.0, top_p=0.95):
    if temperature == 0:
        return int(mx.argmax(logits))  # Greedy: always pick the best

    logits = logits / temperature
    probs  = mx.softmax(logits)

    # Top-p: keep smallest set of tokens whose cumulative prob ≥ top_p
    order    = mx.argsort(-probs)
    sorted_p = probs[order]
    cumsum   = mx.cumsum(sorted_p)
    cutoff   = int(mx.sum(cumsum < top_p)) + 1

    sorted_p = sorted_p[:cutoff]
    sorted_p = sorted_p / mx.sum(sorted_p)
    order    = order[:cutoff]

    idx = mx.random.categorical(mx.log(sorted_p[None, :]))
    return int(order[idx])

Temperature controls randomness. At 0 the highest-probability token always wins (greedy). At 1.0 tokens are drawn in proportion to their probabilities, and at 0.7 the distribution sharpens slightly toward the leaders.

Top-p, or nucleus sampling, keeps the smallest group of tokens whose cumulative probability reaches p, for example 0.95. It trims the long tail of unlikely tokens while leaving some diversity.

Step 9: The generation loop

The complete loop prefills once and then decodes:

def generate(prompt, max_tokens=100, temperature=0.0, top_p=0.95):
    ids = tok.encode(prompt).ids

    # Prefill: process all input tokens at once
    cache  = [None] * N_LAYERS
    logits = forward(ids, cache, 0)
    next_id = sample(logits[-1], temperature, top_p)
    ids.append(next_id)

    # Decode: generate one token at a time
    for _ in range(max_tokens - 1):
        if next_id == EOS_ID:
            break
        pos      = len(ids) - 1
        logits   = forward([next_id], cache, pos)
        next_id  = sample(logits[-1], temperature, top_p)
        ids.append(next_id)
    return tok.decode(ids)

That's the whole inference loop. Serving engines like vLLM, SGLang, TensorRT-LLM and mlx-lm are, at heart, optimisations of this core, and the core fits in about 198 lines.

What the script leaves out

The script covers the full inference pipeline of a standard transformer. Newer models add more:

• Hybrid attention. Qwen3.5 reportedly swaps 75% of its attention layers for Gated DeltaNet, a linear attention variant with a fixed-size recurrent state in place of the full attention matrix. Full attention is like photographic memory, perfect and costly. DeltaNet is like compressed notes, efficient but slightly lossy. A hybrid gets efficient long-context handling plus periodic perfect recall.
• Multi-token prediction. Some models predict several tokens per forward pass, trading a little quality for a large speedup.
• Quantisation. Production inference uses 4-bit or 8-bit weights to save memory and time. This script keeps the original bfloat16 weights.
• Batching. Servers handle many requests at once and share computation across them. This script handles one at a time.

None of these change the core algorithm. They sit on top of the same foundation.

Why bother

Knowing how inference works makes you a sharper practitioner. A "KV cache full" warning from a server stops being cryptic. The memory argument for GQA becomes easy to explain. And a model that rambles and repeats itself points you toward temperature and top-p.

A modern LLM isn't magic. It's 28 layers of normalise, attend and feed forward, repeated, with position woven in by rotation and knowledge held in gated linear layers. The maths is undergraduate level. The engineering is where the cleverness lives, but the maths is the foundation. About 200 lines of Python, real weights and real output is all it takes.

Run it yourself

Everything above is runnable on any Apple Silicon Mac (M1 or later).

1. Download the model

Get Qwen3-0.6B from Hugging Face (the Qwen/Qwen3-0.6B repository). Only four files are needed, so skip the rest:

Qwen3-0.6B/
├── config.json
├── model.safetensors       (1.5 GB — the actual weights)
├── tokenizer.json
└── tokenizer_config.json

Place the model folder next to llm_inference.py, because the script resolves ./Qwen3-0.6B/ from whichever directory it's started in.

2. Create a virtual environment

# Create and activate a fresh venv
python3 -m venv llm_venv
source llm_venv/bin/activate
# Install the minimum packages - just two
pip install mlx tokenizers

No PyTorch, no transformers, no accelerate, no safetensors package. MLX does the array maths and the weight loading, and the tokenizers library does BPE encoding and decoding.

3. Run it

python llm_inference.py

Expect output along these lines:

Model:  Qwen3-0.6B  (28 layers, 1024 dim, 16 Q heads, 8 KV heads)
Prompt: The capital of France is
------------------------------------------------------------
Output: The capital of France is Paris. The capital of Italy is
        Rome. The capital of Spain is Madrid...
Time:   0.89s

If you see ModuleNotFoundError: No module named 'mlx', check that the venv is active (source llm_venv/bin/activate) and that the machine is Apple Silicon. MLX doesn't run on Intel Macs, Linux or Windows.

All of it runs on a laptop, and the answer shows up in under a second.