Skip to content

[Pytorch][Attention][CP][CI] Make 3221 changes robust #3438

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sudhakarsingh27 wants to merge 7 commits into
base: main
Choose a base branch
from
+29 −12
Open

[Pytorch][Attention][CP][CI] Make 3221 changes robust #3438

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
d9d2240
Support no-load-balance THD with FlashAttention
sudhakarsingh27 Aug 28, 2026
9c8ca62
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Aug 28, 2026
fb02936
Clarify no-load-balance backend constraints
sudhakarsingh27 Aug 28, 2026
5e14311
Add no-load-balance FlashAttention regression coverage
sudhakarsingh27 Aug 28, 2026
15af9b8
Use backend selection for FlashAttention CP test
sudhakarsingh27 Aug 28, 2026
11c205c
Clarify forced FA2 regression coverage
sudhakarsingh27 Aug 28, 2026
905de9d
Rely on existing L1 attention runs
sudhakarsingh27 Aug 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 16 additions & 6 deletions tests/pytorch/attention/test_attention_with_cp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -739,12 +739,22 @@ def test_cp_with_fused_attention_no_load_balance(cp_pool):
)


@pytest.mark.skipif(
get_device_compute_capability() != (9, 0) or not FlashAttentionUtils.v3_is_installed,
reason="FlashAttention 3 requires sm90 and an installed FA3 package.",
)
def test_cp_with_flash_attention_3_no_load_balance(cp_pool):
"""Check the supported unpadded FlashAttention 3 path."""
def test_cp_with_flash_attention_no_load_balance(cp_pool):
"""Check the supported unpadded FlashAttention path."""
config = copy.deepcopy(model_configs_flash_attn["cp_2_0"])
config.context_parallel = True
config.cp_comm_type = "all_gather"
config.attn_mask_type = "padding_causal"
available_backends, _, _ = get_available_attention_backends(
config,
qkv_dtype=torch.bfloat16,
qkv_layout="thd_thd_thd",
pad_between_seqs=False,
is_training=True,
deterministic=_deterministic,
)
if not available_backends[0]:
pytest.skip("FlashAttention is unavailable.")
_submit(
cp_pool(2),
dtype="bf16",
Expand Down
19 changes: 13 additions & 6 deletions transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4184,6 +4184,13 @@ def backward(ctx, dout, *_args):
fa_backward_kwargs["deterministic"] = ctx.deterministic
if fa_utils.v2_6_0_plus:
fa_backward_kwargs["softcap"] = 0.0
if (
ctx.qkv_format == "thd"
and ctx.load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE
):
# FA2 GQA backward accumulates into expanded K/V gradient buffers.
# Initialize them before accumulation in this partitioning path.
fa_backward_kwargs["zero_tensors"] = True

local_seq_chunk_ids = (
[rank]
Expand Down Expand Up @@ -4563,6 +4570,7 @@ def backward(ctx, dout, *_args):
None,
None,
None,
None,
)


Expand Down Expand Up @@ -5452,7 +5460,7 @@ def attn_forward_func_with_cp(
assigns one contiguous physical-buffer chunk to each rank and uses one attention
step per rank. Logical sequences remain isolated by ``cu_seqlens``. This strategy
requires THD, all-gather, causal self-attention without a sliding window, and
FusedAttention, or FlashAttention 3 with ``pad_between_seqs=False``. Input producers
FusedAttention, or FlashAttention with ``pad_between_seqs=False``. Input producers
must use the same strategy when partitioning inputs with
:func:`get_batch_on_this_cp_rank` or :func:`get_thd_partitioned_indices`.

Expand Down Expand Up @@ -5540,12 +5548,11 @@ def attn_forward_func_with_cp(
assert (
cp_comm_type == "all_gather"
), "No-load-balance THD partitioning requires cp_comm_type='all_gather'."
# Backend selection is handled by the caller. FlashAttention does not yet
# support inter-sequence padding with this partitioning strategy.
assert (
use_fused_attention or use_flash_attn_3
), "No-load-balance THD partitioning requires FusedAttention or FlashAttention 3."
assert not (
use_flash_attn_3 and pad_between_seqs
), "No-load-balance THD partitioning with FlashAttention 3 does not support padding yet."
not pad_between_seqs or use_fused_attention
), "No-load-balance THD partitioning only supports padding with FusedAttention."
Comment on lines 5553 to +5555

@greptile-apps greptile-apps Bot Aug 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 FA2 lacks regression coverage

This relaxed assertion enables FlashAttention 2 for unpadded no-load-balance THD all-gather context parallelism, but the existing test covers only FlashAttention 3. Add FA2 forward and backward coverage so regressions in output, dQ/dK/dV initialization, backend argument compatibility, and autograd behavior are detected.

Knowledge Base Used: PyTorch attention execution

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

@sudhakarsingh27 sudhakarsingh27 Aug 28, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed. The no-load-balance CP test now queries production backend selection for the BF16 THD all-gather GQA configuration and validates output, dQ, dK, and dV. In L1, FA4 is disabled and FA3 is not installed, so the existing deterministic and non-deterministic suite runs exercise FA2 without a separate forced invocation. The focused path was also validated with FA2 on SM90 and SM100 and with FA3/FA4 on supported hardware.

@cyanguwa cyanguwa Aug 28, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this change because FA3 and pad_between_seqs are incompatible? I think it's better to keep them separate right?

@sudhakarsingh27 sudhakarsingh27 Aug 28, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not a general incompatibility between FA3 and pad_between_seqs. The limitation is specific to the no-load-balance partitioning strategy. Backend selection remains separate: any supported FlashAttention version can run the unpadded path, while inter-sequence padding in this path remains FusedAttention-only. The assertion and comments now reflect that separation.

assert "causal" in attn_mask_type and window_size == (
-1,
0,
Expand Down
Loading