10 октября 2026Команда Сарнас

BitNet на микроконтроллерах: как устроен запуск 1.58-битной LLM на кластере ESP32-S3

Модели с экстремально низкой разрядностью всё чаще рассматривают для устройств без GPU, но их реальная работа упирается в то, как именно ядра выполняют умножение матриц. В статье разбирается внутреннее устройство инференса BitNet b1.58 на NEON: почему «1.58 бита» превращаются в таблицу подстановки, как устроено ядро qgemm_lut и как модель распределяется по кластеру ESP32-S3. Это полезно тем, кто хочет понять, где заканчивается маркетинг низкой разрядности и начинается конкретная арифметика.

Разбираем механику BitNet b1.58 на уровне ядер: как из 1.58-битных весов получается LUT, как ядро qgemm_lut считает умножение матриц на NEON, как устроены жизненный цикл и рабочие буферы на устройстве. Вопрос нетривиален тем, что «1.58 бита» — это не формат хранения, а результат квантования, и в коде он превращается в таблицу заранее посчитанных сумм, а не в прямое умножение.

Типы тензоров и битовые ширины

Функция ggml_bitnet_get_type_bits возвращает битовую ширину по типу тензора: для GGML_TYPE_TL1 и GGML_TYPE_TL2 — 2 бита, для GGML_TYPE_Q4_0 — 4 бита, для всех остальных типов — 0. В исходниках есть две реализации с одинаковым именем: в одной обрабатывается GGML_TYPE_TL1, в другой — GGML_TYPE_TL2, обе возвращают 2 для своего TL-типа и 4 для Q4_0. Полученное значение bits передаётся в ggml_bitnet_mul_mat_task_init и ggml_bitnet_mul_mat_task_compute, то есть определяет разрядность при построении LUT и вычислении.

Формат i2_s упоминается только как имя файла модели в команде запуска инференса.

Архитектура кластера и распределение слоёв

Узел берётся за умножение матриц для пары (src0, src1) только при одновременном выполнении трёх условий: тип src0 должен относиться к поддерживаемым, тип src1 должен быть GGML_TYPE_F32, а тип dst — тоже GGML_TYPE_F32. Если хотя бы одно из условий не выполнено, умножение на этом узле не выполняется. В варианте для ARM (GGML_BITNET_ARM_TL1) дополнительно требуется, чтобы src1->ne[1] было не больше 1, то есть вторая размерность матрицы src1 не превышала единицу.

Поддерживаемыми считаются только типы GGML_TYPE_Q4_0 и GGML_TYPE_TL1. Размеры матриц в самой проверке применимости не участвуют, кроме указанного ограничения src1->ne[1] <= 1 в ARM-варианте; размеры ne10, ne01, ne11 используются только в ggml_bitnet_mul_mat_get_wsize для вычисления размера рабочего буфера.

LUT-компиляция весов: как из 1.58-бит получается таблица

lut_ctor формирует таблицу подстановки для одного блока весов: сначала масштабирует веса b на scales и округляет их до int16, получая две половины вектора весов блока, уже приведённые к целым числам. Из них строятся 8 векторов vec_lut[0..7] как комбинации 0, ± первой половины, ± второй половины: таблица заранее перечисляет все возможные суммы вкладов двух половин блока, чтобы во время умножения оставалось только выбрать нужную запись по индексу активации. После этого Transpose_8_8 переставляет элементы двух групп по 8 векторов, а vqtbl1q_s8 с tbl_mask_q переупорядочивает байты, после чего результат сохраняется в qlut. v0 = q_fin_0.val[0] — это результат vzipq_s16(q0246_0.val[0], q1357_0.val[0]) внутри Transpose_8_8, то есть первый вектор после перестановки.

На один блок весов таблица содержит 16 записей: vec_lut[0..15], из которых первые 8 строятся явно, а вторые 8 заполняются через Transpose_8_8.

Ядро qgemm_lut: пошаговое выполнение

Цикл qgemm_lut_4096_4096

Функция qgemm_lut_4096_4096 сначала обнуляет буфер аккумулятора размером BM4096_4096 (128 элементов) через memset, затем во внешнем цикле по k_outer от 0 до 4096/BBK4096_4096 (BBK4096_4096=64, то есть 64 итерации) вызывает tbl_impl_4096_4096, передавая ей указатель на буфер аккумулятора, сдвиг в LUT на k_outer*BBK4096_4096/2*32 байт и сдвиг в A на k_outer*BBK4096_4096/2/2*BM4096_4096 байт.

Внутри tbl_impl_4096_4096 LUT читается в массив vec_lut[2*KK] (KK=BBK4096_4096/2=32) через vld1q_s8(lut + k*16), то есть по 16 байт на элемент. Для каждого блока i с шагом 32 по BM4096_4096 аккумуляторы vec_c[0..3] обнуляются через vandq_s16(vec_c[i], vec_zero), затем во внутреннем цикле по k от 0 до KK/4 читаются 16-байтовые фрагменты A, из них выделяются старшие полубайты (vshrq_n_u8) и младшие (vandq_u8 с маской 0x0f), по ним выполняется табличный поиск vqtbl1q_s8 по соответствующим vec_lut, результаты попарно объединяются vzipq_s8 и прибавляются к vec_c[0..3].

После завершения циклов значения vec_c расширяются из 16-битных в 32-битные через vmovl_s16/vmovl_high_s16 и прибавляются к c[i+0..i+28] через vst1q_s32. Наконец, результат для каждого i от 0 до BM4096_4096 записывается в C[i].

Размеры блоков и их влияние

Размер блока задаётся макросами: BM14336_4096=256, BM4096_14336=128, BM1024_4096=256, BM4096_4096=128. Внутренний цикл по строкам блока идёт с шагом 32, поэтому при BM=256 выполняется 8 итераций, а при BM=128 — 4 итерации. Число итераций по k одинаково во всех четырёх функциях и равно KK/8, где KK=BBK/3 для three_tbl и KK=BK2/2 для two_tbl; при BBK=96 это KK=32 и 4 итерации k.

Потребление памяти зависит от размера блока: в three_tbl_impl_14336_4096 и three_tbl_impl_1024_4096 (BM=256) объявляются массивы __m256i vec_as[KK / 2]; и __m256i vec_signs[KK / 8];, тогда как в three_tbl_impl_4096_14336 и three_tbl_impl_4096_4096 (BM=128) те же массивы объявляются при том же KK, но цикл по i выполняется вдвое меньше раз. В two_tbl-функциях массив только один: __m256i vec_as[KK / 2];, а vec_signs отсутствует, поэтому two_tbl потребляет меньше памяти, чем three_tbl при том же BM.

Preprocessor и подготовка весов

Функция preprocessor_k<K> (шаблон по длине K) выполняет подготовку весов B к виду, пригодному для LUT-умножения: сначала обнуляет масштаб через partial_max_reset, затем вычисляет масштаб квантования per_tensor_quant(K, LUT_Scales, B), и наконец строит саму таблицу подстановок lut_ctor<K>(QLUT, B, LUT_Scales). ggml_preprocessor(m, k, B, LUT_Scales, QLUT) — это диспетчер: по паре (m, k) он выбирает конкретную инстанциацию preprocessor_k с нужным K (4096 или 14336), например для m==14336 && k==4096 вызывается preprocessor_k<4096>.

Внутри per_tensor_quant по B вычисляется максимум абсолютных значений (через vabsq_f32 и vmaxq_f32 на NEON), и масштаб записывается как scales = 127 / vmaxvq_f32(temp_max) в *lut_scales. lut_ctor затем умножает каждый элемент b на этот масштаб (vmulq_n_f32), округляет к ближайшему целому (vcvtnq_s32_f32), сужает до int16 (vmovn_s32) и из полученных значений строит 8 комбинаций сумм/разностей vec_bs_0 и vec_bs_1 (например vec_lut[0] = -vec_bs_0 - vec_bs_1, vec_lut[6] = vec_bs_0 - vec_bs_1), которые после Transpose_8_8 и перестановки байтов через tbl_mask_q записываются в qlut.

Путь весов: исходные float B → масштабирование на 127/max|B| → округление до целых → построение таблицы сумм/разностей пар → упаковка в int8 LUT, который затем используется в qgemm_lut_14336_4096 через vqtbl1q_s8.

Инициализация и жизненный цикл на устройстве

При первом вызове ggml_bitnet_init проверяет флаг initialized и, если он уже установлен, сразу возвращается; иначе устанавливает initialized = true. Затем, если указатель bitnet_tensor_extras равен nullptr, выделяет массив структур bitnet_tensor_extra размером GGML_BITNET_MAX_NODES через new bitnet_tensor_extra[...], и сбрасывает индекс bitnet_tensor_extras_index в 0. Структура bitnet_tensor_extra хранит кэш квантованных весов и масштабов для тензора.

ggml_bitnet_free при установленном initialized сбрасывает его в false, освобождает массив через delete[] bitnet_tensor_extras и обнуляет указатель. Если initialized не установлен, free ничего не делает.

Барьер в ggml_bitnet_mul_mat

Барьер нужен, чтобы все потоки дождались завершения инициализации общих буферов, которую выполняет только нулевой поток. В ggml_bitnet_mul_mat поток с ith == 0 вызывает ggml_bitnet_mul_mat_task_init, заполняя qlut, lut_scales и lut_biases во wdata. Затем при nth > 1 выполняется ggml_barrier(params->threadpool), синхронизируя все потоки пула. После барьера каждый поток вызывает ggml_bitnet_mul_mat_task_compute, читая уже подготовленные общие данные. Таким образом синхронизируются потоки одного вызова mul_mat: нулевой (инициализатор) и остальные (вычислители).

Переупорядочивание весов

Переупорядочивание весов выполняется в конструкторах LUT: lut_ctor (tl1) и three_lut_ctor/two_lut_ctor (tl2), которые строят таблицу qlut из векторов b, масштабированных на scales.

В lut_ctor для ARM значения b умножаются на scales, округляются до int32, сужаются до int16, из них формируются 16 векторов vec_lut, затем Transpose_8_8 переставляет их, а vqtbl1q_s8 с tbl_mask_q перекладывает байты и результат пишется в qlut по смещениям k*16*8*2 + idx*16*2. В three_lut_ctor/two_lut_ctor для x86 то же самое делается через _mm256_i32gather_ps, _mm256_cvtps_epi32, Transpose_8_8, _mm256_packs_epi32, _mm256_shuffle_epi8 и запись в qlut по смещениям k*256 + g*32. Такая перестановка нужна, чтобы уложить заранее вычисленные суммы вкладов весов в порядке, который затем читается в three_tbl_impl_14336_4096 через _mm256_shuffle_epi8 по индексам активаций.

Память и выравнивание

aligned_malloc выделяет память с выравниванием 64 байта: на Windows через _aligned_malloc(size, 64), на других платформах через posix_memalign(&ptr, 64, size), а aligned_free освобождает её через _aligned_free на Windows и free в остальных случаях. Transpose_8_8 принимает восемь указателей на int16x8_t и переставляет элементы внутри этих 8-элементных векторов, используя только операции vzipq_s16 над уже загруженными регистрами, без собственных требований к выравниванию памяти. В lut_ctor массивы vec_lut[16] и tbl_mask[16] — локальные, а загрузки из b выполняются через vld2q_f32, который не требует выравнивания; запись qlut идёт через vst1_s8.

Размер рабочего буфера

ggml_bitnet_mul_mat_get_wsize вычисляет размер рабочего буфера по размерам матриц src0 и src1: берёт ne01 = src0->ne[1], ne10 = src1->ne[0], ne11 = src1->ne[1]. В ARM-варианте (GGML_BITNET_ARM_TL1) она дополнительно получает bits = ggml_bitnet_get_type_bits(src0->type) и считает:

wsize = ne10 * ne11 * 15 * sizeof(int8_t) + 1 * ne11 * 2 * sizeof(bitnet_float_type);

Если sizeof(bitnet_float_type) == 2, добавляется std::max(ne10, ne01) * ne11 * sizeof(bitnet_float_type). В x86-варианте (GGML_BITNET_X86_TL2) bits не запрашивается, а wsize = ne10 * ne11 * 11 * sizeof(int8_t) + 2 * ne11 * 2 * sizeof(bitnet_float_type), и при sizeof(bitnet_float_type) == 2 добавляется std::max(ne10, ne01) * ne11 * sizeof(bitnet_float_type). В обоих вариантах результат выравнивается вверх до кратного 64: wsize = ((wsize - 1) / 64 + 1) * 64.

bitnet_float_type

Что видит клиент

run_inference.py запускается локально и принимает аргументы командной строки: -m MODEL (путь к файлу модели), -n N_PREDICT (число предсказываемых токенов), -p PROMPT (запрос для генерации), -t THREADS (число потоков), -c CTX_SIZE (размер контекста запроса), -temp TEMPERATURE (температура, управляющая случайностью генерации) и -cnv (включение режима чата). Пример запуска:

python run_inference.py -m models/BitNet-b1.58-2B-4T/ggml-model-i2_s.gguf -p "You are a helpful assistant" -cnv

Компромиссы и почему сделано именно так

LUT-подход реализован через подготовку таблиц: ggml_bitnet_mul_mat_task_init заполняет qlut, lut_scales и lut_biases, а затем ggml_bitnet_mul_mat_task_compute использует их вместе с qweights и scales. Размер рабочей памяти под LUT вычисляется в ggml_bitnet_mul_mat_get_wsize как ne10 * ne11 * 15 * sizeof(int8_t) плюс масштабы и, при 2-байтовом bitnet_float_type, дополнительный буфер. Поддержка ограничена типами GGML_TYPE_TL1 и GGML_TYPE_Q4_0, для которых ggml_bitnet_get_type_bits возвращает 2 и 4 бита соответственно. Условие применимости задано в ggml_bitnet_can_mul_mat: тип src0 поддерживается, src1 и dst имеют тип F32.

Разные размеры блоков

Для пары размерностей 14336×4096 задано BM14336_4096 (число строк результата, обрабатываемых за один вызов), а для 4096×14336 — BM4096_14336; в qgemm_lut_14336_4096 буфер CBits объявляется как alignas(32) uint32_t CBits[BM14336_4096], а в qgemm_lut_4096_14336 — как alignas(32) uint32_t CBits[BM4096_14336]. Аналогично для 3200×8640 используется BM3200_8640, для 3200×3200 — BM3200_3200, для 1024×4096 — BM1024_4096, для 4096×4096 — BM4096_4096.

Размер блока по K задаётся отдельно через BBK: например, BBK3200_8640=64, BBK3200_3200=128, BBK4096_14336=128, BBK1024_4096=64, BBK4096_4096=64, и число внешних итераций вычисляется как 8640 / BBK3200_8640, 4096 / BBK14336_4096, 14336 / BBK4096_14336, 4096 / BBK1024_4096, 4096 / BBK4096_4096. Внутри tbl_impl_* размер BM задаёт шаг внешнего цикла по i (например, for (int i = 0; i < BM3200_8640; i += 32) и for (int i = 0; i < BM3200_3200; i += 64)), а BBK — число итераций внутреннего цикла по k (KK = BBK/2, затем k < KK/4 или k < KK/2).

Как сравнивали и что получилось

Для запуска бенчмарка инференса используется скрипт e2e_benchmark.py. Он принимает обязательный аргумент -m/--model — путь к файлу модели, и необязательные: -n/--n-token (число генерируемых токенов, по умолчанию 128), -p/--n-prompt (число токенов промпта, по умолчанию 512), -t/--threads (число потоков, по умолчанию 2), -h/--help (показать справку и выйти).

Пример запуска: python utils/e2e_benchmark.py -m /path/to/model -n 200 -p 256 -t 4 — это запускает бенчмарк инференса с моделью по пути /path/to/model, генерируя 200 токенов из промпта в 256 токенов с использованием 4 потоков. Для раскладок модели, не поддерживаемых публичными моделями, предоставляются скрипты для генерации фиктивной модели с заданной раскладкой и запуска бенчмарка на вашей машине. Упоминается демонстрация bitnet.cpp с моделью BitNet b1.58 3B на Apple M2.

Сценарии нагрузки

Пример запуска python utils/e2e_benchmark.py -m /path/to/model -n 200 -p 256 -t 4 задаёт генерацию 200 токенов, промпт из 256 токенов и 4 потока. По умолчанию -n равен 128, -p равен 512, -t равен 2.

Воспроизводимость

Чтобы повторить замер, сначала подготовьте окружение скриптом setup_env.py, который принимает параметры --hf-repo (выбор модели для инференса), --model-dir (каталог для сохранения/загрузки модели), --log-dir (каталог для логов), --quant-type (тип квантования), --quant-embd (квантовать эмбеддинги в f16) и --use-pretuned (использовать предварительно настроенные параметры ядра). Затем сгенерируйте тестовую модель командой generate-dummy-bitnet-model.py, указав каталог модели, выходной файл, тип выходных данных и размер модели. После этого запустите замер через e2e_benchmark.py, передав путь к модели флагом -m, число генерируемых токенов флагом -n, размер промпта флагом -p и число потоков флагом -t:

python utils/e2e_benchmark.py -m models/dummy-bitnet-125m.tl1.gguf -p 512 -n 128

Порядок именно такой: setup_env.py, затем generate-dummy-bitnet-model.py, затем e2e_benchmark.py.

Что из этого следует на практике

Умножение матриц в BitNet на ARM — это не арифметика над 1.58-битными числами, а выборка из таблицы: vqtbl1q_s8 по индексам активаций. Поэтому стоимость определяется не разрядностью весов, а объёмом LUT и числом итераций внутренних циклов.
Применимость ядра жёстко ограничена типами: is_type_supported пропускает только GGML_TYPE_Q4_0 и GGML_TYPE_TL1, а src1 и dst обязаны быть F32. В ARM-варианте добавляется ограничение src1->ne[1] <= 1.
Размер рабочего буфера считается по размерам матриц и округляется вверх до кратного 64; при 2-байтовом bitnet_float_type в него добавляется буфер std::max(ne10, ne01) * ne11 * sizeof(bitnet_float_type). Значений по умолчанию при нехватке памяти нет, поэтому запас нужно закладывать самому.
Инициализация общих буферов выполняется только нулевым потоком, а барьер ставится лишь при nth > 1. При одном потоке синхронизация не нужна.
Размеры блоков BM и BBK заданы отдельно для каждой пары размерностей и определяют число итераций циклов; в three_tbl-функциях при том же BM объявляется на один массив больше (vec_signs), чем в two_tbl, что влияет на потребление памяти.
Воспроизводимость бенчмарка опирается на три скрипта в фиксированном порядке; сами цифры производительности на ESP32-S3 не приводятся, а демонстрация bitnet.cpp упоминается только на Apple M2.