RepoJournal
Google

Google

JAX, the GenAI SDK, and the Cloud libs — Google's open source layer

Pick a date

Topics: AI / ML Python Full archive →

The Wire · Showcase

JAX FIXES ROCM TRITON SUPPORT, CLOUD PYTHON TIGHTENS IMPORT PERFORMANCE

By RepoJournal · Filed · About Google

JAX's Triton integration on AMD GPUs is finally working after resolving a device architecture detection bug that broke every pallas_call on ROCm hardware.

The JAX team shipped a critical fix [1] that resolves ROCm gfx architecture detection, which was causing Triton pallas_call to fail with "No supported GPU devices found" on AMD Instinct hardware. Instead of requiring GpuVersion enum members that don't exist for marketing names like MI355X, the fix now falls back to the device's compute capability (gcnArchName) and uses it directly as the Triton target. Alongside this, JAX expanded copyright enforcement [2] to shell scripts and CMake files, catching licensing gaps in ROCm build tooling. Over in google-cloud-python, the team landed an import profiler CI check [3] that tracks cold-start performance across the entire monorepo, targeting sub-500ms imports for Cloud Run and Functions. The profiler work uncovered a setuptools issue on Python 3.15 that required fixing [4], ensuring namespace package resolution works correctly on newer versions. In python-genai, lazy imports of the MCP SDK landed [5], reducing startup overhead for users who don't use model context protocol features.

Action items

References

  1. [1] [pallas:triton] Resolve ROCm gfx arch from the device ISA ↗ google/jax
  2. [2] Expand copyright check to include shell scripts and CMake files. ↗ google/jax
  3. [3] ci: add import profiler check across monorepo ↗ googleapis/google-cloud-python
  4. [4] fix(ci): install setuptools in temporary virtualenv for import profile (#17714) googleapis/google-cloud-python
  5. [5] chore: defer `mcp` SDK import to first MCP usage ↗ googleapis/python-genai

Quick answers

What shipped in Google on July 15, 2026?
JAX's Triton integration on AMD GPUs is finally working after resolving a device architecture detection bug that broke every pallas_call on ROCm hardware. In total, 40 commits and 35 pull requests landed.
Who contributed to Google on July 15, 2026?
5 developers shipped this update, including copybara-service, magaonka-amd, Wanlin Du, noahdietz, and hebaalazzeh.
What were the notable Google updates?
[pallas:triton] Resolve ROCm gfx arch from the device ISA, Expand copyright check to include shell scripts and CMake files, and ci: add import profiler check across monorepo.

More from Google

Daily updates, in your inbox

Follow Google

JAX, the GenAI SDK, and the Cloud libs — Google's open source layer We'll email you a link to confirm first.

Free. Confirm via email. Unsubscribe in one click.

— or follow the whole beat:

Elsewhere on the wire