Conceptio › Archive › arXiv CS
arXiv CSopen access

WaveAlign: Cache-Aware Query-Row Scheduling for Sparse Attention in Long-Video Generation

Zijian Dai et al. · arxiv_cs
arXiv CS · Papers · License: Open Access
Open Source ↗Direct PDF ↓
clouddistributed-computingparallel-computing
distributed computing, parallel computing, cloud

WaveAlign: Cache-Aware Query-Row Scheduling for Sparse Attention in Long-Video Generation Zijian Dai1,2 , Sen Han1 , Youhui Bai1 , Shannon Wang2 , Kan Wu1 Jingkai Huang3 , Yuhang Wang1 , Jing Li1 , Cheng Li1,2 1

University of Science and Technology of China Institute of Artificial Intelligence, Hefei Comprehensive National Science Center 3 South China University of Technology [email protected], [email protected], [email protected], [email protected], [email protected], [email protected], [email protected], [email protected], [email protected]

arXiv:2609.34814v1 [cs.DC] 28 Sep 2026

2

(blocks) for attention calculation, which significantly reduces the attention computation while preserving generation quality. Nevertheless, in practice, the speedup brought by such sparse attention methods is far below the reduction in computation, even with GPU-optimized backend kernels such as FlashAttention4 [48] and FlashInfer [47]. Across Wan2.1 and LTX2.3, the speedup can only reach 80% of the expected level. This gap reveals substantial untapped optimization potential in existing GPU implementations of sparse attention. Our profiling shows that this gap is fundamentally a memory-locality problem. Sparsification increases the L2 miss rate, collapses L2 cache reuse, and drives HBM bandwidth toward saturation, shifting attention from compute-bound to memory-bandwidth-bound execution. We trace this transition to two root causes of poor L2 cache locality. A wave comprises Q blocks executing concurrently across all GPU compute units, together with the K/V blocks they access. First, Q blocks in the same wave select spatially scattered, weakly overlapping K/V blocks, so their union overflows L2 cache and blocks are evicted before reuse. Second, Q blocks select different numbers of K/V blocks, creating workload imbalance that causes adjacent waves to interleave in execution, temporally misaligning accesses to shared blocks and extending their reuse distance. To address these problems, we present WaveAlign, a lightweight high-performance sparse attention framework. Note that the attention outputs w.r.t. the queries are computed independently. Through reordering the query rows, gathering the rows accessing similar KV blocks, WaveAlign improves L2 locality and hence improves the kernel efficiency. However, finding an effective row order is non-trivial. Existing localityaware reordering methods for sparse matrix multiplication [1], [4], [17], [43] cannot be directly applied to online sparse attention. Lower-cost methods provide little improvement in L2 locality, whereas methods that substantially improve locality incur preprocessing overhead that outweighs their kernel-time savings. A practical solution must capture similarities across irregular K/V access patterns and mitigate row-workload imbalance while providing reordering overhead from offsetting the resulting kernel gains.

Abstract—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 lowrank 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 sparseattention 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.25× kernel and 1.17× end-to-end generation speedup without quality loss.

I. I NTRODUCTION Diffusion transformers (DiT) has emerged as a dominant paradigm for high-fidelity video generation, for its strong scalability and high-fidelity synthesis capabilities [28]. Compared with image generation, video generation typically requires handling longer sequences, since the sequence length increases with both spatial resolution and the number of frames, so video generation remains very costly and time-consuming. The bottleneck mainly comes from the quadratic complexity in terms of input sequence length in the attention module [6]. For example, generating a 401-frame 720p video, corresponding to a 364K sequence length, takes 10,678 seconds using Wan2.114B [34] on a single NVIDIA H100 GPU, among which the attention accounts for 85.86%. Similar behaviors are also observed in LTX2.3 [21] and other video DiTs. Utilizing the inherent sparsity of attention [16], [18], [42], [49], each query (block) only selects its most relevant KV

1

To this end, WaveAlign proposes a cache-aware queryrow reordering method, which cast row reordering as an optimization problem and approximate its solution with a lightweight two-stage algorithm. Exploiting the observation that similarities among irregular mask rows are low-rank, Stage 1 embeds the rows via a matrix-free rank-2 SVD [11] and orders them by polar angle, grouping rows with similar K/V-access patterns. Stage 2 then sorts rows within each wave by descending K/V-block count, so that short and long rows interleave across wave boundaries and shared K/V blocks are reused before eviction. The two stages are complementary: Stage 1 creates cross-row reuse opportunities and Stage 2 realizes them before cache eviction, jointly improving L2 locality. WaveAlign only permutes query and mask rows and restores the original output order, preserving sparse-attention semantics without modifying the sparse algorithms or backend kernels. Moreover, WaveAlign introduces an automatic discrimination mechanism that bypasses query-row reordering when the expected gain cannot amortize its overhead, such as for short sequences or highly sparse masks. We implement WaveAlign in vLLM-Omni [32] and evaluate it on two GPU architectures, using FlashAttention4 as backend kernel on NVIDIA H100 and FlashInfer on NVIDIA A40. Our evaluation spans two widely used video DiT families, Wan2.1 [35] and LTX2.3 [21], and four representative dynamic sparse-attention methods. Across these workloads, WaveAlign increases the L2 cache hit ratio from 28.48%– 36.35% to 79.38%–89.06% and reduces HBM read traffic by 80.71%–92.11%, shifting sparse attention from the memorybandwidth-bound regime back into the compute-bound regime. This improved memory efficiency yields up to 1.25× kernel speedup, which translates into up to 1.17× end-to-end generation speedup. Because WaveAlign preserves sparse-attention semantics, it maintains the same generation quality as the corresponding baselines.

with a causal mask, video DiTs apply full, non-causal selfattention over the entire sequence. Inference iteratively refines the latent noise over tens of denoising steps, each requiring a complete forward pass [13]. Because hidden states change at every step, the derived K/V tensors are transient and cannot be cached across steps, forcing full long-sequence attention to be recomputed throughout inference. Computation bottleneck of DiT. Despite their generation quality, DiTs are painfully slow. Using Wan2.1 T2V 1.3B [33], a popular open source DiT model, generating an 81 frame video at a resolution of 720×1280 takes over five minutes on an H100 GPU. At 1920×1080, the generation time increases dramatically to approximately 25 minutes [2]. Our profiling shows that attention dominates the generation time, accounting for 73% and 90% in the two settings above. This overhead stems from the quadratic complexity of attention with respect to sequence length [6]. Video sequences are exceptionally long and grow with both frame count and spatial resolution, reaching 76K and 364K tokens in these settings. Moreover, bi-directional attention and iterative denoising incur this full quadratic cost at every layer and denoising step. Other popular video DiTs, such as LTX2.3 [21], exhibit similar behavior. Therefore, reducing attention cost is essential for efficient long, high-resolution video generation. B. Sparse Attention for Long-Video Generation To mitigate the high cost of full attention, recent studies exploit the inherent sparsity of attention mechanisms [39]. Only a small subset of key/value tokens typically receives substantial attention weights and materially affects the output. Identifying these tokens and computing only the corresponding interactions can substantially reduce computation while preserving generation quality [30]. Consequently, recent works have increasingly favored dynamic-pattern methods, as important interactions vary across inputs. These methods estimate relevance from the online attention statistics and select important K/V tokens for each query at runtime, offering broader applicability and better accuracy preservation. MInference [16] instantiates headspecific A-shape, vertical-slash, or general sparse patterns; FlexPrefill [18] further adapts the pattern and budget to each input and head; XAttention [42] scores tokens via antidiagonal aggregation; and SpargeAttn [49] predicts relevance from pooled Q/K features and refines it via online softmax statistics. Block-sparse attention for modern GPUs. To exploit the GPU memory hierarchy as shown on the left of Figure 1, memory accesses should be contiguous, while on-chip computation is naturally tiled over data blocks [5], [6]. GPU-oriented dynamic sparse attention therefore organizes the Q, K, and V tensors and their computation at block granularity, with each block containing tens to hundreds of consecutive tokens. Based on this organization, dynamic sparse attention encodes block-level computation using a binary mask M . Dense attention computes each query block Qi with all K/V block pairs to produce Oi , whereas sparse attention executes only the interactions retained by M . As illustrated on the right of

II. BACKGROUND A. Video Diffusion Transformers (DiTs) DiT models. Video generation powers a growing range of applications, from entertainment and advertising to virtual reality [37], [41]. Video diffusion transformers (DiTs) [22], [35], [46], which pair diffusion models with Transformer backbones [28], have become the dominant architecture, as they scale well and excel at capturing long-range spatiotemporal dependencies [35], [46]. Popular video DiTs share the same recipe. A video is encoded into a sequence of spatiotemporal tokens, whose length grows with both the frame count and the spatial resolution. Each Transformer layer applies multi-head attention (MHA) to capture dependencies across tokens, followed by a feedforward network (FFN) [31]. MHA projects the hidden states X into queries (Q), keys (K), and values (V ), computes softmax-normalized query–key scores, and aggregates the values accordingly [31]. Unlike autoregressive large language models (LLMs), which restrict attention to preceding tokens

2

wave1 row 5 row SM6 row 4 row 7 0 SM0 SMEM SM1 / SM2 SM3 L1 Q L1 O L1 L1 registers Cache ₀Cache resid ₀Cache Cache K ent V hit₁ hit₁ miss hit strea L2 Cache = 50MBmed (BW ≈ 12 TB/s)

K₀ V₀

SM 1 SMEM / registers Q O ₁ resid row₁ 0 K ent V ₀ strea row₀ 1

L2 cache hit
 K₂ V₂ K₁ V₁ bandwidth≈12 TB/s

K ₁K₀

V V₀₁

O₀

⋯

⋯

⋯

⋯

Q₇

K₇

V₇

O₇

HBM = 80GB

K ₀

Q K V

Q K₀ ₀ V ₀ O

K₀ V₀ Q₀ Q₁

med row 2

Q₂

row 3

Q₃

outer

L2 Cache = 50MB(shared by all loop row 4 SMs) HBM = 80GB (BW = 3.35 TB/s)

Q₀

Upper / Actual Speedup Gap

inner loop

wave0 row 0 row 1 row 2 row 3 K₁

K₂ SM K₃ 2 V₁ SMEM V₂ V₃ / registers Q O resid K ent V

K₄

K₅

K₆

V₄

V₅

V₆

strea med L2 cache mis bandwidth TB/s

K₇ SM 3 V₇SMEM / registers Q O resid O₀ K ent V O₁ strea med O₂

O₃

Q₄

V row 5 ₀

Q₅

row 6

Q₆

row 7 Q₇ Q Q K₁ K V₁ V

K

V

K

O₄ V O₅

Q

Q

Q

Q

K V

K

K

K

K

V

V

V

O₇

V ₁ O O O O O O O GPU architecture, block-sparse attention, and the mapping of query ₀ ₁ O

FlexPrefill

XAttention

SpargeAttn

1.25

1.15

1.05

230

287

344

Seq Length (K) (a) LTX2.3

401 148

220

292

364

Seq Length (K) (b) Wan2.1-T2V-14B

Fig. 2. Upper bound vs. actual attention kernel speedup across sequence lengths on Wan2.1-T2V-14B and LTX2.3, with each sparse-attention method using its default sparsity configuration on H100.

O₆ Q

MInference

III. M OTIVATION

Fig. 1. rows and K/V blocks onto GPU execution units. We use NVIDIA H100 as an example and show 4 of its 132 SMs for clarity.

Sparse attention eliminates a substantial fraction of the FLOPs in dense attention, yet existing implementations often fail to achieve proportional speedups. We first quantify this gap, then show that sparsification shifts the kernel from compute-bound to memory-bandwidth-bound (Section III-B), and finally identify two root causes (Section III-C).

Figure 1, each mask row corresponds to a query block. A purple block with Mij = 1 denotes an executed interaction between Qi and (Kj , Vj ), while a white block with Mij = 0 denotes a skipped interaction. For example, Q0 uses (K0 , V0 ) and (K3 , V3 ) to produce O0 , while Q3 also uses (K3 , V3 ), showing that different P query blocks may share K/V blocks. The row sum ni = j Mij , or its number of nonzero blocks (NNZ), determines the workload of Qi .

A. FLOP Reduction Does Not Yield Proportional Speedup We evaluate four representative dynamic sparse-attention methods: MInference [16], FlexPrefill [18], XAttention [42], and SpargeAttn [49]. We run them on two widely used video DiT models, Wan2.1-T2V-14B [34] and LTX2.3 [12], using an NVIDIA H100 GPU with FlashAttention4 [48] as the backend kernel. The workloads span sequence lengths of 148K–364K tokens for Wan2.1 and 230K–401K tokens for LTX2.3; detailed setups are provided in Section VI-A. For each workload, ρ denotes the retained block density, i.e., the fraction of non-skipped blocks among all dense-attention blocks. Executing a ρ fraction of the blocks gives a computeproportional upper-bound speedup of 1/ρ. We measure the actual sparse-kernel speedup over the dense kernel as Sactual , and define the gap as (1/ρ)/Sactual . Figure 2 shows a consistent gap between the computeproportional and measured speedups. The upper bound exceeds the measured 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, whose masks retain predefined structures, exhibit smaller gaps than XAttention and SpargeAttn, which construct fully contentdependent dynamic masks. These results show that FLOP reduction alone cannot yield sparse-attention kernel speedup.

C. Mapping Sparse Attention onto GPUs After constructing the block mask M , the sparse-attention method passes Q, K, V , and M to a backend kernel such as FlashAttention4 [48] or FlashInfer [47] for attention computation. As illustrated in Figure 1, the QKV tensors and attention output O reside in HBM, while the kernel loads the required blocks on chip for computation. These kernels follow the FlashAttention-style nested-loop structure, with query blocks in the outer loop and the K/V blocks selected by the corresponding mask row in the inner loop. For each Qi , the task represented by mask row i is assigned to one streaming multiprocessor (SM) on the GPU, which maintains the online-softmax state and partial output Oi on chip while iterating over its selected K/V blocks. The completed Oi is then written back to HBM. The query-row tasks concurrently occupying all SMs form a wave [24]. In the four-SM example in Figure 1, rows 0–3 initially form wave 0, followed by rows 4–7 in wave 1. A wave is not a synchronization barrier. If row 2 finishes first because it accesses fewer K/V blocks, the freed SM immediately begins row 4 while the remaining rows of wave 0 continue. Because all SMs share the L2 cache, a K/V block loaded by one row can be reused by another if it remains resident in L2 cache; otherwise, it must be fetched again from HBM. The K/V access overlap, row workloads, and execution order therefore jointly determine the L2 cache hit ratio and HBM traffic.

B. Sparse Attention Becomes Memory-Bandwidth-Bound To locate the source of the speedup gap, we profile the 364K-token Wan2.1 workload from Figure 2 using NVIDIA Nsight Compute (NCU) [26]. Table I reports each method’s default density ρ, HBM bandwidth saturation measured as Speed of Light (SOL) [24], and L2 hit ratio. The results show that sparsification degrades cache behavior. Dense attention achieves the highest L2 hit ratio and the lowest HBM SOL, because its regular K/V access stream enables effective cache

3

TABLE I NCU CHARACTERIZATION ON ATTENTION MODULE OF WAN 2.1-14B MODEL WITH 364K TOKENS ON H100.

Methods Dense MInference FlexPrefill XAttention SpargeAttn Random

Density ρ (%)

HBM SOL (%)

L2 Hit (%)

100.0 18.3 29.4 28.8 28.3 20.0

49.38 79.28 86.43 83.44 84.06 87.77

46.51 21.67 14.72 17.94 17.35 10.64

Memory-bound

row 0 row 1 Wave 0 row 2 row 3

SM0 SM1

row 4 row 5 Wave 1 row 6 row 7

SM2 SM3

(a) Spatially dispersed K/V accesses

Dense

800

time

0 1 2 3 4 5 6

Compute-bound H100 peak (990)

1000

Achieved Arithmetic Rate (TFLOP/s)

KV index legend

Wave 0 Wave 1 row 0 row 4 row 1 row 2

row 6 row 5

row 3

row 7

(b) NNZ-induced temporal misalignment

Fig. 4. Two root causes of poor L2 cache locality.

MInference FlexPrefill

blocks close together in time. Dynamic masks undermine both conditions. As illustrated in Figure 4, we consider eight query rows scheduled on four SMs, where rows 0–3 form Wave 0 and rows 4–7 form Wave 1. Each colored block denotes an accessed K/V block, with its color indicating the block index. Root cause 1: Spatially dispersed K/V accesses. Rows scheduled in the same wave share few K/V blocks, because dynamic masks assign each row an irregular access set (Figure 4(a)). With little overlap, the SMs’ concurrent accesses span a large union of K/V blocks that overflows the working set of L2 cache. Useful blocks are evicted before they are reused, and each eviction converts a potential L2 hit into an extra HBM fetch. Root cause 2: NNZ-induced temporal misalignment. Even rows with overlapping access sets may fail to reuse shared blocks when their accesses are separated in time. As shown in Figure 4(b), rows access different numbers of K/V blocks and therefore take different times to execute. Under streaming scheduling, an SM that finishes a shorter row immediately starts one from the next wave, temporally interleaving the two waves. This workload imbalance separates accesses to shared blocks. Although rows 4–7 all access the block pair (K6 , V6 ), they reach it far enough apart that the cached copy may be evicted between accesses, forcing repeated HBM fetches. These two root causes jointly determine whether K/V reuse can occur. The sparse method fixes which K/V blocks each row needs, but not which rows run together or when they reach the shared blocks. Spatial overlap creates opportunities for crossrow reuse, while temporal alignment determines whether these opportunities can be realized before the blocks are evicted from L2. A failure in either condition leads to poor L2 cache locality and repeated HBM accesses.

XAttention

600

SpargeAttn

450 100

ridge point

150

300 700 Operational intensity (FLOP/B)

Random-CV

2000

Fig. 3. Roofline analysis on attention kernels of Wan2.1-14B model with 364K tokens on H100.

reuse. All four dynamic sparse methods, in contrast, issue irregular and discontinuous K/V accesses, which substantially reduce L2 locality and drive up HBM SOL. Consistent with their smaller speedup gaps in Figure 2, MInference and FlexPrefill retain more structured access patterns and therefore better L2 locality than XAttention and SpargeAttn. Finally, a synthetic random mask approximates the locality floor for fully unstructured K/V accesses, which XAttention and SpargeAttn already approach. To determine whether the increased HBM traffic shifts the kernel bottleneck from computation to memory bandwidth, we further conduct a roofline analysis [38] in Figure 3. We estimate FLOPs from the two matrix multiplications in attention using the number of executed attention blocks and the per-block operation count. Kernel execution time and HBM traffic are measured directly with NCU tool. The achieved arithmetic throughput is calculated by dividing the estimated FLOPs by the measured kernel time, while operational intensity is calculated by dividing the estimated FLOPs by the measured HBM traffic. The compute roof corresponds to the theoretical FP16 peak throughput of H100. The results show that dense attention is compute-bound, whereas all sparse methods become memory-bandwidth-bound. Although sparsification reduces FLOPs, poor L2 locality prevents HBM traffic from decreasing proportionally, limiting kernel speedup and producing the gap observed in Figure 2.

IV. WAVE A LIGN I NTERNALS A. Design Rationale As described in Section II-C, each query row is mapped to an SM, while rows in the same wave execute concurrently or in close temporal proximity and share the L2 cache. Meanwhile, query rows are mutually independent, allowing query and

C. Why Access Is Irregular: Two Root Causes Effective L2 reuse requires rows executing within the same wave to access overlapping K/V blocks and to reach these

4

𝑄, 𝐾, 𝑉

Sparse mask 𝑀

Sparse method

𝑂 Adaptive Skip Sparse reordering attention Skip kernel

𝑂

𝑇 𝑀8×8 ≈ 𝑈8×2 Σ2×2 𝑉2×8

Inverse reorder

0 1 2 3 4 5 6 7

No skip Two-stage reorder

0.1 0.1 0.4 −0.5 0.1 −0.5 0.4 −0.3

𝒖𝟏

①Origin mask

Sparse mask 𝑀′

𝒖𝟐

0.6 −0.4 0.2 −0.1 −0.3 0.4 −0.1 −0.4

29.1° 79.5°

θ = 79.5° 140.7° 186.4° 239.2° 289.4° 298.4° 352.1°

𝒖𝟏

𝒖𝟐

② Rank-2 truncated SVD

③2D polar embedding

2 0 5 3 7 1 4 6

④Sort by θ

⑤Reordered mask

Fig. 6. Example of SVD-Based Global Spatial Ordering.

Fig. 5. Workflow of WaveAlign.

min

mask rows to be jointly reordered and their outputs restored afterward without changing attention semantics. Together, these observations enable us to group rows with overlapping K/V accesses into the same wave, increasing cross-SM L2 reuse and reducing HBM traffic. Building on this opportunity, we propose WaveAlign, a lightweight query-row scheduling framework. As shown in Figure 5, the original workflow of sparse attention in the black path directly passes mask M to the backend kernel. In contrast, WaveAlign introduces the red components as plugin stages between mask construction and the backend kernel. Once receiving M , WaveAlign first uses an adaptive guard to determine whether reordering is profitable (Section IV-E). If reordering is skipped, the workload follows the original path. Otherwise, WaveAlign applies its two-stage row-ordering strategy to construct a reordered mask M ′ (Section IV-C and IV-D). After kernel execution, an inverse reorder restores the output rows to their original positions. Since each query row still attends to exactly the same K/V blocks, WaveAlign changes only the physical execution order and preserves both the semantics of sparse attention. However, realizing this workflow efficiently poses two challenges. First, the irregular masks generated by dynamic sparse attention make it difficult to efficiently identify rows with similar K/V access patterns. Second, reordering must incur low overhead, since it is repeatedly performed across attention heads, layers, and denoising steps and can otherwise offset the resulting kernel speedup.

Ph

s.t.

k X X

Shl [i] − Shl [j]

2

l=1 i,j

Shl = Phl Mh , Ph = [Ph1 ; Ph2 ; . . . ; Phk ],

(1)

Ph ∈ {0, 1}MB ×MB , Ph PhT = I. Here, Ph is the row-permutation matrix for head h, and Phl selects the rows assigned to wave l. Accordingly, Shl denotes the reordered mask rows executed in that wave. The constraint Ph PhT = I ensures that each query row appears exactly once. C. SVD-Based Global Spatial Ordering For each attention head, directly searching the optimal permutation requires jointly partitioning and ordering all mask rows to minimize their intra-wave distances, which is NP-hard, motivating the efficient approximation introduced next. Problem (1) seeks a row partition that keeps mutually close rows in the same wave. This is exactly what SVD exposes. Since statistically several query blocks attend to nearly identical K/V blocks, the rows are highly correlated and the singular spectrum is top-heavy, and a few leading directions already capture the dominant row geometry. In the following, we show that the row partitioning of Problem (1) can be turned into a low-dimensional geometric problem in a principled rather than heuristic way, as illustrated in Figure 6. Before using the SVD, we first apply column centering to the mask,   fh = I − 1 11⊤ Mh . (2) M MB Centering subtracts the same mean row from every row, so fh [i] − all pairwise row distances are preserved, namely, ∥M fh [j]∥ = ∥Mh [i] − Mh [j]∥, so reordering Mh is equivalent M fh . to reordering M fh = P σt ut vt⊤ with σ1 ≥ σ2 ≥ Write the SVD as M t · · · ≥ σR ≥ 0 and R = min(MB , NB ). Since the right singular vectors {vt } are orthonormal, the pairwise row distance decomposes exactly as X 2 fh [i] − M fh [j]∥2 = ∥M σt2 ut [i] − ut [j] . (3)

B. Problem Definition To systematically address these challenges, we first formulate this problem. For attention head h, let Mh ∈ {0, 1}MB ×NB denote its block-sparse mask, where MB and NB are the numbers of query and K/V blocks, respectively, and Mh [i, j] = 1 indicates that query block i accesses K/V block j. WaveAlign aims to maximize the K/V-block overlap among query rows assigned to the same scheduling wave. Specifically, we aim to divide the whole Mh into k waves, and each wave contains at most W query-block rows selected from Mh . In order to make the rows within the same wave as close as possible, we construct the following problem:

t

Assigning row i the coordinate vector ϕr (i) = (σ1 u1 [i], . . . , σr ur [i]) therefore makes the Euclidean

5

Algorithm (MFCSI)

distance in the r-dimensional embedding equal to the rank-r partial sum of the true row distance, with residual P 2 2 ϵij = σ t>r t (ut [i] − ut [j]) , which is small when the singular spectrum decays quickly, as is typical for these masks. By the Eckart–Young theorem [9], [23], this truncation is the fh , so the leading (scaled) left best rank-r approximation of M singular vectors are exactly the r-dimensional coordinates that best preserve the pairwise row distances in the least-squares sense (classical multidimensional scaling). The SVD thus yields the optimal low-rank distance-preserving embedding on which we build the ordering, rather than an ad-hoc proxy, and a few leading directions suffice because the σt2 weights concentrate the row geometry in the top of the spectrum. 1) Polar spectral ordering: Although a larger r preserves the row distances more faithfully, in practice we set r = 2 and use only the rank-2 truncation [50], for two reasons. • Sufficiency. Two directions already capture the structure that matters at this stage: the σt2 weights place most of the row geometry in the leading modes, and we only need a coarse neighborhood order here; the fine intrawave alignment is handled in Section IV-D. • Linearizability. The scheduler ultimately needs a single one-dimensional execution order, so the coordinates must be collapsed onto one axis. Two dimensions form the smallest embedding that a polar angle can turn into a total order: every row then has a well-defined angle and the rows sort directly, whereas r ≥ 3 admits no comparably cheap and rotation-robust linearization and would instead require a clustering or space-filling heuristic. As analyzed above, we can represent each row i by a 2dimensional point (σ1 u1 [i], σ2 u2 [i]), and then sort them by their polar angles. Note that σ1 , σ2 ≥ 0, so they merely reparametrizes the angle monotonically but leaves the resulting order unchanged. We therefore drop the singular-value weights and order the rows by the polar angles  ai = atan2 u2 [i], u1 [i] . (4)

1 Matrix-Free Centered Subspace Iteration

Require: binary mask M ∈ {0, 1}MB ×NB , iterations   T f = I − 1 11⊤ M Ensure: top-2 left singular vectors u1 , u2 of M M B

1: c ← M ⊤ 1 2: m ← MB 3: X ← random orthonormal block of shape MB × 2 4: Subspace power iteration 5: for t = 1 to T do 1 6: Y ← M ⊤X − m c (1⊤ X) 1 7: Z ← MY − m 1 (c⊤ Y ) 8: X ← Q from the QR factorization Z = QR 9: end for 10: Optional: order by singular value 1 1 11: Y ← M ⊤ X − m c (1⊤ X); Z ← M Y − m 1 (c⊤ Y ) 12: B ← X ⊤ Z 13: (Λ, R) ← eig(B) 14: X ← XR 15: u1 , u2 ← columns of X 16: return u1 , u2

K/V block index legend 0

11

wave0 wave1 wave2 SM0 r0 (6) r10 (5) r16 (9) r14 (7) r20 (7) SM1 r1 (8) r11 (9) r21 (9) SM2 r2 (6) r13 (7) r19 (8) SM3 r3 (7) r15 (9) r23 (9) SM4 r4 (10) r18 (9) SM5 r5 (5) r9 (8) r17 (9) SM6 r6 (5) r8 (7) r12 (10) r22 (10) SM7 r7 (6)

(a) Original row order: high overlap, temporal misalign

23

wave0

wave1

wave2

r10 (5) r23 (9) SM0 r4 (10) r14 (7) r19 (8) SM1 r1 (8) r13 (7) r17 (9) SM2 r3 (7) r15 (9) r16 (9) SM3 r0 (6) r9 (8) r18 (9) SM4 r2 (6) r8 (7) r22 (10) SM5 r7 (6) r20 (7) SM6 r5 (5) r12 (10) r21 (9) SM7 r6 (5) r11 (9)

(b) Intra-wave NNZ-desc: high overlap, temporal align

Fig. 7. Effect of Stage 2 intra-wave ordering, ri (n) denotes row i with n nonzero blocks.

from the origin, so on the raw mask the leading singular vector is spent on the row centroid, which tracks the perrow workload (σ1 u1 [i] = ⟨Mh [i], v1 ⟩ ≈ ni ) rather than the K/V-access shape. Centering sends this constant (workload) fh = 0), so both direction to a zero singular value (1⊤ M f retained coordinates u1 , u2 of Mh carry purely structural information.

As shown in Figure 6, each row i is placed at the point (u1 [i], u2 [i]), and the polar angle unrolls the ring into a onedimensional sequence where angularly adjacent rows tend to share more K/V blocks. 2) Matrix-Free Centered Subspace Iteration: To obtain fh and apply u1 , u2 , a direct approach would first form M an SVD, or equivalently an eigendecomposition of the Gram fh M f⊤ ; both are computationally expensive. Instead, matrix M h we compute u1 , u2 with a subspace power iteration, detailed fh or the in Algorithm 1. Crucially, it never materializes M MB × MB Gram matrix: centering enters only as a rankf⊤ X = 1 correction to two sparse matrix–vector products, M h fh Y = Mh Y − 1 1(c⊤ Y ) with Mh⊤ X − M1B c(1⊤ X) and M MB c = Mh⊤ 1. Each iteration only costs O(nnz). It converges geometrically to the top-2 invariant subspace, and a small fixed T (5–15) suffices because the subsequent angular sort is invariant to in-plane rotations and sign flips of (u1 , u2 ). It is worth emphasizing that the centering step is necessary rather than cosmetic: the embedding coordinates are measured

D. NNZ-Guided Intra-Wave Temporal Alignment The SVD-based ordering mentioned above (Stage 1) addresses spatially dispersed K/V accesses by grouping rows with similar access patterns. However, as discussed in Section III-C, it does not account for the varying numbers of nonzero blocks across rows. Figure 7(a) illustrates 24 rows scheduled on 8 SMs over 3 waves. SM0–SM7 initially execute rows r0 –r7 , but rows with fewer NNZ blocks finish earlier and immediately fetch subsequent rows from the Wave 1. For example, after completing r5 and r6 , the corresponding SMs proceed to r9 and r8 , respectively. This progressive execution skew creates temporal misalignment across waves. Consequently, although Stage 1 places rows with overlapping K/V accesses close together, their accesses to the same blocks may occur far apart in time, weakening L2 cache reuse.

6

Algorithm 2 Two-stage reordering

V. IMPLEMENTATION We integrate WaveAlign into vLLM-Omni [32] as a lightweight plug-in between sparse-mask construction and backend-kernel execution. Reordering operator. We introduce two additional operators. The first implements the two-stage row reordering described in Algorithm 2, applying the resulting permutation to the query blocks and sparse mask before attention kernel. The second applies the inverse permutation to restore the attention outputs to their original row order. Both operators are composed from PyTorch tensor primitives and registered in vLLM-Omni, allowing them to be invoked transparently throughout the inference pipeline. FlashInfer per-head scheduling. On the A40 platform, we use FlashInfer [47] as the backend kernel. Its native scheduler prioritizes the K/V-head dimension, allowing Q-block computations from different heads to run concurrently across SMs. Because each head accesses a separate set of K/V blocks without cross-head overlap, concurrent execution only increases contention in the shared L2 cache. Following FlashAttention4 [48], we instead prioritize the query-sequence dimension and execute K/V heads sequentially. This exposes only one head’s K/V working set at a time, reducing L2 contention and improving intra-head locality. Section VI-B and VI-C evaluate the effectiveness of this scheduling modification.

Require: binary mask M ∈ {0, 1}H×MB ×NB , wave size W Ensure: per-head row permutations Ph 1: for h = 1 to H do 2: Stage 1: SVD-based global spatial ordering 3: Compute the top-2 left singular vectors u1 , u2 by MFCSI 4: ai ← atan2(u2 [i], u1 [i]) 5: O ← rows sorted by ai 6: Split O into consecutive buckets B1 , . . . , BK of at most W rows 7: Stage 2: NNZ-based intra-wave alignment P 8: N N Zi ← j Mh [i, j] for all rows i 9: for each bucket Bk do 10: sort rows in Bk by descending N N Zi 11: end for 12: Ph ← concatenation of the refined buckets B1 , . . . , BK 13: end for

To mitigate this temporal misalignment, we first sort the rows within each Stage 1 Wave in descending order of NNZ blocks and then schedule them onto the SMs. As shown in Figure 7(b), the first Wave assigns the longest row r4 to SM0 and the shortest row r6 to SM7. After the shorter rows r5 and r6 complete, their SMs fetch the longest rows in the next Wave, r12 and r11 , respectively. This long-short complementary assignment at successive Wave boundaries limits the accumulation of execution skew and keeps the SMs approximately aligned. Because rows in adjacent waves have similar K/V access patterns after Stage 1, temporal alignment allows their shared K/V blocks to be accessed within shorter reuse distances, thereby improving L2 cache locality.

VI. E VALUATION A. Experimental Setup Hardware and software platforms. We evaluate WaveAlign on two NVIDIA GPU platforms. The H100 GPU provides 80GB of HBM3 memory and 50MB L2 cache shared by 132 SMs. The A40 GPU provides 48GB of GDDR6 memory and 6MB L2 cache shared by 84 SMs. The primary software stack includes CUDA 13.0 [25], PyTorch 2.11 [29], NsightCompute 2025.3.1 [26], vLLM-Omni 0.20.0 [32]. We use FlashAttention4 [48] as the backend kernel on H100 GPU and FlashInfer [47] on A40 GPU, as FlashAttention4 depends on the Hopper and Blackwell architecture and does not support Ampere GPUs. Models. We evaluate WaveAlign on two representative opensource Text-to-Video (T2V) model families, Wan2.1 [35] and LTX2.3 [12]. Our H100 experiments use Wan2.1-T2V14B [34] and LTX2.3-22B, whereas A40 uses the smaller Wan2.1-T2V-1.3B [33] due to the limited memory capacity. Workloads. We use prompts from the Penguin Video Benchmark [15] to evaluate video generation. To construct workloads with different sequence lengths, we vary the number of frames while keeping the resolution fixed. For Wan2.1-T2V-14B, we use a resolution of 720×1280 and vary the frames from 161 to 481, yielding sequences of 148K–436K tokens. For LTX2.322B, the resolution is 1216×1920, and we vary the range of frames from 601 to 1401, yielding sequences of 173K–401K. For Wan2.1-T2V-1.3B, we set resolution as 720×1280 and frames from 121 to 241, yielding sequences of 112K–220K. Baselines and configurations. Our baselines are four sparseattention methods without applying WaveAlign under default

E. Runtime Coordination Two-stage coordination. Sections IV-C and IV-D improve L2 locality from complementary spatial and temporal dimensions, respectively. Algorithm 2 integrates the two stages. Given the sparse mask M and wave size W , WaveAlign reorders each attention head independently (Line 1). It first applies Stage 1 to derive a global row order based on K/V-access similarity (Lines 2–6), and partitions the ordered rows into wave-sized buckets (Line 7). Within each bucket, Stage 2 sorts rows by descending NNZ count to reduce cross-wave execution skew (Lines 8–12). Finally, the refined buckets are concatenated to produce the row permutation that can be directly scheduled by the backend kernel (Line 13). Adaptive reordering skip. Row reordering is not always profitable, because its overhead must be amortized by the resulting kernel-time reduction. We skip reordering in three cases. First, for short sequences, the entire K/V working set fits in the L2 cache, making row order largely irrelevant to cache reuse. Second, when the mask is too sparse, the remaining attention workload is insufficient to amortize the reordering cost. Third, when adjacent rows already exhibit high K/Vaccess similarity, the original order provides adequate L2 locality. To make this decision automatically, our adaptive skip module extracts lightweight statistics from the sparse mask and combines them with offline profiles of kernel latency and reordering cost. Reordering is enabled only when the estimated latency reduction exceeds its overhead.

7

E2E Time (103 s)

Original Wan2.1-T2V-14B (XAttn)

Wan2.1-T2V-14B (Sparge)

WaveAlign LTX2.3 (XAttn)

4.5

15

15

7.5

10

10

5

3

5

5

2.5

1.5

0

0

148 220 292 364 436 Sequence length (K)

148 220 292 364 436 Sequence length (K)

0

173 230 287 344 401 Sequence length (K)

0

LTX2.3 (Sparge)

173 230 287 344 401 Sequence length (K)

Fig. 8. End-to-end video generation latency with and without WaveAlign on Wan2.1-T2V-14B and LTX2.3 across various sequence lengths atop H100.

E2E Time (103 s)

Original

6

O + Seq

Wan2.1-T2V-1.3B (XAttn)

O + WA

Wan2.1-T2V-1.3B (Sparge)

4.5

4

3.0

2

1.5

0

112 148 184 220 Sequence length (K)

speedups across both models, both sparse-attention methods, and all evaluated sequence lengths, reaching up to 1.17× on Wan2.1-T2V-14B and 1.14× on LTX2.3-22B. The gains generally increase with sequence length because the growing K/V working set increasingly exceeds the fixed L2 cache capacity, leading to more severe cache thrashing and heavy HBM accesses under the original row order. A40 results. Figure 9 reports the end-to-end generation latency of XAttention and Sparge on Wan2.1-T2V-1.3B under four configurations: FlashInfer’s native concurrent-head execution (Original), the sequential-head scheduler described in Section V (O+Seq), WaveAlign applied to the native scheduler (O+WA), and the combination of both optimizations (O+Seq+WA). Compared with Original, O+Seq alone achieves a speedup of up to 1.09×, demonstrating the benefit of our sequential-head implementation. Applying WaveAlign directly to Original yields a speedup of up to 1.29×. Combining the two optimizations further improves upon O+WA by up to 11.0%, with O+Seq+WA consistently achieving the lowest latency. The two optimizations improve cache behavior in complementary ways: sequential-head execution prevents different heads from simultaneously competing for the A40’s 6MB L2 cache, while WaveAlign further increases the L2 hit ratio by improving intra-head K/V reuse. At comparable sequence lengths, the gains on A40 are larger than those on H100 because the A40’s much smaller L2 cache becomes capacity-constrained earlier, making cache locality important even for shorter sequences.

O + Seq + WA

0.0

112 148 184 220 Sequence length (K)

Fig. 9. E2E video generation latency of Wan2.1-T2V-1.3B on A40/FlashInfer under different combinations of head scheduling and WaveAlign.

sparse configurations: XAttention [42] with threshold 0.9, SpargeAttn [49] and MInference [16] with pre-autotuned config and FlexPrefill [18] with g = 0.60. We use these configurations for the kernel-level evaluation of four methods. For the end-to-end evaluation, however, we exclude MInference and FlexPrefill because their overhead of mask construction is prohibitively high in our T2V workloads. For FlashAttention4 on H100, we use a block size of 128×128 tokens and one resident cooperative thread array (CTA) per SM across 132 SMs, yielding a wave size of 132. For FlashInfer on A40, we use a block size of 64×64 tokens and two resident CTAs per SM across 84 SMs, yielding a wave size of 168. Metrics. We report end-to-end generation latency and sparseattention kernel latency, together with speedup relative to the corresponding baselines. We separately measure the total reordering overhead, including permutation construction, query and mask gathering, and inverse for output. To explain the latency reduction, we report L2 cache hit ratio and HBM read traffic. Finally, to verify that WaveAlign preserves the sparseattention semantics and generated video quality, we report PSNR, SSIM, LPIPS, and the VBench [14] metrics Imaging Quality and Subject Consistency.

C. Kernel Latency Breakdown WaveAlign optimizes only sparse attention, so its end-toend speedup comes entirely from reduced attention latency. We next isolate the attention computation and analyze its kernellevel latency on H100 and A40. H100 results. Figure 10 illustrates the H100 sparse-attention kernel latency of MInference, FlexPrefill, XAttention, and SpargeAttn using real QKV tensors captured from Wan2.1T2V-14B. Each WaveAlign bar separates the optimized kernel latency from the complete reordering overhead. Including this overhead, WaveAlign consistently outperforms the original kernel across all four methods and sequence lengths, achieving speedups of 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. Reordering accounts for only 1.1%–

B. End-to-End Video Generation Latency H100 results. Figure 8 reports the end-to-end video generation latency of XAttention and SpargeAttn with and without WaveAlign on Wan2.1-T2V-14B and LTX2.3-22B. In the original baselines, sparse attention dominates end-to-end latency, accounting for 65.8%–88.0% on Wan2.1-T2V-14B and 58.1%–84.4% on LTX2.3-22B. Consequently, WaveAlign’s kernel-level improvements translate into consistent end-to-end

8

Attention latency (s)

Original

MInference (H100)

2

1.5 148

220

292

364

0

436

Sequence length(K)

148

220

Attention latency (s)

MInference (A40)

1 0

292

364

148

184

220

Sequence length(K)

2

2 148

O + Seq

0.8

0

220

292

364

436

O + WA

O + Seq + WA

FlexPrefill (A40)

184

0

220

220

292

364

436

Sequence length(K)

XAttention (A40)

SpargeAttn (A40) 3

1.5

148

148

Reorder overhead

3

112

0

Sequence length(K)

0.4

112

4

0

436

SpargeAttn (H100)

4

Sequence length(K)

Original 2

Reorder overhead

XAttention (H100)

4

3

0

WaveAlign

FlexPrefill (H100)

1.5 112

Sequence length(K)

148

184

220

0

112

Sequence length(K)

148

184

220

Sequence length(K)

Fig. 10. Sparse-attention kernel latency on H100 (top) and A40 (bottom). For WaveAlign configurations, the complete reordering overhead is shown separately above the kernel latency.

L2 hit ratio (%)

Original 100

50

0

XAttn Sparge MInf

(a) L2 hit ratio

Flex

WaveAlign

Norm. HBM read

3.2% of the total WaveAlign attention latency and is consistently outweighed by the kernel-time reduction. The gains increase with sequence length as the growing K/V working set exceeds L2 capacity, amplifying cache evictions and repeated HBM accesses that WaveAlign mitigates through shorter reuse distances. The gains are also larger for XAttention and SpargeAttn because their fully content-dependent masks produce more irregular K/V accesses than the structured masks of MInference and FlexPrefill. A40 results. The A40 panels in Figure 10 report the kernel latency of the same four methods using real QKV tensors captured from Wan2.1-T2V-1.3B. We evaluate parallel- and sequential-head execution, each with and without WaveAlign, using FlashInfer’s native parallel-head scheduler as the baseline. The legend follows the naming convention in Section VI-B, and all reported WaveAlign speedups include the complete reordering overhead. Across the four sequence lengths, O+WA achieves maximum speedups over Original of 1.104×, 1.090×, 1.599×, and 1.520× for MInference, FlexPrefill, XAttention, and SpargeAttn, respectively. Relative to O+Seq, O+Seq+WA achieves maximum speedups of 1.133×, 1.077×, 1.614×, and 1.556×, respectively. Compared with Original, O+Seq+WA achieves maximum speedups of 1.158×, 1.090×, 1.635×, and 1.601×, and consistently provides the lowest total attention latency. WaveAlign improves both parallel- and sequential-head execution across all four methods, with its 7.5–29.7ms reordering overhead consistently outweighed by the kernel-time reduction. Sequential-head scheduling reduces cross-head L2 contention, while WaveAlign improves intra-head K/V locality, making O+Seq+WA the most effective configuration. Compared with H100, A40 is more sensitive to row ordering. While WaveAlign achieves up to 1.252× speedup on H100, O+WA reaches 1.599× and 1.520× for XAttention and

1.0 0.5 0.0

XAttn Sparge MInf

Flex

(b) HBM read

Fig. 11. L2 cache hit ratio and HBM read traffic on H100/FlashAttention4 with and without WaveAlign across four sparse-attention methods. HBM reads are normalized to the corresponding no-reorder baseline, and all results report the median of three NCU runs.

SpargeAttn on A40, respectively, even at shorter sequence lengths. This is because the A40’s smaller L2 cache (6MB versus 50MB on H100) becomes capacity-constrained earlier, amplifying the benefits of row reordering for irregular, contentdependent masks. D. L2 Cache Locality and HBM Traffic To identify the source of the kernel speedup, we use NVIDIA Nsight Compute (NCU) [26] to profile FlashAttention4 on H100 using a Wan2.1-T2V-14B workload with 401 frames at 720×1280, corresponding to 364K visual tokens. Figure 11 shows consistent locality improvements across all four sparse methods: the L2 cache hit ratio increases from 28.48%–36.35% to 79.38%–89.06%, while HBM reads decrease by 80.71%–92.11%. In particular, XAttention and SpargeAttn have the lowest baseline hit ratios because their fully content-dependent masks produce more dynamic and irregular K/V access patterns. WaveAlign raises their hit ratios from 36.35% and 28.48% to 84.17% and 80.04%, respectively. Since WaveAlign preserves both the sparse mask and computation, these results show that its speedup comes from

9

Original

O + S1

O + S2

O + S1 + S2

0.5

0.0

XAttn

Sparge

MInf

(a) Kernel time

Flex

Norm. HBM read

L2 hit ratio (%)

Kernel time

100 1.0

50

0

XAttn

Sparge

MInf

Flex

1.0

0.5

0.0

XAttn

(b) L2 hit ratio

Sparge

MInf

Flex

(c) HBM read

Fig. 12. Ablation of WaveAlign’s two-stage reorder across four sparse-attention methods under four cumulative configurations: Original, O+S1, O+S2 and O+S1+S2. S1 and S2 denote global spatial ordering and intra-wave temporal alignment, respectively. Kernel time and HBM reads are normalized to Original.

improved K/V reuse: shorter reuse distances keep more K/V blocks in L2 and avoid repeated HBM accesses.

Normalized latency

Original

E. Ablation Study Figure 12 evaluates the two stages of WaveAlign on H100 with FlashAttention4 using a Wan2.1-T2V-14B workload of 401 frames at 720×1280, corresponding to 364K visual tokens. We evaluate XAttention, SpargeAttn, MInference, and FlexPrefill under four cumulative configurations: Original, O+S1, O+S2, and O+S1+S2. Stage 1 applies the global spatial row ordering described in Section IV-C, while Stage 2 adds the intra-wave NNZ-based refinement described in Section IV-D. We use an execution-wave size of 132 rows, matching the number of SMs on H100. Stage 1 alone provides modest improvements, achieving up to 1.0319× kernel speedup, increasing the L2 hit ratio from 28.48%–36.35% to 34.88%–44.47%, and reducing HBM read traffic by 8.71%–18.61%. Stage 2 alone delivers substantially larger gains. Across the four methods, O+S2 achieves 1.091×– 1.147× kernel speedup, increases the L2 hit ratio to 68.35%– 78.71%, and reduces HBM read traffic by 68.89%–83.67% relative to Original. Combining both stages consistently produces the best results. Relative to Original, O+S1+S2 accelerates the kernel by up to 1.17×, raises the L2 hit ratio to 79.38%– 89.06%, and reduces HBM traffic by up to 92.11%. The two stages improve L2 reuse along complementary dimensions. Stage 1 groups rows with similar K/V access patterns into nearby waves, increasing the likelihood that different SMs request the same blocks. However, uneven row workloads can separate these requests in time and prevent cache reuse. Stage 2 sorts rows by descending NNZ to align CTA progress and shorten reuse distances, keeping shared K/V blocks in L2 for subsequent accesses. Stage 1 establishes spatial overlap, while Stage 2 converts it into temporal reuse; both are therefore necessary for effective L2 reuse.

1.0

Skip

Adaptive path

Reorder overhead

Density = 28.8%

Sequence length = 364K

Skip

Skip

Skip

1.09× 1.18×

Skip

1.14× 1.18× 1.20×

0.5 0.0

25 49 98 197 364 Sequence length (K)

1.7 6.7 14.5 28.8 57.4 Retained density (%)

Fig. 13. Adaptive reordering skip for XAttention with FlashAttention4 on Wan2.1-T2V-14B under different sequence lengths and retained densities.

XAttention

WaveAlign

SpargeAttn

Original

WaveAlign

132×132 zoom Head 0 mask

Original

Fig. 14. Head-0 sparse masks of XAttention and SpargeAttn before and after WaveAlign. The top row shows the full masks, with local regions marked in red; the bottom shows the corresponding 132×132 zoomed windows.

so WaveAlign skips reordering and falls back to the default sparse-attention path, corresponding to a 1.0× speedup over the baseline. From 197K onward, the growing K/V working set makes cache locality increasingly important, and reordering becomes profitable, delivering up to 1.18× speedup. Retained density exhibits a similar trend: WaveAlign skips reordering at 1.7% and 6.7% density, but enables it from 14.5% onward, achieving 1.14×–1.20× speedup. These results show that adaptive skipping prevents reordering overhead from causing regressions while preserving its benefits in profitable regimes. Visualization of sparse mask. Figure 14 visualizes the head-0 sparse masks of XAttention and SpargeAttn on Wan2.1-T2V14B before and after WaveAlign reordering. Each zoomed region spans 132 query rows, corresponding to one execu-

F. Other Factors Adaptive reordering skip. Figure 13 demonstrates the necessity of adaptive reordering proposed in Section IV-E by reporting kernel-level latency, including the reordering overhead, for XAttention on Wan2.1-T2V-14B with FlashAttention4. At short sequence lengths of 25K, 49K, and 98K, the kerneltime reduction is insufficient to amortize the reordering cost,

10

TABLE II C OMPARISON WITH OTHER REORDERING ALGORITHMS . Time (ms)

Speedup w/o overhead

Speedup w/ overhead

Overhead (ms)

HBM SOL (%)

L2 Hit (%)

2723.99 2234.37 2198.40 2257.93 2259.70 2707.63

1.000 1.221 1.236 1.208 1.203 1.008

1.000 1.210 0.005 0.715 0.715 0.961

0.00 19.85 535435.34 1556.40 1540.98 133.07

83.98 15.37 10.38 19.24 20.90 81.65

17.57 76.72 83.36 72.26 70.30 19.29

Method Original WaveAlign Hypergraph MinHash LSH Clustering

Memory-bound

MInference FlexPrefill XAttention SpargeAttn

450 100

150

Sparse method

Density

PSNR

SSIM

LPIPS

ImgQual

SubCons

Wan2.1-T2V-14B

Dense XAttn XAttn+WaveAlign Sparge Sparge+WaveAlign

1.000 0.507 0.507 0.559 0.559

– 17.498 17.498 19.161 19.161

– 0.581 0.581 0.675 0.675

– 0.284 0.284 0.224 0.224

0.695 0.686 0.686 0.694 0.694

0.963 0.839 0.839 0.956 0.956

LTX2.3-22B

Dense XAttn XAttn+WaveAlign Sparge Sparge+WaveAlign

1.000 0.599 0.599 0.588 0.588

– 27.683 27.683 26.485 26.485

– 0.910 0.910 0.885 0.885

– 0.094 0.094 0.115 0.115

0.581 0.594 0.594 0.589 0.589

0.912 0.899 0.899 0.898 0.898

ing Wan2.1-T2V-14B and LTX2.3-22B with XAttention and SpargeAttn. All videos are generated at 480p with 81 frames. Across both models and sparse attention methods, the original and WaveAlign variants achieve identical results on all reported metrics. This is because WaveAlign only reorders the mask rows and their corresponding queries, while preserving the selected K/V blocks and the backend kernel’s per-row computation. The outputs are then restored to their original order, leaving the sparse attention semantics and inference accuracy unchanged.

800

600

Model

Compute-bound H100 peak (990)

1000

Achieved Arithmetic Rate (TFLOP/s)

TABLE III Q UALITY AND CONSISTENCY METRICS BEFORE AND AFTER WAVE A LIGN .

ridge point

300 700 Operational intensity (FLOP/B)

2000

Fig. 15. Roofline analysis on attention kernels with WaveAlign of Wan2.114B model with 364K tokens on H100.

VII. R ELATED W ORK tion wave on H100, and 132 K/V-block columns. Before reordering, the masks are dominated by diagonal patterns, and rows within the same wave exhibit little column-wise overlap, indicating that different SMs rarely access the same K/V blocks. After reordering, more vertically aligned nonzero blocks appear, showing that rows executed within the same wave are more likely to share K/V accesses. Since WaveAlign preserves each row’s selected blocks and changes only the execution order, this increased column-wise overlap directly improves cross-SM L2 reuse. Compare with other reorder algorithms. Table II compares WaveAlign with four representative locality-aware sparse matrix multiplication reordering methods reviewed in Section VII, using the configuration in Section VI-D. These methods expose a trade-off between locality and overhead. Clustering has the lowest overhead among these baselines but barely raises the L2 hit ratio from 17.57% to 19.29%. Hypergraph, MinHash, and LSH substantially improve L2 locality, but their preprocessing costs outweigh the kernel savings, resulting in inference slowdowns. In contrast, WaveAlign raises the L2 hit ratio to 76.72% with only 19.85ms of overhead, retaining a 1.210× net speedup. It is therefore the only evaluated method that combines effective locality improvement with practical online overhead. Roofline analysis. Compared with Figure 3, Figure 15 shows that WaveAlign increases operational intensity and achieved arithmetic rate, moving all four sparse-attention kernels across the ridge point from memory-bandwidth-bound to computebound execution. Video generation accuracy. Table III evaluates video generation quality on the Penguin Video Benchmark [15] us-

Sparse attention for video generation. This topic has attracted growing interest. Fixed-pattern methods reuse predefined layouts across inputs and denoising steps. Rule-based approaches, including LongNet [7], LogSparse [19], STA [51], Radial Attention [20], and LVSA [10], exploit spatiotemporal locality, whereas profile-guided methods, such as SVG [39] and Sparse-vDiT [3], select head- or layer-specific patterns from representative attention profiles. Content-adaptive methods, such as AdaSpa [40] and SVG2 [45], identify important blocks or tokens from the current input to balance sparsity and generation quality. Complementary to these methods, WaveAlign reorders rows after mask construction to improve kernel efficiency without changing the selected blocks. GPU kernels for sparse attention. To accelerate attention, prior kernel optimizations fall into two groups. The FlashAttention series [5], [6], [48] improve kernel efficiency through IO-aware tiling, online softmax, and GPU pipelining. Sparse or flexible attention kernels, including FlashInfer [47], FlexAttention [8], FlashMask [36], Block Sparse FlashAttention [27], and FSA [44], further support diverse sparse masks. WaveAlign complements these backends by reordering the mask to unlock their full performance. Locality-aware sparse matrix reordering. Reordering is widely used to improve locality in sparse matrix multiplication, with representative approaches based on LSH-based row reordering [17], hypergraph partitioning [1], MinHashbased ordering [4], and clustering [43]. However, our prior evaluation shows that the high preprocessing overhead makes them unsuitable for runtime-generated sparse attention masks. To the best of our knowledge, WaveAlign provides a fast and effective reordering method tailored to sparse attention.

11

VIII. CONCLUSION

[17] P. Jiang, C. Hong, and G. Agrawal, “A novel data transformation and execution strategy for accelerating sparse matrix multiplication on GPUs,” in Proceedings of the 25th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, 2020, pp. 376–388. [18] X. Lai, J. Lu, Y. Luo, Y. Ma, and X. Zhou, “FlexPrefill: A context-aware sparse attention mechanism for efficient long-sequence inference,” in International Conference on Learning Representations, 2025. [19] S. Li, X. Jin, Y. Xuan, X. Zhou, W. Chen, Y.-X. Wang, and X. Yan, “Enhancing the locality and breaking the memory bottleneck of transformer on time series forecasting,” in Advances in Neural Information Processing Systems, vol. 32, 2019. [20] X. Li, M. Li, T. Cai, H. Xi, S. Yang, Y. Lin, L. Zhang, S. Yang, J. Hu, K. Peng, M. Agrawala, I. Stoica, K. Keutzer, and S. Han, “Radial attention: o(n log n) sparse attention with energy decay for long video generation,” in Advances in Neural Information Processing Systems, vol. 38, 2025. [21] LTX Team, “LTX-2.3,” LTX Documentation, Mar. 2026, accessed: Jul. 27, 2026. [Online]. Available: https://docs.ltx.video/models [22] X. Ma, Y. Wang, X. Chen, G. Jia, Z. Liu, Y.-F. Li, C. Chen, and Y. Qiao, “Latte: Latent diffusion transformer for video generation,” Transactions on Machine Learning Research, 2025. [23] L. Mirsky, “Symmetric gauge functions and unitarily invariant norms,” The Quarterly Journal of Mathematics, vol. 11, no. 1, pp. 50–59, 1960. [24] NVIDIA Corporation, “NVIDIA Nsight Compute Profiling Guide,” https://docs.nvidia.com/nsight-compute/ProfilingGuide/index.html# compute-model, accessed: Jul. 29, 2026. [25] NVIDIA Corporation, “CUDA Toolkit 13.0 Documentation,” https: //docs.nvidia.com/cuda/archive/13.0.0/, 2025, accessed: Jul. 30, 2026. [26] NVIDIA Corporation, “NVIDIA Nsight Compute 2025.3.1 Release Notes,” https://docs.nvidia.com/nsight-compute/ReleaseNotes/topics/ updates-2025-3-1.html, 2025, accessed: Jul. 30, 2026. [27] D. Ohayon, I. Lamprecht, I. Hubara, I. Cohen, D. Soudry, and N. Elata, “Block sparse FlashAttention,” arXiv preprint arXiv:2512.07011, 2025. [28] W. Peebles and S. Xie, “Scalable diffusion models with transformers,” in Proceedings of the IEEE/CVF International Conference on Computer Vision, 2023, pp. 4195–4205. [29] PyTorch Foundation, “PyTorch 2.11 Release Blog,” https://pytorch.org/ blog/pytorch-2-11-release-blog/, Mar. 2026, accessed: Jul. 30, 2026. [30] Y. Sun, Z. Li, Y. Zhang, T. Pan, B. Dong, Y. Guo, and J. Wang, “Efficient attention mechanisms for large language models,” Patterns, p. 101594, 2026. [31] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin, “Attention is all you need,” in Advances in Neural Information Processing Systems, vol. 30, 2017, pp. 5998–6008. [32] vLLM Project, “vLLM-Omni,” GitHub repository, 2026, accessed: Jul. 27, 2026. [Online]. Available: https://github.com/vllm-project/vllmomni [33] Wan-AI, “Wan2.1-T2V-1.3B model card,” https://huggingface.co/WanAI/Wan2.1-T2V-1.3B, 2025, hugging Face, accessed: Jul. 30, 2026. [34] Wan-AI, “Wan2.1-T2V-14B model card,” https://huggingface.co/WanAI/Wan2.1-T2V-14B, 2025, hugging Face, accessed: Jul. 30, 2026. [35] Wan-Video, “Wan2.1: Open and advanced large-scale video generative models,” GitHub repository, 2025, [Online]. Available: https://github. com/Wan-Video/Wan2.1. Accessed: Jul. 30, 2026. [36] G. Wang, J. Zeng, X. Xiao, S. Wu, J. Yang, L. Zheng, Z. Chen, J. Bian, D. Yu, and H. Wang, “FlashMask: Efficient and rich mask extension of FlashAttention,” in International Conference on Learning Representations, 2025. [37] Y. Wang, X. Liu, W. Pang, L. Ma, S. Yuan, P. E. Debevec, and N. Yu, “Survey of video diffusion models: Foundations, implementations, and applications,” arXiv preprint arXiv:2504.16081, 2025. [38] S. Williams, A. Waterman, and D. A. Patterson, “Roofline: An insightful visual performance model for multicore architectures,” Communications of the ACM, vol. 52, no. 4, pp. 65–76, 2009. [39] H. Xi, S. Yang, Y. Zhao, C. Xu, M. Li, X. Li, Y. Lin, H. Cai, J. Zhang, D. Li, J. Chen, I. Stoica, K. Keutzer, and S. Han, “Sparse VideoGen: Accelerating video diffusion transformers with spatial-temporal sparsity,” in Proceedings of the 42nd International Conference on Machine Learning, ser. Proceedings of Machine Learning Research, vol. 267, 2025, pp. 68 208–68 224. [40] Y. Xia, S. Ling, F. Fu, Y. Wang, H. Li, X. Xiao, and B. Cui, “Trainingfree and adaptive sparse attention for efficient long video generation,” in Proceedings of the IEEE/CVF International Conference on Computer Vision, 2025, pp. 15 982–15 993.

WaveAlign is an efficient reordering framework for accelerating dynamic sparse attention in long-video generation. By jointly improving spatial KV locality and temporal workload alignment, it increases L2 cache reuse and reduces HBM traffic without modifying the sparse attention algorithm. WaveAlign reduces end-to-end video generation latency by up to 1.17× and sparse attention latency by up to 1.25× over strong baselines across multiple models and algorithms. R EFERENCES [1] G. Ballard, A. Druinsky, N. Knight, and O. Schwartz, “Hypergraph partitioning for sparse matrix-matrix multiplication,” ACM Transactions on Parallel Computing, vol. 3, no. 3, pp. 18:1–18:34, 2016. [2] J. Chen, W. He, Y. Gu, Y. Zhao, J. Yu, J. Chen, D. Zou, Y. Lin, Z. Zhang, M. Li, H. Xi, L. Zhu, E. Xie, S. Han, and H. Cai, “DC-VideoGen: Efficient video generation with deep compression video autoencoder,” arXiv preprint arXiv:2509.25182, 2025. [3] P. Chen, X. Zeng, M. Zhao, P. Ye, M. Shen, W. Cheng, G. Yu, and T. Chen, “Sparse-vDiT: Unleashing the power of sparse attention to accelerate video diffusion transformers,” in Proceedings of the AAAI Conference on Artificial Intelligence, vol. 40, no. 4, 2026, pp. 2957– 2965. [4] F. Chierichetti, R. Kumar, S. Lattanzi, M. Mitzenmacher, A. Panconesi, and P. Raghavan, “On compressing social networks,” in Proceedings of the 15th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, 2009, pp. 219–228. [5] T. Dao, “FlashAttention-2: Faster attention with better parallelism and work partitioning,” in International Conference on Learning Representations, 2024. [6] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré, “FlashAttention: Fast and memory-efficient exact attention with IO-awareness,” in Advances in Neural Information Processing Systems, vol. 35, 2022, pp. 16 344– 16 359. [7] J. Ding, S. Ma, L. Dong, X. Zhang, S. Huang, W. Wang, N. Zheng, and F. Wei, “LongNet: Scaling transformers to 1,000,000,000 tokens,” arXiv preprint arXiv:2307.02486, 2023. [8] J. Dong, B. Feng, D. Guessous, Y. Liang, and H. He, “FlexAttention: A programming model for generating fused attention variants,” in Proceedings of Machine Learning and Systems, vol. 7, 2025. [9] C. Eckart and G. Young, “The approximation of one matrix by another of lower rank,” Psychometrika, vol. 1, no. 3, pp. 211–218, 1936. [10] G. Glorian, I. Lamprou, Z. Zhang, Y. Yuan, and H. Liu, “LVSA: Training-free sparse attention for long video diffusion,” arXiv preprint arXiv:2605.31057, 2026. [11] G. H. Golub and C. F. Van Loan, Matrix computations. JHU press, 2013. [12] Y. HaCohen, N. Chiprut, B. Brazowski, D. Shalem, D. Moshe, E. Richardson, E. Levin, G. Shiran, N. Zabari, O. Gordon, P. Panet, S. Weissbuch, V. Kulikov, Y. Bitterman, Z. Melumian, and O. Bibi, “LTX-Video: Realtime video latent diffusion,” arXiv preprint arXiv:2501.00103, 2025. [13] J. Ho, A. Jain, and P. Abbeel, “Denoising diffusion probabilistic models,” in Advances in Neural Information Processing Systems, vol. 33, 2020, pp. 6840–6851. [14] Z. Huang, Y. He, J. Yu, F. Zhang, C. Si, Y. Jiang, Y. Zhang, T. Wu, Q. Jin, N. Chanpaisit, Y. Wang, X. Chen, L. Wang, D. Lin, Y. Qiao, and Z. Liu, “VBench: Comprehensive benchmark suite for video generative models,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2024, pp. 21 807–21 818. [15] Hunyuan Foundation Model Team, “Penguin video benchmark,” GitHub repository, 2025, accessed: Jul. 27, 2026. [Online]. Available: https://github.com/Tencent-Hunyuan/HunyuanVideo/ blob/main/assets/PenguinVideoBenchmark.csv [16] H. Jiang, Y. Li, C. Zhang, Q. Wu, X. Luo, S. Ahn, Z. Han, A. H. Abdi, D. Li, C.-Y. Lin, Y. Yang, and L. Qiu, “MInference 1.0: Accelerating prefilling for long-context LLMs via dynamic sparse attention,” in Advances in Neural Information Processing Systems, vol. 37, 2024.

12

[41] Z. Xing, Q. Feng, H. Chen, Q. Dai, H. Hu, H. Xu, Z. Wu, and Y.-G. Jiang, “A survey on video diffusion models,” ACM Computing Surveys, vol. 57, no. 2, 2025. [42] R. Xu, G. Xiao, H. Huang, J. Guo, and S. Han, “XAttention: Block sparse attention with antidiagonal scoring,” in Proceedings of the 42nd International Conference on Machine Learning, ser. Proceedings of Machine Learning Research, vol. 267, 2025, pp. 69 819–69 831. [43] S. Yadav and B. Asgari, “Boötes: Boosting the efficiency of sparse accelerators using spectral clustering,” in Proceedings of the 58th Annual IEEE/ACM International Symposium on Microarchitecture, 2025. [44] R. Yan, Y. Jiang, and B. Yuan, “Flash sparse attention: An alternative efficient implementation of native sparse attention kernel,” arXiv preprint arXiv:2508.18224, 2025. [45] S. Yang, H. Xi, Y. Zhao, M. Li, J. Zhang, H. Cai, Y. Lin, X. Li, C. Xu, K. Peng, J. Chen, S. Han, K. Keutzer, and I. Stoica, “Sparse VideoGen2: Accelerate video generation with sparse attention via semantic-aware permutation,” in Advances in Neural Information Processing Systems, vol. 38, 2025. [46] Z. Yang, J. Teng, W. Zheng, M. Ding, S. Huang, J. Xu, Y. Yang, W. Hong, X. Zhang, G. Feng, D. Yin, Y. Zhang, W. Wang, Y. Cheng, B. Xu, X. Gu, Y. Dong, and J. Tang, “CogVideoX: Text-to-video diffusion models with an expert transformer,” in International Conference on Learning Representations, 2025. [47] Z. Ye, L. Chen, R. Lai, W. Lin, Y. Zhang, S. Wang, T. Chen, B. Kasikci, V. Grover, A. Krishnamurthy, and L. Ceze, “FlashInfer: Efficient and customizable attention engine for LLM inference serving,” in Proceedings of Machine Learning and Systems, vol. 7, 2025. [48] T. Zadouri, M. Hoehnerbach, J. Shah, V. Thakkar, and T. Dao, “FlashAttention-4: Algorithm and kernel pipelining co-design for asymmetric hardware scaling,” in Proceedings of Machine Learning and Systems, vol. 8, 2026. [49] J. Zhang, C. Xiang, H. Huang, J. Wei, H. Xi, J. Zhu, and J. Chen, “SpargeAttention: Accurate and training-free sparse attention accelerating any model inference,” in Proceedings of the 42nd International Conference on Machine Learning, ser. Proceedings of Machine Learning Research, vol. 267, 2025, pp. 76 397–76 413. [50] L. Zhang, J. S. Marron, H. Shen, and Z. Zhu, “Singular value decomposition and its visualization,” Journal of Computational and Graphical Statistics, vol. 16, no. 4, pp. 833–854, 2007. [51] P. Zhang, Y. Chen, R. Su, H. Ding, I. Stoica, Z. Liu, and H. Zhang, “Fast video generation with sliding tile attention,” in Proceedings of the 42nd International Conference on Machine Learning, ser. Proceedings of Machine Learning Research, vol. 267, 2025, pp. 74 714–74 731.

AI U SE ChatGPT-5.6 Extra High was used solely to polish the manuscript’s language and improve readability. All other aspects of the work, including problem identification, algorithm and system design, implementation, experimental evaluation, and analysis of results, were conducted independently by the authors without AI assistance.

13

Record · ID 1108706 · SHA-256 d480bec0eac440dd
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.