From 95ae9f822a9f60370b1f240645ec3626d791f121 Mon Sep 17 00:00:00 2001 From: Yuzhong Wang Date: Tue, 16 Jun 2026 23:05:09 -0700 Subject: [PATCH] Refactor GDN A2A helper flow Signed-off-by: Yuzhong Wang --- megatron/core/ssm/gated_delta_net.py | 132 ++++++++++++++++----------- 1 file changed, 79 insertions(+), 53 deletions(-) diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 06eb0763e57..f202581105a 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -346,42 +346,9 @@ def forward( qkvzba, _ = self.in_proj(hidden_states) nvtx_range_pop(suffix="in_proj") - # CP All to All: CP to HP - if self.cp_size > 1: - # # Pre-permute head dim so a single unsectioned a2a is equivalent to per-section a2a. - head_perm = _build_head_perm_for_split_sections( - ( - self.qk_dim_local_tp, - self.qk_dim_local_tp, - self.v_dim_local_tp, - self.v_dim_local_tp, - self.num_value_heads // self.tp_size, - self.num_value_heads // self.tp_size, - ), - self.pg_collection.cp.size(), - torch.cuda.current_device(), - ) - qkvzba = qkvzba.index_select(-1, head_perm) - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - qkvzba = tensor_a2a_cp2hp( - qkvzba, - seq_dim=0, - head_dim=-1, - cp_group=self.pg_collection.cp, - undo_attention_load_balancing=False, - ) - if self.cp_size > 1: - # Permute at the seq dim so that a single unsectioned a2a - # is equivalent to per-sequence a2a. - # This also folds the ``_undo_attention_load_balancing`` step. - thd_cp_a2a_idx, thd_cp_a2a_inv = _build_thd_cp_a2a_perm( - cu_seqlens_q, self.cp_size, seq_len - ) - qkvzba = qkvzba.index_select(0, thd_cp_a2a_idx) - else: - qkvzba = tensor_a2a_cp2hp( - qkvzba, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp - ) + qkvzba, thd_cp_a2a_inv = self._a2a_cp_to_hp( + qkvzba, self.cp_size, self.pg_collection.cp, cu_seqlens_q, seq_len, packed_seq_params + ) # Transpose: s b x --> b s x # From sbhd to bshd format @@ -493,23 +460,9 @@ def _gated_norm_and_a2a(core_attn_out: torch.Tensor, gate: torch.Tensor): norm_out_hp = norm_out_hp.reshape(batch, seq_len, -1) norm_out_hp = norm_out_hp.transpose(0, 1).contiguous() - # CP all to all: HP to CP - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - if self.cp_size > 1: - norm_out_hp = norm_out_hp.index_select(0, thd_cp_a2a_inv) - norm_out = tensor_a2a_hp2cp( - norm_out_hp, - seq_dim=0, - head_dim=-1, - cp_group=self.pg_collection.cp, - redo_attention_load_balancing=False, - ) - else: - norm_out = tensor_a2a_hp2cp( - norm_out_hp, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp - ) - - return norm_out + return self._a2a_hp_to_cp( + norm_out_hp, self.cp_size, self.pg_collection.cp, packed_seq_params, thd_cp_a2a_inv + ) if self.recompute_norm_out: self.norm_out_checkpoint = tensor_parallel.CheckpointWithoutOutput() @@ -527,6 +480,79 @@ def _gated_norm_and_a2a(core_attn_out: torch.Tensor, gate: torch.Tensor): return out, out_bias + def _a2a_cp_to_hp( + self, + qkvzba: torch.Tensor, + cp_size: int, + cp_group: torch.distributed.ProcessGroup, + cu_seqlens_q: Optional[torch.Tensor], + seq_len: int, + packed_seq_params: Optional[PackedSeqParams], + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Run GDN context-parallel to hidden-parallel A2A and return its inverse context.""" + if cp_size > 1: + # Pre-permute head dim so a single unsectioned a2a is equivalent to per-section a2a. + head_perm = _build_head_perm_for_split_sections( + ( + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + self.v_dim_local_tp, + self.num_value_heads // self.tp_size, + self.num_value_heads // self.tp_size, + ), + cp_size, + qkvzba.device, + ) + qkvzba = qkvzba.index_select(-1, head_perm) + + thd_cp_a2a_inv = None + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + qkvzba = tensor_a2a_cp2hp( + qkvzba, + seq_dim=0, + head_dim=-1, + cp_group=cp_group, + undo_attention_load_balancing=False, + ) + if cp_size > 1: + # Permute at the seq dim so that a single unsectioned a2a + # is equivalent to per-sequence a2a. + # This also folds the ``_undo_attention_load_balancing`` step. + thd_cp_a2a_idx, thd_cp_a2a_inv = _build_thd_cp_a2a_perm( + cu_seqlens_q, cp_size, seq_len + ) + qkvzba = qkvzba.index_select(0, thd_cp_a2a_idx) + else: + qkvzba = tensor_a2a_cp2hp(qkvzba, seq_dim=0, head_dim=-1, cp_group=cp_group) + + return qkvzba, thd_cp_a2a_inv + + def _a2a_hp_to_cp( + self, + norm_out: torch.Tensor, + cp_size: int, + cp_group: torch.distributed.ProcessGroup, + packed_seq_params: Optional[PackedSeqParams], + thd_cp_a2a_inv: Optional[torch.Tensor], + ) -> torch.Tensor: + """Run GDN hidden-parallel to context-parallel A2A using CP-to-HP context.""" + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + if cp_size > 1: + assert thd_cp_a2a_inv is not None + norm_out = norm_out.index_select(0, thd_cp_a2a_inv) + norm_out = tensor_a2a_hp2cp( + norm_out, + seq_dim=0, + head_dim=-1, + cp_group=cp_group, + redo_attention_load_balancing=False, + ) + else: + norm_out = tensor_a2a_hp2cp(norm_out, seq_dim=0, head_dim=-1, cp_group=cp_group) + + return norm_out + @jit_fuser def _apply_gated_norm(self, x, gate): # Output Norm