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

Re-enable oss paged attn kernel

Browse Source 
PiperOrigin-RevId: 725411244
This commit is contained in:
jax authors 2025-02-10 17:46:40 -08:00
parent ffd3faad72
commit d3ed6ca0cc
2 changed files with 0 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tests/pallas/BUILD
View File

@ -463,7 +463,6 @@ jax_multiplatform_test(
tags = [
"noasan", # Times out.
"nomsan", # Times out.
"notap", # this code has data race issues that XLA improvements unhide. b/392946030
"notsan", # Times out.
],
deps = [

1
tests/pallas/tpu_paged_attention_kernel_test.py
View File

@ -133,7 +133,6 @@ class PagedAttentionKernelTest(jtu.JaxTestCase):
attn_logits_soft_cap,
are_kv_quantized,
):
self.skipTest("This kernel has data races that need to be fixed.")
if not jtu.is_device_tpu_at_least(4):
self.skipTest("Only supports TPU generation 4 or above")
if jtu.is_device_tpu(version=4) and are_kv_quantized: