Refactor custom_splash_attention for improved clarity#431

copybara-service[bot] merged 1 commit into

mainAI-Hypercomputer/maxdiffusion:mainfrom

custom_attn_fixAI-Hypercomputer/maxdiffusion:custom_attn_fixCopy head branch name to clipboard