Remove unnecessary head_dim restriction in TPU Flash Attention.

GitOrigin-RevId: 1f1b365

Author: ds-hwang

axlearn/common/flash_attention/tpu_attention.py

@@ -899,17 +899,10 @@ def is_supported(
             return False
         bias: BaseAttentionBias = input_batch["bias"]
         _, _, explicit_bias = split(bias, MaskFnAttentionBias, SegmentIdAttentionBias)
-        query: Tensor = input_batch["query"]
-        head_dim = query.shape[-1]
 
         if explicit_bias.has_value():
             return self._log_unsupported("explicit bias is not supported.")
 
-        if head_dim % splash_attention_kernel.NUM_LANES != 0:
-            return self._log_unsupported(
-                f"{head_dim=} is not divisible by {splash_attention_kernel.NUM_LANES=}"
-            )
-
         if (
             not self.get_backend_overrides("splash_use_fused_bwd_kernel", True)
             and self.cfg.dropout_rate > 0.0

axlearn/common/flash_attention/tpu_attention_test.py

@@ -126,7 +126,7 @@ def inside_tracing(mask):
         kv_len=[128],
         num_heads=[4],
         mask=[None, causal_mask],
-        per_head_dim=[128],
+        per_head_dim=[64, 128],
         q_dtype=[jnp.float32, jnp.bfloat16],
         kv_dtype=[jnp.float32, jnp.bfloat16],
     )
@@ -209,15 +209,15 @@ def test_gradient(
             self.skipTest("Backward path is broken on CPU")
         if mask not in (None, causal_mask) and query_length_multiplier > 1:
             self.skipTest("Sliding window attention does not make sense when q_len != kv_len.")
-        if dropout_rate > 0.0 and (attention_bias_type is not None or per_head_dim % 128 != 0):
+        if dropout_rate > 0.0 and attention_bias_type is not None:
             self.skipTest(
-                "Dropout is only supported with SplashAttention (which requires \
-                            no bias, and per_head_dim being a multiple of 128.)"
+                "Dropout is only supported with SplashAttention (which requires no bias.)"
             )
         if q_dtype == jnp.bfloat16 and kv_dtype == jnp.float32:
             self.skipTest("Q must have higher precision than KV.")
+
         # pylint: disable=protected-access
-        fallback_to_legacy = per_head_dim % 128 != 0 or (attention_bias_type is not None)
+        fallback_to_legacy = attention_bias_type is not None
         num_kv_heads = num_heads // head_group_size
         q, k, v, bias = generate_attention_data(
             batch_size,

axlearn/common/flash_attention/tpu_splash_attention.py

@@ -138,14 +138,11 @@ def flash_attention_kernel(
     mask_function: MaskFunctionType | None,
     dropout_rate: float,
 ):
+    del head_dim
     float32 = jnp.float32
     # pylint: disable=invalid-name
     HEAD_DIM_MINOR = QKVLayout.HEAD_DIM_MINOR
 
-    head_dim_repeats, rem = divmod(head_dim, NUM_LANES)
-    if rem != 0:
-        raise NotImplementedError(f"{head_dim=} should be a multiple of {NUM_LANES}")
-
     h, i, j = pl.program_id(0), pl.program_id(1), pl.program_id(2)
 
     @pl.when(j == 0)
@@ -171,8 +168,8 @@ def init():
     def body(kv_compute_index, _):
         slice_k = pl.ds(kv_compute_index * bkv_compute, bkv_compute)
         m_prev, l_prev = m_scratch_ref[...], l_scratch_ref[...]
-        assert m_prev.shape == (bq, NUM_LANES)
-        assert l_prev.shape == (bq, NUM_LANES)
+        assert m_prev.shape == (bq, 1)
+        assert l_prev.shape == (bq, 1)
 
         q = q_ref[...] if q_layout == HEAD_DIM_MINOR else q_ref[...].T
         qk_dims = NT_DIM_NUMBERS if k_layout == HEAD_DIM_MINOR else NN_DIM_NUMBERS
@@ -217,17 +214,13 @@ def body(kv_compute_index, _):
         m_curr = qk.max(axis=-1)[:, None]
         assert m_curr.shape == (bq, 1)
         m_next = jnp.maximum(m_prev, m_curr)
-        assert m_next.shape == (bq, NUM_LANES)
-
-        bkv_repeats, rem = divmod(bkv_compute, NUM_LANES)
-        if rem != 0:
-            raise NotImplementedError(f"{bkv_compute=} should be a multiple of {NUM_LANES}")
+        assert m_next.shape == (bq, 1)
 
-        s_curr = jnp.exp(qk - pltpu.repeat(m_next, bkv_repeats, axis=1))
+        s_curr = jnp.exp(qk - m_next)
         assert s_curr.shape == (bq, bkv_compute)
 
         l_curr = jax.lax.broadcast_in_dim(s_curr.sum(axis=-1), l_prev.shape, (0,))
-        assert l_curr.shape == (bq, NUM_LANES)
+        assert l_curr.shape == (bq, 1)
 
         alpha = jnp.exp(m_prev - m_next)
         l_next = l_curr + alpha * l_prev
@@ -254,7 +247,7 @@ def body(kv_compute_index, _):
             s_curr = jnp.where(dropout_mask, 0.0, s_curr) / (1.0 - dropout_rate)
         o_curr = lax.dot_general(s_curr, v, sv_dims)
 
-        alpha_o = pltpu.repeat(alpha, head_dim_repeats, axis=1)
+        alpha_o = alpha
         o_scratch_ref[:] = alpha_o * o_scratch_ref[:] + o_curr
 
     @pl.when(should_run)
@@ -269,10 +262,10 @@ def end():
         if logit_sink_ref is not None:
             sink_value = logit_sink_ref[h].astype(jnp.float32)
             l = l + jnp.exp(sink_value - m_scratch_ref[...])
-        l_inv = pltpu.repeat(1.0 / l, head_dim_repeats, axis=1)
+        l_inv = 1.0 / l  # TODO(dhwang2): dividing directly by `l` is a more stable operation.
         o_ref[...] = (o_scratch_ref[...] * l_inv).astype(o_ref.dtype)
         if logsumexp_ref is not None:
-            assert logsumexp_ref.shape == (bq, NUM_LANES)
+            assert logsumexp_ref.shape == (bq, 1)
             logsumexp_ref[...] = (jnp.log(l) + m_scratch_ref[...]).astype(logsumexp_ref.dtype)
 
         m_scratch_ref[...] = jnp.zeros_like(m_scratch_ref)
@@ -546,28 +539,28 @@ def kv_segment_ids_index_map(
     num_scalar_prefetch = 4
 
     out_shapes = [
-        jax.ShapeDtypeStruct((bq, NUM_LANES), jnp.float32),  # m_scratch
-        jax.ShapeDtypeStruct((bq, NUM_LANES), jnp.float32),  # l_scratch
+        jax.ShapeDtypeStruct((bq, 1), jnp.float32),  # m_scratch
+        jax.ShapeDtypeStruct((bq, 1), jnp.float32),  # l_scratch
         jax.ShapeDtypeStruct((bq, head_dim), jnp.float32),  # o_scratch
         jax.ShapeDtypeStruct((num_q_heads, q_seq_len, head_dim), q.dtype),
     ]
     out_specs = [
         # TODO(sharadmv): convert m/l to be scratch
-        pl.BlockSpec((bq, NUM_LANES), lambda h, i, j, *_: (0, 0)),
-        pl.BlockSpec((bq, NUM_LANES), lambda h, i, j, *_: (0, 0)),
+        pl.BlockSpec((bq, 1), lambda h, i, j, *_: (0, 0)),
+        pl.BlockSpec((bq, 1), lambda h, i, j, *_: (0, 0)),
         pl.BlockSpec((bq, head_dim), lambda h, i, j, *_: (0, 0)),
         pl.BlockSpec((None, bq, head_dim), out_index_map),
     ]
     if save_residuals:
         out_shapes += [
-            jax.ShapeDtypeStruct((num_q_heads, q_seq_len, NUM_LANES), jnp.float32),  # logsumexp
+            jax.ShapeDtypeStruct((num_q_heads, q_seq_len, 1), jnp.float32),  # logsumexp
         ]
 
         def logsumexp_index_map(h, i, *_):
             return h, i, 0
 
         out_specs += [
-            pl.BlockSpec((None, bq, NUM_LANES), logsumexp_index_map),
+            pl.BlockSpec((None, bq, 1), logsumexp_index_map),
         ]
     else:
         out_shapes += [None]