-
Notifications
You must be signed in to change notification settings - Fork 47
Pull requests: openxla/tokamax
Author
Label
Projects
Milestones
Reviews
Assignee
Sort
Pull requests list
Use pl.BoundedSlice instead of integer for pl.BlockSpec allowing dynamic memory allocation.
#1205
opened Aug 12, 2026 by
copybara-service
Bot
Loading…
Move
l reciprocal calculation outside of SM90 attention VJP.
#1200
opened Aug 12, 2026 by
copybara-service
Bot
Loading…
Use default lowering semantics for GLU kernel.
#1125
opened Jul 29, 2026 by
copybara-service
Bot
Loading…
Scale the Pallas-Triton attention VJP block sizes with the head dim.
#1122
opened Jul 29, 2026 by
bzantium
Loading…
Support logits soft-capping in the Pallas-Triton attention VJP.
#1121
opened Jul 29, 2026 by
bzantium
Loading…
Add jax_triton implementation of normalization op
#1071
opened Jul 20, 2026 by
samanklesaria
Loading…
Centralize unsigned power-of-2 arithmetic helpers in Tokamax and optimize GMM V2.
#1022
opened Jul 9, 2026 by
copybara-service
Bot
Loading…
Use sliced refs, instead of indexer parameter, for various
plgpu.loads.
#1007
opened Jul 7, 2026 by
copybara-service
Bot
Loading…
[DO NOT REVIEW] Verify Tokamax OSS workflow
#995
opened Jun 29, 2026 by
vanbasten23
Contributor
Loading…
Add static calibration option to gmm v2.
#993
opened Jun 26, 2026 by
vanbasten23
Contributor
Loading…
Support both FFI and legacy Triton custom call targets in HLO extraction to prevent failures during JAX FFI transition.
#991
opened Jun 26, 2026 by
copybara-service
Bot
Loading…
Add "update" mode to Tokamax autotuner runner to allow incremental cache population.
#990
opened Jun 26, 2026 by
copybara-service
Bot
Loading…
Make
peak_memory_mb a required field in BenchmarkResult.
#988
opened Jun 26, 2026 by
copybara-service
Bot
Loading…
docs: add triangle multiplication (AlphaFold2 Evoformer) documentation
#984
opened Jun 24, 2026 by
Ashutosh0x
Loading…
docs: add comprehensive FlexAttention documentation page
#983
opened Jun 24, 2026 by
Ashutosh0x
Loading…
docs: add dedicated documentation for Gated Linear Unit (SwiGLU) and Layer Normalization (RMSNorm)
#982
opened Jun 24, 2026 by
Ashutosh0x
Loading…
chore: annotate bare TODO/FIXME comments with descriptive context across attention module
#981
opened Jun 24, 2026 by
Ashutosh0x
Loading…
docs: add comprehensive ragged_dot (GMM/MoE) documentation covering V1 and V2 kernels
#980
opened Jun 24, 2026 by
Ashutosh0x
Loading…
Previous Next
ProTip!
Updated in the last three days: updated:>2026-08-12.