[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 NNX migration prep (1/N): pure_nnx flag and init_state_fn scaffolding #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 NNX migration prep (2/N): NNX utils and sharding utilities #3470)
3. ✅ NNX fully supported end-to-end: TrainStateNNX, model creation, gradient accumulation, checkpointing, and training loop dispatch. (PR NNX migration prep (3/N): TrainState, model creation, and end-to-end training loop #3500)
4. ✅ Sharding diagnostics on NNX, plus post-training bugfixes that surfaced once the NNX path got exercised end-to-end. (PR [NNX] NNX migration prep (4/N): sharding tools and post-training fixes #3652)
   4.5. ✅ Linen↔NNX checkpoint converter. (PR [NNX] NNX migration prep (4.5/N): Linen<->NNX checkpoint converter #3843)
   4.6. ❌ Linen↔NNX checkpoint comparator (sibling branch on PR4.5).
5. ✅ NNX correctness fixes, feature enablements, and vocab tiling on NNX.
6. ✅ NNX-native DPO.
7. 🔄 [This PR] NNX-native MaxEngine inference. Drops the route-to-Linen path in maxengine.py; pure_nnx=True now drives a real NNX flow (two-mode abstract Transformer, nnx.merge/nnx.state per JIT body, pure-dict cache adapter so bulk_insert etc. reuse Linen plumbing). Stacks on PR6. Carve-outs (NotImplementedError'd with PR pointers): NNX+quantization → PR9.5; NNX+LoRA → PR8; prefill_multisampling/prefill_concat/stack_prefill_result_cache=True → follow-up.
8. ❌ NNX-native LoRA + GRPO.
9. ❌ NNX-aware QK-Clip + remaining checkpoint utilities.
   9.5. ❌ NNX + AQT in MaxEngine + serve-mode reload + gpt3 prefill fix.
10. ❌ Vocab tiling custom_vjp for NNX.
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 routed maxengine.py inference to the Linen implementation regardless of pure_nnx, so pure_nnx=True users silently got the Linen engine. This PR replaces that route with a real NNX flow.

When config.pure_nnx=True, the engine builds an NNX Transformer, splits it into (params, cache, rest) with nnx.split, and at every JIT body reconstructs the model with nnx.merge(..., copy=True) to run the forward pass. Linen is preserved byte-for-byte: every NNX edit is gated if config.pure_nnx:, and pure_nnx=False stays the default.

Diff: +404 / −33 across 3 files. The only deletion is the "use Linen path regardless of pure_nnx" comment block in MaxEngine.__init__.

Design notes

• Two abstract models. __init__ builds self.model (PREFILL, batch 1) and self.model_ar (AUTOREGRESSIVE, decode-state batch). NNX cache vars take their logical axis names (CACHE_BATCH vs CACHE_BATCH_PREFILL) from the construction model_mode; the decode-state cache must use AR-mode names so bulk_insert's cache_batch substring lookup hits. nnx.eval_shape is called directly (not via create_nnx_abstract_model) to avoid the jax.set_mesh wrap that trips Flax 0.12.6 on logical-only axes like "norm".
• Cache adapter. The engine's cache helpers (bulk_insert, _insert_jit, _maybe_*_prefill_result_cache) walk the cache by path[-1].key. NNX nnx.Cache state would expose a .value accessor there, so the cache flows as a plain dict on both paths: nnx.State.to_pure_dict() after the run, nnx.replace_by_pure_dict() before nnx.merge. The Linen helpers stay unchanged.
• Param loading (_load_params_nnx): accepts user-provided NNX params or loads via from_pretrained, materializes the rest (RNG) state once for nnx.merge, and builds the prefill/AR KV-cache shardings.
• Annotation helpers (maxtext_utils.py): get_prefill_kv_cache_annotations_nnx / get_kv_cache_annotations_nnx mirror the Linen helpers' return shape (a tree of PartitionSpec).

Carve-outs

These raise NotImplementedError on the NNX path — deliberate scope cuts, not silent fallbacks. Each points at its follow-up; the Linen path (pure_nnx=False) remains the workaround.

Site	Tracked in
quantization (quantize_params, quantized-checkpoint load)	PR9.5
LoRA (load_single_adapter)	PR8
prefill_multisampling, prefill_concat, stack_prefill_result_cache=True	PR7.5

Tests

$ pytest tests/integration/maxengine_test.py
$ pytest tests/unit/maxengine_test.py -v

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

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!