NATTEN1D is slower than NATTEN2D?

[Yunge6666]

I'm curious why NATTEN1D is much slower than NATTEN2D when using the same kernel size.
Specifically, my settings are 2D kernel size = 7×7, 1D kernel size = 49, dim = 128, number of head = 4, feature map/sequence length = 28×28=784

[alihassanijr]

Can you please specify what version of NATTEN you're using, on what GPU, and what backend you're using (if selecting a specific one).

[Yunge6666]

Can you please specify what version of NATTEN you're using, on what GPU, and what backend you're using (if selecting a specific one).

I am using RTX3080 and the version of NATTEN is 0.14.4. When I was testing, I used from natten import NeighborhoodAttention2D, NeighborhoodAttention1D and tested them.

[alihassanijr]

We strongly recommend using newer builds, ideally the latest (0.21.0).

We can't continue supporting builds earlier than 0.17.5 anymore.

If I recall correctly, 0.14.X had some specialized kernels just for 2D and for kernel sizes 3x3, 5x5, 7x7, 9x9, 11x11, and 13x13, and only specific head dims. Anything outside of that would use the default naive kernels.
That could potentially explain the performance difference you're seeing, but even if not, those kernels were just a proof of concept intended to be functionally correct, and are not optimized.

0.17.5 and later ship our fused kernels, so changes to your code will be necessary.

Performance is still not the best it can be for your GPU unfortunately, but it's significantly better than what you get with 0.14.X. Our future updates will ship better kernels for RTX cards as well.

[Yunge6666]

We strongly recommend using newer builds, ideally the latest (0.21.0).

We can't continue supporting builds earlier than 0.17.5 anymore.

If I recall correctly, 0.14.X had some specialized kernels just for 2D and for kernel sizes 3x3, 5x5, 7x7, 9x9, 11x11, and 13x13, and only specific head dims. Anything outside of that would use the default naive kernels. That could potentially explain the performance difference you're seeing, but even if not, those kernels were just a proof of concept intended to be functionally correct, and are not optimized.

0.17.5 and later ship our fused kernels, so changes to your code will be necessary.

Performance is still not the best it can be for your GPU unfortunately, but it's significantly better than what you get with 0.14.X. Our future updates will ship better kernels for RTX cards as well.

Thank you for your answer. I'd actually like to try version 0.17.5 or later, but I'm getting a message saying it only supports CUDA 12 and higher. Due to a driver issue, my current GPU can only use CUDA 11.8, but I can use PyTorch 2.7.0 and higher. I'm not sure if the requirement for 0.17.5 and later is specific to PyTorch, or if it also requires CUDA. Is CUDA 12 or higher mandatory?

[alihassanijr]

The CUDA 12.X requirement wasn't enforced until the 0.20.0 release; support for 11.X was just deprecated, but you can still build it from source, and I think 0.17.5 and earlier builds will still work out of the box with CUDA 11.X.

As long as you have PyTorch and the CUDA toolkit available, you should be able to just checkout v0.17.5 and build for your GPU architecture (Ada, SM89):

git clone https://github.com/SHI-Labs/NATTEN

cd NATTEN

git checkout v0.17.5

make CUDA_ARCH="8.9"

But of course I'd recommend upgrading your driver sometime. Older runtimes and compilers can leave performance on the table, or be discontinued with newer software other than NATTEN (i.e. PyTorch).

With CTK 13.0 being released a while back, I suspect support for CTK 11.X will start to fade away soon.

[Yunge6666]

Ok. Thanks for the suggestion. I previously installed directly using pip, but kept getting errors. I'll try compiling from source as you suggested. Also, I saw that someone else discussed the issue with arbitrary masks before. Is the current code still unsupported?

[alihassanijr]

I previously installed directly using pip, but kept getting errors

Please share the install.out log file when you run into build issues and I'll be happy to take a look.

Also, I saw that someone else discussed the issue with arbitrary masks before

Sorry what issue? CUDA 11.X support? Performance differences between 1D and 2D? Arbitrary masks in general?

NATTEN has never supported arbitrary masks.

Is the current code still unsupported?

Can you please clarify what exactly you're referring to by "the current code"?

[Yunge6666]

I'm sorry for the confusion. The arbitrary mask I mentioned is a separate question. It has nothing to do with my CUDA version; it's a feature I might need in my project. I saw someone ask about custom masks. Actually, what I wanted to ask is, is it possible to have different neighborhoods for each token in NATTEN?

[alihassanijr]

I see; unfortunately no, token-level parameterization is not supported in NATTEN. It would be part of that arbitrary mask support, which I might work on at some point in the future.

[Yunge6666]

I see; unfortunately no, token-level parameterization is not supported in NATTEN. It would be part of that arbitrary mask support, which I might work on at some point in the future.

Ok, thanks for your information.

We strongly recommend using newer builds, ideally the latest (0.21.0).

We can't continue supporting builds earlier than 0.17.5 anymore.

If I recall correctly, 0.14.X had some specialized kernels just for 2D and for kernel sizes 3x3, 5x5, 7x7, 9x9, 11x11, and 13x13, and only specific head dims. Anything outside of that would use the default naive kernels. That could potentially explain the performance difference you're seeing, but even if not, those kernels were just a proof of concept intended to be functionally correct, and are not optimized.

0.17.5 and later ship our fused kernels, so changes to your code will be necessary.

Performance is still not the best it can be for your GPU unfortunately, but it's significantly better than what you get with 0.14.X. Our future updates will ship better kernels for RTX cards as well.

I used another device from my lab, which has an RTX3080, CUDA 12.6, and PyTorch 2.7.0. I then installed the latest NATTEN (0.20.1) and ran 1D and 2D tests. I found that at small sizes, such as 56x56=3136, and a kernel size of 7x7=49, head =2, dims=64, there was almost no difference in efficiency between 1D and 2D. However, when I increased the size to 128x128=16384, still using the same kernel size, the number of heads and dims, I found that 1D was about 1.25 times faster than 2D. I'm wondering if this is normal?

[alihassanijr]

In short, yes; 2D being slower than 1D is normal (but 1D being slower than 2D is unexpected).

For starters, with every additional dimension, there is more room for "wasted FLOPs", which means unless you're using something like stride to avoid that, your 2D problem will just always have a worse throughput because it's wasting FLOPs.

However, the bulk of what users typically feel is usually a performance issue in our original FNA kernels.

Your GPU is Ada (SM89), which means you can only our CUTLASS FNA backend (and Flex backend, but with limited support).

With those kernels, it's pretty typical for 2D to be slower than 1D, and 3D to be slower than 2D.
The reason why this is the case is due to the additional overhead of software predication. This is discussed in the original FNA paper, and our recent GNA paper as well.

This overhead is largely eliminated in our new kernels for Hopper and Blackwell (datacenter-class cards only).
We do plan to ship kernels for earlier architectures and RTX cards with the new design as well, which will hopefully minimize the difference between the different dimensionalities.

[Yunge6666]

In short, yes; 2D being slower than 1D is normal (but 1D being slower than 2D is unexpected).

For starters, with every additional dimension, there is more room for "wasted FLOPs", which means unless you're using something like stride to avoid that, your 2D problem will just always have a worse throughput because it's wasting FLOPs.

However, the bulk of what users typically feel is usually a performance issue in our original FNA kernels.

Your GPU is Ada (SM89), which means you can only our CUTLASS FNA backend (and Flex backend, but with limited support).

With those kernels, it's pretty typical for 2D to be slower than 1D, and 3D to be slower than 2D. The reason why this is the case is due to the additional overhead of software predication. This is discussed in the original FNA paper, and our recent GNA paper as well.

This overhead is largely eliminated in our new kernels for Hopper and Blackwell (datacenter-class cards only). We do plan to ship kernels for earlier architectures and RTX cards with the new design as well, which will hopefully minimize the difference between the different dimensionalities.

Ok. This is what I expected, because currently I want to map 2D data to a 1D sequence and then perform attention through NATTEN1D. But I found that the gap of latency/FPS is not noticeable at low resolution, and there will be a significant gap at high resolution. Do you have pretrained weights in high resolution?

[alihassanijr]

I want to map 2D data to a 1D sequence and then perform attention through NATTEN1D

If your data naturally assumes a multi-dimensional layout, you might not want to do that.
The existence of spatial locality is the only reason why methods like neighborhood attention make sense.

At the same time, consider that a future upgrade could make 2D as fast as (possibly even faster than) 1D is for you today.

But I found that the gap of latency/FPS is not noticeable at low resolution, and there will be a significant gap at high resolution.

Yes; smaller and larger problem sizes can behave differently for any operation, not just attention or neighborhood attention.

Do you have pretrained weights in high resolution?

Pre-trained weights for what model exactly?

[Yunge6666]

I want to map 2D data to a 1D sequence and then perform attention through NATTEN1D

If your data naturally assumes a multi-dimensional layout, you might not want to do that. The existence of spatial locality is the only reason why methods like neighborhood attention make sense.

At the same time, consider that a future upgrade could make 2D as fast as (possibly even faster than) 1D is for you today.

But I found that the gap of latency/FPS is not noticeable at low resolution, and there will be a significant gap at high resolution.

Yes; smaller and larger problem sizes can behave differently for any operation, not just attention or neighborhood attention.

Do you have pretrained weights in high resolution?

Pre-trained weights for what model exactly?

Thanks for your suggestion, it helps me a lot!
Do you have any high-resolution pre-trained models trained on Imagenet-1k?

[alihassanijr]

We pre-trained the DiNAT-Large and DiNAT_s-Large variants on ImageNet-21K and then ImageNet-1K at 224x224, then fine-tuned at 384x384, but that's the highest resolution we trained on ImageNet classification. Detection and segmentation were technically trained at higher res. All of these were done almost 3 years ago.

https://github.com/SHI-Labs/Neighborhood-Attention-Transformer/blob/main/classification/DiNAT.md#dinat