Fix Phase Attention Benchmark to simulate GQA 2 KV heads

app.py

@@ -238,18 +238,20 @@ with gr.Blocks(title="Kalpanā API — ZeroGPU NVIDIA A100", theme=gr.themes.Sof
                         res = {"context_length": seq_len}
                         
                         num_layers = 24
-                        num_heads = 14
+                        num_q_heads = 14
+                        num_kv_heads = 2
                         embed_dim = 896
                         head_dim = 64
                         
-                        # 1. Profile Kalpana (Full 24 Layers)
+                        # 1. Profile Kalpana (Full 24 Layers, 2 KV Heads)
                         torch.cuda.empty_cache()
                         torch.cuda.reset_peak_memory_stats()
                         try:
-                            Q_layer = torch.randn(1, num_heads, head_dim, device=device)
+                            # Q is simulated for 14 heads but Kalpana state only needs 2 KV heads
+                            Q_layer = torch.randn(1, num_kv_heads, head_dim, device=device) # using kv heads for simple projection test
                             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 = TrueO1PhaseAttentionLayer(embed_dim=head_dim*num_kv_heads, num_heads=num_kv_heads, bands=2048, device=device)
                                 p_attn.current_t = seq_len
                                 phase_layers.append(p_attn)
                             
@@ -262,18 +264,18 @@ with gr.Blocks(title="Kalpanā API — ZeroGPU NVIDIA A100", theme=gr.themes.Sof
                         except Exception as e:
                             res["kalpana_peak_vram_mb"] = f"OOM: {str(e)[:50]}..."
                         
-                        # 2. Profile Standard SDPA (Full 24 Layers)
+                        # 2. Profile Standard SDPA (Full 24 Layers, 2 KV Heads)
                         torch.cuda.empty_cache()
                         torch.cuda.reset_peak_memory_stats()
                         try:
                             K_list, V_list = [], []
-                            Q_std = torch.randn(1, num_heads, 1, head_dim, device=device)
+                            Q_std = torch.randn(1, num_kv_heads, 1, head_dim, device=device)
                             import torch.nn.functional as F
                             
                             with torch.inference_mode():
                                 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 = torch.randn(1, num_kv_heads, seq_len, head_dim, device=device)
+                                    V = torch.randn(1, num_kv_heads, seq_len, head_dim, device=device)
                                     K_list.append(K)
                                     V_list.append(V)
                                     _ = F.scaled_dot_product_attention(Q_std, K, V)