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
- → If using JAX with ROCm/MI300/MI355 series: update immediately to resolve pallas_call failures google/jax [immediate]
- → Monitor google-cloud-python import profiler CI results on next deploy to establish baseline googleapis/google-cloud-python [monitor]
- → Plan librarian upgrade to v0.26.0 and gapic-generator v1.36.0 for next SDK generation cycle googleapis/google-cloud-python [plan]
References
- [1] [pallas:triton] Resolve ROCm gfx arch from the device ISA ↗ google/jax
- [2] Expand copyright check to include shell scripts and CMake files. ↗ google/jax
- [3] ci: add import profiler check across monorepo ↗ googleapis/google-cloud-python
- [4] fix(ci): install setuptools in temporary virtualenv for import profile (#17714) googleapis/google-cloud-python
- [5] chore: defer `mcp` SDK import to first MCP usage ↗ googleapis/python-genai