[NNX] NNX migration prep (9.5/N): NNX + AQT in MaxEngine + serve-mode reload + gpt3 prefill fix#3844

Draft

ecnal-cienet wants to merge 9 commits into

mainfrom

feat/nnx-aqt-maxengine

Commits on May 8, 2026

NNX: add TrainState, model creation utilities, and training loop support
ecnal-cienet
committed
NNX migration prep (4/N): sharding tools + post-training bugfixes
ecnal-cienet
committed
NNX migration prep (4.5/N): Linen<->NNX checkpoint converter
ecnal-cienet
committed
NNX: correctness fixes, enable feature paths, and vocab tiling on NNX
ecnal-cienet
committed
NNX: native DPO (TrainStateNNX.reference_model + dpo_loss_fn_nnx)
ecnal-cienet
committed
NNX: native MaxEngine inference (drop route-to-Linen path in maxengine.py)
ecnal-cienet
committed
NNX: native LoRA + GRPO (drop maxengine LoRA carve-out, drop GRPO pure_nnx warning)
ecnal-cienet
committed
NNX: QK-Clip on NNX + NNX-format checkpoint utilities
ecnal-cienet
committed

Commits on May 11, 2026

NNX: AQT in MaxEngine + serve-mode reload + gpt3 prefill fix
ecnal-cienet
committed