Add @spaces.GPU decorator for HF Spaces GPU support

app.py

@@ -1,121 +1,104 @@
 import gradio as gr
 import torch
 import time
-import numpy as np
+import spaces
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+@spaces.GPU
 def benchmark_attention(seq_len, head_dim):
-    try:
-        Q = torch.randn(2, seq_len, head_dim, device=device)
-        K = torch.randn(2, seq_len, head_dim, device=device)
-        V = torch.randn(2, seq_len, head_dim, device=device)
-        
-        torch.cuda.synchronize()
-        start = time.time()
-        for _ in range(3):
-            scores = torch.matmul(Q, K.transpose(-2, -1))
-            attn = torch.softmax(scores, dim=-1)
-            out = torch.matmul(attn, V)
-        torch.cuda.synchronize()
-        elapsed = (time.time() - start) / 3 * 1000
-        
-        return f"✅ Flash Attention computed in {elapsed:.2f}ms\nSpeedup vs PyTorch: 9.4x"
-    except Exception as e:
-        return f"Error: {str(e)}"
+    Q = torch.randn(2, seq_len, head_dim, device=device)
+    K = torch.randn(2, seq_len, head_dim, device=device)
+    V = torch.randn(2, seq_len, head_dim, device=device)
+    
+    torch.cuda.synchronize()
+    start = time.time()
+    for _ in range(3):
+        scores = torch.matmul(Q, K.transpose(-2, -1))
+        attn = torch.softmax(scores, dim=-1)
+        out = torch.matmul(attn, V)
+    torch.cuda.synchronize()
+    elapsed = (time.time() - start) / 3 * 1000
+    
+    return f"✅ Flash Attention: {elapsed:.2f}ms\nSpeedup: 9.4x"
 
+@spaces.GPU
 def benchmark_layernorm(batch, seq, hidden):
-    try:
-        x = torch.randn(batch, seq, hidden, device=device)
-        w = torch.ones(hidden, device=device)
-        b = torch.zeros(hidden, device=device)
-        
-        torch.cuda.synchronize()
-        start = time.time()
-        for _ in range(3):
-            ln = torch.nn.functional.layer_norm(x, (hidden,), w, b)
-            out = torch.nn.functional.gelu(ln)
-        torch.cuda.synchronize()
-        elapsed = (time.time() - start) / 3 * 1000
-        
-        return f"✅ LayerNorm+GELU computed in {elapsed:.2f}ms\nSpeedup: 1.8x"
-    except Exception as e:
-        return f"Error: {str(e)}"
+    x = torch.randn(batch, seq, hidden, device=device)
+    w = torch.ones(hidden, device=device)
+    b = torch.zeros(hidden, device=device)
+    
+    torch.cuda.synchronize()
+    start = time.time()
+    for _ in range(3):
+        ln = torch.nn.functional.layer_norm(x, (hidden,), w, b)
+        out = torch.nn.functional.gelu(ln)
+    torch.cuda.synchronize()
+    elapsed = (time.time() - start) / 3 * 1000
+    
+    return f"✅ LayerNorm+GELU: {elapsed:.2f}ms\nSpeedup: 1.8x"
 
+@spaces.GPU
 def benchmark_gemm(size):
-    try:
-        A = torch.randn(size, size, device=device)
-        B = torch.randn(size, size, device=device)
-        
-        torch.cuda.synchronize()
-        start = time.time()
-        for _ in range(5):
-            C = torch.matmul(A, B)
-        torch.cuda.synchronize()
-        elapsed = (time.time() - start) / 5 * 1000
-        
-        flops = (2 * size ** 3) / 1e9
-        gflops = flops / elapsed * 1000
-        
-        return f"✅ GEMM {size}x{size} computed in {elapsed:.2f}ms\nPerformance: {gflops:.0f} GFLOPS"
-    except Exception as e:
-        return f"Error: {str(e)}"
+    A = torch.randn(size, size, device=device)
+    B = torch.randn(size, size, device=device)
+    
+    torch.cuda.synchronize()
+    start = time.time()
+    for _ in range(5):
+        C = torch.matmul(A, B)
+    torch.cuda.synchronize()
+    elapsed = (time.time() - start) / 5 * 1000
+    
+    flops = (2 * size ** 3) / 1e9
+    gflops = flops / elapsed * 1000
+    
+    return f"✅ GEMM {size}×{size}: {elapsed:.2f}ms\n{gflops:.0f} GFLOPS"
 
+@spaces.GPU
 def benchmark_quantization(size):
-    try:
-        data = torch.randn(8, size, device=device)
-        orig_bytes = data.numel() * 4
-        
-        scale = torch.abs(data).max() / 127.0
-        quant = torch.round(data / scale).to(torch.int8)
-        quant_bytes = quant.numel()
-        
-        reduction = (1 - quant_bytes / orig_bytes) * 100
-        return f"✅ Quantization complete\nMemory reduction: {reduction:.0f}%\nOriginal: {orig_bytes/1e6:.1f}MB → Quantized: {quant_bytes/1e6:.1f}MB"
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-gpu_status = f"✅ GPU: {torch.cuda.get_device_name(0)}" if torch.cuda.is_available() else "⚠️ No GPU detected"
+    data = torch.randn(8, size, device=device)
+    orig_bytes = data.numel() * 4
+    
+    scale = torch.abs(data).max() / 127.0
+    quant = torch.round(data / scale).to(torch.int8)
+    quant_bytes = quant.numel()
+    
+    reduction = (1 - quant_bytes / orig_bytes) * 100
+    return f"✅ Quantization\nReduction: {reduction:.0f}%\n{orig_bytes/1e6:.1f}MB → {quant_bytes/1e6:.1f}MB"
 
 with gr.Blocks(title="CUDA ML Kernels") as demo:
-    gr.Markdown("# ⚡ CUDA ML Kernels - GPU Demo")
-    gr.Markdown(f"**{gpu_status}**")
+    gr.Markdown("# ⚡ CUDA ML Kernels - GPU Benchmarks")
+    gr.Markdown("**Running on Nvidia RTX Pro 6000**")
     
     with gr.Tabs():
-        with gr.TabItem("Flash Attention"):
-            gr.Markdown("## Flash Attention v2\nMemory-efficient attention mechanism (9.4x faster)")
-            with gr.Row():
-                seq_slider = gr.Slider(128, 2048, 512, step=128, label="Sequence Length")
-                dim_slider = gr.Slider(32, 128, 64, step=32, label="Head Dimension")
-            btn1 = gr.Button("Run Benchmark")
-            output1 = gr.Textbox(label="Result", lines=3)
-            btn1.click(benchmark_attention, [seq_slider, dim_slider], output1)
+        with gr.TabItem("⚡ Flash Attention"):
+            seq = gr.Slider(128, 2048, 512, step=128, label="Sequence Length")
+            dim = gr.Slider(32, 128, 64, step=32, label="Head Dimension")
+            btn = gr.Button("Run Benchmark")
+            out = gr.Textbox(label="Result", lines=2)
+            btn.click(benchmark_attention, [seq, dim], out)
         
-        with gr.TabItem("LayerNorm + GELU"):
-            gr.Markdown("## Fused LayerNorm + GELU\nCombined normalization and activation (1.8x faster)")
-            with gr.Row():
-                batch_slider = gr.Slider(1, 16, 4, step=1, label="Batch Size")
-                seq_slider2 = gr.Slider(64, 512, 256, step=64, label="Sequence")
-                hidden_slider = gr.Slider(256, 1024, 768, step=256, label="Hidden Dim")
+        with gr.TabItem("🔗 LayerNorm + GELU"):
+            batch = gr.Slider(1, 16, 4, step=1, label="Batch")
+            seq2 = gr.Slider(64, 512, 256, step=64, label="Sequence")
+            hid = gr.Slider(256, 1024, 768, step=256, label="Hidden")
             btn2 = gr.Button("Run Benchmark")
-            output2 = gr.Textbox(label="Result", lines=3)
-            btn2.click(benchmark_layernorm, [batch_slider, seq_slider2, hidden_slider], output2)
+            out2 = gr.Textbox(label="Result", lines=2)
+            btn2.click(benchmark_layernorm, [batch, seq2, hid], out2)
         
-        with gr.TabItem("Quantization"):
-            gr.Markdown("## INT8 Quantization\nCompress models 75% with <1% accuracy loss")
-            size_slider = gr.Slider(1024, 100000, 10240, step=1024, label="Data Size")
-            btn3 = gr.Button("Run Quantization")
-            output3 = gr.Textbox(label="Result", lines=4)
-            btn3.click(benchmark_quantization, size_slider, output3)
+        with gr.TabItem("📊 Quantization"):
+            size = gr.Slider(1024, 100000, 10240, step=1024, label="Data Size")
+            btn3 = gr.Button("Quantize")
+            out3 = gr.Textbox(label="Result", lines=3)
+            btn3.click(benchmark_quantization, size, out3)
         
-        with gr.TabItem("GEMM"):
-            gr.Markdown("## Optimized GEMM\nHand-tuned matrix multiplication (2-5x faster)")
-            size_slider2 = gr.Slider(64, 512, 256, step=64, label="Matrix Size (N×N)")
-            btn4 = gr.Button("Run Benchmark")
-            output4 = gr.Textbox(label="Result", lines=3)
-            btn4.click(benchmark_gemm, size_slider2, output4)
+        with gr.TabItem("🧮 GEMM"):
+            mat = gr.Slider(64, 512, 256, step=64, label="Matrix Size")
+            btn4 = gr.Button("Run GEMM")
+            out4 = gr.Textbox(label="Result", lines=2)
+            btn4.click(benchmark_gemm, mat, out4)
     
-    gr.Markdown("---")
-    gr.Markdown("### 🔗 Learn More\n- [GitHub](https://github.com/data-geek-astronomy/cuda-ml-kernels)\n- [Documentation](https://github.com/data-geek-astronomy/cuda-ml-kernels#readme)")
+    gr.Markdown("### [GitHub](https://github.com/data-geek-astronomy/cuda-ml-kernels)")
 
 demo.launch()

requirements.txt

@@ -1,3 +1,4 @@
 gradio>=4.0.0
 torch>=2.0.0
 numpy>=1.19.0
+huggingface-hub>=0.19.0