Skip to content

Pull requests: openxla/tokamax

New pull request New
30 Open 873 Closed
30 Open 873 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

[Pallas] Make core_map use pl.kernel
#916 opened May 20, 2026 by copybara-service Bot Loading…
Introduce a new tunable param num_stacked_q_heads that specifies how many consecutive Q heads are processed for each KV head in a single block iteration.
#915 opened May 19, 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…
[NFC] Replace no longer needed manual tiling hack.
#910 opened May 13, 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…
Reduce the number of arg specs tested under the ragged_dot API test.
#770 opened Mar 11, 2026 by copybara-service Bot Loading…
Add missing tcgen05.fence instructions to SM100 MGPU attention.
#757 opened Mar 9, 2026 by copybara-service Bot Loading…
Update triangle multiplication benchmark and pin JAX to >=0.9.0.1
#730 opened Mar 2, 2026 by rpachauri Collaborator • Draft
Add cuda-bench group of dependencies to pyproject.toml
#676 opened Feb 19, 2026 by andportnoy Member Loading…
@sbodenstein
Partition and parallelize numpy RNG when initializing large arrays.
#654 opened Feb 13, 2026 by copybara-service Bot Loading…
Add CuTeDSL Blackwell tutorial GEMM
#647 opened Feb 12, 2026 by andportnoy Member Loading…
34
[pallas] Passing backend= in addition to compiler_params= is redundant
#643 opened Feb 12, 2026 by copybara-service Bot Loading…
ProTip! Follow long discussions with comments:>50.