[3outeille]

What does this PR do?

Stack: #50493 → #50496 → #50511 → #50514 → #50515 → this PR → #50499 → #50521 (Trainer) → #50522 (docs)

After the all-to-all, attention sees the whole sequence, so the mask of the whole sequence can be rebuilt there. This PR rebuilds it with the masking_utils mask functions, so it matches what the model builds on an unsharded run. That lifts three refusals:

• Padded batches: a 2D attention_mask is kept for the full sequence. sdpa gets the full mask. Flash kernels get the 2D mask and unpad, as on an unsharded run. 4D masks are still refused.
• Sliding windows with sdpa: each layer's window comes from the sliding_window the model passes to the attention call, else module.sliding_window, else the config for layers that layer_types marks as sliding.
• Chunked attention (llama4) with sdpa. It is still refused with flash kernels.

With sdpa, the whole-sequence mask takes batch × seq × seq booleans, so a flash kernel is the better choice for long sequences (the docs PR says so).

Tests

tests/context_parallel, CPU (gloo). Loss, logits slice and every gradient match an unsharded run for: right and left padding, a sliding window with sdpa (unpacked and packed), a gemma3-style layer_types mix, chunked attention, and padding combined with the automatic padding to a multiple of cp_size.

CP test from #50499 on every causal LM, with the tester configs unchanged:

result	classes
pass (plain and packed)	60, now including the sliding-window models (gemma3, mistral, cohere2, ministral, olmo3, exaone4, mellum, gemma4_unified, ...)
fail	afmoe packed and hrm_text (both fail on an unsharded or FSDP-only run too, see #50499), gemma4_text (float32 rounding of the embedding gradient; passes at atol=1e-4)

On 2× H100 (Qwen3-0.6B, bf16), CP gives the same loss and logits as an unsharded run for batches with left and right padding (sdpa and the FA3 hub kernel), and for alternating sliding-window layers with sdpa, unpacked and packed.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[3outeille]

Not needed any more: with the plan-based rewrite the model builds its own masks over the whole sequence, so padding masks, sliding windows and chunked attention work without context parallel code (now tested in #50496).