The Wire · Showcase
JAX ROLLS BACK PALLAS TRITON CHANGES, ADDS CONFIG GUARDRAILS
By RepoJournal · Filed · About Google
JAX is partially reverting a Pallas Triton PTX compilation change while shipping new configuration controls to prevent auto-pvary insertion from breaking downstream code.
The team partially rolled forward bits needed for Pallas Triton to compile to PTX [1], but this required reverting a prior commit — the change doesn't alter existing Pallas Triton behavior, just restores compilation infrastructure. Parallel to that, JAX is shipping a new config option to disable auto-pvary insertion [2], a necessary escape hatch for teams hitting unexpected behavior from automatic parallelization. The move suggests the auto-pvary feature is too aggressive in its current state and catching real production patterns it shouldn't touch. Minor documentation fixes landed in Pallas pipelining [3], cleaning up typos in the guides. These aren't flashy changes, but they're the friction-reduction work that keeps large compute workloads from breaking between releases.
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
References
- [1] Partially rolling forward the bits needed to compile Pallas Triton to PTX ↗ google/jax
- [2] add config option to disable auto-pvary insertion ↗ google/jax
- [3] Fix typo in Pallas pipelining documentation ↗ google/jax