XLA Serving v0.0.12

XLA Serving v0.0.12

The release that finishes the move from Neptune to XLA Serving. Expect renamed packages, modules and binaries; see the migration notes below. Alongside the rename: ahead-of-time (AOT) serving without a JIT, priority-aware batching, faster host-to-device transfers, and a new round-trip from XLA SavedModel back to JAX.

⚠️ Breaking changes: Neptune → XLA Serving

Before	After	PRs
neptune_model / neptune-model wheel	xla_saved_model / xla-saved-model	#470
neptune_model_runner wheel	xla_saved_model_runner	#528
neptune_data	xla_serving_data	#487
npt-cli	xsm-cli (npt-cli kept as an alias)	#470
neptune_model_* / neptune_bundle_* protos	xla_saved_model_* / xla_serving_bundle_*	#484, #511
neptune Bazel module, runtime/neptune/…	xla_serving, runtime/xla_serving/…	#483, #476
neptune_main_* server targets and Docker images	xla_serving_main_*	#504
neptune_warmup_requests	xla_serving_warmup_requests	#536
NEPTUNE_ header guards, neptune_stream FFI target	XLA_SERVING_, xla_serving_stream	#482, #534

Remaining NPT/Neptune references in docs, metrics and flags were cleaned up in #533, #535, #548, #576, #577.

✨ Highlights

Ahead-of-time serving

• AOT StableHLO execution no longer depends on JIT compilation or LLVM codegen (#518); runtime/interpreter is split into :interpreter_core plus opt-in capabilities (#517).
• New xsm-cli recompile subcommand for offline compilation, with automatic deviceless-compilation detection (#412, #499).
• Target-topology routing wired into ModelRunner, with on-disk AOT tests (#465).
• End-to-end AOT benchmark suite under examples/aot_benchmark (#500, #549) and AOT user and deployment guides (#501).

Batching and queues

• PriorityAwareQueue with PriorityBatchingPolicy, plumbed through QueueOptions (#407, #473, #475, #485, #492); SharedBatchScheduler can switch to QueueBase behind a flag (#463).
• pad_variable_length_inputs batch option: export, interpreter and batching-kernel support (#496, #508, #552).

Host-to-device transfers

• BatchedH2DTransferExecutor supports partitioned and tiled sharding, custom byte strides and single-device HloSharding normalization (#469, #543, #546).
• JAX host-resident placement (memory_kind='pinned_host' / 'host') is bridged through (#464).
• [N, H, W, C] inputs with small C are transposed to [N, C, H, W] during R1 staging (#571); kImmutableZeroCopy host buffers (#585); a distinct IFRT stream per StableHLO execution (#580); no more per-input SingleDeviceSharding rebuilds (#586).

Export and JAX round-tripping

• Load exported functions and weights back into JAX: ExportedApplyFn, xsm_fn_to_apply_fn, and the stitching tools (#481, #489, #598).
• The PyTree structure of an exported function is recorded in the JAX supplemental and used when rebuilding, including custom nodes without registration (#531, #532, #551).
• TF functions returning tracked containers keep their output structure (#553, #565).
• Input aliasing (#567), float4e2m1fn / float8e8m0fnu dtypes (#530), and a single model-version file (#495).

Runtime and server

• Multi-model serving in OrbaxCallbackGrpcServer (#468) and --model_config_file in the serving binary (#497).
• NanoRt is the default CPU runtime (#581); server binaries link tcmalloc (#584).
• Faster CPU embedding lookups: mapped tables, grouped scheduling, SVE, inline small batches (#570, #573, #574).
• enable_input_r1_optimization / enable_output_r1_optimization on the Python ModelRunner (#568).

CLI

• xsm-cli update for changing a pipeline's max batch size in place (#494).

🐛 Fixes

• Predict RPCs now finish off the completing interpreter thread (#526) — thanks @hyunminch!
• Data race in ThreadSafeStatus::status() (#578).
• Cache key construction in AotStableHloRunner::PopulateExecutableBundles (#498).
• Hardcoded PredictionService gRPC method path (#539).
• Segfault when TensorFlow is imported before xla_serving_data (#556).
• nanobind module initialization in composite_processor_pybind (#521).
• Legacy Orbax JAX supplemental MIME type accepted by the stitching loader (#522).

🔧 Build, CI and packaging

• Python 3.13 and 3.14 in presubmits and the release matrix (#579, #583).
• Every example runs end to end in the OSS export and in CI (#569); build-time AOT test models and capability suites (#520).
• XLA pin moved to third_party/xla/xla.MODULE.bazel and bumped to a recent green commit, dropping the bzlmod patch (#589, #590, #563).
• Bazel remote cache improvements (#486); ml-build container pinned by digest (#529); jax<0.11.2 pinned in CI pending a Flax fix (#544).

New Contributors

• @hyunminch made their first contribution in #526

Full Changelog: v0.0.11...v0.0.12