Benjamin Chetioui a9ab614123 [Pallas/Mosaic GPU] Add an abstraction to obtain a slice of dynamic shared memory when using waprgroup semantics. ...

Explicitly make the assumption that `runtime_smem` starts at `0` in the Pallas
module context---which should be enforced by Mosaic GPU.

This is in preparation of changes implementing transform inference.

PiperOrigin-RevId: 732091266

2025-02-28 04:38:25 -08:00

..

mosaic

[Pallas TPU] Use grid_env for pipeline body so we can query num_programs/program_id inside the block spec

2025-02-27 12:53:02 -08:00

mosaic_gpu

[Pallas/Mosaic GPU] Add an abstraction to obtain a slice of dynamic shared memory when using waprgroup semantics.

2025-02-28 04:38:25 -08:00

triton

Change int4 packing from big-endian to little-endian

2025-02-27 08:13:43 -08:00

__init__.py

[Pallas] Upstream pallas to JAX

2023-08-01 16:43:13 -07:00

BUILD

[JAX] Allow pallas to accept scalar shape semaphores.

2025-02-14 23:20:47 -08:00

core.py

[Pallas] Add name parameter to core_map

2025-02-26 18:59:01 -08:00

cost_estimate.py

[better_errors] Continue adding debug info to Jaxprs (step 7)

2025-02-09 18:14:33 +02:00

helpers.py

[Pallas TPU] Add helpers for writing collectives

2025-02-04 15:39:10 -08:00

hlo_interpreter.py

[better_errors] Port the Pallas debug info mechanisms to the new JAX DebugInfo.

2025-02-25 14:43:17 +01:00

pallas_call.py

[Pallas] Add option for async DMAs in the new TPU interpret mode

2025-02-25 18:19:20 -08:00

primitives.py

Refactor Jax FFI lowering to prepare for implementing CPU/GPU callbacks using XLA's FFI.

2025-02-21 09:45:59 -08:00

utils.py

[Pallas TPU] Add helpers for writing collectives

2025-02-04 15:39:10 -08:00