diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py
index 35f1ce38ca..45460513c2 100644
--- a/src/maxtext/checkpoint_conversion/to_maxtext.py
+++ b/src/maxtext/checkpoint_conversion/to_maxtext.py
@@ -346,41 +346,62 @@ def _build_multi_axis_stacked_tensor(
     hook_fns: Any,
     target_shape: tuple,
     config,
+    mt_key: str = "",
 ) -> np.ndarray:
-  """Builds a MaxText tensor by stacking HF weights along two axes (experts and layers).
+  """Builds a MaxText tensor by stacking HF weights along two axes.
 
-  This function handles the complex case for scanned MoE layers, producing a tensor
-  with the shape (num_experts, num_layers, ...).
+  Two cases share this helper:
+    * MoE expert stacking: outer=experts, inner=layers, placed at the LEADING two
+      axes -> shape (num_experts, num_layers, ...).
+    * Gemma4 nested block scan (``scanned_blocks-local_layers``): the block's local
+      layers are an inner scan nested inside the outer block scan, so the two
+      stacked axes live at ``(param_scan_axis, param_scan_axis + 1)`` rather than
+      the leading axes, e.g. (emb, blocks, local, ...).
 
   Args:
       hf_source_keys: A nested (2D) list of Hugging Face parameter names.
-                      Outer list iterates experts, inner list iterates layers.
+                      The outer list is the outer stack axis, the inner list the
+                      inner stack axis.
       tensor_getter_fn: A callable that takes a HF key and returns the tensor (as numpy array).
       hook_fns: The hook function(s) to apply to each individual weight.
       target_shape: The final shape of the target MaxText tensor.
       config: The MaxText pyconfig object.
+      mt_key: The MaxText parameter key, used to detect the gemma4 nested-scan case.
 
   Returns:
       The final, assembled NumPy array for the MaxText parameter.
   """
-  all_expert_tensors = []
-  # The hook function needs the shape of an individual slice, not the full stacked tensor.
-  # For multi-axis stacking (experts, layers, ...), the slice shape is target_shape[2:]
-  mt_slice_shape = target_shape[2:]
-
-  # Outer loop iterates through experts
-  for layer_keys_for_expert in hf_source_keys:
-    layer_tensors_for_expert = []
-    # Inner loop iterates through layers for the current expert
-    for hf_key_single in layer_keys_for_expert:
+  # Gemma4's local layers are an inner scan nested inside the block scan, so the
+  # two stacked axes go at (param_scan_axis, param_scan_axis + 1); everything else
+  # (MoE experts/layers) stacks at the leading axes.
+  if isinstance(mt_key, str) and "scanned_blocks-local_layers" in mt_key:
+    outer_axis, inner_axis = config.param_scan_axis, config.param_scan_axis + 1
+  else:
+    outer_axis, inner_axis = 0, 1
+
+  # The hook function needs the shape of an individual slice: target_shape with
+  # the two stacked axes removed.
+  stacked_axes = {outer_axis, inner_axis}
+  mt_slice_shape = tuple(d for i, d in enumerate(target_shape) if i not in stacked_axes)
+
+  all_outer_tensors = []
+  # Outer loop iterates the outer stack axis (experts, or blocks for gemma4 local)
+  for inner_keys in hf_source_keys:
+    inner_tensors = []
+    # Inner loop iterates the inner stack axis (layers, or local layers for gemma4)
+    for hf_key_single in inner_keys:
       if isinstance(hf_key_single, (list, tuple)):
         hf_tensor_numpy = tuple(tensor_getter_fn(k) for k in hf_key_single)
       else:
         hf_tensor_numpy = tensor_getter_fn(hf_key_single)
-      processed_hf_tensor = apply_hook_fns(hf_tensor_numpy, mt_slice_shape, hook_fns)
-      layer_tensors_for_expert.append(processed_hf_tensor)
-    all_expert_tensors.append(np.stack(layer_tensors_for_expert, axis=0))
-  return np.stack(all_expert_tensors, axis=0)
+      inner_tensors.append(apply_hook_fns(hf_tensor_numpy, mt_slice_shape, hook_fns))
+    all_outer_tensors.append(np.stack(inner_tensors, axis=0))
+  stacked = np.stack(all_outer_tensors, axis=0)  # (outer, inner, *slice)
+
+  # Move the (outer, inner) axes from leading positions to their targets.
+  if (outer_axis, inner_axis) != (0, 1):
+    stacked = np.moveaxis(stacked, (0, 1), (outer_axis, inner_axis))
+  return stacked
 
 
 def _build_single_axis_stacked_tensor(
@@ -432,7 +453,7 @@ def _build_single_axis_stacked_tensor(
   return np.stack(tensors_to_stack, axis=axis_to_stack)
 
 
-def _get_hf_loading_function(hf_source_keys_or_key, tensor_getter, hook_fn, mt_target_shape_or_shapes, config):
+def _get_hf_loading_function(hf_source_keys_or_key, tensor_getter, hook_fn, mt_target_shape_or_shapes, config, mt_key=""):
   """Determine the loading function for HF keys.
 
   This function natively supports `composite_hf_key` mapping (where multiple HF keys
@@ -483,6 +504,7 @@ def _loader(getter, key, shape, hook):
         hook_fn,
         mt_target_shape_or_shapes,
         config,
+        mt_key,
     )
   return load_fn
 
@@ -1032,6 +1054,7 @@ def _eager_getter(key):
           hook_fn,
           mt_target_shape_or_shapes,
           config,
+          mt_param_key_or_keys,
       )
 
       # Step 3: Load hf keys and convert to maxtext keys
diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py
index 05432174ec..66a3b7f03a 100644
--- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py
+++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py
@@ -2824,112 +2824,76 @@ def GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False
     num_experts = tcfg.get("num_experts")
     num_experts = num_experts if num_experts is not None else 1
 
-    # Main scanned blocks
-    for layer_in_block in range(attention_pattern_length):
-      hf_indices = list(range(layer_in_block, num_scanned, attention_pattern_length))
-      prefix = f"params-decoder-scanned_blocks-layers_{layer_in_block}"
-      mapping.update(  # pyrefly: ignore[no-matching-overload]
-          {
-              f"{prefix}-self_attention-query-kernel": [
-                  f"{text_base}.layers.{i}.self_attn.q_proj.weight" for i in hf_indices
-              ],
-              f"{prefix}-self_attention-key-kernel": [
-                  f"{text_base}.layers.{i}.self_attn.k_proj.weight" for i in hf_indices
-              ],
-              f"{prefix}-self_attention-value-kernel": (
-                  None
-                  if share_kv_projections and layer_in_block == 5
-                  else [f"{text_base}.layers.{i}.self_attn.v_proj.weight" for i in hf_indices]
-              ),
-              f"{prefix}-self_attention-out-kernel": [
-                  f"{text_base}.layers.{i}.self_attn.o_proj.weight" for i in hf_indices
-              ],
-              f"{prefix}-self_attention-query_norm-scale": [
-                  f"{text_base}.layers.{i}.self_attn.q_norm.weight" for i in hf_indices
-              ],
-              f"{prefix}-self_attention-key_norm-scale": [
-                  f"{text_base}.layers.{i}.self_attn.k_norm.weight" for i in hf_indices
-              ],
-          }
-      )
-      if maxtext_config.v_norm_with_scale:
-        mapping.update(  # pyrefly: ignore[no-matching-overload]
-            {
-                f"{prefix}-self_attention-value_norm-scale": [
-                    f"{text_base}.layers.{i}.self_attn.v_norm.weight" for i in hf_indices
-                ]
-            }
-        )
-      mapping.update(  # pyrefly: ignore[no-matching-overload]
-          {
-              f"{prefix}-pre_self_attention_norm-scale": [
-                  f"{text_base}.layers.{i}.input_layernorm.weight" for i in hf_indices
-              ],
-              f"{prefix}-post_self_attention_norm-scale": [
-                  f"{text_base}.layers.{i}.post_attention_layernorm.weight" for i in hf_indices
-              ],
-              f"{prefix}-pre_ffw_norm-scale": [
-                  f"{text_base}.layers.{i}.pre_feedforward_layernorm.weight" for i in hf_indices
-              ],
-              f"{prefix}-post_ffw_norm-scale": [
-                  f"{text_base}.layers.{i}.post_feedforward_layernorm.weight" for i in hf_indices
-              ],
-              f"{prefix}-mlp-pre_feedforward_layernorm_2-scale": [
-                  f"{text_base}.layers.{i}.pre_feedforward_layernorm_2.weight" if num_experts > 1 else None
-                  for i in hf_indices
-              ],
-              f"{prefix}-mlp-post_feedforward_layernorm_1-scale": [
-                  f"{text_base}.layers.{i}.post_feedforward_layernorm_1.weight" if num_experts > 1 else None
-                  for i in hf_indices
-              ],
-              f"{prefix}-mlp-post_feedforward_layernorm_2-scale": [
-                  f"{text_base}.layers.{i}.post_feedforward_layernorm_2.weight" if num_experts > 1 else None
-                  for i in hf_indices
-              ],
-              f"{prefix}-mlp-pre_forward_scale_2": [
-                  f"{text_base}.layers.{i}.router.scale" if num_experts > 1 else None for i in hf_indices
-              ],
-              f"{prefix}-mlp-wi_0-kernel": [
-                  f"{text_base}.layers.{i}.mlp.gate_proj.weight" if num_experts == 1 else None for i in hf_indices
-              ],
-              f"{prefix}-mlp-wi_1-kernel": [
-                  f"{text_base}.layers.{i}.mlp.up_proj.weight" if num_experts == 1 else None for i in hf_indices
-              ],
-              f"{prefix}-mlp-wo-kernel": [
-                  f"{text_base}.layers.{i}.mlp.down_proj.weight" if num_experts == 1 else None for i in hf_indices
-              ],
-              f"{prefix}-mlp-moe_block-MoeBlock_0-gate-kernel": [
-                  f"{text_base}.layers.{i}.router.proj.weight" if num_experts > 1 else None for i in hf_indices
-              ],
-              f"{prefix}-mlp-moe_block-MoeBlock_0-wi_0": [
-                  f"{text_base}.layers.{i}.experts.gate_up_proj" if num_experts > 1 else None for i in hf_indices
-              ],
-              f"{prefix}-mlp-moe_block-MoeBlock_0-wi_1": [
-                  f"{text_base}.layers.{i}.experts.gate_up_proj" if num_experts > 1 else None for i in hf_indices
-              ],
-              f"{prefix}-mlp-moe_block-MoeBlock_0-wo": [
-                  f"{text_base}.layers.{i}.experts.down_proj" if num_experts > 1 else None for i in hf_indices
-              ],
-              f"{prefix}-mlp-moe_block-MoeBlock_0-per_expert_scale": [
-                  f"{text_base}.layers.{i}.router.per_expert_scale" if num_experts > 1 else None for i in hf_indices
-              ],
-              f"{prefix}-mlp-moe_block-shared_experts-wi_0-kernel": [
-                  f"{text_base}.layers.{i}.mlp.gate_proj.weight" if num_experts > 1 else None for i in hf_indices
-              ],
-              f"{prefix}-mlp-moe_block-shared_experts-wi_1-kernel": [
-                  f"{text_base}.layers.{i}.mlp.up_proj.weight" if num_experts > 1 else None for i in hf_indices
-              ],
-              f"{prefix}-mlp-moe_block-shared_experts-wo-kernel": [
-                  f"{text_base}.layers.{i}.mlp.down_proj.weight" if num_experts > 1 else None for i in hf_indices
-              ],
-              f"{prefix}-layer_scalar": [f"{text_base}.layers.{i}.layer_scalar" for i in hf_indices],
-          }
-      )
-      mapping = {
-          k: v
-          for k, v in mapping.items()
-          if (isinstance(v, list) and len(v) > 0 and v[0] is not None) or (not isinstance(v, list) and v is not None)
-      }
+    # Main scanned blocks. The block runs 5 local (sliding-window) layers as an
+    # inner scan, then a single global layer, so the scanned params are:
+    #   scanned_blocks-local_layers-*  -> nested [block][local]  (doubly scanned)
+    #   scanned_blocks-global_layer-*  -> flat [block]           (block scan only; the
+    #     length-1 _scan_global_layer scan is a runtime memory boundary, not a param stack)
+    num_blocks = num_scanned // attention_pattern_length
+
+    def hf_layer(idx, suffix):
+      return f"{text_base}.layers.{idx}.{suffix}"
+
+    # (maxtext subkey, hf suffix, gate) with gate in {"always", "dense", "moe"}.
+    # The attention value projection and value_norm are handled separately below.
+    param_specs = [
+        ("self_attention-query-kernel", "self_attn.q_proj.weight", "always"),
+        ("self_attention-key-kernel", "self_attn.k_proj.weight", "always"),
+        ("self_attention-out-kernel", "self_attn.o_proj.weight", "always"),
+        ("self_attention-query_norm-scale", "self_attn.q_norm.weight", "always"),
+        ("self_attention-key_norm-scale", "self_attn.k_norm.weight", "always"),
+        ("pre_self_attention_norm-scale", "input_layernorm.weight", "always"),
+        ("post_self_attention_norm-scale", "post_attention_layernorm.weight", "always"),
+        ("pre_ffw_norm-scale", "pre_feedforward_layernorm.weight", "always"),
+        ("post_ffw_norm-scale", "post_feedforward_layernorm.weight", "always"),
+        ("mlp-pre_feedforward_layernorm_2-scale", "pre_feedforward_layernorm_2.weight", "moe"),
+        ("mlp-post_feedforward_layernorm_1-scale", "post_feedforward_layernorm_1.weight", "moe"),
+        ("mlp-post_feedforward_layernorm_2-scale", "post_feedforward_layernorm_2.weight", "moe"),
+        ("mlp-pre_forward_scale_2", "router.scale", "moe"),
+        ("mlp-wi_0-kernel", "mlp.gate_proj.weight", "dense"),
+        ("mlp-wi_1-kernel", "mlp.up_proj.weight", "dense"),
+        ("mlp-wo-kernel", "mlp.down_proj.weight", "dense"),
+        ("mlp-moe_block-MoeBlock_0-gate-kernel", "router.proj.weight", "moe"),
+        ("mlp-moe_block-MoeBlock_0-wi_0", "experts.gate_up_proj", "moe"),
+        ("mlp-moe_block-MoeBlock_0-wi_1", "experts.gate_up_proj", "moe"),
+        ("mlp-moe_block-MoeBlock_0-wo", "experts.down_proj", "moe"),
+        ("mlp-moe_block-MoeBlock_0-per_expert_scale", "router.per_expert_scale", "moe"),
+        ("mlp-moe_block-shared_experts-wi_0-kernel", "mlp.gate_proj.weight", "moe"),
+        ("mlp-moe_block-shared_experts-wi_1-kernel", "mlp.up_proj.weight", "moe"),
+        ("mlp-moe_block-shared_experts-wo-kernel", "mlp.down_proj.weight", "moe"),
+        ("layer_scalar", "layer_scalar", "always"),
+    ]
+    if maxtext_config.v_norm_with_scale:
+      param_specs.append(("self_attention-value_norm-scale", "self_attn.v_norm.weight", "always"))
+
+    def _spec_active(gate):
+      return gate == "always" or (gate == "dense" and num_experts == 1) or (gate == "moe" and num_experts > 1)
+
+    # Local layers (block positions 0..4): nested [block][local] -> inner scan.
+    local_prefix = "params-decoder-scanned_blocks-local_layers"
+    local_positions = list(range(attention_pattern_length - 1))
+    for subkey, suffix, gate in param_specs:
+      if _spec_active(gate):
+        mapping[f"{local_prefix}-{subkey}"] = [  # pyrefly: ignore[no-matching-overload]
+            [hf_layer(b * attention_pattern_length + l, suffix) for l in local_positions] for b in range(num_blocks)
+        ]
+    mapping[f"{local_prefix}-self_attention-value-kernel"] = [  # pyrefly: ignore[no-matching-overload]
+        [hf_layer(b * attention_pattern_length + l, "self_attn.v_proj.weight") for l in local_positions]
+        for b in range(num_blocks)
+    ]
+
+    # Global layer (block position 5): single layer scanned over blocks only.
+    global_prefix = "params-decoder-scanned_blocks-global_layer"
+    global_position = attention_pattern_length - 1
+    for subkey, suffix, gate in param_specs:
+      if _spec_active(gate):
+        mapping[f"{global_prefix}-{subkey}"] = [  # pyrefly: ignore[no-matching-overload]
+            hf_layer(b * attention_pattern_length + global_position, suffix) for b in range(num_blocks)
+        ]
+    if not share_kv_projections:
+      mapping[f"{global_prefix}-self_attention-value-kernel"] = [  # pyrefly: ignore[no-matching-overload]
+          hf_layer(b * attention_pattern_length + global_position, "self_attn.v_proj.weight") for b in range(num_blocks)
+      ]
 
     # Remainder layers
     if num_remaining > 0:
@@ -3444,10 +3408,10 @@ def reshape_pos_emb(x, target_shape):
     attention_pattern_length = 6
     num_remaining = nlayers % attention_pattern_length
 
-    # Scanned sub-layer prefixes
-    for layer_in_block in range(attention_pattern_length):
-      is_global = layer_in_block % 6 == 5
-      prefix = f"params-decoder-scanned_blocks-layers_{layer_in_block}"
+    # Scanned block prefixes. The hooks are per-layer-slice, so they are identical
+    # regardless of scan nesting; only the prefix (and is_global for the shared-kv
+    # skip) differ between the local layers and the single global layer.
+    def _attach_block_hooks(prefix, is_global):
       for key in kernel_keys:
         if share_kv_projections and is_global and key == "self_attention-value-kernel":
           continue
@@ -3456,8 +3420,6 @@ def reshape_pos_emb(x, target_shape):
         hooks[f"{prefix}-{key}"] = reshape_kernel
       for key in norm_keys:
         hooks[f"{prefix}-{key}"] = scale_rmsnorm_layer
-
-      # Add these specialized 3D tensor hooks inside the loop
       if saving_to_hf and num_experts > 1:
         wi0_key = f"{prefix}-mlp-moe_block-MoeBlock_0-wi_0"
         wi1_key = f"{prefix}-mlp-moe_block-MoeBlock_0-wi_1"
@@ -3467,6 +3429,9 @@ def reshape_pos_emb(x, target_shape):
         hooks[f"{prefix}-mlp-moe_block-MoeBlock_0-wi_1"] = split_moe_wi1
       hooks[f"{prefix}-mlp-moe_block-MoeBlock_0-wo"] = reshape_moe_wo
 
+    _attach_block_hooks("params-decoder-scanned_blocks-local_layers", is_global=False)
+    _attach_block_hooks("params-decoder-scanned_blocks-global_layer", is_global=True)
+
     # Remainder sub-layer prefixes
     if num_remaining > 0:
       for rem_idx in range(num_remaining):
diff --git a/src/maxtext/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py
index 94ed14935e..f02bdc5f43 100644
--- a/src/maxtext/checkpoint_conversion/utils/utils.py
+++ b/src/maxtext/checkpoint_conversion/utils/utils.py
@@ -313,36 +313,42 @@ def process_maxtext_param(
 
     return output_weights
 
-  # Case 4: Multi-axis stacked (Scanned MoE layer)
-  # The tensor is stacked on expert and layer axes. We slice experts first, then layers.
-  # MaxText format is (experts, layers, ...), so expert axis is 0, layer axis is 1.
-  max_logging.log("\tscan moe")
-  expert_axis_to_slice = 0
-
-  # Outer loop for experts
-  for expert_idx, expert_paths_for_layer in enumerate(hf_target_paths):
-    # Slice along the expert axis to get the tensor for the current expert across all layers.
+  # Case 4: Multi-axis stacked. Two sub-cases (the inverse of _build_multi_axis_stacked_tensor):
+  #   - Scanned MoE: the tensor is stacked on (experts, layers) at the LEADING two axes, so we
+  #     slice axis 0 (experts) then axis 0 again (layers, after the expert axis is removed).
+  #   - Gemma4 nested block scan (scanned_blocks-local_layers): the block's local layers are an
+  #     inner scan nested in the block scan, so the two axes are at (param_scan_axis,
+  #     param_scan_axis + 1) -- outer = blocks, inner = local. We slice param_scan_axis (blocks),
+  #     then param_scan_axis again (local shifts down into that slot once blocks is removed).
+  key_str = maxtext_param_key[0] if isinstance(maxtext_param_key, tuple) else maxtext_param_key
+  if isinstance(key_str, str) and "scanned_blocks-local_layers" in key_str:
+    max_logging.log("\tscan gemma4 local")
+    outer_axis_to_slice = maxtext_config.param_scan_axis
+    inner_axis_to_slice = maxtext_config.param_scan_axis
+  else:
+    max_logging.log("\tscan moe")
+    outer_axis_to_slice = 0
+    inner_axis_to_slice = 0
+
+  # Outer loop (experts for MoE, blocks for gemma4 local)
+  for outer_idx, inner_paths in enumerate(hf_target_paths):
     if isinstance(maxtext_param_weight, list):
-      expert_tensor_slice = [
-          jax.lax.index_in_dim(x, expert_idx, axis=expert_axis_to_slice, keepdims=False) for x in maxtext_param_weight
+      outer_slice = [
+          jax.lax.index_in_dim(x, outer_idx, axis=outer_axis_to_slice, keepdims=False) for x in maxtext_param_weight
       ]
     else:
-      expert_tensor_slice = jax.lax.index_in_dim(
-          maxtext_param_weight, expert_idx, axis=expert_axis_to_slice, keepdims=False
-      )
+      outer_slice = jax.lax.index_in_dim(maxtext_param_weight, outer_idx, axis=outer_axis_to_slice, keepdims=False)
 
-    # Inner loop for layers
-    for layer_idx, hf_path in enumerate(expert_paths_for_layer):
-      # Slice the expert tensor along the layer axis to get the final individual weight.
-      # axis is 0 on the new sliced tensor
-      if isinstance(expert_tensor_slice, list):
-        layer_tensor_slice = [jax.lax.index_in_dim(x, layer_idx, axis=0, keepdims=False) for x in expert_tensor_slice]
+    # Inner loop (layers for MoE, local layers for gemma4)
+    for inner_idx, hf_path in enumerate(inner_paths):
+      if isinstance(outer_slice, list):
+        inner_slice = [jax.lax.index_in_dim(x, inner_idx, axis=inner_axis_to_slice, keepdims=False) for x in outer_slice]
       else:
-        layer_tensor_slice = jax.lax.index_in_dim(expert_tensor_slice, layer_idx, axis=0, keepdims=False)
+        inner_slice = jax.lax.index_in_dim(outer_slice, inner_idx, axis=inner_axis_to_slice, keepdims=False)
 
       _process(
           hf_path,
-          layer_tensor_slice,
+          inner_slice,
           output_weights,
           current_hook_fns,
           hf_shape_map,
diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 73f6ddc4a6..0f2960effb 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -584,6 +584,7 @@ logical_axis_rules: [
                       ['engram_dim', ['tensor']],
                       ['dense_layers', []],
                       ['moe_layers', []],
+                      ['local_layers', []],
                       ['mhc', []],
                       # ==========================================
                       # Inference(Prefill, Decode, Cache)
diff --git a/src/maxtext/configs/inference/vllm.yml b/src/maxtext/configs/inference/vllm.yml
index e977ee41fc..82046b7ee8 100644
--- a/src/maxtext/configs/inference/vllm.yml
+++ b/src/maxtext/configs/inference/vllm.yml
@@ -93,6 +93,7 @@ logical_axis_rules: [
                       ['embed', []],
                       ['norm', []],
                       ['layers', []],
+                      ['local_layers', []],
                       ['dense_layers', []],
                       ['moe_layers', []],
                       # ==========================================
diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py
index 25cbea73c4..d722270c4c 100644
--- a/src/maxtext/layers/decoders.py
+++ b/src/maxtext/layers/decoders.py
@@ -1426,7 +1426,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
@@ -1449,7 +1454,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,
@@ -1460,6 +1465,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,
@@ -1470,6 +1476,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 012c5c5155..9d076abceb 100644
--- a/src/maxtext/layers/nnx_decoders.py
+++ b/src/maxtext/layers/nnx_decoders.py
@@ -650,9 +650,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
 
@@ -862,9 +872,24 @@ 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.
 
+    This is the fuller of the two NNX scan appliers: it also threads external
+    (vLLM) KV caches via a static unroll and re-applies scan-axis metadata.
+    ``nnx_scan.apply_scanned_layers`` is a leaner, model-agnostic alternative
+    (currently used only by Gemma4); unifying the two is a follow-up cleanup.
+
     Args:
       layers: The stacked NNX module whose params are scanned over.
       x_in: The carry (hidden state) fed into the first layer.
@@ -873,6 +898,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:
@@ -958,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
@@ -990,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
 
       # Move the scan axis to each variable's param_scan_axis and restore its name
@@ -1970,13 +2005,25 @@ def _apply_gemma4_scanned_blocks(
     attention_pattern_length = len(gemma4.GEMMA4_ATTENTION_PATTERN)
     scan_length = cfg.num_decoder_layers // attention_pattern_length
 
-    # 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
index 1bd9667d94..e43a4d47bf 100644
--- a/src/maxtext/layers/nnx_scan.py
+++ b/src/maxtext/layers/nnx_scan.py
@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Utilities for constructing stacks of scanned NNX layers."""
+"""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
@@ -87,3 +88,57 @@ def update_leaf(leaf):
   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.
+
+  Note:
+    This is a minimal, model-agnostic scan primitive. The other NNX decoder
+    paths instead go through ``NNXDecoder._apply_layers_sequentially``, a heavier
+    applier that also threads external (vLLM) KV caches via a static unroll and
+    re-applies scan-axis metadata. Gemma4 is currently the only caller of this
+    function; unifying the two appliers is a follow-up cleanup.
+  """
+  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 6bb705cb23..14fc68cebf 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,200 @@ 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}.")
+
+    # The block runs its local (sliding-window) layers first, then a single
+    # global layer, matching GEMMA4_ATTENTION_PATTERN. Derive the per-type
+    # counts from the pattern (well, its first num_of_layers entries) rather
+    # than hardcoding the 5-local / 1-global split.
+    active_pattern = GEMMA4_ATTENTION_PATTERN[:num_of_layers]
+    self.num_local = sum(1 for attn_type in active_pattern if attn_type == AttentionType.LOCAL_SLIDING)
+    self.num_global = sum(1 for attn_type in active_pattern if attn_type == AttentionType.GLOBAL)
+
+    # num_of_layers can be 0: the decoders always construct a "remainder" block
+    # for num_decoder_layers % pattern_length layers, which is 0 whenever the
+    # layer count divides evenly (e.g. 31b = 60 layers). That block is built but
+    # never applied, so num_local or num_global may legitimately be 0 here.
+    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
+      )
+    else:
+      self.global_layer = None
+
+  def _run_layer(self, layer, y, layer_kwargs, kv_cache=None):
+    """Invokes one ``Gemma4DecoderLayer``, returning ``(output, updated_kv_cache)``.
+
+    This is the shared leaf used by the local scan, the global length-1 scan,
+    and the external kv-cache unroll, so it runs in every mode (train / prefill
+    / autoregressive). ``updated_kv_cache`` is ``None`` when the layer emits a bare
+    output rather than an ``(output, kv_cache)`` tuple.
+    """
+    out = layer(y, **layer_kwargs, kv_cache=kv_cache)
+    return out if isinstance(out, tuple) else (out, None)
+
+  @property
+  def _remat_enabled(self):
+    """Whether the block rematerializes its own layers.
+
+    False when the caller applies block-level remat instead
+    (``apply_internal_remat=False``) or when remat is disabled
+    (``remat_policy == "none"``). Note that ``remat_policy_fn``
+    is ``None`` for both ``"none"`` and ``"full"``, so it
+    cannot distinguish "no remat" from "full remat" on its own.
+    """
+    return self.apply_internal_remat and self.config.remat_policy != "none"
+
+  def _scan_local_layers(self, y, layer_kwargs):
+    """Runs the local (sliding-window) layers via a per-layer rematerialized ``jax.lax.scan``."""
+    remat = self._remat_enabled
+    return nnx_scan.apply_scanned_layers(
+        self.local_layers,
+        y,
+        length=self.num_local,
+        param_scan_axis=self.config.param_scan_axis,
+        apply_fn=lambda layer, carry: self._run_layer(layer, carry, layer_kwargs)[0],
+        remat=remat,
+        remat_policy=self.remat_policy_fn if remat else None,
+        # prevent_cse is only consulted by jax.checkpoint, i.e. when remat=True;
+        # its value is irrelevant otherwise.
+        prevent_cse=maxtext_utils.should_prevent_cse_in_remat(self.config) if remat else True,
+    )
+
+  def _scan_global_layer(self, y, layer_kwargs):
+    """Runs the single global-attention layer inside a length-1 ``jax.lax.scan``.
+
+    The length-1 scan is guarded by a trip-count-one while boundary and wraps
+    the layer in its own ``jax.checkpoint``, which keeps only one layer's
+    full-sequence-attention working set live at a time; without the boundary
+    (blocks are unrolled) XLA co-schedules every block's backward working set
+    and OOMs.
+    """
+    cfg = self.config
+    # Split the state into Intermediates and everything else. Non-Intermediate
+    # state (the large persistent weights/residuals) is carried through the scan
+    # so it stays off the offload-bitcast-prone ys path. Intermediates instead go
+    # in as scan xs and come out as ys: a sow can create or grow an Intermediate
+    # during the call (e.g. MoE moe_lb_loss accumulates into a tuple), which would
+    # break a carry's fixed pytree, and closing them over would mutate state from
+    # the wrong trace level (nnx.merge aliases the variables). Routing through
+    # xs/ys sidesteps both -- xs/ys have no matching-structure constraint and xs is
+    # trace-local. For a dense layer that sows nothing (31b) the Intermediate
+    # partition is empty and this is a no-op.
+    graphdef_g, intermediate_g, other_g = nnx.split(self.global_layer, nnx.Intermediate, ...)
+    intermediate_xs = jax.tree.map(lambda x: x[None], intermediate_g)
+
+    def run_global_layer(carry, intermediate_slice):
+      hidden_states, other = carry
+      layer = nnx.merge(graphdef_g, intermediate_slice, other)
+      new_hidden_states = self._run_layer(layer, hidden_states, layer_kwargs)[0]
+      _, new_intermediate, new_other = nnx.split(layer, nnx.Intermediate, ...)
+      return (new_hidden_states, new_other), new_intermediate
+
+    # Offloaded (pinned-host) residuals can't cross the trip-count-one boundary,
+    # so save would-be-offloaded tensors on device for the global layer 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._remat_enabled:
+      prevent_cse = maxtext_utils.should_prevent_cse_in_remat(self.config)
+      run_global_layer = jax.checkpoint(
+          run_global_layer,
+          policy=global_remat_policy,
+          prevent_cse=prevent_cse,
       )
-      setattr(self, layer_name, layer)
+
+    # Carry the non-Intermediate state through the loop instead of returning it as
+    # a stacked [1, ...] result: slicing that result previously introduced a bitcast
+    # between device and pinned-host memory under offload remat. Only the (tiny)
+    # Intermediates ride the xs/ys path.
+    with xla_metadata.set_xla_metadata(**{"skip-simplify-while-loops_trip-count-one": "true"}):
+      (y, final_other), stacked_intermediate = jax.lax.scan(
+          run_global_layer,
+          (y, other_g),
+          intermediate_xs,
+          length=1,
+      )
+
+    # Squeeze the length-1 scan axis off the updated Intermediate state and write
+    # it back to the module along with the carried non-Intermediate state.
+    intermediate_state = jax.tree.map(lambda x: x[0], stacked_intermediate)
+    nnx.update(self.global_layer, final_other, intermediate_state)
+    return y
+
+  def _forward_with_external_kv_cache(self, y, kv_cache, layer_kwargs):
+    """Runs the block with externally-supplied per-layer kv caches (vLLM PagedAttention).
+
+    Scanning would stack the kv-cache list, which copies it and breaks the
+    in-place PagedAttention updates, so the layers are unrolled statically. The
+    block's ``kv_cache`` is a per-layer list: the first ``num_local`` entries
+    feed the local layers, followed by the single global layer. Returns
+    ``(y, updated_kvs)`` with one updated cache per layer.
+    """
+    updated_kvs = []
+
+    if self.local_layers is not None:
+      # Slice the scanned local stack per layer, run it, collect the updated kv
+      # caches, and re-stack the per-layer state.
+      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 = self._run_layer(layer, y, layer_kwargs, kv_cache[i])
+        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)
+
+    if self.global_layer is not None:
+      y, new_kv = self._run_layer(self.global_layer, y, layer_kwargs, kv_cache[self.num_local])
+      updated_kvs.append(new_kv)
+
+    return y, tuple(updated_kvs)
 
   def __call__(
       self,
@@ -468,32 +660,36 @@ def __call__(
     cfg = self.config
     inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed"))
     inputs = checkpoint_name(inputs, "decoder_layer_input")
-    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}")(
-          y,
-          decoder_segment_ids,
-          decoder_positions,
-          deterministic,
-          model_mode,
-          previous_chunk=previous_chunk,
-          slot=slot,
-          bidirectional_mask=bidirectional_mask,
-          kv_cache=current_kv,
-          attention_metadata=attention_metadata,
-      )
-      if kv_cache is not None:
-        updated_kvs.append(new_kv)
 
+    # Arguments shared by every layer in the block. model_mode differentiates
+    # train / prefill / autoregressive inside each layer; the block itself does
+    # not branch on it.
+    layer_kwargs = {
+        "decoder_segment_ids": decoder_segment_ids,
+        "decoder_positions": decoder_positions,
+        "deterministic": deterministic,
+        "model_mode": model_mode,
+        "slot": slot,
+        "previous_chunk": previous_chunk,
+        "bidirectional_mask": bidirectional_mask,
+        "attention_metadata": attention_metadata,
+    }
+
+    # Externally-supplied per-layer caches (vLLM PagedAttention) force a static
+    # unroll; otherwise attention manages its own cache and we take the scanned
+    # path (train and standard prefill/autoregressive alike).
     if kv_cache is not None:
-      return y, tuple(updated_kvs)
-    elif cfg.scan_layers:
+      return self._forward_with_external_kv_cache(inputs, kv_cache, layer_kwargs)
+
+    y = inputs
+    if self.local_layers is not None:
+      y = self._scan_local_layers(y, layer_kwargs)
+    if self.global_layer is not None:
+      y = self._scan_global_layer(y, layer_kwargs)
+
+    if cfg.scan_layers:
       return y, None
-    else:
-      return y
+    return y
 
 
 Gemma4ScannableBlockToLinen = nnx_wrappers.to_linen_class(
diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py
index 022bfb8413..fcd5acb5cc 100644
--- a/tests/unit/nnx_decoders_test.py
+++ b/tests/unit/nnx_decoders_test.py
@@ -22,7 +22,9 @@
 """
 
 import sys
+from types import SimpleNamespace
 import unittest
+from unittest import mock
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -36,8 +38,10 @@
 
 from maxtext.common.common_types import (
     DECODING_ACTIVE_SEQUENCE_INDICATOR,
+    MODEL_MODE_AUTOREGRESSIVE,
     MODEL_MODE_PREFILL,
     MODEL_MODE_TRAIN,
+    AttentionType,
     DecoderBlockType,
     MultimodalInput,
 )
@@ -47,7 +51,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
@@ -716,6 +720,133 @@ 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 _SowingGemma4DecoderLayer(nnx.Module):
+  """Stand-in whose global layer sows an accumulating Intermediate, like MoE moe_lb_loss."""
+
+  def __init__(self, *, attention_type, **unused_kwargs):
+    self.is_global = attention_type == AttentionType.GLOBAL
+    # A trivial variable so the local layers have state for apply_scanned_layers to
+    # scan over (a bare module has nothing to scan and lax.scan can't infer length).
+    self.marker = nnx.Intermediate(jnp.zeros(()))
+
+  def __call__(self, inputs, *unused_args, kv_cache=None, **unused_kwargs):
+    output = inputs + 1
+    if self.is_global:
+      # nnx.sow appends into a tuple by default, so it grows across calls -- the
+      # MoE moe_lb_loss pattern that must not enter the global length-1 scan carry.
+      self.sow(nnx.Intermediate, "moe_lb_loss", jnp.sum(output))
+    if kv_cache is None:
+      return output
+    return output, kv_cache
+
+
+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_global_layer_sown_intermediate_accumulates_across_calls(self):
+    """A global layer that sows an accumulating Intermediate (e.g. MoE moe_lb_loss)
+    must not break the length-1 scan carry, even when the Intermediate already
+    exists from a previous call and the sow grows its tuple (1 -> 2 elements)."""
+    call_kwargs = {
+        "decoder_segment_ids": None,
+        "decoder_positions": None,
+        "deterministic": True,
+        "model_mode": MODEL_MODE_AUTOREGRESSIVE,
+    }
+    with mock.patch.object(gemma4, "Gemma4DecoderLayer", _SowingGemma4DecoderLayer):
+      block = self._make_block()
+      # First call creates moe_lb_loss on the global layer (1-tuple).
+      block(jnp.zeros((1, 1, 1)), **call_kwargs)
+      # Second call: moe_lb_loss already exists and the sow appends -> 2-tuple.
+      # Carrying it in the scan would change the carry pytree; the type-based
+      # split keeps Intermediates on the ys path instead.
+      block(jnp.zeros((1, 1, 1)), **call_kwargs)
+
+    self.assertEqual(len(block.global_layer.moe_lb_loss.value), 2)
+
+  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."""
 
@@ -880,6 +1011,75 @@ def test_gemma4_scanned_layers(self):
         (cfg.global_batch_size_to_train_on, cfg.max_target_length, cfg.vocab_size),
     )
 
+  def test_gemma4_block_external_kv_cache_matches_scanned_path(self):
+    """External-kv-cache path must numerically match the scanned (kv=None) path.
+
+    Guards the real stacked-parameter slice/re-stack in
+    ``Gemma4ScannableBlock._forward_with_external_kv_cache`` (``nnx.split`` of the
+    scanned local stack, ``param_scan_axis`` moveaxis, per-layer merge, re-stack,
+    ``nnx.update``) against the ``jax.lax.scan`` path. The mock-based
+    ``TestGemma4ScannableBlock`` tests cover call ordering / cache collection but
+    use trivial params, so they never exercise the real stacked-param mechanics.
+
+    Uses ``model_mode=TRAIN`` with ``dot_product`` attention: attention then
+    ignores the external caches (passing them straight through), so both paths
+    compute the same forward and only the loop mechanism (scan vs static unroll)
+    differs -- any mismatch is a slice/re-stack bug.
+    """
+    cfg = _make_config(
+        decoder_block="gemma4",
+        scan_layers=True,
+        num_decoder_layers=len(gemma4.GEMMA4_ATTENTION_PATTERN),  # exactly one full 5-local + 1-global block
+        base_emb_dim=128,
+        base_num_query_heads=4,
+        base_num_kv_heads=4,
+        # float32 + high matmul precision so the two paths agree to tight tolerance;
+        # the paths are mathematically identical, so any bf16-level rounding drift
+        # between scan and static-unroll compilation would only mask a real bug.
+        dtype="float32",
+        weight_dtype="float32",
+        matmul_precision="highest",
+    )
+    mesh = _make_mesh(cfg)
+
+    def make_block():
+      # Same seed => identical params in both blocks, so any output difference is
+      # attributable to the scan-vs-unroll code path, not initialization.
+      return gemma4.Gemma4ScannableBlock(
+          config=cfg,
+          mesh=mesh,
+          model_mode=MODEL_MODE_TRAIN,
+          quant=None,
+          num_of_layers=len(gemma4.GEMMA4_ATTENTION_PATTERN),
+          remat_policy_fn=None,
+          apply_internal_remat=False,
+          rngs=nnx.Rngs(params=0),
+      )
+
+    batch = cfg.global_batch_size_to_train_on
+    seq = cfg.max_target_length
+    inputs = jax.random.normal(self.rng, (batch, seq, cfg.emb_dim), dtype=cfg.dtype)
+    segment_ids = jnp.full((batch, seq), DECODING_ACTIVE_SEQUENCE_INDICATOR)
+    positions = jnp.broadcast_to(jnp.arange(seq)[None], (batch, seq))
+    call_kwargs = {
+        "decoder_segment_ids": segment_ids,
+        "decoder_positions": positions,
+        "deterministic": True,
+        "model_mode": MODEL_MODE_TRAIN,
+    }
+
+    # Scanned path (kv_cache=None); scan_layers=True => returns (y, None).
+    y_scanned, _ = make_block()(inputs, **call_kwargs)
+
+    # External-kv path: one cache per layer, ordered local[0..4] then global.
+    num_layers = len(gemma4.GEMMA4_ATTENTION_PATTERN)
+    external_kv = tuple(jnp.zeros((batch, seq), dtype=cfg.dtype) for _ in range(num_layers))
+    y_external, updated_kvs = make_block()(inputs, **call_kwargs, kv_cache=external_kv)
+
+    self.assertEqual(y_external.shape, inputs.shape)
+    self.assertEqual(len(updated_kvs), num_layers)
+    np.testing.assert_allclose(y_external, y_scanned, rtol=1e-5, atol=1e-5)
+
 
 @pytest.mark.tpu_only
 class TestGemma4SmallNNXDecoder(unittest.TestCase):
diff --git a/tests/unit/nnx_scan_test.py b/tests/unit/nnx_scan_test.py
index b307abe935..fcbe975af6 100644
--- a/tests/unit/nnx_scan_test.py
+++ b/tests/unit/nnx_scan_test.py
@@ -15,9 +15,12 @@
 """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
@@ -61,5 +64,64 @@ def test_create_zero_length_returns_none(self):
     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()
diff --git a/tests/unit/param_mapping_test.py b/tests/unit/param_mapping_test.py
index 8a18c64b9d..16ddcc7530 100644
--- a/tests/unit/param_mapping_test.py
+++ b/tests/unit/param_mapping_test.py
@@ -18,7 +18,9 @@
 from unittest import mock
 import numpy as np
 
+from maxtext.checkpoint_conversion.to_maxtext import _build_multi_axis_stacked_tensor
 from maxtext.checkpoint_conversion.utils import param_mapping
+from maxtext.checkpoint_conversion.utils.utils import process_maxtext_param
 
 
 class ParamMappingTest(unittest.TestCase):
@@ -184,7 +186,59 @@ def test_gemma4_mapping_scanned(self):
     maxtext_config.use_multimodal = False
     maxtext_config.v_norm_with_scale = False
     mapping = param_mapping.GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True)
-    self.assertIn("params-decoder-scanned_blocks-layers_0-self_attention-query-kernel", mapping)
+    # The block scans its 5 local layers (nested [block][local]) then a single global layer.
+    self.assertIn("params-decoder-scanned_blocks-local_layers-self_attention-query-kernel", mapping)
+    self.assertIn("params-decoder-scanned_blocks-global_layer-self_attention-query-kernel", mapping)
+    # local_layers value is nested [block][local]; global_layer is flat over blocks.
+    num_blocks = config["num_hidden_layers"] // 6
+    local_val = mapping["params-decoder-scanned_blocks-local_layers-self_attention-query-kernel"]
+    global_val = mapping["params-decoder-scanned_blocks-global_layer-self_attention-query-kernel"]
+    self.assertEqual(len(local_val), num_blocks)
+    self.assertEqual(len(local_val[0]), 5)
+    self.assertEqual(len(global_val), num_blocks)
+
+  def test_gemma4_local_layers_stack_unstack_roundtrip(self):
+    """Stacking HF weights into the nested [block][local] MaxText layout (to_maxtext) and
+    un-stacking them back (to_huggingface) must be identity, with the two scan axes placed at
+    (param_scan_axis, param_scan_axis + 1) -- not the leading axes used for MoE expert stacking."""
+    num_blocks, num_local = 2, 5
+    slice_shape = (4, 3)  # per-(block, local) HF weight shape
+    cfg = mock.Mock()
+    cfg.param_scan_axis = 1
+    cfg.scan_layers = True
+    cfg.weight_dtype = "float32"
+    cfg.rope_type = ""
+    cfg.model_name = "gemma4-31b"
+
+    mt_key = "params-decoder-scanned_blocks-local_layers-self_attention-query-kernel"
+
+    def value_of(b, l):
+      return np.full(slice_shape, b * 100 + l, dtype=np.float32)
+
+    hf_names = [[f"b{b}_l{l}" for l in range(num_local)] for b in range(num_blocks)]
+
+    def getter(name):
+      block_idx, local_idx = name.split("_")
+      return value_of(int(block_idx[1:]), int(local_idx[1:]))
+
+    # target: per-slice shape (4, 3) with (blocks, local) inserted at axes (1, 2)
+    target_shape = (slice_shape[0], num_blocks, num_local, slice_shape[1])
+
+    # Forward (to_maxtext): stack. Blocks must land at axis 1, local at axis 2.
+    stacked = _build_multi_axis_stacked_tensor(hf_names, getter, None, target_shape, cfg, mt_key)
+    self.assertEqual(stacked.shape, target_shape)
+    for b in range(num_blocks):
+      for l in range(num_local):
+        np.testing.assert_array_equal(stacked[:, b, l, :], value_of(b, l))
+
+    # Backward (to_huggingface): un-stack and check it round-trips to the originals.
+    param_map = {mt_key: hf_names}
+    hf_shape_map = {f"b{b}_l{l}": slice_shape for b in range(num_blocks) for l in range(num_local)}
+    out = dict(process_maxtext_param(mt_key, stacked, param_map, {}, hf_shape_map, cfg))
+    self.assertEqual(len(out), num_blocks * num_local)
+    for b in range(num_blocks):
+      for l in range(num_local):
+        np.testing.assert_array_equal(out[f"b{b}_l{l}"], value_of(b, l))
 
   # Specific tests with assertions
   def test_reshape_kernel_hook(self):