Cleaning up custom_splash_attention