Fix varlen seqused_k behaviour with 0-length sequences.

Author: nshepperd

src/flash_attn_jax/varlen.py

@@ -34,8 +34,8 @@ def _flash_mha_varlen_vjp_bwd(config, pack, dout):
         ends = seqlens_k[1:]
         lens = ends - starts
         zero = jnp.zeros(q.shape[0], dtype=jnp.int32)
-        ixl = jnp.arange(q.shape[0]) - jnp.cumsum(zero.at[ends].set(lens))
-        limits = jnp.cumsum(zero.at[starts].set(seqused_k-jnp.concatenate([jnp.array([0]), seqused_k[:-1]])))
+        ixl = jnp.arange(q.shape[0]) - jnp.cumsum(zero.at[ends].add(lens))
+        limits = jnp.cumsum(zero.at[starts].add(seqused_k-jnp.concatenate([jnp.array([0]), seqused_k[:-1]])))
         mask = (ixl < limits)
         v = v * mask[:, None, None]
         k = k * mask[:, None, None]

tests/test_varlen.py

@@ -42,7 +42,7 @@ def check1(ref_out, jax_out, out):
 @pytest.mark.parametrize("m", [1, 2]) # for MQA/GQA
 def test_varlen_flash_fwd(m, h, d, causal, local, dtype, seqused_k_limit):
     window_size = (3,3) if local else (-1,-1)
-    lens = [1, 2, 6, 10]
+    lens = [1, 2, 0, 6, 10]
     b = len(lens)
     total_seqlen = sum(lens)
 
@@ -91,7 +91,7 @@ def ref(q,k,v):
 @pytest.mark.parametrize("m", [1, 2]) # for MQA/GQA
 def test_varlen_flash_bwd(m, h, d, causal, local, dtype, seqused_k_limit):
     window_size = (3,3) if local else (-1,-1)
-    lens = [1, 2, 4, 6, 8, 10]
+    lens = [1, 2, 0, 6, 10]
     b = len(lens)
     total_seqlen = sum(lens)
     if seqused_k_limit is not None and (causal or local):