RT-Lynx: Правильное применение разреженности GEMM для диффузионных моделей

Аннотация

Диффузионные трансформеры (Diffusion Transformers, DiT) демонстрируют высокую производительность в задачах генерации изображений, однако сопряжены со значительными вычислительными затратами при инференсе. Предшествующие работы снижали эти затраты с помощью квантования и дистилляции, но полуструктурированная разреженность, способная почти вдвое сократить количество FLOPs, остаётся малоизученной. Ключевая причина этого заключается в том, что большинство существующих подходов сосредоточены на разреживании весов, а прореживание 50% весов может привести к потере критической ёмкости модели и ухудшению качества генерации. Однако наше исследование показывает, что активации DiT по своей природе разрежены и значительно более устойчивы к N:M полуструктурированному разреживанию, чем веса. Основываясь на этом наблюдении, мы предлагаем сдвиг парадигмы от разреживания весов к разреживанию активаций. Мы представляем RT-Lynx, который применяет N:M разреживание к активациям и включает методы компенсации ошибок для уменьшения потери точности. Кроме того, мы реализуем высокооптимизированные ядра CUDA, адаптированные для данного сценария, что позволяет достичь в среднем ускорения до 1,55x в линейных слоях. Обширные эксперименты на нескольких диффузионных моделях демонстрируют, что наш метод сохраняет качество генерации исходных моделей, одновременно существенно ускоряя инференс.

English

Diffusion Transformers (DiT) achieve strong performance in image generation but incur substantial inference costs. While prior work has reduced this cost via quantization and distillation, semi-structured sparsity, which can nearly halve FLOPs, remains underexplored. A key reason is that most existing approaches focus on weight sparsification, and pruning 50% of the weights can remove critical model capacity and degrade generation quality. Our study, however, shows that DiT activations are intrinsically sparse and significantly more robust to N:M semi-structured sparsification than weights. Motivated by this observation, we advocate a paradigm shift from weight sparsification to activation sparsification. We propose RT-Lynx, which applies N:M sparsification to activations and incorporates error-compensation techniques to mitigate accuracy loss. We further implement highly optimized CUDA kernels tailored to this setting, achieving up to a 1.55x speedup on average in linear layers. Extensive experiments across multiple diffusion models demonstrate that our method preserves the generation quality of the original models while substantially accelerating inference.