Линейный слой изнутри: как написать свой GEMM на CUDA и перестать бояться нейросетей

Зачем вам этот гайд, если есть PyTorch?

Потому что копировать nn.Linear из библиотеки – это как водить машину, не зная, как работает двигатель. Когда что-то сломается, вы будете беспомощны. А в нейросетях ломается все, особенно когда вы пытаетесь ускорить инференс на своей видеокарте. Я видел десятки проектов, где люди часами тюнингают гиперпараметры, но не могут объяснить, как их тензор умножается на матрицу весов. Это позор.

Проблема: ваш линейный слой – это бутылочное горлышко

Весь Transformer, вся LLM, вся ваша любимая архитектура держится на одной операции: GEMM (General Matrix Multiply). Линейный слой – это просто обертка вокруг нее. Вы подаете на вход тензор (batch, input_dim), умножаете на матрицу весов (input_dim, output_dim), добавляете bias. Звучит просто. Пока вы не попробуете сделать это на GPU для батча из 32 примеров с размерностью 4096. Тогда вы узнаете, что такое настоящая боль.

1 Математика, которую нельзя пропустить

Пишем формулу, но без фанатизма: Y = X * W^T + B. Да, именно W^T (транспонированная). В PyTorch и TensorFlow по историческим причинам веса хранятся в транспонированном виде. X – матрица батча, W – веса, B – bias. Наша задача – реализовать это умножение максимально быстро.

// Псевдокод того, что нам нужно сделать
void linear_layer_cpu(float* output, const float* input, const float* weight, const float* bias,
                      int batch_size, int input_dim, int output_dim) {
    for (int b = 0; b < batch_size; ++b) {
        for (int o = 0; o < output_dim; ++o) {
            float sum = bias != nullptr ? bias[o] : 0.0f;
            for (int i = 0; i < input_dim; ++i) {
                sum += input[b * input_dim + i] * weight[o * input_dim + i]; // W уже хранится транспонированной
            }
            output[b * output_dim + o] = sum;
        }
    }
}

Этот код на CPU будет работать в 1000 раз медленнее, чем на GPU. Вы готовы к этому?

2 Готовим почву: базовый C++ и память

Прежде чем бросаться на GPU, нужно понять, как работать с памятью. Забудьте про new и delete. Мы используем cudaMallocManaged для Unified Memory или раздельные выделения. Вот как выглядит структура нашего тензорного объекта (упрощенно):

struct Tensor {
    float* data;
    int dims[2]; // [rows, cols]
    Tensor(int rows, int cols) {
        cudaMallocManaged(&data, rows * cols * sizeof(float));
        dims[0] = rows; dims[1] = cols;
        cudaMemset(data, 0, rows * cols * sizeof(float));
    }
    ~Tensor() { cudaFree(data); }
};

3 Первое ядро CUDA: наивный GEMM

Теперь самое интересное. Пишем ядро умножения матриц. Базовый вариант: каждый поток вычисляет один элемент выходной матрицы. Если вы не читали мою статью "CUDA с нуля: пишем ядро для сложения векторов", сделайте это сейчас. Иначе следующий код покажется вам китайской грамотой.

__global__ void gemm_naive(float* C, const float* A, const float* B,
                           int M, int N, int K) {
    // M - строки A и C, N - столбцы B и C, K - столбцы A / строки B
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < M && col < N) {
        float sum = 0.0f;
        for (int k = 0; k < K; ++k) {
            sum += A[row * K + k] * B[k * N + col];
        }
        C[row * N + col] = sum;
    }
}

// Вызов ядра для линейного слоя: A=input, B=weight^T, C=output
dim3 blockDim(16, 16);
dim3 gridDim((output_dim + 15) / 16, (batch_size + 15) / 16);
gemm_naive<<>>(output, input, weight, batch_size, output_dim, input_dim);

Этот код уже работает на GPU. И он уже в 50 раз быстрее CPU-версии. Но это только начало.

4 Оптимизация: shared memory и tiling

Наивное ядро гоняет данные из глобальной памяти для каждого вычисления. Это медленно. Решение – загрузить плитки (tiles) матриц в shared memory, которая быстрее в сотни раз. Вот схема:

__global__ void gemm_tiled(float* C, const float* A, const float* B,
                           int M, int N, int K) {
    __shared__ float As[TILE_SIZE][TILE_SIZE];
    __shared__ float Bs[TILE_SIZE][TILE_SIZE];

    int bx = blockIdx.x, by = blockIdx.y;
    int tx = threadIdx.x, ty = threadIdx.y;

    int row = by * TILE_SIZE + ty;
    int col = bx * TILE_SIZE + tx;

    float sum = 0.0f;

    for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; ++t) {
        // Загружаем плитки в shared memory
        if (row < M && t * TILE_SIZE + tx < K)
            As[ty][tx] = A[row * K + t * TILE_SIZE + tx];
        else
            As[ty][tx] = 0.0f;

        if (t * TILE_SIZE + ty < K && col < N)
            Bs[ty][tx] = B[(t * TILE_SIZE + ty) * N + col];
        else
            Bs[ty][tx] = 0.0f;

        __syncthreads();

        for (int k = 0; k < TILE_SIZE; ++k) {
            sum += As[ty][k] * Bs[k][tx];
        }
        __syncthreads();
    }

    if (row < M && col < N) {
        C[row * N + col] = sum;
    }
}

Это классический алгоритм. TILE_SIZE обычно 16 или 32. На современных GPU (архитектура Hopper на 2026 год) есть тензорные ядра, но для понимания основ shared memory – must have.

Нюансы, которые всех ломают

Вы написали ядро, оно работает. Но производительность в 3 раза ниже, чем у cuBLAS. Почему?

• Memory coalescing: ваши потоки должны читать из глобальной памяти последовательные адреса. Иначе шина памяти простаивает.
• Bank conflicts в shared memory: если несколько потоков обращаются к одной банке памяти, они становятся в очередь. Решение – padding (добавление лишнего столбца в массив).
• Зануление bias: если bias нет, не заставляйте ядро прибавлять ноль. Создайте отдельное ядро или передавайте nullptr.

И да, всегда проверяйте ошибки CUDA после каждого вызова ядра. cudaDeviceSynchronize() и cudaGetLastError() – ваши лучшие друзья.

А что с backward pass?

Для обучения нужны градиенты. Backward pass – это еще два умножения матриц: градиент по весам = input.T * grad_output, градиент по входу = grad_output * weight. Реализуйте их как отдельные ядра. Или используйте технику, описанную в статье "Тензорный параллелизм в Llama.cpp", чтобы распределить вычисления.

Собираем все вместе: линейный слой как класс

class LinearLayer {
private:
    Tensor weight, bias, weight_grad, bias_grad;
    int input_dim, output_dim;
    bool has_bias;
public:
    LinearLayer(int in_dim, int out_dim, bool with_bias = true)
        : input_dim(in_dim), output_dim(out_dim), has_bias(with_bias),
          weight(out_dim, in_dim), bias(1, out_dim),
          weight_grad(out_dim, in_dim), bias_grad(1, out_dim) {
        // Инициализация весов (например, Xavier)
        initialize_weight(weight.data, out_dim, in_dim);
        if (has_bias) cudaMemset(bias.data, 0, output_dim * sizeof(float));
    }

    void forward(Tensor& output, const Tensor& input) {
        // output = input * weight^T + bias
        gemm_tiled<<<...>>>(output.data, input.data, weight.data,
                            input.dims[0], output_dim, input_dim);
        if (has_bias) add_bias<<<...>>>(output.data, bias.data, input.dims[0], output_dim);
    }

    void backward(...) { /* Реализация backward */ }
};

Что дальше? Вы построили фундамент

Теперь вы понимаете, из чего состоят линейные слои в Llama 3, GPT-4.5 (или какой там актуальной модели на 2026 год). Вы можете:

• Добавить поддержку float16 или bfloat16 для экономии памяти.
• Реализовать весь Transformer с нуля, используя свой GEMM.
• Оптимизировать под конкретную архитектуру GPU, используя тензорные ядра через WMMA API.
• Собрать свою маленькую библиотеку для инференса, как описанный здесь движок.

Когда вы в следующий раз будете использовать torch.nn.Linear, вы почувствуете разницу. Вы будете знать, что под капотом. И когда вам понадобится ускорить модель для масштабирования на кластере, у вас будет фундамент, чтобы не просто копировать код из Stack Overflow, а понимать его.

Подписаться на канал