Sergei Lebedev 498e81ab10 Pallas now exclusively uses XLA for compiling kernels on GPU ...

The old lowering pass via Triton Python APIs has been removed and the
JAX_TRITON_COMPILE_VIA_XLA environment variable no longer has any effect.

PiperOrigin-RevId: 621857046

2024-04-04 07:47:26 -07:00

..

ops

Introduce an "inline_seq_dim" mode to paged attention kernel. The mode will fuse kernel instances along the sequence dimension into one kernel, hence reducing the number of kernels.

2024-04-03 13:02:19 -07:00

__init__.py

Register GPU/TPU lowering for pallas_call_p lazily

2024-04-01 14:40:33 -07:00

gpu.py

Pallas now exclusively uses XLA for compiling kernels on GPU

2024-04-04 07:47:26 -07:00

tpu.py

[Pallas TPU Pipeline] Add support for arbitrary manual prefetches in pipeline callbacks + compute execution.

2024-02-26 18:44:35 -08:00