Унутры LLM: інферэнс Qwen3-0.6B на чыстым Python і MLX
Пакрокава будуем інферэнс LLM прыкладна ў 200 радках Python з MLX: токенізатар, RMSNorm, RoPE, grouped query attention, SwiGLU, KV-кэш і сэмплаванне.
Крыніцы: How Modern LLMs Actually Generate Text: A Pure Python Walkthrough
Уводзіце «Століця Франціі» — і мадэль за апошніяя 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: Эмбеддынг ператварае цэлыкі лічыбельныя на вектары
Першая рэальная вычысленне прысвячае кожны ID токэна 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
Кожны вектар масштабуецца па сваёй сабе велічыне, а пасля корэгуецца за дапамою вагоў, выучаных для кожнай характэрыстыкі. Не выконваецца вырахунак сярэдньяго значэння і няма падкорэкцыі, што робіць гэты падход простым, не паслабляючы якасці. 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"
У первасправнай схеме трэція стрэлка не роціруецца больш, чым другая, таму спрыяйце схему як прыблізныя зображэння, а не як точныя.
Механізм Attention пораўнюе стрэлку запіту з стрэлкай ключа за дапамой скалярнага добутку. Стрэлкі, якія знаходзяцца пад аднаковым кутам, даюць вялікі добутак, а стрэлкі пад дужа разнымі кутамі — маленькі. Такім чынам добутак відбівае адстань між двума токенамі, пры чым ніякія явныя номеры пазіцый не зберагаюцца. RoPE не дадае пазіцыю як окалечны сігнал; роціяванне включае пазіцыю ў кожны вектар.
Qwen3 задае THETA = 1,000,000 — высокую базовую частоту, якая дазволяе модэлю справляцца з вельмі дзяўжымі кантэкстамі, аж до 262K токенав без дадатковых расшырэнняў.
Шаг 5: Attention з групаванымі галавамі запіту
Гэта ключовая частка. Якшы вы разумееце механізм Attention, вы разумееце сам трансфармер.
З 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 галавак ключ-значэння, таму кожная галавка ключ-значэння служыць двум галавкам запыткаў. эта падвайнае кеш KV з мінімальныя втратамі ў якасці.
Вся функцыя адзірвання
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, маска не патрэбна, адтуды ўнічыжны новы токен можа бачыць усё, што знаходзіцца раней.
Формула. Функцыя Attention вычысляе 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 супаняецца з стыснутымі нотаткамі: эфектывная, але з незначнымі втратамі. Гібрыдны падход дазволяе эфектывна працаваць з дзяўжынным кантэкстам і адночасна атрымваць ідэальныя рэзультаты.
- Прагнозаванне калькольнікаў. Дзеякія мадэлі прагнозуюць калькольнікі ў колькасці калькольнікаў за адну перадачу, жэртвуючы часткай якосці зарады значнага прышвартавання.
Ніякі з этых методаў не зміняе основнага алгорытму. Яны працуюць на той самый фундамент.
За што гэта трэба?
Розумеўшы, як функцыонуе выканання задач, вы станавіцеся лепшым спецыялістам. Паведамленні сервера пра „праўераную кеш-памяць KV“ больш не будуць загадкамі. Аргументы, стосаваныя памяці для GQA, станавяцца простымі да пояснення. А мадэль, якая прагалодзіцца і павтарае ся, паказвае направленне для налаштавання параметраў temperature і top-p.
Сучасны ЛЛМ не ёсць чарамі. Гэта 28 слоёў процэсаў нормалізацыі, атэнцыі і перэдачы дадзенняў, якія павтараюцца, пры чым пазіцыі дадзенняў урахоўваныя за дапамогою ротацыі, а знанні зберагваныя ў слоях лінейных функцый з кантролюемым выходам. Матэматыка на рэвэле бакалавра. Самае важнае — гэта інжынерыя, але матэматыка являецца фундаментам. Для гэтага патрэбна проста апошнё 200 рэйкаў коду на Python, справжнія вагі і справжні выходны рэзультат.
Запускайце самі
Усё вышэзгаданае можна запусціць на будзь-яму камп’ютары з процэсорам 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', пераканаўцеся, што сэрвэр з актывацыям віртуальнай асяроддзя ўвімкнуты (source llm_venv/bin/activate) і што прыстрой падтрымлівае тэхнологію Apple Silicon. MLX не працюе на камп’ютарах з процэсорамі Intel, у сістэмах Linux чы Windows.
Усё гэта працюе на ноутбуку, і адпаведны адказ практычна за секунду.