Implementing NATTEN for 4K image generation with Flux as mentioned in GNA paper

[xuanmingShang]

Hello,

I'm working on implementing the 4K image generation approach using NATTEN as described in the GNA paper. The paper demonstrates accelerating Flux for high-resolution image generation.

Environment
Hardware: A800 GPU
NATTEN version: 0.20.0
CUDA: 12.6
PyTorch: 2.7.0

Implementation Approach

I'm modifying the attention processor of URAE (which enables Flux to generate 4K images) to leverage na1d for efficient attention computation. For image generation tasks, each image token needs to interact with text tokens - a pattern that might be described as "self-cross attention" as mentioned in issue #82.

My implementation leverages the additional_key/value parameters of na1d to enable this cross-domain interaction. I'm basing my approach on the URAE attention processor implementation found here: URAE/attention_processor.py#L82

hidden_states = na1d(query.transpose(1, 2)[:,512:,:,:], key.transpose(1, 2)[:,512:,:,:], value.transpose(1, 2)[:,512:,:,:], \
                        kernel_size=80, stride=16, \
                        additional_keys=key.transpose(1, 2)[:,:512,:,:], additional_values=value.transpose(1, 2)[:,:512,:,:], \
                        backend="cutlass-fna", \
                        attention_kwargs={"backend": "cutlass-fmha"} )
text_hidden_states = na1d(query.transpose(1, 2)[:,:512,:,:], key.transpose(1, 2)[:,:512,:,:], value.transpose(1, 2)[:,:512,:,:], \
                        kernel_size=512, \
                        additional_keys=key.transpose(1, 2)[:,512:,:,:], additional_values=value.transpose(1, 2)[:,512:,:,:], \
                        backend="cutlass-fna", \
                        attention_kwargs={"backend": "cutlass-fmha"} )
hidden_states = torch.cat([text_hidden_states, hidden_states], dim=1)
hidden_states = hidden_states.reshape(batch_size, -1, attn.heads * head_dim)

When running inference with this implementation, I'm getting noisy/corrupted images rather than proper high-resolution outputs. I've experimented with different parameter configurations but haven't been able to achieve the results described in the paper.

Questions

• Is there any official implementation or example of using NATTEN for 4K image generation with Flux, as mentioned in the GNA paper?
• Are there specific considerations or configurations needed when using na1d with additional keys/values for text-image cross-attention?
• Could you identify any issues in my implementation that might be causing the noisy outputs?

I appreciate any guidance or pointers to relevant resources. Thank you for your time!

[alihassanijr]

Thank you for your interest.

We will work on open sourcing our implementation on top of FLUX soon, but yes there's a few issues I see with the above implementation.

The first is that we apply 2D Neighborhood Attention with kernel size 80x80 and not 1D Neighborhood Attention with kernel size 80.

In our setup, when generating 4K images the feature map ends up being 256x256, or 65K tokens.
The 1D implementation above introduces more than 99.8% sparsity (1 - 80/65K), and because it is 1D spatial locality is broken, so that probably explains part of the problem.

The second issue is that the text hidden states must go through self attention, not neighborhood attention, in MMDiT-like architectures.
Sparsity is introduced only into the interaction between image tokens and themselves.

Your use of additional_{keys,values} appears correct; with it each image token attends to all text tokens, again as expected.

In short, these are the changes I'd make:

• hidden_states = na1d -> hidden_states = na2d

  • note that you should also reshape query and key and value from [batch, 65K, heads, head_dim] to [batch, 256, 256, heads, head_dim], as na2d expects 5D tensors for query, key, and value. additional_{keys,values} must remain 4D and NATTEN will handle necessary reshapes internally.

• text_hidden_states = na1d -> text_hidden_states = attention

  • You should also remove kernel_size, stride, and additional_{keys,values} from this, and pass in the entire key and value tensors without breaking them into text and image tokens.
  • For attention, you can use SDPA, or your choice of attention. NATTEN also ships standard dot product attention: from natten import attention.

CC @Rbrq03 @AdityaKane2001 since they're more familiar with the FLUX implementation.

[alihassanijr]

By the way I do have one suggestion for your implementation.

You correctly set the backends for your GPU architecture, which is Ampere, and cutlass-fna/cutlass-fmha are the best backends for Ampere. However, those backends are not as fast as Flash Attention 2, the default baseline, so your speedup may be more limited than we observed in the paper.

For now, only our Hopper FNA/FMHA and Blackwell FNA/FMHA can provide you with FLOP-proportional speedups.

I don't have access to A800s, but on the A100 SXM, here's what the standard self attention with FA2 looks like:

python -m natten.profiler --backend fav2 -n 24 -d 128 -i 256 256

|===========================================================================================|
|                                      Profiler results                                     |
|===========================================================================================|
+===========+=================+======+================================+=========+===========+
| Framework | Kernel category | Arch | Operation                      | # calls | Runtime   |
+===========+=================+======+================================+=========+===========+
| CUTLASS   | -               | -    | pytorch_flash.flash_fwd_kernel | 1       | 256.208ms |
+===========+=================+======+================================+=========+===========+
|           |                 |      | Total                          |         | 256.208ms |
+===========+=================+======+================================+=========+===========+

And here's FNA with tile shape auto-tuning, which you can easily do using our profiling toolkit:

python -m natten.profiler -n 24 -d 128 -i 256 256 -w 80 80 -s 16 16 --optimize

Searching 27 forward pass configs

|==============================|
|      Best configuration      |
|==============================|
+===============+==============+
| Parameter     | Value        |
+===============+==============+
| backend       | cutlass-fna  |
| fmha_backend  | cutlass-fmha |
| q_tile_shape  | (4, 16)      |
| kv_tile_shape | (8, 16)      |
+===============+==============+


|======================================================================|
|                           Profiler results                           |
|======================================================================|
+===========+=================+======+============+=========+==========+
| Framework | Kernel category | Arch | Operation  | # calls | Runtime  |
+===========+=================+======+============+=========+==========+
| CUTLASS   | attention       | Sm80 | FnaForward | 1       | 56.074ms |
+===========+=================+======+============+=========+==========+
|           |                 |      | Total      |         | 56.074ms |
+===========+=================+======+============+=========+==========+

Of course I would expect Attention to be 80+% of the workload with Ampere, so even with the above 4.6X speedup, you could see at least 2.7X speedup in each diffusion step you apply GNA. At the same time, FLOP-proportional speedup with this level of sparsity is 10.2X, so the expected speedup from that would be ~3.6, which is unattainable at the moment.

[xuanmingShang]

Wow, thank you for the incredibly detailed explanation! This is exactly what I was looking for and it solves my problem perfectly.

[xuanmingShang]

Hello,

Thank you so much for your detailed responses in the paper and previous issues. They have been incredibly helpful.

Following your guidance and the approach described in the GNA paper, I successfully implemented training-free sparse 4K text-to-image generation by using the original self-attention for the first 9 steps and sparse attention for the remaining steps.

However, to further improve inference speed, I attempted to use sparse attention for all denoising steps. This leads to a new problem: the generated images exhibit repeating patterns, where the target content appears N times within a single image.

To address this issue, I trained the model on my own 4K dataset using flow matching loss, but the training did not eliminate the repeating content problem.

I would like to ask:

1. Did you encounter similar repeating pattern issues when training your 4K text-to-image model with sparse attention?
2. What strategies or techniques did you employ during training to mitigate or eliminate this repetition problem?
3. Do you have any suggestions or references for training with "full-sparse attention" (sparse attention at all steps)?

Any insights or guidance would be greatly appreciated. Thank you for your time and help!
@alihassanijr @Rbrq03 @AdityaKane2001

[alihassanijr]

We didn't train any models for GNA, it was all out-of-the-box inference.

And yes, if we use it for all diffusion steps with kernel size 80x80, even standard NA will exhibit this problem:

By introducing the full self attention steps in the GNA work, this problem can be largely mitigated -- for FLUX 4k we did 9 out of 28 steps:

[alihassanijr]

However, it's always useful to choose a finer-grained target.
Using not just self attention, but different levels of sparsity across diffusion steps, layers, even attention heads, is likely going to provide the best trade-off.

One instance I can point out that I worked on was Cosmos Predict2 (this is a world model and not just an image gen model).

Over there each DiT layer has different GNA parameters, based on the error introduced with respect to self attention, and according to some threshold which we set by conducting small experiments.

In the end, some layers ended up with as high as 98% sparsity, some with 95, some 90, some 75, and some with only 50%.
We kept at most 1 self attention layer.

However, this was only possible because we fine-tuned the model for a bit using these settings.