PaperlayerView exact revision on arXiv

Source-linked AI summary

RT-Lynx: Putting GEMM Sparsity in the Right Place for Diffusion Models

Xing Cong, Hanlin Tang, Kan Liu, Tao Lan, Lin Qu, Chenhao Xie

arXiv:2605.26632v3cs.LG

TL;DR

DiT inference is expensive, while weight sparsification can damage generation quality and activation sparsification introduces quality and runtime challenges. RT-Lynx sparsifies intermediate activations, compensates errors, and uses optimized CUDA kernels; experiments report preserved quality and substantial acceleration, including 1.55× linear-layer speedup.

  • Problem

    DiT inference has substantial cost, while weight sparsification can degrade generation quality and online activation sparsification can impose high overhead.

  • Method

    RT-Lynx applies N:M sparsification to intermediate activations, adds norm-based compensation and a lightweight LoRA branch, and uses optimized CUDA kernels for online sparse execution.

  • Results

    RT-Lynx preserves generation quality across multiple diffusion models while accelerating inference, achieving up to 1.55× linear-layer speedup.

  • Takeaways & Limitations

    Activation sparsification offers a practical alternative to weight sparsification for accelerating DiT inference while retaining model quality.

  • Takeaways & Limitations

    On single-stream paths in Z-Image and FLUX, certain linear layers had to be skipped because LoRA could not fully close the performance gap.

Abstract

from arXiv · show

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.

1 Introduction

DiT inference is costly because each generation requires many compute-intensive diffusion steps. RT-Lynx shifts semi-structured sparsification from weights to activations and combines compensation with optimized execution to accelerate inference while preserving quality.

  • DiT inference faces severe latency and energy overhead because generation requires tens of iterative diffusion steps.
  • 50% weight pruning can exceed 3% accuracy loss, and weight sparsification severely compromises DiT image-generation capability.
  • Token-wise activations are intrinsically sparse, with only a small subset of channels substantially activated within each token.
  • 1.88× Sparse GEMM speedup and 1.55× linear-layer speedup are achieved through a unified CUDA pipeline combining online N:M sparsification with sparse Tensor Core execution.
  • RT-Lynx applies activation sparsification with norm-based compensation and LoRA adaptation to recover model performance after sparsification.
  • Evaluations on mainstream DiT models demonstrate consistent acceleration with negligible quality degradation.

2 Preliminary

N:M semi-structured sparsity retains a fixed number of nonzeros in each local group, enabling hardware-efficient computation. Linear layers can apply this pattern statically to weights or dynamically to input-adaptive activations.

  • N:M sparsity retains N nonzero elements in each group of M numbers, creating regular patterns suitable for hardware decoding.
  • Sparse Tensor Cores support N:M computations and provide up to 2× theoretical acceleration.
  • A linear layer computes Y = X · W^T, where X is the input, W^T is the weight matrix, and Y is the output.
  • Weight sparsity produces a static sparse matrix fixed at inference, enabling direct use of hardware sparse kernels.
  • Activation sparsity applies N:M sparsity to activations at runtime, making the computation dynamic and input-adaptive.

3 Motivation

The motivation compares weight and activation sparsity in DiT. Weights lack intrinsic structured patterns and incur substantial reconstruction error, whereas activations concentrate near zero and exhibit strong intrinsic sparsity.

  • Weight-sparse DiT models perform markedly worse than the dense baseline, while native activation sparsity shows more promising generation quality.
  • DiT weights follow a quasi-Gaussian distribution spread across the normalized range rather than an intrinsic 2:4 structure.
  • Weight sparsification incurs substantial reconstruction errors across all layers because structured pruning removes salient parameters.
  • Only about 5% ∼10% of activation neurons are active, contrasting with the dense distribution of model weights.
  • The TopK operator keeps 2 elements per group to impose a 2:4 structured pattern on activations.

4 Methodology

RT-Lynx addresses activation-sparsification quality loss and system overhead through norm compensation, LoRA residual recovery, selective layer handling, and fused sparse execution. Its pipeline combines online pruning with optimized CUDA processing while retaining dense computation where needed.

  • Motivation: Activation sparsification can cause mild blurring and exceed 40% of runtime without optimization, motivating compensation and kernel-level acceleration.These issues hinder a lossless, end-to-end sparsification pipeline.
  • 4.1 Norm-Compensated Sparsification: Norm compensation rescales pruned activations to preserve the original l2-norm with negligible computational overhead.The method retains the largest entries under a 2:4 constraint and restores their magnitude to match the dense activation.
  • 4.2 LoRA Adaptation and Fine-Tuning: A lightweight LoRA branch models high-frequency residuals from low-magnitude activations and adds them to the sparse output.The final output combines the norm-compensated sparse result with a low-rank residual using LoRA matrices, with R = 64.
  • 4.3 Selective Layer Skipping in Single-stream DiT: Single-stream DiT paths retain a performance gap that LoRA cannot fully address, so RT-Lynx skips selected linear layers in Z-Image and FLUX.The skipped layers differ by model: attn.o_proj and mlp.up for Z-Image, and attn.o_proj and mlp.down for FLUX.
  • 4.4 CUDA Kernel Optimization: The proposed pipeline integrates sparse computation with dense processing, avoiding isolated LoRA execution and its associated memory traffic and synchronization.Figure 5 contrasts separated conventional pruning, formatting, and computation with the unified proposed flow.
  • 4.4 CUDA Kernel Optimization: The online execution framework fuses pattern determination, Top-K selection, and compression into one CUDA path that produces SpTC-compatible 2:4 activations.Its sparse GEMM uses a streamK-style block-parallel pipeline to improve sparse compute efficiency.

5 Experiments

Experiments across diverse DiT architectures and datasets show that RT-Lynx preserves generation quality while accelerating sparse inference. Ablations and compatibility tests further examine its components, overhead, and integration with other acceleration methods.

  • Setups: RT-Lynx is evaluated on Qwen-Image, FLUX.1, and Z-Image using 20k training prompts and 5,000 evaluation prompts each from MJHQ-30K and sDCI.The evaluation covers heterogeneous model scales and structures, with FID, Image Reward, CLIP-IQA, and CLIP-Score measuring generation quality.
  • Accuracy: Weight sparsification substantially degrades Qwen-Image quality, whereas RT-Lynx with R = 64 consistently outperforms competing methods and can surpass the FP16 model on both benchmarks.Its generated results are described as nearly indistinguishable from the original model.
  • Accuracy: FID 21.25 is recovered by combining Sparse Activation, Norm Compensation, and LoRA, while Layer Skipping further optimizes generation for single-stream models.Qualitative results report artifact mitigation across model scales.
  • Compatibility: RT-Lynx preserves comparable quality when combined with step distillation, W8A8 quantization, feature caching, and sparse attention; with 8-step Z-Image, FID is 26.17 versus 25.70 for the distilled baseline.Traditional weight pruning produces FID 360.2 in the same distilled setting.
  • SpeedUp: Online sparsification overhead is limited to less than 10%, enabling up to 1.88× speedup; aggregate Qwen-Image linear-layer time falls from 6.74ms to 4.83ms.The MLP Down projection decreases from 2.22ms to 1.39ms.
  • End-to-End: End-to-end speedups are about 1.2×, with linear layers reaching 1.43×–1.55× and Qwen-Image generation time decreasing from 0.75s to 0.62s.The reported gains are attributed primarily to acceleration of computation-dominant Linear layers.

6 Related Work

Related work spans weight structured sparsity, under-explored activation sparsity, and broader diffusion acceleration techniques. Existing activation methods often target settings that do not directly apply to DiTs, motivating RT-Lynx’s focus.

  • Model Sparsity for LLMs: LLM sparsification methods such as SparseGPT, Wanda, RIA, BaWA, and Slim primarily impose structured constraints on weights for hardware-efficient execution.Training-free approaches commonly determine masks through one-shot calibration statistics.
  • Activation Sparsity: Activation sparsity remains less explored, and prior decoding-focused methods are limited to cases with fewer than 4 tokens, unlike DiTs with typically more than 1,000.The cited related works most relevant to this setting are An et al. and Haziza et al.
  • Diffusion Acceleration: Diffusion acceleration research includes quantization, distillation, sparse attention, and feature caching, with representative work spanning 4-bit and 1.58-bit quantization.Other approaches replace or approximate attention or cache features across diffusion steps.

7 Conclusion

RT-Lynx shifts DiT semi-structured sparsification from weights to activations and combines compensation methods with optimized CUDA kernels. Across multiple diffusion models, it preserves generation quality while accelerating inference.

  • Conclusion: RT-Lynx applies N:M sparsification to intermediate activations, using norm-based compensation, a lightweight LoRA branch, and layer skipping to reduce sparsification error.The approach is designed as an end-to-end DiT sparsification solution.
  • Conclusion: Optimized CUDA kernels generate and exploit online activation sparsity during inference, achieving up to a 1.55× average speedup in linear layers.The system targets practical execution of activation sparsity rather than only theoretical FLOP reduction.
  • Conclusion: Extensive experiments show that RT-Lynx preserves generation quality while substantially accelerating inference across multiple diffusion models.The paper characterizes this as lossless N:M sparsification for DiT models with substantial speedups.

Impact Statement

The paper positions RT-Lynx as an efficiency-oriented use of activation sparsity and hardware-supported semi-structured execution. The supplied passages describe its compatibility with other acceleration methods and its evaluation setup.

  • Impact Statement: RT-Lynx aims to reduce diffusion-model computational cost, latency, and energy consumption while preserving generation quality.The stated scope includes combination with quantization and distillation.
  • Sparse Tensor Core: Sparse Tensor Cores support semi-structured sparsity through fixed, hardware-predictable patterns rather than arbitrary sparse layouts.The 2:4 pattern is identified as the most representative practically deployed format.
  • Setups: Figure 10 compares weight sparsity with naive activation sparsity using Top-K selection for 2:4 sparsity and reports Image Reward for generated images.The comparison concerns visual impact on generation quality.
  • Setups: The method’s image-branch sparsity configuration is used throughout the paper, while applying sparsity to both image and text branches shows a clearer advantage over weight sparsity.The experimental setup excludes dynamic feature-modulation modules and focuses on text and image data flows.

C Additional Results of Rethinking Sparsity

Additional analyses find that DiT activations are more concentrated near zero and incur lower sparsification error than weights, while LoRA rank 64 offers a favorable quality–overhead trade-off.

  • Activations remain highly concentrated around zero across most Linear layer types, whereas weights generally have more uniform distributions.
  • Under the same 2:4 sparsity, activation sparsification removes mainly low-magnitude components and produces lower instantaneous representation error than weight sparsification.
  • Weight sparsification yields consistently higher mean squared errors across layer types and depths, while activation sparsification produces lower, more stable errors.
  • LoRA rank selection balances generative fidelity against adaptation cost, making rank a central hyperparameter for sparse inference.
  • 0.7306% parameter overhead and 9.808% relative time overhead make R = 64 the reported sweet spot for Qwen-Image.At R = 64, the method reports FID (21.25), Image_Reward (1.304), and CLIP.score (27.33) on MJHQ after 2k steps.

E Relative Error Analysis of Activation Sparsity in a Text–Image Hybrid DiT

The study evaluates activation-sparsification sensitivity across layer types in FLUX.1-dev and Z-Image under full text-to-image inference. Mixed text–image layers are especially sensitive, with Out-Down[mix] reaching RFE = 0.275 in FLUX.1-dev.

  • Relative Error Analysis: Activation sparsification produces pronounced layer-wise sensitivity in both FLUX.1-dev and Z-Image during full text-to-image generation.The analysis measures average relative Frobenius error across different Transformer linear layer types.
  • Relative Error Analysis: Mixed text–image layers consistently exhibit higher RFE than other layer types, indicating vulnerability in cross-modal shared representations.This pattern appears across the evaluated hybrid double/single-stream and single-stream architectures.
  • Relative Error Analysis: 0.275 RFE is attained by the Out-Down[mix] layer in FLUX.1-dev, the highest reported value in that model.The value is measured under the full text-to-image inference pipeline.

F.5 Implementation Details

The implementation fixes inference settings across models, applies sparsification to selected image-related layers, and evaluates compatibility with LoRA, FP8 quantization, and distillation. Across Qwen-Image-2512, FLUX.1-dev, and Z-Image, the proposed method outperforms weight-sparsification baselines while preserving visual quality.

  • Implementation Details: All models use seed 42 and 1024 × 1024 images unless otherwise specified, with model-specific guidance scales and inference-step counts.Qwen-Image uses cfg_scale 4.0 and 40 steps; FLUX.1-dev uses 2.0 and 50; Z-Image uses 1.0 and 8.
  • Implementation Details: The sparsification replacement rule targets image-related layers in Qwen-Image and the initial double-stream stage of FLUX.1-dev.The restriction to image branches is discussed with its empirical implications in Appendix B.
  • Implementation Details: LoRA fine-tuning uses a teacher–student objective with rank 64 to balance accuracy and inference efficiency.The student is trained to align its outputs with the teacher by minimizing their discrepancy.
  • Compatibility Validation: The method is evaluated for compatibility with symmetric FP8 activation quantization, FP8 weights, and the already distilled Z-Image-turbo model.The distillation experiment assesses whether activation sparsification can add acceleration on top of existing distillation gains.
  • Experimental Results: The proposed method consistently outperforms Wanda, RIA, BaWA, and Slim across Qwen-Image, FLUX.1-dev, and Z-Image.The paper attributes weight-pruning sensitivity to model-weight fluctuations and architectural differences, with significant degradation reported for some models such as Z-Image.

H Supplemental Results for SpeedUp

RT-Lynx improves sparse GEMM efficiency while reducing online sparse overhead, and activation sparsity lowers latency across major linear components in multiple diffusion models. The speedup pattern generalizes across aspect ratios and architectures.

  • GEMM Backend Efficiency: 1.79× is the maximum speedup achieved by RT-Lynx over PyTorch-SpMM and cuSparseLt on RTX 4090 GPUs.The kernel shows acceleration patterns similar to those on enterprise-grade H20 hardware.
  • GEMM Backend Efficiency: 4.38% is the lowest reported Sparse-Cost for RT-Lynx, substantially reducing online sparse execution overhead.This reduction addresses metadata-processing overhead that can otherwise bottleneck sparse computation.
  • Model-Level Speedup: Activation sparsity systematically reduces latency in MLP projections and attention-related QKV and output layers across FLUX.1-dev and Z-Image.The sparse execution path outperforms dense execution at every evaluated layer.
  • Model-Level Speedup: The latency advantage persists across varying aspect ratios and the evaluated layer configurations of Qwen-Image, FLUX.1-dev, and Z-Image.Figure 19 covers 60 Qwen-Image layers, 56 FLUX.1-dev layers, and 30 Z-Image layers.
Loading 2605.26632v3…