Commits

Breadcrumbs

History for

onmain

User selector

All users

All time

Commit History

Commits on Mar 25, 2026

• Remove mentions to megablox.

  Show description for 117a409
  

  sbodenstein

  authored and

  copybara-github

  committed
  117a409

  Copy full SHA for 117a409

• Remove hardcoded heuristics for Pallas Mosaic TPU ragged dot.

  Show description for 424ec7d
  

  niyatic21

  authored and

  copybara-github

  committed
  424ec7d

  Copy full SHA for 424ec7d

Commits on Mar 24, 2026

• Use reshape for triangle multiplication weights.

  Show description for ad0fc7d
  

  sbodenstein

  authored and

  copybara-github

  committed
  ad0fc7d

  Copy full SHA for ad0fc7d

• No public description

  Show description for fe070cb
  

  niyatic21

  authored and

  copybara-github

  committed
  fe070cb

  Copy full SHA for fe070cb

• Prepare GLU op for backend-specific heuristics or tuning spaces.

  Show description for 196c3a5
  

  Giorgio Arena

  authored and

  copybara-github

  committed
  196c3a5

  Copy full SHA for 196c3a5

Commits on Mar 23, 2026

• Fix benchmarks.

  Show description for 9b95a08
  

  sbodenstein

  authored and

  copybara-github

  committed
  9b95a08

  Copy full SHA for 9b95a08

• Remove hardcoded heuristics for Pallas Mosaic TPU ragged dot.

  Show description for 2dac768
  

  sbodenstein

  authored and

  copybara-github

  committed
  2dac768

  Copy full SHA for 2dac768

• Add benchmark for attention autotuning.

  Show description for c630b22
  

  sbodenstein

  authored and

  copybara-github

  committed
  c630b22

  Copy full SHA for c630b22

Commits on Mar 20, 2026

• Switch benchmarks to use tokamax public benchmarking functions.

  Show description for e2c26f0
  

  sbodenstein

  authored and

  copybara-github

  committed
  e2c26f0

  Copy full SHA for e2c26f0

• Add 'mlcompass' tag to some AlphaFold attention benchmarks.

  Show description for 9bbd688
  

  niyatic21

  authored and

  copybara-github

  committed
  9bbd688

  Copy full SHA for 9bbd688

Commits on Mar 19, 2026

• Fix ragged_dot benchmarks

  Show description for 9643540
  

  niyatic21

  authored and

  copybara-github

  committed
  9643540

  Copy full SHA for 9643540

• Update tokamax version to 0.0.12

  Show description for 4936e75
  

  niyatic21

  authored and

  copybara-github

  committed
  4936e75

  Copy full SHA for 4936e75

• Expose `ragged_dot_general` in the top-level tokamax API

  Show description for a577399
  

  Tokamax team

  authored and

  copybara-github

  committed
  a577399

  Copy full SHA for a577399

• [tokamax] Moving GLU to OSS.

  Show description for 7266f68
  

  cperivol

  authored and

  copybara-github

  committed
  7266f68

  Copy full SHA for 7266f68

• [tokamax] Add support for tokamax SM100 attention bias.

  Show description for 0e2dce2
  

  cperivol

  authored and

  copybara-github

  committed
  0e2dce2

  Copy full SHA for 0e2dce2

Commits on Mar 18, 2026

• Switch ragged dot arg specs to use default group sizes.

  Show description for c24a9b8
  

  niyatic21

  authored and

  copybara-github

  committed
  c24a9b8

  Copy full SHA for c24a9b8

Commits on Mar 17, 2026

• Suppressed pytype errors in uses of MLIR APIs

  Show description for 8cba6a6
  

  superbobry

  authored and

  copybara-github

  committed
  8cba6a6

  Copy full SHA for 8cba6a6

Commits on Mar 13, 2026

• Enable Pallas kernel to take both bfloat16 or float32 as input and accumulate in float32.

  Show description for 076d001
  

  Tokamax team

  authored and

  copybara-github

  committed
  076d001

  Copy full SHA for 076d001

• [Pallas:MGPU] Remove partitioned_axis from copy_gmem_to_smem

  Show description for 869c892
  

  apaszke

  authored and

  copybara-github

  committed
  869c892

  Copy full SHA for 869c892

Commits on Mar 12, 2026

• No public description

  Show description for 51396a3
  

  niyatic21

  authored and

  copybara-github

  committed
  51396a3

  Copy full SHA for 51396a3

• No public description

  Show description for 66c9a6d
  

  niyatic21

  authored and

  copybara-github

  committed
  66c9a6d

  Copy full SHA for 66c9a6d

• Enable Tokamax test on TPU.

  Show description for 694f2a1
  

  niyatic21

  authored and

  copybara-github

  committed
  694f2a1

  Copy full SHA for 694f2a1

Commits on Mar 11, 2026

• Add ragged_dot arg_spec for MaxText.

  Show description for 967cd5b
  

  sbodenstein

  authored and

  copybara-github

  committed
  967cd5b

  Copy full SHA for 967cd5b

• Don't canonicalize `group_sizes` to `GroupSizes` in ragged dot.

  Show description for cd7b7d7
  

  chr1sj0nes

  authored and

  copybara-github

  committed
  cd7b7d7

  Copy full SHA for cd7b7d7

• Make tokamax version programmatically queryable.

  Show description for f0cf073
  

  sbodenstein

  authored and

  copybara-github

  committed
  f0cf073

  Copy full SHA for f0cf073

• Support f32 output from SM100 MGPU attention.

  Show description for a7e9525
  

  chr1sj0nes

  authored and

  copybara-github

  committed
  a7e9525

  Copy full SHA for a7e9525

Commits on Mar 10, 2026

• Minor fix in MGPU attention test.

  Show description for cebb519
  

  chr1sj0nes

  authored and

  copybara-github

  committed
  cebb519

  Copy full SHA for cebb519

• No public description

  Show description for e4e889c
  

  niyatic21

  authored and

  copybara-github

  committed
  e4e889c

  Copy full SHA for e4e889c

• Cleanup mask code in SM100 MGPU attention.

  Show description for 3a225ac
  

  chr1sj0nes

  authored and

  copybara-github

  committed
  3a225ac

  Copy full SHA for 3a225ac

• Re-enable f32 inputs in SM100 MGPU attention tests.

  Show description for 1cf2280
  

  chr1sj0nes

  authored and

  copybara-github

  committed
  1cf2280

  Copy full SHA for 1cf2280

• Pipeline epilogue store to gmem in SM100 MGPU attention.

  Show description for 59a2948
  

  chr1sj0nes

  authored and

  copybara-github

  committed
  59a2948

  Copy full SHA for 59a2948

• Remove redundant warpgroup barrier in SM100 MGPU attention.

  Show description for d7ab1fd
  

  chr1sj0nes

  authored and

  copybara-github

  committed
  d7ab1fd

  Copy full SHA for d7ab1fd

Commits on Mar 9, 2026

• Improve the experimental TPU ragged dot tuning script.

  Show description for 1cb77fc
  

  rdyro

  authored and

  copybara-github

  committed
  1cb77fc

  Copy full SHA for 1cb77fc

• Fix shape annotations in Tokamax linear softmax cross entropy loss.

  Show description for f0f5eab
  

  sashabu

  authored and

  copybara-github

  committed
  f0f5eab

  Copy full SHA for f0f5eab

Commits on Mar 7, 2026

• Support a per-batch, per-head additive attention bias.

  Show description for a574b66
  

  rdyro

  authored and

  copybara-github

  committed
  a574b66

  Copy full SHA for a574b66