From 006ead9dcbec360c0a656ee84d34446dc80c88c7 Mon Sep 17 00:00:00 2001
From: Brayden Zhong <b8zhong@uwaterloo.ca>
Date: Mon, 26 May 2025 04:27:48 -0400
Subject: [PATCH] [FA][Test] Fix Sparse FA test (#6306)

---
 sgl-kernel/tests/test_sparse_flash_attn.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/sgl-kernel/tests/test_sparse_flash_attn.py b/sgl-kernel/tests/test_sparse_flash_attn.py
index 4ddb6d7f5..4aa0a7c19 100644
--- 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(