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

Temporarily disable GQA attention tests on GPU, which were broken by a Triton integrate.

Browse Source 
PiperOrigin-RevId: 715516188
This commit is contained in:
Peter Hawkins 2025-01-14 13:47:56 -08:00 committed by jax authors
parent c78487d23d
commit d1810b42cb
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/pallas/gpu_attention_test.py
View File

@ -14,6 +14,7 @@
import os
import sys
import unittest
from absl.testing import absltest
from absl.testing import parameterized
@ -168,6 +169,7 @@ class DecodeAttentionTest(PallasBaseTest):
for return_residuals in [False, True]
])
@jax.numpy_dtype_promotion("standard")
@unittest.skip("TODO(b/389925439): gqa tests started failing after triton integrate")
def test_gqa(
self,
batch_size,