Update app.py

app.py

@@ -8,7 +8,12 @@ import contextlib
 # ⚙️ Setup
 # ============================================================
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-flash_attn = get_kernel("kernels-community/vllm-flash-attn3")
+
+if torch.cuda.is_available():
+    flash_attn = get_kernel("kernels-community/vllm-flash-attn3")
+    flash_attn_func = flash_attn.flash_attn_func
+else:
+    flash_attn_func = None  # CPU fallback
 
 # ============================================================
 # 🧠 Reference attention (PyTorch SDPA)
@@ -19,20 +24,6 @@ def reference_attention(query, key, value, causal=False):
         out = torch.nn.functional.scaled_dot_product_attention(query, key, value, is_causal=causal)
     return out.transpose(1, 2).contiguous()
 
-def var_reference_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, causal=False):
-    batch_size = cu_seqlens_q.shape[0] - 1
-    total_tokens_q = q.shape[0]
-    out = torch.zeros((total_tokens_q, q.shape[1], q.shape[2]), device=q.device, dtype=q.dtype)
-    for b in range(batch_size):
-        start_q, end_q = cu_seqlens_q[b], cu_seqlens_q[b + 1]
-        start_k, end_k = cu_seqlens_k[b], cu_seqlens_k[b + 1]
-        q_slice = q[start_q:end_q].unsqueeze(0)
-        k_slice = k[start_k:end_k].unsqueeze(0)
-        v_slice = v[start_k:end_k].unsqueeze(0)
-        attn_out = reference_attention(q_slice, k_slice, v_slice, causal=causal)
-        out[start_q:end_q] = attn_out.squeeze(0)
-    return out
-
 # ============================================================
 # 🎛️ Test function
 # ============================================================
@@ -48,25 +39,30 @@ def run_flash_attention(B=2, S=5, H=4, D=8, seed=42):
         print(f"Running FlashAttention Tests on device: {device}")
         print(f"Input shape: B={B}, S={S}, H={H}, D={D}\n")
 
-        # Standard attention
+        # Always run PyTorch reference attention
         out_ref = reference_attention(q, k, v)
-        try:
-            out_flash = flash_attn["flash_attn_func"](q, k, v, causal=False)
-        except TypeError:
-            out_flash, _ = flash_attn["flash_attn_func"](q, k, v, causal=False)
-        print("1. Standard attention:")
-        print(f"  Reference: {out_ref.shape}, Flash: {out_flash.shape}")
-        print(f"  Outputs close: {torch.allclose(out_flash, out_ref, atol=1e-2, rtol=1e-3)}\n")
-
-        # Causal attention
-        out_ref_causal = reference_attention(q, k, v, causal=True)
-        try:
-            out_causal = flash_attn["flash_attn_func"](q, k, v, causal=True)
-        except TypeError:
-            out_causal, _ = flash_attn["flash_attn_func"](q, k, v, causal=True)
-        print("2. Causal attention:")
-        print(f"  Reference: {out_ref_causal.shape}, Flash: {out_causal.shape}")
-        print(f"  Outputs close: {torch.allclose(out_causal, out_ref_causal, atol=1e-2, rtol=1e-3)}\n")
+        print(f"✅ Reference attention OK: {out_ref.shape}\n")
+
+        # Run FlashAttention if CUDA available
+        if flash_attn_func is not None:
+            try:
+                out_flash, _ = flash_attn_func(q, k, v, causal=False)
+                print("⚡ FlashAttention (non-causal):")
+                print(f"  Output: {out_flash.shape}")
+                print(f"  Close to reference: {torch.allclose(out_flash, out_ref, atol=1e-2, rtol=1e-3)}\n")
+
+                out_flash_causal, _ = flash_attn_func(q, k, v, causal=True)
+                out_ref_causal = reference_attention(q, k, v, causal=True)
+                print("⚡ FlashAttention (causal):")
+                print(f"  Output: {out_flash_causal.shape}")
+                print(f"  Close to reference: {torch.allclose(out_flash_causal, out_ref_causal, atol=1e-2, rtol=1e-3)}\n")
+
+            except Exception as e:
+                print("❌ FlashAttention test failed:")
+                print(str(e))
+        else:
+            print("⚠️ CUDA not available — FlashAttention test skipped.")
+            print("Using PyTorch SDPA (reference) only.\n")
 
     return log.getvalue()
 
@@ -75,7 +71,7 @@ def run_flash_attention(B=2, S=5, H=4, D=8, seed=42):
 # ============================================================
 with gr.Blocks(title="Flash Attention Kernel Tester") as demo:
     gr.Markdown("## ⚡ Flash Attention Kernel Tester")
-    gr.Markdown("Compare PyTorch SDPA vs FlashAttention implementations interactively.")
+    gr.Markdown("Compare PyTorch SDPA vs FlashAttention implementations interactively. Works on both CPU and CUDA.")
 
     with gr.Row():
         B = gr.Slider(1, 8, value=2, step=1, label="Batch Size (B)")
@@ -90,4 +86,4 @@ with gr.Blocks(title="Flash Attention Kernel Tester") as demo:
 
     run_btn.click(run_flash_attention, inputs=[B, S, H, D, seed], outputs=output)
 
-demo.launch(share=True)
+demo.launch()