Aligns benchmarks with sparse attn imports

LoserCheems · LoserCheems · commit a0ed87d916ec · 2025-11-09T23:36:20.000+08:00
Updates benchmark integrations to load the flash_sparse_attn implementations so the renamed package continues to back the CUDA, Triton, and Flex runs.
Renames the availability guards and status messages to keep diagnostic output aligned with the new module namespace.
diff --git a/benchmarks/backward_equivalence.py b/benchmarks/backward_equivalence.py
@@ -21,33 +21,33 @@
 
 # Import the compiled CUDA extension
 try:
-    from flash_dmattn.flash_dmattn_interface import flash_dmattn_func
-    print("✅ Successfully imported flash_dmattn interface")
+    from flash_sparse_attn.flash_sparse_attn_interface import flash_sparse_attn_func
+    print("✅ Successfully imported flash_sparse_attn interface")
 except ImportError as e:
-    print(f"❌ Failed to import flash_dmattn interface: {e}")
+    print(f"❌ Failed to import flash_sparse_attn interface: {e}")
     print("Please make sure the package is properly installed with: pip install .")
     # Don't exit here, just warn
-    flash_dmattn_func = None
+    flash_sparse_attn_func = None
 
 # Import the Triton implementation
 try:
-    from flash_dmattn.flash_dmattn_triton import triton_dmattn_func
-    print("✅ Successfully imported flash_dmattn_triton")
+    from flash_sparse_attn.flash_sparse_attn_triton import triton_sparse_attn_func
+    print("✅ Successfully imported flash_sparse_attn_triton")
 except ImportError as e:
-    print(f"❌ Failed to import flash_dmattn_triton: {e}")
+    print(f"❌ Failed to import flash_sparse_attn_triton: {e}")
     print("Please make sure the Triton implementation is available.")
     # Don't exit here, just warn
-    triton_dmattn_func = None
+    triton_sparse_attn_func = None
 
 # Import the Flex Attention implementation
 try:
-    from flash_dmattn.flash_dmattn_flex import flex_dmattn_func
-    print("✅ Successfully imported flash_dmattn_flex")
+    from flash_sparse_attn.flash_sparse_attn_flex import flex_sparse_attn_func
+    print("✅ Successfully imported flash_sparse_attn_flex")
 except ImportError as e:
-    print(f"❌ Failed to import flash_dmattn_flex: {e}")
+    print(f"❌ Failed to import flash_sparse_attn_flex: {e}")
     print("Please make sure the Flex Attention implementation is available.")
     # Don't exit here, just warn
-    flex_dmattn_func = None
+    flex_sparse_attn_func = None
 
 
 def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
@@ -189,7 +189,7 @@ def dynamic_mask_attention_cuda(
     Returns:
         tuple: (attn_outputs, dq, dk, dv, dbias)
     """
-    if flash_dmattn_func is None:
+    if flash_sparse_attn_func is None:
         raise ImportError("CUDA implementation not available")
 
     query_states_leaf = query_states
@@ -210,8 +210,8 @@ def dynamic_mask_attention_cuda(
     key_states = key_states.transpose(1, 2).contiguous()            # [batch, key_len, num_kv_heads, head_dim]
     value_states = value_states.transpose(1, 2).contiguous()        # [batch, key_len, num_kv_heads, head_dim]
 
-    # Call the flash_dmattn_func interface
-    attn_outputs = flash_dmattn_func(
+    # Call the flash_sparse_attn_func interface
+    attn_outputs = flash_sparse_attn_func(
         query=query_states,
         key=key_states,
         value=value_states,
@@ -256,7 +256,7 @@ def dynamic_mask_attention_triton(
     Returns:
         tuple: (attn_outputs, dq, dk, dv, dbias)
     """
-    if triton_dmattn_func is None:
+    if triton_sparse_attn_func is None:
         raise RuntimeError("Triton implementation not available")
     
     _, num_heads, _, _ = query_states.shape
@@ -288,7 +288,7 @@ def dynamic_mask_attention_triton(
     value_states = value_states.transpose(1, 2)         # [batch, key_len, num_heads, head_dim]  
 
     # Call the Triton implementation
-    attn_outputs = triton_dmattn_func(
+    attn_outputs = triton_sparse_attn_func(
         query=query_states,
         key=key_states,
         value=value_states,
@@ -330,7 +330,7 @@ def dynamic_mask_attention_flex(
     Returns:
         tuple: (attn_outputs, dq, dk, dv, dbias)
     """
-    if flex_dmattn_func is None:
+    if flex_sparse_attn_func is None:
         raise RuntimeError("Flex Attention implementation not available")
     
     _, num_heads, _, _ = query_states.shape
@@ -359,7 +359,7 @@ def dynamic_mask_attention_flex(
     attn_bias = attn_bias.contiguous()                              # [batch, num_heads, seqlen_q, seqlen_k]
     
     # Call the Flex Attention implementation
-    attn_outputs = flex_dmattn_func(
+    attn_outputs = flex_sparse_attn_func(
         query_states,
         key_states,
         value_states,
@@ -474,7 +474,7 @@ def test_cuda_backward_equivalence(accuracy_threshold=0.95):
     print("🚀" + "=" * 76 + "🚀")
 
     # Check if CUDA implementation is available
-    if flash_dmattn_func is None:
+    if flash_sparse_attn_func is None:
         print("❌ CUDA implementation not available, skipping test.")
         return False
     
@@ -734,7 +734,7 @@ def test_triton_backward_equivalence(accuracy_threshold=0.95):
     print("🚀" + "=" * 76 + "🚀")
 
     # Check if Triton implementation is available
-    if triton_dmattn_func is None:
+    if triton_sparse_attn_func is None:
         print("❌ Triton implementation not available, skipping test.")
         return False
     
diff --git a/benchmarks/backward_performance.py b/benchmarks/backward_performance.py
@@ -28,33 +28,33 @@
 
 # Import the compiled CUDA extension
 try:
-    from flash_dmattn.flash_dmattn_interface import flash_dmattn_func
-    print("✅ Successfully imported flash_dmattn interface")
+    from flash_sparse_attn.flash_sparse_attn_interface import flash_sparse_attn_func
+    print("✅ Successfully imported flash_sparse_attn interface")
 except ImportError as e:
-    print(f"❌ Failed to import flash_dmattn interface: {e}")
+    print(f"❌ Failed to import flash_sparse_attn interface: {e}")
     print("Please make sure the package is properly installed with: pip install .")
     # Don't exit here, just warn
-    flash_dmattn_func = None
+    flash_sparse_attn_func = None
 
 # Import the Triton implementation
 try:
-    from flash_dmattn.flash_dmattn_triton import triton_dmattn_func
-    print("✅ Successfully imported flash_dmattn_triton")
+    from flash_sparse_attn.flash_sparse_attn_triton import triton_sparse_attn_func
+    print("✅ Successfully imported flash_sparse_attn_triton")
 except ImportError as e:
-    print(f"❌ Failed to import flash_dmattn_triton: {e}")
+    print(f"❌ Failed to import flash_sparse_attn_triton: {e}")
     print("Please make sure the Triton implementation is available.")
     # Don't exit here, just warn
-    triton_dmattn_func = None
+    triton_sparse_attn_func = None
 
 # Import the Flex Attention implementation
 try:
-    from flash_dmattn.flash_dmattn_flex import flex_dmattn_func
-    print("✅ Successfully imported flash_dmattn_flex")
+    from flash_sparse_attn.flash_sparse_attn_flex import flex_sparse_attn_func
+    print("✅ Successfully imported flash_sparse_attn_flex")
 except ImportError as e:
-    print(f"❌ Failed to import flash_dmattn_flex: {e}")
+    print(f"❌ Failed to import flash_sparse_attn_flex: {e}")
     print("Please make sure the Flex Attention implementation is available.")
     # Don't exit here, just warn
-    flex_dmattn_func = None
+    flex_sparse_attn_func = None
 
 
 def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
@@ -207,7 +207,7 @@ def dynamic_mask_attention_backward_cuda(
     Returns:
         tuple: (output_tensor, timing_ms) or ("OOM", 0) or ("Not Available", 0)
     """
-    if flash_dmattn_func is None:
+    if flash_sparse_attn_func is None:
         return "Not Available", 0
 
     attn_bias, attn_mask = prepare_mask(
@@ -223,7 +223,7 @@ def dynamic_mask_attention_backward_cuda(
     value_states = value_states.transpose(1, 2).contiguous()        # [batch, key_len, num_kv_heads, head_dim]
 
     try:
-        attn_outputs = flash_dmattn_func(
+        attn_outputs = flash_sparse_attn_func(
             query=query_states,
             key=key_states,
             value=value_states,
@@ -277,7 +277,7 @@ def dynamic_mask_attention_backward_triton(
     Returns:
         tuple: (output_tensor, timing_ms) or ("OOM", 0) or ("Not Available", 0)
     """
-    if triton_dmattn_func is None:
+    if triton_sparse_attn_func is None:
         return "Not Available", 0
     
     _, num_heads, _, _ = query_states.shape
@@ -305,7 +305,7 @@ def dynamic_mask_attention_backward_triton(
     attn_bias = attn_bias.contiguous()                              # [batch, num_heads, seqlen_q, seqlen_k]
 
     try:
-        attn_outputs = triton_dmattn_func(
+        attn_outputs = triton_sparse_attn_func(
             query=query_states,
             key=key_states,
             value=value_states,
@@ -356,7 +356,7 @@ def dynamic_mask_attention_backward_flex(
     Returns:
         tuple: (output_tensor, timing_ms) or ("OOM", 0) or ("Not Available", 0)
     """
-    if flex_dmattn_func is None:
+    if flex_sparse_attn_func is None:
         return "Not Available", 0
     
     _, num_heads, _, _ = query_states.shape
@@ -384,7 +384,7 @@ def dynamic_mask_attention_backward_flex(
     attn_bias = attn_bias.contiguous()                              # [batch, num_heads, seqlen_q, seqlen_k]
 
     try:
-        attn_outputs = flex_dmattn_func(
+        attn_outputs = flex_sparse_attn_func(
             query_states,
             key_states,
             value_states,
diff --git a/benchmarks/forward_equivalence.py b/benchmarks/forward_equivalence.py
@@ -21,33 +21,33 @@
 
 # Import the compiled CUDA extension
 try:
-    from flash_dmattn.flash_dmattn_interface import flash_dmattn_func
-    print("✅ Successfully imported flash_dmattn interface")
+    from flash_sparse_attn.flash_sparse_attn_interface import flash_sparse_attn_func
+    print("✅ Successfully imported flash_sparse_attn interface")
 except ImportError as e:
-    print(f"❌ Failed to import flash_dmattn interface: {e}")
+    print(f"❌ Failed to import flash_sparse_attn interface: {e}")
     print("Please make sure the package is properly installed with: pip install .")
     # Don't exit here, just warn
-    flash_dmattn_func = None
+    flash_sparse_attn_func = None
 
 # Import the Triton implementation
 try:
-    from flash_dmattn.flash_dmattn_triton import triton_dmattn_func
-    print("✅ Successfully imported flash_dmattn_triton")
+    from flash_sparse_attn.flash_sparse_attn_triton import triton_sparse_attn_func
+    print("✅ Successfully imported flash_sparse_attn_triton")
 except ImportError as e:
-    print(f"❌ Failed to import flash_dmattn_triton: {e}")
+    print(f"❌ Failed to import flash_sparse_attn_triton: {e}")
     print("Please make sure the Triton implementation is available.")
     # Don't exit here, just warn
-    triton_dmattn_func = None
+    triton_sparse_attn_func = None
 
 # Import the Flex Attention implementation
 try:
-    from flash_dmattn.flash_dmattn_flex import flex_dmattn_func
-    print("✅ Successfully imported flash_dmattn_flex")
+    from flash_sparse_attn.flash_sparse_attn_flex import flex_sparse_attn_func
+    print("✅ Successfully imported flash_sparse_attn_flex")
 except ImportError as e:
-    print(f"❌ Failed to import flash_dmattn_flex: {e}")
+    print(f"❌ Failed to import flash_sparse_attn_flex: {e}")
     print("Please make sure the Flex Attention implementation is available.")
     # Don't exit here, just warn
-    flex_dmattn_func = None
+    flex_sparse_attn_func = None
 
 
 def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
@@ -181,8 +181,8 @@ def dynamic_mask_attention_cuda(
     Returns:
         attn_outputs: [batch_size, query_len, num_heads, head_dim]
     """
-    if flash_dmattn_func is None:
-        raise RuntimeError("flash_dmattn_func not available")
+    if flash_sparse_attn_func is None:
+        raise RuntimeError("flash_sparse_attn_func not available")
 
     attn_bias, attn_mask = prepare_mask(
         query_states,
@@ -196,8 +196,8 @@ def dynamic_mask_attention_cuda(
     key_states = key_states.transpose(1, 2)         # [batch, key_len, num_kv_heads, head_dim]
     value_states = value_states.transpose(1, 2)     # [batch, key_len, num_kv_heads, head_dim]
 
-    # Call the flash_dmattn_func interface
-    attn_outputs = flash_dmattn_func(
+    # Call the flash_sparse_attn_func interface
+    attn_outputs = flash_sparse_attn_func(
         query_states,
         key_states,
         value_states,
@@ -239,7 +239,7 @@ def dynamic_mask_attention_triton(
     Returns:
         attn_outputs: [batch_size, query_len, num_heads, head_dim]
     """
-    if triton_dmattn_func is None:
+    if triton_sparse_attn_func is None:
         raise RuntimeError("Triton implementation not available")
     
     _, num_heads, _, _ = query_states.shape
@@ -267,7 +267,7 @@ def dynamic_mask_attention_triton(
     attn_bias = attn_bias.contiguous()                          # [batch, num_heads, seqlen_q, seqlen_k]
     
     # Call the Triton implementation
-    attn_outputs = triton_dmattn_func(
+    attn_outputs = triton_sparse_attn_func(
         query_states,
         key_states,
         value_states,
@@ -306,7 +306,7 @@ def dynamic_mask_attention_flex(
     Returns:
         attn_outputs: [batch_size, query_len, num_heads, head_dim]
     """
-    if flex_dmattn_func is None:
+    if flex_sparse_attn_func is None:
         raise RuntimeError("Flex Attention implementation not available")
     
     _, num_heads, _, _ = query_states.shape
@@ -334,7 +334,7 @@ def dynamic_mask_attention_flex(
     attn_bias = attn_bias.contiguous()                              # [batch, num_heads, seqlen_q, seqlen_k]
 
     # Call the Flex Attention implementation
-    attn_outputs = flex_dmattn_func(
+    attn_outputs = flex_sparse_attn_func(
         query_states,
         key_states,
         value_states,
@@ -446,7 +446,7 @@ def test_cuda_forward_equivalence(accuracy_threshold=0.95):
     print("🚀" + "=" * 76 + "🚀")
     
     # Check if CUDA implementation is available
-    if flash_dmattn_func is None:
+    if flash_sparse_attn_func is None:
         print("❌ CUDA implementation not available, skipping test.")
         return False
     
@@ -653,7 +653,7 @@ def test_triton_forward_equivalence(accuracy_threshold=0.95):
     print("🔬 Testing Forward Pass Equivalence: Python vs Triton 🔬")
     print("🔥" + "=" * 76 + "🔥")
     
-    if triton_dmattn_func is None:
+    if triton_sparse_attn_func is None:
         print("❌ Triton implementation not available, skipping Triton tests")
         return False
     
@@ -859,7 +859,7 @@ def test_flex_forward_equivalence(accuracy_threshold=0.95):
     print("🔬 Testing Forward Pass Equivalence: Python vs Flex Attention 🔬")
     print("🌟" + "=" * 76 + "🌟")
     
-    if flex_dmattn_func is None:
+    if flex_sparse_attn_func is None:
         print("❌ Flex Attention implementation not available, skipping Flex Attention tests")
         return False
     
diff --git a/benchmarks/forward_performance.py b/benchmarks/forward_performance.py
