From 71b2d0169b951a5f80b1a7162bdf376afce5aed6 Mon Sep 17 00:00:00 2001
From: shiaho <222622008+shiaho777@users.noreply.github.com>
Date: Fri, 9 Oct 2026 15:45:18 +0800
Subject: [PATCH] fix: zero left-padded tokens before the GDN recurrence

GatedDeltaNet.forward received attention_mask and did not use it. A causal conv and the delta-rule update both carry state forward, so left-padding tokens change the real tokens that follow them. The loss mask removes those positions from the loss and does not clear the state.

Padding keys are the positions masked for every query, the same reduction gpt_model uses (attention_mask.all(dim=(1, 2))). Those hidden states are multiplied by zero before the projection. A mask whose shape is not [batch, seq] is left unchanged. Packed sequences keep using cu_seqlens and are not masked here.

Checked with python3 -m pytest tests/test_gdn_padding_keep.py. A [1, 1, 3, 3] mask with the first key padded produces keep [0, 1, 1]. A mask of the wrong length returns None.
---
 .../model/modules/gated_delta_net.py          | 29 ++++++++++++++++-
 tests/test_gdn_padding_keep.py                | 32 +++++++++++++++++++
 2 files changed, 60 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_gdn_padding_keep.py

diff --git a/src/mcore_bridge/model/modules/gated_delta_net.py b/src/mcore_bridge/model/modules/gated_delta_net.py
index 965e54be..e7e4a7ad 100644
--- a/src/mcore_bridge/model/modules/gated_delta_net.py
+++ b/src/mcore_bridge/model/modules/gated_delta_net.py
@@ -33,6 +33,30 @@
     _GatedDeltaNet = object
 
 
+def gdn_padding_keep(attention_mask, seq_len: int, batch: int):
+    """Return ``[seq, batch, 1]`` with 0 on padding positions, or None.
+
+    Megatron's padding mask is ``attention_mask.all(dim=(1, 2))`` on a
+    ``[batch, 1, seq, seq]`` mask, True where a key is padding. Left padding
+    has to be zero before the causal conv and the delta-rule update. The loss
+    mask does not clear that state. A mask whose length does not match the
+    hidden states is left alone rather than applied to the wrong axis.
+    """
+    if isinstance(attention_mask, dict):
+        attention_mask = attention_mask.get('full_attention')
+    if not torch.is_tensor(attention_mask):
+        return None
+    if attention_mask.dim() == 4 and attention_mask.shape[0] == batch:
+        pad = attention_mask.bool().all(dim=(1, 2))
+    elif attention_mask.dim() == 2 and tuple(attention_mask.shape) == (batch, seq_len):
+        pad = attention_mask.bool()
+    else:
+        return None
+    if tuple(pad.shape) != (batch, seq_len) or not bool(pad.any()):
+        return None
+    return (~pad).transpose(0, 1).unsqueeze(-1)
+
+
 # Code borrowed from NVIDIA/Megatron-LM
 def _unpack_sequence(x, cu_seqlens, dim=1):
     unpacked_x = []
@@ -209,8 +233,11 @@ def forward(
             (Tuple[Tensor, Tensor]) GDN output and bias.
 
         """
-        # TODO: Deal with attention_mask (There is an issue when left padding is used.)
         inference_context = deprecate_inference_params(inference_context, inference_params)
+        if packed_seq_params is None:
+            keep = gdn_padding_keep(attention_mask, hidden_states.shape[0], hidden_states.shape[1])
+            if keep is not None:
+                hidden_states = hidden_states * keep.to(dtype=hidden_states.dtype, device=hidden_states.device)
 
         use_sp = self.config.sequence_parallel and self.tp_size > 1
         tp_group = self.pg_collection.tp
diff --git a/tests/test_gdn_padding_keep.py b/tests/test_gdn_padding_keep.py
new file mode 100644
index 00000000..5a75ba25
--- /dev/null
+++ b/tests/test_gdn_padding_keep.py
@@ -0,0 +1,32 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Left-padding positions are zeroed before the GDN recurrence."""
+import ast
+import torch
+from pathlib import Path
+
+
+def _keep():
+    path = Path(__file__).resolve().parents[1] / 'src/mcore_bridge/model/modules/gated_delta_net.py'
+    tree = ast.parse(path.read_text())
+    fn = next(node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == 'gdn_padding_keep')
+    module = ast.Module(body=[fn], type_ignores=[])
+    ast.fix_missing_locations(module)
+    namespace = {'torch': torch}
+    exec(compile(module, str(path), 'exec'), namespace)
+    return namespace['gdn_padding_keep']
+
+
+def test_left_pad_column_is_dropped():
+    keep = _keep()
+    # [batch, 1, seq, seq], True = masked. Batch 0 has the first key padded.
+    mask = torch.zeros(1, 1, 3, 3, dtype=torch.bool)
+    mask[0, :, :, 0] = True
+    out = keep(mask, seq_len=3, batch=1)
+    assert out.shape == (3, 1, 1)
+    assert torch.equal(out.squeeze(), torch.tensor([0, 1, 1]))
+
+
+def test_unrelated_shape_is_ignored():
+    keep = _keep()
+    mask = torch.ones(2, 3, dtype=torch.bool)
+    assert keep(mask, seq_len=4, batch=2) is None