[chillee]

Hi, one of the authors of this blog post (Horace He), along with Driss Guessous, Yanbo Liang, and Joy Dong.

We’re quite happy with this abstraction - happy to answer any questions about it!

[zaptrem]

For those of us using the 2D NATTEN kernel from their library along with torch.compile, is this faster? Especially given all their tricks (e.g., the non-deterministic KV-parallelism)

[chillee]

In my (very amateurish) testing, I think the performance seemed pretty comparable (for non-dilated natten). I need to do some proper benchmarking though!

[imjonse]

Is this for Ampere and newer only as FA2?

[chillee]

I believe it should run on V100 as well (although definitely not tested as well), and an user reported that they got it running on T4 too.