Pallas: Writing TPU kernels in JAX
repo
Fields
id
sources/pallas-writing-tpu-kernels-in-jaxPublisher
Google / JAX
Kind
repo
Accessed
2026-10-04
notes
Pallas TPU documentation. Covers memory spaces (HBM/VMEM), pl.kernel, pipelining via emit_pipeline, matmul with 91% FLOP/s utilization at 8192^3 bf16, SparseCore support, distributed computing with RDMA.
W10 LOWCONFIDENCE PROVENANCE AUDIT 2026-10-10: re-classified UNCHECKED after an actual re-fetch. Fetched from this host: HTTP 200 with substantive body text. WHAT WAS READ: JAX Pallas TPU docs page, live.. The page is first-party for the named publisher and live. Raised 0.5 -> 0.85, status draft -> verified. This record remains a SELF-DESCRIPTION page, not an independent measurement, so it does not reach the 0.9+ tier reserved for verified against-primary-source numbers.
Notes
Pallas TPU documentation. Covers memory spaces (HBM/VMEM), pl.kernel, pipelining via emit_pipeline, matmul with 91% FLOP/s utilization at 8192^3 bf16, SparseCore support, distributed computing with RDMA.
W10 LOWCONFIDENCE PROVENANCE AUDIT 2026-10-10: re-classified UNCHECKED after an actual re-fetch. Fetched from this host: HTTP 200 with substantive body text. WHAT WAS READ: JAX Pallas TPU docs page, live.. The page is first-party for the named publisher and live. Raised 0.5 -> 0.85, status draft -> verified. This record remains a SELF-DESCRIPTION page, not an independent measurement, so it does not reach the 0.9+ tier reserved for verified against-primary-source numbers.
Cross-references
This record references no other records.