From c6072502b7e59316375e5a54e9ed63c876ffdf21 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Wed, 29 Jul 2026 18:27:17 -0700 Subject: [PATCH] Enable THD with hierarchical context parallelism The THD-aware A2A and P2P paths already compose correctly and match the non-CP reference in forward and backward passes. Remove the stale runtime and backend gates, and enable the existing FlashAttention and FusedAttention test coverage for this combination. Signed-off-by: Sudhakar Singh --- tests/pytorch/attention/test_attention_with_cp.py | 11 ----------- .../dot_product_attention/context_parallel.py | 3 --- .../pytorch/attention/dot_product_attention/utils.py | 7 ------- 3 files changed, 21 deletions(-) diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 681ee5e6e0..994f2e198f 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -336,11 +336,6 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type if config.attn_bias_type != "no_bias" and cp_comm_type in ["all_gather", "a2a", "a2a+p2p"]: pytest.skip("No support for bias with cp_comm_type={all_gather, a2a, a2a+p2p}!") - if qkv_format == "thd" and cp_comm_type == "a2a+p2p": - pytest.skip( - "CP implementation with QKVO A2A+P2P (Hierarchical A2A) does not support THD format" - " yet!" - ) if ( qkv_format == "thd" and cp_comm_type == "all_gather" @@ -558,12 +553,6 @@ def test_cp_with_fused_attention( if config.attn_bias_type != "no_bias" and cp_comm_type in ["all_gather", "a2a", "a2a+p2p"]: pytest.skip("No support for bias with cp_comm_type={all_gather, a2a, a2a+p2p}!") - if qkv_format == "thd" and cp_comm_type == "a2a+p2p": - pytest.skip( - "CP implementation with QKVO A2A+P2P (Hierarchical A2A) does not support THD format" - " yet!" - ) - if (config.window_size[0] != -1 or config.window_size[1] not in [-1, 0]) and cp_comm_type in [ "p2p", "a2a+p2p", 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 995dd2e90a..b444f034c8 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -5018,9 +5018,6 @@ def attn_forward_func_with_cp( assert ( isinstance(cp_group, list) and len(cp_group) == 2 ), "CP implementation a2a+p2p requires cp_group = [a2a_cp_group, p2p_cp_group]!" - assert ( - qkv_format != "thd" - ), f"{qkv_format} format is not supported with hierarchical CP implementation yet!" assert ( attn_bias_type == "no_bias" ), f"{attn_bias_type} bias type is not supported with hierarchical CP implementation yet!" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 8a4c6ce286..4bd2e5fd84 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1231,13 +1231,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt cp_comm_type, ) use_fused_attention = False - elif qkv_format == "thd" and cp_comm_type in ["a2a+p2p"]: - logger.debug( - "Disabling FusedAttention as it does not support context parallelism with THD" - " format and cp_comm_type = %s", - cp_comm_type, - ) - use_fused_attention = False elif ( window_size is not None and (window_size[0] != -1 or window_size[1] not in [-1, 0])