Skip to content

Commit 0f84433

Browse files
apaszkeGoogle-ML-Automation
authored andcommitted
Disable tests for the paged attention kernel
I ran the TPU race checker on it and it did report a number of races that were uncovered by recent Mosaic compiler changes. PiperOrigin-RevId: 761447182
1 parent 4d6e39e commit 0f84433

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

tests/pallas/tpu_paged_attention_kernel_test.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -265,6 +265,8 @@ def test_paged_attention(
265265
attn_logits_soft_cap,
266266
are_kv_quantized,
267267
):
268+
# TODO(mvoz, skyewm): Re-enable this test once the data race is fixed.
269+
self.skipTest("This kernel has data races that need to be fixed.")
268270
if not jtu.is_device_tpu_at_least(4):
269271
self.skipTest("Only supports TPU generation 4 or above")
270272
if jtu.is_device_tpu(version=4) and are_kv_quantized:

0 commit comments

Comments
 (0)