Как работает L1/L2 cache hierarchy в A100/H100 и как ее использовать для LLM?

Краткий тезис

Иерархия кэшей L1/L2 в GPU A100/H100 критически влияет на производительность LLM-инференса. L1 (192/256 КБ на SM) используется как shared memory для временных данных (например, блоки матриц при attention), а L2 (40/50 МБ) служит буфером между SM и глобальной памятью. Для LLM ключевая оптимизация — размещение KV cache в L2 при длинных контекстах и использование tiling в shared memory для эффективного вычисления attention, что снижает latency и повышает throughput.

1. Зачем нужна иерархия кэша в GPU для LLM

LLM (Large Language Models) требуют огромного количества операций с памятью: веса модели, активации, KV cache. Скорость доступа к разным уровням памяти различается на порядки:

Для LLM инференса основное узкое место — memory-bound операции (загрузка весов и KV cache). Эффективное использование L1/L2 позволяет сократить обращения к HBM, ускоряя генерацию.

2. Архитектура GPU: SM, L1, L2, Global Memory

Streaming Multiprocessor (SM) — вычислительный блок GPU. A100 содержит 108 SM, H100 — 132 SM. Каждый SM имеет собственный L1 кэш (также используется как shared memory). Все SM разделяют L2 кэш и HBM (High Bandwidth Memory).

• L1 cache / shared memory: физически один и тот же блок SRAM, программист может распределять его между кэшем данных и shared memory (например, 50/50 или 75/25).
• L2 cache: единый для всех SM, когерентный, служит для кэширования данных из HBM.
• Global memory (HBM): основная память GPU, большая ёмкость (40/80 ГБ), но высокая задержка.

3. L1 cache / Shared Memory в A100 и H100

Shared memory — это быстрая память, управляемая программистом. В контексте LLM она используется для:

• Хранения блоков матриц при tiled matrix multiplication (GEMM).
• Буферизации частей Q, K, V при вычислении attention.
• Аккумуляции частичных результатов (например, softmax).

L1 cache автоматически кэширует обращения к global memory, но для LLM обычно предпочтительнее явно использовать shared memory, чтобы гарантировать попадание.

4. L2 cache в A100 и H100

L2 кэш является общим для всех SM. Он кэширует данные из HBM, уменьшая latency при повторных обращениях. Для LLM:

• KV cache (ключи и значения прошлых токенов) может помещаться в L2 при длине контекста до ~50К токенов (для H100 с 50 МБ L2, если KV cache занимает ~1 МБ на токен? На самом деле KV cache для одного слоя: 2 * hidden_dim * num_heads * precision. Для Llama 70B: hidden_dim=8192, num_heads=64, FP16 → 2819264*2 = 2 МБ на слой. 80 слоёв → 160 МБ на один токен. Значит, L2 не вмещает весь KV cache. Но для небольших моделей или при малом batch size часть может поместиться.)
• Веса модели (если модель небольшая) также могут кэшироваться в L2, но обычно веса большие и не помещаются.

5. Как LLM использует память: веса, KV cache, активации

При инференсе LLM выделяют три основных типа данных:

1. Веса модели — загружаются один раз и остаются в HBM (если не помещаются в L2).
2. KV cache — растёт с каждым новым токеном, хранит прошлые ключи и значения для каждого слоя.
3. Активации (промежуточные результаты) — временные, живут только в рамках одного forward pass.

Для каждого шага декодирования:

• Загружаются веса текущего слоя (memory-bound).
• Загружается KV cache предыдущих токенов (memory-bound).
• Вычисляется attention (compute-bound, если данные уже в shared memory).

Оптимизация иерархии кэша направлена на то, чтобы максимально использовать L1/shared для вычислений и L2 для повторного чтения KV cache.

6. Использование L1 / Shared Memory для LLM: FlashAttention и Tiling

FlashAttention — алгоритм, который явно использует shared memory для вычисления точного attention без сохранения полной матрицы S = QK^T в HBM. Основные шаги:

1. Разбить Q, K, V на блоки (tiles), которые помещаются в shared memory.
2. Загрузить блок K и V из HBM в shared memory.
3. Вычислить частичное QK^T в регистрах/shared memory.
4. Применить softmax по частям (online softmax).
5. Записать результат в HBM.

Пример кода (псевдо-CUDA):

__global__ void flash_attention_kernel(float* Q, float* K, float* V, float* O, int N, int d) {
    extern __shared__ float shared[];
    float* K_block = shared;
    float* V_block = shared + block_size * d;
    
    int tx = threadIdx.x;
    float acc = 0.0f;
    float m = -INFINITY, l = 0.0f;
    
    for (int j = 0; j < N; j += block_size) {
        // загрузка блока K и V в shared memory
        if (tx < block_size * d) {
            K_block[tx] = K[j * d + tx];
            V_block[tx] = V[j * d + tx];
        }
        __syncthreads();
        
        // вычисление частичного S = Q * K_block^T
        float s = 0.0f;
        for (int k = 0; k < d; ++k) {
            s += Q[blockIdx.x * d + k] * K_block[tx * d + k];
        }
        s *= scale;
        
        // online softmax
        float m_prev = m;
        m = max(m, s);
        l = l * expf(m_prev - m) + expf(s - m);
        acc = acc * expf(m_prev - m) + expf(s - m) * V_block[tx * d + blockIdx.y];
        __syncthreads();
    }
    O[blockIdx.x * d + blockIdx.y] = acc / l;
}

Этот подход минимизирует обращения к HBM, используя shared memory (L1) для временных блоков.

7. Использование L2 для KV Cache

KV cache — это массив ключей и значений для каждого слоя. При генерации каждого нового токена нужно прочитать весь предыдущий KV cache (длины T). Если T достаточно мало, весь KV cache может поместиться в L2, что даёт значительное ускорение.

Пример оценки

• Модель: Llama 7B (hidden_dim=4096, num_heads=32, FP16).
• KV cache на один токен: 2 * 4096 * 32 * 2 байта = 524 288 байт ≈ 0.5 МБ.
• Для T=100 токенов: 50 МБ — как раз размер L2 H100.
• Для T=200 токенов: 100 МБ — уже не помещается, часть будет читаться из HBM.

Оптимизация Использовать PagedAttention (vLLM) — разбивать KV cache на блоки (pages) размером с кэш-линию L2, чтобы увеличить локальность и уменьшить промахи.

8. Оптимизации: увеличение shared memory, уменьшение global memory access

8.1 Балансировка L1/Shared

В A100/H100 можно настроить распределение между L1 и shared memory через cudaFuncSetAttribute:

cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 164*1024); // для A100

Увеличение shared memory позволяет использовать большие блоки в FlashAttention, но уменьшает L1 cache для автоматического кэширования.

8.2 Использование Tensor Cores

Tensor Cores (в H100 четвёртого поколения) выполняют умножение матриц с половинной точностью (FP16/BF16) очень быстро. Для LLM attention можно использовать wmma (warp matrix multiply-accumulate) в shared memory, что ещё больше ускоряет вычисления.

8.3 Fused kernels

Объединение нескольких операций (например, attention + feed-forward) в один kernel позволяет сохранять промежуточные данные в shared memory, избегая записи в HBM.

8.4 Оптимизация batch size и sequence length

• Маленький batch (1-4): KV cache может помещаться в L2, latency низкая.
• Большой batch: KV cache вытесняется, но throughput растёт за счёт лучшей утилизации SM.

9. Практические советы для LLM инженера

• Используйте FlashAttention (реализации в xformers, flash-attn) — они уже оптимизированы под shared memory.
• Для long context (>32K токенов) рассмотрите Ring Attention или Blockwise Parallel Transformer, которые распределяют KV cache между несколькими GPU.
• Мониторьте occupancy (загрузку SM) и cache hit rate через Nsight Compute.
• Настройте распределение L1/Shared под конкретный kernel: для attention нужно больше shared, для GEMM — больше L1.
• Используйте FP8 (H100) для KV cache — уменьшает объём в 2 раза, позволяя поместить больше токенов в L2.

Пет-проект для закрепления

Задача Написать простой CUDA kernel для attention, использующий shared memory (L1) для блоков K и V, и сравнить его производительность с версией, работающей напрямую с HBM.

Инструменты CUDA Toolkit, Python с PyTorch (для вызова kernel), Nsight Compute для профилирования.

Шаги:

1. Реализовать базовый attention kernel без shared memory (чтение K, V из global memory).
2. Реализовать tiled attention с shared memory (как в FlashAttention, но упрощённо).
3. Замерить время выполнения для разных размеров (sequence length 512, 1024, 2048) на A100/H100.
4. Построить график ускорения.
5. Добавить использование Tensor Cores через wmma (опционально).

Ожидаемый результат Ускорение в 2-5 раз для длинных последовательностей, понимание влияния размера shared memory на производительность.

Связь с другими вопросами

Навигация

• Предыдущий: 702
• Следующий: 704
• Индекс: 00. Индекс разборов