[RL][qwen3_5] GDN: reset conv/recurrent state at packed sample boundaries

tests/unit_tests/test_attention.py

@@ -0,0 +1,35 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+from torchtitan.models.common.attention import create_varlen_metadata_for_document
+
+
+def test_single_document():
+    # One document per row: positions never reset after 0, so the only boundary
+    # is the sequence end.
+    positions = torch.arange(50).unsqueeze(0)
+    meta = create_varlen_metadata_for_document(positions)
+    assert meta.cu_seq_q.tolist() == [0, 50]
+    assert meta.cu_seq_k.tolist() == [0, 50]
+    assert meta.max_q == 50
+
+
+def test_packed_documents():
+    # Two documents packed in one row (positions restart at 0 at the boundary).
+    positions = torch.cat([torch.arange(30), torch.arange(40)]).unsqueeze(0)
+    meta = create_varlen_metadata_for_document(positions)
+    assert meta.cu_seq_q.tolist() == [0, 30, 70]
+    assert meta.max_q == 40
+
+
+def test_batched_rows_are_flattened():
+    # Boundaries are cumulative over the flattened [batch * seq_len] layout, so
+    # each row start is also a boundary.
+    positions = torch.arange(20).unsqueeze(0).repeat(2, 1)
+    meta = create_varlen_metadata_for_document(positions)
+    assert meta.cu_seq_q.tolist() == [0, 20, 40]

torchtitan/experiments/rl/tests/test_qwen3_5_gdn_packing.py

@@ -0,0 +1,68 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""GatedDeltaNet packed-sequence correctness.
+
+When several samples are packed into one row (``positions`` restart at 0 per
+sample), GatedDeltaNet is given ``cu_seqlens`` marking the boundaries so its
+recurrent state and causal conv reset per sample. This test checks that a packed
+row processed with ``cu_seqlens`` matches processing each sample on its own.
+
+Needs a GPU: the FLA gated-delta kernels are Triton/CUDA only.
+"""
+
+import pytest
+import torch
+
+from torchtitan.models.common.attention import create_varlen_metadata_for_document
+from torchtitan.models.qwen3_5 import _qwen35_deltanet_config
+
+
+def _max_abs_diff(x: torch.Tensor, y: torch.Tensor) -> float:
+    return (x - y).abs().max().item()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="FLA GDN kernels need CUDA")
+def test_gated_delta_net_packing_matches_separate():
+    torch.manual_seed(0)
+    dev, dt, dim = "cuda", torch.bfloat16, 256
+    cfg = _qwen35_deltanet_config(
+        dim=dim,
+        n_key_heads=2,
+        n_value_heads=4,
+        key_head_dim=64,
+        value_head_dim=64,
+        layer_id=0,
+    )
+    gdn = cfg.build().to(dev)
+    with torch.no_grad():
+        for name, param in gdn.named_parameters():
+            if "A_log" in name or "dt_bias" in name:
+                param.zero_()
+            elif "norm" in name and "weight" in name:
+                param.fill_(1.0)
+            else:
+                param.normal_(0, 0.02)
+    gdn = gdn.to(dt)
+
+    len_a, len_b = 30, 40
+    x_a = torch.randn(1, len_a, dim, device=dev, dtype=dt)
+    x_b = torch.randn(1, len_b, dim, device=dev, dtype=dt)
+    x_packed = torch.cat([x_a, x_b], dim=1)
+    # Packed positions restart at 0 for the second sample.
+    positions = torch.cat(
+        [torch.arange(len_a, device=dev), torch.arange(len_b, device=dev)]
+    ).unsqueeze(0)
+    cu_seqlens = create_varlen_metadata_for_document(positions).cu_seq_q
+
+    with torch.no_grad():
+        out_a = gdn(x_a)  # each sample on its own (no packing -> cu_seqlens None)
+        out_b = gdn(x_b)
+        out_packed = gdn(x_packed, cu_seqlens)
+
+    # Each packed sample matches its standalone output (to bf16 tolerance).
+    assert _max_abs_diff(out_packed[:, :len_a], out_a) < 1e-3
+    assert _max_abs_diff(out_packed[:, len_a:], out_b) < 1e-3

torchtitan/models/qwen3_5/model.py

@@ -11,6 +11,7 @@
 import torch
 import torch.nn.functional as F
 
+from fla.modules.convolution import causal_conv1d as _fla_causal_conv1d
 from fla.ops.gated_delta_rule import (
     chunk_gated_delta_rule as _fla_chunk_gated_delta_rule,
     fused_recurrent_gated_delta_rule as _fla_fused_recurrent_gated_delta_rule,
@@ -20,7 +21,11 @@
 from torch.distributed.tensor.experimental import local_map
 
 from torchtitan.models.common import Conv1d, FeedForward, Linear
-from torchtitan.models.common.attention import AttentionMasksType, BaseAttention
+from torchtitan.models.common.attention import (
+    AttentionMasksType,
+    BaseAttention,
+    create_varlen_metadata_for_document,
+)
 from torchtitan.models.common.decoder import Decoder
 from torchtitan.models.utils import get_moe_model_nparams_and_flops
 from torchtitan.protocols.module import Module
@@ -41,6 +46,7 @@ def _torch_native_gated_delta(
     v: torch.Tensor,
     g: torch.Tensor,
     beta: torch.Tensor,
+    cu_seqlens: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Standalone math reference for the gated delta rule recurrence.
 
@@ -51,10 +57,36 @@ def _torch_native_gated_delta(
         v: (bs, seqlen, n_heads, value_head_dim)
         g: (bs, seqlen, n_heads) — log-space decay, always negative
         beta: (bs, seqlen, n_heads) — update gate ∈ (0, 1)
+        cu_seqlens: optional varlen boundaries over a flattened [1, total] input;
+            when set, the recurrent state resets at each boundary (packed samples).
 
     Returns:
         output: (bs, seqlen, n_heads, value_head_dim)
     """
+    # Packed: run the recurrence per segment with a fresh state, then restore shape.
+    if cu_seqlens is not None:
+        B, L = q.shape[0], q.shape[1]
+
+        def _flatten(t: torch.Tensor) -> torch.Tensor:
+            return t.reshape(1, B * L, *t.shape[2:])
+
+        qf, kf, vf, gf, bf = (
+            _flatten(q),
+            _flatten(k),
+            _flatten(v),
+            _flatten(g),
+            _flatten(beta),
+        )
+        bounds = cu_seqlens.tolist()
+        segs = [
+            _torch_native_gated_delta(
+                qf[:, s:e], kf[:, s:e], vf[:, s:e], gf[:, s:e], bf[:, s:e]
+            )
+            for s, e in zip(bounds[:-1], bounds[1:])
+        ]
+        out = torch.cat(segs, dim=1)
+        return out.reshape(B, L, *out.shape[2:])
+
     B, L, H, D_k = q.shape
     D_v = v.shape[-1]
     dtype = q.dtype
@@ -185,6 +217,10 @@ def forward(
         v: torch.Tensor,
         g: torch.Tensor,
         beta: torch.Tensor,
+        *,
+        # Keyword-only so the TP sharding's local_map treats it as replicated
+        # pass-through metadata (like attention_masks), not a sharded input.
+        cu_seqlens: torch.Tensor | None = None,
     ) -> torch.Tensor:
         # Expand Q/K heads to match V when n_value_heads > n_key_heads
         if q.shape[2] != v.shape[2]:
@@ -194,7 +230,18 @@ def forward(
             k = k.repeat_interleave(repeat, dim=2)
 
         if self.backend == "torch_native":
-            return _torch_native_gated_delta(q, k, v, g, beta)
+            return _torch_native_gated_delta(q, k, v, g, beta, cu_seqlens=cu_seqlens)
+
+        # The FLA varlen path takes ONE [1, total] sequence with cu_seqlens marking
+        # the sample boundaries, so the recurrent state resets per packed sample.
+        # Flatten [B, L, ...] -> [1, B*L, ...]; cu_seqlens is None when not packed.
+        bs, seqlen = q.shape[0], q.shape[1]
+        if cu_seqlens is not None:
+            q = q.reshape(1, bs * seqlen, *q.shape[2:])
+            k = k.reshape(1, bs * seqlen, *k.shape[2:])
+            v = v.reshape(1, bs * seqlen, *v.shape[2:])
+            g = g.reshape(1, bs * seqlen, *g.shape[2:])
+            beta = beta.reshape(1, bs * seqlen, *beta.shape[2:])
 
         if self.backend == "fla_chunked":
             result = _fla_chunk_gated_delta_rule(
@@ -204,6 +251,7 @@ def forward(
                 g,
                 beta,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=cu_seqlens,
             )
         elif self.backend == "fla_fused_recurrent":
             result = _fla_fused_recurrent_gated_delta_rule(
@@ -213,6 +261,7 @@ def forward(
                 g,
                 beta=beta,
                 use_qk_l2norm_in_kernel=True,
+                cu_seqlens=cu_seqlens,
             )
         else:
             raise ValueError(
@@ -221,7 +270,10 @@ def forward(
             )
 
         # FLA kernels return (output, final_state); we only need output
-        return result[0]
+        out = result[0]
+        if cu_seqlens is not None:
+            out = out.reshape(bs, seqlen, *out.shape[2:])
+        return out
 
 
 class GatedDeltaNet(Module):
@@ -279,7 +331,49 @@ def __init__(self, config: Config):
         self.norm = config.norm.build()
         self.out_proj = config.out_proj.build()
 
-    def _causal_conv(self, x: torch.Tensor, conv: nn.Module) -> torch.Tensor:
+    def _causal_conv(
+        self, x: torch.Tensor, conv: nn.Module, cu_seqlens: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        # Packed samples: the causal conv window must not cross sample boundaries
+        # (else the first conv_kernel_size-1 tokens of each sample would see the
+        # previous one). F.conv1d cannot express a per-segment reset, so the packed
+        # path uses fla.causal_conv1d, which resets at cu_seqlens. The conv is
+        # depthwise (per-channel), so under TP (channel-sharded) we run it on local
+        # shards via local_map -- same pattern as the F.conv1d path below. GDN conv
+        # is bias-free.
+        if cu_seqlens is not None:
+            bs, seqlen, _ = x.shape
+
+            def _fla_conv(x_local: torch.Tensor, w_local: torch.Tensor) -> torch.Tensor:
+                # x_local: [bs, seqlen, C_local]; depthwise weight [C_local, 1, k].
+                # fla wants one [1, total, C] sequence; cu_seqlens spans the
+                # flattened [bs*seqlen] layout (offsets already include batch rows).
+                # fla's causal_conv1d is untyped (pyrefly reads it as a Tensor).
+                y = _fla_causal_conv1d(
+                    x_local.reshape(1, bs * seqlen, -1),
+                    weight=w_local.squeeze(1),
+                    bias=None,
+                    activation="silu",
+                    cu_seqlens=cu_seqlens,
+                )
+                if isinstance(y, tuple):
+                    y = y[0]
+                return y.reshape(bs, seqlen, -1)
+
+            if isinstance(x, DTensor):
+                # Channel-sharded depthwise conv: run per shard, restore DTensor-ness.
+                x_plc = x.placements
+                w_plc = conv.weight.placements  # pyrefly: ignore [missing-attribute]
+                conv_dt = local_map(
+                    _fla_conv,
+                    out_placements=(x_plc,),
+                    in_placements=(x_plc, w_plc),
+                    in_grad_placements=(x_plc, w_plc),
+                    device_mesh=x.device_mesh,
+                )
+                return conv_dt(x, conv.weight)  # pyrefly: ignore
+            return _fla_conv(x, conv.weight)  # pyrefly: ignore [bad-argument-type]
+
         x = F.pad(x.transpose(1, 2), [self.conv_kernel_size - 1, 0])
         if isinstance(x, DTensor):
             # TODO: Remove once the DTensor Conv1d dispatch fix for sharded
@@ -315,16 +409,18 @@ def _conv(x_local: torch.Tensor, w_local: torch.Tensor) -> torch.Tensor:
             x = conv(x)
         return F.silu(x).transpose(1, 2)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, x: torch.Tensor, cu_seqlens: torch.Tensor | None = None
+    ) -> torch.Tensor:
         bs, seqlen, _ = x.shape
 
         # Shapes:
         #   xq, xk: (bs, seqlen, n_key_heads * key_head_dim)
         #   xv, xz: (bs, seqlen, n_value_heads * value_head_dim)
         #   xa, xb: (bs, seqlen, n_value_heads)
-        xq = self._causal_conv(self.in_proj_q(x), self.conv_q)
-        xk = self._causal_conv(self.in_proj_k(x), self.conv_k)
-        xv = self._causal_conv(self.in_proj_v(x), self.conv_v)
+        xq = self._causal_conv(self.in_proj_q(x), self.conv_q, cu_seqlens)
+        xk = self._causal_conv(self.in_proj_k(x), self.conv_k, cu_seqlens)
+        xv = self._causal_conv(self.in_proj_v(x), self.conv_v, cu_seqlens)
         xz = self.in_proj_z(x)
         xa = self.in_proj_a(x)
         xb = self.in_proj_b(x)
@@ -339,7 +435,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         g = -torch.exp(self.A_log.float()) * F.softplus(xa.float() + self.dt_bias)
         beta = torch.sigmoid(xb)
 
-        output = self.kernel(xq, xk, xv, g, beta)
+        output = self.kernel(xq, xk, xv, g, beta, cu_seqlens=cu_seqlens)
 
         xz = xz.view(bs, seqlen, -1, self.value_head_dim)
         output = self.norm(output, xz)
@@ -483,12 +579,13 @@ def forward(
         x: torch.Tensor,
         attention_masks: AttentionMasksType | None,
         positions: torch.Tensor | None = None,
+        cu_seqlens: torch.Tensor | None = None,
     ) -> torch.Tensor:
         h = self.attention_norm(x)
         if self.full_attn:
             h = self.attn(h, attention_masks, positions)
         else:
-            h = self.attn(h)
+            h = self.attn(h, cu_seqlens)
         x = x + h
 
         h = self.ffn_norm(x)
@@ -769,8 +866,19 @@ def forward(  # pyrefly: ignore [bad-override]
         # 3D MRoPE positions for multimodal batches, else 2D text positions.
         rope_positions = mrope_positions if mrope_positions is not None else positions
         assert rope_positions is not None
+        # cu_seqlens holds the packed sample boundaries (positions == 0), computed
+        # once and shared across GatedDeltaNet layers so their recurrent state and
+        # causal conv reset per sample. Reuses the same document-varlen metadata as
+        # the full-attention masks. None for a single unpacked sequence (unchanged
+        # path). NOTE: not context-parallel aware (CP would split the sequence and
+        # break these boundaries); CP is unsupported for GatedDeltaNet.
+        cu_seqlens = None
+        if positions is not None:
+            cu_seqlens = create_varlen_metadata_for_document(positions).cu_seq_q
+            if cu_seqlens.numel() <= 2:  # single unpacked sample -> non-varlen path
+                cu_seqlens = None
         for layer in self.layers.values():
-            x = layer(x, attention_masks, rope_positions)
+            x = layer(x, attention_masks, rope_positions, cu_seqlens)
 
         x = self.norm(x) if self.norm is not None else x
         if self._skip_lm_head: