Внутри LLM: инференс на чистом Python для Qwen3-0.6B на MLX

Вводится текст «Столица Франции — » и модель за примерно 900 миллисекунд заполняет его словом «Париж». Что происходит за это краткое время? Какая последовательность преобразований превращает пять слов в прогноз?

Лучший способ узнать это — самостоятельно создать весь процесс инференса вместо того, чтобы использовать библиотеку, скрывающую детали. Ниже приведено пошаговое руководство с использованием реальных весов Qwen3-0.6B, загруженных в обычный Python через MLX — фреймворк для массивов от Apple. Полный скрипт состоит примерно из 200 строк и работает на любом компьютере с процессором 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: Эмбеддинг преобразует целые числа в векторы

Первая реальная операция заключается в преобразовании каждого идентификатора токена в 1024-мерный вектор.

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 перед применением алгоритма RoPE выполняется операция RMSNorm над векторами Q и K каждого блока, что является более современным методом, способствующим стабильности обучения. Их веса 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, настоящих весов и реального вывода.

Запустите сами

Всё, что описано выше, можно запустить на любом компьютере Apple Silicon Mac (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 не работает на Mac с процессорами Intel, в Linux или Windows.

Всё это работает на ноутбуке, и ответ появляется менее чем за секунду.