Победа в хакатоне PyTorch: как оптимизировать GPU kernel для инференса LLM — разбор техник из первых рук

Когда стандартные оптимизации бессильны

Типичная история: вы берете свежую модель, скажем, LFM2-350M с архитектурой Gated DeltaNet, запускаете инференс на NVIDIA B200, а получаете 50 токенов в секунду. Вместо ожидаемых 500. Знакомо?

Мы столкнулись с этим на хакатоне PyTorch в марте 2026 года. Задача - ускорить инференс кастомной модели Liquid AI на 8 GPU. Стандартные трюки - quantization, torch.compile, vLLM - не работали. Архитектура Gated DeltaNet ломала все привычные оптимизации.

Проблема глубже. Современные LLM в 2026 году - это не просто трансформеры. Это гибриды с state space моделями, линейными вниманиями, динамической архитектурой. MoE-слои, sparse активации, conditional computation. Стандартные ядра PyTorch не заточены под эту экзотику.

Почему ваш kernel на B200 работает в 10 раз медленнее возможного

NVIDIA B200, вышедшая в конце 2025, принесла не только терафлопсы, но и новые головные боли. Архитектура Blackwell с chiplet design требует переосмысления memory access patterns. Старые трюки с shared memory и warp shuffles больше не работают так эффективно.

Но проблема не только в железе. Посмотрите на свой код инференса. Скорее всего, он выглядит так:

# Типичный медленный инференс (как НЕ делать)
for layer in model.layers:
    hidden_states = layer.attention(hidden_states, use_cache=True)
    hidden_states = layer.mlp(hidden_states)
    # Тут еще 5-10 операций активации, нормализации, residual

Каждая операция - отдельный kernel launch. Каждый launch - синхронизация, overhead драйвера, неоптимальное использование памяти. На B200 с её 8 TB/s памяти это убийственно.

В статье про свой vLLM для LFM я уже показывал, как гибридный KV cache решает часть проблем. Но для хакатона нужны были более агрессивные оптимизации.

1 Диагностика: находим реальное узкое место

Первое правило оптимизации: не угадывать, а измерять. На хакатоне у нас было 48 часов. Тратить их на оптимизацию не тех частей кода - самоубийство.

Мы использовали PyTorch Profiler с последним расширением для B200. Ключевые метрики на 2026 год:

• Memory bandwidth utilization - на B200 должно быть >80%, у нас было 35%
• Tensor Core occupancy - новые TF32-tc ядра использовались на 15%
• Kernel launch latency - 7 мкс на вызов, при том что многие ядра работали 2-3 мкс

Оказалось, главная проблема - не вычисления, а организация данных и overhead вызовов. Gated DeltaNet требует частого переключения между разными режимами внимания, что приводило к constant cache thrashing.

2 Fusion: убиваем kernel launch overhead

Каждый вызов ядра в PyTorch 2.4+ на B200 стоит примерно 5-10 мкс. Когда у вас 20 операций на слой и 24 слоя, это складывается в 5ms только на overhead. При генерации 50 токенов в секунду это 25% времени!

Решение - fusion. Но не тот наивный fusion, который делает torch.compile. Нам нужен был custom fusion под специфику Gated DeltaNet.

Вот как выглядел наш fusion kernel для одного слоя:

// Псевдокод fused kernel для Gated DeltaNet layer
__global__ void gated_deltanet_layer_fused(
    float* hidden_states,
    float* kv_cache,
    float* gates,
    int seq_len,
    int batch_size) {
    
    // 1. Загружаем все данные в shared memory за один коалисцированный доступ
    __shared__ float shmem[8192];  // 16KB на блок
    load_coalesced(hidden_states, shmem, threadIdx.x, blockDim.x);
    
    // 2. Вычисляем attention и MLP параллельно, пока данные в кэше
    float attention_out = compute_gated_attention(shmem, gates);
    float mlp_out = compute_mlp(shmem);
    
    // 3. Dynamic gating - но без branch divergence!
    float gate_value = shmem[gate_index];
    float output = gate_value * attention_out + (1 - gate_value) * mlp_out;
    
    // 4. LayerNorm fused with residual
    output = fused_layernorm_residual(output, shmem);
    
    // 5. Сохраняем результат
    store_coalesced(output, hidden_states);
}

Ключевые трюки:

1. Coalesced memory access - вместо случайных чтений/записей, организуем данные так, чтобы 32 потока в warp читали 32 последовательных float
2. Persistent threads - ядро не завершается после обработки одного элемента, а обрабатывает несколько, уменьшая launch overhead
3. Predicated execution - вместо if-else ветвлений в Gated DeltaNet используем predicate регистры

3 Адаптивный KV caching для Gated DeltaNet

Стандартный KV cache не работает для динамических архитектур. В Gated DeltaNet attention weights меняются в зависимости от gate значений. Кэшировать ключи и значения после attention - бессмысленно, они устаревают к следующему токену.

Наше решение - two-level adaptive cache:

• Level 1: Кэшируем raw projections (Q, K, V) до gating
• Level 2: Кэшируем частично вычисленные attention scores для стабильных gates

Алгоритм работал так:

class AdaptiveKVCache:
    def __init__(self, num_layers, num_heads, head_dim):
        self.raw_cache = []  # Raw Q,K,V projections
        self.score_cache = []  # Precomputed scores
        self.gate_history = []  # Gate values для анализа стабильности
        
    def get(self, layer_idx, position, current_gates):
        # Если gates не изменились больше чем на 10% - используем score cache
        if self.is_stable(layer_idx, position, current_gates):
            return self.score_cache[layer_idx][position]
        # Иначе пересчитываем из raw cache
        raw = self.raw_cache[layer_idx][position]
        return compute_gated_attention(raw, current_gates)

Стабильность gates определяли через экспоненциальное скользящее среднее. Если gates колебались меньше порога - использовали кэш второго уровня. Это снижало вычисления на 60% для стабильных контекстов.

4 Использование новых возможностей B200: TMA и асинхронность

Blackwell B200 принесла Tensor Memory Accelerator (TMA) - аппаратный блок для эффективных тензорных копий. Стандартные cudaMemcpy его не используют. Нужны новые API из CUDA 13+.

Настройка TMA для нашего KV cache:

// Инициализация TMA для кэша
cudaTensorMap_t tma_map;
cudaTensorMapObject_t map_obj;

cudaTensorMapObjectParams params = {};
params.tensorPtr = kv_cache_ptr;
params.tensorShape = {seq_len, num_heads, head_dim};
params.tensorStrides = {head_dim * num_heads, head_dim, 1};
params.tensorDataType = CUDA_R_16F;  // BF16 cache
params.tmemAccessSize = CUDA_TMA_ACCESS_64B;  // Для B200

cudaTensorMapCreate(&tma_map, &params);

// В ядре используем TMA загрузку
asm volatile(
    "cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], [%2];\n"
    : : "r"(shmem_ptr), "l"(tma_map), "r"(mbarrier_ptr));

TMA уменьшил latency загрузки кэша с 200 нс до 40 нс. Но была загвоздка: TMA требует specific memory alignment и swizzling pattern. Если ваши тензоры не выровнены по 64B границам, производительность падает в 10 раз.

Еще одна фишка B200 - асинхронные копии между chiplet'ами. У B200 четыре chiplet'а, и data movement между ними раньше был синхронным. В CUDA 13 появились асинхронные барьеры и копии:

// Асинхронная копия между chiplet'ами
cudaMemcpyAsync3DParams params = {};
params.srcPtr = make_cudaPitchedPtr(src, width, height, depth);
params.dstPtr = make_cudaPitchedPtr(dst, width, height, depth);
params.extent = make_cudaExtent(width, height, depth);
params.kind = cudaMemcpyDeviceToDevice;
params.flags = cudaMemcpyInterChipletAsync;  // Новый флаг!

cudaMemcpy3DAsync(&params, stream);

Типичные ошибки, которые сведут на нет все оптимизации

После 48 часов хакинга мы сделали все эти оптимизации. И получили... 10% ускорение вместо ожидаемых 5x. Почему?

Оказалось, мы наступили на все грабли, описанные в статье про тихий провал кастомных ядер. Вот что сломалось:

1. Silent fallback на CPU: наш fused kernel требовал 128KB shared memory. По умолчанию PyTorch ограничивает 48KB. Ядро компилировалось, но при запуске падало в fallback на CPU. Без ошибок!
2. Bank conflicts в shared memory: мы оптимизировали coalesced access к глобальной памяти, но создали bank conflicts в shared. На B200 bank width 4 байта, а не 8 как в Hopper.
3. Неверное использование tensor cores: новые TF32-tc ядра требуют specific data layout. Наши матрицы были выровнены для старых tensor cores, что давало 1/4 производительности.

Решение - агрессивный profiling и валидация:

# Команды для отладки на B200
nsys profile --stats=true ./inference_script
# Смотрим REAL occupancy, не theoretical
nv-nsight-cu-cli --metrics "sm__maximum_warps_avg,sm__warps_active_avg"
# Проверяем использование TMA
ncu --metrics "tma__bytes_queried,tma__throughput"

Финальный результат: от 50 до 480 токенов в секунду

После фикса всех багов и тонкой настройки:

• Baseline (чистый PyTorch): 52 токена/сек на 8×B200
• После fusion: 145 токена/сек (2.8x)
• + Adaptive KV cache: 240 токена/сек (4.6x)
• + TMA и асинхронность: 380 токена/сек (7.3x)
• + Оптимизация под B200 специфики: 480 токена/сек (9.2x)

9.2x ускорение на тех же железе и модели. Это принесло нам первое место на хакатоне.

Но главное не цифры. Главное - инсайты:

1. Современные GPU (особенно B200) - это не просто более быстрые старые GPU. Это новая архитектура, требующая переписывания ядер с нуля.
2. Fusion важен, но умный fusion важнее. Слепое объединение операций может замедлить из-за register pressure.
3. Профайлеры врут. Или точнее, показывают упрощенную картину. Нужно смотреть на hardware counters, а не на высокоуровневые метрики.

Если беретесь за оптимизацию kernel'ов в 2026 году, начинайте не с кода, а с документации на конкретную GPU. Архитектурные изменения между поколениями теперь кардинальны. То, что работало на A100, будет тормозить на B200. То, что летало на H100, не запустится на Blackwell.

И последнее: наши оптимизации для Gated DeltaNet оказались полезны и для других dynamic architectures. Liquid AI уже интегрировали некоторые техники в свою кодобазу. А мы... мы просто рады, что не спали двое суток и победили.

Что дальше? Прогноз на 2027 год

Если думаете, что B200 - это предел, ждите Blackwell Ultra в конце 2026. Слухи: 12 chiplet'ов, unified memory до 192GB, и новый challenge - оптимизация для 3D-стэка памяти.

Уже сейчас готовьтесь к тому, что стандартные подходы к инференсу умрут. MoE модели с тысячами экспертов, модели с динамической архитектурой в runtime, нейросети, которые меняют граф вычислений для каждого батча - это будущее 2027 года.

Наш код и слайды с learnings мы выложили в открытый доступ. Но предупреждаю: это не готовое решение, а набор техник. Вашу модель придется оптимизировать с нуля. Как мы это сделали для разгона LLM до предела.

Удачи в следующих хакатонах. И да пребудет с вами низкая latency.

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