ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Сравнить Hyena vs FlashAttention на 128k

1. Цель задачи

Эмпирически сравнить две архитектуры внимания — Hyena (O(n) по памяти) и FlashAttention (O(n²) по памяти) — на длине последовательности 128k токенов. Измерить пиковое потребление GPU-памяти и время прямого прохода (forward pass). Сформулировать практические рекомендации: при каких длинах контекста стоит переходить с FlashAttention на Hyena, как это влияет на throughput.

Ключевой результат Таблица сравнения Memory(N) и Time(N) для N = 4k, 8k, 16k, 32k, 64k, 128k с наглядными графиками.

2. Исходные данные

Если нет реального GPU с 24+ GB — симулируем:

1. Установить ограничение памяти через torch.cuda.set_per_process_memory_fraction(0.5) (например, для GPU 16 GB эмулировать 8 GB).
2. Запускать только до длины, помещающейся в лимит; для FlashAttention на 128k (B=1, D=1024) нужно минимум ~16 GB (реально ~12-14 GB). Если нет GPU совсем — выполнить расчёт теоретической сложности и экстраполировать на основе замеров на малых длинах (4k, 8k) на CPU. Это даст качественное понимание.
3. Для симуляции можно воспользоваться Google Colab с A100 (бесплатно ~40-80 GB) — лимит 12 часов.

3. Технологический стек

4. Этапы выполнения

Этап 1: Настройка окружения и импорт моделей (45 мин)

Действия

1. Создать виртуальное окружение Python 3.10+:

   python -m venv venv_hyena
   source venv_hyena/bin/activate
   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
   pip install flash-attn --no-build-isolation
   pip install git+https://github.com/HazyResearch/safari.git
   pip install matplotlib pandas wandb
   

2. Проверить установку:

   import torch
   from flash_attn import flash_attn_func
   from safari import HyenaOperator
   
   print("FlashAttention и HyenaOperator импортированы успешно")
   

3. Определить конфигурации для бенчмарка:

   • batch_size = 1 (чтобы избежать дополнительных эффектов)
   • head_dim = 64 (стандарт для FlashAttention)
   • num_heads = 4 (или 8, в зависимости от hidden_dim)
   • hidden_dim = num_heads * head_dim (256 или 512)
   • Длины: [4096, 8192, 16384, 32768, 65536, 131072]

Ожидаемый результат этапа Работающее окружение, все зависимости установлены, готовый скрипт-заготовка.

Этап 2: Реализация бенчмарка для FlashAttention (2 часа)

Действия

1. Написать функцию benchmark_flash_attention(L, B, H, D_head):

   def benchmark_flash_attention(L, B, H, D_head):
       q = torch.randn(B, H, L, D_head, device='cuda', dtype=torch.float16)
       k = torch.randn(B, H, L, D_head, device='cuda', dtype=torch.float16)
       v = torch.randn(B, H, L, D_head, device='cuda', dtype=torch.float16)
       torch.cuda.reset_peak_memory_stats()
       start = torch.cuda.Event(enable_timing=True)
       end = torch.cuda.Event(enable_timing=True)
       start.record()
       out = flash_attn_func(q, k, v, causal=True)  # causal для авторегрессии
       end.record()
       torch.cuda.synchronize()
       time_ms = start.elapsed_time(end)
       peak_mem = torch.cuda.max_memory_allocated() / (1024**3)  # GB
       return time_ms, peak_mem
   

2. Запустить цикл по всем длинам. Для длин, на которых возникает CUDA out of memory, записать "OOM". Обработать исключения:

   for L in lengths:
       try:
           t, m = benchmark_flash_attention(L, 1, 4, 64)
           results.append({'L': L, 'time_ms': t, 'mem_GB': m})
       except RuntimeError as e:
           if 'out of memory' in str(e):
               results.append({'L': L, 'time_ms': None, 'mem_GB': 'OOM'})
               torch.cuda.empty_cache()
           else:
               raise
   

3. Сохранить результаты в results_flash.csv.

Ожидаемый результат этапа CSV-файл с временем и памятью для FlashAttention на всех длинах (где помещается), с метками OOM.

Этап 3: Реализация бенчмарка для Hyena (1.5 часа)

Действия

1. Создать модель Hyena с аналогичной размерностью (hidden_dim = 256, num_heads=4, filter_order=2... см. документацию safari):

   from safari import HyenaOperator
   model = HyenaOperator(
       d_model=256,
       l_max=131072,   # максимальная длина
       num_heads=4,
       filter_order=2,
       post_order=0,
       local_window_size=256,
       causal=True,
   ).to('cuda').half()
   

   Примечание: HyenaOperator может потребовать дополнительных параметров (layer_idx, residual). Уточнить по документации safari.

2. Написать функцию benchmark_hyena(L, B, H, D):

   def benchmark_hyena(L, B, hidden_dim):
       x = torch.randn(B, L, hidden_dim, device='cuda', dtype=torch.float16)
       torch.cuda.reset_peak_memory_stats()
       start = torch.cuda.Event(enable_timing=True)
       end = torch.cuda.Event(enable_timing=True)
       start.record()
       out = model(x)  # shape: (B, L, hidden_dim)
       end.record()
       torch.cuda.synchronize()
       time_ms = start.elapsed_time(end)
       peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
       return time_ms, peak_mem
   

3. Запустить те же длины. Зафиксировать результаты в results_hyena.csv.

Ожидаемый результат этапа CSV-файл с данными по Hyena. Обратить внимание, что Hyena может выдать ошибку на длинах > l_max — если в конструкторе указано l_max=131072, то должно работать на всех.

Этап 4: Построение графиков и анализ (1 час)

Действия

1. Загрузить два CSV в pandas:

   df_flash = pd.read_csv('results_flash.csv')
   df_hyena = pd.read_csv('results_hyena.csv')
   

2. Построить два графика (отдельные figure):

   • Memory vs Sequence Length (log-log scale)
   • Time vs Sequence Length (log-log scale)

   На каждом графике — две линии: FlashAttention и Hyena. Добавить аппроксимацию O(n²) и O(n) пунктиром.

3. Вычислить среднее ускорение памяти для Hyena на максимальной измеренной длине (где обе не OOM). Например: на 64k FlashAttention: 8.1 GB, Hyena: 1.2 GB → в 6.75x меньше.

4. Вывести таблицу с результатами в консоль и сохранить в comparison_table.tex (или .md).

Ожидаемый результат этапа Два PNG-графика (memory_vs_length.png, time_vs_length.png) и сводная таблица.

Этап 5: Оформление отчёта (30 мин)

Действия

1. Написать краткий markdown-отчёт report.md с:
   • Описание эксперимента
   • Таблица результатов (L, Flash Mem, Flash Time, Hyena Mem, Hyena Time)
   • Графики (вставлены как изображения)
   • Выводы: при L < 16k разница незначительна; при L > 32k Hyena лучше по памяти; при L = 128k FlashAttention OOM, Hyena работает.
   • Ограничения: использовался малый hidden_dim, не замеряли backward pass, не учитывали throughput при batched inference.

Ожидаемый результат этапа Файл report.md с результатами.

5. Критерии приемки (Definition of Done)

• Установлены и корректно импортируются пакеты flash-attn и safari.
• Запущен бенчмарк для FlashAttention для длин 4k–128k; для длин, где OOM, зафиксирован факт OOM.
• Запущен бенчмарк для Hyena для тех же длин; все длины прошли без OOM.
• Получены CSV-файлы с колонками L, time_ms, mem_GB.
• Построены логарифмические графики memory(L) и time(L) с двумя кривыми.
• Визуально подтверждено: FlashAttention — O(n²), Hyena — O(n) (наклон линии 2 vs 1).
• Написаны выводы с указанием порога, с которого Hyena становится выгоднее по памяти (например, для 16k+).
• Отчёт report.md содержит все таблицы и графики.
• Эксперимент воспроизводим (указаны команды установки, версии).

6. Ожидаемый результат

Основной артефакт Папка hyena_vs_flash_benchmark/, содержащая:

• benchmark.py — скрипт для запуска бенчмарка (объединяет этапы 2-3)
• results_flash.csv — сырые данные FlashAttention
• results_hyena.csv — сырые данные Hyena
• memory_vs_length.png, time_vs_length.png — графики
• report.md — отчёт с выводами

Опционально

• Таблица в LaTeX (table.tex) для вставки в статью.
• Логи wandb (если использовался).

7. Возможные сложности и их решение

8. Бюджет времени (оценка)

Примечание: Для первого раза (освоение API) заложите +2 часа на отладку. Итого ~8 часов.

9. Связанные вопросы из базы знаний

10. Чек-лист самопроверки

• Я установил flash-attn через pip с правильной версией CUDA и проверил импорт.
• Я установил safari (или другую реализацию Hyena) и проверил, что модель создаётся без ошибок.
• Я написал функции замеров с torch.cuda.reset_peak_memory_stats() до запуска и torch.cuda.max_memory_allocated() после.
• Я использовал одинаковые параметры (hidden_dim, batch_size, dtype) для обеих архитектур.
• Я обработал исключения CUDA OOM и записал их в результаты, не прерывая цикл.
• Я построил графики в логарифмическом масштабе и вижу наклон 2 для FlashAttention и 1 для Hyena.
• Я описал выводы: при какой длине Hyena становится предпочтительнее по памяти, а также отметил, что по времени Hyena может быть медленнее на коротких длинах.
• Отчёт сохранён в .md с вставленными изображениями (относительные пути).
• Все файлы лежат в одной папке, скрипт воспроизводим командой python benchmark.py.