Skip to content

Pull requests: openxla/tokamax

New pull request New
37 Open 921 Closed
37 Open 921 Closed
Author
Filter by author
Loading
Label
Filter by label
Loading
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Loading
Milestones
Filter by milestone
Loading
Reviews
Filter by reviews
No reviews Review required Approved review Changes requested
Assignee
Filter by who’s assigned
Assigned to nobody Loading
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

Update autotuning tests to use Pallas Mosaic GPU for GLU.
#972 opened Jun 23, 2026 by copybara-service Bot Loading…
NotImplementedError to workaround a heisenbug.
#970 opened Jun 22, 2026 by copybara-service Bot Loading…
NotImplementedError to workaround a heisenbug.
#968 opened Jun 19, 2026 by copybara-service Bot Loading…
Internal
#966 opened Jun 18, 2026 by copybara-service Bot Loading…
[trace_to_jaxpr] Create trace_to_jaxpr_nocache
#964 opened Jun 18, 2026 by copybara-service Bot Loading…
Internal
#962 opened Jun 17, 2026 by copybara-service Bot Loading…
1
Fix ragged dot kernels under WG semantics.
#960 opened Jun 17, 2026 by copybara-service Bot Loading…
Made minor kernel change to make sure the kernel is run by internal test framework.
#959 opened Jun 16, 2026 by copybara-service Bot Loading…
Add B200 autotuning cache for Triton and Mosaic ragged dot
#954 opened Jun 16, 2026 by copybara-service Bot Loading…
[Pallas] Make core_map use pl.kernel
#916 opened May 20, 2026 by copybara-service Bot Loading…
Expose ignore_autotuning_cache publicly.
#914 opened May 19, 2026 by copybara-service Bot Loading…
Switch to pyrefly for Python type checking for many Tokamax python libraries.
#912 opened May 18, 2026 by copybara-service Bot Loading…
[pallas:triton] Rolling forward compilation to PTX
#909 opened May 13, 2026 by copybara-service Bot Loading…
Internal
#907 opened May 11, 2026 by copybara-service Bot Loading…
Internal
#905 opened May 8, 2026 by copybara-service Bot Loading…
Fix wrong rhs gradient for empty groups in Pallas-Triton ragged dot.
#889 opened Apr 28, 2026 by renos Loading…
Perform bucketed compute in SM90 MGPU quant ragged dot.
#882 opened Apr 27, 2026 by copybara-service Bot Loading…
Tokamax autotuning cache update. Auto-generated by xmanager/247825154
#878 opened Apr 22, 2026 by copybara-service Bot Loading…
Avoid usage of deprecated jax.core APIs.
#860 opened Apr 15, 2026 by copybara-service Bot Loading…
Change benchmarking timers to run multiple iterations together.
#841 opened Apr 8, 2026 by stzeng Collaborator Loading…
Test
#838 opened Apr 8, 2026 by sbodenstein Collaborator Loading…
Clarify block_mask values in documentation
#804 opened Mar 27, 2026 by Guitaricet Loading…
1
Triton and Mosaic for linear_softmax_cross_entropy_loss
#801 opened Mar 27, 2026 by captainpete Loading…
fix: nested masks
#799 opened Mar 25, 2026 by emergenz Loading…
[Mosaic GPU] Remove the partitioned from MGPU APIs
#777 opened Mar 13, 2026 by copybara-service Bot Loading…
ProTip! Mix and match filters to narrow down what you’re looking for.