GSPN-2: Efficient Parallel Sequence Modeling

Hongxu Yin (NVIDIA) · Xinhao Li (University of California, San Diego) · Kai Han (The University of Hong Kong) · Hongjun Wang (The University of Hong Kong) · yitong jiang (The Chinese University of Hong Kong) · Collin McCarthy (University of California, Davis) · David Wehr (NVIDIA) · Hanrong Ye (NVIDIA) · Ka Chun Cheung (NVIDIA) · Wonmin Byeon (NVIDIA Research) · Jinwei Gu (Department of Computer Science and Engineering, The Chinese University of Hong Kong) · Ke Chen (NVIDIA) · Pavlo Molchanov (NVIDIA) · Jan Kautz (NVIDIA) · Sifei Liu (NVIDIA)
2d kernelchannel-shared weightscomputational costdata transfersgeneralized spatial propagation networkglobal spatial contextgpu kernelsimage classificationline-scan propagationpropagation weightsself-attentionshared memorystructured matrix transformationstext-to-image synthesisvision transformer

Efficient vision transformer remains a bottleneck for high-resolution images and long-video related real-world applications. Generalized Spatial Propagation Network (GSPN) \cite{wang2025parallel} addresses this by replacing quadratic self-attention with a line-scan propagation scheme, bringing the cost close to linear in the number of rows or columns, while retaining accuracy. Despite this advancement, the existing GSPN implementation still suffers from (i) heavy overhead due to repeatedly launching GPU kernels, (ii) excessive data transfers from global GPU memory, and (iii) redundant computations caused by maintaining separate propagation weights for each channel. We introduce GSPN-2, a joint algorithm–system redesign. In particular, we eliminate thousands of micro-launches from the previous implementation into one single 2D kernel, explicitly pin one warp to each channel slice, and stage the previous column's activations in shared memory. On the model side, we introduce a set of channel-shared propagation weights that replace per-channel matrices, trimming parameters, and align naturally with the affinity map used in transformer attention. Experiments demonstrate GSPN-2's effectiveness across image classification and text-to-image synthesis tasks, matching transformer-level accuracy with significantly lower computational cost. GSPN-2 establishes a new efficiency frontier for modeling global spatial context in vision applications through its unique combination of structured matrix transformations and GPU-optimized implementation.