Fix Phase Attention Benchmark to simulate 24 layers

app.py

@@ -237,34 +237,49 @@ with gr.Blocks(title="Kalpanā API — ZeroGPU NVIDIA A100", theme=gr.themes.Sof
                     for seq_len in ctx_lengths:
                         res = {"context_length": seq_len}
                         
-                        # 1. Profile Phase Attention
+                        num_layers = 24
+                        num_heads = 14
+                        embed_dim = 896
+                        head_dim = 64
+                        
+                        # 1. Profile Kalpana (Full 24 Layers)
                         torch.cuda.empty_cache()
                         torch.cuda.reset_peak_memory_stats()
-                        phase_attn = TrueO1PhaseAttentionLayer(embed_dim=64, num_heads=2, bands=2048, device=device)
-                        Q = torch.randn(1, 2, 32, device=device) # dim=32
-                        phase_attn.current_t = seq_len
-                        # Simulate forward pass
                         try:
-                            _ = phase_attn.forward(Q)
+                            Q_layer = torch.randn(1, num_heads, head_dim, device=device)
+                            phase_layers = []
+                            for i in range(num_layers):
+                                p_attn = TrueO1PhaseAttentionLayer(embed_dim=embed_dim, num_heads=num_heads, bands=2048, device=device)
+                                p_attn.current_t = seq_len
+                                phase_layers.append(p_attn)
+                            
+                            for p_attn in phase_layers:
+                                _ = p_attn.forward(Q_layer)
+                            
                             res["kalpana_peak_vram_mb"] = round(torch.cuda.max_memory_allocated() / (1024*1024), 2)
+                            del phase_layers, Q_layer
                         except Exception as e:
-                            res["kalpana_peak_vram_mb"] = f"OOM or Error: {str(e)}"
-                        del phase_attn, Q
+                            res["kalpana_peak_vram_mb"] = f"OOM: {str(e)[:50]}..."
                         
-                        # 2. Profile Standard SDPA (Generation Step)
+                        # 2. Profile Standard SDPA (Full 24 Layers)
                         torch.cuda.empty_cache()
                         torch.cuda.reset_peak_memory_stats()
                         try:
-                            # Standard KV cache: [batch, heads, seq_len, dim]
-                            K = torch.randn(1, 2, seq_len, 32, device=device)
-                            V = torch.randn(1, 2, seq_len, 32, device=device)
-                            Q_std = torch.randn(1, 2, 1, 32, device=device)
+                            K_list, V_list = [], []
+                            Q_std = torch.randn(1, num_heads, 1, head_dim, device=device)
                             import torch.nn.functional as F
-                            _ = F.scaled_dot_product_attention(Q_std, K, V)
+                            
+                            for i in range(num_layers):
+                                K = torch.randn(1, num_heads, seq_len, head_dim, device=device)
+                                V = torch.randn(1, num_heads, seq_len, head_dim, device=device)
+                                K_list.append(K)
+                                V_list.append(V)
+                                _ = F.scaled_dot_product_attention(Q_std, K, V)
+                                
                             res["standard_peak_vram_mb"] = round(torch.cuda.max_memory_allocated() / (1024*1024), 2)
-                            del K, V, Q_std
+                            del K_list, V_list, Q_std
                         except Exception as e:
-                            res["standard_peak_vram_mb"] = f"OOM or Error: {str(e)}"
+                            res["standard_peak_vram_mb"] = f"OOM: {str(e)[:50]}..."
                             
                         results.append(res)