[PyTorch] Extend no-load-balance CP to the a2a comm type

[Rudin6]

Flat (no-load-balance) sharding gives CP rank r the contiguous global token range [r * s_local, (r + 1) * s_local). For cp_comm_type='a2a' the exchange already restores the full sequence on every rank before attention, so the dual-chunk reordering is a no-op that costs a full extra copy of q, k and v.

a2a never divides cu_seqlens or max_seqlen: after the exchange each rank holds the whole sequence and only its own head slice, so the sequence metadata, the mask and the attention math are identical to cp_size = 1. Every attn_mask_type and window_size therefore stays valid, and THD needs no new per-step metadata.

Skipping the reorder also removes its view(cp_size * 2, ...), which is the only reason the local sequence length had to be even. The per-sequence divisibility requirement drops from 2 * cp_size to cp_size; for THD no individual sequence needs padding at all, only the packed total.

Measured on 8xH200 for one before_attn exchange of q+k+v (bshd, cp=8, h=40, d=128), max over ranks: -12.0% at s=8192, -29.5% at s=32768 and -31.7% at s=75600. The gain grows with payload because the removed index_select is bandwidth bound.

Scope matches the existing all_gather strategy: FP8 and CUDA graph capture stay rejected, and a2a+p2p is not covered. The p2p path is unchanged.

Tests add flat a2a coverage over bshd, sbhd and THD, including a config whose sequence length is divisible by cp_size but not by 2 * cp_size, and a causal config. All 9 pass; dual-chunk a2a on the same configs and the existing all_gather no-load-balance tests still pass.

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Confidence Score: 5/5

The production changes appear safe to merge, with one non-blocking test-harness limitation still outstanding from the previous review.

Findings

1.  Flat bias tests crash ▶

Summary

This PR extends no-load-balance context parallelism to the A2A communication path.

• Skips dual-chunk sequence reordering when tokens are already flat-sharded in global sequence order.
• Relaxes the local even-sequence-length restriction for flat A2A operation.
• Preserves global THD sequence metadata after the exchange reconstructs the full packed sequence.
• Adds fused-attention coverage for bshd, sbhd, and thd, including causal attention and sequence lengths not divisible by twice the CP size.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[Flat local token shard] --> B[A2A exchange]
    B --> C[Full global sequence<br/>local head shard]
    C --> D[Attention]
    D --> E[Split sequence by CP rank]
    E --> F[A2A exchange]
    F --> G[Local token shard<br/>full heads]

Loading

_{Reviews (2) · Last reviewed commit: "Merge branch 'main' into cp-a2a-no-load-..."}

[greptile-apps]

Flat bias tests crash

This flat bshd/sbhd branch does not initialize seq_idx, but the later attention-bias setup still reads it and assumes a 2 * world_size dual-chunk layout. As a result, any flat-a2a test configured with attention bias will fail with an unbound-variable error instead of validating the feature. The new configurations all use no_bias, so they do not cover this path.