Multi-GPU CPT and other padding-free or packed training runs could fail when attention moved to a model shard on another GPU:

Attention bias and Query/Key/Value should be on the same device
query.device: cuda:1
attn_bias: cuda:0

xFormers creates BlockDiagonalCausalMask metadata on CUDA's current device when no device is supplied. Under model parallelism, that can remain cuda:0 while later layers run elsewhere. The process-wide cache also omitted the target device from its key, so packed and inherited block masks could reach the kernel on the wrong device.

Fix

• Scope reusable xFormers masks by device and move generated or inherited masks to the query device before dispatch.
• Skip reconstruction when the mask is already on the correct device.
• Preserve compatibility with older xFormers versions and clear the global mask cache with the other packed-attention caches.
• Add isolated regressions plus a real two-GPU xFormers kernel test.

Testing

• Attention-mask suite: 16 passed, including the real CUDA 0 to CUDA 1 regression.
• Packing and attention-dispatch suites: 57 passed.
• TRL padding-free compatibility suite: 122 passed, 26 skipped.
• Raw-text suite: 9 passed.
• Ruff, Python compilation, and diff checks passed.

Runtime validation

On a two-GPU CUDA host, packed and inherited masks created on cuda:0 were correctly moved for query, key, and value tensors on cuda:1. xFormers validation passed and memory-efficient attention completed on cuda:1 with finite output.