Fix GPU Attention Tests (#1314)

* Fix if statement and ensure config is_supported is run

* Add to other test

Author: andersensam

axlearn/common/flash_attention/gpu_attention.py

@@ -811,7 +811,7 @@ def is_supported(
             # key/value to be even.
             if not self._check_block_size(input_batch, block_size=2):
                 return False
-        if kv_cache_type == KVCache:
+        elif kv_cache_type == KVCache:
             if query.shape[1] > 1:
                 return self._log_unsupported("multi-step decoding is not supported.")
             if not key.shape[1] % 2 == 0:

axlearn/common/flash_attention/gpu_attention_test.py

@@ -398,6 +398,13 @@ def test_cudnn_dropout_against_xla_dropout(
     test_fn = CuDNNGPUFlashAttention.default_config().set(**cfg).instantiate()
     ref_fn = ReferenceMHA.default_config().set(**cfg).instantiate()
 
+    k1, k2, k3 = jax.random.split(jax.random.PRNGKey(0), 3)
+    q = jax.random.normal(k1, qkv_shape, dtype=dtype)
+    k = jax.random.normal(k2, qkv_shape, dtype=dtype)
+    v = jax.random.normal(k3, qkv_shape, dtype=dtype)
+    input_batch = dict(query=q, key=k, value=v, bias=bias, logit_sink=None)
+    chex.assert_equal(test_fn.is_supported(input_batch, kv_cache_type=None), True)
+
     dropout_mask = (
         test_fn(
             dict(
@@ -416,13 +423,6 @@ def test_cudnn_dropout_against_xla_dropout(
     # the same mask.
     jax.clear_caches()
 
-    k1, k2, k3 = jax.random.split(jax.random.PRNGKey(0), 3)
-    q = jax.random.normal(k1, qkv_shape, dtype=dtype)
-    k = jax.random.normal(k2, qkv_shape, dtype=dtype)
-    v = jax.random.normal(k3, qkv_shape, dtype=dtype)
-    input_batch = dict(query=q, key=k, value=v, bias=bias, logit_sink=None)
-    chex.assert_equal(test_fn.is_supported(input_batch, kv_cache_type=None), True)
-
     ref_fn = functools.partial(
         ref_fn,
         dropout_mask=dropout_mask,
@@ -492,6 +492,7 @@ def test_cudnn_dropout_determinism():
         logit_sink=None,
     )
     fn = CuDNNGPUFlashAttention.default_config().set(dropout_rate=0.1).instantiate()
+    chex.assert_equal(fn.is_supported(input_batch, kv_cache_type=None), True)
 
     outputs = []
     grads = []