From d9d22408aee8d7fa30c288313e0056cf41a8ad9e Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 28 Aug 2026 12:13:38 -0700 Subject: [PATCH 1/7] Support no-load-balance THD with FlashAttention FlashAttention 2 backward needs zeroed gradient outputs for deterministic grouped-query attention in this path. Permit unpadded FlashAttention backends because inter-sequence padding requires no backend-specific handling, and align the autograd return arity with the forward inputs. Signed-off-by: Sudhakar Singh --- .../dot_product_attention/context_parallel.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 78e599199f..f65375f418 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -4184,6 +4184,11 @@ 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 + ): + fa_backward_kwargs["zero_tensors"] = True local_seq_chunk_ids = ( [rank] @@ -4563,6 +4568,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, ) @@ -5540,12 +5546,9 @@ def attn_forward_func_with_cp( assert ( cp_comm_type == "all_gather" ), "No-load-balance THD partitioning requires cp_comm_type='all_gather'." - 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." + assert use_fused_attention or not pad_between_seqs, ( + "No-load-balance THD partitioning only supports padding with FusedAttention." + ) assert "causal" in attn_mask_type and window_size == ( -1, 0, From 9c8ca620a6eb2929f64360b498970bae5aa98231 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 19:16:10 +0000 Subject: [PATCH 2/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../attention/dot_product_attention/context_parallel.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index f65375f418..07faa4ee39 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -5546,9 +5546,9 @@ def attn_forward_func_with_cp( assert ( cp_comm_type == "all_gather" ), "No-load-balance THD partitioning requires cp_comm_type='all_gather'." - assert use_fused_attention or not pad_between_seqs, ( - "No-load-balance THD partitioning only supports padding with FusedAttention." - ) + assert ( + use_fused_attention or not pad_between_seqs + ), "No-load-balance THD partitioning only supports padding with FusedAttention." assert "causal" in attn_mask_type and window_size == ( -1, 0, From fb029367c63fbbfab93bb6245ac826929006e7aa Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 28 Aug 2026 13:37:13 -0700 Subject: [PATCH 3/7] Clarify no-load-balance backend constraints Document why FA2 initializes backward gradient buffers and why inter-sequence padding remains limited to FusedAttention. This makes the supported backend matrix explicit without changing behavior. Signed-off-by: Sudhakar Singh --- .../attention/dot_product_attention/context_parallel.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 07faa4ee39..e9948cb9d8 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -4188,6 +4188,8 @@ def backward(ctx, dout, *_args): 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 = ( @@ -5458,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`. @@ -5546,8 +5548,10 @@ 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 not pad_between_seqs + not pad_between_seqs or use_fused_attention ), "No-load-balance THD partitioning only supports padding with FusedAttention." assert "causal" in attn_mask_type and window_size == ( -1, From 5e14311b1b478c90e8f59712b48c644fb6c4bd7c Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 28 Aug 2026 14:12:45 -0700 Subject: [PATCH 4/7] Add no-load-balance FlashAttention regression coverage Reuse the focused CP test across supported FlashAttention generations and force FA2 from L1 in a fresh process. Backend controls are read at import time, so a separate invocation covers deterministic FA2 GQA backward without adding runner plumbing. Signed-off-by: Sudhakar Singh --- qa/L1_pytorch_distributed_unittest/test.sh | 6 ++++ .../attention/test_attention_with_cp.py | 28 ++++++++++++++----- 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index e0c92849f8..63c5ad9a01 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -43,6 +43,12 @@ else NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py" fi +# Force FA2 in a fresh process to cover deterministic no-load-balance GQA backward. +NVTE_FLASH_ATTN_V3=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s \ + --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp_fa2_no_load_balance.xml \ + $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py::test_cp_with_flash_attention_no_load_balance \ + || test_fail "FlashAttention 2 no-load-balance deterministic GQA" + python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/distributed/test_sanity.py || test_fail "test_sanity.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/distributed/test_numerics.py || test_fail "test_numerics.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics_exact.xml $TE_PATH/tests/pytorch/distributed/test_numerics_exact.py || test_fail "test_numerics_exact.py" diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 32a54a7884..07633f6200 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -46,9 +46,10 @@ test_essential = bool(int(os.getenv("NVTE_TEST_ESSENTIAL", "1"))) # An installed FA4 package must not select tests when the backend is explicitly disabled. -fa4_enabled = bool(int(os.getenv("NVTE_FLASH_ATTN", "1"))) and bool( - int(os.getenv("NVTE_FLASH_ATTN_V4", "1")) -) +flash_attn_enabled = bool(int(os.getenv("NVTE_FLASH_ATTN", "1"))) +fa2_enabled = flash_attn_enabled and bool(int(os.getenv("NVTE_FLASH_ATTN_V2", "1"))) +fa3_enabled = flash_attn_enabled and bool(int(os.getenv("NVTE_FLASH_ATTN_V3", "1"))) +fa4_enabled = flash_attn_enabled and bool(int(os.getenv("NVTE_FLASH_ATTN_V4", "1"))) model_configs_flash_attn = { # test: ModelConfig(b, sq, hq, dqk) @@ -740,11 +741,24 @@ 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.", + not ( + ( + get_device_compute_capability() == (9, 0) + and ( + (fa2_enabled and FlashAttentionUtils.v2_plus) + or (fa3_enabled and FlashAttentionUtils.v3_is_installed) + ) + ) + or ( + get_device_compute_capability() >= (10, 0) + and fa4_enabled + and FlashAttentionUtils.v4_is_installed + ) + ), + reason="FlashAttention 2 or 3 on sm90, or FlashAttention 4 on sm100+, is required.", ) -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.""" _submit( cp_pool(2), dtype="bf16", From 15af9b81b8d99df54bf286b2a4dd9b819b590c70 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 28 Aug 2026 15:22:51 -0700 Subject: [PATCH 5/7] Use backend selection for FlashAttention CP test Query the existing backend selector for the exact THD context-parallel configuration instead of duplicating architecture and version gates in pytest. This keeps capability decisions in one place while allowing version-forced test invocations to exercise production selection logic. Signed-off-by: Sudhakar Singh --- .../attention/test_attention_with_cp.py | 38 +++++++++---------- 1 file changed, 17 insertions(+), 21 deletions(-) diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 07633f6200..3b44c31a2f 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -46,10 +46,9 @@ test_essential = bool(int(os.getenv("NVTE_TEST_ESSENTIAL", "1"))) # An installed FA4 package must not select tests when the backend is explicitly disabled. -flash_attn_enabled = bool(int(os.getenv("NVTE_FLASH_ATTN", "1"))) -fa2_enabled = flash_attn_enabled and bool(int(os.getenv("NVTE_FLASH_ATTN_V2", "1"))) -fa3_enabled = flash_attn_enabled and bool(int(os.getenv("NVTE_FLASH_ATTN_V3", "1"))) -fa4_enabled = flash_attn_enabled and bool(int(os.getenv("NVTE_FLASH_ATTN_V4", "1"))) +fa4_enabled = bool(int(os.getenv("NVTE_FLASH_ATTN", "1"))) and bool( + int(os.getenv("NVTE_FLASH_ATTN_V4", "1")) +) model_configs_flash_attn = { # test: ModelConfig(b, sq, hq, dqk) @@ -740,25 +739,22 @@ def test_cp_with_fused_attention_no_load_balance(cp_pool): ) -@pytest.mark.skipif( - not ( - ( - get_device_compute_capability() == (9, 0) - and ( - (fa2_enabled and FlashAttentionUtils.v2_plus) - or (fa3_enabled and FlashAttentionUtils.v3_is_installed) - ) - ) - or ( - get_device_compute_capability() >= (10, 0) - and fa4_enabled - and FlashAttentionUtils.v4_is_installed - ) - ), - reason="FlashAttention 2 or 3 on sm90, or FlashAttention 4 on sm100+, is required.", -) 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", From 11c205c7f87787c0e69e8ff9f65e326810e3297f Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 28 Aug 2026 15:33:09 -0700 Subject: [PATCH 6/7] Clarify forced FA2 regression coverage Document why the focused deterministic case runs in a fresh process and how the version flags force FA2, while the full runs retain normal backend selection. Signed-off-by: Sudhakar Singh --- qa/L1_pytorch_distributed_unittest/test.sh | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index 63c5ad9a01..2ec058bf06 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -43,7 +43,9 @@ else NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py" fi -# Force FA2 in a fresh process to cover deterministic no-load-balance GQA backward. +# The full runs above exercise both modes with normal backend selection. +# Determinism is captured when the test module is imported, so use a fresh process and +# disable FA3 (FA4 is disabled above) to force FA2 for the deterministic GQA regression. NVTE_FLASH_ATTN_V3=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s \ --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp_fa2_no_load_balance.xml \ $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py::test_cp_with_flash_attention_no_load_balance \ From 905de9d7a46053f7c9c5f7bc01e42998215e850b Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 28 Aug 2026 15:50:29 -0700 Subject: [PATCH 7/7] Rely on existing L1 attention runs The general context-parallel test is already collected by both L1 attention invocations. Since this suite selects FA2, a separate forced deterministic invocation duplicates coverage without exercising a distinct path. Signed-off-by: Sudhakar Singh --- qa/L1_pytorch_distributed_unittest/test.sh | 8 -------- 1 file changed, 8 deletions(-) diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index 2ec058bf06..e0c92849f8 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -43,14 +43,6 @@ else NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py" fi -# The full runs above exercise both modes with normal backend selection. -# Determinism is captured when the test module is imported, so use a fresh process and -# disable FA3 (FA4 is disabled above) to force FA2 for the deterministic GQA regression. -NVTE_FLASH_ATTN_V3=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s \ - --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp_fa2_no_load_balance.xml \ - $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py::test_cp_with_flash_attention_no_load_balance \ - || test_fail "FlashAttention 2 no-load-balance deterministic GQA" - python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/distributed/test_sanity.py || test_fail "test_sanity.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/distributed/test_numerics.py || test_fail "test_numerics.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics_exact.xml $TE_PATH/tests/pytorch/distributed/test_numerics_exact.py || test_fail "test_numerics_exact.py"