mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-22 19:46:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
4b49c03523
rocm_jax/tests
History
Jevin Jiang 4b49c03523 Open source TPU-friendly ragged paged attention kernel. 
Key features:
* ***Support mixed prefill and decode*** to increase throughput for inference. (eg., ***5x*** speedup compared to padded Muti-Queries Paged Attention implementation for llama-3-8b.)
* ***No explicit `swapaxes`*** for `seq_len` and `num_head` in pre/post kernel. The kernel takes `num_head` in 2nd minor as it naturally was. We fold swapaxes to strided load/store in the kernel and apply transpose on the fly.
* ***No GMM (Grouped Matmul) Metadata required!*** We calculate the metadata on the fly in the kernel. This can speed up ***10%***!
* ***Increase MXU utilization 8x in GQA*** by grouping shared q heads for MXU in decode.
* ***Minimize recompilation:*** The only factors can cause recompilation are model specs, `max_num_batched_tokens` and `max_num_seqs` in the setting of mixed engine.

PiperOrigin-RevId: 734269519
2025-03-06 13:36:45 -08:00
..
filecheck
Add public APIs for jax.lax monoidal reductions
2025-02-11 16:00:03 -08:00
mosaic
[Mosaic GPU] Remove support for large tiles on Blackwell
2025-03-05 10:34:53 -08:00
notebooks
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
pallas
Open source TPU-friendly ragged paged attention kernel.
2025-03-06 13:36:45 -08:00
testdata
…
third_party/scipy
Fix CI test failure in line_search_test.
2024-05-15 07:38:20 -07:00
ann_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
aot_test.py
[aot] Add support for as_text(debug_info=True).
2025-01-10 07:59:56 +02:00
api_test.py
vmap mismatch size error message: handle *args
2025-03-05 13:08:54 -08:00
api_util_test.py
Consolidate handling of input argument resolution in custom_* APIs.
2024-08-13 00:30:23 -07:00
array_api_skips.txt
array API: improve test coverage
2024-12-05 12:56:07 -08:00
array_api_test.py
Remove deprecated jax.experimental.array_api
2025-01-06 15:19:02 -08:00
array_interoperability_test.py
Increase the minimum NumPy version to v1.25.
2024-12-18 08:18:57 -05:00
array_test.py
Remove _parsed_pspec from everywhere in JAX except for NamedSharding constructor. I'll do that in the next CL since that has a dependency on C++ so needs guards.
2025-02-27 10:17:06 -08:00
attrs_test.py
ref errors
2024-12-18 07:46:14 +00:00
batching_test.py
Make JAX test suite pass (at least most of the time) with multiple threads enabled.
2025-01-10 06:58:46 -08:00
blocked_sampler_test.py
[Pallas] Implement block-invariant sampling.
2024-06-24 11:20:39 -07:00
BUILD
[JAX] Disable msan and asan for the profiler test running on nvidia gpu
2025-03-04 10:34:11 -08:00
cache_key_test.py
Remove code that supported jaxlib < 0.5.
2025-01-17 14:22:27 -05:00
checkify_test.py
Don't monkey-patch functions in test_utils to count events for tests.
2024-12-12 09:58:14 -05:00
cholesky_update_test.py
Bump minimum jaxlib version to v0.4.30.
2024-06-18 12:35:08 -04:00
clear_backends_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
colocated_python_test.py
Remove code present to support jaxlib < 0.5.1.
2025-02-26 07:40:40 -05:00
compilation_cache_test.py
#sdy unskip JAX Shardy tests that are already passing
2025-01-23 09:26:38 -08:00
config_test.py
Set the value of config.jax_platforms directly instead of setting the env variable JAX_PLATFORMS. Setting JAX_PLATFORMS doesn’t do anything because config.jax_platforms is already initialized at that point.
2024-07-22 15:09:15 -07:00
core_test.py
Add public APIs for jax.lax monoidal reductions
2025-02-11 16:00:03 -08:00
cudnn_fusion_test.py
Re-enable cudnn_fusion_test on A100.
2024-11-01 15:48:07 +00:00
custom_linear_solve_test.py
Handle effects in lax.custom_linear_solve.
2025-02-03 11:14:48 -05:00
custom_partitioning_sharding_rule_test.py
[jax:custom_partitioning] Support SdyShardingRule with multiple leading
2025-01-23 08:20:38 -08:00
custom_root_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
debug_info_test.py
[better_errors] Port the Pallas debug info mechanisms to the new JAX DebugInfo.
2025-02-25 14:43:17 +01:00
debug_nans_test.py
Fix shard_map debug_nan leakage of manual out_avals in the impl rules of jit i.e. impl rule of jit saw a manual out_aval which is not expected. This is a band-aid for now with a TODO to do a proper fix
2025-02-24 10:15:21 -08:00
debugger_test.py
[JAX] Add a new jax_num_cpu_devices flag that allows the user to specify the number of CPU directly.
2025-01-08 06:37:44 -08:00
debugging_primitives_test.py
Remove code present to support jaxlib < 0.5.1.
2025-02-26 07:40:40 -05:00
deprecation_test.py
[internal] add deprecation test utilities
2025-01-10 11:54:09 -08:00
device_test.py
Bump minimum jaxlib version to 0.4.31. The corresponding xla_extension_version is 279 and mlir_api_version is 57
2024-07-29 18:44:31 -07:00
distributed_test.py
add distributed.is_initialized
2025-02-18 16:47:19 -08:00
dtypes_test.py
Internal: avoid adding _DimExpr to dtypes._weak_types
2025-01-29 09:11:02 -08:00
dynamic_api_test.py
Port tests away from setUpClass and setUpModule to setUp alone.
2025-01-08 08:14:50 -08:00
error_check_test.py
Minor bug fixes in error checking
2025-03-06 06:57:52 -08:00
errors_test.py
Merge pull request from MichaelHudgins:tsan-resultstore
2025-02-06 14:55:57 +00:00
experimental_rnn_test.py
Update RNN kernels to use FFI.
2025-02-06 18:27:58 -08:00
export_back_compat_test.py
#sdy Add JAX backwards compatibility test.
2025-03-03 06:01:34 -08:00
export_harnesses_multi_platform_test.py
[export] Fixes for export_harnesses_multi_platform_test.
2025-01-29 06:12:33 -08:00
export_test.py
#sdy support JAX export tests when Shardy is enabled.
2025-03-03 04:57:06 -08:00
extend_test.py
[better_errors] Make it explicit that debug_info is not None.
2025-02-13 22:07:04 -08:00
ffi_test.py
Remove code present to support jaxlib < 0.5.1.
2025-02-26 07:40:40 -05:00
fft_test.py
Relax some test tolerances in N-D FFT tests.
2024-12-20 03:21:53 -08:00
for_loop_test.py
Make trace dispatch purely a function of context rather than a function of both context and data. This lets us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind and so on.
2024-10-29 11:04:31 -07:00
fused_attention_stablehlo_test.py
Removed unused import
2025-03-05 15:18:43 +00:00
garbage_collection_guard_test.py
Make JAX test suite pass (at least most of the time) with multiple threads enabled.
2025-01-10 06:58:46 -08:00
generated_fun_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
gpu_memory_flags_test.py
Add CUDA custom call example as a JAX test
2024-06-17 15:21:49 -04:00
heap_profiler_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
image_test.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
infeed_test.py
Make JAX test suite pass (at least most of the time) with multiple threads enabled.
2025-01-10 06:58:46 -08:00
jax_jit_test.py
Add test case for jit compile of vmap on gpu.
2025-01-23 17:50:19 -08:00
jax_to_ir_test.py
[jax2tf] Remove non-native serialization test from jax_to_ir_test
2024-10-07 04:21:38 -07:00
jaxpr_effects_test.py
Refactor Jax FFI lowering to prepare for implementing CPU/GPU callbacks using XLA's FFI.
2025-02-21 09:45:59 -08:00
jaxpr_util_test.py
add pjit forwarding rule
2024-05-25 17:46:01 +00:00
jet_test.py
Adding support for copy_p primitive to jet.
2024-11-07 17:06:02 -08:00
key_reuse_test.py
[key reuse] fix signature for device_put
2025-01-17 09:47:50 -08:00
lax_autodiff_test.py
Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0
2024-12-17 10:12:34 -08:00
lax_control_flow_test.py
Update references to JAX's GitHub repo
2025-03-04 18:31:09 -08:00
lax_metal_test.py
adopt jax.extend.backend
2024-10-07 12:27:35 -07:00
lax_numpy_einsum_test.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
lax_numpy_indexing_test.py
Improved errors when indexing with floats
2025-02-28 15:04:07 -08:00
lax_numpy_operators_test.py
Implement jax.numpy.spacing
2024-10-03 10:40:39 -07:00
lax_numpy_reducers_test.py
jax.numpy reductions: avoid upcast of f16 when dtype is specified by user
2025-02-12 11:49:39 -08:00
lax_numpy_test.py
jax.numpy ndim/shape/size: deprecate non-array input
2025-03-04 10:42:32 -08:00
lax_numpy_ufuncs_test.py
Make logaddexp and logaddexp2 into ufuncs
2024-11-18 09:27:36 -08:00
lax_numpy_vectorize_test.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
lax_scipy_sparse_test.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
lax_scipy_special_functions_test.py
Update jax.scipy.special.gamma and gammasgn to return NaN for negative integer inputs.
2024-11-18 20:33:27 -05:00
lax_scipy_spectral_dac_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
lax_scipy_test.py
Merge pull request from mattjj:logsumexp-where-grad-nan
2025-01-10 11:27:35 -08:00
lax_test.py
[sharding_in_types] Make the typing checks and sharding rule checks a little bit less strict when the current or aval mesh is empty/unset. Also some more changes as listed below:
2025-02-12 10:03:01 -08:00
lax_vmap_op_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
lax_vmap_test.py
Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0
2024-12-17 10:12:34 -08:00
layout_test.py
Make sure default layout is None for input and output layout in all codepaths
2025-02-27 14:26:25 -08:00
lazy_loader_test.py
…
linalg_sharding_test.py
Temporarily skip some more linalg sharding checks.
2025-02-28 12:26:18 -08:00
linalg_test.py
Fix linalg.norm to return zero for proper norms of empty matrices.
2025-03-03 15:02:34 -05:00
lobpcg_test.py
Make JAX test suite pass (at least most of the time) with multiple threads enabled.
2025-01-10 06:58:46 -08:00
logging_test.py
logging fixes
2024-11-05 13:28:17 -08:00
lru_cache_test.py
Don't write atime file if JAX_COMPILATIION_CACHE_MAX_SIZE == -1
2025-02-14 12:01:55 -08:00
magma_linalg_test.py
Remove code present to support jaxlib < 0.5.1.
2025-02-26 07:40:40 -05:00
memories_test.py
Remove the skip for test_output_streaming_inside_scan
2025-03-03 14:54:03 -08:00
mesh_utils_test.py
[Take 2] Raise a better error message if anything other than a sequence of ints is passed to make_mesh or create_device_mesh
2024-11-28 09:24:20 -08:00
metadata_test.py
Disable test_source_file_prefix_removal test because there is cross-contamination of metadata information from different call sites because of cached jitted functions
2024-07-12 12:07:24 -07:00
mock_gpu_test.py
Annotate several tests as thread-unsafe.
2025-01-10 11:24:39 -08:00
mock_gpu_topology_test.py
Annotate several tests as thread-unsafe.
2025-01-10 11:24:39 -08:00
monitoring_test.py
Allow duration event listeners to take extra keyword arguments.
2023-10-04 13:15:49 -07:00
mosaic_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
multi_device_test.py
Make JAX test suite pass (at least most of the time) with multiple threads enabled.
2025-01-10 06:58:46 -08:00
multibackend_test.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
multiprocess_gpu_test.py
add distributed.is_initialized
2025-02-18 16:47:19 -08:00
mutable_array_test.py
[better_errors] Cleanup use of DebugInfo.arg_names and result_paths
2025-02-23 08:27:56 +02:00
name_stack_test.py
Just use jit as the string in error messages instead of jit and pjit based on resource_env. This is to start deprecating the need for with mesh and replace it with use_mesh(mesh).
2025-03-05 20:09:30 -08:00
nn_test.py
Improve after review # 2
2025-02-26 04:48:25 +00:00
ode_test.py
[XLA:CPU] Use central difference to calculate numerical gradient
2025-01-22 07:49:43 -08:00
optimizers_test.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
package_structure_test.py
C++ tree with path API
2024-11-30 21:26:48 -08:00
pgle_test.py
Fix test failure in PGLE tests.
2025-02-18 09:44:16 +00:00
pickle_test.py
Canonicalize PartitionSpec so that we can delete ParsedPartitionSpec. We need to do this after sharding-in-types to speed up NamedSharding construction and remove a lot of tech debt and unnecessary complexity.
2025-02-27 08:59:25 -08:00
pjit_test.py
Just use jit as the string in error messages instead of jit and pjit based on resource_env. This is to start deprecating the need for with mesh and replace it with use_mesh(mesh).
2025-03-05 20:09:30 -08:00
pmap_test.py
[better_errors] Refactor debug info tests
2025-01-20 20:21:01 +01:00
polynomial_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
pretty_printer_test.py
Compute source maps when pretty-printing jaxprs.
2024-05-06 15:45:25 -04:00
profiler_test.py
Add a profiler test for gpu run
2025-02-28 13:45:46 -08:00
python_callback_test.py
[xla:cpu] Implement XLA FFI handlers for CPU Jax callbacks.
2025-02-12 13:53:36 -08:00
pytorch_interoperability_test.py
Set PYTHONWARNINGS=error in bazel tests.
2024-09-24 12:30:11 -07:00
qdwh_test.py
Bump up tolerance for tests using _testQdwh.
2024-06-14 23:18:20 -07:00
ragged_collective_test.py
Remove code present to support jaxlib < 0.5.1.
2025-02-26 07:40:40 -05:00
random_lax_test.py
Add use_high_dynamic_range_gumbel flag which allows sampling gumbel such
2025-03-02 19:42:40 -08:00
random_test.py
Don't monkey-patch functions in test_utils to count events for tests.
2024-12-12 09:58:14 -05:00
roofline_test.py
Add unfused_hbm usage to binary ops and dot_general.
2025-02-25 16:10:25 -08:00
scaled_matmul_stablehlo_test.py
Add apache header.
2025-02-26 15:26:56 -06:00
scipy_fft_test.py
Fix dtype bug in jax.scipy.fft.idct
2024-09-25 12:55:43 -04:00
scipy_interpolate_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
scipy_ndimage_test.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
scipy_optimize_test.py
Fix some warnings causing CI failures on ARM.
2024-09-24 17:25:26 -07:00
scipy_signal_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
scipy_spatial_test.py
test: change random matrix generation for Rotation
2025-01-31 09:21:46 -08:00
scipy_stats_test.py
feat(gh-13291): Add exponential distribution functions: cdf, logcdf, sf, logsf, and ppf
2025-02-01 12:51:11 +05:00
shape_poly_test.py
Make api_test.py work when test cases are run using multiple threads.
2025-01-08 14:09:07 -08:00
shard_alike_test.py
Use the mesh of out_aval when converting GSPMDSharding to NamedSharding. This makes sure that the axis types of the corresponding output is correct.
2025-02-20 17:13:24 -08:00
shard_map_test.py
PRNGKeyArray.aval should have the correct logical sharding. This required refactoring code so that we don't hit recursion errors.
2025-03-01 18:18:19 -08:00
source_info_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
source_mapper_test.py
Merge pull request from justinjfu:sourcemap_win
2025-02-06 12:32:10 -08:00
sourcemap_test.py
Merge pull request from gspschmid:gschmid/sourcemaps
2024-05-08 05:44:08 -07:00
sparse_bcoo_bcsr_test.py
[JAX] disable flaky parameter permutations for sparse_bcoo_bcsr test.
2025-02-03 16:02:06 -08:00
sparse_nm_test.py
[XLA:GPU] Skip small tile sizes for sparse gemms on Ampere as well. Enable the JAX test again that has been failing.
2024-11-11 08:58:15 -08:00
sparse_test.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
sparsify_test.py
[better_errors] Improve error message for lax.switch branches output structure mismatch
2025-01-10 08:03:33 +02:00
stack_test.py
Avoid 'from jax import config' imports
2024-04-11 13:23:27 -07:00
state_test.py
[better_errors] Continue adding debug info to Jaxprs (step 7)
2025-02-09 18:14:33 +02:00
stax_test.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
string_array_test.py
Remove code present to support jaxlib < 0.5.1.
2025-02-26 07:40:40 -05:00
svd_test.py
Bump some SVD test tolerances.
2024-07-15 12:54:00 -07:00
transfer_guard_test.py
fix some additional warnings related to
2024-06-13 16:06:14 -07:00
tree_util_test.py
Add an option to simplify keystr output and use a custom separator.
2025-01-21 10:18:42 -08:00
typing_test.py
Upgrade remaining sources to Python 3.9
2023-12-13 10:29:45 +00:00
util_test.py
[better_errors] Continue adding debug info to Jaxprs (step 7)
2025-02-09 18:14:33 +02:00
version_test.py
Refactor JAX wheel build rules to control the wheel filename and maintain reproducible wheel content and filename results.
2025-02-05 10:01:23 -08:00
warnings_util_test.py
Switch to a new thread-safe utility for catching warnings.
2025-01-09 11:58:34 -05:00
x64_context_test.py
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
2024-09-20 07:52:33 -07:00
xla_bridge_test.py
Use backend._get_all_devices() to validate devices.
2025-01-24 11:09:16 -08:00
xla_interpreter_test.py
…
xla_metadata_test.py
Make trace dispatch purely a function of context rather than a function of both context and data. This lets us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind and so on.
2024-10-29 11:04:31 -07:00