

And for Pytorch DeepSpeed has built one as well: DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale, Mixture of Experts - blog posts:  1, 2 and specific deployment with large transformer-based natural language generation models: blog post, Megatron-Deepspeed branch.
Using PyTorch native attention and Flash Attention
PyTorch's torch.nn.functional.scaled_dot_product_attention (SDPA) can also call FlashAttention and memory-efficient attention kernels under the hood. SDPA support is currently being added natively in Transformers and is used by default for torch>=2.1.1 when an implementation is available. Please refer to PyTorch scaled dot product attention for a list of supported models and more details.
Check out this blogpost to learn more about acceleration and memory-savings with SDPA.