mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix sparse dot metadata loader

Browse Source 
Metadata loader was using incorrect warp assignment, which resulted in incorrect addresses with num_warps>4. This was previously missed, as the autotuner rarely selected such configs.

PiperOrigin-RevId: 633513110
This commit is contained in:
Sergey Kozub 2024-05-14 03:07:12 -07:00 committed by jax authors
parent e735a00cdc
commit 84774b39e6
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tests/BUILD
View File

@ -999,8 +999,7 @@ jax_test(
],
enable_configs = [
"gpu_a100",
# TODO(b/337303303): re-enable the test
# "gpu_h100",
"gpu_h100",
],
deps = [
"//jax:experimental_sparse",