PaperlayerView exact revision on arXiv

Source-linked AI summary

Parameterized Stripe Attention for Efficient Video Generation

Xingyu Jia, Baole Ai, Ang Wang, Kang Zhao, Yong Li

arXiv:2609.37001v1cs.CVcs.AI

TL;DR

Video DiTs incur substantial inference latency, with full spatio-temporal attention as a major computational cost and existing sparse methods facing a flexibility-efficiency trade-off. PSA encodes recurring stripe patterns in a unified mask and pairs it with an efficient kernel and offline search; experiments report end-to-end speedups over FlashAttention-3 on two video models.

  • Problem

    Video DiT inference is costly, and existing sparse methods lack a unified structural characterization that can cover diverse attention patterns while maintaining efficiency.

    Source: §1 · pp. 1–2

  • Method

    PSA parameterizes periodic diagonal stripe patterns in a unified mask, processes them with one hardware-efficient kernel, and uses training-free offline search to select per-head masks.

    Source: §1 · p. 2, §3.2 · p. 5

  • Results

    1.57× and 1.37× end-to-end speedups over FlashAttention-3 on HunyuanVideo and Wan 2.1, respectively, with acceptable quality degradation (PSNR > 23).

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    The parameterized representation and offline search support head-adapted sparse patterns while avoiding runtime search overhead, and the kernel achieves up to 61% MFU.

    Source: §5 · p. 10

    The current work is limited to video-generation models, and PSA-Search takes 2.5 hours on 8 H100 GPUs to evaluate candidate masks for every head.

    Source: A Limitations · p. 14

Abstract

from arXiv · show

Diffusion Transformers (DiTs) enable high-quality video generation but suffer from substantial inference latency, primarily attributable to the computationally expensive full spatio-temporal attention. While sparse attention methods offer potential solutions, existing approaches face an inherent flexibility--efficiency dilemma: predefined masks lack the flexibility to capture diverse attention patterns, while runtime-determined masks introduce overheads and sacrifice hardware efficiency. We identify the lack of a unified structural characterization of DiT attention as a key limitation of existing methods, and establish that video DiT attention exhibits \textbf{periodic diagonal stripe structures} along both temporal and spatial dimensions. To formally encode these structured patterns within a single efficient kernel, we present {\bf PSA}, a parameterized stripe attention that formalizes the observed stripe regularity, unifying diverse attention patterns for efficient mask generation. This unified representation enables a single hardware-efficient CUDA kernel to process all sparse patterns, achieving FlashAttention-3-level Model FLOPs Utilization. To determine optimal sparsity configurations, we propose a training-free offline search algorithm that automatically maximizes sparsity under a specified error tolerance for each attention head. Experiments on HunyuanVideo and Wan~2.1 demonstrate that PSA achieves 1.57$\times$ and 1.37$\times$ end-to-end speedups over FlashAttention-3 baselines, with acceptable visual quality degradation.

1 Introduction

Video DiTs face costly inference dominated by full spatio-temporal attention, while existing sparse methods trade pattern coverage for efficiency. PSA characterizes recurring attention structure and combines a unified mask, hardware-efficient kernel, and offline per-head search.

1 Introduction

  • DiT denoising is computationally costly, and full spatio-temporal attention accounts for over 60% of total inference time in HunyuanVideo.Source: §1 · p. 1
  • Predefined masks miss diverse patterns, whereas runtime mask generation can sacrifice hardware efficiency, motivating a unified structural characterization and hardware-aware implementation.Source: §1 · pp. 1–2
  • 1.57× and 1.37× end-to-end speedups over FlashAttention-3 on HunyuanVideo and Wan 2.1, respectively, with acceptable quality degradation (PSNR > 23).Source: §1 · pp. 2–3
  • Video DiT attention exhibits periodic diagonal stripe structures along temporal and spatial dimensions, with intra-frame, inter-frame, hybrid, and uniform patterns.Source: §1 · p. 2
  • PSA's six-parameter mask unifies diverse sparse patterns for a single hardware-efficient CUDA kernel, reaching 61% MFU and adding 3D-padding overhead of about 5% for 720p.Source: §1 · pp. 2–3
  • A training-free offline search selects each head's highest-sparsity mask under an error threshold, with distributed parallel search support and no inference-time search overhead.Source: §1 · pp. 2–3

2 Related Work

Related work on efficient video-generation attention uses either predefined sparse masks or dynamically determined masks. The paper characterizes these paradigms as limited in pattern coverage or efficiency, respectively.

2 Related Work

  • Existing video-generation attention methods use predefined mask formats or determine sparse masks dynamically from attention scores during inference.Source: §2 · p. 3
  • Predefined approaches are constrained by limited pattern sets, while dynamic mask generation adds overhead and can yield limited speedups over optimized dense attention.Source: §2 · p. 3
  • PSA is presented as uniformly parameterizing diverse sparse patterns with hardware-aware kernels to combine flexibility and efficiency.Source: §2 · p. 3

3 Method

The method characterizes video DiT attention as periodic diagonal stripes, encodes diverse patterns in a shared mask and hardware-aware kernel, and searches offline for per-head configurations.

3.1 Structural Analysis of Attention Patterns

  • Video DiT attention forms periodic diagonal stripes because local spatio-temporal concentration combines with RoPE-3D periodicity, yielding intra-frame, inter-frame, hybrid, and uniform patterns.Locality determines stripe widths, while periodicity determines recurring intervals and offsets.Source: §3.1 · pp. 3–4

3.2 Efficient and Flexible Kernel

  • The six-parameter mask represents diverse stripe patterns for one hardware-efficient kernel, with masking performed at PSABlock granularity.The representation covers patterns that account for over 60% of heads but are not captured by restricted SVG and STA configurations; the paper reports 0.9831 average cosine similarity with full attention.Source: §3.2 · pp. 5–6

3.3 Automatic and Efficient Sparsity Discovery

  • PSA-Search uses prompt-stable attention to select each head’s maximum-sparsity mask offline under a user-set L2 error threshold, avoiding runtime search overhead.The search can instead target a sparsity level and select the mask with minimum error.Source: §3.3 · pp. 6–7

4 Experiments and Results

Across Wan 2.1 and HunyuanVideo configurations, the experiments evaluate PSA against sparse-attention baselines and test its speed, quality, robustness to sparsity, and structural premise. The results include faster inference, strong generation quality, and evidence that the stripe structure and searched masks hold across settings.

4.2 PSA Performance and Video Quality

  • PSA’s kernel is approximately 1.59× and 1.63× faster than SVG and SVG2 on HunyuanVideo (768×1280), and competitive with STA.The kernel comparison uses the same mask pattern across methods; PSA also has an end-to-end advantage over STA.Source: Table 2 · p. 8, §4.2 · p. 8
  • In end-to-end inference on HunyuanVideo at 768 × 1280, PSA achieves relative speedups of 1.13×, 1.10×, and 1.06× over SVG, SVG2, and STA, respectively.Source: §4.2 · p. 8
  • PSA achieves the best VBench and frame-level fidelity metrics among sparse methods across configurations; on HunyuanVideo (720×1280), PSNR is 30.35 at 59.74% sparsity.The reported VBench metrics are overall consistency and temporal style adherence; frame-level metrics include MSE, PSNR, SSIM, and LPIPS.Source: §4.2 · p. 8

4.3 Sensitivity Tests

  • As sparsity rises from 30% to 60%, quality degrades smoothly and PSNR remains above 30 at 60%; PSA’s kernel beats SVG across settings, while STA is marginally faster but fixed-resolution.PSA’s 3D-padding supports arbitrary video resolutions with negligible overhead.Source: §4.3 · p. 8, Table 3 · p. 9

4.4 Bracketing Comparison

  • At both sparsity-bracketing endpoints, PSA dominates every baseline in speed and quality: Ours-min is fastest at 1.51×, while Ours-max has the best PSNR at 26.76.The comparison uses HunyuanVideo (768 × 1280); Ours-min computes the most and Ours-max the least among the compared operating points.Source: Table 4 · p. 9, §4.4 · p. 9

4.5 Generality of the Stripe Structure

  • Dominant-channel heads exceed 78.9% in all nine settings (median ≈90%), and the attention maps show periodic diagonal stripes across the tested models.The settings span resolutions, model sizes, architectures, modalities, tokenizations, and denoising schedules; Open-Sora v1.2’s RoPE-free spatial attention also shows locality-induced stripes.Source: Table 5 · p. 9, §4.5 · p. 9

4.6 Fidelity of Searched Masks

  • At 60% sparsity, searched masks achieve 0.9831 average cosine similarity with full attention, and over 93.6% of heads have L2 loss below 10^-2.These results assess the fidelity of PSA-Search’s sparse approximations.Source: §4.6 · p. 9

4.7 Ablation Studies

  • At fixed 60% sparsity, the full six-parameter mask reaches PSNR 30.35, 2.5 dB above the next-best inter-frame-only variant.Inter-frame-only reaches PSNR 27.81 versus 24.01 for intra-frame-only; the study attributes the full model’s gain to offset and stride parameters capturing multi-diagonal and hybrid patterns.Source: Table 6 · p. 10, §4.7 · p. 10

5 Conclusion

PSA represents video DiT attention as periodic diagonal stripes and uses that structure to unify sparse masks in one efficient kernel. Its per-head offline search maximizes sparsity under an error constraint, and experiments report end-to-end speedups over FlashAttention-3 with minimal quality degradation.

5 Conclusion

  • Video DiT attention exhibits periodic diagonal stripes along temporal and spatial dimensions, which PSA encodes with a unified parameterized mask.The mask supports a single hardware-efficient CUDA kernel, reported to achieve up to 61% MFU.Source: §5 · p. 10
  • A training-free offline search assigns each head its maximum-sparsity mask under an error constraint, without runtime overhead.Source: §5 · p. 10
  • PSA achieves end-to-end speedups of 1.57× on HunyuanVideo and 1.37× on Wan 2.1 over FlashAttention-3, with minimal quality degradation.Source: §5 · p. 10

A Limitations

The paper identifies limits in PSA’s demonstrated scope and search cost, then describes its kernel, mask-search procedure, candidate initialization, and selected sparsity levels for the evaluated models.

A Limitations

  • PSA is evaluated only for video generation, so transfer to other model families or domains remains uncertain; searching all candidate masks takes 2.5 hours on 8 H100 GPUs.The authors propose a future two-stage search to filter candidates before detailed evaluation.Source: A Limitations · p. 14

B Pseudocode for PSA-Kernel

  • PSA-Kernel splits each GPU block’s work between producer and consumer warpgroups, with double-buffering to hide memory latency behind computation.The producer loads required KV blocks while consumers compute attention and accumulate results.Source: B Pseudocode for PSA-Kernel · p. 14

C.1 Pseudocode for PSA-Search

  • PSA-Search evaluates candidate masks by L2 loss per head, then chooses the highest-sparsity mask within the user’s loss threshold.Previously recorded losses allow threshold changes without repeating the costly loss evaluation.Source: §C.1 · pp. 14–15

C.2 Initialization of Mask Candidates.

  • Candidate masks are initialized from visualized attention heatmaps, with ranges for inter-frame and intra-frame widths, offsets, and strides set for the model’s dimensions.Infinite stride corresponds to a single diagonal on the relevant axis.Source: §C.2 · pp. 15–16

C.3 Wan 2.1 14B

  • For Wan 2.1 14B, PSA-Search reaches 57.57% overall sparsity using loss tolerance α = 0.007.Source: §C.3 · p. 16

C.4 HunyuanVideo 768×1280

  • For HunyuanVideo at 768×1280, loss threshold α = 0.018 yields 60.33% overall sparsity.Source: §C.4 · p. 16

C.5 HunyuanVideo 720×1280

  • For HunyuanVideo at 720×1280, loss threshold α = 0.03 yields 59.74% overall sparsity.Source: §C.5 · p. 16

D Locality-Induced Periodic Diagonal Stripes

Video spatio-temporal locality produces periodic diagonal stripes in the two-dimensional attention map. Their strides follow token layout, while their widths follow the locality radii; RoPE can further modulate the pattern for some heads.

D Locality-Induced Periodic Diagonal Stripes

  • Under row-major token ordering, temporal locality creates diagonal bands with stride HW, while column-wise locality creates bands with stride W.The corresponding half-widths are εt·HW and εh·W.Source: D Locality-Induced Periodic Diagonal… · p. 17
  • Row-wise locality creates a continuous band along the main diagonal, and composing the locality conditions yields periodic stripes whose strides and widths follow token layout and locality radii.Source: D Locality-Induced Periodic Diagonal… · p. 17
  • Locality alone can produce stripes without RoPE periodicity; for heads with a dominant RoPE channel, RoPE further modulates stride and introduces nonzero offset.The paper attributes stripes in the 9.70% of heads without a dominant RoPE channel to locality alone.Source: D Locality-Induced Periodic Diagonal… · pp. 17–18

E Empirical Support for m∗ λ in RoPE-3D

RoPE-channel analyses support periodic diagonal stripe structures in video DiT attention, with dominant channels associated with intra-frame and inter-frame patterns. Most Wan 2.1 heads show a dominant channel, while locality-driven and uniform heads are also handled within the stripe framework.

E Empirical Support for m∗ λ in RoPE-3D

  • Intra-frame dominant RoPE channels localize within Mw, and increasing channel index corresponds to wider spacing between adjacent stripes.This is reported as corroborating the periodicity formulation.Source: E Empirical Support for m∗ λ in RoPE-3D · p. 18
  • Inter-frame dominant channels localize within Mt, with higher channel indices corresponding to wider spacing between adjacent stripes.The observed spacing is consistent with the periodicity formulation.Source: E Empirical Support for m∗ λ in RoPE-3D · p. 18
  • Hybrid-pattern heads show pronounced dominant channels across Mt, Mh, and Mw sub-groups.The passage describes this as a superposition of distinct attention patterns.Source: E Empirical Support for m∗ λ in RoPE-3D · p. 18
  • 87.55% of Wan 2.1 attention heads have a dominance ratio exceeding 3×, supporting the single-channel dominance assumption.Among the remaining heads, 2.74% are uniform and 9.70% still exhibit diagonal stripes.Source: E Empirical Support for m∗ λ in RoPE-3D · p. 18
  • At 60% sparsity, the 9.70% of heads without a dominant RoPE channel achieve 0.9872 average cosine similarity with full attention.The passage attributes these heads' stripes to spatio-temporal locality and reports that PSA captures them.Source: E Empirical Support for m∗ λ in RoPE-3D · pp. 18–19
  • Uniform-pattern heads lack a dominant channel, but PSA-Search finds wide-stripe or checkerboard masks with cosine similarity above 0.98.These masks maintain meaningful sparsity, supporting coverage of all four observed pattern types.Source: E Empirical Support for m∗ λ in RoPE-3D · pp. 18–19

F Heatmaps of Different Heads in CogVideoX

CogVideoX heatmaps show intra-frame, inter-frame, hybrid, and uniform patterns analogous to those in HunyuanVideo and Wan 2.1. Heatmaps for the same head across different prompts have cosine similarity above 0.99.

F Heatmaps of Different Heads in CogVideoX

G Heatmap Visualization Across Different Prompts

H Loss Threshold Sensitivity Analysis

Across 10–70% sparsity, PSA consistently outperforms SVG2 in speed and quality and has a better speed–quality trade-off than SVG. The α setting controls tolerated loss and resulting sparsity, with model-specific values used to target approximately 60% sparsity.

H Loss Threshold Sensitivity Analysis

I End-to-End Performance and Quality Across Varying Sparsity Levels

  • Across 10–70% sparsity, PSA consistently outperforms SVG2 in speed/quality and achieves a better speed–quality trade-off than SVG.SVG achieves only marginal speedup at 70% sparsity while suffering severe quality degradation.Source: I End-to-End Performance and Quality… · p. 21

J Qualitative Comparison with Baselines

Across multiple prompts and models, the extracted video frames are described as nearly lossless relative to the FlashAttention-3 baseline. The checklist states that the abstract and introduction's three PSA contributions are consistent with the theoretical and experimental results in Sections 3–4.

J Qualitative Comparison with Baselines

NeurIPS Paper Checklist

  • The checklist says PSA's parameterized representation, efficient kernel, and automatic search contributions align with Sections 3–4.Source: NeurIPS Paper Checklist · p. 26

2. Limitations

The paper says its limitations are discussed in Appendix A and explains why it reports no error bars. It also describes the information and compute resources used to reproduce its experiments.

2. Limitations

  • Limitations are discussed in Appendix A.Source: §2 · p. 26
  • Reproduction materials include Algorithms 1–3, model configurations, search-parameter ranges, and kernel implementation details.Source: §2 · p. 26
  • The paper reports no error bars, explaining that evaluation randomness comes mainly from initial noise while sparse- and full-attention outputs are compared deterministically.The paper says this makes error bars unnecessary for that comparison.Source: §2 · p. 26
  • Experiments use a single NVIDIA H100 GPU, while the search takes 2.5 hours on 8 H100 GPUs.Source: §2 · p. 27

Related papers

Loading 2609.37001v1…