Skip to content

CI: 06/06/25 upstream sync #460

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 1,907 commits into from
Closed

CI: 06/06/25 upstream sync #460

wants to merge 1,907 commits into from

Conversation

rocm-repo-management-api-2[bot]
Copy link

Daily sync with upstream

jakevdp and others added 30 commits May 23, 2025 10:11
@jakevdp
[tree_util] raise more informative error when pytree equality check f…
aa63a15 
…ails
@justinjfu @Google-ML-Automation
[Mosaic GPU] Add support for copy_gmem_to_smem in Warp semantics.
d4ab826 
PiperOrigin-RevId: 762475094
@Google-ML-Automation
Merge pull request jax-ml#28978 from matthiasdiener:patch-1
f429162 
PiperOrigin-RevId: 762490961
@justinjfu @Google-ML-Automation
[Mosaic GPU] Add barrier transformation support to tcgen05_mma.
c4a90c1 
Also fix accumulator argument when it's dynamic.

PiperOrigin-RevId: 762509416
@Google-ML-Automation
Move jax/_src/custom_dce.py to its own BUILD rule
f5a9d46 
Creating smaller build rules enforces better organized dependency graphs in the JAX project, helps pytype propagate annotations correctly, and leads to improved build and iteration times.

This was unblocked by moving batching & ad to their own rules in prior changes.

PiperOrigin-RevId: 762527517
@pschuh @Google-ML-Automation
Transfer library: poison outstanding buffer fetches upon connection f…
9153ab7 
…ailure.

PiperOrigin-RevId: 762546985
@ZacCranko @Google-ML-Automation
This is a change to patch some internal Google builds while we comple…
57d07e1 
…te a refactor.

PiperOrigin-RevId: 762547026
@vfdev-5
TSAN CI, make jax buid/test step fail if missing deps wheels
ae2f943
@Google-ML-Automation
Merge pull request jax-ml#28979 from jakevdp:pytree-err
704c3c6 
PiperOrigin-RevId: 762562570
@rdyro @Google-ML-Automation
Simplify attention VJP definition
292dea6 
PiperOrigin-RevId: 762567722
@Google-ML-Automation
[ragged-paged-attn] Implement static kv cache quantization. (The scal…
966bcb9 
…e of kv cache is a scalar float value)

PiperOrigin-RevId: 762576286
@Google-ML-Automation
Move jax/_src/tree.py to its own build rule
2b9d7c8 
Creating smaller build rules enforces better organized dependency graphs in the JAX project, helps pytype propagate annotations correctly, and leads to improved build and iteration times.

PiperOrigin-RevId: 762589488
@dlwh
Update index.rst
4eb3220 
add marin to ecosystem
@Google-ML-Automation
Move jax/_src/sourcemap to its own build rule
d0195f2 
Creating smaller build rules enforces better organized dependency graphs in the JAX project, helps pytype propagate annotations correctly, discourages private imports downstream, and leads to improved build and iteration times.

PiperOrigin-RevId: 762621491
@jburnim @Google-ML-Automation
Use block_until_ready to fix races in TPU interpret mode tests
0833cc2 
PiperOrigin-RevId: 762621608
@Google-ML-Automation
Merge pull request jax-ml#28992 from dlwh:marin-link
29f9905 
PiperOrigin-RevId: 762623893
@mrguenther @Google-ML-Automation
Update XlaCallModule so tests are compatible with DCE.
f28565d 
Move the logic verifying that `shape_assertion` custom calls have side effects to run before MLIR optimizations are applied instead of after. Any `shape_assertion` custom call violating this condition (i.e. declared as pure) is likely to be removed by dead-code elimination, making it undetectable after optimizations. (Until recently, the test passed because DCE wasn't correctly applied to `custom_call` ops.)

PiperOrigin-RevId: 762634648
@Google-ML-Automation
Update XLA dependency to use revision
8b54a6d 
http://github.com/openxla/xla/commit/bc63ed41661e939ced1b3d3bfd7d3083eaabd747.

PiperOrigin-RevId: 762804332
@yashk2810 @Google-ML-Automation
Enter into the right mesh context during shmap DCE
f6b8cb6 
PiperOrigin-RevId: 762956847
@Google-ML-Automation
Update XLA dependency to use revision
37068b6 
http://github.com/openxla/xla/commit/2c35b7a03388e9ee81b3a71037c372595232ff84.

PiperOrigin-RevId: 763076180
Clarify that upper bound takes precedence in jnp.clip where bounds ar…
c22bba2 
…e incongruent
@apaszke @Google-ML-Automation
[Mosaic GPU] Add missing allocator config and skips in one of our dis…
f9c7a14 
…tributed tests

I added them in all other files, but forgot about this one.

PiperOrigin-RevId: 763352483
@apaszke @Google-ML-Automation
[Pallas:MGPU] Support remote async copies and use them in the collect…
fae05bd 
…ive matmul

PiperOrigin-RevId: 763353415
@Google-ML-Automation
Update XLA dependency to use revision
4bfd163 
http://github.com/openxla/xla/commit/5a5e232f7bb9a2fa0d79f461f86a3cfa2c78f2cf.

PiperOrigin-RevId: 763372229
@mooskagh @Google-ML-Automation
Fix a test which blocks the openxla change.
444e952 
The C128 matmuls will be routed to cuBLAS rather than to be handled by the loop emitter, causing a very slight numerical difference.
Therefore, don't be very strict in the comparison.

PiperOrigin-RevId: 763397887
@andportnoy
[Mosaic GPU] Use PTX ISA version = min(ptxas, LLVM)
c1e8f25
@superbobry @Google-ML-Automation
[pallas] The cf dialect is now always available
f35d708 
PiperOrigin-RevId: 763697379
@Google-ML-Automation
Merge pull request jax-ml#28595 from andportnoy:mosaic-gpu-ptx-isa-fr…
2cbec58 
…om-ptxas-and-llvm

PiperOrigin-RevId: 763701410
@Google-ML-Automation
Update XLA dependency to use revision
3aa4e36 
http://github.com/openxla/xla/commit/cb67f2f7ce4787f63f5fc80dc5c30cd3dee8f4e3.

PiperOrigin-RevId: 763710186
@apaszke @Google-ML-Automation
[Mosaic GPU] Work around MLIR recognizing strided<[1]> as identity la…
f68aab1 
…yout in some ops

I can't explain it, but if we don't do it then the verifier sometimes fails...
I'm not even sure how to properly trigger this in a test right now, but worst case it
would result in more verifier failures to fix, so I think it's fine to merge as is.

PiperOrigin-RevId: 763711454
Google-ML-Automation and others added 26 commits June 5, 2025 01:46
@Google-ML-Automation
Update XLA dependency to use revision
c5ef4e5 
http://github.com/openxla/xla/commit/da63f2971676b6cd97a72ba52883717cf48e13d8.

PiperOrigin-RevId: 767494983
@bchetioui @Google-ML-Automation
[Pallas/Mosaic GPU] Expose the new TCGEN05_COL layout.
986c411 
`TCGEN05_COL` is to `TCGEN05` what `WGMMA_ROW` is to `WGMMA`.

PiperOrigin-RevId: 767500979
@YousefElbrolosy
doc: clarified lack of gpu support for schur and sqrtm
af8e2e3
@hawkinsp @Google-ML-Automation
Don't repeatedly recompute a tuple of axis names for a membership test.
b7a3250 
PiperOrigin-RevId: 767596548
@dimitar-asenov @Google-ML-Automation
[Mosaic GPU] Move should_have_transforms to inference_utils.
59c0171 
PiperOrigin-RevId: 767610707
@justinjfu @Google-ML-Automation
[Pallas][Mosaic GPU] Use separate allocations for collective TMEM.
f8ab209 
The distinction between collective/non-collective TMEM is done at allocation time so we need to allocate two separate blocks.

PiperOrigin-RevId: 767611353
@Google-ML-Automation
Merge pull request jax-ml#29224 from jakevdp:array-impl
9c16c6b 
PiperOrigin-RevId: 767611487
@yashk2810 @Google-ML-Automation
Cache get_vma because it's the same thing we do for get_sharding an…
1c4bb50 
…d there's no need to do the same calculation if the keys don't change.

PiperOrigin-RevId: 767612102
@jakevdp
lax_numpy: move array and asarray to their own submodule
0b899a9 
This leads to a better dependency graph and avoids local imports.
@mattjj @Google-ML-Automation
skip pytype on slow file
ef9d3f8 
PiperOrigin-RevId: 767687785
@mattjj @dougalm
[cleanup] remove core.gensym, and Var.suffix
9fc670e 
Co-authored-by: Dougal Maclaurin <[email protected]>
@Google-ML-Automation
Merge pull request jax-ml#29087 from YousefElbrolosy:doc/sqrtm-schur-…
f51effa 
…not-supported-on-gpu

PiperOrigin-RevId: 767693519
@Google-ML-Automation
Merge pull request jax-ml#29246 from jakevdp:array-refactor
4353f34 
PiperOrigin-RevId: 767694982
@emilyfertig @Google-ML-Automation
Make sure unsupported transfers between multi-process CPU arrays and …
18d0da9 
…TPU/GPU arrays raise a helpful error and don't take the new experimental cross-host device transfer path.

PiperOrigin-RevId: 767696126
@Google-ML-Automation
Merge pull request jax-ml#29273 from mattjj:no-more-gensym
5c33588 
PiperOrigin-RevId: 767702568
@hawkinsp @Google-ML-Automation
Don't recompute source_info.current() in DynamicJaxprTracer.
960f9c5 
PiperOrigin-RevId: 767751945
@justinjfu @Google-ML-Automation
[Pallas][Mosaic GPU] Support column slicing on TMEM.
ff50b5f 
PiperOrigin-RevId: 767755766
@bchetioui @Google-ML-Automation
[Mosaic GPU] Fix 2xf32 -> 2xf8e4m3fn conversion.
013b1b1 
I messed up the first time, the output register is naturally 16-bit wide.
The PTX looks much cleaner now.

PiperOrigin-RevId: 767756118
@justinjfu @Google-ML-Automation
[Pallas][Mosaic GPU] Skip tcgen05 reduce test on WG semantics.
fc43122 
PiperOrigin-RevId: 767767487
@justinjfu @Google-ML-Automation
[Pallas][Mosaic GPU] Add support for load/broadcast using TCGEN05 ROW…
c30ed91 
…/COL layouts.

PiperOrigin-RevId: 767772068
@jakevdp
lax.top_k: raise error if indices will overflow
a5ce3ad
@Google-ML-Automation
Merge pull request jax-ml#29254 from jakevdp:top-k-overflow
f4597d3 
PiperOrigin-RevId: 767827963
@pschuh @Google-ML-Automation
Bring back tree concat optimization for np.array(...)
50d93ee 
PiperOrigin-RevId: 767902551
@Google-ML-Automation
[Mosaic] Adds both direct (where hardware can) support for int8 Trans…
fd43650 
…pose, and canonicalization for where it cannot.

To support Transpose of int8 we need to extend to i32, run transpose with bf16 and truncate.

PiperOrigin-RevId: 767915759
@bchetioui @Google-ML-Automation
[Mosaic GPU] Add support for lowering 2xbf16 -> 2xf8e4m3fn converts.
20d641f 
PiperOrigin-RevId: 767920329
@emilyfertig @Google-ML-Automation
Fix logic for checking supported cross-host device transfers, since t…
e66745d 
…he (unsupported) TPU PjRt C API client also contains the string "TFRT TPU".

PiperOrigin-RevId: 767921346
@rocm-repo-management-api-2 rocm-repo-management-api-2 bot requested a review from a team as a code owner June 6, 2025 06:02
@rocm-repo-management-api-2 rocm-repo-management-api-2 bot enabled auto-merge (rebase) June 6, 2025 06:02
auto-merge was automatically disabled June 19, 2025 16:33

Pull request was closed

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.