Sergei Lebedev 4f00249aa8 [pallas:mosaic_gpu] Do not specify the default index_map in tests ...

PiperOrigin-RevId: 743816110

2025-04-03 22:31:54 -07:00

..

BUILD

Reduced sharding in various tests.

2025-03-24 15:03:33 -07: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

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07: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

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07: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] Do not specify the default index_map in tests

2025-04-03 22:31:54 -07:00

ops_test.py

[Mosaic GPU] Add warpgroup lowering for Exp2 in Pallas.

2025-03-28 07:22:24 -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

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -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

[pallas] Index Pallas refs instead of using pl.load and pl.store

2025-03-24 13:51:48 -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

Fixed failing ExcessPrecisionTest.test_matmul_f32_out_simple test.

2025-03-28 09:56:09 -07:00

tpu_gmm_test.py

Make gmm TPU kernel tests significantly cheaper

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

tpu_ops_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

tpu_paged_attention_kernel_test.py

Clean up: num_groups = num_q_heads // num_kv_heads

2025-03-28 10:19:16 -07: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

[pallas] Removed pl.device_id. Use lax.axis_index instead.

2025-04-02 15:45:52 -07:00

tpu_pallas_interpret_distributed_test.py

[Pallas] In TPU interpret mode, add initial barrier for kernels without one.

2025-03-24 08:38:02 -07:00

tpu_pallas_interpret_test.py

[Pallas] Skip reads/writes from/to slices of kernel input/output buffers when the slices do not change between iterations of the grid loop that interprets kernels on CPU.

2025-03-27 03:03:25 -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

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

tpu_pallas_state_test.py

Removed trivial docstrings from JAX tests

2025-03-17 07:49:37 -07:00

tpu_pallas_test.py

[Pallas TPU] Support DMA priority in async copy start

2025-04-03 22:26:40 -07:00

tpu_ragged_paged_attention_test.py

Use f32 scratch for output so we only need to transfer output with desired dtype back to HBM.

2025-03-28 15:13:33 -07:00

tpu_splash_attention_kernel_sharded_test.py

Adding sharding support to dynamic masks

2025-03-20 20:19:04 +00:00

tpu_splash_attention_kernel_test.py

Adding sharding support to dynamic masks

2025-03-20 20:19:04 +00:00

tpu_splash_attention_mask_test.py

Adding sharding support to dynamic masks

2025-03-20 20:19:04 +00:00