Need to check if the ESM2 layer handles the chain breaks properly.

[murrellb]

Can't recall if this is done right or not, and I'm not sure how the flash attention kernel will handle this. Not sure what the official version does either though.