Record · corpus.blog/posts/01a08763-66da-707f-adad-3b424c608a25

Thinking in Pallas - Sharded MatMuls

considerthebulldog.com · published 31 March 2026

A record is what was published and who pointed at it. The text of the post is not held here: read it at the source, or at its Wayback capture.

Read the post at considerthebulldog.com Wayback capture

0blogs citing

First cited: —
Most recent: —
Rank this month: not ranked
External links: 32
Words captured: 18,796

Cited by

No blog we hold has cited this post yet.

Links from this post

Link

Host

https://assets-prd.ignimgs.com/2022/04/14/3cce6b1b-387a-4dd7-991d-8b8104dd6c36-1649971120552.png

assets-prd.ignimgs.com

https://github.com/R-Oliver/sfp

github.com

https://github.com/jax-ml/jax/blob/main/jax/experimental/pallas/ops/tpu/paged_attention/paged_attention_kernel.py

github.com

https://github.com/jax-ml/jax/blob/main/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py

github.com

https://github.com/jax-ml/jax/blob/main/jax/experimental/pallas/ops/tpu/megablox/gmm.py

github.com

https://github.com/openxla/tokamax/tree/main

github.com

https://github.com/AI-Hypercomputer/maxtext

github.com

https://tensorpro.github.io/matmul-visuals-site/visuals.html

tensorpro.github.io

https://jax-ml.github.io/scaling-book/tpus/

jax-ml.github.io

https://github.com/jax-ml/jax/blob/fc1b3213cf2dd1c7ed1143f3521c8e682df630a4/jax/_src/pallas/mosaic/primitives.py

github.com

https://docs.jax.dev/en/latest/pallas/tpu/matmul.html

docs.jax.dev

https://github.com/jax-ml/jax/blob/main/jax/_src/pallas/mosaic/tpu_info.py

github.com

https://github.com/jax-ml/jax/blob/aa9977f069d30d918cf229f4f5dec3ba8d30f9fd/jax/_src/pallas/mosaic/core.py

github.com

https://docs.jax.dev/en/latest/pallas/tpu/distributed.html

docs.jax.dev

https://jax-ml.github.io/scaling-book/roofline/

jax-ml.github.io

https://jax-ml.github.io/scaling-book/sharding/

jax-ml.github.io

https://jax-ml.github.io/scaling-book/transformers/

jax-ml.github.io

https://jax-ml.github.io/scaling-book/profiling/

jax-ml.github.io

https://patricktoulme.substack.com/p/from-jax-to-vliw-tracing-a-computation

patricktoulme.substack.com

https://patricktoulme.substack.com/p/when-xla-isnt-enough-from-pallas

patricktoulme.substack.com

https://www.youtube.com/watch?v=Q1Y8B_XQn9k&t=1827s

youtube.com

https://www.youtube.com/watch?v=XIVPdSkLbT4

youtube.com

https://accelerated-computing.academy/fall25/labs/lab11/

accelerated-computing.academy

https://github.com/jax-ml/jax/blob/main/jax/experimental/pallas/g3doc/debugging.md

github.com

https://henryhmko.github.io/posts/nsa_tpu/nsa_tpu.html

henryhmko.github.io

https://sdbuchanan.com/blog/jax-1/

sdbuchanan.com

https://github.com/AI-Hypercomputer/maxtext/blob/ed517cf80d9aa81f76e236c5516dacebfe39e96d/benchmarks/xla_flags_library.py

github.com

https://docs.cloud.google.com/tpu/docs/performance-guide

docs.cloud.google.com

https://openxla.org/xla/flags_guidance

openxla.org

https://maxtext.readthedocs.io/en/latest/guides/optimization/benchmark_and_performance.html

maxtext.readthedocs.io

https://www.youtube.com/watch?v=liKrhX2gm44

youtube.com

https://developers.googleblog.com/a-developers-guide-to-debugging-jax-on-cloud-tpus-essential-tools-and-techniques/

developers.googleblog.com