The Wire · Showcase
JAX SWITCHES TO BZLMOD, SHIPS AOT COMPILATION FOR TRITON KERNELS
By RepoJournal · Filed · About Google
JAX's build system migration to Bzlmod is live, paired with major performance improvements for autotuned GPU kernels that compile at instantiation rather than runtime.
The JAX team switched to Bzlmod by default [1], a significant infrastructure move that restructures how pip dependencies are exposed in the build system. Pip packages now route through `@pypi//<package>` rather than standalone repositories, requiring updates to lockfiles and dependency declarations. Simultaneously, JAX shipped AOT compilation for autotuned Triton kernels [2], moving compilation of candidate configurations from JIT time into the instantiate phase. This eliminates serialization overhead since instantiate and execute now happen in the same process. Supporting infrastructure landed across three commits: Triton now handles host-side TMA descriptors [4], the pallas multi-GPU layer auto-infers `PROMISE_IN_BOUNDS` in pipeline emission [5], and requirements lock files were regenerated [3] to accommodate the Bzlmod transition. On the GenAI front, the Python SDK cleaned deprecated Turn types [7], dropped unused auth imports [6], and fixed automatic function calling bugs in streaming contexts [8]. BigFrames hit v2.48.0 with expanded string operation support in the transpiler [9], while BigTable standardized client-side metrics across implementations [10]. The google-cloud-python repo saw two reverts [11] [12] of generator and librarian changes, suggesting recent commits needed correction.
One email a day. Unsubscribe in one click.
Keep up with Google in about 3 minutes a day: what actually shipped — the commits, pull requests, releases, and security advisories that matter.
One email a day. Unsubscribe in one click. Read a past issue →
Action items
- → Review JAX Bzlmod migration impact on your build pipeline if using JAX internally google/jax [plan]
- → Test JAX AOT compilation changes if using Triton kernels with autotuning google/jax [monitor]
- → Update to BigFrames v2.48.0 if you depend on string transpilation googleapis/google-cloud-python [plan]
- → Verify BigTable client-side metrics are aligned with your monitoring expectations googleapis/google-cloud-python [monitor]
References
- [1] Switch to bzlmod by default. ↗ google/jax
- [2] Implement AOT compilation for autotuned JAX-Triton kernels. ↗ google/jax
- [3] Regenerate requirements lock files. ↗ google/jax
- [4] [triton] Added support for host-side TMA descriptors ↗ google/jax
- [5] [pallas:mgpu] Automatically infer `PROMISE_IN_BOUNDS` in `emit_pipeline` ↗ google/jax
- [6] chore: drop unused auth imports from the MCP utils module googleapis/python-genai
- [7] chore: remove deprecated Turn types googleapis/python-genai
- [8] fix: improve AFC(automatic function calling) in chat including fixing bugs in AFC for generate_content_stream. also log warnings in generate_content, generate_content_stream and their async variants t ↗ googleapis/python-genai
- [9] bigframes: v2.48.0 ↗ googleapis/google-cloud-python
- [10] fix(bigtable): standardize client side metrics (#17899) googleapis/google-cloud-python
- [11] Revert "fix(generator): use flat_ref_types in test templates and delete remove-unused-imports" (#18088) googleapis/google-cloud-python
- [12] Revert "chore(librarian): remove post processing script for ces" (#18094) googleapis/google-cloud-python