Pallas currently lowers to Mosaic on TPUs and Mosaic GPU on newer NVIDIA GPUs. There is also a Triton GPU backend, but it is maintained on a best-effort basis and not recommended for new use. Mosaic GPU is currently aimed at Hopper and newer GPUs.
In operating systems, “kernel” means the core program managing the machine. In numerical programming, a kernel is different, it is a small function intended to run close to the hardware, often many times in parallel.
A good mental model is:
A kernel is not “the whole computation”. A kernel is the work done by one worker on one piece of the computation.
For example, if you add two vectors of length 1024, ordinary JAX lets you write:
