fix

Author: Iamleos

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

@@ -115,6 +115,7 @@ def ref_ragged_paged_attention(
         if custom_mask != None:
             raise ValueError(f"use causal mask, custom_mask is not None")
     else:
+        print(f"######### {custom_mask=} {kv_lens=}")
         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

@@ -405,6 +405,7 @@ def run_test(self, mode, lens, mode_args):
             cache_loc_list.append(padded_page_indices)
         page_table = jnp.stack(cache_loc_list)
 
+        print(f"@@@@@ {forward_batch.spec_info=}")
         expected = ref_ragged_paged_attention(
             q.reshape(q.shape[0], num_heads, head_dim),
             k.reshape(k.shape[0] // page_size, page_size, num_kv_heads, head_dim),