wan2.2 is faster than FA2 but consumes more VRAM.

[dz1iang]

Hello, thank you for the open-source SLA project. My experience with using it on wan2.2 is as follows:

1. The training performance has improved by over 30%;
2. In scenarios where FA2 previously did not encounter OOM (Out of Memory), the SLA project now experiences OOM issues.

Finally, I would like to ask two questions:

1. Why does this (OOM issue) occur?
2. Are there any optimization techniques for SLA, particularly regarding computing efficiency and VRAM (Video Random Access Memory) usage?

[Lyxien]

Hello, I'd like to ask, how is the video quality?

[dz1iang]

Hello, I'd like to ask, how is the video quality?

Currently, video generation with small-scale training works normally, and no specific evaluation has been conducted.

[whx1003]

@dz1iang You’re right that SLA consumes more memory, since we need to precompute and store the results of the linear components (as described in Algorithms 1 and 2 of the paper). We estimate the additional memory required to be roughly $\frac{d}{B}$ times of the original, where $d$ is the head dimension ($128$ for Wan models) and $B$ is the block size (either $64$ or $128$). In practice, SLA may use around twice as much memory as FA2.

Unfortunately, the only tunable parameter is the block size $B$, and currently we only support values of 64 or 128, as shown here:

SLA/sparse_linear_attention/kernel.py

Lines 388 to 389 in 283c5b3

	assert BLOCK_M == 64 or BLOCK_M == 128
	assert BLOCK_N == 64

In general, a larger block size requires more shared memory, and based on our tests on an RTX 5090, 128 is the largest block size we can support.

[dz1iang]

@dz1iang You’re right that SLA consumes more memory, since we need to precompute and store the results of the linear components (as described in Algorithms 1 and 2 of the paper). We estimate the additional memory required to be roughly d B times of the original, where d is the head dimension ( 128 for Wan models) and B is the block size (either 64 or 128 ). In practice, SLA may use around twice as much memory as FA2.

Unfortunately, the only tunable parameter is the block size B , and currently we only support values of 64 or 128, as shown here:

SLA/sparse_linear_attention/kernel.py

Lines 388 to 389 in 283c5b3

assert BLOCK_M == 64 or BLOCK_M == 128
assert BLOCK_N == 64
In general, a larger block size requires more shared memory, and based on our tests on an RTX 5090, 128 is the largest block size we can support.

@whx1003 Thank you for your reply. Additionally, I would like to ask: for GPUs with other architectures, is it sufficient to adjust the above-mentioned block size, or are there other aspects that can be debugged? For example, GPUs with Architecture Ampere or Architecture Hopper.

[jt-zhang]

We’ve just updated the Triton implementation of SLA. Training is now more stable, faster, and typically achieves better training results.

We’ve released the code for SageSLA, a very fast SLA (Sparse-Linear Attention) forward pass based on SageAttention. It uses some code from SpargeAttn. Please refer to the SageSLA/ directory for how to use SageSLA.

Feel free to try it out! Now SageSLA is very fast on RTX and Hopper GPUs.