diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py
index a915f7f85d..234f494bb8 100644
--- a/src/maxtext/layers/decoders.py
+++ b/src/maxtext/layers/decoders.py
@@ -411,7 +411,12 @@ def get_remat_policy(self):
             "kv_proj",
             "qkv_proj",
         )
-      elif cfg.remat_policy == "qkv_proj_offloaded":
+      elif cfg.remat_policy in ("qkv_proj_offloaded", "minimal_offloaded", "custom"):
+        # minimal_offloaded offloads all except context. All three share a single
+        # source of truth for their save/offload name lists (see
+        # maxtext_utils.get_save_and_offload_names) so that offloading configured via
+        # `custom` resolves identically to the named presets.
+        save_names, offload_names = maxtext_utils.get_save_and_offload_names(cfg)
         policy = jax.checkpoint_policies.save_and_offload_only_these_names(
             names_which_can_be_saved=[],
             names_which_can_be_offloaded=["query_proj", "value_proj", "key_proj", "kv_proj"],
@@ -1433,7 +1438,12 @@ def _apply_gemma4_scanned_blocks(
     if num_full_blocks > 0:
       ScannableBlockToLinen = gemma4.Gemma4ScannableBlockToLinen
       policy = self.get_remat_policy()
-      RemattedGemma4Block = self.set_remat_policy([ScannableBlockToLinen], policy)[0]
+      # Gemma4ScannableBlock rematerializes its own local (scanned) and global
+      # layers when apply_internal_remat=True, so we do NOT wrap it in
+      # block-level remat here (that would double-rematerialize and make XLA
+      # treat the whole block as one unit). Unrolling the block scan lets XLA
+      # free each block's activations across iterations instead of keeping the
+      # block live as a unit.
 
       kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan(
           kv_caches, num_full_blocks, block_pattern_len, stack=True
@@ -1456,7 +1466,7 @@ def _apply_gemma4_scanned_blocks(
 
       # For a fully scanned block, apply it inside a nn.scan over the calculated number of full blocks
       y, returned_kv_cache = nn.scan(
-          RemattedGemma4Block,
+          ScannableBlockToLinen,
           variable_axes={
               "params": cfg.param_scan_axis,
               "cache": 0,
@@ -1467,6 +1477,7 @@ def _apply_gemma4_scanned_blocks(
           split_rngs={"params": True, "dropout": cfg.enable_dropout},
           in_axes=in_axes_tuple,
           length=num_full_blocks,
+          unroll=num_full_blocks,
           metadata_params={
               nn.PARTITION_NAME: "layers",
               "abstract_init": False,
@@ -1477,6 +1488,8 @@ def _apply_gemma4_scanned_blocks(
           quant=self.quant,
           model_mode=model_mode,
           num_of_layers=block_pattern_len,
+          remat_policy_fn=policy,
+          apply_internal_remat=True,
           name="scanned_blocks",
       )(
           y, *broadcast_args
diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py
index 528523e438..8f9525feac 100644
--- a/src/maxtext/layers/nnx_decoders.py
+++ b/src/maxtext/layers/nnx_decoders.py
@@ -37,7 +37,7 @@
     ShardMode,
 )
 from maxtext.layers import initializers, linears, mhc, normalizations, quantizations
-from maxtext.layers import nnx_wrappers
+from maxtext.layers import nnx_scan, nnx_wrappers
 from maxtext.layers.attentions import Attention
 from maxtext.layers.embeddings import Embed, PositionalEmbedding, attend_on_embedding
 from maxtext.layers.normalizations import RMSNorm
@@ -646,9 +646,19 @@ def _init_scanned_gemma4(self, decoder_block_classes, rngs, mesh):
     attention_pattern_length = len(gemma4.GEMMA4_ATTENTION_PATTERN)
     scan_length = config.num_decoder_layers // attention_pattern_length
     num_remaining_layers = config.num_decoder_layers % attention_pattern_length
-    layer_kwargs = {"num_of_layers": attention_pattern_length}
-
-    rem_layer_kwargs = {"num_of_layers": num_remaining_layers}
+    policy = self.get_remat_policy()
+    # The pure-NNX decoder skips block-level remat (skip_block_remat=True below),
+    # so the block rematerializes its own local/global layers instead.
+    layer_kwargs = {
+        "num_of_layers": attention_pattern_length,
+        "remat_policy_fn": policy,
+        "apply_internal_remat": True,
+    }
+    rem_layer_kwargs = {
+        "num_of_layers": num_remaining_layers,
+        "remat_policy_fn": policy,
+        "apply_internal_remat": True,
+    }
 
     RemattedGemma4Block = gemma4.Gemma4ScannableBlock
 
@@ -839,94 +849,20 @@ def _create_scanned_layers(
       **layer_kwargs,
   ):
     """Creates a scanned stack of layers using jax.lax.scan for memory-efficient initialization."""
-    if length == 0:
-      return None
-    scan_axis = self.config.param_scan_axis
-
-    # Fork rngs to get per-layer RNG states for scanning
-    try:
-      forked_rngs = rngs.fork(split=length)
-    except:  # pylint: disable=bare-except
-      pass
-
-    rngs_graphdef, rngs_state = nnx.split(forked_rngs)
-
-    first_rng_state = jax.tree.map(lambda x: x[0], rngs_state)
-    ref_rngs = nnx.merge(rngs_graphdef, first_rng_state)
-    ref_layer = decoder_layer_class(
-        config=self.config,
-        mesh=self.mesh,
-        quant=self.quant,
-        model_mode=self.model_mode,
-        rngs=ref_rngs,
-        **layer_kwargs,
+    return nnx_scan.create_scanned_layers(
+        lambda layer_rngs: decoder_layer_class(
+            config=self.config,
+            mesh=self.mesh,
+            model_mode=self.model_mode,
+            quant=self.quant,
+            rngs=layer_rngs,
+            **layer_kwargs,
+        ),
+        length=length,
+        param_scan_axis=self.config.param_scan_axis,
+        metadata_axis_name=metadata_axis_name,
+        rngs=rngs,
     )
-    layer_graphdef, _, _ = nnx.split(ref_layer, nnx.Param, ...)
-    del ref_layer
-
-    def scan_body(carry, rng_state_slice):
-      layer_rngs = nnx.merge(rngs_graphdef, rng_state_slice)
-      layer = decoder_layer_class(
-          config=self.config,
-          mesh=self.mesh,
-          quant=self.quant,
-          model_mode=self.model_mode,
-          rngs=layer_rngs,
-          **layer_kwargs,
-      )
-      _, params, rest = nnx.split(layer, nnx.Param, ...)
-      return carry, (params, rest)
-
-    _, (stacked_params, stacked_rest) = jax.lax.scan(scan_body, None, rngs_state)
-
-    if scan_axis != 0:
-      stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0), stacked_params)
-
-    def _add_scan_metadata(state, axis):
-      def _update_leaf(leaf):
-        if hasattr(leaf, "replace") and hasattr(leaf, "value"):
-          replace_kwargs = {}
-          if hasattr(leaf, "get_metadata"):
-            replace_kwargs.update(leaf.get_metadata())
-
-          replace_kwargs[nnx.PARTITION_NAME] = metadata_axis_name
-          replace_kwargs["param_scan_axis"] = axis
-
-          for key in [
-              "sharding",
-              "out_sharding",
-              "kernel_axes",
-              "sharding_names",
-          ]:
-            val = getattr(leaf, key, None)
-            if val is None and key in replace_kwargs:
-              val = replace_kwargs[key]
-
-            if val is not None:
-              if isinstance(val, str):
-                val = (val,)
-              if isinstance(val, tuple):
-                l = list(val)
-                # Safely insert the scan axis into the logical axes string
-                if metadata_axis_name not in l:
-                  insert_idx = min(axis, len(l))
-                  l.insert(insert_idx, metadata_axis_name)
-                  replace_kwargs[key] = tuple(l)
-
-          return leaf.replace(**replace_kwargs)
-        return leaf
-
-      # We must use a custom is_leaf to catch the VariableState instances
-      return jax.tree.map(
-          _update_leaf,
-          state,
-          is_leaf=lambda x: hasattr(x, "replace") and hasattr(x, "value"),
-      )
-
-    stacked_params = _add_scan_metadata(stacked_params, scan_axis)
-    stacked_rest = _add_scan_metadata(stacked_rest, 0)
-
-    return nnx.merge(layer_graphdef, stacked_params, stacked_rest)
 
   def _apply_layer_with_remat(self, layer: nnx.Module, y: jax.Array, policy: Any, prevent_cse: bool, **kwargs):
     """Helper to cleanly apply jax.checkpoint to a single unscanned layer or block."""
@@ -944,7 +880,17 @@ def pure_layer_fn(state_in, y_in):
 
     return out
 
-  def _apply_layers_sequentially(self, layers, x_in, *args, length: int, kv_caches_stacked=None, **kwargs):
+  def _apply_layers_sequentially(
+      self,
+      layers,
+      x_in,
+      *args,
+      length: int,
+      kv_caches_stacked=None,
+      skip_block_remat: bool = False,
+      unroll: int = 1,
+      **kwargs,
+  ):
     """Runs the layer stack using nnx.scan.
 
     Args:
@@ -955,6 +901,11 @@ def _apply_layers_sequentially(self, layers, x_in, *args, length: int, kv_caches
       kv_caches_stacked: Optional pytree whose leaves have shape [num_layers, ...].
         When provided, the i-th slice is passed as `kv_cache=` to layer i and the
         updated caches are returned as a third element of the tuple.
+      skip_block_remat: When True, do not wrap the scanned body in jax.checkpoint.
+        Used when the scanned module already applies its own (finer-grained,
+        e.g. per-layer) remat internally, to avoid double rematerialization.
+      unroll: Number of scan iterations to unroll into straight-line code
+        (forwarded to jax.lax.scan). unroll >= length fully unrolls the loop.
       **kwargs: Keyword args forwarded to the layer (filtered by the layer signature).
 
     Returns:
@@ -1037,7 +988,12 @@ def layer_fn(carry, scanned_vars):
         return new_carry, (new_current_state, updated_kv)
       return new_carry, new_current_state
 
-    layer_fn_wrapped = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
+    if skip_block_remat:
+      # The scanned module applies its own remat internally; wrapping the whole
+      # body again would double-remat and recompute the entire block.
+      layer_fn_wrapped = layer_fn
+    else:
+      layer_fn_wrapped = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
 
     if use_kv:
       # If kv_caches is provided (e.g., from vLLM), we CANNOT use jax.lax.scan
@@ -1069,7 +1025,7 @@ def layer_fn(carry, scanned_vars):
       params = maxtext_utils_nnx.nnx_ensure_scan_leading_axis(params, length)
       state = maxtext_utils_nnx.nnx_ensure_scan_leading_axis(state, length)
 
-      final_carry, scanned_state = jax.lax.scan(layer_fn_wrapped, x_in, (params, state))
+      final_carry, scanned_state = jax.lax.scan(layer_fn_wrapped, x_in, (params, state), unroll=unroll)
       returned_kv_stacked = None
 
       # Ensure metadata rank matches the stacked values
@@ -2032,13 +1988,25 @@ def _apply_gemma4_scanned_blocks(
     if attention_metadata is not None:
       layer_kwargs["attention_metadata"] = attention_metadata
 
-    # Apply the main scan over the full blocks
+    # Apply the main scan over the full blocks. Gemma4ScannableBlock applies
+    # per-layer remat internally (local scan + global layer), so skip the
+    # block-level remat here to avoid double rematerialization. Unrolling the
+    # block loop (one iteration per repeated block) lets XLA pipeline/free block
+    # activations across iterations (memory + overlap knob).
+    block_unroll = max(1, scan_length)
     if scan_length > 0:
       grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan(
           kv_caches, scan_length, attention_pattern_length, stack=False
       )
       y, self.scanned_blocks, _ = self._apply_layers_sequentially(
-          self.scanned_blocks, y, *layer_args, length=scan_length, kv_caches_stacked=grouped_kv_caches, **layer_kwargs
+          self.scanned_blocks,
+          y,
+          *layer_args,
+          length=scan_length,
+          kv_caches_stacked=grouped_kv_caches,
+          skip_block_remat=True,
+          unroll=block_unroll,
+          **layer_kwargs,
       )
       maxtext_utils.update_kv_caches_after_scan(
           kv_caches, grouped_kv_caches, scan_length, attention_pattern_length, stacked=False
diff --git a/src/maxtext/layers/nnx_scan.py b/src/maxtext/layers/nnx_scan.py
new file mode 100644
index 0000000000..b675c65357
--- /dev/null
+++ b/src/maxtext/layers/nnx_scan.py
@@ -0,0 +1,137 @@
+# Copyright 2023–2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for constructing and applying stacks of scanned NNX layers."""
+
+from collections.abc import Callable
+from typing import Any
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+
+
+def create_scanned_layers(
+    layer_factory: Callable[[nnx.Rngs], nnx.Module],
+    *,
+    length: int,
+    param_scan_axis: int,
+    metadata_axis_name: str,
+    rngs: nnx.Rngs,
+) -> nnx.Module | None:
+  """Constructs an NNX layer whose variables are stacked for a layer scan."""
+  if length == 0:
+    return None
+
+  forked_rngs = rngs.fork(split=length)
+  rngs_graphdef, rngs_state = nnx.split(forked_rngs)
+
+  first_rng_state = jax.tree.map(lambda x: x[0], rngs_state)
+  reference_layer = layer_factory(nnx.merge(rngs_graphdef, first_rng_state))
+  layer_graphdef, _, _ = nnx.split(reference_layer, nnx.Param, ...)
+  del reference_layer
+
+  def scan_body(carry, rng_state_slice):
+    layer = layer_factory(nnx.merge(rngs_graphdef, rng_state_slice))
+    _, params, rest = nnx.split(layer, nnx.Param, ...)
+    return carry, (params, rest)
+
+  _, (stacked_params, stacked_rest) = jax.lax.scan(scan_body, None, rngs_state)
+
+  if param_scan_axis != 0:
+    stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, param_scan_axis), stacked_params)
+
+  def add_scan_metadata(state, axis):
+    def update_leaf(leaf):
+      if hasattr(leaf, "replace") and hasattr(leaf, "value"):
+        replace_kwargs = {}
+        if hasattr(leaf, "get_metadata"):
+          replace_kwargs.update(leaf.get_metadata())
+
+        replace_kwargs[nnx.PARTITION_NAME] = metadata_axis_name
+        replace_kwargs["param_scan_axis"] = axis
+
+        for key in ["sharding", "out_sharding", "kernel_axes", "sharding_names"]:
+          value = getattr(leaf, key, None)
+          if value is None and key in replace_kwargs:
+            value = replace_kwargs[key]
+
+          if value is not None:
+            if isinstance(value, str):
+              value = (value,)
+            if isinstance(value, tuple):
+              logical_axes = list(value)
+              if metadata_axis_name not in logical_axes:
+                logical_axes.insert(min(axis, len(logical_axes)), metadata_axis_name)
+                replace_kwargs[key] = tuple(logical_axes)
+
+        return leaf.replace(**replace_kwargs)
+      return leaf
+
+    return jax.tree.map(
+        update_leaf,
+        state,
+        is_leaf=lambda x: hasattr(x, "replace") and hasattr(x, "value"),
+    )
+
+  stacked_params = add_scan_metadata(stacked_params, param_scan_axis)
+  stacked_rest = add_scan_metadata(stacked_rest, 0)
+  return nnx.merge(layer_graphdef, stacked_params, stacked_rest)
+
+
+def apply_scanned_layers(
+    layers: nnx.Module,
+    carry: Any,
+    *,
+    length: int,
+    param_scan_axis: int,
+    apply_fn: Callable[[nnx.Module, Any], Any],
+    remat: bool = False,
+    remat_policy: Callable[..., Any] | None = None,
+    prevent_cse: bool = True,
+    unroll: int = 1,
+) -> Any:
+  """Applies stacked NNX layers using ``jax.lax.scan``.
+
+  This helper owns the generic NNX state and scan-axis mechanics. ``apply_fn``
+  defines the model-specific module invocation and must return the next carry.
+  ``remat`` is separate from ``remat_policy`` because ``None`` is JAX's full
+  rematerialization policy, not an indication that rematerialization is off.
+
+  Externally managed per-layer state, such as KV caches, is not supported by
+  this scan path.
+  """
+  if length <= 0:
+    return carry
+
+  layer_graphdef, params, state = nnx.split(layers, nnx.Param, ...)
+  if param_scan_axis != 0:
+    params = jax.tree.map(lambda x: jnp.moveaxis(x, param_scan_axis, 0), params)
+
+  def scan_body(current_carry, scanned_state):
+    current_params, current_state = scanned_state
+    current_layer = nnx.merge(layer_graphdef, current_params, current_state)
+    next_carry = apply_fn(current_layer, current_carry)
+    return next_carry, nnx.state(current_layer)
+
+  scan_fn = jax.checkpoint(scan_body, policy=remat_policy, prevent_cse=prevent_cse) if remat else scan_body
+  final_carry, scanned_state = jax.lax.scan(scan_fn, carry, (params, state), unroll=unroll)
+
+  if param_scan_axis != 0:
+    scanned_params, scanned_other = scanned_state.split(nnx.Param, ...)
+    scanned_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, param_scan_axis), scanned_params)
+    scanned_state = nnx.State.merge(scanned_params, scanned_other)
+
+  nnx.update(layers, scanned_state)
+  return final_carry
diff --git a/src/maxtext/models/gemma4.py b/src/maxtext/models/gemma4.py
index dbf8efc4de..4c1fe98bdb 100644
--- a/src/maxtext/models/gemma4.py
+++ b/src/maxtext/models/gemma4.py
@@ -15,18 +15,19 @@
 """Specialized layers for Gemma 4."""
 
 import jax
+from jax.experimental import xla_metadata
 from jax.ad_checkpoint import checkpoint_name
 from jax.sharding import Mesh
 import jax.numpy as jnp
 
 from flax import linen as nn
 from flax import nnx
-from typing import Optional
+from typing import Optional, Any
 
 from maxtext.common.common_types import Config, AttentionType, MODEL_MODE_PREFILL
 from maxtext.layers import initializers
 from maxtext.layers import moe
-from maxtext.layers import nnx_wrappers
+from maxtext.layers import nnx_scan, nnx_wrappers
 from maxtext.layers import quantizations
 from maxtext.layers.attentions import Attention
 from maxtext.layers.linears import MlpBlock
@@ -35,6 +36,7 @@
 from maxtext.layers.normalizations import RMSNorm
 from maxtext.layers.quantizations import AqtQuantization as Quant
 from maxtext.utils import max_utils
+from maxtext.utils import maxtext_utils
 
 
 GEMMA4_ATTENTION_PATTERN = (
@@ -409,7 +411,7 @@ def update_cache(cache, val):
 
 
 class Gemma4ScannableBlock(nnx.Module):
-  """A repeatable block of Gemma4 decoder layers."""
+  """A repeatable block of Gemma4 decoder layers, scanning local layers."""
 
   def __init__(
       self,
@@ -418,7 +420,9 @@ def __init__(
       model_mode: str,
       rngs: nnx.Rngs,
       quant: None | Quant = None,
-      num_of_layers: int = 1,
+      num_of_layers: int = 6,
+      remat_policy_fn: Any = None,
+      apply_internal_remat: bool = False,
   ):
     """Initializes the instance.
 
@@ -429,6 +433,14 @@ def __init__(
       rngs: The random number generators for initialization.
       quant: The quantization configuration.
       num_of_layers: The number of layers in the model.
+      remat_policy_fn: The resolved rematerialization policy function.
+      apply_internal_remat: When True, the block rematerializes its own local
+        (scanned) and global layers, and the caller must NOT also apply
+        block-level remat (that would double-rematerialize and make XLA treat the
+        whole block as one unit). Both the pure-NNX and linen decoders set this
+        and skip block-level remat, so remat happens per layer rather than over
+        the whole block. When False, the block does not self-remat and relies on
+        the caller's block-level remat instead.
     """
     self.config = config
     self.mesh = mesh
@@ -436,20 +448,95 @@ def __init__(
     self.quant = quant
     self.rngs = rngs
     self.num_of_layers = num_of_layers
+    self.remat_policy_fn = remat_policy_fn
+    self.apply_internal_remat = apply_internal_remat
+
+    pattern_length = len(GEMMA4_ATTENTION_PATTERN)
+    if not 0 <= num_of_layers <= pattern_length:
+      raise ValueError(f"Gemma4ScannableBlock must contain between 0 and {pattern_length} layers; got {num_of_layers}.")
+
+    # Pattern is 5 local, 1 global.
+    self.num_local = min(5, num_of_layers)
+    self.num_global = max(0, num_of_layers - 5)
+
+    if self.num_local > 0:
+      self.local_layers = nnx_scan.create_scanned_layers(
+          lambda layer_rngs: Gemma4DecoderLayer(
+              config=self.config,
+              mesh=self.mesh,
+              model_mode=self.model_mode,
+              quant=self.quant,
+              rngs=layer_rngs,
+              attention_type=AttentionType.LOCAL_SLIDING,
+              layer_idx=0,  # layer_idx is not used in the class
+          ),
+          length=self.num_local,
+          param_scan_axis=self.config.param_scan_axis,
+          metadata_axis_name="local_layers",
+          rngs=self.rngs,
+      )
+    else:
+      self.local_layers = None
 
-    for layer_id in range(self.num_of_layers):
-      attention_type = get_attention_type(layer_id)
-      layer_name = f"layers_{layer_id}"
-      layer = Gemma4DecoderLayer(
+    if self.num_global > 0:
+      self.global_layer = Gemma4DecoderLayer(
           config=self.config,
           mesh=self.mesh,
           model_mode=self.model_mode,
           rngs=self.rngs,
           quant=self.quant,
-          attention_type=attention_type,
-          layer_idx=layer_id,
+          attention_type=AttentionType.GLOBAL,
+          layer_idx=5,  # layer_idx is not used in the class
       )
-      setattr(self, layer_name, layer)
+    else:
+      self.global_layer = None
+
+  def _apply_local_layers(
+      self,
+      y,
+      decoder_segment_ids,
+      decoder_positions,
+      deterministic,
+      model_mode,
+      slot=None,
+      previous_chunk=None,
+      bidirectional_mask=None,
+      attention_metadata=None,
+  ):
+    """Applies the block's local-attention layers via a per-layer rematerialized scan."""
+
+    def apply_layer(layer, carry):
+      layer_out = layer(
+          carry,
+          decoder_segment_ids,
+          decoder_positions,
+          deterministic,
+          model_mode,
+          slot=slot,
+          previous_chunk=previous_chunk,
+          bidirectional_mask=bidirectional_mask,
+          attention_metadata=attention_metadata,
+      )
+      return layer_out[0] if isinstance(layer_out, tuple) else layer_out
+
+    apply_remat = self.apply_internal_remat and self.config.remat_policy != "none"
+    if apply_remat:
+      prevent_cse = maxtext_utils.should_prevent_cse_in_remat(self.config)
+      remat_policy = self.remat_policy_fn
+    else:
+      prevent_cse = True
+      remat_policy = None
+
+    return nnx_scan.apply_scanned_layers(
+        self.local_layers,
+        y,
+        length=self.num_local,
+        param_scan_axis=self.config.param_scan_axis,
+        apply_fn=apply_layer,
+        remat=apply_remat,
+        remat_policy=remat_policy,
+        prevent_cse=prevent_cse,
+    )
 
   def __call__(
       self,
@@ -471,22 +558,120 @@ def __call__(
     y = inputs
 
     updated_kvs = []
-    for layer_id in range(self.num_of_layers):
-      current_kv = kv_cache[layer_id] if kv_cache is not None else None
-      y, new_kv = getattr(self, f"layers_{layer_id}")(
+
+    if kv_cache is not None:
+      # External per-layer KV caches require statically slicing the scanned local layers.
+      if self.local_layers is not None:
+        graphdef, params, state = nnx.split(self.local_layers, nnx.Param, ...)
+        scan_axis = self.config.param_scan_axis
+        if scan_axis != 0:
+          params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0), params)
+        per_layer_states = []
+        for i in range(self.num_local):
+          current_params = jax.tree.map(lambda x, i=i: x[i], params)
+          current_state = jax.tree.map(lambda x, i=i: x[i], state)
+          layer = nnx.merge(graphdef, current_params, current_state)
+          y, new_kv = layer(
+              y,
+              decoder_segment_ids,
+              decoder_positions,
+              deterministic,
+              model_mode,
+              slot=slot,
+              previous_chunk=previous_chunk,
+              bidirectional_mask=bidirectional_mask,
+              kv_cache=kv_cache[i],
+              attention_metadata=attention_metadata,
+          )
+          updated_kvs.append(new_kv)
+          per_layer_states.append(nnx.state(layer))
+
+        stacked_state = jax.tree.map(lambda *xs: jnp.stack(xs), *per_layer_states)
+        if scan_axis != 0:
+          stacked_params, stacked_other = stacked_state.split(nnx.Param, ...)
+          stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), stacked_params)
+          stacked_state = nnx.State.merge(stacked_params, stacked_other)
+        nnx.update(self.local_layers, stacked_state)
+    elif self.local_layers is not None:
+      y = self._apply_local_layers(
           y,
           decoder_segment_ids,
           decoder_positions,
           deterministic,
           model_mode,
-          previous_chunk=previous_chunk,
           slot=slot,
+          previous_chunk=previous_chunk,
           bidirectional_mask=bidirectional_mask,
-          kv_cache=current_kv,
           attention_metadata=attention_metadata,
       )
+
+    if self.global_layer is not None:
       if kv_cache is not None:
+        y, new_kv = self.global_layer(
+            y,
+            decoder_segment_ids,
+            decoder_positions,
+            deterministic,
+            model_mode,
+            previous_chunk=previous_chunk,
+            slot=slot,
+            bidirectional_mask=bidirectional_mask,
+            kv_cache=kv_cache[self.num_local],
+            attention_metadata=attention_metadata,
+        )
         updated_kvs.append(new_kv)
+      else:
+        graphdef_g, state_g = nnx.split(self.global_layer)
+
+        def scan_global_layer(carry, _):
+          current_y, current_state = carry
+          layer = nnx.merge(graphdef_g, current_state)
+          layer_out = layer(
+              current_y,
+              decoder_segment_ids,
+              decoder_positions,
+              deterministic,
+              model_mode,
+              previous_chunk=previous_chunk,
+              slot=slot,
+              bidirectional_mask=bidirectional_mask,
+              attention_metadata=attention_metadata,
+          )
+          new_y = layer_out[0] if isinstance(layer_out, tuple) else layer_out
+          return (new_y, nnx.state(layer)), None
+
+        # Remat the global layer behind the trip-count-one while boundary below,
+        # which keeps only one layer's full-sequence-attention activations live;
+        # without it (blocks are unrolled) XLA co-schedules every block's backward
+        # working set and OOMs. Offloaded (pinned-host) residuals can't cross that
+        # boundary, so the global layer saves would-be-offloaded tensors on device
+        # instead; the local-layer scan (a real multi-iteration scan) still offloads.
+        global_remat_policy = self.remat_policy_fn
+        offload_names = maxtext_utils.get_save_and_offload_names(cfg)
+        if offload_names[0] or offload_names[1]:
+          save_names, offload_to_device = offload_names
+          global_remat_policy = jax.checkpoint_policies.save_only_these_names(*(save_names + offload_to_device))
+
+        if self.apply_internal_remat and self.config.remat_policy != "none":
+          prevent_cse = maxtext_utils.should_prevent_cse_in_remat(self.config)
+          scan_global_layer = jax.checkpoint(
+              scan_global_layer,
+              policy=global_remat_policy,
+              prevent_cse=prevent_cse,
+          )
+
+        # Carry state through the loop instead of returning a stacked [1, ...]
+        # scan result: slicing that result previously introduced a bitcast
+        # between device and pinned-host memory under offload remat.
+        with xla_metadata.set_xla_metadata(**{"skip-simplify-while-loops_trip-count-one": "true"}):
+          (y, global_state), _ = jax.lax.scan(
+              scan_global_layer,
+              (y, state_g),
+              xs=None,
+              length=1,
+          )
+
+        nnx.update(self.global_layer, global_state)
 
     if kv_cache is not None:
       return y, tuple(updated_kvs)
diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py
index 92a50dc887..4b89242455 100644
--- a/src/maxtext/utils/maxtext_utils.py
+++ b/src/maxtext/utils/maxtext_utils.py
@@ -195,6 +195,47 @@ def should_prevent_cse_in_remat(config):
   return True
 
 
+def get_save_and_offload_names(config) -> tuple[list[str], list[str]]:
+  """Returns the ``(save_names, offload_names)`` split for remat policies built via
+  ``jax.checkpoint_policies.save_and_offload_only_these_names``.
+
+  ``save_names`` are checkpointed tensors kept in device HBM; ``offload_names`` are moved to
+  pinned host. This is the single source of truth shared by ``Decoder.get_remat_policy`` (which
+  builds the save-and-offload policy) and by models that use custom ways to handle offload (
+  e.g. Gemma4's global-layer with scan). It also makes ``remat_policy=custom`` with tensors marked
+  ``offload`` resolve to the same name sets as the named presets.
+
+  Returns a ``(save_names, offload_names)`` tuple:
+    * ``custom``: ``(config.tensors_on_device, config.tensors_to_offload)`` -- the per-tensor
+      assignments. Either list may be empty: all tensors set to ``device`` gives an empty offload
+      list, all set to ``remat`` gives ``([], [])``.
+    * ``qkv_proj_offloaded`` / ``minimal_offloaded``: ``([], <hardcoded offload names>)`` -- presets
+      that only offload and save nothing on device.
+    * any other policy (``full``, ``minimal``, ``save_*``, ``none``, ...): ``([], [])`` -- these do
+      not use ``save_and_offload_only_these_names``, so they contribute no names to this split.
+      Note ``([], [])`` here means "no names for this split", not that the policy saves nothing
+      overall (e.g. ``save_out_proj`` still saves ``out_proj`` via ``save_only_these_names``).
+  """
+  if config.remat_policy == "qkv_proj_offloaded":
+    return [], ["query_proj", "value_proj", "key_proj", "kv_proj"]
+  if config.remat_policy == "minimal_offloaded":
+    return [], [
+        "query_proj",
+        "value_proj",
+        "key_proj",
+        "kv_proj",
+        "qkv_proj",
+        "out_proj",
+        "mlpwi_0",
+        "mlpwi_1",
+        "mlpwi",
+        "mlpwo",
+    ]
+  if config.remat_policy == "custom":
+    return list(config.tensors_on_device or []), list(config.tensors_to_offload or [])
+  return [], []
+
+
 def load_compiled(config, partial_train, state, execution_devices):
   """# Loading a serialized compiled train step function."""
 
diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py
index cce1da530d..16525cae33 100644
--- a/tests/unit/maxtext_utils_test.py
+++ b/tests/unit/maxtext_utils_test.py
@@ -1730,5 +1730,56 @@ def test_update_kv_caches_after_scan_invalid_type(self):
       maxtext_utils.update_kv_caches_after_scan(kv_caches_tuple, returned_kv_cache, scan_length=1, block_len=2)
 
 
+@pytest.mark.cpu_only
+class TestGetSaveAndOffloadNames(unittest.TestCase):
+  """Tests for maxtext_utils.get_save_and_offload_names (pure config logic, no device needed)."""
+
+  @staticmethod
+  def _cfg(remat_policy, tensors_on_device=None, tensors_to_offload=None):
+    return SimpleNamespace(
+        remat_policy=remat_policy,
+        tensors_on_device=tensors_on_device,
+        tensors_to_offload=tensors_to_offload,
+    )
+
+  def test_named_preset_matches_equivalent_custom(self):
+    """qkv_proj_offloaded's offload names resolve identically to an equivalent custom config.
+
+    A real custom config keeps decoder_layer_input on device by default, so its full tuple
+    differs from the preset by that (benign, boundary) save entry -- assert only the offload halves.
+    """
+    _, preset_offload = maxtext_utils.get_save_and_offload_names(self._cfg("qkv_proj_offloaded"))
+    _, custom_offload = maxtext_utils.get_save_and_offload_names(
+        self._cfg(
+            "custom",
+            tensors_on_device=["decoder_layer_input"],
+            tensors_to_offload=["query_proj", "value_proj", "key_proj", "kv_proj"],
+        )
+    )
+    self.assertEqual(custom_offload, preset_offload)
+
+  def test_kv_proj_retained_in_offload_presets(self):
+    """Regression guard: kv_proj must stay in the offload presets (it is a real checkpoint name)."""
+    _, qkv_offload = maxtext_utils.get_save_and_offload_names(self._cfg("qkv_proj_offloaded"))
+    _, minimal_offload = maxtext_utils.get_save_and_offload_names(self._cfg("minimal_offloaded"))
+    self.assertIn("kv_proj", qkv_offload)
+    self.assertIn("kv_proj", minimal_offload)
+
+  def test_custom_reads_config_lists(self):
+    save, offload = maxtext_utils.get_save_and_offload_names(
+        self._cfg("custom", tensors_on_device=["context"], tensors_to_offload=["out_proj"])
+    )
+    self.assertEqual(save, ["context"])
+    self.assertEqual(offload, ["out_proj"])
+
+  def test_custom_handles_none_lists(self):
+    self.assertEqual(maxtext_utils.get_save_and_offload_names(self._cfg("custom")), ([], []))
+
+  def test_non_offloading_policies_return_empty(self):
+    """Policies that don't use the save/offload split contribute no names to it."""
+    for policy in ("full", "minimal", "save_out_proj", "save_qkv_proj", "none"):
+      self.assertEqual(maxtext_utils.get_save_and_offload_names(self._cfg(policy)), ([], []))
+
+
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py
index 4585c9c3ee..873ef32992 100644
--- a/tests/unit/nnx_decoders_test.py
+++ b/tests/unit/nnx_decoders_test.py
@@ -22,7 +22,10 @@
 """
 
 import sys
+from types import SimpleNamespace
 import unittest
+from unittest import mock
+from unittest.mock import MagicMock, patch
 import pytest
 
 import jax
@@ -34,8 +37,10 @@
 
 from maxtext.common.common_types import (
     DECODING_ACTIVE_SEQUENCE_INDICATOR,
+    MODEL_MODE_AUTOREGRESSIVE,
     MODEL_MODE_PREFILL,
     MODEL_MODE_TRAIN,
+    AttentionType,
     DecoderBlockType,
     MultimodalInput,
 )
@@ -45,7 +50,7 @@
 from maxtext.layers.embeddings import Embed
 from maxtext.layers.nnx_decoders import NNXDecoder, NNXDecoderLayer, deepstack_process
 from maxtext.layers.normalizations import RMSNorm
-from maxtext.models import gemma4_small
+from maxtext.models import gemma4, gemma4_small
 from maxtext.models.gpt3 import Gpt3LayerNorm
 from maxtext.models.llama2 import LlamaDecoderLayer
 from maxtext.utils import maxtext_utils
@@ -714,6 +719,92 @@ def test_scan_layers(self):
   unittest.main()
 
 
+class _StatefulGemma4DecoderLayer(nnx.Module):
+  """Small stand-in that exposes cache ordering and mutable-state updates."""
+
+  def __init__(self, *, attention_type, **unused_kwargs):
+    self.increment = 10 if attention_type == AttentionType.GLOBAL else 1
+    self.call_count = nnx.Intermediate(jnp.array(0, dtype=jnp.int32))
+    self.received_attention_metadata = nnx.Intermediate(jnp.array(False))
+
+  def __call__(
+      self,
+      inputs,
+      *unused_args,
+      kv_cache=None,
+      attention_metadata=None,
+      **unused_kwargs,
+  ):
+    self.call_count.value += 1
+    self.received_attention_metadata.value = attention_metadata is not None
+    output = inputs + self.increment
+    if kv_cache is None:
+      return output
+    return output, kv_cache + self.increment
+
+
+class TestGemma4ScannableBlock(unittest.TestCase):
+  """Tests Gemma4's nested local/global decoder block behavior."""
+
+  def setUp(self):
+    super().setUp()
+    self.config = SimpleNamespace(
+        dtype=jnp.float32,
+        param_scan_axis=1,
+        remat_policy="none",
+        scan_layers=True,
+    )
+
+  def _make_block(self):
+    return gemma4.Gemma4ScannableBlock(
+        config=self.config,
+        mesh=None,
+        model_mode=MODEL_MODE_AUTOREGRESSIVE,
+        rngs=nnx.Rngs(0),
+    )
+
+  def test_updates_state_through_global_single_iteration_scan(self):
+    with mock.patch.object(gemma4, "Gemma4DecoderLayer", _StatefulGemma4DecoderLayer):
+      block = self._make_block()
+      output, updated_kvs = block(
+          jnp.zeros((1, 1, 1)),
+          decoder_segment_ids=None,
+          decoder_positions=None,
+          deterministic=True,
+          model_mode=MODEL_MODE_AUTOREGRESSIVE,
+      )
+
+    np.testing.assert_array_equal(output, jnp.full((1, 1, 1), 15))
+    self.assertIsNone(updated_kvs)
+    np.testing.assert_array_equal(block.local_layers.call_count.value, jnp.ones(5, dtype=jnp.int32))
+    np.testing.assert_array_equal(block.global_layer.call_count.value, 1)
+
+  def test_restores_local_state_and_preserves_kv_order(self):
+    attention_metadata = object()
+
+    with mock.patch.object(gemma4, "Gemma4DecoderLayer", _StatefulGemma4DecoderLayer):
+      block = self._make_block()
+      output, updated_kvs = block(
+          jnp.zeros((1, 1, 1)),
+          decoder_segment_ids=None,
+          decoder_positions=None,
+          deterministic=True,
+          model_mode=MODEL_MODE_AUTOREGRESSIVE,
+          kv_cache=tuple(jnp.array(i) for i in range(6)),
+          attention_metadata=attention_metadata,
+      )
+
+    np.testing.assert_array_equal(output, jnp.full((1, 1, 1), 15))
+    np.testing.assert_array_equal(jnp.stack(updated_kvs), jnp.array([1, 2, 3, 4, 5, 15]))
+    np.testing.assert_array_equal(block.local_layers.call_count.value, jnp.ones(5, dtype=jnp.int32))
+    np.testing.assert_array_equal(
+        block.local_layers.received_attention_metadata.value,
+        jnp.ones(5, dtype=jnp.bool_),
+    )
+    np.testing.assert_array_equal(block.global_layer.call_count.value, 1)
+    np.testing.assert_array_equal(block.global_layer.received_attention_metadata.value, True)
+
+
 class TestNNXDecoderDeepseekAndGemma4(unittest.TestCase):
   """Tests for Deepseek and Gemma4 specific decoder logic."""
 
diff --git a/tests/unit/nnx_scan_test.py b/tests/unit/nnx_scan_test.py
new file mode 100644
index 0000000000..fcbe975af6
--- /dev/null
+++ b/tests/unit/nnx_scan_test.py
@@ -0,0 +1,127 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for NNX scan utilities."""
+
+import unittest
+from unittest import mock
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from maxtext.layers import nnx_scan
+
+
+class _LinearLayer(nnx.Module):
+
+  def __init__(self, rngs: nnx.Rngs):
+    self.kernel = nnx.Param(jax.random.normal(rngs.params(), (2, 2)))
+
+  def __call__(self, inputs):
+    return inputs @ self.kernel.value
+
+
+@pytest.mark.cpu_only
+class TestCreateScannedLayers(unittest.TestCase):
+  """Tests for nnx_scan.create_scanned_layers."""
+
+  def test_create_stacks_params_at_param_scan_axis(self):
+    """Per-layer params are stacked along param_scan_axis."""
+    length = 3
+    for axis, expected_shape in ((0, (length, 2, 2)), (1, (2, length, 2))):
+      layers = nnx_scan.create_scanned_layers(
+          _LinearLayer,
+          length=length,
+          param_scan_axis=axis,
+          metadata_axis_name="layers",
+          rngs=nnx.Rngs(0),
+      )
+      self.assertEqual(layers.kernel.value.shape, expected_shape)
+
+  def test_create_zero_length_returns_none(self):
+    """A zero-length stack short-circuits to None."""
+    layers = nnx_scan.create_scanned_layers(
+        _LinearLayer,
+        length=0,
+        param_scan_axis=0,
+        metadata_axis_name="layers",
+        rngs=nnx.Rngs(0),
+    )
+    self.assertIsNone(layers)
+
+
+@pytest.mark.cpu_only
+class TestApplyScannedLayers(unittest.TestCase):
+  """Tests for nnx_scan.apply_scanned_layers."""
+
+  def test_nonzero_param_scan_axis_round_trip(self):
+    """The scan dimension is stored at param_scan_axis and restored for application."""
+    length = 3
+    layers = nnx_scan.create_scanned_layers(
+        _LinearLayer,
+        length=length,
+        param_scan_axis=1,
+        metadata_axis_name="layers",
+        rngs=nnx.Rngs(0),
+    )
+
+    self.assertEqual(layers.kernel.value.shape, (2, length, 2))
+
+    inputs = jnp.array([1.0, -1.0])
+    kernels = jnp.moveaxis(layers.kernel.value, 1, 0)
+    expected = inputs
+    for kernel in kernels:
+      expected = expected @ kernel
+
+    actual = nnx_scan.apply_scanned_layers(
+        layers,
+        inputs,
+        length=length,
+        param_scan_axis=1,
+        apply_fn=lambda layer, carry: layer(carry),
+    )
+
+    np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5)
+    self.assertEqual(layers.kernel.value.shape, (2, length, 2))
+
+  def test_full_remat_checkpoints_scan_body_with_none_policy(self):
+    """A None policy means full remat and must not disable jax.checkpoint."""
+    layers = nnx_scan.create_scanned_layers(
+        _LinearLayer,
+        length=2,
+        param_scan_axis=1,
+        metadata_axis_name="layers",
+        rngs=nnx.Rngs(0),
+    )
+
+    with mock.patch.object(nnx_scan.jax, "checkpoint", wraps=jax.checkpoint) as checkpoint:
+      nnx_scan.apply_scanned_layers(
+          layers,
+          jnp.array([1.0, -1.0]),
+          length=2,
+          param_scan_axis=1,
+          apply_fn=lambda layer, carry: layer(carry),
+          remat=True,
+          remat_policy=None,
+      )
+
+    checkpoint.assert_called_once()
+    self.assertIsNone(checkpoint.call_args.kwargs["policy"])
+
+
+if __name__ == "__main__":
+  unittest.main()