Pull requests: jax-ml/jax
Pull requests list
[Pallas:MGPU] Use a single barrier to wait for A and B transfers in Blackwell matmul
#31748
by copybara-service
bot
was merged Sep 11, 2025
Loading…
[Pallas:MGPU] Don't use delay_release in Hopper MGPU matmul example
#31736
by copybara-service
bot
was merged Sep 11, 2025
Loading…
[Pallas:MGPU] Use the planar_snake CTA order to improve L2 cache hits in Hopper matmul
#31735
by copybara-service
bot
was merged Sep 11, 2025
Loading…
[Pallas:MGPU] Add support for tiling the TMA epilogue
#31699
by copybara-service
bot
was merged Sep 11, 2025
Loading…
[Pallas:MGPU] Use multiple iterations when benchmarking the Hopper matmul
#31690
by copybara-service
bot
was closed Sep 11, 2025
Loading…
[Mosaic GPU] Fix the incorrect type annotations in the Mosaic GPU profiler
#31743
by copybara-service
bot
was merged Sep 11, 2025
Loading…
[Mosaic GPU] Add support for performing multiple measurements in a single cupti session
#31691
by copybara-service
bot
was closed Sep 11, 2025
Loading…
[Pallas:MGPU] Add a helper for a 2D grid traversal with better locality
#31705
by copybara-service
bot
was merged Sep 11, 2025
Loading…
[Mosaic GPU] Use all SMs to send the data
#31624
by copybara-service
bot
was merged Sep 8, 2025
Loading…
[Pallas:MGPU] Add async_prefetch support in warp thread semantics
#31629
by copybara-service
bot
was closed Sep 8, 2025
Loading…
[Mosaic GPU] Only send from one SM in each group of SMs that share the same M coordinate
#31623
by copybara-service
bot
was closed Sep 8, 2025
Loading…
[Pallas:MGPU] Support async_prefetch in warpgroup lowering
#31621
by copybara-service
bot
was merged Sep 8, 2025
Loading…
[Mosaic GPU] Improve the async_prefetch test
#31622
by copybara-service
bot
was merged Sep 8, 2025
Loading…
[Pallas:MGPU] Add a primitive for TMA async prefetch
#31620
by copybara-service
bot
was merged Sep 8, 2025
Loading…
[Mosaic GPU] Add support for async TMA prefetch
#31589
by copybara-service
bot
was merged Sep 8, 2025
Loading…
[Pallas:MGPU] Use lax.empty to allocate output buffers for plgpu.kernel
#31594
by copybara-service
bot
was merged Sep 8, 2025
Loading…
[Mosaic GPU] Add support for dumping LLVM IR for the host code
#31588
by copybara-service
bot
was merged Sep 5, 2025
Loading…
[Mosaic GPU][NFC] Refactor async_copy in preparation for async_prefetch
#31587
by copybara-service
bot
was merged Sep 5, 2025
Loading…
[Pallas] Allow semaphore waits that don't decrement the semaphore value
#31512
by copybara-service
bot
was merged Sep 5, 2025
Loading…
[Pallas:MGPU] Use a protocol to clearly indicate the new method on the WS pipeline
#31559
by copybara-service
bot
was merged Sep 5, 2025
Loading…
[Pallas:MGPU] Make the MGPU kernel safe no matter how many SMs are mapped to N
#31511
by copybara-service
bot
was closed Sep 4, 2025
Loading…
[Pallas:MGPU] Allow manual allocations in the WS pipeline
#31549
by copybara-service
bot
was merged Sep 4, 2025
Loading…
Skip ASAN/MSAN in linalg_test due to new memory leaks in SciPy
#31548
by copybara-service
bot
was merged Sep 4, 2025
Loading…
[Pallas:SC] Add SparseCore tests to presubmit, fix the currently broken targets
#31542
by copybara-service
bot
was merged Sep 4, 2025
Loading…
ProTip!
Type g p on any issue or pull request to go back to the pull request listing page.