perf(gated_delta_net): fold q/k L2-norm into the gated_delta_rule kernel

megatron/core/ssm/gated_delta_net/common.py

@@ -389,13 +389,16 @@ def _prepare_input_for_gated_delta_rule(
         batch: int,
         seq_len: int,
         *gate_feats: tuple[torch.Tensor],
+        use_qk_l2norm_in_kernel: bool = False,
     ) -> dict[str, torch.Tensor]:
         """
         Prepare all gated delta rule kernel inputs.
 
         Fuses split, reshape, L2 norm, decay/gate activations, repeat_interleave, and
         contiguous operations. ``gate_feats`` holds the variant-specific in_proj
         sections, which ``_compute_gates`` turns into the decay and gating tensors.
+        When ``use_qk_l2norm_in_kernel`` is true, q/k normalization is deferred to
+        the gated delta rule kernel to avoid materializing normalized q/k here.
 
         Returns:
             (dict[str, Tensor]): Kernel inputs keyed by kernel argument name (``q``,
@@ -413,8 +416,9 @@ def _prepare_input_for_gated_delta_rule(
         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:
+        # Let a supporting kernel own normalization so caller autograd does not retain an
+        # additional pre-split, normalized query_key activation alongside the kernel q/k.
+        if self.use_qk_l2norm and not use_qk_l2norm_in_kernel:
             query_key = l2norm(query_key.contiguous())
 
         # Split query and key

megatron/core/ssm/gated_delta_net/gdn.py

@@ -248,7 +248,15 @@ def forward(
         # Prepare all kernel inputs (split, reshape, L2 norm, gates, contiguous)
         nvtx_range_push(suffix="prepare_input_for_gated_delta_rule")
         kernel_inputs = self._prepare_input_for_gated_delta_rule(
-            qkv, gate, A_log_local_cp, dt_bias_local_cp, batch, seq_len, beta, alpha
+            qkv,
+            gate,
+            A_log_local_cp,
+            dt_bias_local_cp,
+            batch,
+            seq_len,
+            beta,
+            alpha,
+            use_qk_l2norm_in_kernel=self.use_qk_l2norm,
         )
         gate = kernel_inputs.pop("gate")
         nvtx_range_pop(suffix="prepare_input_for_gated_delta_rule")
@@ -258,7 +266,7 @@ def forward(
             **kernel_inputs,
             initial_state=None,
             output_final_state=False,
-            use_qk_l2norm_in_kernel=False,
+            use_qk_l2norm_in_kernel=self.use_qk_l2norm,
             cu_seqlens=cu_seqlens_q,
         )
         nvtx_range_pop(suffix="gated_delta_rule")
@@ -452,8 +460,8 @@ def torch_chunk_gated_delta_rule(
     query, key, value = q, k, v
     initial_dtype = query.dtype
     if use_qk_l2norm_in_kernel:
-        query = l2norm(query, dim=-1, eps=1e-6)
-        key = l2norm(key, dim=-1, eps=1e-6)
+        query = l2norm(query)
+        key = l2norm(key)
     query, key, value, beta, g = [
         x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
     ]

megatron/core/ssm/gated_delta_net/gdn2.py

@@ -377,8 +377,8 @@ def torch_chunk_gdn2(
 
     initial_dtype = q.dtype
     if use_qk_l2norm_in_kernel:
-        q = l2norm(q, dim=-1, eps=1e-6)
-        k = l2norm(k, dim=-1, eps=1e-6)
+        q = l2norm(q)
+        k = l2norm(k)
 
     # b s h d -> b h s d, and compute the whole recurrence in fp32
     query, key, value, g, b, w = [

tests/unit_tests/ssm/test_gated_delta_net.py

@@ -25,6 +25,7 @@
 from megatron.core.ssm.gated_delta_net.common import (
     _build_head_perm_for_split_sections,
     _build_thd_cp_a2a_perm,
+    l2norm,
     tensor_a2a_cp2hp,
     tensor_a2a_hp2cp,
 )
@@ -54,6 +55,229 @@ def _unpack_sequence(x: torch.Tensor, cu_seqlens: torch.Tensor, dim=1) -> list[t
     return unpacked_x
 
 
+@pytest.mark.parametrize("use_gdn2", [False, True], ids=["gdn", "gdn2"])
+@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.")
+@pytest.mark.internal
+def test_torch_gdn_l2norm_flag_matches_explicit_normalization(use_gdn2):
+    """The torch fallback flag must match explicit q/k normalization."""
+    torch.manual_seed(123)
+    torch.cuda.set_device(Utils.local_rank % torch.cuda.device_count())
+    device = torch.cuda.current_device()
+    batch_size, sequence_length, num_heads, head_dim = 1, 8, 2, 16
+    shape = (batch_size, sequence_length, num_heads, head_dim)
+    q = torch.randn(shape, device=device)
+    k = torch.randn(shape, device=device)
+    v = torch.randn(shape, device=device)
+
+    if use_gdn2:
+        kernel = torch_chunk_gdn2
+        kernel_kwargs = {
+            "g": torch.full(shape, -0.01, device=device),
+            "b": torch.full(shape, 0.1, device=device),
+            "w": torch.full(shape, 0.1, device=device),
+        }
+    else:
+        kernel = torch_chunk_gated_delta_rule
+        gate_shape = (batch_size, sequence_length, num_heads)
+        kernel_kwargs = {
+            "g": torch.full(gate_shape, -0.01, device=device),
+            "beta": torch.full(gate_shape, 0.1, device=device),
+        }
+
+    actual_output, actual_state = kernel(
+        q=q,
+        k=k,
+        v=v,
+        chunk_size=8,
+        output_final_state=True,
+        use_qk_l2norm_in_kernel=True,
+        **kernel_kwargs,
+    )
+    expected_output, expected_state = kernel(
+        q=l2norm(q),
+        k=l2norm(k),
+        v=v,
+        chunk_size=8,
+        output_final_state=True,
+        use_qk_l2norm_in_kernel=False,
+        **kernel_kwargs,
+    )
+
+    torch.testing.assert_close(actual_output, expected_output, rtol=0, atol=0)
+    torch.testing.assert_close(actual_state, expected_state, rtol=0, atol=0)
+
+
+@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.")
+@pytest.mark.internal
+def test_gdn_in_kernel_l2norm_matches_caller_end_to_end():
+    """GDN's BF16 in-kernel route must match the former caller-normalized route."""
+    Utils.initialize_model_parallel(
+        tensor_model_parallel_size=1, pipeline_model_parallel_size=1, context_parallel_size=1
+    )
+    try:
+        tp_group = parallel_state.get_tensor_model_parallel_group()
+        cp_group = parallel_state.get_context_parallel_group()
+        pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group)
+        config = TransformerConfig(
+            hidden_size=128,
+            linear_conv_kernel_dim=4,
+            linear_key_head_dim=32,
+            linear_value_head_dim=32,
+            linear_num_key_heads=2,
+            linear_num_value_heads=4,
+            num_layers=1,
+            normalization="RMSNorm",
+            use_cpu_initialization=True,
+            layernorm_zero_centered_gamma=True,
+            num_attention_heads=4,
+            num_query_groups=2,
+            activation_func=F.silu,
+            bf16=True,
+            tensor_model_parallel_size=1,
+            sequence_parallel=False,
+            context_parallel_size=1,
+            experimental_attention_variant="gated_delta_net",
+            linear_attention_freq=[1],
+            transformer_impl="transformer_engine",
+        )
+        gdn_spec = get_experimental_attention_variant_module_spec(config=config)
+
+        def build_gdn():
+            return (
+                gdn_spec.module(
+                    config,
+                    submodules=gdn_spec.submodules,
+                    layer_number=1,
+                    bias=False,
+                    conv_bias=False,
+                    conv_init=1.0,
+                    use_qk_l2norm=True,
+                    A_init_range=(1, 16),
+                    pg_collection=pg_collection,
+                )
+                .cuda()
+                .bfloat16()
+            )
+
+        model_parallel_cuda_manual_seed(42)
+        torch.manual_seed(42)
+        kernel_norm_gdn = build_gdn()
+        caller_norm_gdn = build_gdn()
+        caller_norm_gdn.load_state_dict(kernel_norm_gdn.state_dict())
+
+        def install_route_probe(module, caller_normalizes):
+            original_prepare = module._prepare_input_for_gated_delta_rule
+            original_kernel = module.gated_delta_rule
+            observed = {}
+
+            def wrapped_prepare(
+                qkv,
+                gate,
+                A_log_local_cp,
+                dt_bias_local_cp,
+                batch,
+                seq_len,
+                *gate_feats,
+                use_qk_l2norm_in_kernel=False,
+            ):
+                observed["prepare_flag"] = use_qk_l2norm_in_kernel
+                effective_flag = False if caller_normalizes else use_qk_l2norm_in_kernel
+                kernel_inputs = original_prepare(
+                    qkv,
+                    gate,
+                    A_log_local_cp,
+                    dt_bias_local_cp,
+                    batch,
+                    seq_len,
+                    *gate_feats,
+                    use_qk_l2norm_in_kernel=effective_flag,
+                )
+
+                # The new route must pass the raw post-convolution q/k to FLA. This
+                # explicit check catches accidental caller normalization (and hence
+                # double normalization) even when BF16 rounding hides the difference.
+                if not caller_normalizes:
+                    query_key, _ = torch.split(
+                        qkv.detach(),
+                        [
+                            2 * module.qk_dim_local_tp // module.cp_size,
+                            module.v_dim_local_tp // module.cp_size,
+                        ],
+                        dim=-1,
+                    )
+                    query_key = query_key.reshape(batch, seq_len, -1, module.key_head_dim)
+                    split_size = module.qk_dim_local_tp // module.key_head_dim // module.cp_size
+                    raw_query, raw_key = torch.split(query_key, [split_size, split_size], dim=2)
+                    repeat_factor = module.num_value_heads // module.num_key_heads
+                    if repeat_factor > 1:
+                        raw_query = raw_query.repeat_interleave(repeat_factor, dim=2)
+                        raw_key = raw_key.repeat_interleave(repeat_factor, dim=2)
+                    observed["prepare_kept_qk_raw"] = torch.equal(
+                        kernel_inputs["q"].detach(), raw_query.contiguous()
+                    ) and torch.equal(kernel_inputs["k"].detach(), raw_key.contiguous())
+                return kernel_inputs
+
+            def wrapped_kernel(*args, **kwargs):
+                observed["kernel_flag"] = kwargs.get("use_qk_l2norm_in_kernel", False)
+                if caller_normalizes:
+                    kwargs = {**kwargs, "use_qk_l2norm_in_kernel": False}
+                return original_kernel(*args, **kwargs)
+
+            module._prepare_input_for_gated_delta_rule = wrapped_prepare
+            module.gated_delta_rule = wrapped_kernel
+            return observed
+
+        kernel_route = install_route_probe(kernel_norm_gdn, caller_normalizes=False)
+        caller_route = install_route_probe(caller_norm_gdn, caller_normalizes=True)
+
+        torch.manual_seed(123)
+        hidden_states = torch.randn(
+            (16, 2, config.hidden_size), device=torch.cuda.current_device(), dtype=torch.bfloat16
+        )
+
+        def run(module):
+            module.zero_grad(set_to_none=True)
+            module_input = hidden_states.detach().clone().requires_grad_(True)
+            output, _ = module(module_input, None)
+            output.float().square().mean().backward()
+            parameter_grads = {
+                name: parameter.grad.detach().clone()
+                for name, parameter in module.named_parameters()
+                if parameter.grad is not None
+            }
+            return output.detach(), module_input.grad.detach().clone(), parameter_grads
+
+        kernel_output, kernel_input_grad, kernel_parameter_grads = run(kernel_norm_gdn)
+        caller_output, caller_input_grad, caller_parameter_grads = run(caller_norm_gdn)
+
+        assert kernel_route["prepare_flag"] is True
+        assert kernel_route["kernel_flag"] is True
+        assert kernel_route["prepare_kept_qk_raw"] is True
+        # The compatibility route still goes through GatedDeltaNet.forward, then
+        # overrides only where normalization is performed to model the former code.
+        assert caller_route["prepare_flag"] is True
+        assert caller_route["kernel_flag"] is True
+
+        atol = rtol = 2e-2
+        torch.testing.assert_close(
+            kernel_output.float(), caller_output.float(), atol=atol, rtol=rtol
+        )
+        torch.testing.assert_close(
+            kernel_input_grad.float(), caller_input_grad.float(), atol=atol, rtol=rtol
+        )
+        assert set(kernel_parameter_grads) == set(caller_parameter_grads)
+        for name in kernel_parameter_grads:
+            torch.testing.assert_close(
+                kernel_parameter_grads[name].float(),
+                caller_parameter_grads[name].float(),
+                atol=atol,
+                rtol=rtol,
+                msg=lambda msg, name=name: f"Parameter grad mismatch for {name}: {msg}",
+            )
+    finally:
+        Utils.destroy_model_parallel()
+
+
 @pytest.mark.parametrize("use_gdn2", [False, True], ids=["gdn", "gdn2"])
 @pytest.mark.parametrize(
     ("tp_size", "sp", "cp_size"),