Wewnątrz LLM: czysta inferyencja w Pythonie dla Qwen3-0.6B na MLX

Wpisz „Stolica Francji to” i model uzupełnia to słowem „Paryż” w około 900 milisekund. Co dzieje się w tym krótkim czasie? Która seria przekształceń zamienia pięć słów w przewidywanie?

Najlepszym sposobem na dowiedzenie się tego jest samodzielne stworzenie całego procesu inferencji zamiast korzystania z biblioteki, która ukrywa szczegóły. Poniższy przewodnik pokazuje to przy użyciu rzeczywistych wag Qwen3-0.6B, załadowanych do zwykłego Pythona za pośrednictwem MLX, frameworku do zarządzania tablicami firmy Apple. Pełny skrypt składa się z około 200 linii i działa na każdym komputerze Apple Silicon Mac. Instrukcje instalacji znajdują się na końcu.

Zarys procesu

Qwen3, GPT-4, Llama oraz ich odpowiedniki przechodzą przez ten sam wspólny proces:

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

Tekst trafia do modelu, przekształca się w liczby, przechodzi przez 28 bloków transformatorowych i wychodzi jako rozkład prawdopodobieństwa dla 151 936 możliwych kolejnych tokenów. Wybierany jest jeden token i do tekstu jest on dodawany, po czym pętla zaczyna się od nowa.

Ta pętla składa się z dwóch faz:

• Prefill. Wszystkie tokeny wejściowe są przetwarzane jednocześnie. Model czyta zadanie i buduje swój wewnętrzny stan.
• Decode. Tokeny są generowane po jednym. Nowy token musi uwzględniać tylko te wcześniejsze, których klucze i wartości znajdują się w pamięci KV, więc nic nie jest przetwarzane ponownie.

Faza Prefill przebiega równolegle, natomiast faza Decode sekwencyjnie. To właśnie pamięć KV sprawia, że proces dekodowania jest szybki. Bez niej generowanie tokena N oznaczałoby ponowne obliczanie mechanizmu uwagi dla wszystkich N-1 tokenów poprzedzających go.

A teraz kolejno omówimy każdą z tych faz.

Krok 1: Tokenizacja przekształca tekst w liczby całkowite

Model nigdy nie widzi słów, tylko liczb całkowitych.

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]

Każda liczba całkowita jest indeksem w słowniku zawierającym 151 936 tokenów. Tokenizer używa kodowania par bajtowych, więc granice tokenów nie zawsze pokrywają się z granicami słów. „capital” to jeden token, a „France” – inny, przy czym spacja na początku należy do tego tokenu. Tworzenie BPE od zera nie jest tutaj najciekawszą częścią. Tokenizer to tabela wyszukiwania, a skrypt po prostu wykorzystuje tę już dostępną.

Krok 2: Embedding przekształca liczby całkowite w wektory

Pierwsza rzeczywista operacja przyporządkowuje każdy identyfikator tokena do wektora o wymiarze 1 024.

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

To nic więcej niż wyszukiwanie w tabeli: macierz embeddingów ma kształt (151936, 1024) – jedna wiersz na token i jedna kolumna na wymiar ukryty, a kod odwołuje się do niej. Można traktować tę macierz jako słownik modelu: każdy token ma nauczoną 1024-cyfrową „odpowiedź palcową”, która koduje jego znaczenie. Odpowiedzi palcowe dla „Paris” i „France” wskazują w podobnych kierunkach, natomiast „Paris” i „computer” – w zupełnie różnych.

Krok 3: RMSNorm

Przed zastosowaniem mechanizmu uwagi stany ukryte są normalizowane. Współczesne modele używają RMSNorm zamiast starszego LayerNorm.

Mówiąc prościej, wzór dzieli każdy wektor przez jego wartość średniej kwadratowej – czyli pierwiastek kwadratowy średniej kwadratów jego elementów plus mała stała – a następnie mnoży przez nauczony wagę dla każdej cechy.

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

Każdy wektor jest skalowany według swojej własnej wartości magnitudy, a następnie dostosowywany za pomocą wag nauczonych dla poszczególnych cech. Nie ma tu odejmowania wartości średniej ani biasu, co sprawia, że jest to prostsze od LayerNorm bez straty jakości. EPS (1e-6) zapobiega dzieleniu przez zero, a waga w, nauczona podczas treningu, umożliwia modelowi niezależne skalowanie każdej z 1 024 wymiarów.

Krok 4: RoPE włącza informacje o pozycji do wektorów

W tym miejscu wiele artykułów o „LLM od zera” traci czytelników, ale jest to prostsze, niż się wydaje.

Problem. Mechanizm uwagi jest niezmienny pod względem permutacji. Jeśli przetasować tokeny wejściowe, wynik będzie taki sam, więc model nie może stwierdzić, który token pojawił się pierwszy.

Stary sposób rozwiązania. Dodanie numeru pozycji do embeddingu każdego tokena. To działa, ale nie ma zastosowania w bardziej złożonych przypadkach, a model ma trudności z rozumieniem względnych odległości.

Nowoczesny sposób rozwiązania. Przekształcenie wektorów zapytania i klucza o kąt, który zależy od pozycji tokena. To właśnie Rotary Position Embedding, czyli RoPE. Każda para wymiarów jest traktowana jako punkt na płaszczyźnie i jest obracana o kąt określony przez pozycję tokena.

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)

Wyobraź sobie każdą parę wymiarów jako płaszczyznę 2D. Tokeny znajdujące się w różnych pozycjach obracają się o różne kąty, więc gdy mechanizm uwagi oblicza iloczyn skalarny zapytania i klucza, odpowiedź zależy od odległości między nimi, a nie tylko od tego, co zawierają. Oto trzy tokeny w pozycjach 0, 1 i 2, każdy przedstawiony jako strzała 2D obrócona nieco bardziej niż poprzednia:

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"

W oryginalnym rysunku trzeci strzałek nie jest obracany bardziej niż drugi, więc należy traktować diagram jako przybliżony obraz, a nie dokładny.

Mechanizm Attention porównuje strzałkę zapytania ze strzałką klucza za pomocą iloczynu skalarnego. Strzałki o tym samym kącie dają duży iloczyn, natomiast strzałki o bardzo różnych kątach dają mały iloczyn. Dlatego ten iloczyn odzwierciedla odległość między dwoma tokenami, przy czym żaden numer pozycji nie jest przechowywany w żadnym miejscu. RoPE nie dodaje pozycji jako odrębnego sygnału. Obrót wprowadza informację o pozycji do każdego wektora.

Qwen3 ustawia THETA = 1,000,000 – wysoką bazową częstotliwość, która pozwala modelowi radzić sobie z bardzo długimi kontekstami, nawet do 262K tokenów bez żadnych rozszerzeń.

Krok 5: Attention z grupowanymi głowami zapytania

To jest istota całego rozwiązania. Zrozumienie mechanizmu Attention oznacza zrozumienie modelu transformer.

Od MHA do MQA do 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 wykorzystuje GQA, posiadając 16 głowic zapytań oraz zaledwie 8 głowic klucz-wartość, dzięki czemu każda głowica klucz-wartość obsługuje dwie głowice zapytań. To zmniejsza rozmiar pamięci cache KV o połowę przy niewielkiej utracie jakości.

Cała funkcja uwagi

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

Kilka elementów tego wymaga wyjaśnienia.

QK-Norm. Qwen3 stosuje RMSNorm do wektorów Q i K każdej głowy przed procesem RoPE – to nowszy trik, który ułatwia trening. Jego wagi q_norm i k_norm składają się z 128 liczb, co odpowiada skali dla każdego wymiaru w obrębie danej głowy.

Bufor KV. Jest to najważniejsza optymalizacja podczas wykonywania modelu. Podczas fazy prefill obliczane są wartości K i V dla każdego tokena wejściowego i są one przechowywane. Podczas dekodowania obliczane są one tylko dla nowego tokena; następnie są dodawane do bufora, a mechanizm uwagi korzysta z całej historii tych wartości. Pominięcie bufora zmusiłoby do ponownego obliczania K i V dla każdego poprzedniego tokena, co zamieniłoby pracę liniową w kwadratową.

Maska przyczynowa. Podczas wypełniania wstępnego token na pozycji 3 nie może widzieć tokena na pozycji 4, który znajduje się w jego przyszłości. Kod dodaje -1e9 do górnej trójkąta macierzy ocen, aby to zagwarantować. Podczas dekodowania, gdy długość sekwencji wynosi 1, nie potrzebna jest żadna maska, ponieważ pojedynczy nowy token może patrzeć na wszystko, co znajduje się przed nim.

Wzór. Mechanizm uwagi oblicza funkcję softmax z wartości (Q pomnożone przez transponowaną K, podzielone przez pierwiastek kwadratowy wymiaru głowy), a następnie mnoży wynik przez V. Iloczyn skalarny mierzy, jak dobrze każda zapytanie pasuje do każdego klucza. Dzielenie przez pierwiastek kwadratowy wielkości głowy zapobiega tak dużemu wzrostowi ocen, który mogłby spowodować nasycenie funkcji softmax. Funkcja softmax przekształca oceny w prawdopodobieństwa, a mnożenie przez V daje ważoną średnią wektorów wartości.

Krok 6: Warstwa feed-forward SwiGLU

Po przetworzeniu każdy token przechodzi przez sieć typu feed-forward. Współczesne modele używają SwiGLU zamiast starszej wersji opartej na ReLU. Mówiąc prościej: z wejścia obliczane są dwie projekcje, jedna przechodzi przez funkcję SiLU i służy jako brama, obie są mnożone element po elemencie, a końcowa projekcja przywraca wynik do pierwotnego stanu.

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)

Sieć feed-forward to miejsce, w którym model przechowuje większość swojej wiedzy, co jest jednym z powodów, dla których duże, gęste modele wydają się tak skuteczne. Rozszerza wektor o 1024 wymiarów do 3072 wymiarów, stosuje bramkę SiLU, a następnie projektuje wynik z powrotem do 1024 wymiarów. Bramka ta uczy się, jakie informacje przepuszczać, a jakie blokować, nieco jak ochroniarz dla neuronów. Rozszerzanie z 1024 na 3072 i z powrotem daje modelowi możliwość przechowywania i przetwarzania informacji, a właśnie mechanizm bramek sprawia, że nowoczesne sieci feed-forward przewyższają te oparte na ReLU.

Krok 7: Blok transformera

Blok ten łączy wszystko powyższe:

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

Sekwencja jest normalizowana, następnie dodawany jest resztkowy sygnał, ponownie normalizowana, a potem przekazywana dalej, po czym ponownie dodawany jest resztkowy sygnał.

Te połączenia resztkowe (x = x + ...) są niezbędne. Umożliwiają one przepływ gradientów prosto przez sieć podczas treningu, a także pozwalają informacjom ominąć poszczególne bloki podczas inferencji. Model 28-warstwowy bez nich byłby praktycznie nietrenowalny.

Qwen3-0.6B składa się z 28 takich bloków, z których każdy nieco bardziej udoskonala reprezentację danych. Po 28. warstwie stan ukryty zawiera już wystarczająco dużo informacji o wejściu, aby przewidzieć następny token.

Krok 8: Część LM i próbkowanie

Po ostatnim bloku następuje jeszcze jedna normalizacja, a potem projekcja na rozmiar słownika:

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

Wynikiem jest 151 936 wartości, zwanych logitami, po jednej dla każdego możliwego tokena. Wyższa wartość oznacza bardziej prawdopodobny token. Dla zapytania „Stolica Francji to” pięć najwyższych prawdopodobieństw po zastosowaniu funkcji softmax wygląda w ten sposób:

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

„Paris” wygrywa z 65,8%, co jest pewną prognozą. Następnie proces próbkowania przekształca te prawdopodobieństwa w jeden identyfikator tokena:

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])

Temperatura kontroluje losowość. Przy wartości 0 token o najwyższej prawdopodobieństwie zawsze wygrywa (metoda chciwa). Przy 1,0 tokeny są wybierane proporcjonalnie do ich prawdopodobieństw, a przy 0,7 rozkład nieco się wyostrza na korzyść tokenów o najwyższych szansach.

Top-p, czyli próbkowanie jądrowe, zachowuje najmniejszą grupę tokenów, których kumulatywne prawdopodobieństwo osiąga wartość p, na przykład 0,95. Pozwala to usunąć mniej prawdopodobne tokeny z długiego „ogona”, zachowując jednocześnie pewną różnorodność.

Krok 9: Pętla generowania

Pełna pętla najpierw dokonuje wstępnego wypełnienia, a następnie koduje:

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)

To jest cały cykl inferencji. Silniki serwujące takie jak vLLM, SGLang, TensorRT-LLM i mlx-lm to w istocie optymalizacje tego rdzenia, a sam rdzeń mieści się w około 198 liniach kodu.

Co pomija skrypt

Skrypt obejmuje cały proces inferencji standardowego transformera. Nowsze modele dodają jeszcze więcej elementów:

• Hybrydowa uwaga. Podobno Qwen3.5 zastępuje 75% swoich warstw uwagi modelem Gated DeltaNet – wariantem uwagi liniowej z stałym stanem rekurencyjnym zamiast pełnej macierzy uwagi. Pełna uwaga przypomina pamięć fotograficzną: doskonała, ale kosztowna. DeltaNet przypomina skompresowane notatki – efektywny, ale z niewielkimi stratami. Hybryda zapewnia efektywne obsługiwanie długich kontekstów w połączeniu z możliwością dokładnego przywoływania informacji w określonych momentach.
• Przewidywanie wielu tokenów. Niektóre modele przewidują kilka tokenów na każdą iterację, oferując niewielką stratę jakości w zamian za znaczne przyspieszenie.
• Kwantyzacja. Proces inferencji wykorzystuje wagi 4-bitowe lub 8-bitowe w celu zaoszczędzenia pamięci i czasu. Ten skrypt zachowuje oryginalne wagi typu bfloat16.
• Pracowanie w partiach. Serwery obsługują wiele żądań jednocześnie i dzielą się obliczeniami między nimi. Ten skrypt obsługuje tylko jedno żądanie na raz.

Żadna z tych metod nie zmienia podstawowego algorytmu. Wszystkie opierają się na tej samej bazie.

Dlaczego to ważne?

Znajomość sposobu działania inferencji czyni cię lepszym praktykiem. Ostrzeżenie „pamięć KV pełna” od serwera przestaje być niejasne. Argument dotyczący zużycia pamięci w przypadku GQA staje się łatwy do wyjaśnienia. A model, który mówi bezładnie i powtarza się, wskazuje na konieczność użycia parametrów temperatury i top-p.

Współczesny LLM to nie magia. To 28 warstw operacji normalizacji, uwagi i przekazywania sygnału, powtarzanych wielokrotnie, z uwzględnieniem informacji o pozycji poprzez rotację, a wiedza przechowywana jest w warstwach liniowych z mechanizmami kontrolnymi. Matematyka wymagana to poziom studiów licencjackich. To inżynieria stanowi źródło pomysłowości, ale matematyka jest fundamentem. Wystarczy około 200 linijek kodu w Pythonie, rzeczywiste wagi oraz rzeczywisty wynik.

Zrób to sam

Wszystko powyżej można uruchomić na dowolnym komputerze z procesorem Apple Silicon (M1 lub nowszym).

1. Pobierz model

Pobierz Qwen3-0.6B z Hugging Face (repositorium Qwen/Qwen3-0.6B). Potrzebne są tylko cztery pliki, więc pomijaj pozostałe:

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

Umieść folder z modelem obok pliku llm_inference.py, ponieważ skrypt odnajduje ścieżkę ./Qwen3-0.6B/ w zależności od katalogu, z którego został uruchomiony.

2. Stwórz środowisko wirtualne

# 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

Brak PyTorch, brak transformerów, brak accelerate, brak pakietu safetensors. MLX zajmuje się obliczeniami na macierzach i ładowaniem wag, natomiast biblioteka tokenizerów realizuje kodowanie i dekodowanie BPE.

3. Uruchom to

python llm_inference.py

Oczekuj wyniku w następującym kształcie:

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

Jeśli zobaczysz ModuleNotFoundError: No module named 'mlx', sprawdź, czy środowisko venv jest aktywne (source llm_venv/bin/activate) oraz czy komputer ma procesor Apple Silicon. MLX nie działa na komputerach Mac z procesorami Intel, systemie Linux ani Windows.

Wszystko to działa na laptopie, a odpowiedź pojawia się w mniej niż sekundę.