diff --git a/.ci/docker/requirements.txt b/.ci/docker/requirements.txt
index a2cec99127..7fb150c9ac 100644
--- a/.ci/docker/requirements.txt
+++ b/.ci/docker/requirements.txt
@@ -8,4 +8,4 @@ safetensors
 einops
 pillow
 spmd_types==0.2.5
-attn-gym[linear]==0.0.5
+attn-gym[linear]==0.0.6
diff --git a/pyproject.toml b/pyproject.toml
index 8adfbd9fe2..b6b9bf0f2f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -28,7 +28,7 @@ dependencies = [
     "einops",
     "pillow",
     "spmd_types==0.2.5",
-    "attn-gym[linear]==0.0.5",
+    "attn-gym[linear]==0.0.6",
 ]
 dynamic = ["version"]
 
diff --git a/tests/unit_tests/gpu/test_qwen3_5_deltanet.py b/tests/unit_tests/gpu/test_qwen3_5_deltanet.py
index adba767f25..ad02042897 100644
--- a/tests/unit_tests/gpu/test_qwen3_5_deltanet.py
+++ b/tests/unit_tests/gpu/test_qwen3_5_deltanet.py
@@ -9,7 +9,9 @@
 
 import torch
 import torch.nn.functional as F
+from attn_gym.linear import l2norm, recurrent_gdn
 from torch import nn
+
 from torchtitan.models.common.attention import (
     create_varlen_metadata_for_document,
     VarlenMetadata,
@@ -103,8 +105,8 @@ def _reference_causal_conv1d_varlen(
     cu_seqlens: torch.Tensor,
     cu_seqlens_cpu: torch.Tensor,
 ) -> torch.Tensor:
-    """Per-document depthwise causal conv + silu, matching the model's FLA
-    varlen conv (which is triton/CUDA-only). Patched over
+    """Per-document depthwise causal conv + silu, matching the model's Attention
+    Gym varlen conv (which is CUDA-only). Patched over
     ``gdn._causal_conv1d_varlen`` for CPU runs.
     """
     conv_kernel_size = weight.shape[-1]
@@ -399,7 +401,7 @@ def test_extracted_forward_matches_main(self):
     def _assert_packed_run_matches_per_document(self, model, x, positions, masks):
         """Packed forward under ``masks`` must equal stitched per-doc forwards.
 
-        The model's varlen conv is FLA (triton/CUDA-only); substitute the
+        The model's varlen conv is Attention Gym (CUDA-only); substitute the
         per-document torch reference for these CPU runs. The per-document
         forwards below take the non-varlen conv path, which runs on CPU.
         """
@@ -600,8 +602,112 @@ def test_fla_fused_recurrent_varlen_matches_independent_document_forwards(self):
             "fla_fused_recurrent", atol=2e-2, rtol=2e-2
         )
 
+    def test_batch_invariant_recurrent_matches_paged_attention_gym(self):
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA is unavailable")
+
+        from torchtitan.models.qwen3_5.gdn import _recurrent_gdn_fwd
+
+        torch.manual_seed(42)
+        num_tokens, num_heads, key_dim, value_dim = 12, 4, 64, 64
+        q = torch.randn(
+            1,
+            num_tokens,
+            num_heads,
+            key_dim,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        k = torch.randn_like(q)
+        v = torch.randn(
+            1,
+            num_tokens,
+            num_heads,
+            value_dim,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        decay = -torch.rand(
+            1,
+            num_tokens,
+            num_heads,
+            device="cuda",
+            dtype=torch.float32,
+        )
+        update_gate = torch.rand(
+            1,
+            num_tokens,
+            num_heads,
+            device="cuda",
+            dtype=torch.float32,
+        )
+        cu_seqlens = torch.tensor([0, 5, 12], device="cuda", dtype=torch.int32)
+
+        actual = _recurrent_gdn_fwd(
+            q,
+            k,
+            v,
+            decay,
+            update_gate,
+            cu_seqlens,
+            cu_seqlens.cpu(),
+        )
+
+        normalized_q = l2norm(q, cu_seqlens=cu_seqlens)
+        normalized_k = l2norm(k, cu_seqlens=cu_seqlens)
+        state_cache = torch.randn(
+            5,
+            num_heads,
+            value_dim,
+            key_dim,
+            device="cuda",
+            dtype=torch.float32,
+        )
+        prefix_end = 2
+        prefix_cu_seqlens = torch.tensor(
+            [0, prefix_end], device="cuda", dtype=torch.int32
+        )
+        prefix_output, _ = recurrent_gdn(
+            normalized_q[:, :prefix_end],
+            normalized_k[:, :prefix_end],
+            v[:, :prefix_end],
+            decay[:, :prefix_end],
+            update_gate[:, :prefix_end],
+            state_cache,
+            cu_seqlens=prefix_cu_seqlens,
+            scale=key_dim**-0.5,
+            state_indices=torch.tensor([3], device="cuda", dtype=torch.int32),
+            has_initial_state=torch.tensor([False], device="cuda"),
+        )
+
+        state_indices = torch.tensor([3, 1], device="cuda", dtype=torch.int32)
+        has_initial_state = torch.tensor([True, False], device="cuda")
+        remaining_cu_seqlens = torch.tensor(
+            [0, 5 - prefix_end, num_tokens - prefix_end],
+            device="cuda",
+            dtype=torch.int32,
+        )
+        remaining_output, _ = recurrent_gdn(
+            normalized_q[:, prefix_end:],
+            normalized_k[:, prefix_end:],
+            v[:, prefix_end:],
+            decay[:, prefix_end:],
+            update_gate[:, prefix_end:],
+            state_cache,
+            cu_seqlens=remaining_cu_seqlens,
+            scale=key_dim**-0.5,
+            state_indices=state_indices,
+            has_initial_state=has_initial_state,
+        )
+        expected = torch.cat(
+            (prefix_output, remaining_output),
+            dim=1,
+        )
+
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
     def test_varlen_offsets_are_fresh_per_deltanet_invocation(self):
-        """Successive DeltaNet invocations must not share FLA's cache key."""
+        """Successive DeltaNet invocations must not share convolution metadata."""
         torch.manual_seed(42)
         model = self._make_deltanet()
         x_TD = torch.randn(8, 4)
diff --git a/torchtitan/experiments/rl/batch_invariance.py b/torchtitan/experiments/rl/batch_invariance.py
index 3dd6f5bfb8..20ce18cda5 100644
--- a/torchtitan/experiments/rl/batch_invariance.py
+++ b/torchtitan/experiments/rl/batch_invariance.py
@@ -72,7 +72,15 @@ def patch_bmm_for_batch_invariance() -> None:
     global _batch_invariant_bmm_lib
     if _batch_invariant_bmm_lib is not None:
         return
-    from vllm.model_executor.determinism.batch_invariant import bmm_batch_invariant
+    try:
+        from vllm.model_executor.determinism.batch_invariant import bmm_batch_invariant
+    except ModuleNotFoundError as error:
+        if error.name not in {
+            "vllm.model_executor.determinism",
+            "vllm.model_executor.determinism.batch_invariant",
+        }:
+            raise
+        from vllm.model_executor.layers.batch_invariant import bmm_batch_invariant
 
     _batch_invariant_bmm_lib = torch.library.Library("aten", "IMPL")
     _batch_invariant_bmm_lib.impl("bmm", bmm_batch_invariant, "CUDA")
diff --git a/torchtitan/experiments/rl/models/gdn.py b/torchtitan/experiments/rl/models/gdn.py
index 6544b5ee18..4982dfd033 100644
--- a/torchtitan/experiments/rl/models/gdn.py
+++ b/torchtitan/experiments/rl/models/gdn.py
@@ -6,44 +6,42 @@
 
 """vLLM paged-cache adapter for TorchTitan's Gated DeltaNet.
 
-The enclosing Qwen3.5 module owns all parameters. This adapter runs the same FLA
-convolution and recurrence kernels as training while reading and updating vLLM's
-paged convolution and SSM states.
+The enclosing Qwen3.5 module owns all parameters. This adapter runs Attention
+Gym's paging-aware convolution kernels. Batch-invariant recurrence and decode use
+Attention Gym, while ordinary prefill uses FLA's parallel chunk kernel.
 
 Batch-invariant execution has two additional requirements:
 
 * The accumulated SSM cache state uses float32. Decode otherwise rounds the state
   through bfloat16 after every token, unlike a single prefill call. The convolution
   cache stays in model dtype because it only stores trailing input columns.
-* Every recurrence receives a materialized initial state and ``cu_seqlens``. This
-  keeps FLA's Triton specialization identical for fresh prefill and resumed state.
+* Batch-invariant recurrence uses the same Attention Gym scan as the trainer.
 
-The recurrent kernel and materialized state are also used outside batch-invariant
-mode because vLLM generation must support prefix continuation. SSM cache precision
-and prefill kernel selection differ by mode.
+Attention Gym updates the paged state pool directly. The non-batch-invariant FLA
+prefill path retains its state gather and scatter.
 """
 
 from dataclasses import dataclass
 
 import torch
 import torch.nn.functional as F
-
-# Using the trainer's FLA kernels is required for batch-invariant parity.
-from fla.modules.convolution import (
-    causal_conv1d as _fla_causal_conv1d,
-    causal_conv1d_update as _fla_causal_conv1d_update,
+from attn_gym.linear import (
+    causal_conv1d as _attn_gym_causal_conv1d,
+    causal_conv1d_decode as _attn_gym_causal_conv1d_decode,
+    l2norm as _attn_gym_l2norm,
+    recurrent_gdn as _attn_gym_recurrent_gdn,
+    recurrent_gdn_decode as _attn_gym_recurrent_gdn_decode,
 )
 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,
 )
+
 from torchtitan.distributed.utils import is_in_batch_invariant_mode
 from torchtitan.protocols.module import Module
 
 # The recurrence mutates paged state and must run eager at a breakable cudagraph
 # split point. This decorator is inert when breakable capture is disabled.
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
-
 from vllm.config import get_current_vllm_config
 from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.mamba.abstract import MambaBase
@@ -60,7 +58,7 @@ class VLLMInnerGatedDeltaNet(Module, MambaBase):
     """Paged-cache inner GDN implementation.
 
     The enclosing ``qwen3_5.gdn.GatedDeltaNet`` owns all parameters. This
-    module owns only vLLM cache plumbing and the FLA kernels.
+    module owns only vLLM cache plumbing and kernel dispatch.
 
     The enclosing module and vLLM cache are both head-sharded under tensor
     parallelism. Speculative decoding is not supported.
@@ -86,6 +84,8 @@ def __init__(self, config: Config) -> None:
         self.num_speculative_tokens = (
             speculative_config.num_speculative_tokens if speculative_config else 0
         )
+        if self.num_speculative_tokens != 0:
+            raise ValueError("Attention Gym GDN does not support speculative decoding.")
 
         self.num_k_heads = config.num_k_heads
         self.num_v_heads = config.num_v_heads
@@ -108,11 +108,13 @@ def __init__(self, config: Config) -> None:
         self.local_num_v_heads = self.num_v_heads // self.tensor_parallel_size
         self.local_key_dim = self.local_num_k_heads * self.head_k_dim
 
-        if is_in_batch_invariant_mode():
-            # The recurrent state accumulates in float32 inside FLA, so preserving
-            # it across scheduler calls avoids per-token bfloat16 rounding. The
-            # convolution cache only stores model-dtype input columns verbatim.
-            self.cache_config.mamba_ssm_cache_dtype = "float32"
+        if is_conv_state_dim_first():
+            raise ValueError(
+                "Attention Gym GDN requires VLLM_SSM_CONV_STATE_LAYOUT=SD so "
+                "the paged convolution history has contiguous channels."
+            )
+
+        self.cache_config.mamba_ssm_cache_dtype = "float32"
 
         # vLLM populates this via the KV-cache allocator: (conv_state, ssm_state).
         self.kv_cache = (torch.tensor([]), torch.tensor([]))
@@ -158,7 +160,7 @@ def get_state_shape(self) -> tuple[tuple[int, ...], ...]:
     def _split_qkv(
         self, mixed_qkv: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Split local fused channels and add FLA's singleton batch dim."""
+        """Split local fused channels and add a singleton batch dim."""
         num_tokens = mixed_qkv.shape[0]
         local_key_dim = self.local_key_dim
         query = (
@@ -185,17 +187,13 @@ def _run_recurrence(
         b: torch.Tensor,
         negative_exp_A: torch.Tensor,
         dt_bias: torch.Tensor,
-        initial_state: torch.Tensor,
         ssm_state: torch.Tensor,
         slot_indices: torch.Tensor,
         cu_seqlens: torch.Tensor,
-        use_chunk: bool,
+        has_initial_state: torch.Tensor | None,
+        batch_invariant: bool,
     ) -> torch.Tensor:
-        """Run FLA's gated-delta rule and update the selected paged SSM slots.
-
-        ``use_chunk`` selects the parallel chunk kernel over the sequential
-        recurrent kernel; the caller sets it for non-batch-invariant prefill.
-        """
+        """Run the selected recurrence and update the paged SSM slots."""
         query, key, value = self._split_qkv(conv_output)
         # Grouped-value heads: expand q/k to match the value head count.
         if query.shape[2] != value.shape[2]:
@@ -206,17 +204,41 @@ def _run_recurrence(
         decay = (negative_exp_A * F.softplus(a.float() + dt_bias)).unsqueeze(0)
         update_gate = torch.sigmoid(b).unsqueeze(0)
 
-        # Non-batch-invariant prefill/mixed batches use the parallel chunk kernel
-        # (far faster than the O(seqlen) sequential recurrence). Batch-invariant
-        # mode keeps the recurrent kernel so the generator matches the trainer's
-        # recurrent forward bitwise; pure decode is a single token, where chunking
-        # gains nothing and the recurrent kernel stays cudagraph-capturable.
-        gated_delta_rule = (
-            _fla_chunk_gated_delta_rule
-            if use_chunk
-            else _fla_fused_recurrent_gated_delta_rule
+        if batch_invariant:
+            query = _attn_gym_l2norm(query, cu_seqlens=cu_seqlens)
+            key = _attn_gym_l2norm(key, cu_seqlens=cu_seqlens)
+            output, _ = _attn_gym_recurrent_gdn(
+                query,
+                key,
+                value,
+                decay,
+                update_gate,
+                ssm_state,
+                cu_seqlens=cu_seqlens,
+                scale=self.head_k_dim**-0.5,
+                state_indices=slot_indices,
+                has_initial_state=has_initial_state,
+                # Triton autotuning breaks batch invariance.
+                autotune=False,
+            )
+            return output
+
+        num_sequences = slot_indices.numel()
+        initial_state = conv_output.new_zeros(
+            num_sequences,
+            self.local_num_v_heads,
+            self.head_k_dim,
+            self.head_v_dim,
+            dtype=torch.float32,
         )
-        output, final_state = gated_delta_rule(
+        if has_initial_state is None:
+            initial_state.copy_(ssm_state[slot_indices].transpose(-1, -2))
+        else:
+            resumed_slots = slot_indices[has_initial_state]
+            initial_state[has_initial_state] = ssm_state[resumed_slots].transpose(
+                -1, -2
+            )
+        output, final_state = _fla_chunk_gated_delta_rule(
             query,
             key,
             value,
@@ -227,6 +249,7 @@ def _run_recurrence(
             cu_seqlens=cu_seqlens,
             use_qk_l2norm_in_kernel=True,
         )
+        assert final_state is not None
         ssm_state[slot_indices] = final_state.transpose(-1, -2).to(ssm_state.dtype)
         return output
 
@@ -269,13 +292,8 @@ def _forward(
         state_indices = gdn_metadata.non_spec_state_indices_tensor
         assert state_indices is not None
         ssm_state = self.kv_cache[1]
-        # vLLM's default conv-state layout keeps channels innermost (.., W-1, C);
-        # FLA needs channels first (.., C, W-1), so transpose that case.
-        conv_state = (
-            self.kv_cache[0]
-            if is_conv_state_dim_first()
-            else self.kv_cache[0].transpose(-1, -2)
-        )
+        conv_state = self.kv_cache[0]
+        assert conv_bias is None
         negative_exp_A = -torch.exp(A_log.float())
         dt_bias = dt_bias.float()
         num_decodes = gdn_metadata.num_decodes
@@ -293,21 +311,13 @@ def _forward(
 
         decode_slots = state_indices[:num_decodes]
         if num_decodes > 0:  # pure decode, or mixed prefill decode
-            decode_conv_state = conv_state[decode_slots]
-            # vLLM stores the trailing W - 1 inputs. FLA expects W entries but
-            # ignores the first, so prepend a zero column and drop it on write.
-            zero_padding = decode_conv_state.new_zeros(
-                decode_conv_state.shape[0], decode_conv_state.shape[1], 1
-            )
-            conv_cache = torch.cat([zero_padding, decode_conv_state], dim=-1)
-            decode_conv_output, conv_cache = _fla_causal_conv1d_update(
-                mixed_qkv[:num_decode_tokens],
-                conv_cache,
-                weight=conv_weight,
-                bias=conv_bias,
+            decode_conv_output = _attn_gym_causal_conv1d_decode(
+                mixed_qkv[:num_decode_tokens].contiguous(),
+                conv_weight,
+                conv_state,
                 activation="silu",
+                state_indices=decode_slots,
             )
-            conv_state[decode_slots] = conv_cache[..., 1:].to(conv_state.dtype)
             conv_output[:num_decode_tokens] = decode_conv_output
 
         prefill_slots = None
@@ -316,6 +326,7 @@ def _forward(
             assert gdn_metadata.prefill_state_indices is not None
             prefill_slots = gdn_metadata.prefill_state_indices
             prefill_has_initial_state = gdn_metadata.prefill_has_initial_state
+            assert prefill_has_initial_state is not None
             prefill_start = num_decode_tokens if num_decodes > 0 else 0
             # cu_seqlens must be 0-based within the prefill slice that the conv
             # kernel receives (mixed_qkv[prefill_start:]).
@@ -340,85 +351,74 @@ def _forward(
                 prefill_has_initial_state.any()
             )
             conv_initial_state = mixed_qkv.new_zeros(
-                num_prefill_sequences, mixed_qkv.shape[1], self.conv_kernel_size
+                num_prefill_sequences,
+                self.conv_kernel_size - 1,
+                mixed_qkv.shape[1],
             )
             # Fresh prefills keep zero state; prefix-cache continuations restore
             # only the sequence slots identified by vLLM metadata.
             if has_continuations:
                 resumed_slots = prefill_slots[prefill_has_initial_state]
-                conv_initial_state[prefill_has_initial_state, :, 1:] = conv_state[
+                conv_initial_state[prefill_has_initial_state] = conv_state[
                     resumed_slots
                 ]
-            prefill_conv_output, conv_final_state = _fla_causal_conv1d(
+            prefill_conv_output, conv_final_state = _attn_gym_causal_conv1d(
                 mixed_qkv[prefill_start:num_actual_tokens].unsqueeze(0),
-                weight=conv_weight,
-                bias=conv_bias,
+                conv_weight,
                 activation="silu",
                 cu_seqlens=prefill_cu_seqlens,
                 initial_state=conv_initial_state,
-                output_final_state=True,
+                return_final_state=True,
             )
-            conv_state[prefill_slots] = conv_final_state[..., 1:].to(conv_state.dtype)
+            conv_state[prefill_slots] = conv_final_state.to(conv_state.dtype)
             conv_output[prefill_start:num_actual_tokens] = prefill_conv_output.squeeze(
                 0
             )
 
-        # Recurrence over the whole batch in one call. Decode (T=1) and prefill
-        # (T>1) sequences run together, delimited by cu_seqlens; each gets a
-        # materialized fp32 initial state. The kernel processes each sequence
-        # independently from its cu_seqlens entry and initial-state row, so a
-        # single call covers the whole batch (chunk or recurrent, chosen below).
+        # Recurrence over the whole batch in one call. The batch-invariant path
+        # addresses the paged SSM pool directly; ordinary prefill uses FLA.
         cu_seqlens = gdn_metadata.non_spec_query_start_loc[: num_sequences + 1]
         if num_prefills == 0:
-            # Pure decode is captured in a CUDA graph, which forbids host syncs and
-            # data-dependent (boolean-mask) indexing. Every decode sequence resumes,
-            # so gather its paged SSM state directly with the integer slot indices.
             all_slots = decode_slots
-            initial_state = (
-                ssm_state[all_slots].transpose(-1, -2).to(torch.float32).contiguous()
-            )
+            has_initial_state = None
         else:
-            # Prefill / mixed runs eager at the graph break, so boolean-mask gather
-            # and the host sync it implies are allowed. A sequence resumes from
-            # paged state iff it is a decode (always) or a prefill prefix-cache
-            # continuation; fresh prefills keep a zero initial state.
             all_slots = (
                 torch.cat([decode_slots, prefill_slots])
                 if num_decodes > 0
                 else prefill_slots
             )
-            resumes_from_cache = torch.zeros(
+            has_initial_state = torch.ones(
                 num_sequences, dtype=torch.bool, device=mixed_qkv.device
             )
-            resumes_from_cache[:num_decodes] = True
             if prefill_has_initial_state is not None:
-                resumes_from_cache[num_decodes:] = prefill_has_initial_state
-            initial_state = conv_output.new_zeros(
-                num_sequences,
-                self.local_num_v_heads,
-                self.head_k_dim,
-                self.head_v_dim,
-                dtype=torch.float32,
-            )
-            initial_state[resumes_from_cache] = (
-                ssm_state[all_slots[resumes_from_cache]]
-                .transpose(-1, -2)
-                .to(torch.float32)
+                has_initial_state[num_decodes:] = prefill_has_initial_state
+
+        batch_invariant = is_in_batch_invariant_mode()
+        if num_prefills == 0 and not batch_invariant:
+            _attn_gym_recurrent_gdn_decode(
+                conv_output[:num_decode_tokens],
+                a[:num_decode_tokens].unsqueeze(0),
+                b[:num_decode_tokens].unsqueeze(0),
+                A_log.float(),
+                dt_bias,
+                ssm_state,
+                all_slots,
+                scale=self.head_k_dim**-0.5,
+                out=output[:num_decode_tokens].unsqueeze(0),
             )
-        # Non-batch-invariant prefill/mixed batches take the parallel chunk kernel;
-        # batch-invariant mode and pure decode take the recurrent kernel.
-        use_chunk = num_prefills > 0 and not is_in_batch_invariant_mode()
+            return
+
         recurrent_output = self._run_recurrence(
             conv_output,
             a,
             b,
             negative_exp_A,
             dt_bias,
-            initial_state,
             ssm_state,
             all_slots,
             cu_seqlens,
-            use_chunk,
+            has_initial_state,
+            batch_invariant=batch_invariant,
         )
         output[:num_actual_tokens] = recurrent_output[0, :num_actual_tokens].to(
             output.dtype
diff --git a/torchtitan/experiments/rl/tests/test_bitwise_parity.py b/torchtitan/experiments/rl/tests/test_bitwise_parity.py
index 2c404c6dd0..0e7630b8f8 100644
--- a/torchtitan/experiments/rl/tests/test_bitwise_parity.py
+++ b/torchtitan/experiments/rl/tests/test_bitwise_parity.py
@@ -632,16 +632,17 @@ def setUpClass(cls):
         if hf_path:
             config.hf_assets_path = hf_path
 
-        from torchtitan.tools.utils import has_cuda_capability
+        from torchtitan.tools.utils import get_cuda_flash_attention_impl
 
-        if has_cuda_capability(9, 0):
+        flash_attention_impl = get_cuda_flash_attention_impl()
+        if flash_attention_impl is not None:
             from torch.nn.attention import (
                 activate_flash_attention_impl,
                 current_flash_attention_impl,
             )
 
-            if current_flash_attention_impl() != "FA3":
-                activate_flash_attention_impl("FA3")
+            if current_flash_attention_impl() != flash_attention_impl:
+                activate_flash_attention_impl(flash_attention_impl)
 
         # Enable batch-invariant mode BEFORE init_distributed
         set_batch_invariance(config.trainer.debug.batch_invariant)
diff --git a/torchtitan/models/kimi_k3/kda.py b/torchtitan/models/kimi_k3/kda.py
index ac53adc104..5148bad76f 100644
--- a/torchtitan/models/kimi_k3/kda.py
+++ b/torchtitan/models/kimi_k3/kda.py
@@ -12,7 +12,7 @@
 import torch.nn.functional as F
 from attn_gym.linear.kda import bound_gate, chunk_kda
 from attn_gym.linear.kda.fwd.triton.l2norm_fwd import l2norm
-from attn_gym.linear.kda.short_conv import causal_conv1d
+from attn_gym.linear.short_conv import causal_conv1d
 from torch import nn
 
 from torchtitan.models.common.attention import AttentionMasksType, VarlenMetadata
@@ -91,7 +91,6 @@ def forward(
             A_log_N.float(),
             dt_bias_NK.float(),
             lower_bound=self.lower_bound,
-            impl="fused",
         )
         output_BTNV, _ = chunk_kda(
             l2norm(q_BTNK),
diff --git a/torchtitan/models/qwen3_5/gdn.py b/torchtitan/models/qwen3_5/gdn.py
index 5831629242..91d21f0b75 100644
--- a/torchtitan/models/qwen3_5/gdn.py
+++ b/torchtitan/models/qwen3_5/gdn.py
@@ -6,13 +6,20 @@
 
 """Gated DeltaNet modules for Qwen3.5."""
 
+# Tensor dimensions: B = batch, T = tokens, N = heads, K = key dimension,
+# V = value dimension, S = state slots, D = channels.
+
 from dataclasses import dataclass
 from typing import Literal
 
 import spmd_types as spmd
 import torch
 import torch.nn.functional as F
-from fla.modules.conv.triton.ops import CausalConv1dFunction
+from attn_gym.linear import (
+    causal_conv1d as _attn_gym_causal_conv1d,
+    l2norm as _attn_gym_l2norm,
+    recurrent_gdn as _attn_gym_recurrent_gdn,
+)
 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,
@@ -30,7 +37,6 @@
 
 spmd.register_local_autograd_function(ChunkGatedDeltaRuleFunction)
 spmd.register_local_autograd_function(FusedRecurrentFunction)
-spmd.register_local_autograd_function(CausalConv1dFunction)
 
 
 @spmd.local_map(
@@ -48,28 +54,19 @@ def _causal_conv1d_varlen(
     cu_seqlens: torch.Tensor,
     cu_seqlens_cpu: torch.Tensor | None,
 ) -> torch.Tensor:
-    """FLA depthwise causal conv with per-document resets (CUDA-only).
+    """Depthwise causal conv with per-document resets (CUDA-only).
 
     A pure-torch per-document reference lives in
     ``tests/unit_tests/gpu/test_qwen3_5_deltanet.py``.
     """
-    if cu_seqlens_cpu is None:
-        raise ValueError(
-            "Qwen3.5 FLA varlen conv requires a CPU cu_seqlens tensor. "
-            "Build VarlenMetadata with include_host_offsets=True."
-        )
-
-    from fla.modules.conv.causal_conv1d import causal_conv1d as _fla_causal_conv1d
-
-    out_BTD, _ = _fla_causal_conv1d(
-        x=x_TD.unsqueeze(0),
-        weight=weight.squeeze(1),
-        bias=None,
+    del cu_seqlens_cpu
+    out_BTD = _attn_gym_causal_conv1d(
+        x_TD.unsqueeze(0),
+        weight.squeeze(1),
         activation="silu",
-        backend="triton",
         cu_seqlens=cu_seqlens,
-        cu_seqlens_cpu=cu_seqlens_cpu,
     )
+    assert isinstance(out_BTD, torch.Tensor)
     return out_BTD.squeeze(0)
 
 
@@ -104,9 +101,9 @@ def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
     "torchtitan::recurrent_gdn_fwd", mutates_args=(), device_types="cuda"
 )
 def _recurrent_gdn_fwd(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
+    q_BTNK: torch.Tensor,
+    k_BTNK: torch.Tensor,
+    v_BTNV: torch.Tensor,
     g: torch.Tensor,
     beta: torch.Tensor,
     cu_seqlens: torch.Tensor,
@@ -114,43 +111,62 @@ def _recurrent_gdn_fwd(
 ) -> torch.Tensor:
     """Run the batch-invariant GDN recurrent forward kernel.
 
-    The vLLM generator must use the recurrent kernel for per-token decode. The
-    trainer uses the same kernel with a materialized float32 initial state and
-    varlen metadata so its forward is bitwise identical to generation.
+    The vLLM generator uses Attention Gym's paging-aware recurrent kernel for
+    per-token decode. The trainer uses the same recurrence with a materialized
+    float32 initial state and varlen metadata so its forward is bitwise identical
+    to generation.
     """
     num_sequences = int(cu_seqlens.numel()) - 1
-    initial_state = q.new_zeros(
-        num_sequences,
-        q.shape[2],
-        q.shape[3],
-        v.shape[3],
+    # state_cache_SNVK: [num_sequences + 1, N, V, K].
+    state_cache_SNVK = q_BTNK.new_empty(
+        num_sequences + 1,
+        q_BTNK.shape[2],
+        v_BTNV.shape[3],
+        q_BTNK.shape[3],
         dtype=torch.float32,
     )
-    output, _ = _fla_fused_recurrent_gated_delta_rule(
-        q,
-        k,
-        v,
+    state_indices = torch.arange(
+        1,
+        num_sequences + 1,
+        dtype=torch.int32,
+        device=q_BTNK.device,
+    )
+    has_initial_state = torch.zeros(
+        num_sequences,
+        dtype=torch.bool,
+        device=q_BTNK.device,
+    )
+    # FLA normalizes Q/K inside its kernels. Attention Gym's recurrent GDN
+    # expects normalized inputs, so apply the same normalization explicitly.
+    normalized_q_BTNK = _attn_gym_l2norm(q_BTNK, cu_seqlens=cu_seqlens)
+    normalized_k_BTNK = _attn_gym_l2norm(k_BTNK, cu_seqlens=cu_seqlens)
+    out_BTNV, _ = _attn_gym_recurrent_gdn(
+        normalized_q_BTNK,
+        normalized_k_BTNK,
+        v_BTNV,
         g,
-        beta=beta,
-        initial_state=initial_state,
-        output_final_state=True,
-        use_qk_l2norm_in_kernel=True,
+        beta,
+        state_cache_SNVK,
         cu_seqlens=cu_seqlens,
+        scale=q_BTNK.shape[-1] ** -0.5,
+        state_indices=state_indices,
+        has_initial_state=has_initial_state,
+        autotune=False,
     )
-    return output.to(q.dtype)
+    return out_BTNV.to(q_BTNK.dtype)
 
 
 @_recurrent_gdn_fwd.register_fake
 def _recurrent_gdn_fwd_fake(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
+    q_BTNK: torch.Tensor,
+    k_BTNK: torch.Tensor,
+    v_BTNV: torch.Tensor,
     g: torch.Tensor,
     beta: torch.Tensor,
     cu_seqlens: torch.Tensor,
     cu_seqlens_cpu: torch.Tensor,
 ) -> torch.Tensor:
-    return torch.empty_like(v, dtype=q.dtype)
+    return torch.empty_like(v_BTNV, dtype=q_BTNK.dtype)
 
 
 @torch.library.custom_op(