RT-Lynx: Aplicando la Esparsidad de GEMM de la Manera Correcta para Modelos de Difusión

Resumen

Los Transformadores de Difusión (DiT) logran un alto rendimiento en la generación de imágenes, pero incurren en costos de inferencia sustanciales. Si bien trabajos previos han reducido este costo mediante cuantización y destilación, la escasez semiestructurada, que puede casi reducir a la mitad los FLOPs, sigue siendo poco explorada. Una razón clave es que la mayoría de los enfoques existentes se centran en la esparsificación de pesos, y podar el 50% de los pesos puede eliminar la capacidad crítica del modelo y degradar la calidad de generación. Sin embargo, nuestro estudio muestra que las activaciones de DiT son intrínsecamente dispersas y significativamente más robustas a la esparsificación semiestructurada N:M que los pesos. Motivados por esta observación, abogamos por un cambio de paradigma de la esparsificación de pesos a la esparsificación de activaciones. Proponemos RT-Lynx, que aplica esparsificación N:M a las activaciones e incorpora técnicas de compensación de errores para mitigar la pérdida de precisión. Adicionalmente, implementamos núcleos CUDA altamente optimizados y adaptados a este contexto, logrando una aceleración promedio de hasta 1.55x en las capas lineales. Experimentos exhaustivos en múltiples modelos de difusión demuestran que nuestro método preserva la calidad de generación de los modelos originales mientras acelera sustancialmente la inferencia.

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.