RT-Lynx: Die GEMM-Sparsity auf die richtige Weise für Diffusionsmodelle einsetzen

Zusammenfassung

Diffusionstransformatoren (DiT) erzielen eine starke Leistung in der Bildgenerierung, verursachen jedoch erhebliche Inferenzkosten. Während frühere Arbeiten diese Kosten durch Quantisierung und Destillation gesenkt haben, ist die halbstrukturierte Sparsität, die die FLOPs nahezu halbieren kann, noch wenig erforscht. Ein Hauptgrund ist, dass sich die meisten bestehenden Ansätze auf die Gewichtsverdünnung konzentrieren und das Entfernen von 50 % der Gewichte kritische Modellkapazität eliminieren und die Generierungsqualität beeinträchtigen kann. Unsere Studie zeigt jedoch, dass DiT-Aktivierungen intrinsisch dünnbesetzt sind und deutlich robuster gegenüber der N:M-halbstrukturierten Verdünnung sind als Gewichte. Motiviert durch diese Beobachtung plädieren wir für einen Paradigmenwechsel von der Gewichtsverdünnung hin zur Aktivierungsverdünnung. Wir schlagen RT-Lynx vor, das eine N:M-Verdünnung auf Aktivierungen anwendet und Fehlerkompensationstechniken integriert, um den Genauigkeitsverlust zu mindern. Darüber hinaus implementieren wir hochoptimierte, auf diese Umgebung zugeschnittene CUDA-Kerne, die in linearen Schichten durchschnittlich eine bis zu 1,55-fache Beschleunigung erzielen. Umfangreiche Experimente mit mehreren Diffusionsmodellen zeigen, dass unsere Methode die Generierungsqualität der Originalmodelle bewahrt und gleichzeitig die Inferenz erheblich beschleunigt.

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.