fix

Author: Iamleos

python/sgl_jax/srt/layers/attention/flash_attn_kernel/flash_attention.py

@@ -115,7 +115,7 @@ def ref_ragged_paged_attention(
         if custom_mask != None:
             raise ValueError(f"use causal mask, custom_mask is not None")
     else:
-        if custom_mask == None or custom_mask.size() < jnp.cumsum(kv_lens)[-1]:
+        if custom_mask == None or custom_mask.size < jnp.cumsum(kv_lens)[-1]:
             raise ValueError(
                 f"use custom_mask, custom_mask length must larger than total kv length"
             )

python/sgl_jax/test/test_flashattention.py

@@ -247,6 +247,7 @@ def align_to_size(l, size, value=0):
     attention_backend = FlashAttentionBackend(
         num_heads, num_kv_heads, head_dim, page_size=page_size, mesh=mesh
     )
+    print(f"!!!!!!!! {causal=}")
     if not causal:
         forward_mode = ForwardMode.TARGET_VERIFY
         custom_mask = create_custom_mask(lens)
@@ -307,8 +308,10 @@ def align_to_size(l, size, value=0):
         cache_loc=cache_loc,
         extend_prefix_lens=extend_prefix_lens,
         extend_seq_lens=extend_seq_lens,
+        spec_info=spec_info,
     )
     fb.attn_backend.forward_metadata = attention_backend.get_forward_metadata(mwb)
+
     return fb, q, k, v