VMoBA: Mixture-of-Block Attention for Video Diffusion Models

1 Introduction

• Full attention scales quadratically with sequence length, while sparse attention for VDMs has primarily targeted inference and remains underexplored for training.The paper notes that a 720p video can exceed 76K tokens.Source: §1 · p. 2
• The authors attribute vanilla MoBA’s VBench drop from 68.25 to 56.88 to its text-oriented 1D locality design mismatching video’s 3D locality.Flattening the video latent space can separate spatially adjacent tokens across blocks, diluting block means.Source: §1 · p. 2
• Video attention exhibits 1D-2D-3D locality, varying query importance, and head-specific similarity concentration, challenging uniform block selection.The paper notes that equal block counts may under-allocate to queries with stronger key affinities, while fixed top-k ignores differences between heads.Source: §1 · p. 2
• VMoBA combines layer-wise recurrent 1D-2D-3D partitions, global selection of salient query-key blocks, and cumulative-similarity threshold selection.The partition scheme adapts across layers; global selection pools block interactions across a head, and thresholding avoids a fixed top-k.Source: §1 · pp. 2–3
• VMoBA enables training VDMs on longer sequences, outperforming vanilla MoBA while achieving comparable or superior generation quality to full attention.It also demonstrates competitive performance in training-free settings, offering speedups for long sequences.Source: §1 · p. 3

2 Related Work

• Diffusion inference acceleration reduces denoising steps through samplers, noise schedules, or reuse of features, while distillation trains a student to use fewer steps.The passage notes that training-free methods can be hyperparameter-sensitive and degrade fidelity, while distillation needs extra large-scale training and may struggle to scale resolution.Source: §2 · p. 3
• Linear attention can achieve linear complexity but may require architectural changes and can underperform full attention on some tasks.Sparse methods such as MoBA, NSA, and Block-Attention can switch between full and sparse attention modes without those architectural changes.Source: §2 · p. 3
• Video-specific sparse attention methods mainly accelerate inference, whereas VMoBA is presented as supporting VDM training on longer sequences and inference acceleration.The passage describes VMoBA inference as lossless in quality.Source: §2 · p. 3

3.1 Overview

• VMoBA partitions key blocks, selects blocks per attention head, then computes attention only over selected key-value pairs.Global selection pools query-key block similarities across queries; threshold selection sets the number of selected blocks using cumulative similarity.Source: §3.1 · pp. 3–4, Figure 2 · p. 4

3.2 Layer-wise Recurrent Block Partition

• A layer-wise recurrent 1D-2D-3D partition scheme matches blocks to temporal, spatial, or combined spatio-temporal patterns.The scheme recurs across layers and reduces block count compared with uniform 3D partitioning.Source: §3.2 · pp. 4–5, Figure 3 · p. 4

3.3 Global Block Selection

• Global Block Selection pools query-key block similarities across queries and selects the most significant interactions for each attention head.This replaces per-query fixed block selection, addressing variation in query importance.Source: §3.3 · pp. 5–6, Figure 4 · p. 5

3.4 Threshold-based Block Selection

• Threshold-based selection sorts global similarities and dynamically selects blocks until their normalized cumulative similarity exceeds a threshold.This accommodates head-specific concentration differences rather than using a fixed top-k.Source: §3.4 · p. 6, Figure 5 · p. 6

4.1 Setup

• Experiments use Wan 2.1 1.3B on Koala-36M, assess five VBench aspects plus PSNR, and compare MoBA, DiTFastAttn, SVG, and Full Attention.Models in each setting use identical data; the training-based evaluation uses prompts after optimization.Source: §4.1 · pp. 6–7

4.2 Quantitative Results

• At 76K tokens, VMoBA achieves 1.35x speedup, 19% above vanilla MoBA, and leads BGConsis while ranking second in TextConsis and ImageQual.At original resolution, its speedup over Full Attention is 1.01x, compared with MoBA's 0.83x.Source: Table 1 · p. 7, §4.2 · p. 7
• At 55K-token spatial extension, VMoBA reaches VBench mean 68.34 vs Full Attention 68.25 and 1.48x speedup; temporal extension yields 1.44x speedup.The spatial result uses 187 GPU Hours for VMoBA versus 276 GPU Hours for Full Attention.Source: Table 2 · p. 7, §4.2 · p. 8

4.3 Qualitative Results

• Figure 6 contrasts settings: VMoBA adds prompt-aligned cooking details at high resolution and shows more realistic scenes after temporal extension.For spatially extended training, the coastal scene generated by VMoBA is described as aesthetically pleasing and consistent with the prompt, unlike Full Attention's unclear coast depiction.Source: Figure 6 · p. 8, §4.3 · pp. 8–9

4.4 Ablation Study

• Removing any one partition pattern generally lowers performance; 2-3D scores slightly higher but takes longer, indicating 1D contributes to efficiency.The authors interpret the results as supporting diverse partitioning schemes for capturing spatio-temporal locality.Source: Table 3 · p. 9, §4.4 · p. 9
• The threshold-plus-global block-selection strategy performs best across all metrics; removing either component lowers performance.Threshold selection adapts the number of blocks, while global selection chooses blocks from a global pool.Source: Table 3 · p. 9, §4.4 · p. 9
• Larger thresholds generally improve performance but increase training time; the authors identify 0.25 as a balance.The passage attributes the trade-off to selecting more blocks at larger thresholds.Source: Table 3 · p. 9, §4.4 · p. 9
• More blocks tend to improve performance while increasing training time; the default 8-48-72 setting offers comparable scores at lower cost.Source: Table 3 · p. 9, §4.4 · p. 9

• VMoBA is a mixture-of-block sparse attention mechanism designed to address computational challenges in training long-sequence VDMs.Source: §5 · p. 9
• Its three innovations are layer-wise recurrent block partitioning, global block selection, and threshold-based block selection.Source: §5 · p. 9
• VMoBA achieves speedups of up to 1.48x while maintaining or improving generation quality compared with full attention.The paper also reports competitive speedups in training-free inference settings.Source: §5 · p. 9

Appendix A

• Appendix A covers implementation details of the experiments.Source: Appendix · p. 12

Appendix B

• Appendix B contains additional experiments.Source: Appendix · p. 12

Appendices C–D

• Appendix C compares pre-training loss between VMoBA and Full Attention, and Appendix D covers limitations and future work.Source: Appendix · p. 12

• Table 4 presents implementation details for the main-paper experiments.Source: Table 4 · p. 12, A Implementation Details · p. 12
• Training-free experiments use 3D block partitioning with the original MoBA local and topk query-key block selection options.Source: A Implementation Details · p. 12
• Training experiments use the full VMoBA setting, which the authors report has stronger performance and efficiency gains than MoBA.Source: A Implementation Details · p. 12

B Additional Experiments

• Table 5 reports training results for the original and shorter sequence lengths.The passage does not provide the table's individual values.Source: Table 5 · p. 12
• VMoBA consistently outperforms MoBA in both performance and efficiency across the two additional sequence-length experiments.The experiments cover original-length 33K and shorter-length 13K sequences.Source: B Additional Experiments · p. 12
• At 33K, VMoBA's speed is very close to the pre-trained full-attention model, while at 13K it is slower.The authors attribute this pattern to VMoBA's acceleration becoming more significant at longer sequence lengths.Source: B Additional Experiments · p. 12

C Pre-training Loss Comparison

• At 46K, VMoBA and full attention have closely overlapping pre-training validation-loss curves.The authors conclude VMoBA is a reliable alternative to full attention for video-generation-model pre-training.Source: C Pre-training Loss Comparison · p. 13, Figure 7 · p. 13
• At 11K, VMoBA's validation loss remains higher, but the gap from full attention shrinks as model size increases.The passage says VMoBA performs closely with full attention at larger model scales.Source: C Pre-training Loss Comparison · p. 13, Figure 7 · p. 13
• The experiments train at 11K and 46K tokens, with 1-2-3D block counts of 10, 48, and 60 and threshold τ = 0.25.The model uses the Wan 2.1 architecture, with VMoBA replacing all self-attention layers; Table 6 gives the configuration.Source: C Pre-training Loss Comparison · p. 13, Table 6 · p. 13

D Limitations and Future Work

• At short sequence lengths, VMoBA's speedup is suboptimal even though its FLOPs are lower.The authors say this may reflect inconsistent memory distribution among selected blocks.Source: D Limitations and Future Work · p. 13
• The current FlashAttention-based implementation may not fully optimize the speed of sparse attention.This is presented as a possible explanation for the short-sequence speedup limitation.Source: D Limitations and Future Work · p. 13
• Future work may focus on more efficient, hardware-friendly implementations to bring practical speed closer to theoretical FLOPs.Source: D Limitations and Future Work · p. 13