Всередині LLM: чисте виконання на Python для Qwen3-0.6B на MLX

Вводиться типове запитання «Столиця Франції — це», і модель завершує його відповіддю «Париж» приблизно за 900 мілісекунд. Що відбувається протягом цього короткого проміжку часу? Яка послідовність перетворень перетворює п’ять слів на прогноз?

Найкращий спосіб дізнатися про це — самостійно створити весь процес інференсу замість того, щоб використовувати бібліотеку, яка приховує деталі. У наведеному нижче огляді це робиться за допомогою справжніх ваг Qwen3-0.6B, завантажених у звичайний Python через MLX — фреймворк масивів від Apple. Повний скрипт складається приблизно з 200 рядків та працює на будь-якому Mac з процесором Apple Silicon. Інструкції з налаштування наведені в кінці.

Огляд процесу

Qwen3, GPT-4, Llama та їхні аналоги всі працюють за одним спільним алгоритмом:

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

Текст потрапляє всередину, перетворюється на числа, проходить через 28 блоків трансформаторів та виходить у вигляді розподілу ймовірностей між 151 936 можливими наступними токенами. Вибирається один токен та додається до результату, після чого цикл починається знову.

Цей цикл складається з двох фаз:

• Попереднє заповнення. Усі вхідні токени обробляються одночасно. Модель читає запит та формує свій внутрішній стан.
• Декодування. Токени генеруються по одному. Новий токен має враховувати лише попередні токени, ключі та значення яких зберігаються у кеші KV, тому нічого не обробляється заново.

Попереднє заповнення відбувається паралельно, а декодування — послідовно. Саме кеш KV забезпечує швидке декодування. Без нього для створення токена N довелося б заново обчислювати взаємодію з усіма токенами N-1, що передували йому.

Тепер розглянемо кожну з фаз по черзі.

Крок 1: Токенізація перетворює текст на цілі числа

Модель ніколи не бачить слів, лише цілих чисел.

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]

Кожне ціле число є індексом у словнику, що містить 151,936 токенів. Токенізатор використовує кодування пар байтів, тому межі токенів не завжди збігаються з межами слів. „capital“ — це один токен, а „France“ — інший, причому пробіл на початку належить до токену. Створення BPE з нуля тут не є цікавою частиною. Токенізатор — це таблиця пошуку, і скрипт просто використовує вже опубліковану.

Крок 2: Ембеддинг перетворює цілі числа на вектори

Перша справжня обробка даних — це прив’язка кожного ідентифікатора токену до 1,024-вимірного вектора.

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

Це просто пошук у таблиці: матриця вбудовування має форму (151936, 1024) — по одному рядку на кожен токен та по одній стовпцю на кожний прихований вимір, і код здійснює індексацію в неї. Уявіть собі цю матрицю як словник моделі: кожен токен має навчений 1024-значний «відбиток пальця», який кодує його значення. Відбитки пальців для „Париж“ та „Франція“ вказують у схожих напрямках, тоді як відбитки „Париж“ та „комп’ютер“ — у зовсім різних.

Крок 3: RMSNorm

Перед обробкою за допомогою механізму уваги приховані стани нормалізуються. Сучасні моделі використовують RMSNorm замість старішого LayerNorm.

Простіше кажучи, формула ділить кожен вектор на його квадратичне середнє значення — це квадратний корінь від середнього квадратів його елементів плюс дуже мала константа — а потім множить на навчений коефіцієнт для кожної ознаки.

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

Кожен вектор масштабується за допомогою своєї власної величини, а потім коригується за допомогою ваг, навчених для кожної ознаки окремо. Не використовується віднімання середнього значення та немає поправки на смещення, що робить цей підхід простішим за LayerNorm без втрати якості. EPS (1e-6) запобігає діленню на нуль, а вага w, навчена під час тренування, дозволяє моделі незалежно масштабувати кожен з 1 024 вимірів.

Крок 4: RoPE вбудовує інформацію про позицію у вектори

Саме тут багато статей про „LLM з нуля“ втрачають читачів, але насправді все простіше, ніж здається.

Проблема. Механізм уваги є інваріантним до перестановок. Якщо перетасувати вхідні токени, результат буде таким самим, тому модель не може визначити, який токен був спочатку.

Старий спосіб вирішення. Додати номер позиції до ембеддингу кожного токена. Це працює, але не є універсальним, і модель має труднощі з розумінням відносних відстаней.

Сучасний спосіб вирішення. Перетворити вектори запиту та ключа під кут, який залежить від позиції токена. Це і є Rotary Position Embedding, або RoPE. Кожна пара вимірів розглядається як точка в площині, яка повертається під кут, визначений позицією токена.

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)

Уявіть кожну пару вимірів як 2D-площину. Токени, розташовані в різних позиціях, обертаються на різну кількість, тому коли механізм уваги обчислює скалярний добуток запиту та ключа, результат залежить від відстані між ними, а не лише від того, що вони містять. Ось три токени на позиціях 0, 1 та 2, кожен з яких зображений як 2D-стрілка, обернута ще більше, ніж попередня:

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"

У початковому малюнку третя стрілка не повернута далі, ніж друга, тому сприймайте діаграму як приблизний зображення, а не точний варіант.

Механізм уваги порівнює стрілку запиту зі стрілкою ключа за допомогою скалярного добутку. Стрілки, що знаходяться під однаковим кутом, дають великий добуток, а стрілки з дуже різними кутами — малий. Отже, добуток відображає відстань між двома токенами, причому жоден номер позиції не зберігається окремо. RoPE не додає позицію як окремий сигнал. Поворот вбудовує інформацію про позицію в кожен вектор.

Qwen3 встановлює THETA = 1,000,000 — високу базову частоту, яка дозволяє моделі працювати з дуже довгими контекстами, аж до 262 тисяч токенів без додаткових засобів.

Крок 5: Увага з групованими головами запиту

Це ключова частина. Якщо ви зрозумієте механізм уваги, ви зрозумієте принцип роботи трансформера.

Від MHA до MQA до 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 використовує GQA, має 16 голів запитів та лише 8 голів ключ-значення, тож кожна голова ключ-значення обслуговує дві голови запитів. Це скорочує об’єм кешу ключ-значення вдвічі з мінімальними втратами якості.

Вся функція уваги

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

Деякі частини цього потребують пояснення.

QK-Norm. Qwen3 застосовує RMSNorm до векторів Q та K кожної голови перед процесом RoPE — це новий підхід, який стабілізує процес навчання. Його ваги q_norm та k_norm складаються з 128 чисел, що відповідає масштабу кожного виміру в межах голови.

Кеш KV. Це найважливіша оптимізація під час інференсу. Під час попереднього заповнення кешу обчислюються та зберігаються вектори K та V для кожного вхідного токена. Під час декодування обчислюються K та V лише для нового токена; вони додаються до кешу, і процес уваги працює з усією історією. Якби кеш було пропущено, довелося б заново обчислювати K та V для кожного попереднього токена, що перетворює лінійну роботу на квадратичну.

Маска причинності. Під час попереднього заповнення токен на позиції 3 не повинен бачити токен на позиції 4, який знаходиться у його майбутньому. Для цього код додає -1e9 до верхньої трикутної частини матриці оцінок. Під час декодування, коли довжина послідовності дорівнює 1, маска не потрібна, оскільки єдиний новий токен може бачити все, що було до нього.

Формула. Механізм уваги обчислює softmax від значення (Q помножене на транспоноване K, поділене на квадратний корінь розміру головки), а потім множить отриманий результат на V. Скалярний добуток вимірює ступінь збігу кожного запиту з кожним ключем. Ділення на квадратний корінь розміру головки запобігає тому, щоб оцінки ставали настільки великими, що softmax досягає свого максимуму. Функція softmax перетворює оцінки на ймовірності, а множення на V дає зважену середню величину векторів значень.

Крок 6: Шар прямого передавання SwiGLU

Після обробки кожен токен проходить через мережу з прямим передаванням. Сучасні моделі використовують SwiGLU замість старішої версії, заснованої на ReLU. Простіше кажучи: з вхідних даних обчислюються два проекційні вектори, один з них проходить через функцію SiLU та використовується як «брама», ці два вектори множаться елемент за елементом, а останній проекційний вектор повертає результат назад.

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)

Мережа з прямим передаванням є тією частиною, де модель зберігає більшу частину своїх знань, що є однією з причин такої високої ефективності великих щільних моделей. Вона розширює 1024-вимірний вектор до 3072 вимірів, застосовує білок SiLU та проєктує його назад до 1024 вимірів. Цей білок вчиться, яку інформацію допускати, а яку блокувати, щось на кшталт охоронця для нейронів. Розширення від 1024 до 3072 та назад надає моделі можливість зберігати та перетворювати інформацію, а саме білки є причиною того, чому сучасні мережі з прямим передаванням перевершують моделі з ReLU.

Крок 7: Блок трансформера

Блок поєднує все вищезазначене:

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

Послідовність нормалізується, додається резидуальний сигнал, знову нормалізується, виконується передавання даних вперед, і знову додається резидуальний сигнал.

Ці резидуальні з’єднання (x = x + ...) є життєво важливими. Вони дозволяють градієнтам проходити прямо через мережу під час навчання, а також дозволяють інформації оминати окремі блоки під час інтерпретації. Модель з 28 шарами без них була б майже ненавчуваною.

Qwen3-0.6B складається з 28 таких блоків, кожен з яких ще більше удосконалює представлення даних. До 28-го шару прихований стан містить достатньо інформації про вхідні дані для прогнозування наступного токена.

Крок 8: Голова LM та вибір зразків

Після останнього блоку відбувається ще одна нормалізація, а потім проєкція на розмір словника:

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

Результатом є 151 936 оцінок, які називаються логітами — по одній для кожного можливого токена. Вища оцінка означає більш ймовірний токен. Для фрази „Столиця Франції —“ п’ять найвищих ймовірностей після застосування алгоритму softmax виглядають так:

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

„Париж“ перемагає з 65,8% — це впевнений прогноз. Потім процес вибірки перетворює ймовірності на єдиний ідентифікатор токена:

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

Температура контролює випадковість. При значенні 0 токен із найвищою ймовірністю завжди перемагає (метод жадібності). При значенні 1,0 токени обираються пропорційно їхнім ймовірностям, а при значенні 0,7 розподіл трохи загострюється у бік лідерів.

Top-p, або вибірка за ядром, зберігає найменшу групу токенів, чия кумулятивна ймовірність досягає значення p, наприклад 0,95. Цей метод скорочує довгий хвіст малоймовірних токенів, зберігаючи при цьому певну різноманітність.

Крок 9: Цикл генерації

Повний цикл спочатку заповнює дані один раз, а потім здійснює декодування:

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)

Це і є весь цикл інференсу. Мотори обробки даних, такі як vLLM, SGLang, TensorRT-LLM та mlx-lm, по суті є оптимізаціями цього ядра, причому саме ядро складається приблизно з 198 рядків коду.

Що пропускає скрипт

Скрипт охоплює повний процес інференсу стандартного трансформера. Новіші моделі додають ще більше елементів:

• Гібридна увага. За повідомленнями, Qwen3.5 замінює 75% своїх шарів уваги на Gated DeltaNet — варіант лінійної уваги з постійним рекурентним станом замість повної матриці уваги. Повна увага схожа на фотографічну пам’ять: ідеальна, але дорога. DeltaNet схожа на стиснуті нотатки: ефективна, але з невеликими втратами. Гібридний підхід забезпечує ефективну обробку довгих контекстів разом із періодичним ідеальним відтворенням інформації.
• Прогнозування кількох токенів. Деякі моделі прогнозують кілька токенів за один прохід, жертвуючи невеликою якістю заради значного прискорення.
• Квантація. Для економії пам’яті та часу під час інференсу використовуються ваги 4-бітної або 8-бітної точності. Цей скрипт зберігає оригінальні ваги формату bfloat16.
• Батчування. Сервери обробляють багато запитів одночасно та діляться обчисленнями між ними. Цей скрипт обробляє лише один запит за раз.

Жоден з цих факторів не змінює основного алгоритму. Вони працюють на основі тієї самої структури.

Навіщо це потрібно?

Розуміння принципів інференсу робить вас кращим фахівцем. Попередження сервера про „повну кеш-пам’ять KV“ більше не є загадковим. Аргументи щодо використання пам’яті у моделі GQA стають простими для пояснення. А якщо модель говорить безладно та повторюється, це вказує на необхідність налаштувань температури та параметра top-p.

Сучасний LLM — це не магія. Це 28 шарів процедур нормалізації, уваги та передавання даних, які повторюються, з врахуванням позиції через ротацію, а знання зберігаються у шарах лінійних функцій із контролем доступу. Математика тут на рівні бакалаврських курсів. Саме інженерія є джерелом кмітливості, але математика є основою. Для цього потрібно лише близько 200 рядків коду на Python, справжні ваги та справжній результат.

Запустіть самі

Усе вищезазначене можна запустити на будь-якому Mac з процесором Apple Silicon (M1 або новішим).

1. Завантажте модель

Завантажте Qwen3-0.6B з Hugging Face (репозиторій Qwen/Qwen3-0.6B). Потрібні лише чотири файли, тож можна пропустити решту:

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

Розмістіть папку з моделлю поруч із llm_inference.py, оскільки скрипт шукає ./Qwen3-0.6B/ у тій директорії, з якої він запущений.

2. Створіть віртуальне середовище

# 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

Жодного PyTorch, жодних transformers, жодного accelerate, жодного пакету safetensors. MLX виконує обчислення з масивами та завантажує ваги, а бібліотека tokenizers здійснює кодування та декодування за алгоритмом BPE.

3. Запустіть його

python llm_inference.py

Очікуйте результат у такому форматі:

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

Якщо ви бачите ModuleNotFoundError: No module named 'mlx', переконайтеся, що середовище venv активне (source llm_venv/bin/activate) та що пристрій має процесор Apple Silicon. MLX не працює на комп’ютерах з процесорами Intel, Linux чи Windows.

Усе це працює на ноутбуці, і відповідь з’являється менш ніж за секунду.