diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py
index 62c5fe0345..4785b8bd92 100644
--- a/src/maxtext/inference/maxengine/maxengine.py
+++ b/src/maxtext/inference/maxengine/maxengine.py
@@ -31,18 +31,14 @@
 else:
   from jax.experimental.layout import DeviceLocalLayout as DLL  # type: ignore # pylint: disable=no-name-in-module
 
-from flax import linen as nn
 from flax import nnx
 from flax import struct
 from flax.linen import partitioning as nn_partitioning
-import flax
 
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import MAXTEXT_PKG_DIR
-from maxtext.models import models
 from maxtext.layers import quantizations
 from maxtext.inference import inference_utils
-from maxtext.multimodal import processor as mm_processor
 from maxtext.utils import lora_utils
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
@@ -127,45 +123,39 @@ def __init__(self, config: Any, devices: Any | None = None):
 
     # Model and Optimizer definition.
     quant = quantizations.configure_quantization(config)
-    if config.pure_nnx:
-      # `serve` only when the on-disk checkpoint already carries `qrhs.frozen`
-      # (no full-precision kernel). For `checkpoint_is_quantized=False` with
-      # quant enabled we stay in `train` mode and let AQT quantize per-forward
-      # against the full-precision kernel — same numerical result as `serve`
-      # for absmax calibration, just slower.
-      nnx_quant_mode_str = "serve" if (quant is not None and config.checkpoint_is_quantized) else "train"
-      # We need both PREFILL and AR abstract models because the cache vars inherit
-      # CACHE_BATCH_PREFILL vs CACHE_BATCH from the construction model_mode, and
-      # bulk_insert searches for the substring "cache_batch" in the AR-mode names.
-      # Calling nnx.eval_shape directly (instead of create_nnx_abstract_model) avoids
-      # the jax.set_mesh wrap that trips Flax 0.12.6 on logical-only axes like "norm".
-      _create_model = model_creation_utils.get_nnx_create_model_fn(
-          config, mesh=self._mesh, model_mode=MODEL_MODE_PREFILL, quant_mode_str=nnx_quant_mode_str
-      )
-      _create_model_ar = model_creation_utils.get_nnx_create_model_fn(
-          config, mesh=self._mesh, model_mode=MODEL_MODE_AUTOREGRESSIVE, quant_mode_str=nnx_quant_mode_str
-      )
-      self._nnx_quant_mode_str = nnx_quant_mode_str
-      with nn_partitioning.axis_rules(config.logical_axis_rules):
-        abstract_model = nnx.eval_shape(_create_model)
-        abstract_model_ar = nnx.eval_shape(_create_model_ar)
-      self.model = abstract_model
-      self.model_ar = abstract_model_ar
-      # 3-way split so JIT bodies can pass (params, cache, rest) separately to
-      # nnx.merge. `rest` (RNG state etc.) is materialized in load_params.
-      graphdef, _, _, _ = nnx.split(abstract_model, nnx.Param, nnx.Cache, ...)
-      self.graphdef = graphdef
-      # Layers may bake their construction-time model_mode into static attributes,
-      # so a call must be merged with the graphdef built for that same mode.
-      graphdef_ar, _, _, _ = nnx.split(abstract_model_ar, nnx.Param, nnx.Cache, ...)
-      self.graphdef_ar = graphdef_ar
-      self._create_model_fn = _create_model
-      self._nnx_rest_state = None
-    else:
-      self.model = models.transformer_as_linen(config, mesh=self._mesh, quant=quant, model_mode=MODEL_MODE_PREFILL)
-      self.graphdef = None
-      self.graphdef_ar = None
-      self._create_model_fn = None
+    # `serve` only when the on-disk checkpoint already carries `qrhs.frozen`
+    # (no full-precision kernel). For `checkpoint_is_quantized=False` with
+    # quant enabled we stay in `train` mode and let AQT quantize per-forward
+    # against the full-precision kernel — same numerical result as `serve`
+    # for absmax calibration, just slower.
+    nnx_quant_mode_str = "serve" if (quant is not None and config.checkpoint_is_quantized) else "train"
+    # We need both PREFILL and AR abstract models because the cache vars inherit
+    # CACHE_BATCH_PREFILL vs CACHE_BATCH from the construction model_mode, and
+    # bulk_insert searches for the substring "cache_batch" in the AR-mode names.
+    # Calling nnx.eval_shape directly (instead of create_nnx_abstract_model) avoids
+    # the jax.set_mesh wrap that trips Flax 0.12.6 on logical-only axes like "norm".
+    _create_model = model_creation_utils.get_nnx_create_model_fn(
+        config, mesh=self._mesh, model_mode=MODEL_MODE_PREFILL, quant_mode_str=nnx_quant_mode_str
+    )
+    _create_model_ar = model_creation_utils.get_nnx_create_model_fn(
+        config, mesh=self._mesh, model_mode=MODEL_MODE_AUTOREGRESSIVE, quant_mode_str=nnx_quant_mode_str
+    )
+    self._nnx_quant_mode_str = nnx_quant_mode_str
+    with nn_partitioning.axis_rules(config.logical_axis_rules):
+      abstract_model = nnx.eval_shape(_create_model)
+      abstract_model_ar = nnx.eval_shape(_create_model_ar)
+    self.model = abstract_model
+    self.model_ar = abstract_model_ar
+    # 3-way split so JIT bodies can pass (params, cache, rest) separately to
+    # nnx.merge. `rest` (RNG state etc.) is materialized in load_params.
+    graphdef, _, _, _ = nnx.split(abstract_model, nnx.Param, nnx.Cache, ...)
+    self.graphdef = graphdef
+    # Layers may bake their construction-time model_mode into static attributes,
+    # so a call must be merged with the graphdef built for that same mode.
+    graphdef_ar, _, _, _ = nnx.split(abstract_model_ar, nnx.Param, nnx.Cache, ...)
+    self.graphdef_ar = graphdef_ar
+    self._create_model_fn = _create_model
+    self._nnx_rest_state = None
     self.replicated_sharding = jax.sharding.NamedSharding(self._mesh, P(None))
 
     self.abstract_params = None
@@ -351,65 +341,7 @@ def load_params(self, *args, params=None, rng: PRNGKeyType | None = None, **kwar
     if rng is None:
       rng = jax.random.PRNGKey(0)
 
-    if self.config.pure_nnx:
-      return self._load_params_nnx(params=params, rng=rng)
-
-    if self.model.quant and self.config.checkpoint_is_quantized:
-      print("Loading from the quantized checkpoint...")
-      self.model.quant.quant_mode = quantizations.get_quant_mode("serve")
-
-    rng1, rng2, rng3 = jax.random.split(rng, 3)
-    if params:
-      print("Resharding given params")
-      init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.config, False, rng)
-      _, self.state_mesh_annotations, state_mesh_shardings = maxtext_utils.get_abstract_state(
-          self.config, self._mesh, init_state_fn, False
-      )
-      # reshard given params based on shardings from config in MaxEngine
-      params = jax.device_put(params, state_mesh_shardings.params)
-      state = maxtext_utils.init_decode_state(None, params)
-      state = max_utils.unbox_logicallypartioned(state)
-    else:
-      init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.config, False, rng1)
-      state, self.state_mesh_annotations = maxtext_utils.setup_decode_state(self.config, self._mesh, None, init_state_fn)
-    # pylint: disable=isinstance-second-argument-not-valid-type
-    self.abstract_params = jax.tree_util.tree_map(
-        lambda x: jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype, sharding=x.sharding)
-        if isinstance(x, jax.Array)
-        else None,
-        state.params,
-    )
-
-    self.prefill_kv_cache_annotations = maxtext_utils.get_prefill_kv_cache_annotations(
-        self.model, self.config, rng2, self._mesh
-    )
-    self.prefill_kv_cache_shardings = jax.tree_util.tree_map(
-        lambda x: jax.sharding.NamedSharding(self._mesh, x),
-        self.prefill_kv_cache_annotations,
-    )
-
-    if self.config.stack_prefill_result_cache:
-      # Add extra axis for the axis generated by the stack.
-      self.prefill_kv_cache_shardings = jax.tree_util.tree_map(
-          lambda x: jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec(None, *x.spec)),
-          self.prefill_kv_cache_shardings,
-      )
-      self.prefill_kv_cache_shardings = self.prefill_kv_cache_shardings["decoder"]["layers_0"]
-
-    self.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations(self.model, self.config, rng2, self._mesh)
-    self.kv_cache_shardings = jax.tree_util.tree_map(
-        lambda x: jax.sharding.NamedSharding(self._mesh, x),
-        self.kv_cache_annotations,
-    )
-
-    if self.model.quant and not self.config.checkpoint_is_quantized:
-      params = self.quantize_params(state, rng3)
-    else:
-      params = state.params
-
-    self.print_stats("After load_params")
-
-    return params
+    return self._load_params_nnx(params=params, rng=rng)
 
   def _load_params_nnx(self, params, rng):
     """NNX equivalent of load_params: returns an nnx.Param state and populates KV cache shardings.
@@ -558,142 +490,77 @@ def apply_adapter(self, base_params, adapter_config, adapter_params):
 
     lora_rank = int(adapter_config["r"])
     lora_scale_factor = float(adapter_config["lora_alpha"]) / lora_rank
-    if self.config.pure_nnx:
-      lora_utils.apply_lora_on_base_params_nnx(base_params, adapter_params, lora_scale_factor)
-    else:
-      lora_utils.apply_lora_on_base_params(base_params, adapter_params, lora_scale_factor)
+    lora_utils.apply_lora_on_base_params_nnx(base_params, adapter_params, lora_scale_factor)
 
   def unapply_adapter(self, base_params, adapter_config, adapter_params):
     """Unapply the adapter params from the merged params to get back the base params."""
 
     lora_rank = int(adapter_config["r"])
     lora_scale_factor = float(adapter_config["lora_alpha"]) / lora_rank
-    if self.config.pure_nnx:
-      lora_utils.unapply_lora_from_base_params_nnx(base_params, adapter_params, lora_scale_factor)
-    else:
-      lora_utils.unapply_lora_from_base_params(base_params, adapter_params, lora_scale_factor)
+    lora_utils.unapply_lora_from_base_params_nnx(base_params, adapter_params, lora_scale_factor)
 
   def quantize_params(self, state, rng: PRNGKeyType | None = None):
     """Forward pass to quantize decode params."""
     if rng is None:
       rng = jax.random.PRNGKey(0)
-    if self.config.pure_nnx:
-      # NNX takes a different code path: convert-on-load lives in `_load_params_nnx`
-      # via `_convert_and_quantize_nnx`, which runs the dummy forward against a
-      # CONVERT-mode model and transfers `qrhs.frozen` into the SERVE model.
-      # The standalone `quantize_params(state, rng)` API expects a Linen-shape
-      # `state.params` dict and isn't reachable on the NNX pathway in maxengine
-      # (load_params already dispatched to _load_params_nnx).
-      raise NotImplementedError(
-          "Use load_params() on NNX — the convert step runs inside _load_params_nnx via "
-          "_convert_and_quantize_nnx. quantize_params(state, rng) is the Linen API."
-      )
-
-    self.model.quant.quant_mode = quantizations.get_quant_mode("convert")
-
-    @jax.jit
-    def model_apply(_p, _rng):
-      image_shape = mm_processor.get_dummy_image_shape_for_init(
-          model_name=self.config.model_name,
-          batch_size=self.config.micro_batch_size_to_train_on,
-      )
-      audio_shape = mm_processor.get_dummy_audio_shape_for_init(self.config)
-      return self.model.apply(
-          _p | {"aqt": {}},
-          jnp.ones((1, self.config.max_prefill_predict_length), dtype=jnp.int32),
-          jnp.ones((1, self.config.max_prefill_predict_length), dtype=jnp.int32),
-          encoder_images=jnp.ones(image_shape, dtype=jnp.float32) if self.config.use_multimodal else None,
-          # encoder_image_masks indicates valid tiles if image tiling + padding is used in vision encoder input.
-          encoder_image_masks=jnp.ones(image_shape[:2], dtype=jnp.int32)
-          if self.config.use_multimodal and "llama4" in self.config.model_name
-          else None,
-          encoder_audios=jnp.ones(audio_shape, dtype=jnp.float32) if self.config.use_audio else None,
-          decoder_segment_ids=jnp.zeros((1, self.config.max_prefill_predict_length), dtype=jnp.int32),
-          enable_dropout=False,
-          model_mode=MODEL_MODE_PREFILL,
-          rngs={"params": _rng},
-          mutable=True,
-      )
-
-    _, new_vars = model_apply(state.params, rng)
-    # Remove param values which have corresponding qtensors in aqt to save memory.
-    params = {}
-    params["aqt"] = new_vars["aqt"]
-    params["params"] = quantizations.remove_quantized_params(state.params["params"], new_vars["aqt"])
-    self.abstract_params = jax.tree_util.tree_map(
-        lambda x: jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype, sharding=x.sharding),
-        params,
+    # NNX takes a different code path: convert-on-load lives in `_load_params_nnx`
+    # via `_convert_and_quantize_nnx`, which runs the dummy forward against a
+    # CONVERT-mode model and transfers `qrhs.frozen` into the SERVE model.
+    # The standalone `quantize_params(state, rng)` API expects a Linen-shape
+    # `state.params` dict and isn't reachable on the NNX pathway in maxengine
+    # (load_params already dispatched to _load_params_nnx).
+    raise NotImplementedError(
+        "Use load_params() on NNX — the convert step runs inside _load_params_nnx via "
+        "_convert_and_quantize_nnx. quantize_params(state, rng) is the Linen API."
     )
-    maxtext_utils.save_quantized_checkpoint_if_configured(self.config, params)
-    self.model.quant.quant_mode = quantizations.get_quant_mode("serve")
-    return params
 
   def _maybe_stack_prefill_result_cache(self, cache):
     """Stack the caches across the layers."""
     if not self.config.stack_prefill_result_cache:
       return cache
 
-    if self.config.pure_nnx:
-      if self.config.scan_layers:
-        # scan_layers already stacks the per-layer KV cache on axis 0; nothing to restack.
-        return cache
-      # scan_layers=False: stack the per-layer subtrees under decoder into one
-      # subtree with a leading layer axis (matching the scan_layers=True shape).
-      if "dense_layers_0" in cache["decoder"] or "moe_layers_0" in cache["decoder"]:
-        first_dense = self.config.first_num_dense_layers
-        num_moe = self.config.num_decoder_layers - first_dense
-        layer_keys = [f"dense_layers_{i}" for i in range(first_dense)] + [f"moe_layers_{i}" for i in range(num_moe)]
-      else:
-        layer_keys = [f"layers_{i}" for i in range(self.config.num_decoder_layers)]
-
-      layer_cache = [cache["decoder"][key] for key in layer_keys]
-      stacked = jax.tree.map(lambda *c: jnp.stack(c), *layer_cache)
-      return {"decoder": {"layers": stacked}}
-
-    layer_keys = []
-    for i in range(self.config.num_decoder_layers):
-      layer_keys.append(f"layers_{i}")
-
-    layer_cache = [cache["decoder"][layer_key] for layer_key in layer_keys]
+    if self.config.scan_layers:
+      # scan_layers already stacks the per-layer KV cache on axis 0; nothing to restack.
+      return cache
+    # scan_layers=False: stack the per-layer subtrees under decoder into one
+    # subtree with a leading layer axis (matching the scan_layers=True shape).
+    if "dense_layers_0" in cache["decoder"] or "moe_layers_0" in cache["decoder"]:
+      first_dense = self.config.first_num_dense_layers
+      num_moe = self.config.num_decoder_layers - first_dense
+      layer_keys = [f"dense_layers_{i}" for i in range(first_dense)] + [f"moe_layers_{i}" for i in range(num_moe)]
+    else:
+      layer_keys = [f"layers_{i}" for i in range(self.config.num_decoder_layers)]
 
-    return jax.tree.map(lambda *c: jnp.stack(c), *layer_cache)
+    layer_cache = [cache["decoder"][key] for key in layer_keys]
+    stacked = jax.tree.map(lambda *c: jnp.stack(c), *layer_cache)
+    return {"decoder": {"layers": stacked}}
 
   def _maybe_unstack_prefill_result_cache(self, cache):
     """Unstack the caches across the layers."""
     if not self.config.stack_prefill_result_cache:
       return cache
 
-    if self.config.pure_nnx:
-      if self.config.scan_layers:
-        # Mirror _maybe_stack_prefill_result_cache: the cache already carries the
-        # layer axis, so there is nothing to unstack.
-        return cache
-      # scan_layers=False: split the leading layer axis back into per-layer subtrees.
-      stacked = cache["decoder"]["layers"]
-      res_cache = {"decoder": {}}
-      is_deepseek = (
-          getattr(self.model, "is_deepseek", False)
-          or (hasattr(self.model, "decoder") and getattr(self.model.decoder, "is_deepseek", False))
-          or (hasattr(self.config, "decoder_block") and str(self.config.decoder_block).lower() == "deepseek")
-      )
-      if is_deepseek:
-        first_dense = self.config.first_num_dense_layers
-        num_moe = self.config.num_decoder_layers - first_dense
-        layer_keys = [f"dense_layers_{i}" for i in range(first_dense)] + [f"moe_layers_{i}" for i in range(num_moe)]
-      else:
-        layer_keys = [f"layers_{i}" for i in range(self.config.num_decoder_layers)]
-
-      for idx, key in enumerate(layer_keys):
-        res_cache["decoder"][key] = jax.tree.map(lambda x, i=idx: x[i], stacked)
-      return res_cache
-
-    flat_cache, treedef = jax.tree.flatten(cache)
-    layer_cache = [jax.tree.unflatten(treedef, flat_cache_vars) for flat_cache_vars in zip(*flat_cache, strict=True)]
+    if self.config.scan_layers:
+      # Mirror _maybe_stack_prefill_result_cache: the cache already carries the
+      # layer axis, so there is nothing to unstack.
+      return cache
+    # scan_layers=False: split the leading layer axis back into per-layer subtrees.
+    stacked = cache["decoder"]["layers"]
     res_cache = {"decoder": {}}
+    is_deepseek = (
+        getattr(self.model, "is_deepseek", False)
+        or (hasattr(self.model, "decoder") and getattr(self.model.decoder, "is_deepseek", False))
+        or (hasattr(self.config, "decoder_block") and str(self.config.decoder_block).lower() == "deepseek")
+    )
+    if is_deepseek:
+      first_dense = self.config.first_num_dense_layers
+      num_moe = self.config.num_decoder_layers - first_dense
+      layer_keys = [f"dense_layers_{i}" for i in range(first_dense)] + [f"moe_layers_{i}" for i in range(num_moe)]
+    else:
+      layer_keys = [f"layers_{i}" for i in range(self.config.num_decoder_layers)]
 
-    for i in range(self.config.num_decoder_layers):
-      res_cache["decoder"][f"layers_{i}"] = layer_cache[i]
-
+    for idx, key in enumerate(layer_keys):
+      res_cache["decoder"][key] = jax.tree.map(lambda x, i=idx: x[i], stacked)
     return res_cache
 
   def prefill_aot(  # pylint: disable=too-many-positional-arguments
@@ -786,10 +653,7 @@ def _prefill_jit(
     if existing_prefix is not None:
       if not self.use_chunked_prefill:
         raise ValueError("Using chunked prefill is needed for existing_prefix.")
-      # NNX threads existing_prefix.cache via the nnx_cache local below; only
-      # the Linen path merges cache into input_params (params is a dict there).
-      if not self.config.pure_nnx:
-        input_params = params | {"cache": existing_prefix.cache}
+      # NNX threads existing_prefix.cache via the nnx_cache local below.
       start_position = existing_prefix.common_prefix_tokens.shape[0]
       # TODO(yuyanpeng): rename previous_chunk
       previous_chunk = jnp.expand_dims(existing_prefix.common_prefix_tokens, 0)
@@ -821,52 +685,30 @@ def _prefill_jit(
     sequence_indicator = jnp.expand_dims(one_d_output, 0)
 
     rng, new_rng = jax.random.split(rng)  # pyrefly: ignore[bad-argument-type]
-    if self.config.pure_nnx:
-      # Prefill always operates on batch=1 (one padded prompt at a time).
-      nnx_cache = (
-          existing_prefix.cache if existing_prefix is not None else self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL)
+    # Prefill always operates on batch=1 (one padded prompt at a time).
+    nnx_cache = (
+        existing_prefix.cache if existing_prefix is not None else self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL)
+    )
+    with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      flat_logits, new_cache_dict = self._nnx_run_model(
+          params=input_params,
+          cache_dict=nnx_cache,
+          decoder_input_tokens=input_tokens,
+          decoder_positions=positions,
+          decoder_segment_ids=sequence_indicator,
+          encoder_images=images,
+          encoder_image_masks=image_masks,
+          encoder_videos=videos,
+          encoder_video_masks=video_masks,
+          encoder_video_grid_thw=video_grid_thw,
+          encoder_audios=audio_values,
+          enable_dropout=False,
+          model_mode=MODEL_MODE_PREFILL,
+          previous_chunk=previous_chunk,
+          true_length=true_length,
+          slot=slot,
       )
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        flat_logits, new_cache_dict = self._nnx_run_model(
-            params=input_params,
-            cache_dict=nnx_cache,
-            decoder_input_tokens=input_tokens,
-            decoder_positions=positions,
-            decoder_segment_ids=sequence_indicator,
-            encoder_images=images,
-            encoder_image_masks=image_masks,
-            encoder_videos=videos,
-            encoder_video_masks=video_masks,
-            encoder_video_grid_thw=video_grid_thw,
-            encoder_audios=audio_values,
-            enable_dropout=False,
-            model_mode=MODEL_MODE_PREFILL,
-            previous_chunk=previous_chunk,
-            true_length=true_length,
-            slot=slot,
-        )
-      new_vars = {"cache": new_cache_dict}
-    else:
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        flat_logits, new_vars = self.model.apply(
-            input_params,
-            input_tokens,
-            positions,
-            encoder_images=images,
-            encoder_image_masks=image_masks,
-            encoder_videos=videos,
-            encoder_video_masks=video_masks,
-            encoder_video_grid_thw=video_grid_thw,
-            encoder_audios=audio_values,
-            decoder_segment_ids=sequence_indicator,
-            enable_dropout=False,
-            model_mode=MODEL_MODE_PREFILL,
-            rngs={"params": new_rng},
-            mutable=["cache"],
-            previous_chunk=previous_chunk,
-            true_length=true_length,
-            slot=slot,
-        )
+    new_vars = {"cache": new_cache_dict}
     if return_prompt_logp:
       prompt_logp = inference_utils.prompt_logprobs_from_prefill(flat_logits, input_tokens, true_length)
     else:
@@ -1089,32 +931,19 @@ def _prefill_multisampling_jit(
     sequence_indicator = jnp.expand_dims(one_d_output, 0)
 
     rng, new_rng = jax.random.split(rng)  # pyrefly: ignore[bad-argument-type]
-    if self.config.pure_nnx:
-      # Prefill is batch=1 (one prompt); multi-sampling only draws several first
-      # tokens from the shared logits below. Mirror the _prefill_jit NNX branch.
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        flat_logits, new_cache_dict = self._nnx_run_model(
-            params=params,
-            cache_dict=self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL),
-            decoder_input_tokens=input_tokens,
-            decoder_positions=positions,
-            decoder_segment_ids=sequence_indicator,
-            enable_dropout=False,
-            model_mode=MODEL_MODE_PREFILL,
-        )
-      new_vars = {"cache": new_cache_dict}
-    else:
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        flat_logits, new_vars = self.model.apply(
-            params,
-            input_tokens,
-            positions,
-            decoder_segment_ids=sequence_indicator,
-            enable_dropout=False,
-            model_mode=MODEL_MODE_PREFILL,
-            rngs={"params": new_rng},
-            mutable=["cache"],
-        )
+    # Prefill is batch=1 (one prompt); multi-sampling only draws several first
+    # tokens from the shared logits below. Mirror the _prefill_jit NNX branch.
+    with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      flat_logits, new_cache_dict = self._nnx_run_model(
+          params=params,
+          cache_dict=self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL),
+          decoder_input_tokens=input_tokens,
+          decoder_positions=positions,
+          decoder_segment_ids=sequence_indicator,
+          enable_dropout=False,
+          model_mode=MODEL_MODE_PREFILL,
+      )
+    new_vars = {"cache": new_cache_dict}
 
     next_pos = jnp.full((1, 1), true_length, dtype=jnp.int32)
     selected_logits = jax.lax.dynamic_slice(
@@ -1225,33 +1054,20 @@ def prefill_concat(
     input_tokens = jnp.expand_dims(padded_tokens, 0)  # [BATCH, SEQUENCE]
     decoder_positions = jnp.expand_dims(decoder_positions, 0)
     decoder_segment_ids = jnp.expand_dims(decoder_segment_ids, 0)
-    rng, new_rng = jax.random.split(rng)
-    if self.config.pure_nnx:
-      # Packed prompts run as a single batch=1 prefill; the packed positions and
-      # segment ids keep the prompts separated. Mirror the _prefill_jit NNX branch.
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        flat_logits, new_cache_dict = self._nnx_run_model(
-            params=params,
-            cache_dict=self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL),
-            decoder_input_tokens=input_tokens,
-            decoder_positions=decoder_positions,
-            decoder_segment_ids=decoder_segment_ids,
-            enable_dropout=False,
-            model_mode=MODEL_MODE_PREFILL,
-        )
-      new_vars = {"cache": new_cache_dict}
-    else:
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        flat_logits, new_vars = self.model.apply(
-            params,
-            input_tokens,
-            decoder_positions,
-            decoder_segment_ids=decoder_segment_ids,
-            enable_dropout=False,
-            model_mode=MODEL_MODE_PREFILL,
-            rngs={"params": new_rng},
-            mutable=["cache"],
-        )
+    rng, _ = jax.random.split(rng)
+    # Packed prompts run as a single batch=1 prefill; the packed positions and
+    # segment ids keep the prompts separated. Mirror the _prefill_jit NNX branch.
+    with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      flat_logits, new_cache_dict = self._nnx_run_model(
+          params=params,
+          cache_dict=self._nnx_init_cache_dict(mode=MODEL_MODE_PREFILL),
+          decoder_input_tokens=input_tokens,
+          decoder_positions=decoder_positions,
+          decoder_segment_ids=decoder_segment_ids,
+          enable_dropout=False,
+          model_mode=MODEL_MODE_PREFILL,
+      )
+    new_vars = {"cache": new_cache_dict}
     cache = new_vars["cache"]
     cache = self._maybe_stack_prefill_result_cache(cache)
     if return_prompt_logp:
@@ -1399,28 +1215,16 @@ def _generate_jit(
     previous_token = decode_state["tokens"]
     rng, new_rng = jax.random.split(rng)  # pyrefly: ignore[bad-argument-type]
     # run one step generation
-    if self.config.pure_nnx:
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        out_logits, new_cache_dict = self._nnx_run_model(
-            params=params,
-            cache_dict=decode_state["cache"],
-            decoder_input_tokens=previous_token,
-            decoder_positions=decode_state["next_pos"],
-            enable_dropout=False,
-            model_mode=MODEL_MODE_AUTOREGRESSIVE,
-        )
-      new_vars = {"cache": new_cache_dict}
-    else:
-      with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        out_logits, new_vars = self.model.apply(
-            params | {"cache": decode_state["cache"]},
-            previous_token,
-            decode_state["next_pos"],
-            enable_dropout=False,
-            model_mode=MODEL_MODE_AUTOREGRESSIVE,
-            rngs={"params": new_rng},
-            mutable=["cache"],
-        )
+    with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+      out_logits, new_cache_dict = self._nnx_run_model(
+          params=params,
+          cache_dict=decode_state["cache"],
+          decoder_input_tokens=previous_token,
+          decoder_positions=decode_state["next_pos"],
+          enable_dropout=False,
+          model_mode=MODEL_MODE_AUTOREGRESSIVE,
+      )
+    new_vars = {"cache": new_cache_dict}
     out_logits = jax.lax.with_sharding_constraint(out_logits, self.replicated_sharding)
     new_cache = jax.lax.with_sharding_constraint(new_vars["cache"], self.kv_cache_shardings)
     # sampling tokens
@@ -1937,107 +1741,7 @@ def init_decode_state(
       **kwargs,  # pylint: disable=unused-argument
   ) -> DecodeState:
     """Initialises any state which a generation step transforms."""
-    if rng is None:
-      rng = jax.random.PRNGKey(0)
-
-    if self.config.pure_nnx:
-      return self._init_decode_state_nnx()
-
-    # pylint: disable=unused-argument
-    def init(abstract_params):
-      x = jnp.ones(
-          (int(self.config.per_device_batch_size * self.mesh.size), 1),
-          dtype=jnp.int32,
-      )
-      dummy_image = jnp.ones(
-          mm_processor.get_dummy_image_shape_for_init(
-              model_name=self.config.model_name, batch_size=int(self.config.per_device_batch_size * self.mesh.size)
-          ),
-          dtype=jnp.int32,
-      )
-      dummy_audio = jnp.ones(
-          mm_processor.get_dummy_audio_shape_for_init(self.config),
-          dtype=jnp.float32,
-      )
-      _, cache = self.model.apply(
-          abstract_params,
-          x,
-          x,
-          encoder_images=dummy_image if self.config.use_multimodal else None,
-          encoder_audios=dummy_audio if self.config.use_audio else None,
-          enable_dropout=False,
-          model_mode=MODEL_MODE_AUTOREGRESSIVE,
-          rngs={"params": rng},
-          mutable=["cache"],
-          slot=0,
-      )
-
-      next_pos = jnp.zeros(
-          (int(self.config.per_device_batch_size * self.mesh.size), 1),
-          dtype=jnp.int32,
-      )
-      generated_tokens = jnp.zeros(
-          (int(self.config.per_device_batch_size * self.mesh.size), 1),
-          dtype=jnp.int32,
-      )
-      tokens = jnp.zeros(
-          (int(self.config.per_device_batch_size * self.mesh.size), 1),
-          dtype=jnp.int32,
-      )
-      token_logp = jnp.zeros(
-          (int(self.config.per_device_batch_size * self.mesh.size), 1),
-          dtype=jnp.float32,
-      )
-      return {
-          "logits": jnp.zeros(
-              (
-                  int(self.config.per_device_batch_size * self.mesh.size),
-                  1,
-                  self.config.vocab_size,
-              )
-          ),
-          "cache": cache["cache"],
-          "next_pos": next_pos,
-          "generated_tokens": generated_tokens,
-          "tokens": tokens,
-          "token_logp": token_logp,
-      }
-
-    with nn_partitioning.axis_rules(self.config.logical_axis_rules):
-      abstract_outputs = jax.eval_shape(init, self.abstract_params)
-    logical_annotations = nn.get_partition_spec(abstract_outputs)
-
-    with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-      mesh_annotations = nn.logical_to_mesh(logical_annotations)
-
-    shardings = jax.tree_util.tree_map(
-        lambda mesh_annotation: jax.sharding.NamedSharding(self._mesh, mesh_annotation),
-        mesh_annotations,
-    )
-
-    if self._compiled_initialize_fn is None:
-
-      @functools.partial(jax.jit, out_shardings=shardings)
-      def initialize():
-        return jax.tree_util.tree_map(lambda x: jnp.zeros(x.shape, x.dtype), abstract_outputs)
-
-      self._compiled_initialize_fn = initialize
-
-    init_state = self._compiled_initialize_fn()
-    cache = init_state["cache"]
-
-    def is_lp(k):
-      return isinstance(k, flax.linen.spmd.LogicallyPartitioned)
-
-    self.kv_cache_annotations_named = jax.tree_util.tree_map(
-        lambda x: tuple(x.logical_axes)
-        if hasattr(x, "logical_axes")
-        else (tuple(x.names) if hasattr(x, "names") else ()),
-        cache,
-        is_leaf=is_lp,
-    )
-    zeroed = max_utils.unbox_logicallypartioned(init_state)
-    return zeroed
+    return self._init_decode_state_nnx()
 
   def _init_decode_state_nnx(self) -> DecodeState:
     """NNX equivalent of init_decode_state. Returns a decode_state dict with a pure-dict cache."""
@@ -2130,18 +1834,10 @@ def set_engine_vars_from_base_engine(
   """Set internal vars from base_engine, which has already loaded the checkpoint and has sharding,
   mesh, and kv cache related vars set.
   """
-  if not engine.config.pure_nnx and base_engine.model.quant:
-    # NNX bakes the quant mode in at construction (via _nnx_quant_mode_str) rather
-    # than mutating model.quant.quant_mode, so there's nothing to copy on that path.
-    engine.model.quant.quant_mode = base_engine.model.quant.quant_mode
   engine.state_mesh_annotations = base_engine.state_mesh_annotations
   engine.abstract_params = base_engine.abstract_params
-  if engine.config.pure_nnx:
-    # Linen's get_kv_cache_annotations calls model.init(); NNX modules have no
-    # .init, so use the abstract-model variant (mirrors _load_params_nnx).
-    engine.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations_nnx(engine.model_ar, engine.config, engine.mesh)
-  else:
-    engine.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations(engine.model, engine.config, rng, engine.mesh)  # pylint: disable=protected-access
+  # NNX modules have no .init, so use the abstract-model variant (mirrors _load_params_nnx).
+  engine.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations_nnx(engine.model_ar, engine.config, engine.mesh)
   engine.kv_cache_shardings = jax.tree_util.tree_map(
       lambda x: jax.sharding.NamedSharding(engine.mesh, x),
       engine.kv_cache_annotations,  # pylint: disable=protected-access
diff --git a/src/maxtext/inference/vllm_decode.py b/src/maxtext/inference/vllm_decode.py
index c70cb714f8..4fcdfe2b3c 100644
--- a/src/maxtext/inference/vllm_decode.py
+++ b/src/maxtext/inference/vllm_decode.py
@@ -114,8 +114,6 @@ def decode_with_vllm(config: Config) -> None:
               "debug_sharding": config.debug_sharding,
               "prefuse_moe_weights": config.prefuse_moe_weights,
               "scan_layers": config.scan_layers,
-              "enable_nnx": config.enable_nnx,
-              "pure_nnx_decoder": config.pure_nnx_decoder,
           },
           "sharding": {
               "sharding_strategy": {
@@ -271,8 +269,6 @@ def decode_with_tunix(
           "debug_sharding": config.debug_sharding,
           "prefuse_moe_weights": config.prefuse_moe_weights,
           "scan_layers": config.scan_layers,
-          "enable_nnx": config.enable_nnx,
-          "pure_nnx_decoder": config.pure_nnx_decoder,
       }
   }
 
diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py
index 08de8ed8e0..f304873122 100644
--- a/src/maxtext/utils/lora_utils.py
+++ b/src/maxtext/utils/lora_utils.py
@@ -15,7 +15,6 @@
 """Common LoRA utils needed to support LoRA adapters."""
 
 from collections.abc import Mapping
-from functools import partial
 import json
 import os
 import re
@@ -171,10 +170,7 @@ def load_adapter(config, base_abstract_state_params, adapter_config_path, adapte
     if not gcs_utils.gcs_path_exists(f"{adapter_weights_path}/commit_success.txt"):
       raise FileNotFoundError(f"Failed to read lora_weights from {adapter_weights_path}.")
 
-    if config.pure_nnx:
-      lora_state, _ = get_lora_abstract_state_nnx(base_abstract_state_params, lora_config)
-    else:
-      lora_state, _ = get_lora_abstract_state(base_abstract_state_params, lora_config)
+    lora_state, _ = get_lora_abstract_state_nnx(base_abstract_state_params, lora_config)
 
     with nn_partitioning.axis_rules(config.logical_axis_rules):
       lora_params = checkpointing.load_params_from_path(
@@ -218,37 +214,31 @@ def setup_initial_lora_state(model, data_iterator, tx, config, rng, mesh, checkp
 
   if lora_adapter_path:
     max_logging.log(f"Setting initial state of LoRA with lora_adapter_path = {lora_adapter_path}")
-    if config.pure_nnx:
-      # pylint: disable=import-outside-toplevel
-      from maxtext.common import train_state_nnx
-      from maxtext.utils import model_creation_utils
-
-      _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh)
-
-      def create_train_state_fn():
-        nnx_model = _create_model_partial()
-        wrt = (
-            getattr(nnx, "LoRAParam", nnx.Param)
-            if getattr(getattr(config, "lora", None), "enable_lora", False)
-            else nnx.Param
-        )
-        optimizer = nnx.Optimizer(nnx_model, tx, wrt=wrt)
-        return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
+    # pylint: disable=import-outside-toplevel
+    from maxtext.common import train_state_nnx
+    from maxtext.utils import model_creation_utils
+
+    _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh)
+
+    def create_train_state_fn():
+      nnx_model = _create_model_partial()
+      wrt = (
+          getattr(nnx, "LoRAParam", nnx.Param)
+          if getattr(getattr(config, "lora", None), "enable_lora", False)
+          else nnx.Param
+      )
+      optimizer = nnx.Optimizer(nnx_model, tx, wrt=wrt)
+      return train_state_nnx.TrainStateNNX(nnx_model, optimizer)
 
-      init_state_fn = create_train_state_fn
-    else:
-      init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, True, rng)
+    init_state_fn = create_train_state_fn
     unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, True)
 
     lora_config_path = lora_adapter_path + "adapter_config.json"
 
     lora_config = gcs_utils.read_json_from_gcs(lora_config_path)
 
-    if config.pure_nnx:
-      base_abstract_params = _nnx_param_subtree(unboxed_abstract_state)
-      lora_state, lora_state_annotations = get_lora_abstract_state_nnx(base_abstract_params, lora_config)
-    else:
-      lora_state, lora_state_annotations = get_lora_abstract_state(unboxed_abstract_state.params, lora_config)
+    base_abstract_params = _nnx_param_subtree(unboxed_abstract_state)
+    lora_state, lora_state_annotations = get_lora_abstract_state_nnx(base_abstract_params, lora_config)
 
     lora_weights_path = f"{lora_adapter_path}/0/items"
 
diff --git a/tests/integration/maxengine_test.py b/tests/integration/maxengine_test.py
index f17a523d52..e49bc8b7a9 100644
--- a/tests/integration/maxengine_test.py
+++ b/tests/integration/maxengine_test.py
@@ -96,8 +96,8 @@ def test_stack_and_unstack_prefill_cache_nnx(self):
   # default. test_basic_prefill_nnx / test_basic_decode_nnx below cover the NNX path.
 
   def _init_nnx_pyconfig(self, **kwargs):
-    """Same as init_pyconfig but with the NNX flags turned on."""
-    return self.init_pyconfig(pure_nnx=True, enable_nnx=True, pure_nnx_decoder=True, **kwargs)
+    """Same as init_pyconfig (NNX is the only path now)."""
+    return self.init_pyconfig(**kwargs)
 
   def _build_nnx_params(self, cfg, mesh):
     """Materialize an NNX Transformer and return its nnx.Param state."""
@@ -173,7 +173,7 @@ def test_basic_decode_nnx(self):
       )
 
   def test_quantize_passes_gate_for_nnx(self):
-    """pure_nnx + quantization (convert-on-load) reaches the actual machinery in train mode."""
+    """NNX + quantization (convert-on-load) reaches the actual machinery in train mode."""
     # checkpoint_is_quantized defaults to False — full-precision on disk, AQT
     # quantizes per-forward against the loaded kernel (train mode).
     cfg = self._init_nnx_pyconfig(quantization="int8")
@@ -187,7 +187,7 @@ def test_quantize_passes_gate_for_nnx(self):
       pass  # any other failure (e.g. checkpoint not found) is fine for this test
 
   def test_load_pre_quantized_nnx_passes_quant_gate(self):
-    """pure_nnx + quantization + checkpoint_is_quantized=True clears the load gate."""
+    """NNX + quantization + checkpoint_is_quantized=True clears the load gate."""
     cfg = self._init_nnx_pyconfig(quantization="int8", checkpoint_is_quantized=True)
     engine = maxengine.MaxEngine(cfg, jax.devices())
     self.assertEqual(engine._nnx_quant_mode_str, "serve")  # pylint: disable=protected-access
@@ -219,7 +219,7 @@ def test_quantized_prefill_nnx_train_mode(self):
     self.assertTrue(jnp.all(jnp.isfinite(prefill_result["logits"])))
 
   def test_lora_load_single_adapter_reaches_loader_on_nnx(self):
-    """pure_nnx + LoRA: load_single_adapter dispatches to the NNX loader.
+    """NNX + LoRA: load_single_adapter dispatches to the NNX loader.
 
     A nonexistent adapter path should raise FileNotFoundError from the
     loader itself. A NotImplementedError here would mean the dispatch
diff --git a/tests/post_training/unit/lora_utils_test.py b/tests/post_training/unit/lora_utils_test.py
index 5bbf6f737c..e1110035c5 100644
--- a/tests/post_training/unit/lora_utils_test.py
+++ b/tests/post_training/unit/lora_utils_test.py
@@ -55,8 +55,6 @@
     "base_mlp_dim": 256,
     "max_prefill_predict_length": 4,
     "model_name": "llama2-7b",
-    "enable_nnx": True,
-    "pure_nnx_decoder": True,
     "override_model_config": True,
     "weight_dtype": "bfloat16",
 }
diff --git a/tests/unit/aqt_serve_roundtrip_nnx_test.py b/tests/unit/aqt_serve_roundtrip_nnx_test.py
index 00d6825de7..006a4580d2 100644
--- a/tests/unit/aqt_serve_roundtrip_nnx_test.py
+++ b/tests/unit/aqt_serve_roundtrip_nnx_test.py
@@ -74,9 +74,6 @@ def _init_cfg(self, ckpt_path, *, checkpoint_is_quantized):
         sys.argv[0],
         base_yml,
         "model_name=gpt3-52k",
-        "pure_nnx=true",
-        "enable_nnx=true",
-        "pure_nnx_decoder=true",
         "max_target_length=64",
         "max_prefill_predict_length=16",
         "per_device_batch_size=1",
diff --git a/tests/unit/maxengine_nnx_test.py b/tests/unit/maxengine_nnx_test.py
index 08303f817a..074db10c22 100644
--- a/tests/unit/maxengine_nnx_test.py
+++ b/tests/unit/maxengine_nnx_test.py
@@ -41,9 +41,6 @@ def _nnx_config(self, **kwargs):
         "max_target_length": 8,
         "per_device_batch_size": 1,
         "enable_checkpointing": False,
-        "pure_nnx": True,
-        "enable_nnx": True,
-        "pure_nnx_decoder": True,
     } | kwargs
     return pyconfig.initialize([sys.argv[0], get_test_config_path()], **init_kwargs)