fix(gdn): make fused gate backward gradients repeatable

Signed-off-by: Pinoeer-kingxi <13022943007@163.com>

Author: Pinoeer-kingxi

megatron/core/fusions/fused_pre_gated_delta_rule.py

@@ -672,14 +672,6 @@ def _z_layout_to_qkvzba_kernel(
     tl.store(d_qkvzba_ptrs, dgate_val, mask=s_mask[:, None])
 
 
-@triton.autotune(
-    configs=_g_beta_autotune_configs(),
-    key=["seq_len", "num_v_heads"],
-    # Each autotune trial atomic-adds partial sums into these accumulators.
-    # Without reset_to_zero the trials would stack on top of one another and
-    # produce values that are ``num_trials`` × the correct result.
-    reset_to_zero=["d_A_log_ptr", "d_dt_bias_ptr"],
-)
 @triton.jit
 def _g_beta_backward_kernel(
     qkvzba_ptr,
@@ -721,8 +713,8 @@ def _g_beta_backward_kernel(
         d_A_log[h]   = Σ_{b,s} d_g[b,s,h] * g[b,s,h]
 
     ``d_alpha`` and ``d_beta_raw`` are written into the matching channel slices
-    of ``d_qkvzba``. ``d_A_log`` and ``d_dt_bias`` are reduced via per-element
-    atomic_add to fp32 buffers; the caller casts those to the parameter dtype.
+    of ``d_qkvzba``. Each program writes its parameter-gradient partial sums
+    to a separate FP32 slot; the caller reduces them in a fixed order.
     """
 
     pid_b = tl.program_id(0)
@@ -802,9 +794,10 @@ def _g_beta_backward_kernel(
     tl.store(d_alpha_ptrs, d_alpha.to(d_qkvzba_ptr.dtype.element_ty), mask=mask)
     tl.store(d_beta_ptrs, d_beta_raw.to(d_qkvzba_ptr.dtype.element_ty), mask=mask)
 
-    # ----- Atomic-add (b, s) partials into per-head accumulators -----
-    tl.atomic_add(d_A_log_ptr + h_offs, d_A_log_partial, mask=h_mask)
-    tl.atomic_add(d_dt_bias_ptr + h_offs, d_dt_bias_partial, mask=h_mask)
+    # ----- Store one partial per (batch, sequence block, head) -----
+    partial_offset = (pid_b * tl.num_programs(1) + pid_s) * num_v_heads + h_offs
+    tl.store(d_A_log_ptr + partial_offset, d_A_log_partial, mask=h_mask)
+    tl.store(d_dt_bias_ptr + partial_offset, d_dt_bias_partial, mask=h_mask)
 
 
 # ---------------------------------------------------------------------------
@@ -996,25 +989,23 @@ def _triton_g_beta_backward(
 
     device = qkvzba.device
 
-    g_beta_grid = lambda meta: (
-        batch,
-        triton.cdiv(seq_len, meta["BLOCK_S"]),
-        triton.cdiv(num_value_heads, meta["BLOCK_H"]),
-    )
+    # Fixed tiles keep both the partial sums and their final reduction repeatable.
+    block_s, block_h = 128, 16
+    sequence_blocks = triton.cdiv(seq_len, block_s)
+    g_beta_grid = (batch, sequence_blocks, triton.cdiv(num_value_heads, block_h))
     with _launch_context(device, stream):
-        d_param_grads = torch.empty((2, num_value_heads), dtype=torch.float32, device=device)
-        d_param_grads.zero_()
-        d_A_log = d_param_grads[0]
-        d_dt_bias = d_param_grads[1]
+        d_param_grads = torch.empty(
+            (2, batch, sequence_blocks, num_value_heads), dtype=torch.float32, device=device
+        )
         _g_beta_backward_kernel[g_beta_grid](
             qkvzba,
             A_log,
             dt_bias,
             d_g,
             d_beta_out,
             d_qkvzba_out,
-            d_A_log,
-            d_dt_bias,
+            d_param_grads[0],
+            d_param_grads[1],
             seq_len,
             num_value_heads,
             beta_channel_offset,
@@ -1028,8 +1019,13 @@ def _triton_g_beta_backward(
             d_beta_out.stride(0),
             d_beta_out.stride(1),
             d_beta_out.stride(2),
+            BLOCK_S=block_s,
+            BLOCK_H=block_h,
+            num_warps=4,
+            num_stages=2,
         )
-    return d_qkvzba_out, d_A_log, d_dt_bias
+        d_param_grads = d_param_grads.sum(dim=(1, 2))
+    return d_qkvzba_out, d_param_grads[0], d_param_grads[1]
 
 
 def _launch_context(device: torch.device, stream: Optional["torch.cuda.Stream"]):

tests/unit_tests/determinism/kernels/manifest.py

@@ -480,9 +480,18 @@ class KernelEntry:
         name="gdn_pre_gated_delta_rule_fusion",
         sources=("megatron/core/fusions/fused_pre_gated_delta_rule.py",),
         kind="triton",
-        exempt_reason="Pre-GDR fusion is explicitly rejected with deterministic_mode=True. "
-        "Its backward uses atomic parameter-gradient reductions and timing-based autotuning; "
-        "numerical parity and the rejection guard are covered by the GDN fusion unit tests.",
+        exempt_reason="Full pre-GDR fusion still rejects deterministic_mode=True. "
+        "Other kernels use timing-based autotuning and convolution backward needs its external "
+        "determinism setting. Gate backward is tested separately; full-path numerical parity "
+        "and the rejection guard remain covered by the GDN fusion unit tests.",
+    ),
+    KernelEntry(
+        name="gdn_pre_gated_delta_rule_gate_backward",
+        sources=("megatron/core/fusions/fused_pre_gated_delta_rule.py",),
+        tests=(K + "test_gdn_gate_backward.py",),
+        kind="triton",
+        notes="Gate backward uses fixed tiles and ordered FP32 parameter-gradient reduction; "
+        "replay compares input and FP32 parameter gradients before the caller's BF16 cast.",
     ),
     KernelEntry(
         name="gdn_gated_output_norm_fusion",

tests/unit_tests/determinism/kernels/test_gdn_gate_backward.py

@@ -0,0 +1,81 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Replay the pre-GDR gate gradients before casting away FP32 reduction differences."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from megatron.core.fusions.fused_pre_gated_delta_rule import _triton_g_beta_backward
+from tests.unit_tests.determinism.kernels.harness import assert_replays_bit_exact, seeded
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is unavailable")
+
+
+@pytest.mark.parametrize(
+    "batch,length,heads", [(1, 33, 8), (2, 2048, 16), (1, 4096, 16), (3, 4103, 24)]
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("strided", [False, True])
+def test_gate_backward_replays(
+    batch: int, length: int, heads: int, dtype: torch.dtype, strided: bool
+) -> None:
+    """Compare all gradients bit-exactly under contention, and against Torch autograd."""
+    seeded()
+    key_heads, dim = 4, 32
+    beta_offset = 2 * key_heads * dim + 2 * heads * dim
+    channels = beta_offset + 2 * heads
+    qkvzba = torch.randn(length, batch, channels, dtype=dtype, device="cuda")
+    A_log = torch.randn(heads, dtype=dtype, device="cuda") * 0.2
+    dt_bias = torch.randn(heads, dtype=dtype, device="cuda") * 0.2
+    d_g = torch.randn(batch, length, heads * (2 if strided else 1), device="cuda")
+    d_beta = torch.randn_like(d_g)
+    d_g, d_beta = d_g[..., :heads], d_beta[..., :heads]
+    stream = torch.cuda.Stream() if strided else None
+
+    def backward(
+        qkvzba: torch.Tensor,
+        A_log: torch.Tensor,
+        dt_bias: torch.Tensor,
+        d_g: torch.Tensor,
+        d_beta: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        result = _triton_g_beta_backward(
+            qkvzba,
+            A_log,
+            dt_bias,
+            d_g,
+            d_beta,
+            num_value_heads=heads,
+            key_head_dim=dim,
+            value_head_dim=dim,
+            num_key_heads=key_heads,
+            d_qkvzba_out=torch.zeros_like(qkvzba) if strided else None,
+            stream=stream,
+        )
+        if stream is not None:
+            torch.cuda.current_stream().wait_stream(stream)
+        return result
+
+    outputs, _ = assert_replays_bit_exact(
+        backward,
+        (qkvzba, A_log, dt_bias, d_g, d_beta),
+        backward=False,
+        replays=5,
+        contention=True,
+        what="pre-GDR gate backward",
+    )
+
+    projection = qkvzba.float().detach().requires_grad_()
+    a = A_log.float().detach().requires_grad_()
+    bias = dt_bias.float().detach().requires_grad_()
+    beta = projection[..., beta_offset : beta_offset + heads].permute(1, 0, 2).sigmoid()
+    alpha = projection[..., beta_offset + heads :].permute(1, 0, 2)
+    g = -a.exp() * F.softplus(alpha + bias)
+    reference = torch.autograd.grad((g, beta), (projection, a, bias), (d_g, d_beta))
+    reference = (reference[0].to(dtype), reference[1], reference[2])
+    torch.testing.assert_close(outputs["out[0]"], reference[0])
+    for index in (1, 2):
+        # Torch's exp/sigmoid approximations and reduction tree differ from
+        # Triton's. This numerical reference is separate from the byte replay.
+        torch.testing.assert_close(outputs[f"out[{index}]"], reference[index], atol=1e-4, rtol=2e-5)