В мире генеративного искусственного интеллекта скорость имеет решающее значение. Когда вы разрабатываете инструменты для дизайнеров, редакторы изображений или визуальные производственные конвейеры, задержка (latency) при выводе модели становится главным барьером. Генерация одного изображения требует множества последовательных проходов через модель с миллиардами параметров, и чем выше разрешение, тем дороже вычислительные затраты механизма внимания (attention). NVIDIA представила новое решение этой проблемы: оптимизацию модели FLUX 3 Image Klein на архитектуре Blackwell с использованием TensorRT LLM VisualGen. Это не просто обновление драйверов, а комплексный подход к ускорению инференса, который позволяет сократить время генерации в несколько раз без критической потери качества.
В этой статье мы подробно разберем, как именно работают эти оптимизации, какие технические решения применяются для ускорения диффузионных трансформеров (DiT) и как вы можете применить эти знания на практике. Мы рассмотрим квантование NVFP4, оптимизацию ядер GEMM, разреженное внимание и многопроцессорное масштабирование. Цель — дать вам полное понимание того, как получить максимальную производительность от современных GPU NVIDIA, таких как B200 и RTX PRO 6000, запуская FLUX 3 локально или в облаке.
01Архитектура FLUX 3 Image Klein и вызовы производительности
FLUX 3 Image Klein, разработанный компанией Black Forest Labs (BFL), представляет собой сложную архитектуру, обученную на кластерах NVIDIA GB200 NVL72. Модель состоит из трех ключевых компонентов: текстового кондиционера Qwen3-VL-4B-Instruct, диффузионного трансформера (DiT) для уточнения латентных представлений изображения и автокодировщика для декодирования латентов в итоговое изображение. Важно отметить, что модель предлагает два варианта весов: Base (базовая версия с высокой точностью) и distilled (дистиллированная версия, оптимизированная для скорости с меньшим количеством шагов).
Основная проблема при работе с такими моделями заключается в том, что процесс денормализации (denoising) доминирует над общим временем генерации. На GPU NVIDIA B200 именно этот этап занимает подавляющую часть вычислительных ресурсов. Например, в базовом режиме (Base) с 50 шагами денормализации и коэффициентом guidance (CFG) 3.5, на разрешение 2048x2048, этап денормализации занимает около 14.24 секунд из 14.70 секунд общего времени, что составляет 96.8% всех вычислений. В режиме Image-to-Image (I2I) этот показатель достигает 97.8%. Это означает, что любые оптимизации, направленные на ускорение DiT, дадут максимальный эффект.
02Первый шаг: NVFP4 и оптимизация ядер GEMM
TensorRT LLM VisualGen, специализированный рантайм в составе TensorRT LLM, поддерживает несколько форматов квантования, включая BF16, FP8 и новый формат NVFP4. Использование предварительно квантованных весов NVFP4 для DiT-чекпоинтов FLUX 3 Image Klein является отправной точкой для большинства оптимизаций. По сравнению с классическим форматом BF16, NVFP4 обеспечивает ускорение до 1.62x для генерации Text-to-Image (T2I) и до 1.58x для Image-to-Image (I2I) на одном GPU B200.
Но само по себе квантование — это только половина дела. Вторая половина — это эффективное выполнение матричных умножений. TensorRT LLM VisualGen использует ядра GEMM (General Matrix Multiplication) от FlashInfer, построенные на базе CUTLASS Primitives и Task Scheduling. Эти ядра, называемые PrimsTS, специально настроены под тензорные формы FLUX 3. Они объединяют оптимизированные GEMM с «слитыми» (fused) эпиклогами, такими как GEMM со SwiGLU, а также проекции QKV (query, key, value) с нормализацией и поворотными позиционными эмбеддингами (RoPE).
На момент запуска, зафиксированные ядра NVFP4 PrimsTS обеспечивали ускорение T2I генерации на 2-10% по сравнению с автонастроенной реализацией GEMM в TensorRT LLM. Это кажется небольшим приростом, но в контексте высоконагруженных систем он важен. Для активации этого режима в конфигурационном файле pipeline_config необходимо установить параметр variant в значение nvfp4:
pipeline_config:
variant: nvfp403Оптимизация механизма внимания: квантование и разреженность
Механизм внимания (attention) является главным узким местом при генерации изображений высокого разрешения. В тестах на разрешении 2048x2048 с форматом NVFP4, стандартное (VANILLA) внимание в BF16 занимало около 60% времени выполнения ядра DiT на GPU. Чтобы решить эту проблему, VisualGen предлагает два подхода: квантованное внимание и разреженное внимание.
Квантованное внимание: cuDNN FP8 и VC-Attention
Первый подход использует квантование для снижения вычислительной нагрузки. VisualGen поддерживает два бэкенда: cuDNN FP8 и VC-Attention FP8. VC-Attention является вариантом, реализованным в TensorRT LLM, основанным на методе, описанном в научной статье о VC-Attention. Этот метод использует «бюджет ремонта» (repair budget) для выбора доли токенов с наибольшей ошибкой квантования, чтобы применить остаточную коррекцию.
Результаты тестирования показывают, что по сравнению с VANILLA BF16 вниманием, cuDNN FP8 дает ускорение 1.08x для Base-модели и 1.09x для дистиллированной версии. Однако лидером здесь выступает VC-Attention FP8 с параметром бюджета ремонта ρ=0.005. Оно обеспечивает ускорение 1.18x для Base и 1.16x для дистиллированной модели. При этом метрика LPIPS (Learned Perceptual Image Patch Similarity), измеряющая воспринимаемые различия между изображениями, остается на приемлемом уровне (0.086 для Base и 0.151 для дистиллированной, где меньшее значение означает лучшее сходство с оригиналом).
Для включения VC-Attention FP8 используйте следующую конфигурацию:
attention_config:
backend: TRTLLM
quant_attention_config:
qk_dtype: fp8
v_dtype: fp8
vc_attention: true
vc_repair_budget: 0.005Разреженное внимание: Sol-Attention
Второй подход к ускорению внимания — Sol-Attention, который вводит разреженность (sparsity) в вычисления DiT. Этот метод контролируется параметром τ (тау). Тесты использовали ядро блочно-разреженного внимания на базе PrimsTS. Sol-Attention позволяет пропускать вычисления для менее значимых токенов, что значительно снижает нагрузку.
Для FLUX 3 Image Klein Base на разрешении 2048x2048, Sol-Attention с τ=1.0 обеспечило ускорение денормализации и декодирования автокодировщика на 1.24x по сравнению с плотным VANILLA вниманием с 16 начальными плотными шагами, и на 1.30x с 12 плотными шагами. Выбор между 16 и 12 плотными шагами зависит от ваших приоритетов: 16 шагов обеспечивают лучшую верность (fidelity) изображения, а 12 шагов — максимальную скорость. Более высокие значения τ снижают качество вывода без существенного дополнительного ускорения.
attention_config:
backend: TRTLLM
sparse_attention_config:
algorithm: sol_attn
tau: 1.0
disabled_until_timestep: 0.9604Кэширование шагов: Cache-DiT
Одной из самых эффективных техник для снижения задержки является Cache-DiT. Этот метод позволяет переиспользовать промежуточные выходные данные трансформатора между шагами денормализации, не требуя переобучения или изменения весов модели. Порог остаточной разницы (residual_diff_threshold) контролирует, когда кэшированные данные могут быть повторно использованы, позволяя вам балансировать между скоростью и качеством.
Результаты впечатляют: при пороге 0.24, Cache-DiT почти вдвое сократил задержку NVFP4 Base T2I генерации на обоих протестированных разрешениях. На разрешении 1024x1024 ускорение составило 1.97x для Base-модели и 1.26x для дистиллированной. На разрешении 2048x2048 ускорение было 1.94x для Base и 1.20x для дистиллированной модели. При этом LPIPS scores остались очень низкими (0.033 для Base и 0.022 для дистиллированной), что свидетельствует о минимальной потере качества.
Конфигурация Cache-DiT выглядит следующим образом:
cache_config:
cache_backend: cache_dit
Fn_compute_blocks: 1
Bn_compute_blocks: 0
max_warmup_steps: 4
max_continuous_cached_steps: 3
residual_diff_threshold: 0.2405Многопроцессорное масштабирование: CFG и Ulysses
Для еще более высокой производительности, особенно в серверных средах, TensorRT LLM VisualGen поддерживает масштабирование на нескольких GPU. Это достигается за счет использования параллелизма последовательностей Ulysses и guidance без классификатора (CFG). В режиме Base положительные и отрицательные ветви CFG назначаются на отдельные группы GPU, а последовательность каждой ветви разделяется с помощью Ulysses. Дистиллированная модель, не использующая CFG, применяет Ulysses ко всей последовательности.
На примере восьми GPU B200 в конфигурации DGX B200, масштабирование с одного до восьми GPU сократило медианную задержку T2I для Base-модели с 14.90 до 3.00 секунд, а для дистиллированной — с 2.73 до 1.08 секунд. Ускорение этапа денормализации было еще более значительным: 6.01x для Base и 4.44x для дистиллированной. Однако стоит отметить, что примерно половина времени генерации дистиллированной модели на восьми GPU остается вне этапа денормализации, что ограничивает общий прирост.
Для настройки многопроцессорного вывода на восьми GPU используйте следующую конфигурацию:
parallel_config:
cfg_size: 2
ulysses_size: 4Для дистиллированной модели установите cfg_size: 1 и ulysses_size: 8. Также важно отметить, что на меньших разрешениях (например, 1024x1024) CUDA Graphs дополняют параллелизм, снижая накладные расходы на запуск ядер. Добавление CUDA Graphs к компиляции сократило задержку Base с 1.823 до 0.809 секунд (ускорение 2.25x) и дистиллированной с 0.791 до 0.366 секунд (ускорение 2.16x) на восьми GPU.
cuda_graph_config:
enable: true06Итоговая производительность и практическое применение
Комбинация всех описанных оптимизаций на одном GPU B200 при разрешении 2048x2048 дает колоссальный результат. По сравнению с неоптимизированным инференсом BF16 VANILLA, компиляция, NVFP4 с ядрами PrimsTS, Sol-Attention и Cache-DiT обеспечили ускорение генерации T2I в 3.46 раза (сокращение времени с 22.36 до 6.46 секунд) и ускорение I2I в 4.45 раза (сокращение с 54.64 до 12.28 секунд). Для дистиллированной модели 16-шаговый вариант достиг ускорения до 2.34x.
Эти оптимизации также работают на потребительских и профессиональных видеокартах. На одном GPU NVIDIA RTX PRO 6000 Blackwell Server Edition, использование компиляции, NVFP4 с ядрами CuTe DSL и Cache-DiT дало ускорение от 1.73x до 3.25x, что расширяет применимость этих методов за пределы дата-центров.
07Что это значит на практике
Для разработчиков и энтузиастов, работающих с FLUX 3, эти данные означают, что теперь возможно запускать высококачественную генерацию изображений в реальном времени или близко к нему. Если вы используете локальную машину с мощным GPU (например, RTX PRO 6000 или будущими потребительскими Blackwell), вы можете ожидать ускорения в 2-3 раза по сравнению с базовой реализацией, используя только NVFP4 и Cache-DiT. Для серверных развертываний с несколькими GPU B200, вы можете обрабатывать десятки запросов в секунду, обеспечивая отзывчивый интерфейс для пользователей.
Рекомендуемый путь для начала:
- Запустите предварительно собранный контейнер TensorRT LLM с поддержкой FLUX 3.
- Начните с конфигурации NVFP4 и включите torch.compile.
- Поэкспериментируйте с Cache-DiT, так как оно дает наибольший прирост для Base-модели с минимальной настройкой.
- Если качество критично, добавьте Sol-Attention с 16 плотными шагами. Если важна скорость — 12 шагов.
- Для максимального ускорения на нескольких GPU настройте Ulysses parallelism и CUDA Graphs.
Вы можете запустить пример FLUX 3, используя предоставленные YAML-конфигурации, и исследовать API рантайма TensorRT LLM VisualGen для тонкой настройки под свои задачи. Все инструменты доступны в репозитории TensorRT LLM GitHub, где вы можете сообщать об ошибках и делиться результатами с сообществом.
Оптимизация генерации изображений — это не просто техническая задача, это ключ к созданию доступных и быстрых AI-инструментов. Используя архитектуру Blackwell и библиотеку TensorRT LLM, NVIDIA предоставляет разработчикам мощный набор средств для вывода AI-моделей на новый уровень производительности.
Источник: NVIDIA Developer ↗