Sergei Lebedev 64230d1c93 [pallas:mosaic_gpu] WG lowering now supports while_p ...

PiperOrigin-RevId: 736996154

2025-03-14 14:59:29 -07:00

..

BUILD

Open source TPU-friendly ragged paged attention kernel.

2025-03-06 13:36:45 -08:00

export_back_compat_pallas_test.py

[Mosaic TPU] Guard tests for new features by the libtpu version

2024-12-19 05:04:09 -08:00

export_pallas_test.py

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

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

fuser_block_spec_test.py

[Pallas/Fuser] Add support for custom_call_jvp/pjit to push_block_spec

2025-03-03 17:43:13 -08:00

gpu_attention_test.py

Add a param to not normalize the attention weights

2025-02-25 00:55:58 +00:00

gpu_ops_test.py

Increase the absolute error tolerance to fix flaky tests.

2025-02-07 11:59:13 -08:00

gpu_paged_attention_test.py

Add win32 guard to fix imports on Windows

2025-01-28 10:32:19 -08:00

indexing_test.py

Always suppress the differing_executors Hypothesis health check

2024-12-19 08:00:53 -08:00

mgpu_attention_test.py

[pallas:mgpu] Change FA3 kernel bc lax.div doesn't like mixed types anymore.

2025-02-14 05:10:49 -08:00

mosaic_gpu_test.py

[pallas:mosaic_gpu] WG lowering now supports while_p

2025-03-14 14:59:29 -07:00

ops_test.py

[Pallas] Increase #rows when casting to x2.

2025-03-14 14:32:33 -07:00

pallas_cost_estimate_test.py

[Pallas] Add support for run_state to cost estimator.

2024-12-06 10:36:02 -08:00

pallas_error_handling_test.py

[XLA:Mosaic][Pallas] Enable vector.ExtractOp for non-zero indices.

2024-09-26 13:57:45 -07:00

pallas_jumble_test.py

[pallas] fix jumble test flakiness

2024-12-05 15:31:23 -08:00

pallas_shape_poly_test.py

Reverts 0d058ce86f04a44a51abba1261768fb46edf69d9

2024-07-25 01:50:36 -07:00

pallas_test.py

Add explicit support for float8_e4m3b11fnuz in pl.dot

2025-03-14 02:51:55 -07:00

pallas_vmap_test.py

[pallas] Do not skip vmap tests on GPU when x64 is enabled

2024-11-20 05:08:23 -08:00

tpu_all_gather_test.py

[pallas] Renamed platform-specific tests.

2024-07-10 08:23:06 -07:00

tpu_fusable_matmul_test.py

[Pallas] Add experimental (private for now) API for manual fusion into Pallas kernels

2025-03-03 17:05:51 -08:00

tpu_gmm_test.py

Make gmm TPU kernel tests significantly cheaper

2024-12-17 08:38:59 -08:00

tpu_ops_test.py

[Pallas TPU] Fix preferred_element_type propagation in dot_general with const

2025-03-11 15:06:07 -07:00

tpu_paged_attention_kernel_test.py

Re-enable oss paged attn kernel

2025-02-10 17:47:22 -08:00

tpu_pallas_async_test.py

Disable certain tests on V4 and below.

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

tpu_pallas_distributed_test.py

tpu_pallas_distributed_test: skip unless on TPU

2025-01-27 12:41:28 -08:00

tpu_pallas_interpret_distributed_test.py

[Pallas] Adds a simple dynamic race detector for TPU interpret mode.

2025-03-05 15:15:21 -08:00

tpu_pallas_interpret_test.py

[Pallas] Enable skipping of floating-point operations when interpreting Pallas TPU kernels on CPU.

2025-03-10 15:14:00 -07:00

tpu_pallas_pipeline_test.py

[Pallas TPU] Temporarily strengthen restrictions on Pallas tests

2025-01-15 02:32:14 -08:00

tpu_pallas_random_test.py

[Pallas] Add legacy PRNG key support to Pallas PRNG

2025-03-14 12:30:04 -07:00

tpu_pallas_state_test.py

[Pallas] Use core_map instead of shard_map for Shmallas

2024-10-15 03:26:58 -07:00

tpu_pallas_test.py

Remove spurious zero_to_zero conversion used exclusively for backcompat types as a way of supporting (best effort) unsupported types on hardware to make it easier to debug.

2025-02-27 19:21:37 -08:00

tpu_ragged_paged_attention_test.py

Fix the assumption that pages_per_seq is already a multiple of num_kv_pages_per_blk.

2025-03-11 12:36:33 -07:00

tpu_splash_attention_kernel_test.py

Allow query and keys that aren’t multiples of 128

2025-02-25 19:13:24 +00:00

tpu_splash_attention_mask_test.py

Support dynamic masks in splash attention

2025-01-28 00:14:53 +00:00