Warning

Pallas is experimental and is changing frequently.See thePallas Changelog for the recent changes.

You can expect to encounter errors and unimplemented cases, e.g., whenlowering of high-level JAX concepts that would require emulation,or simply because Pallas is still under development.

Guides

• Pallas Quickstart
  • Hello world in Pallas
  • Pallas programming model
• Software Pipelining
  • Memory Hierarchies
  • Pipelining Basics
  • Pallas Pipelining API
  • Sharp edges
  • Analyzing the performance
• Grids and BlockSpecs
  • grid, a.k.a. kernels in a loop
  • BlockSpec, a.k.a. how to chunk up inputs

TPU backend guide

• Pallas TPU
  • Writing TPU kernels with Pallas
  • TPU Pipelining
  • Matrix Multiplication
  • Scalar Prefetch and Block-Sparse Computation
  • Distributed Computing in Pallas for TPUs
  • Pseudo-Random Number Generation

Mosaic GPU backend guide

• Pallas:Mosaic GPU
  • Writing Mosaic GPU kernels with Pallas
  • Mosaic GPU Pipelining
  • Writing high-performance matrix multiplication kernels for Blackwell
  • Collective matrix multiplication

Design Notes

• Pallas Design Notes
  • Pallas Design
  • Pallas Async Operations

Other

• Pallas Changelog