It seems like position_ids must be combined with padding_free, otherwise it will silently fall back to flash_attn_varlen_func. I'll need to update the collator to flatten the batch. Ill do this tomorrow

EDIT: not exactly true, it will try to build the position_ids from the completion mask, and this is not what we want here.

ah, technically it should be possible to return position_ids and data collator will then only flatten batch of prepared position ids! good idea

Awesome and quick work! Just tested it with this branch + transformers main, seems to be correct