feat(gmm): JIT-compilable JAX GMM for TPU with forward/backward and tokamax comparison#161

Open

sii-xinglong wants to merge 15 commits intomainfrom

worktree-wise-swinging-kahan

Commits on Apr 6, 2026

docs: add GMM Pallas kernel design for MoE support

sii-xinglong
and
claude
committed
docs: add GMM Pallas kernel implementation plan

sii-xinglong
and
claude
committed
feat(gmm): add CPU reference implementations for gmm and tgmm

sii-xinglong
and
claude
committed
feat(gmm): add make_group_metadata for kernel scheduling

sii-xinglong
and
claude
committed
feat(gmm): implement GMM and TGMM Pallas TPU kernels

sii-xinglong
and
claude
committed
fix(gmm): add @jax.jit decorators, fix traced grid dimensions

sii-xinglong
and
claude
committed
feat(gmm): add differentiable gmm with custom_vjp and public API

sii-xinglong
and
claude
committed
chore: remove planning docs from PR

sii-xinglong
and
claude
committed
fix(gmm): correct transpose_rhs docstring in gmm_ref

sii-xinglong
and
claude
committed
refactor(gmm): remove Pallas kernels, use bf16 mul + f32 accumulation in refs

sii-xinglong
and
claude
committed
test(gmm): use bfloat16 inputs in CPU reference tests

sii-xinglong
and
claude
committed
chore(gmm): add tokamax as optional TPU dependency

sii-xinglong
and
claude
committed
feat(gmm): add JIT-compilable gmm/tgmm using lax.scan

sii-xinglong
and
claude
committed
feat(gmm): add custom_vjp for differentiable gmm with forward/backward

sii-xinglong
and
claude
committed
test(gmm): add JAX vs tokamax comparison tests for bf16 GMM

sii-xinglong
and
claude
committed