[FA][Test] Fix Sparse FA test (#6306)

2025-05-26 04:27:48 -04:00
parent 0d503090aa
commit 006ead9dcb
1 changed files with 1 additions and 0 deletions
--- a/sgl-kernel/tests/test_sparse_flash_attn.py
+++ b/sgl-kernel/tests/test_sparse_flash_attn.py
@@ -10,6 +10,7 @@ from sgl_kernel.sparse_flash_attn import (
    sparse_attn_func,
    sparse_attn_varlen_func,
 )
+from test_flash_attention import construct_local_mask


 def ref_attn(