RT-Lynx: Colocando a Esparsidade GEMM da Maneira Correta para Modelos de Difusão

Resumo

Transformadores de Difusão (DiT) alcançam desempenho robusto na geração de imagens, mas incorrem em custos substanciais de inferência. Embora trabalhos anteriores tenham reduzido esse custo por meio de quantização e destilação, a esparsidade semi-estruturada, que pode reduzir quase pela metade os FLOPs, permanece pouco explorada. Uma razão principal é que a maioria das abordagens existentes foca na esparsificação de pesos, e podar 50% dos pesos pode remover capacidade crítica do modelo e degradar a qualidade da geração. Nosso estudo, no entanto, mostra que as ativações do DiT são intrinsecamente esparsas e significativamente mais robustas à esparsificação semi-estruturada N:M do que os pesos. Motivados por essa observação, defendemos uma mudança de paradigma da esparsificação de pesos para a esparsificação de ativações. Propomos o RT-Lynx, que aplica esparsificação N:M às ativações e incorpora técnicas de compensação de erro para mitigar a perda de precisão. Além disso, implementamos kernels CUDA altamente otimizados adaptados a essa configuração, alcançando uma aceleração de até 1,55x em média nas camadas lineares. Extensos experimentos em múltiplos modelos de difusão demonstram que nosso método preserva a qualidade de geração dos modelos originais enquanto acelera substancialmente a inferência.

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.