Source-linked AI summary
WaveAlign: Cache-Aware Query-Row Scheduling for Sparse Attention in Long-Video Generation
Zijian Dai, Sen Han, Youhui Bai, Shannon Wang, Kan Wu, Jingkai Huang, Yuhang Wang, Jing Li, Cheng Li
TL;DR
Long-video DiTs make attention costly over very long token sequences, and dynamic sparse attention often achieves less speedup than its computation reduction suggests. WaveAlign reorders query rows to improve K/V locality while preserving sparse-attention semantics. It substantially improves cache hit ratios and reduces HBM traffic, with kernel and end-to-end speedups across the evaluated workloads.
Problem
Dynamic sparse attention reduces computation while preserving generation quality, but its realized speedup falls short of the reduction in computation.
Source: §I · p. 1Method
WaveAlign combines low-rank SVD grouping of rows with similar K/V accesses and within-wave sorting by K/V-block count, skipping reordering when its expected gain cannot cover its overhead.
Source: §I · p. 2Results
L2 cache hit ratio rises from 28.48%–36.35% to 79.38%–89.06%, HBM read traffic falls by 80.71%–92.11%, and speedup reaches 1.25× for kernels and 1.17× end-to-end.
Source: §I · p. 2Takeaways & Limitations
Across multiple models and algorithms, WaveAlign reduces video-generation latency by up to 1.17× and sparse-attention latency by up to 1.25× without modifying the sparse-attention algorithm.
Source: §VIII · p. 12
Abstract
from arXiv · showhide
Long-video generation with diffusion transformers (DiTs) produces extremely long token sequences, making attention a dominant inference bottleneck. Dynamic sparse attention reduces computation, but its realized speedup remains limited because irregular query-row execution degrades L2 cache locality and increases HBM traffic. We present WaveAlign, a lightweight, cache-aware query-row reordering framework for dynamic sparse attention. WaveAlign formulates row ordering as an optimization problem and approximates it with two stages. The first stage derives a low-rank SVD representation of sparse-mask rows and groups query rows with similar K/V access patterns, increasing K/V overlap among concurrently scheduled rows. The second stage exploits streaming GPU scheduling by sorting rows within each wave in descending order of their K/V-block counts, so that short rows from the current wave are followed by long rows from the next. This aligns K/V accesses across wave boundaries and enables shared blocks to be reused before eviction. An adaptive skip module avoids unprofitable reordering. By only permuting query and mask rows, WaveAlign preserves sparse-attention semantics and requires no changes to existing methods or backend kernels. Across two GPU architectures, two video DiTs, and four sparse-attention methods, WaveAlign raises the L2 cache hit ratio from 28.48%--36.35% to 79.38%--89.06%, reduces HBM read traffic by up to 92.11%, and achieves up to 1.25x kernel and 1.17x end-to-end generation speedup without quality loss.
I. INTRODUCTION
Long-video DiTs face costly attention over very long token sequences, while dynamic sparsity often realizes less speedup than its computation reduction predicts. WaveAlign reorders query rows to improve locality, yielding substantial cache, traffic, and speedup gains while preserving quality.
A. Video Diffusion Transformers (DiTs)
- A 401-frame 720p Wan2.1-14B video takes 10,678 seconds on one H100, with attention accounting for 85.86% of generation time.Longer video sequences increase with both spatial resolution and frame count, while attention has quadratic complexity in sequence length.Source: §I · p. 1
I. INTRODUCTION
- Dynamic sparse attention reduces computation while preserving generation quality, but its speedup reaches only 80% of the expected level across Wan2.1 and LTX2.3.The gap indicates untapped optimization potential in existing GPU implementations.Source: §I · p. 1
- Irregular K/V accesses and workload imbalance reduce L2 reuse, saturate HBM bandwidth, and shift sparse attention from compute-bound to memory-bandwidth-bound execution.Rows in a wave can access scattered blocks whose union overflows L2; uneven row workloads can interleave adjacent waves and separate accesses to shared blocks.Source: §I · p. 1
- Existing locality-aware sparse-matrix reordering methods do not directly fit online sparse attention: cheaper methods improve locality little, while stronger ones incur excessive preprocessing overhead.The design challenge is to capture irregular K/V-pattern similarities and mitigate workload imbalance without letting reordering overhead outweigh kernel savings.Source: §I · p. 1
- WaveAlign combines low-rank SVD row grouping by K/V-access similarity with within-wave sorting by descending K/V-block count, and skips reordering when its expected gain cannot amortize overhead.It only permutes query and mask rows, restoring output order to preserve sparse-attention semantics without changing sparse methods or backend kernels.Source: §I · p. 2
- Across evaluated workloads, WaveAlign raises L2 cache hit ratio from 28.48%–36.35% to 79.38%–89.06%, reduces HBM read traffic by 80.71%–92.11%, and yields up to 1.25× kernel and 1.17× end-to-end speedup.Generation quality matches the corresponding baselines.Source: §I · p. 2
B. Sparse Attention for Long-Video Generation
Dynamic sparse attention skips query–key/value interactions to reduce computation while retaining generation quality. Its GPU execution maps query rows to SMs, where shared L2-cache reuse depends on overlapping accesses and execution timing.
B. Sparse Attention for Long-Video Generation
- Sparse attention executes only the query-block and K/V-block interactions retained by a binary mask, reducing computation while preserving generation quality.Each mask row represents a query block, and each retained entry denotes an executed interaction.Source: B. Sparse Attention for Long-Video… · pp. 2–3
C. Mapping Sparse Attention onto GPUs
- Each query row is assigned to an SM, and concurrent query-row tasks across SMs form a wave; shared K/V blocks can be reused through L2 if they remain resident.The illustrated H100 mapping shows four of its 132 SMs.Source: C. Mapping Sparse Attention onto GPUs · p. 3, Figure 1 · p. 3
- Rows with different K/V-block counts finish at different times, allowing freed SMs to start rows from the next wave before the current wave completes.Access overlap, row workloads, and execution order jointly affect L2 hit ratio and HBM traffic.Source: C. Mapping Sparse Attention onto GPUs · p. 3
- Sparse attention eliminates many dense-attention FLOPs, yet implementations often fail to achieve proportional speedups, motivating measurement of the gap and its causes.The analysis proceeds from the speedup gap to the memory-bandwidth bottleneck and its root causes.Source: C. Mapping Sparse Attention onto GPUs · p. 3
A. FLOP Reduction Does Not Yield Proportional Speedup
Dynamic sparse attention does not convert FLOP reduction into proportional kernel speedup. Profiling attributes the gap to degraded cache locality and increased HBM traffic, which shift sparse kernels into a memory-bandwidth-bound regime.
A. FLOP Reduction Does Not Yield Proportional Speedup
- The evaluation compares four dynamic sparse-attention methods on Wan2.1 and LTX2.3 workloads, using H100 with FlashAttention4 and sequence lengths of 148K–401K tokens across the models.The analyzed 364K-token Wan2.1 workload is profiled with Nsight Compute.Source: A. FLOP Reduction Does Not Yield… · p. 3
- The compute-proportional upper bound exceeds measured sparse-kernel speedup by 1.12–1.26× on LTX2.3 and 1.09–1.26× on Wan2.1, with the gap generally increasing at longer sequence lengths.MInference and FlexPrefill have smaller gaps than fully content-dependent XAttention and SpargeAttn.Source: A. FLOP Reduction Does Not Yield… · p. 3
B. Sparse Attention Becomes Memory-Bandwidth-Bound
- All four dynamic sparse methods show reduced L2 locality and increased HBM bandwidth saturation relative to dense attention; MInference and FlexPrefill retain better locality than XAttention and SpargeAttn.XAttention and SpargeAttn approach the locality floor of a synthetic random mask.Source: B. Sparse Attention Becomes… · pp. 3–4, Table I · p. 4
- Roofline analysis finds dense attention compute-bound but all sparse methods memory-bandwidth-bound, because poor L2 locality prevents HBM traffic from falling proportionally with FLOPs.This limits kernel speedup and produces the gap between compute-proportional and measured speedups.Source: B. Sparse Attention Becomes… · p. 4, Figure 3 · p. 4
C. Why Access Is Irregular: Two Root Causes
Poor L2 locality arises from spatially scattered K/V accesses and workload imbalance that temporally separates accesses to shared blocks. WaveAlign uses query-row reordering to group overlapping accesses while preserving attention semantics.
C. Why Access Is Irregular: Two Root Causes
- Rows in one wave may access few common K/V blocks, so their combined working set can overflow L2 and evict blocks before reuse.Rows with different nonzero-block counts also finish at different times; streaming scheduling interleaves waves and can separate shared-block accesses enough for eviction.Source: C. Why Access Is Irregular: Two Root… · p. 4, Figure 4 · p. 4
- Spatial overlap creates cross-row reuse opportunities, while temporal alignment determines whether shared blocks are reached before eviction; failure of either condition increases HBM accesses.The sparse method fixes each row’s K/V blocks but not which rows execute together or when they reach shared blocks.Source: C. Why Access Is Irregular: Two Root… · p. 4
A. Design Rationale
- Query rows are mutually independent, so query and mask rows can be jointly reordered and outputs restored without changing attention semantics.Grouping rows with overlapping K/V accesses in the same wave increases cross-SM L2 reuse and reduces HBM traffic.Source: A. Design Rationale · pp. 4–5
- WaveAlign uses an adaptive guard before its two-stage row ordering, then inversely reorders outputs after kernel execution while retaining each row’s selected K/V blocks.Reordering is skipped when expected gains cannot amortize its overhead; the challenges include finding similar access patterns efficiently and keeping repeated reordering overhead low.Source: A. Design Rationale · p. 5
B. Problem Definition
WaveAlign formulates query-row scheduling as dividing mask rows into waves to maximize K/V overlap within each wave. It uses low-rank SVD geometry to organize similar rows for that partitioning.
B. Problem Definition
- WaveAlign partitions each head’s mask rows into waves of at most W rows to maximize K/V-block overlap within each wave.The optimization jointly chooses the row partition and ordering.Source: B. Problem Definition · p. 5
C. SVD-Based Global Spatial Ordering
- Because mask rows are correlated and their singular spectrum is top-heavy, leading SVD directions capture the dominant row geometry for low-dimensional ordering.The paper presents this as a principled way to turn the partitioning problem into a geometric one.Source: C. SVD-Based Global Spatial Ordering · p. 5
- The rank-r embedding from scaled leading left singular vectors gives the best rank-r approximation of pairwise row distances in the least-squares sense.The omitted distance contribution is small when the singular spectrum decays quickly, as is typical for these masks.Source: C. SVD-Based Global Spatial Ordering · pp. 5–6
1) Polar spectral ordering:
WaveAlign uses a rank-2 SVD embedding for coarse row neighborhoods, then sorts rows by polar angle to produce a one-dimensional order. This order places rows with similar K/V access patterns near one another.
1) Polar spectral ordering:
- Rank-2 coordinates capture the coarse row neighborhoods, while fine alignment within each wave is handled later.The paper chooses two dimensions as the smallest embedding that can be linearized by polar angle.Source: 1) Polar spectral ordering · p. 6
- Sorting rows by polar angle converts their two-dimensional SVD coordinates into a single execution order.The nonnegative singular-value weights only reparameterize the angle and do not change the resulting order.Source: 1) Polar spectral ordering · p. 6
- Angularly adjacent rows tend to share more K/V blocks, so the ordering groups rows with similar access patterns near one another.Figure 6 places each row at (u1[i], u2[i]) and uses angle to unroll the points into a sequence.Source: 1) Polar spectral ordering · p. 6
2) Matrix-Free Centered Subspace Iteration:
WaveAlign computes the top two left singular vectors of the centered sparse mask using a matrix-free subspace power iteration. Centering preserves pairwise row distances while removing workload information from the retained coordinates.
2) Matrix-Free Centered Subspace Iteration:
- The matrix-free subspace power iteration computes the top two left singular vectors without materializing the centered mask or its dense Gram matrix.Centering is applied as a rank-1 correction to sparse matrix–vector products.Source: 2) Matrix-Free Centered Subspace… · p. 6
- Each iteration costs O(nnz), and a small fixed 5–15 iterations suffices because the angular sort is invariant to in-plane rotations and sign flips.The iteration converges geometrically to the top-2 invariant subspace.Source: 2) Matrix-Free Centered Subspace… · p. 6
- Centering preserves pairwise row distances while sending the constant workload direction to a zero singular value, leaving the retained coordinates to represent mask structure.On the raw mask, the leading singular vector instead tracks per-row workload.Source: 2) Matrix-Free Centered Subspace… · p. 6
D. NNZ-Guided Intra-Wave Temporal Alignment
Stage 1 groups rows with similar K/V access patterns, but variable row lengths can desynchronize their execution across waves and weaken cache reuse. Stage 2 sorts rows by nonzero-block count within each wave to align accesses across wave boundaries.
D. NNZ-Guided Intra-Wave Temporal Alignment
- Stage 1’s spatial ordering does not account for differing row lengths, so short rows can advance to later waves while long rows are still running.This temporal skew can separate accesses to shared K/V blocks and weaken L2 reuse.Source: D. NNZ-Guided Intra-Wave Temporal… · p. 6
- Stage 2 sorts rows within each Stage 1 wave by descending NNZ-block count before scheduling them onto the SMs.The two-stage algorithm forms consecutive buckets of at most W rows, then sorts each bucket by NNZ.Source: D. NNZ-Guided Intra-Wave Temporal… · p. 7
- Assigning long rows early and short rows late lets the latter finish near the start of the next wave, shortening reuse distances for shared K/V blocks.The paper describes this as a long-short complementary assignment across successive wave boundaries.Source: D. NNZ-Guided Intra-Wave Temporal… · p. 7
E. Runtime Coordination
WaveAlign coordinates row reordering at runtime and integrates it as a plug-in between sparse-mask construction and backend execution. The evaluation covers two GPU platforms, two video DiT families, and workloads with varied sequence lengths.
E. Runtime Coordination
- WaveAlign first groups rows by K/V-access similarity, then sorts each wave by descending nonzero count before concatenating the resulting row permutation.The adaptive skip module bypasses reordering when its estimated latency reduction does not exceed the overhead.Source: E. Runtime Coordination · p. 7
V. IMPLEMENTATION
- Two PyTorch-based operators apply the permutation to query and mask rows before attention, then restore outputs to their original row order.The operators are registered in vLLM-Omni as a plug-in between mask construction and backend-kernel execution.Source: §V · p. 7
A. Experimental Setup
- Evaluation uses FlashAttention4 on H100 and FlashInfer on A40, with Wan2.1 and LTX2.3 workloads spanning varied sequence lengths.End-to-end evaluation excludes MInference and FlexPrefill because their mask-construction overhead is prohibitively high in the T2V workloads.Source: A. Experimental Setup · pp. 7–8
B. End-to-End Video Generation Latency
WaveAlign reduces end-to-end and kernel latency across the evaluated models, sparse-attention methods, and sequence lengths. Gains are larger where cache capacity is more constrained, and reordering overhead is outweighed by kernel-time reductions.
B. End-to-End Video Generation Latency
- 1.17× and 1.14× are the maximum end-to-end speedups on Wan2.1-T2V-14B and LTX2.3-22B, respectively, across evaluated methods and sequence lengths.Sparse attention accounts for 65.8%–88.0% of baseline latency on Wan2.1-T2V-14B and 58.1%–84.4% on LTX2.3-22B.Source: B. End-to-End Video Generation Latency · p. 8
- On A40, combining sequential-head scheduling with WaveAlign yields up to 11.0% lower latency than WaveAlign with the native scheduler and consistently achieves the lowest latency.WaveAlign alone reaches up to 1.29× speedup over the native-scheduler baseline.Source: B. End-to-End Video Generation Latency · p. 8
C. Kernel Latency Breakdown
- 1.092×–1.252× kernel speedups across four methods on H100 include reordering overhead, which accounts for only 1.1%–3.2% of total WaveAlign attention latency.The reported speedup ranges are 1.092×–1.181× for MInference, 1.088×–1.179× for FlexPrefill, 1.124×–1.252× for XAttention, and 1.125×–1.234× for SpargeAttn.Source: C. Kernel Latency Breakdown · pp. 8–9, Figure 10 · p. 9
- On A40, O+Seq+WA reaches maximum speedups over Original of 1.158×, 1.090×, 1.635×, and 1.601× for MInference, FlexPrefill, XAttention, and SpargeAttn, respectively.The combined configuration consistently has the lowest total attention latency; its reordering overhead is 7.5–29.7 ms.Source: C. Kernel Latency Breakdown · p. 9
- At comparable sequence lengths, A40 gains exceed H100 gains because its 6MB L2 cache becomes capacity-constrained earlier than H100’s 50MB cache.O+WA reaches 1.599× and 1.520× speedups for XAttention and SpargeAttn on A40, respectively.Source: B. End-to-End Video Generation Latency · p. 8, C. Kernel Latency Breakdown · p. 9
D. L2 Cache Locality and HBM Traffic
WaveAlign increases L2 reuse and reduces HBM reads, with its two stages contributing complementary spatial and temporal alignment. Additional evaluations examine adaptive skipping, mask patterns, roofline behavior, and generated-video quality.
D. L2 Cache Locality and HBM Traffic
- Across four methods, L2 hit ratios rise from 28.48%–36.35% to 79.38%–89.06%, while HBM reads fall by 80.71%–92.11%.XAttention and SpargeAttn have the lowest baseline hit ratios because their content-dependent masks produce more irregular K/V accesses.Source: D. L2 Cache Locality and HBM Traffic · p. 9
- Up to 1.17× kernel speedup, 79.38%–89.06% L2 hit ratios, and up to 92.11% lower HBM traffic result from combining both reordering stages.Stage 1 establishes spatial overlap between rows, while Stage 2 aligns CTA progress to shorten reuse distances; the authors report both are necessary for effective L2 reuse.Source: Figure 12 · p. 10, D. L2 Cache Locality and HBM Traffic · p. 10
- Stage 1 alone reaches up to 1.0319× kernel speedup, while Stage 2 alone reaches 1.091×–1.147× across methods and reduces HBM reads by 68.89%–83.67%.Stage 1 raises L2 hit ratios to 34.88%–44.47%; Stage 2 raises them to 68.35%–78.71%.Source: D. L2 Cache Locality and HBM Traffic · p. 10
F. Other Factors
- Reordering is skipped at short sequence lengths and low retained densities when overhead cannot be amortized; when profitable, it reaches up to 1.18× and 1.14×–1.20× speedup, respectively.The evaluated sequence lengths were 25K, 49K, and 98K for skipping, with reordering enabled from 197K onward; density skipping occurred at 1.7% and 6.7%, with enabling from 14.5% onward.Source: Figure 13 · p. 10, F. Other Factors · p. 10
- After reordering, mask rows within an execution wave show more vertically aligned nonzero blocks, indicating greater sharing of K/V accesses across SMs.The zoomed regions span 132 query rows and 132 K/V-block columns.Source: Figure 14 · p. 10, F. Other Factors · p. 11
- WaveAlign raises L2 hit ratio to 76.72% with 19.85ms overhead and 1.210× net speedup, while the compared methods either improve locality little or incur inference slowdowns.Roofline analysis reports that WaveAlign moves all four sparse-attention kernels from memory-bandwidth-bound to compute-bound execution; quality metrics are identical before and after WaveAlign across the tested models and methods.Source: F. Other Factors · p. 11, Figure 15 · p. 11
VII. RELATED WORK
The related work situates WaveAlign alongside sparse-attention methods, GPU kernels, and sparse-matrix reordering, emphasizing its role after mask construction. The conclusion highlights locality, traffic, and latency improvements; the paper also states that AI assistance was limited to language polishing.
VII. RELATED WORK
- WaveAlign complements fixed-pattern and content-adaptive attention methods by reordering rows after mask construction without changing selected blocks.It also complements GPU attention backends, while the authors distinguish it from sparse-matrix reordering methods whose preprocessing overhead makes them unsuitable for runtime-generated masks.Source: §VII · p. 11
VIII. CONCLUSION
- WaveAlign jointly improves spatial K/V locality and temporal workload alignment, reducing end-to-end video-generation latency by up to 1.17× and sparse-attention latency by up to 1.25×.The conclusion frames these as results across multiple models and algorithms.Source: §VIII · p. 12
AI USE
- AI assistance was used only to polish language and improve readability; the authors report conducting the other listed research activities independently.Those activities include problem identification, algorithm and system design, implementation, evaluation, and results analysis.Source: AI USE · p. 13