# How Linear Attention Improves Efficiency in DiT for High‑Resolution Images

> Discover how Linear Attention slashes DiT complexity from O(N²) to O(N), enabling 4K image generation on consumer GPUs by overcoming standard self-attention's memory bottleneck.

- Repository: [NVIDIA Research Projects/Sana](https://github.com/NVlabs/Sana)
- Tags: deep-dive
- Published: 2026-05-19

---

**Linear Attention reduces DiT complexity from O(N²) to O(N) using kernel‑based factorization, enabling 4K image generation on consumer GPUs by eliminating the quadratic memory bottleneck of standard self‑attention.**

Sana, the open‑source text‑to‑image and video generation suite from NVlabs, replaces vanilla Diffusion Transformer (DiT) blocks with a linear‑complexity attention mechanism. This architectural choice directly addresses the prohibitive cost of processing high‑resolution latents that contain tens of thousands of tokens after auto‑encoder compression.

## Kernel‑Based Linearization in WanLinearAttention

Standard transformer self‑attention scales quadratically with sequence length, making high‑resolution image generation computationally expensive. Linear Attention rewrites the softmax operation as a separable kernel product that permits associative reordering of matrix multiplications.

### Mathematical Foundation

The attention computation transforms from:

\[
\text{Attn}(Q,K,V)=\frac{(\phi(Q)\,\phi(K)^{\top})V}{\phi(Q)\,\phi(K)^{\top}\mathbf{1}}
\]

Where \(\phi(\cdot)\) represents a positive kernel function. Because the kernel is separable, the matrix multiplication rearranges to compute cumulative statistics rather than pairwise token interactions. This reduces memory and compute requirements from O(N²) to O(N), where N represents the number of latent tokens.

### Implementation Details

The `WanLinearAttention` class in [`diffusion/model/wan/model.py`](https://github.com/NVlabs/Sana/blob/main/diffusion/model/wan/model.py) implements this mechanism using ReLU as the kernel function:

```python

# diffusion/model/wan/model.py

class WanLinearAttention(WanSelfAttention):
    def __init__(self, dim, num_heads, ...):
        super().__init__(dim, num_heads, ...)
        self.kernel_func = nn.ReLU(inplace=False)   # φ(·)

    @torch.autocast(device_type="cuda", enabled=False)
    def attn_matmul(self, q, k, v):
        # pad V to keep dimensions compatible

        v = F.pad(v.float(), (0, 0, 0, 1), mode="constant", value=self.PAD_VAL)
        # (V @ K)   →  (d_v, d_k)   cumulative sum over tokens

        vk = torch.matmul(v, k)
        # (vk @ Q) → (d_v, d_q)   per‑query projection

        out = torch.matmul(vk, q)
        # normalise by the extra padding term

        norm_out = out[:, :, :-1] / (out[:, :, -1:] + self.eps)
        return norm_out

```

The `forward` method applies `self.kernel_func` to queries and keys before invoking `attn_matmul`. This kernel application happens per‑token, while the matrix products operate globally across the sequence, maintaining linear scaling with token count.

## Memory and Compute Efficiency

Linear Attention delivers three critical advantages for high‑resolution synthesis:

- **Linear Complexity**: Processing time scales proportionally with token count rather than quadratically. A 4K × 4K image generating ~25,000 tokens after 32× compression remains feasible, whereas vanilla attention would create a 625‑million‑element attention matrix.
- **Reduced Memory Footprint**: The implementation stores only cumulative statistics (sum‑vectors of size `(d_k, d_v)`) rather than the full N × N attention matrix. This frees GPU memory for larger batch sizes or higher‑resolution latents.
- **Inference Speed**: According to the Sana source documentation, Linear DiT achieves approximately **10× speed‑up** over vanilla DiT when generating 1024 × 1024 images, reducing generation time from minutes to seconds.

## Video Generation and Constant‑Memory KV Cache

For video synthesis, token counts explode further as each frame adds a temporal dimension. Sana extends the linear formulation to **Block‑Linear Attention**, which maintains only cumulative statistics per temporal block. This creates a constant‑memory KV cache that enables minute‑length video generation without the memory explosion typical of standard attention mechanisms. The block‑wise implementation resides in the scheduler components under `diffusion/scheduler/`.

## Using Linear DiT in Practice

### Loading the Model with Diffusers

The `SanaPipeline` automatically loads the Linear DiT architecture when using the official checkpoints:

```python
import torch
from diffusers import SanaPipeline

# This pipeline internally loads the linear‑attention DiT architecture.

pipe = SanaPipeline.from_pretrained(
    "Efficient-Large-Model/SANA-Image_0.6B_diffusers",
    torch_dtype=torch.bfloat16,
)
pipe.vae.to(torch.float32)               # VAE runs in FP32 for stability

pipe.text_encoder.to(torch.bfloat16)     # Small decoder‑only LLM

pipe.to("cuda")

prompt = "A hyper‑realistic 4K portrait of a tiger in a jungle, photorealistic, intricate details."
image = pipe(prompt=prompt, height=4096, width=4096, guidance_scale=7.5,
             num_inference_steps=30).images[0]
image.save("tiger_4k.png")

```

The model registration in [`diffusion/model/wan/model_wrapper.py`](https://github.com/NVlabs/Sana/blob/main/diffusion/model/wan/model_wrapper.py) aliases `WanLinearAttentionModel` as `SanaWanLinearAttentionModel`, ensuring the pipeline instantiates the linear‑attention variant rather than standard quadratic attention.

### Verifying Linear Attention Layers

You can confirm that a loaded model uses Linear Attention by inspecting the transformer blocks:

```python
from diffusers import SanaPipeline

pipe = SanaPipeline.from_pretrained("Efficient-Large-Model/SANA-Image_0.6B_diffusers")

# Grab the first DiT block

block = pipe.transformer.blocks[0]
print(type(block.attn))   # → <class 'diffusion.model.wan.model.WanLinearAttention'>

```

This verification confirms the presence of `WanLinearAttention`, the linear‑complexity implementation that enables efficient high‑resolution generation.

## Summary

- **Linear Attention** replaces the O(N²) self‑attention in DiT with an O(N) kernel‑based factorization, implemented in [`diffusion/model/wan/model.py`](https://github.com/NVlabs/Sana/blob/main/diffusion/model/wan/model.py) as `WanLinearAttention`.
- The **ReLU kernel** enables separable matrix operations that compute cumulative statistics rather than pairwise attention maps, reducing memory usage from N × N matrices to constant‑size sum‑vectors.
- **High‑resolution viability**: Linear DiT generates 4096 × 4096 images in under a second on 16 GB GPUs, compared to the prohibitive costs of vanilla attention.
- **Video extension**: Block‑Linear Attention provides constant‑memory KV caching for temporal sequences, supporting extended video generation without memory constraints.

## Frequently Asked Questions

### What is the computational complexity of Linear Attention in Sana?

Linear Attention operates in **O(N)** time and memory complexity relative to token count, compared to the O(N²) scaling of standard self‑attention. This linear scaling results from kernel factorization that eliminates the need to compute and store the full N × N attention matrix.

### How does Linear Attention differ from standard self‑attention?

Standard self‑attention computes softmax‑scaled dot‑products between all token pairs, requiring quadratic memory. Linear Attention uses a **ReLU kernel** to decompose the attention operation into separable components, allowing the model to compute cumulative statistics (KᵀV) once and project all queries against these statistics.

### Can Linear Attention handle 4K image generation on consumer hardware?

Yes. According to the NVlabs/Sana documentation, Linear DiT can generate **4096 × 4096 resolution images** on GPUs with 16 GB VRAM. The linear memory scaling prevents the out‑of‑memory errors that would occur with quadratic attention when processing the ~25,000 tokens produced by a 32× compressed 4K latent.

### Is Linear Attention used for both images and video in Sana?

Yes. While standard Linear DiT handles image generation, Sana extends the mechanism to video through **Block‑Linear Attention**. This variant maintains cumulative statistics per temporal block, creating a constant‑memory KV cache that supports minute‑length video generation without the memory explosion typical of full attention across space‑time tokens.