[NNX] NNX migration prep (7/N): NNX-native MaxEngine inference

[ecnal-cienet]

NNX Migration Route Map

1. ✅ Add NNX scaffolding: pure_nnx flag, init_state_fn, TrainStateNNX, NNX utils. Linen workflow unchanged. (PR #3427)
2. ✅ NNX sharding utilities: get_abstract_state_nnx, get_named_sharding_nnx, set_named_sharding_nnx, get_partition_spec_nnx, get_mesh_from_config. (PR #3470)
3. ✅ NNX fully supported end-to-end: TrainStateNNX, model creation, gradient accumulation, checkpointing, and training loop dispatch. (PR #3500)
4. ✅ NNX sharding diagnostics, bidirectional Linen↔NNX checkpoint conversion utilities, and post-training fixes. (PR #3652)
5. ✅ NNX correctness fixes, feature enablements, and vocab tiling on NNX. No-op while pure_nnx=False stays default. (PR #3766)
6. ✅ NNX-native DPO. Closes the only remaining hard NotImplementedError on the NNX path; pure_nnx=True + use_dpo=True is now supported.
7. 🔄 [This PR] NNX-native MaxEngine inference. Drops the route-to-Linen path in maxengine.py; pure_nnx=True now drives a real NNX inference flow end-to-end (prefill, generate, KV cache).
8. ❌ NNX-native LoRA + GRPO.
9. ❌ NNX-aware QK-Clip + remaining checkpoint utilities (standalone_checkpointer, generate_param_only_checkpoint, layerwise_quantization, convert_gpt3_ckpt_from_paxml).
10. ❌ Vocab tiling custom_vjp for NNX (perf optimization, not correctness).
11. ❌ Set NNX defaults to True; regenerate sharding goldens; flip back integration-test pure_nnx=False annotations.
12. ❌ Delete Linen-specific code paths and NNX compatibility flags.

Description

PR5 audited maxengine.py and routed three call sites to the Linen path (state.params shape, state_mesh_shardings.params, and Linen-format inference checkpoint serving), because transformer_as_linen returns a Linen module regardless of pure_nnx. That preserved Linen serving but meant pure_nnx=True users silently got the Linen engine — the cited rationale was "the flag affects training, not inference serving."

This PR replaces those routes with a real NNX inference flow. When config.pure_nnx=True, the engine now builds an NNX Transformer via from_config(rngs=...), splits params/cache/rest with nnx.split(model, nnx.Param, nnx.Cache, ...), and at every JIT body merges the model concretely with nnx.merge(graphdef, params, cache, rest, copy=True) to run the forward pass. Linen is byte-for-byte preserved: every NNX edit is gated if config.pure_nnx:.

The diff is +651 / −34 across 3 files, with the only physical deletion being the "use Linen path regardless of pure_nnx" comment block at the top of MaxEngine.__init__.

Part 1: Engine construction — two abstract NNX models

File	Change
inference/maxengine/maxengine.py (__init__)	Builds two abstract NNX Transformer instances on the NNX path: self.model with model_mode=PREFILL (batch=1, single padded prompt) and self.model_ar with model_mode=AUTOREGRESSIVE (batch=micro_batch_size_to_train_on, decode-state shape). One graphdef (self.graphdef) is cached from the PREFILL model; both abstract models contain initialized nnx.Cache vars (KVCache __init__ allocates both prefill and AR caches when model_mode in {PREFILL, AUTOREGRESSIVE}).
inference/maxengine/maxengine.py (__init__)	Uses nnx.eval_shape directly inside nn_partitioning.axis_rules(config.logical_axis_rules) rather than create_nnx_abstract_model — the latter wraps the trace in jax.set_mesh(mesh) which makes Flax 0.12.6's _to_variable resolve logical axis names against the global mesh and raise on logical-only names like "norm". Same reason get_abstract_state_nnx avoids set_mesh.

The two-model split is required because NNX cache vars take their logical axis names (CACHE_BATCH vs CACHE_BATCH_PREFILL) from the model's construction model_mode. The decode_state cache must use AR-mode annotations so that bulk_insert's cache_batch substring lookup hits.

Part 2: NNX param loading

File	Change
inference/maxengine/maxengine.py (_load_params_nnx)	Accepts user-provided NNX-shape params or loads via from_pretrained. For user-provided params, materializes a concrete model once via _create_model_fn() to capture a real _nnx_rest_state (RNG vars) for nnx.merge. For checkpoint loads, extracts (graphdef, params, cache, rest) from the loaded nnx_model. Refreshes self.graphdef from the concrete model so subsequent merges line up exactly.
inference/maxengine/maxengine.py (_load_params_nnx)	Builds self.abstract_params as ShapeDtypeStructs of the param state, populates self.prefill_kv_cache_annotations and self.kv_cache_annotations (using AR-mode for the latter so cache_batch appears in the per-leaf logical-axes tuple), and wraps both in NamedSharding.

Part 3: NNX init_decode_state

File	Change
inference/maxengine/maxengine.py (_init_decode_state_nnx)	Zero-initializes a pure-dict cache from the AR-mode abstract model (so the leading batch dim matches generate's input shape (per_device_batch_size * mesh.size, 1)). Builds kv_cache_annotations_named per leaf by reading nnx.Cache.metadata["out_sharding" / "sharding" / "sharding_names"] (Flax 0.12.6 renames vary). Returns the same decode_state dict shape as the Linen path.

Part 4: Cache adapter — NNX state ↔ pure dict

NNX cache state is an nnx.State with nnx.Cache-wrapped leaves; tree paths include a .value accessor at the end. Linen's mutable cache is a plain nested dict with the same leaf names but no .value. The engine's bulk_insert / _insert_jit / insert_partial / _maybe_*_prefill_result_cache helpers walk the cache via tree_map_with_path and switch on path[-1].key (the cache var name like "cached_prefill_key").

File	Change
inference/maxengine/maxengine.py (_nnx_cache_state_template / _nnx_init_cache_dict / _nnx_run_model)	Cache flows through the engine as a pure dict on both paths. NNX state is converted at the JIT boundary via nnx.State.to_pure_dict() (after the model run) and nnx.replace_by_pure_dict(template, pure_dict) (before nnx.merge). The cache plumbing helpers see the same shape on Linen and NNX, requiring zero modifications.

Part 5: NNX prefill / generate JIT branches

File	Change
inference/maxengine/maxengine.py (_prefill_jit, _generate_jit)	NNX path replaces self.model.apply(params | {"cache": cache}, ..., mutable=["cache"]) with a call to _nnx_run_model, which does nnx.merge(graphdef, params, cache, rest, copy=True) → model(...) → nnx.state(model, nnx.Cache).to_pure_dict(). copy=True mirrors the train.py diff_wrapper workaround for TraceContextError from reused Variable objects. Linen branch (the else:) is the prior code, unchanged.
inference/maxengine/maxengine.py (_prefill_jit)	When existing_prefix is not None, the NNX branch threads existing_prefix.cache (already a pure dict) directly into _nnx_run_model's cache_dict= arg instead of the Linen params | {"cache": ...} dict-merge — params is an nnx.State, not a dict.

Part 6: KV-cache annotation helpers

File

Change

utils/maxtext_utils.py

New get_prefill_kv_cache_annotations_nnx(abstract_model, config, mesh) and get_kv_cache_annotations_nnx(abstract_model, config, mesh) mirror the Linen helpers' return shape (a tree of PartitionSpec). Both delegate to _nnx_cache_partition_specs, which extracts nnx.Cache state via nnx.split(model, nnx.Cache, ...), calls get_nnx_named_sharding_with_scan_axis inside nn_partitioning.axis_rules(config.logical_axis_rules) so logical axes (layers, cache_batch, norm, ...) resolve to physical mesh axes, then converts the resulting nnx.State to a pure-dict tree via to_pure_dict().

Part 7: Carve-outs (raise NotImplementedError on the NNX path with PR pointers)

Site	Linen workaround	Tracked in
quantize_params (NNX + quantization)	pure_nnx=False	PR9 (NNX-aware checkpoints + quantization)
load_params when checkpoint_is_quantized=True	pure_nnx=False	PR9
load_single_adapter / apply_adapter / unapply_adapter (NNX + LoRA)	pure_nnx=False	PR8 (NNX-native LoRA + GRPO)
_prefill_multisampling_jit (NNX)	pure_nnx=False	Follow-up
prefill_concat (NNX)	pure_nnx=False	Follow-up
load_params when stack_prefill_result_cache=True (NNX)	pure_nnx=False	Follow-up — with scan_layers=True the NNX cache leaves are already stacked on axis 0; the engine's manual-stack helper assumes an unstacked Linen tree shape

These are deliberate scope cuts, not silent fallbacks. Each raises with a message pointing at the responsible follow-up PR. The Linen path remains the workaround for users who hit any of these sites.

Deferred (intentionally out of scope)

• AOT compilation on the NNX path (aot_compile, _compile_generate_and_get_layouts): functions are not gated on pure_nnx and are structure-agnostic, so they likely work as-is — but they are not exercised by PR7's tests. Recommended manual validation when a JetStream-style serving setup first hits this.
• End-to-end gemma2-2b decode parity on a real checkpoint: the in-PR parity test (below) demonstrates equivalence on a 2-layer test config. A python3 -m MaxText.decode pure_nnx=True ... smoke run against gemma2-2b (with an NNX-format ckpt or a converted Linen one) is recommended manual validation before flipping defaults in PR11.

Tests

• New unit tests (tests/unit/maxengine_test.py, 6 tests):
  • test_init_nnx: engine constructs and exposes graphdef + abstract Transformer.
  • test_basic_prefill_nnx: prefill produces the same dict shape as the Linen path; logits + every cache leaf are finite (catches silent NaN/inf from a bad nnx.merge or cache round-trip); per-layer cache leading axis equals num_decoder_layers (catches scan-axis misalignment).
  • test_basic_decode_nnx: prefill → insert → 4-step generate. Verifies next_pos advances by exactly 1 each step (catches off-by-one cache pointer bugs) and logits stay finite at every step.
  • test_linen_nnx_parity_prefill: strongest evidence — same Linen-init weights are bridged into the NNX engine via linen_nnx_converter.py (convert_linen_to_nnx → _strip_value_wrappers → nnx.replace_by_pure_dict), then prefill is run on both engines with the same prompt + greedy. Asserts logits match within bf16 tolerance (rtol=0.05, atol=0.1; the test config uses bf16 compute) and the greedy first-token argmax is exactly the same. A failure here means the NNX forward pass diverged from Linen on identical weights.
  • test_quantize_raises_for_nnx, test_lora_raises_for_nnx: assert the carve-outs fire with NotImplementedError, so users hitting them get a clear "tracked in PR8/PR9" message instead of an opaque crash.
• Existing Linen tests untouched and still pass (test_basic_prefill, test_basic_decode, test_stack_and_unstack_prefill_cache): every NNX edit is gated on if config.pure_nnx: and pure_nnx=False is the default, so the Linen branch is byte-for-byte the prior code.
• Lint: bash lint.sh — codespell + pylint + pyink all green.
• Comprehensive verification doc: PR7_VERIFICATION.md — TL;DR table, full test-run output, design notes (two abstract models, cache adapter, carve-outs), per-test rationale, what the evidence does not cover (manual validation suggestions), and reproduction commands.
• Test summary:

$ pytest tests/unit/maxengine_test.py -v
test_basic_decode             PASSED
test_basic_decode_nnx         PASSED
test_basic_prefill            PASSED
test_basic_prefill_nnx        PASSED
test_chunked_prefill          SKIPPED  (CPU-only; pre-existing skip)
test_init_nnx                 PASSED
test_linen_nnx_parity_prefill PASSED
test_lora_raises_for_nnx      PASSED
test_quantize_raises_for_nnx  PASSED
test_stack_and_unstack_prefill_cache PASSED
================= 9 passed, 1 skipped =================

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 70.00000% with 402 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...ckpoint_conversion/compare_linen_nnx_checkpoint.py	64.61%	99 Missing and 10 partials ⚠️
src/maxtext/trainers/pre_train/train.py	48.06%	69 Missing and 25 partials ⚠️
src/maxtext/utils/vocabulary_tiling.py	2.43%	40 Missing ⚠️
...xtext/checkpoint_conversion/linen_nnx_converter.py	87.82%	22 Missing and 16 partials ⚠️
src/maxtext/utils/maxtext_utils.py	83.76%	14 Missing and 5 partials ⚠️
src/maxtext/layers/nnx_decoders.py	36.00%	15 Missing and 1 partial ⚠️
src/maxtext/utils/sharding.py	72.22%	6 Missing and 9 partials ⚠️
src/maxtext/trainers/diloco/diloco.py	45.45%	10 Missing and 2 partials ⚠️
src/maxtext/trainers/pre_train/train_compile.py	26.66%	10 Missing and 1 partial ⚠️
src/maxtext/utils/train_utils.py	67.64%	11 Missing ⚠️
... and 12 more

📢 Thoughts on this report? Let us know!