Are there plans to support sparse attention inference on 4090 and 5090?

[helloyongyang]

I found that flex_flash_attn only supports sm90. Are there plans to support sparse attention inference on 4090 and 5090?

[Strivin0311]

Hi, and thanks for your interest.

First of all, multi-GPU support for both sm_8x (A100, RTX 4090) and sm_10x (B200, RTX 5090) architectures is on our roadmap, and the sm_8x branch is in progress and is expected to land in the next month or two.

Second, specific for sparse attention, we are actively optimizing performance right now, through several features including swap-AB (#158), packed-GQA (#185), sparse-load (#179), etc.

So please stay tuned for the upcoming release and detailed performance report!

[helloyongyang]

We find that swap-AB do not support blocksize (128, 128).

swap-AB = False and blocksize = (128, 128) is faster than swap-AB = True and blocksize = (64, 64)

So, why swap-AB do not support blocksize (128, 128)？
@Strivin0311