Charles Hofer c7b407c9f0 Merge branch 'rocm-main' into ci-upstream-sync-151_1

2025-03-18 15:27:35 +00: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

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] jnp.reduce_sum now works for >1D arrays

2025-03-17 05:32:07 -07:00

ops_test.py

Removed trivial docstrings from JAX tests

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

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

Removed trivial docstrings from JAX tests

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

Merge branch 'rocm-main' into ci-upstream-sync-151_1

2025-03-18 15:27:35 +00: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

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

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

Removed trivial docstrings from JAX tests

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

tpu_splash_attention_mask_test.py

Removed trivial docstrings from JAX tests

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