Fix block size handling for cross attention in TPU flash attention#301

michelle-yooh wants to merge 1 commit intomainAI-Hypercomputer/maxdiffusion:mainfrom

yooh/correct-default-block-q-dqAI-Hypercomputer/maxdiffusion:yooh/correct-default-block-q-dqCopy head branch name to clipboard