Пишем inference engine на чистом C: разбор кода и бенчмарки для LFM2-350M

Зачем вообще это нужно? (Проблема)

Ты скачал LFM2-350M, запустил через какую-то обертку на Python и ждешь ответа три минуты. Вентилятор ноутбука воет, как реактивный двигатель. Ты открываешь htop и видишь, что процесс съел 4 ГБ RAM и 150% CPU. Знакомая картина?

Современные фреймворки для инференса — это монстры. Они тащат за собой десятки зависимостей, виртуальные окружения, слои абстракции. А все, что нам нужно — просто умножить матрицы и применить внимание. Быстро. На том железе, которое есть. Если твоя задача — запустить модель на слабом сервере или в embedded-среде, толстые обертки не подходят.

Чистый C, а не Rust или Python? (Выбор технологии)

Rust безопаснее. Python проще. C — быстрее и без накладных расходов. Мы пишем не production-систему, а исследовательский движок, где каждый цикл процессора на счету. Rust со своим borrow checker добавляет overhead в разработке (хотя R3-Engine доказывает, что на Rust тоже можно выжимать максимум). C дает полный контроль над памятью и вычислениями. И да, это больно. Но зато быстро.

Анатомия движка: от бинарного файла до токена

Весь движок умещается в одном файле на 1500 строк. Он делает три вещи: загружает веса, выполняет forward pass, выдает токены. Звучит просто? Сейчас увидишь, где собака зарыта.

1 Загрузка весов: бинарный формат против safetensors

Мы не будем парсить .safetensors. Конвертируем веса в простой бинарный формат: сначала заголовок с размерностями, потом данные float32 подряд. Чтение — один вызов fread. Никакого JSON, никаких потоков.

typedef struct {
    int dim;         // размерность эмбеддинга
    int hidden_dim;  // размерность скрытого слоя
    int n_layers;    // количество слоев
    int n_heads;     // количество голов внимания
    int vocab_size;  // размер словаря
} Config;

Config config;
fread(&config, sizeof(Config), 1, fp);
// Дальше читаем веса последовательно
float* token_embedding_table = malloc(config.vocab_size * config.dim * sizeof(float));
fread(token_embedding_table, sizeof(float), config.vocab_size * config.dim, fp);

2 Ядро вычислений: CBLAS GEMM и почему твой процессор не греется

90% времени инференса — матричные умножения. Мы используем CBLAS интерфейс для вызова оптимизированных под твой CPU функций. На macOS это Accelerate framework, на Linux — OpenBLAS или BLIS.

#include 
// Умножение матрицы на матрицу: C = alpha * A * B + beta * C
cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
            M, N, K,
            alpha,
            A, lda,
            B, ldb,
            beta,
            C, ldc);

Если просто вызывать sgemm, получишь ускорение в 5-10 раз по сравнению с наивной реализацией на циклах. Но можно выжать больше. Например, для небольших матриц (как в attention) overhead вызова CBLAS становится заметен. Здесь нужна ручная оптимизация с интринсиками. Но для LFM2-350M с размерностью 1024, CBLAS работает отлично.

3 RoPE реализация: где все ошибаются

Rotary Positional Embedding — это не просто синусы-косинусы. Нужно применять вращение к парах элементов в query и key. Частая ошибка — вычислять RoPE на лету для каждого токена. Мы предрассчитываем матрицу вращений для максимальной длины контекста (2048 для LFM2-350M) и кэшируем.

void apply_rope(float *q, float *k, int pos, int dim, const float *freqs) {
    for (int i = 0; i < dim; i += 2) {
        float q0 = q[i];
        float q1 = q[i + 1];
        float k0 = k[i];
        float k1 = k[i + 1];
        float fcr = freqs[pos * dim + i];     // cos
        float fci = freqs[pos * dim + i + 1]; // sin
        q[i] = q0 * fcr - q1 * fci;
        q[i + 1] = q0 * fci + q1 * fcr;
        k[i] = k0 * fcr - k1 * fci;
        k[i + 1] = k0 * fci + k1 * fcr;
    }
}

Обрати внимание: мы не используем double для частот, хотя в оригинальной статье RoPE используют double. Для инференса float достаточно. Это экономит память и ускоряет вычисления.

4 Attention без мягких массивов

Мы реализуем multi-head attention с нуля. Каждая голова считает attention score для своего куска эмбеддинга. Ключевой момент — мы не используем отдельные массивы для query, key, value каждой головы. Вместо этого работаем с одним большим массивом и смещениями (offsets). Это уменьшает количество аллокаций и копирований.

После получения logits применяем softmax. Здесь тоже ловушка: наивный softmax нестабилен численно. Используем трюк с вычитанием максимума.

void softmax(float *x, int n) {
    float max_val = x[0];
    for (int i = 1; i < n; i++) {
        if (x[i] > max_val) max_val = x[i];
    }
    float sum = 0.0f;
    for (int i = 0; i < n; i++) {
        x[i] = expf(x[i] - max_val);
        sum += x[i];
    }
    for (int i = 0; i < n; i++) {
        x[i] /= sum;
    }
}

Цифры не врут: бенчмарки на M2 Pro и старом i5

Я прогнал движок на двух машинах: Apple MacBook Pro с M2 Pro (12 ядер) и старом десктопе с Intel i5-10400 (6 ядер, без AVX-512). Обе системы с 32 ГБ RAM. Для сравнения, я также запустил инференс через Hugging Face transformers (с оптимизацией) и llama.cpp в CPU-режиме.

Выводы? Наш движок быстрее чистого PyTorch в 6 раз на старом i5. llama.cpp с квантизацией Q4_K_M все еще быстрее (спасибо оптимизациям под AVX2 и квантизации), но наш движок работает с полной точностью fp32. И главное — он проще. Если тебе нужно запустить модель на чистой CPU без квантизации, наш подход выигрывает.

Ошибки, которые съедят ваше время

• Не выравнивание памяти. CBLAS функции требуют, чтобы данные в матрицах были выровнены по границам, соответствующим SIMD-регистрам (16 или 32 байта). Используйте posix_memalign вместо malloc.
• Проверка размерностей в GEMM. Перепутать M, N, K — самая частая ошибка. Я потратил два часа, пока не понял, что attention scores получаются размером не [seq_len, seq_len], а [seq_len, dim]. Все из-за неправильного порядка аргументов в cblas_sgemm.
• Утечки памяти в цикле инференса. Каждый токен генерирует промежуточные массивы. Если не освобождать их, память закончится через 1000 токенов. Используйте стековые массивы (если размер известен) или arena allocator.
• Игнорирование кэша процессора. При работе с attention матрицами [seq_len, dim] доступ к памяти становится случайным. Это убивает производительность. Попробуйте переупорядочить вычисления, чтобы работать с блоками данных, которые помещаются в L1 кэш. Техники из llama.cpp здесь очень помогают.

Что дальше? Куда развивать движок

Этот движок — основа. Его можно расширить в нескольких направлениях:

1. Поддержка батчинга. Сейчас мы обрабатываем один промпт. Добавив dimension для batch size в матричных операциях, можно ускорить инференс для нескольких запросов сразу. Принципы из статьи про vLLM пригодятся.
2. Квантизация. Добавить поддержку INT8 или даже 4-битных весов. Это сократит использование памяти в 4-8 раз и ускорит вычисления (если процессор поддерживает INT8 инструкции).
3. Поддержка более крупных моделей. LFM2-350M — только начало. LFM2-2.6B уже требует больше памяти и оптимизаций. Возможно, придется добавить paging внимания (как в vLLM) для длинных контекстов.
4. Биндинги для Python. Чтобы использовать движок из Python, оберни ключевые функции в Cython или используй ctypes. Это даст удобный интерфейс без потери производительности.

И последнее: этот движок — образовательный проект. Он показывает, что даже в эпоху гигантских фреймворков можно написать эффективный код на простом C. Если ты хочешь понять, как работают LLM изнутри, нет лучшего способа, чем написать свой инференс-движок. С нуля.

Частые вопросы (FAQ)