XLA Serving v0.0.13

Highlights

• 78 Elementwise Math & Bitwise Data Ops: Added C++ kernels and Python bindings (data.ops.math and data.ops.bitwise) for 78 elementwise math, bitwise, special-function, Bessel, floating-point predicate, and matmul transforms with broadcasting (#610, #611, #612, #634).
• Priority-Aware Batch Scheduling: Added priority-aware request queuing (PriorityAwareQueue, SchedulerConfig, and QueueSelectionPolicy) to the serving runtime's batch scheduler (#599, #602, #615, #619).
• Zero-Copy & Pre-Allocated Buffers: Added support for TPU premapped buffers, pinned host buffers, partial/mixed pre-allocated outputs, input buffer aliasing, and zero-copy tensorflow::Tensor host conversion (#594, #628, #639, #655, #662, #686).
• CLI History & Rollback: Added xsm-cli history and xsm-cli rollback for inspecting and reverting model configuration updates, plus full BatchOptions support in xsm-cli update (#597, #660).

---

What's Changed

Serving Runtime & Performance

• Priority-aware batching: Integrate PriorityAwareQueue, SchedulerScope, QueueSelectionPolicy, and SchedulerConfig into SharedBatchScheduler and BatchRunnerOp (#599, #602, #615, #619).
• Pre-allocated & zero-copy buffers:
  • Support TPU premapped buffers in the serving runtime (#639).
  • Support pinned host buffers and partial/mixed pre-allocated outputs with load-time output name resolution (#628, #655, #662).
  • Honor exported input buffer aliasing in the AOT serving runtime (#594).
  • Add zero-copy tensorflow::Tensor to HostBuffer conversion via a type-keyed converter registry (#686).
• Model warmup improvements: Bypass the batch queue for warmup requests and record model warmup latency metrics (#673, #699, #701).
• Runtime overhead reductions: Cache transform factory lookups in TransformBase (#605), execute RunWithMapReturnOnWorker as a single worker-pool task (#683), and drop unused MLIR dialects from StableHloLoader (#663).
• Self-contained TPU serving image: Bundle libtpu.so directly into the TPU serving container (#640).
• TensorFlow SavedModel interop: Add bundle detection and HostBuffer conversion in TfSavedModelFunctionRunnerOp (#547, #629).

Data Processing (xla_serving_data)

• Math & bitwise transforms: Add 78 elementwise math, bitwise, special-function, Bessel, predicate, and matmul C++ transforms and Python bindings (#610, #611, #612).
• Dynamic shapes: Add PartialShape to represent shapes with unknown dimensions in data/processor and HostBufferSpec (#681, #688).
• CPU embeddings: Unify CPU embedding method signatures with structured input configs (#595, #633).
• Transform registry: Simplify transform deserializer registration with SetDeserializerWithDetail and fall back to std::type_info::name() matching across shared libraries (#635, #649, #664, #665).

CLI (xsm-cli)

• Configuration history & rollback: Add xsm-cli history and xsm-cli rollback commands to inspect and revert model configuration updates (#597).
• Batching configuration: Support full BatchOptions flags in xsm-cli update (#660).
• Model inspection: Improve xsm-cli show and xsm-cli show --details output for CompositeProcessor functions (#607).

Examples & Benchmarks

• DLRM benchmark: Add reproducible DLRM serving benchmark and Cloud TPU v6e results (#606).
• Gemma 3 example: Add Gemma 3 (keras_hub) export example and prediction client by @gaurides (#601).
• Math preprocessor example: Add examples/math_preprocessor demonstrating data.ops.math and data.ops.bitwise (#634).

New Contributors

• @gaurides made their first contribution in #601

Full Changelog: v0.0.12...v0.0.13

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