Add support for non-power-of-2 head size in flash attention

Introduce checks on sequences being divisible by block sizes to address #27224

PiperOrigin-RevId: 762051831

Author: rdyro

jax/experimental/pallas/ops/gpu/attention.py

@@ -86,28 +86,29 @@ def mha_forward_kernel(
     segment_ids_ref: jax.Array | None,  # segment_id arrays
     o_ref: Any,  # Output
     *residual_refs: Any,  # Residual outputs
-    num_heads: int,
     sm_scale: float,
     causal: bool,
     block_q: int,
-    block_d: int,
     block_k: int,
+    head_dim: int,
 ):
   seq_len = k_ref.shape[0]
   start_q = pl.program_id(0)
+  head_dim_padded = q_ref.shape[-1]
 
   # o is the buffer where we accumulate the output on sram.
   # m_i and l_i (see FlashAttention paper) are updated during the k,v loop.
   m_i = jnp.zeros(block_q, dtype=jnp.float32) - float('inf')
   l_i = jnp.zeros(block_q, dtype=jnp.float32)
   # acc is the buffer where we accumulate the output on sram.
-  o = jnp.zeros((block_q, block_d), dtype=jnp.float32)
+  o = jnp.zeros((block_q, head_dim_padded), dtype=jnp.float32)
 
   # Load q: it will stay in L1 throughout. Indices form a matrix because we
   # read, compute, and write all in 2d chunks. 1 element ~= 1 CUDA thread index.
-  # q tile has shape [block_q, block_d], block_d == head_dim.
+  # q tile has shape [block_q, head_dim_padded], head_dim_padded >= head_dim.
   curr_q_slice = pl.dslice(start_q * block_q, block_q)
-  q = q_ref[...]
+  head_mask = (jnp.arange(head_dim_padded) < head_dim)[None, :]
+  q = pl.load(q_ref, (slice(None), slice(None)), mask=head_mask, other=0.0)
   q_segment_ids = (
       None
       if segment_ids_ref is None
@@ -121,7 +122,7 @@ def body(start_k, carry):
     o_prev, m_prev, l_prev = carry
     curr_k_slice = pl.dslice(start_k * block_k, block_k)
 
-    k = pl.load(k_ref, (curr_k_slice, slice(None)))
+    k = pl.load(k_ref, (curr_k_slice, slice(None)), mask=head_mask, other=0.0)
     qk = pl.dot(q, k.T)   # [block_q, block_k]
 
     # Scale logits to convert from base-2 to the natural log domain.
@@ -161,7 +162,7 @@ def body(start_k, carry):
     l_curr = s_curr.sum(axis=-1)
     l_next = l_prev_corr + l_curr
     o_prev_corr = correction[:, None] * o_prev
-    v = pl.load(v_ref, (curr_k_slice, pl.dslice(block_d)))
+    v = pl.load(v_ref, (curr_k_slice, slice(None)), mask=head_mask)
     o_curr = pl.dot(s_curr.astype(v.dtype), v)
 
     o_next = o_prev_corr + o_curr
@@ -182,7 +183,8 @@ def body(start_k, carry):
     lse_ref = residual_refs[0]
     lse_ref[...] = m_i + jnp.log2(l_i)
   # Write output to dram.
-  o_ref[...] = o.astype(o_ref.dtype)
+  pl.store(o_ref, (slice(None), slice(o.shape[-1])), o.astype(o_ref.dtype),
+           mask=head_mask)
 
 def segment_mask(
     q_segment_ids: jax.Array,
@@ -235,6 +237,17 @@ def mha(
   kv_seq_len = k.shape[1]
   block_q = min(block_sizes.block_q, q_seq_len)
   block_k = min(block_sizes.block_k, kv_seq_len)
+  head_dim_padded = pl.next_power_of_2(head_dim)
+  if (q.shape[-1] != k.shape[-1]) or (q.shape[-1] != v.shape[-1]):
+    raise ValueError(
+        f"This kernel expects q, k, and v to have the same head dimension, but"
+        f" found {q.shape=}, {k.shape=}, {v.shape=}."
+    )
+  if q_seq_len % block_q != 0:
+    raise ValueError(f"{q_seq_len=} must be a multiple of {block_q=}")
+  if kv_seq_len % block_k != 0:
+    raise ValueError(f"{kv_seq_len=} must be a multiple of {block_k=}")
+
   # Heuristics.
   grid_ = grid
   if grid_ is None:
@@ -243,21 +256,17 @@ def mha(
   num_warps_ = num_warps
   if num_warps_ is None:
     num_warps_ = 4 if head_dim <= 64 else 8
-  kernel = functools.partial(mha_forward_kernel, num_heads=num_heads,
-                             sm_scale=sm_scale, block_q=block_q,
-                             block_k=block_k, block_d=head_dim,
-                             causal=causal)
+  kernel = functools.partial(mha_forward_kernel, sm_scale=sm_scale,
+                             block_q=block_q, block_k=block_k,
+                             head_dim=head_dim, causal=causal)
 
   in_specs = [
-      pl.BlockSpec(
-          (None, block_q, None, head_dim), lambda i, j, k: (j, i, k, 0)
-      ),
-      pl.BlockSpec(
-          (None, kv_seq_len, None, head_dim), lambda _, j, k: (j, 0, k, 0)
-      ),
-      pl.BlockSpec(
-          (None, kv_seq_len, None, head_dim), lambda _, j, k: (j, 0, k, 0)
-      ),
+      pl.BlockSpec((None, block_q, None, head_dim_padded),
+                   lambda i, j, k: (j, i, k, 0)),
+      pl.BlockSpec((None, kv_seq_len, None, head_dim_padded),
+                   lambda _, j, k: (j, 0, k, 0)),
+      pl.BlockSpec((None, kv_seq_len, None, head_dim_padded),
+                   lambda _, j, k: (j, 0, k, 0)),
   ]
   in_specs.append(
       None  # type: ignore[arg-type]
@@ -270,7 +279,7 @@ def mha(
       grid=grid_,
       in_specs=in_specs,
       out_specs=pl.BlockSpec(
-          (None, block_q, None, head_dim), lambda i, j, k: (j, i, k, 0)
+          (None, block_q, None, head_dim_padded), lambda i, j, k: (j, i, k, 0)
       ),
       compiler_params=plgpu.TritonCompilerParams(
           num_warps=num_warps_, num_stages=num_stages),
@@ -301,6 +310,17 @@ def _mha_forward(
   kv_seq_len = k.shape[1]
   block_q = min(block_sizes.block_q, q_seq_len)
   block_k = min(block_sizes.block_k, kv_seq_len)
+  if (q.shape[-1] != k.shape[-1]) or (q.shape[-1] != v.shape[-1]):
+    raise ValueError(
+        f"This kernel expects q, k, and v to have the same head dimension, but"
+        f" found {q.shape=}, {k.shape=}, {v.shape=}."
+    )
+  if q_seq_len % block_q != 0:
+    raise ValueError(f"{q_seq_len=} must be a multiple of {block_q=}")
+  if kv_seq_len % block_k != 0:
+    raise ValueError(f"{kv_seq_len=} must be a multiple of {block_k=}")
+  head_dim_padded = pl.next_power_of_2(head_dim)
+
   # Heuristics.
   grid_ = grid
   if grid_ is None:
@@ -309,25 +329,22 @@ def _mha_forward(
   num_warps_ = num_warps
   if num_warps_ is None:
     num_warps_ = 4 if head_dim <= 64 else 8
-  kernel = functools.partial(mha_forward_kernel, num_heads=num_heads,
-                             sm_scale=sm_scale, causal=causal, block_q=block_q,
-                             block_k=block_k, block_d=head_dim)
+  kernel = functools.partial(mha_forward_kernel, sm_scale=sm_scale,
+                             causal=causal, block_q=block_q, block_k=block_k,
+                             head_dim=head_dim)
   out_shape = [
       jax.ShapeDtypeStruct(shape=q.shape, dtype=q.dtype),  # out
       jax.ShapeDtypeStruct(
           shape=(batch_size, num_heads, q_seq_len), dtype=jnp.float32  # lse
       ),
   ]
   in_specs = [
-      pl.BlockSpec(
-          (None, block_q, None, head_dim), lambda i, j, k: (j, i, k, 0)
-      ),
-      pl.BlockSpec(
-          (None, kv_seq_len, None, head_dim), lambda _, j, k: (j, 0, k, 0)
-      ),
-      pl.BlockSpec(
-          (None, kv_seq_len, None, head_dim), lambda _, j, k: (j, 0, k, 0)
-      ),
+      pl.BlockSpec((None, block_q, None, head_dim_padded),
+                   lambda i, j, k: (j, i, k, 0)),
+      pl.BlockSpec((None, kv_seq_len, None, head_dim_padded),
+                   lambda _, j, k: (j, 0, k, 0)),
+      pl.BlockSpec((None, kv_seq_len, None, head_dim_padded),
+                   lambda _, j, k: (j, 0, k, 0)),
   ]
   in_specs.append(
       None  # type: ignore[arg-type]
@@ -339,9 +356,8 @@ def _mha_forward(
       grid=grid_,
       in_specs=in_specs,
       out_specs=[
-          pl.BlockSpec(
-              (None, block_q, None, head_dim), lambda i, j, k: (j, i, k, 0)
-          ),
+          pl.BlockSpec((None, block_q, None, head_dim_padded),
+                       lambda i, j, k: (j, i, k, 0)),
           pl.BlockSpec((None, None, block_q), lambda i, j, k: (j, k, i)),
       ],
       compiler_params=plgpu.TritonCompilerParams(
@@ -355,10 +371,11 @@ def _mha_forward(
   return out, (q, k, v, segment_ids, out, lse)
 
 
-def _preprocess_backward_kernel(out_ref, dout_ref, delta_ref):
+def _preprocess_backward_kernel(out_ref, dout_ref, delta_ref, head_dim: int):
   # load
-  o = out_ref[...].astype(jnp.float32)
-  do = dout_ref[...].astype(jnp.float32)
+  head_mask = (jnp.arange(out_ref.shape[-1]) < head_dim)[None, :]
+  o = pl.load(out_ref, (slice(None), slice(None)), mask=head_mask, other=0.0)
+  do = pl.load(dout_ref, (slice(None), slice(None)), mask=head_mask, other=0.0)
   # compute
   delta = jnp.sum(o * do, axis=1)
   # write-back
@@ -368,17 +385,16 @@ def _preprocess_backward_kernel(out_ref, dout_ref, delta_ref):
 def _preprocess_backward(out, do, lse, block_q: int,
                          debug: bool, interpret: bool):
   batch_size, seq_len, num_heads, head_dim = out.shape
+  head_dim_padded = pl.next_power_of_2(head_dim)
   out_shape = jax.ShapeDtypeStruct(lse.shape, lse.dtype)
   delta = pl.pallas_call(
-      _preprocess_backward_kernel,
+      functools.partial(_preprocess_backward_kernel, head_dim=head_dim),
       grid=(pl.cdiv(seq_len, block_q), batch_size, num_heads),
       in_specs=[
-          pl.BlockSpec(
-              (None, block_q, None, head_dim), lambda i, j, k: (j, i, k, 0)
-          ),
-          pl.BlockSpec(
-              (None, block_q, None, head_dim), lambda i, j, k: (j, i, k, 0)
-          ),
+          pl.BlockSpec((None, block_q, None, head_dim_padded),
+                       lambda i, j, k: (j, i, k, 0)),
+          pl.BlockSpec((None, block_q, None, head_dim_padded),
+                       lambda i, j, k: (j, i, k, 0)),
       ],
       out_specs=pl.BlockSpec((None, None, block_q), lambda i, j, k: (j, k, i)),
       compiler_params=plgpu.TritonCompilerParams(num_warps=4, num_stages=3),
@@ -414,7 +430,7 @@ def mha_backward_kernel(
     block_kv_dkv: int,
     block_q_dq: int,
     block_kv_dq: int,
-    block_d: int,
+    head_dim: int,
 ):
   del out_ref  # Not needed
   q_seq_len = q_ref.shape[0]
@@ -427,11 +443,13 @@ def mha_backward_kernel(
   start_k = pl.program_id(2)
   curr_k_slice = pl.dslice(start_k * block_kv_dkv, block_kv_dkv)
 
-  dv = jnp.zeros([block_kv_dkv, block_d], dtype=jnp.float32)
-  dk = jnp.zeros([block_kv_dkv, block_d], dtype=jnp.float32)
+  head_dim_padded = q_ref.shape[-1]
+  dv = jnp.zeros([block_kv_dkv, head_dim_padded], dtype=jnp.float32)
+  dk = jnp.zeros([block_kv_dkv, head_dim_padded], dtype=jnp.float32)
 
-  v = pl.load(v_ref, (curr_k_slice, slice(None)))
-  k = pl.load(k_ref, (curr_k_slice, slice(None)))
+  head_mask = (jnp.arange(head_dim_padded) < head_dim)[None, :]
+  v = pl.load(v_ref, (curr_k_slice, slice(None)), mask=head_mask, other=0.0)
+  k = pl.load(k_ref, (curr_k_slice, slice(None)), mask=head_mask, other=0.0)
   span_k = start_k * block_kv_dkv + jnp.arange(block_kv_dkv)
   kv_segment_ids = (
       None
@@ -443,7 +461,7 @@ def inner_loop_dkdv(start_q, carry):
     dv, dk = carry
     curr_q_slice = pl.dslice(start_q * block_q_dkv, block_q_dkv)
 
-    q = pl.load(q_ref, (curr_q_slice, slice(None)))
+    q = pl.load(q_ref, (curr_q_slice, slice(None)), mask=head_mask, other=0.0)
     qk = pl.dot(q, k.T)
     qk_scale = math.log2(math.e)
     if sm_scale != 1.:
@@ -466,7 +484,8 @@ def inner_loop_dkdv(start_q, carry):
 
     lse = pl.load(lse_ref, (curr_q_slice,))
     di = pl.load(delta_ref, (curr_q_slice,))
-    do = pl.load(do_scaled_ref, (curr_q_slice, slice(None)))
+    do = pl.load(do_scaled_ref, (curr_q_slice, slice(None)), mask=head_mask,
+                 other=0.0)
 
     p = jnp.exp2(qk - lse[:, None])
     dv = dv + pl.dot(p.astype(do.dtype).T, do)
@@ -483,8 +502,10 @@ def inner_loop_dkdv(start_q, carry):
   dv, dk = lax.fori_loop(
       lower_bound, pl.cdiv(q_seq_len, block_q_dkv), inner_loop_dkdv, (dv, dk)
   )
-  dv_ref[...] = dv.astype(dv_ref.dtype)
-  dk_ref[...] = dk.astype(dk_ref.dtype)
+  pl.store(dv_ref, (slice(None), slice(dv.shape[-1])), dv.astype(dv_ref.dtype),
+           mask=head_mask)
+  pl.store(dk_ref, (slice(None), slice(dk.shape[-1])), dk.astype(dk_ref.dtype),
+           mask=head_mask)
 
   # Scan #2: dQ
   #   1. Load a block of Q of size (block_q_dq, head_dim) in SMEM.
@@ -493,22 +514,23 @@ def inner_loop_dkdv(start_q, carry):
   start_q = pl.program_id(2)
   curr_q_slice = pl.ds(start_q * block_q_dq, block_q_dq)
   span_q = start_q * block_q_dq + jnp.arange(block_q_dq)
-  dq = jnp.zeros([block_q_dq, block_d], dtype=jnp.float32)
+  dq = jnp.zeros([block_q_dq, head_dim_padded], dtype=jnp.float32)
 
-  q = pl.load(q_ref, (curr_q_slice, slice(None)))
+  q = pl.load(q_ref, (curr_q_slice, slice(None)), mask=head_mask, other=0.0)
   q_segment_ids = (
       None
       if segment_ids_ref is None
       else pl.load(segment_ids_ref, (curr_q_slice,))
   )
   lse = pl.load(lse_ref, (curr_q_slice,))
-  do = pl.load(do_scaled_ref, (curr_q_slice, slice(None)))
+  do = pl.load(do_scaled_ref, (curr_q_slice, slice(None)), mask=head_mask,
+               other=0.0)
   di = pl.load(delta_ref, (curr_q_slice,))
 
   def inner_loop_dq(start_k, dq):
     curr_k_slice = pl.dslice(start_k * block_kv_dq, block_kv_dq)
-    k = pl.load(k_ref, (curr_k_slice, slice(None)))
-    v = pl.load(v_ref, (curr_k_slice, slice(None)))
+    k = pl.load(k_ref, (curr_k_slice, slice(None)), mask=head_mask, other=0.0)
+    v = pl.load(v_ref, (curr_k_slice, slice(None)), mask=head_mask, other=0.0)
 
     qk = pl.dot(q, k.T)
     qk_scale = math.log2(math.e)
@@ -547,7 +569,8 @@ def inner_loop_dq(start_k, dq):
     upper_bound = pl.cdiv(kv_seq_len, block_kv_dq)
 
   dq = lax.fori_loop(0, upper_bound, inner_loop_dq, (dq))
-  dq_ref[...] = dq.astype(dq_ref.dtype)
+  pl.store(dq_ref, (slice(None), slice(dq.shape[-1])), dq.astype(dq_ref.dtype),
+           mask=head_mask)
 
 
 def _mha_backward(sm_scale: float, causal: bool, block_sizes: BlockSizes,
@@ -576,6 +599,7 @@ def _mha_backward(sm_scale: float, causal: bool, block_sizes: BlockSizes,
     block_kv_dkv = min(block_sizes.block_kv_dkv, kv_seq_len)
     block_q_dq = min(block_sizes.block_q_dq, q_seq_len)
     block_kv_dq = min(block_sizes.block_kv_dq, kv_seq_len)
+    head_dim_padded = pl.next_power_of_2(head_dim)
 
     if q_seq_len // block_q_dq != kv_seq_len // block_kv_dkv:
       raise ValueError(
@@ -591,28 +615,24 @@ def _mha_backward(sm_scale: float, causal: bool, block_sizes: BlockSizes,
     ]
 
     in_specs = [
-        pl.BlockSpec(
-            (None, q_seq_len, None, head_dim), lambda i, j, _: (i, 0, j, 0)
-        ),
-        pl.BlockSpec(
-            (None, kv_seq_len, None, head_dim), lambda i, j, _: (i, 0, j, 0)
-        ),
-        pl.BlockSpec(
-            (None, kv_seq_len, None, head_dim), lambda i, j, _: (i, 0, j, 0)
-        ),
-        pl.BlockSpec(
-            (None, q_seq_len, None, head_dim), lambda i, j, _: (i, 0, j, 0)
-        ),
-        pl.BlockSpec(
-            (None, q_seq_len, None, head_dim), lambda i, j, _: (i, 0, j, 0)
-        ),
+        pl.BlockSpec((None, q_seq_len, None, head_dim_padded),
+                     lambda i, j, _: (i, 0, j, 0)),
+        pl.BlockSpec((None, kv_seq_len, None, head_dim_padded),
+                     lambda i, j, _: (i, 0, j, 0)),
+        pl.BlockSpec((None, kv_seq_len, None, head_dim_padded),
+                     lambda i, j, _: (i, 0, j, 0)),
+        pl.BlockSpec((None, q_seq_len, None, head_dim_padded),
+                     lambda i, j, _: (i, 0, j, 0)),
+        pl.BlockSpec((None, q_seq_len, None, head_dim_padded),
+                     lambda i, j, _: (i, 0, j, 0)),
         pl.BlockSpec((None, None, q_seq_len), lambda i, j, _: (i, j, 0)),
         pl.BlockSpec((None, None, q_seq_len), lambda i, j, _: (i, j, 0)),
     ]
     if segment_ids is None:
       in_specs.insert(3, None)  # type: ignore[arg-type]
     else:
-      in_specs.insert(3, pl.BlockSpec((None, kv_seq_len), lambda i, j, _: (i, 0)))
+      in_specs.insert(3, pl.BlockSpec((None, kv_seq_len),
+                                      lambda i, j, _: (i, 0)))
 
     grid = (batch_size, num_heads, pl.cdiv(kv_seq_len, block_kv_dkv))
     num_warps_ = num_warps
@@ -635,22 +655,22 @@ def _mha_backward(sm_scale: float, causal: bool, block_sizes: BlockSizes,
             block_kv_dkv=block_kv_dkv,
             block_q_dq=block_q_dq,
             block_kv_dq=block_kv_dq,
-            block_d=head_dim,
+            head_dim=head_dim,
         ),
         out_shape=out_shapes,
         in_specs=in_specs,
         grid=grid,
         out_specs=[
             pl.BlockSpec(
-                (None, block_q_dq, None, head_dim),
+                (None, block_q_dq, None, head_dim_padded),
                 lambda i, j, k: (i, k, j, 0),  # dq
             ),
             pl.BlockSpec(
-                (None, block_kv_dkv, None, head_dim),
+                (None, block_kv_dkv, None, head_dim_padded),
                 lambda i, j, k: (i, k, j, 0),  # dk
             ),
             pl.BlockSpec(
-                (None, block_kv_dkv, None, head_dim),
+                (None, block_kv_dkv, None, head_dim_padded),
                 lambda i, j, k: (i, k, j, 0),  # dv
             ),
         ],