splash_attention for the block-diffusion mask#1

Beomi wants to merge 3 commits into

mainBeomi/dllm-jax:mainfrom

splash_attnBeomi/dllm-jax:splash_attnCopy head branch name to clipboard