PulseCol: Periodically Refreshed Column-Sparse Attention for Accelerating Diffusion Language Models

TL;DR

PulseCol introduces a novel column-sparse attention mechanism that reuses sparse patterns across denoising steps, significantly accelerating diffusion language models while preserving quality.

Contribution

It proposes a finer-grained, periodically refreshed column-sparse attention method that improves efficiency and speedup in diffusion language models compared to prior sparse attention techniques.

Findings

PulseCol achieves up to 1.95× speedup over FlashAttention.

It maintains model quality while increasing sparsity and efficiency.

Reuses sparse patterns across iterations for better acceleration.

Abstract

Inference in diffusion large language models (dLLMs) is computationally expensive, as full self-attention must be repeatedly executed at each step of the denoising process without KV cache. Recent sparse attention methods for dLLMs mitigate this cost via block-sparse computation, which is applied only in later iterations when model performance is less sensitive to coarse-grained sparse approximation, but yields limited improvements in computational efficiency and acceleration. This motivates a finer-grained sparsification strategy that can be applied from earlier iterations and leverages reusable sparsity patterns, enabling further efficiency gains. In this work, we introduce PulseCol, a periodically refreshed column-sparse attention method for accelerating diffusion language models. PulseCol replaces coarse block-level sparsity with a finer-grained column-sparse structure, allowing important attention interactions to be retained more precisely while exposing greater sparsity. Built on this column-level formulation, PulseCol further identifies sparse patterns at the early denoising step and reuses them across subsequent iterations, refreshing them only at a small number of intermediate steps to track the evolution of sparse attention patterns during denoising. Experiments show that PulseCol achieves higher sparsity and greater practical speedup than prior sparse attention methods for dLLMs, while maintaining model quality. Enabled by optimized GPU kernels for column-sparse attention, PulseCol delivers up to 1.95$\times$ end-to-end speedup over FlashAttention across several context lengths.