From 5e6f23b22f6be3125aef59ebdc65f51c0d76ddf2 Mon Sep 17 00:00:00 2001 From: lit Date: Mon, 2 Feb 2026 02:21:31 -0800 Subject: [PATCH 1/3] Move some processing into a function so can be compiled --- megatron/core/ssm/gated_delta_net.py | 89 ++++++++++++++++++---------- 1 file changed, 57 insertions(+), 32 deletions(-) diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 601a72a4356..9cb50ba6953 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -388,37 +388,12 @@ def forward( ) nvtx_range_pop(suffix="conv1d") - # Split qkv into query_key, and value - query_key, value = torch.split( - qkv, - [2 * self.qk_dim_local_tp // self.cp_size, self.v_dim_local_tp // self.cp_size], - dim=-1, - ) - query_key = query_key.reshape(batch, seq_len, -1, self.key_head_dim) - value = value.reshape(batch, seq_len, -1, self.value_head_dim) - # Apply L2 norm to query and key - if self.use_qk_l2norm: - query_key = l2norm(query_key.contiguous()) - # Split query and key. - query, key = torch.split( - query_key, - [ - self.qk_dim_local_tp // self.key_head_dim // self.cp_size, - self.qk_dim_local_tp // self.key_head_dim // self.cp_size, - ], - dim=2, + # Prepare QKV tensors (split, reshape, L2 norm, repeat_interleave, contiguous) + nvtx_range_push(suffix="prepare_qkv_for_gated_delta_rule") + query, key, value, gate, beta, alpha = self._prepare_qkv_for_gated_delta_rule( + qkv, gate, beta, alpha, batch, seq_len ) - if self.num_value_heads // self.num_key_heads > 1: - query = query.repeat_interleave(self.num_value_heads // self.num_key_heads, dim=2) - key = key.repeat_interleave(self.num_value_heads // self.num_key_heads, dim=2) - - # Make contiguous - query = query.contiguous() - key = key.contiguous() - value = value.contiguous() - gate = gate.contiguous() - beta = beta.contiguous() - alpha = alpha.contiguous() + nvtx_range_pop(suffix="prepare_qkv_for_gated_delta_rule") # Calculate g and beta nvtx_range_push(suffix="g_and_beta") @@ -426,8 +401,7 @@ def forward( dt_bias_local_cp = get_parameter_local_cp( self.dt_bias, dim=0, cp_group=self.pg_collection.cp ) - g = -A_log_local_cp.exp() * F.softplus(alpha.float() + dt_bias_local_cp) # In fp32 - beta = beta.sigmoid() + g, beta = self._compute_g_and_beta(A_log_local_cp, dt_bias_local_cp, alpha, beta) nvtx_range_pop(suffix="g_and_beta") nvtx_range_push(suffix="gated_delta_rule") @@ -477,6 +451,57 @@ def _apply_gated_norm(self, x, gate): y = y.to(x_dtype) return y + @jit_fuser + def _prepare_qkv_for_gated_delta_rule(self, qkv, gate, beta, alpha, batch, seq_len): + """ + Prepare query, key, value, gate, beta, alpha tensors for gated delta rule. + Fuses split, reshape, L2 norm, repeat_interleave, and contiguous operations. + """ + # Split qkv into query_key and value + query_key, value = torch.split( + qkv, + [2 * self.qk_dim_local_tp // self.cp_size, self.v_dim_local_tp // self.cp_size], + dim=-1, + ) + + # Reshape query_key and value + query_key = query_key.reshape(batch, seq_len, -1, self.key_head_dim) + value = value.reshape(batch, seq_len, -1, self.value_head_dim) + + # Apply L2 norm to query and key + if self.use_qk_l2norm: + query_key = l2norm(query_key.contiguous()) + + # Split query and key + split_size = self.qk_dim_local_tp // self.key_head_dim // self.cp_size + query, key = torch.split(query_key, [split_size, split_size], dim=2) + + # Expand query and key if needed (grouped query attention) + if self.num_value_heads // self.num_key_heads > 1: + repeat_factor = self.num_value_heads // self.num_key_heads + query = query.repeat_interleave(repeat_factor, dim=2) + key = key.repeat_interleave(repeat_factor, dim=2) + + # Make all tensors contiguous + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + gate = gate.contiguous() + beta = beta.contiguous() + alpha = alpha.contiguous() + + return query, key, value, gate, beta, alpha + + @jit_fuser + def _compute_g_and_beta(self, A_log_local_cp, dt_bias_local_cp, alpha, beta): + """ + Compute g (decay) and beta (sigmoid) for gated delta rule. + Fuses exp, softplus, mul, neg, and sigmoid operations. + """ + g = -A_log_local_cp.exp() * F.softplus(alpha.float() + dt_bias_local_cp) # In fp32 + beta = beta.sigmoid() + return g, beta + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_group=None): """Provide a sharded state dictionary for distributed checkpointing.""" # Guard for cases metadata is not provided From 7edef03a11d62266adb858c80881d8deb4735e3f Mon Sep 17 00:00:00 2001 From: lit Date: Wed, 4 Mar 2026 02:23:05 -0800 Subject: [PATCH 2/3] try to fix codecov CI failure --- tests/unit_tests/ssm/test_gated_delta_net.py | 59 ++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index 81f8eed0574..f53f169031b 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -138,6 +138,65 @@ def test_gpu_forward(self): output.dtype == hidden_states.dtype ), f"Output dtype {output.dtype=} mismatch with {hidden_states.dtype=}" + def test_jit_compiled_helpers(self): + import torch._dynamo + + gdn = self.gdn + batch = 2 + seq_len = 16 + + num_v_heads_local = gdn.num_value_heads // gdn.tp_size // gdn.cp_size + + qkv_last_dim = (2 * gdn.qk_dim_local_tp + gdn.v_dim_local_tp) // gdn.cp_size + qkv = torch.randn( + batch, seq_len, qkv_last_dim, + device=torch.cuda.current_device(), dtype=torch.bfloat16, + ) + gate = torch.randn( + batch, seq_len, num_v_heads_local, gdn.value_head_dim, + device=torch.cuda.current_device(), dtype=torch.bfloat16, + ) + beta = torch.randn( + batch, seq_len, num_v_heads_local, + device=torch.cuda.current_device(), dtype=torch.bfloat16, + ) + alpha = torch.randn( + batch, seq_len, num_v_heads_local, + device=torch.cuda.current_device(), dtype=torch.bfloat16, + ) + + # Disable dynamo so coverage.py can trace through the method bodies, + # which are normally wrapped by @jit_fuser (torch.compile). + with torch._dynamo.config.patch(disable=True): + query, key, value, gate_out, beta_out, alpha_out = ( + gdn._prepare_qkv_for_gated_delta_rule( + qkv, gate, beta, alpha, batch, seq_len + ) + ) + + assert query.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim) + assert key.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim) + assert value.shape == (batch, seq_len, num_v_heads_local, gdn.value_head_dim) + assert query.is_contiguous() + assert key.is_contiguous() + assert value.is_contiguous() + + A_log_mock = torch.randn( + num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16, + ) + dt_bias_mock = torch.randn( + num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16, + ) + + with torch._dynamo.config.patch(disable=True): + g, beta_sig = gdn._compute_g_and_beta( + A_log_mock, dt_bias_mock, alpha, beta + ) + + assert g.dtype == torch.float32 + assert g.shape == alpha.shape + assert beta_sig.shape == beta.shape + @pytest.mark.parametrize( ("tp", "sp", "cp"), From 4ae1a027a4dd8cbc9080c68f5240f67205a38f5d Mon Sep 17 00:00:00 2001 From: lit Date: Wed, 4 Mar 2026 17:41:18 -0800 Subject: [PATCH 3/3] Fix lint --- tests/unit_tests/ssm/test_gated_delta_net.py | 37 +++++++++++--------- 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index f53f169031b..8f3c59b3d43 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -149,29 +149,36 @@ def test_jit_compiled_helpers(self): qkv_last_dim = (2 * gdn.qk_dim_local_tp + gdn.v_dim_local_tp) // gdn.cp_size qkv = torch.randn( - batch, seq_len, qkv_last_dim, - device=torch.cuda.current_device(), dtype=torch.bfloat16, + batch, seq_len, qkv_last_dim, device=torch.cuda.current_device(), dtype=torch.bfloat16 ) gate = torch.randn( - batch, seq_len, num_v_heads_local, gdn.value_head_dim, - device=torch.cuda.current_device(), dtype=torch.bfloat16, + batch, + seq_len, + num_v_heads_local, + gdn.value_head_dim, + device=torch.cuda.current_device(), + dtype=torch.bfloat16, ) beta = torch.randn( - batch, seq_len, num_v_heads_local, - device=torch.cuda.current_device(), dtype=torch.bfloat16, + batch, + seq_len, + num_v_heads_local, + device=torch.cuda.current_device(), + dtype=torch.bfloat16, ) alpha = torch.randn( - batch, seq_len, num_v_heads_local, - device=torch.cuda.current_device(), dtype=torch.bfloat16, + batch, + seq_len, + num_v_heads_local, + device=torch.cuda.current_device(), + dtype=torch.bfloat16, ) # Disable dynamo so coverage.py can trace through the method bodies, # which are normally wrapped by @jit_fuser (torch.compile). with torch._dynamo.config.patch(disable=True): query, key, value, gate_out, beta_out, alpha_out = ( - gdn._prepare_qkv_for_gated_delta_rule( - qkv, gate, beta, alpha, batch, seq_len - ) + gdn._prepare_qkv_for_gated_delta_rule(qkv, gate, beta, alpha, batch, seq_len) ) assert query.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim) @@ -182,16 +189,14 @@ def test_jit_compiled_helpers(self): assert value.is_contiguous() A_log_mock = torch.randn( - num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16, + num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16 ) dt_bias_mock = torch.randn( - num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16, + num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16 ) with torch._dynamo.config.patch(disable=True): - g, beta_sig = gdn._compute_g_and_beta( - A_log_mock, dt_bias_mock, alpha, beta - ) + g, beta_sig = gdn._compute_g_and_beta(A_log_mock, dt_bias_mock, alpha, beta) assert g.dtype == torch.float32 assert g.shape == alpha.shape