temporally causal performance worse than full attention for natten3d

[coconutruben]

Hi there

We were benchmarking natten3d for our needs, and noticed that turning on temporally causal is worse than running full attention (vs sdpa on cudnn).

Is this something you have noticed too? for inputs like this

• [1, t, h, w]
• num heads: 24
• head dim: 128
• for height/width we're trying 32 and 64
• for t (temporal) we're trying 30 and 60

Prefix	GPUs	Median (ms)	time	height	width
attention_	1	8.582	30	32	32
natten_ima	1	17.338	30	32	32
attention_	1	33.817	60	32	32
natten_ima	1	61.646	60	32	32
attention_	1	134.908	30	64	64
natten_ima	1	261.328	30	64	64
attention_	1	548.231	60	64	64
natten_ima	1	965.665	60	64	64

other information

• blackwell gb200
• bf16 dtype
• natten 0.21.1
• GPU Driver 580.126.09
• PyTorch 2.9.0+cu130
• CUDA13.0
• cuDNN9.13.00
• NCCL2.27.7
• Triton3.5.0
• Python3.12.3

[AdityaKane2001]

Could you please share the following information to help debug this issue:

1. Kernel size, stride, dilation settings across all dimensions.
2. Are both attention_ and natten_ima from NATTEN? If not, what implementation are you using for self-attention?
3. Are these numbers for the forward or backward pass or both? How are you measuring these times?

[coconutruben]

Thanks for the quick feedback!

1. kernels are "full" size so (30, 32, 32)
2. stride (1,1,1) dilation (1,1,1)
3. attention_ is from pytorch default sdpa with cudnn context manager (we flatten T, H, W to just a single sequence dimension)
4. forward pass
5. We capture a cudagraph of 20 forward passes, and replay it 10 or so times to get median runtime

[alihassanijr]

I think this is just due to the fine-grained masking overhead.
We only disable fine-grained masking for fully block-sparse cases, and not per-dimension, because that will most likely have to be done statically, which means stamping out more kernels than we already do.

Causal masking is never block-sparse, but you're only causal along one dimension (T), so you'd need the mask for just that one dimension, but the other two dimensions can be safely ignored.

I'm not sure if we can implement this right away, but something that might help is using 2-D NA instead of 3-D, by just flattening H and W.
The overhead for 2-D masking is definitely better than 3-D, so that should at least partially improve things.

But if your baseline is cuDNN, also note that the latest version of cuDNN kernels has some additional optimization that haven't been ported back to the CUTLASS C++ kernels we use (we plan to port some of the FA4 features at some point but it'll take time whether we also make the switch to CuTeDSL or port the features to C++ successfully.)

I'll run your use case in the profiler too just to make sure.

[coconutruben]

Thanks for the feedback!! I'll definitely try the 2D example next and check whether this gets us closer.

In the analysis I also noticed that the memory allocation is much larger with this than with cudnn.

Ultimately we care about having a bounded temporal window, of ~some frames, maybe with some dilation sprinkled in there. For that use case, are there some dilations/kernel sizes that or math we can do to hit certain sparsity targets? I ask because we're seeing similar ~260TFLOP performance for a simple (6,h,w) kernel with and without temporal causal mask.

[alihassanijr]

Actually I was wrong. Here's my profiler numbers with the 60x64x64 use case:

cuDNN (via SDPA):

python3 -m natten.profiler -b 1 -n 24 -d 128 -i 60 64 64 -c True False False --backend cudnn

| 2026-02-04 13:01:47,951 | [[ natten.profiling_utils.profiling ]] [ WARNING ]: Running standard self attention; ignoring NA parameters window size, stride, dilation, causal, etc.
                                                Profiler results
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━┓
┃ Framework ┃ Kernel category ┃ Arch  ┃                    Operation                     ┃ # calls ┃  Runtime  ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━┩
│   cuDNN   │    attention    │ Sm100 │ sm100_flash_fprop_f16_knob_7_128x128x128_4x1x... │    1    │ 554.279ms │
│           │                 │       │                      Total                       │         │ 554.279ms │
└───────────┴─────────────────┴───────┴──────────────────────────────────────────────────┴─────────┴───────────┘

Blackwell FNA:

python3 -m natten.profiler -b 1 -n 24 -d 128 -i 60 64 64 -c True False False --backend blackwell-fna

                              Profiler results
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━┓
┃ Framework ┃ Kernel category ┃ Arch  ┃  Operation   ┃ # calls ┃  Runtime  ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━┩
│  CUTLASS  │    attention    │ Sm100 │  FnaForward  │    1    │ 155.381ms │
│  CUTLASS  │  token permute  │   -   │ TokenPermute │    5    │  5.729ms  │
│           │                 │       │    Total     │         │ 161.110ms │
└───────────┴─────────────────┴───────┴──────────────┴─────────┴───────────┘

But even if something's going wrong and NATTEN is picking cutlass-fna for you, then it'll be about 50% worse than what you reported:

python3 -m natten.profiler -b 1 -n 24 -d 128 -i 60 64 64 -c True False False --backend cutlass-fna

                           Profiler results
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━┓
┃ Framework ┃ Kernel category ┃ Arch ┃ Operation  ┃ # calls ┃ Runtime ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━┩
│  CUTLASS  │    attention    │ Sm80 │ FnaForward │    1    │ 1.586s  │
│           │                 │      │   Total    │         │ 1.586s  │
└───────────┴─────────────────┴──────┴────────────┴─────────┴─────────┘

These are all on TOT natten (we haven't done any major perf optimizations since 0.21.1), on a GB200, driver version 580.105.08. NATTEN was compiled with CTK 13.0.

[alihassanijr]

So just to confirm for my own sanity, the numbers you report are just op-level, or end-to-end as part of your model?
If the latter, that could explain the discrepancy.

[coconutruben]

So just to confirm for my own sanity, the numbers you report are just op-level, or end-to-end as part of your model?
If the latter, that could explain the discrepancy.

This should be mostly op level, maybe some minor view op before the actual op but nothing drastic.

Thanks for sharing the profiling numbers, let me see if I can reproduce the profiling and then inspect my code accordingly!

[alihassanijr]

Ultimately we care about having a bounded temporal window, of ~some frames, maybe with some dilation sprinkled in there. For that use case, are there some dilations/kernel sizes that or math we can do to hit certain sparsity targets? I ask because we're seeing similar ~260TFLOP performance for a simple (6,h,w) kernel with and without temporal causal mask.

Certainly! I think just reducing partial tiles alone can take you a long way, even with the fine-grained mask present.
I'll share some information on how to optimize for that without having to run the kernel.
We were supposed to open-source NATTENSim, which basically tells you your upper bound speedup, assuming no implementation overhead, but we're still behind on integrating it into NATTEN.
I can probably just push the reference we have right now somewhere with some documentation (it's just a few pure python scripts).

But I think let's figure out the speed issue.

[alihassanijr]

You can also set this environment variable, and it should log what backend it picks: NATTEN_LOG_LEVEL="DEBUG"

If it's not Blackwell FNA, then that's the problem, and if you could share a minimal repro, it'll help us figure out why it's picking an inferior kernel for you and fix it.

[coconutruben]

I can reproduce your profiler results, and this looks promising overall (pending memory footprint), though I can't reproduce them within my setup.
With the debug, I'm seeing

| 2026-02-04 21:42:30,089 | [[ natten.backends ]] [ DEBUG ]: Backend not set; picked Blackwell FNA kernel.                                                                                                                                                                               

I suspect it's something in how I'm calling the kernel, but the cluster I'm on right now has issues capturing a proper profile for me to check if I'm introducing some odd overhead. I'll update again when I have a clear profile of what's going on, but no further help needed right now, and thank you so much for your quick and insightful comments.

[coconutruben]

Okay, I was able to collect a trace for

(.venv) ruben@worker-103:~/dreisam$ python3 -m natten.profiler -b 1 -n 24 -d 128 -i 180 32 32 -c True False False --backend blackwell-fna -w 6 32 32                                                                                                                                     
|===========================================================================================|
|                                      Profiler results                                     |
|===========================================================================================|
+===========+=================+=======+===============================+=========+===========+
| Framework | Kernel category | Arch  | Operation                     | # calls | Runtime   |
+===========+=================+=======+===============================+=========+===========+
| CUTLASS   | attention       | Sm100 | FnaForward                    | 1       | 16.800ms  |
| PyTorch   | elementwise     | -     | vectorized_elementwise_kernel | 1       | 293.020us |
| PyTorch   | elementwise     | -     | elementwise_kernel            | 5       | 7.799ms   |
+===========+=================+=======+===============================+=========+===========+
|           |                 |       | Total                         |         | 24.892ms  |
+===========+=================+=======+===============================+=========+===========+

but within my own benchmarking setup, and I see that the kernel is actually just 46ms (the main kernel) whereas the profiler claims

void cutlass::device_kernel_sm100<cutlass::fna::kernel::Sm100FnaFwdKernelTmaWarpspecialized<cute::tuple<int, int, int, cute::tuple<cute::tuple<int, int>, int> >, cutlass::fna::collective::Sm100FnaFwdMainloopTmaWarpspecialized<cutlass::bfloat16_t, float, float, cute::tuple<cute::C<256>, cute::C<128>, cute::C<128> >, cute::tuple<int, cute::C<1>, cute::tuple<cute::tuple<int, int>, int> >, cute::tuple<int, cute::C<1>, cute::tuple<cute::tuple<cute::C<0>, int>, int> >, cute::tuple<int, cute::C<1>, cute::tuple<cute::tuple<cute::C<0>, int>, int> >, cutlass::fna::collective::NeighborhoodAttentionMask<cute::tuple<cute::C<true>, cute::C<false>, cute::C<false> > >, cute::tuple<cute::C<8>, cute::C<4>, cute::C<8> >, cute::tuple<cute::C<4>, cute::C<4>, cute::C<8> >, cute::tuple<int, int, int>, cute::tuple<cute::C<2>, cute::C<1>, cute::C<1> > >, cutlass::fna::collective::Sm100FnaFwdEpilogueTmaWarpspecialized<cutlass::bfloat16_t, float, cute::tuple<cute::C<128>, cute::C<128>, cute::C<128> >, cute::tuple<int, cute::C<1>, cute::tuple<cute::tuple<int, int>, int> >, cute::tuple<int, cute::tuple<cute::tuple<cute::C<1>, int>, int> > >, cutlass::fna::kernel::PersistentTileScheduler, cutlass::fna::kernel::Sm100FnaCtxKernelWarpspecializedSchedule> >(cutlass::fna::kernel::Sm100FnaFwdKernelTmaWarpspecialized<cute::tuple<int, int, int, cute::tuple<cute::tuple<int, int>, int> >, cutlass::fna::collective::Sm100FnaFwdMainloopTmaWarpspecialized<cutlass::bfloat16_t, float, float, cute::tuple<cute::C<256>, cute::C<128>, cute::C<128> >, cute::tuple<int, cute::C<1>, cute::tuple<cute::tuple<int, int>, int> >, cute::tuple<int, cute::C<1>, cute::tuple<cute::tuple<cute::C<0>, int>, int> >, cute::tuple<int, cute::C<1>, cute::tuple<cute::tuple<cute::C<0>, int>, int> >, cutlass::fna::collective::NeighborhoodAttentionMask<cute::tuple<cute::C<true>, cute::C<false>, cute::C<false> > >, cute::tuple<cute::C<8>, cute::C<4>, cute::C<8> >, cute::tuple<cute::C<4>, cute::C<4>, cute::C<8> >, cute::tuple<int, int, int>, cute::tuple<cute::C<2>, cute::C<1>, cute::C<1> > >, cutlass::fna::collective::Sm100FnaFwdEpilogueTmaWarpspecialized<cutlass::bfloat16_t, float, cute::tuple<cute::C<128>, cute::C<128>, cute::C<128> >, cute::tuple<int, cute::C<1>, cute::tuple<cute::tuple<int, int>, int> >, cute::tuple<int, cute::tuple<cute::tuple<cute::C<1>, int>, int> > >, cutlass::fna::kernel::PersistentTileScheduler, cutlass::fna::kernel::Sm100FnaCtxKernelWarpspecializedSchedule>::Params)

are you able to trace from the cutlass instantiation whether the kernel is actually doing what I think it should be doing (window sizes, etc) or maybe I'm just propagating the parameters wrong through my pipeline?

[alihassanijr]

Sorry can you clarify?
How is your benchmarking setup timing the main kernel?
Or is your benchmarking setup just using ncu, or torch profiler or something?

The kernel you shared and the profiler result look okay to me.

[coconutruben]

Sorry the kernel I pulled out from my benchmarking, the results from the nattan profiler. In my benchmarking that kernel, which I assume has the same inputs/outputs is 45ms and the whole thing including element wise around 55ms. I haven't compared the kernel that the matten profiler is using to this one yet.

Our setup

• some view manipulation
• natten3d call

all of this wrapped in torch.compile and run in a loop of 20 times or so

that compiled loop captured on a cudagraph

that graph replayed 20 or so times

and then some math

so that's the setup, I'll fish around the natten profiler to see if everything looks the same ie shapes kernels inputs etc and whether I can see the full kernel name and they match

[alihassanijr]

Have you tried profiling your setup using the torch profiler and comparing the numbers there?
We're basically using the torch profiler too in natten (both use the driver / runtime API anyway), with some additional symbol mapping and pretty printing stuff.

"FnaForward" in natten profiler corresponds to void cutlass::device_kernel_sm100<cutlass::fna::kernel::... basically.

And just to confirm, your setup is doing just forward pass right?

Anyway, I think getting a full torch compile report of your setup can shine a light on what's going on here.