app.py

import gradio as gr
from transformers import AutoConfig
from huggingface_hub import list_models
import asyncio
from typing import List
import time
from functools import lru_cache
import json
from datetime import datetime, timedelta
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed

# Credits: This implementation is derived from and builds upon the excellent work by gaunernst
# Original implementation: https://huggingface.co/spaces/gaunernst/kv-cache-calculator

search_cache = {}

POPULAR_MODELS = [
    "Qwen/Qwen3-30B-A3B",
    "meta-llama/Llama-3.1-8B-Instruct",
    "meta-llama/Llama-3.1-70B-Instruct", 
    "microsoft/DialoGPT-medium",
    "microsoft/DialoGPT-large",
    "mistralai/Mistral-7B-Instruct-v0.3",
    "mistralai/Mixtral-8x7B-Instruct-v0.1",
    "deepseek-ai/DeepSeek-V2-Chat",
    "deepseek-ai/DeepSeek-V3-Base",
    "google/gemma-2-9b",
    "google/gemma-2-27b",
    "Qwen/QwQ-32B-Preview",
    "Qwen/Qwen2.5-72B-Instruct",
    "anthropic/claude-3-haiku-20240307",
]

# Static GPU specifications (performance specs don't change, only prices do)
# All GPUs with SM_80+ compute capability (Flash Attention support)
GPU_SPECS = {
    # Consumer RTX 30 Series (Ampere - GA102/GA104/GA106) - SM_8.6
    "RTX 3060": {"memory_gb": 12, "compute_capability": "8.6", "tflops_fp32": 13.0, "category": "Consumer"},
    "RTX 3060 Ti": {"memory_gb": 8, "compute_capability": "8.6", "tflops_fp32": 16.2, "category": "Consumer"},
    "RTX 3070": {"memory_gb": 8, "compute_capability": "8.6", "tflops_fp32": 20.3, "category": "Consumer"},
    "RTX 3070 Ti": {"memory_gb": 8, "compute_capability": "8.6", "tflops_fp32": 21.7, "category": "Consumer"},
    "RTX 3080": {"memory_gb": 10, "compute_capability": "8.6", "tflops_fp32": 29.8, "category": "Consumer"},
    "RTX 3080 Ti": {"memory_gb": 12, "compute_capability": "8.6", "tflops_fp32": 34.1, "category": "Consumer"},
    "RTX 3090": {"memory_gb": 24, "compute_capability": "8.6", "tflops_fp32": 35.6, "category": "Consumer"},
    "RTX 3090 Ti": {"memory_gb": 24, "compute_capability": "8.6", "tflops_fp32": 40.0, "category": "Consumer"},
    
    # Consumer RTX 40 Series (Ada Lovelace - AD102/AD103/AD104/AD106/AD107) - SM_8.9
    "RTX 4060": {"memory_gb": 8, "compute_capability": "8.9", "tflops_fp32": 15.1, "category": "Consumer"},
    "RTX 4060 Ti": {"memory_gb": 8, "compute_capability": "8.9", "tflops_fp32": 22.1, "category": "Consumer"},
    "RTX 4060 Ti 16GB": {"memory_gb": 16, "compute_capability": "8.9", "tflops_fp32": 22.1, "category": "Consumer"},
    "RTX 4070": {"memory_gb": 12, "compute_capability": "8.9", "tflops_fp32": 29.1, "category": "Consumer"},
    "RTX 4070 Super": {"memory_gb": 12, "compute_capability": "8.9", "tflops_fp32": 35.5, "category": "Consumer"},
    "RTX 4070 Ti": {"memory_gb": 12, "compute_capability": "8.9", "tflops_fp32": 40.1, "category": "Consumer"},
    "RTX 4070 Ti Super": {"memory_gb": 16, "compute_capability": "8.9", "tflops_fp32": 44.1, "category": "Consumer"},
    "RTX 4080": {"memory_gb": 16, "compute_capability": "8.9", "tflops_fp32": 48.7, "category": "Consumer"},
    "RTX 4080 Super": {"memory_gb": 16, "compute_capability": "8.9", "tflops_fp32": 52.2, "category": "Consumer"},
    "RTX 4090": {"memory_gb": 24, "compute_capability": "8.9", "tflops_fp32": 83.0, "category": "Consumer"},
    
    # Consumer RTX 50 Series (Blackwell - GB202/GB203/GB205/GB206/GB207) - SM_10.0
    "RTX 5060": {"memory_gb": 8, "compute_capability": "10.0", "tflops_fp32": 18.5, "category": "Consumer"},
    "RTX 5060 Ti": {"memory_gb": 16, "compute_capability": "10.0", "tflops_fp32": 28.2, "category": "Consumer"},
    "RTX 5070": {"memory_gb": 12, "compute_capability": "10.0", "tflops_fp32": 35.1, "category": "Consumer"},
    "RTX 5070 Ti": {"memory_gb": 16, "compute_capability": "10.0", "tflops_fp32": 48.3, "category": "Consumer"},
    "RTX 5080": {"memory_gb": 16, "compute_capability": "10.0", "tflops_fp32": 60.5, "category": "Consumer"},
    "RTX 5090": {"memory_gb": 32, "compute_capability": "10.0", "tflops_fp32": 125.0, "category": "Consumer"},
    
    # Professional/Workstation RTX A Series (Ampere) - SM_8.6
    "RTX A2000": {"memory_gb": 12, "compute_capability": "8.6", "tflops_fp32": 8.0, "category": "Workstation"},
    "RTX A4000": {"memory_gb": 16, "compute_capability": "8.6", "tflops_fp32": 19.2, "category": "Workstation"},
    "RTX A4500": {"memory_gb": 20, "compute_capability": "8.6", "tflops_fp32": 23.7, "category": "Workstation"},
    "RTX A5000": {"memory_gb": 24, "compute_capability": "8.6", "tflops_fp32": 27.8, "category": "Workstation"},
    "RTX A6000": {"memory_gb": 48, "compute_capability": "8.6", "tflops_fp32": 38.7, "category": "Workstation"},
    
    # Professional RTX 6000 Ada (Ada Lovelace) - SM_8.9  
    "RTX 6000 Ada": {"memory_gb": 48, "compute_capability": "8.9", "tflops_fp32": 91.1, "category": "Workstation"},
    
    # Datacenter A100 Series (Ampere) - SM_8.0
    "A100 40GB": {"memory_gb": 40, "compute_capability": "8.0", "tflops_fp32": 19.5, "category": "Datacenter"},
    "A100 80GB": {"memory_gb": 80, "compute_capability": "8.0", "tflops_fp32": 19.5, "category": "Datacenter"},
    
    # Datacenter H100 Series (Hopper) - SM_9.0
    "H100 80GB": {"memory_gb": 80, "compute_capability": "9.0", "tflops_fp32": 67.0, "category": "Datacenter"},
    "H100 94GB": {"memory_gb": 94, "compute_capability": "9.0", "tflops_fp32": 67.0, "category": "Datacenter"},
    
    # Datacenter H200 (Hopper) - SM_9.0
    "H200 141GB": {"memory_gb": 141, "compute_capability": "9.0", "tflops_fp32": 67.0, "category": "Datacenter"},
    
    # Datacenter B200 (Blackwell) - SM_10.0
    "B200 180GB": {"memory_gb": 180, "compute_capability": "10.0", "tflops_fp32": 80.0, "category": "Datacenter"},
    
    # Datacenter L40/L40S (Ada Lovelace) - SM_8.9
    "L40": {"memory_gb": 48, "compute_capability": "8.9", "tflops_fp32": 91.6, "category": "Datacenter"},
    "L40S": {"memory_gb": 48, "compute_capability": "8.9", "tflops_fp32": 91.6, "category": "Datacenter"},
}

# Price cache with timestamp
price_cache = {}
PRICE_CACHE_DURATION = timedelta(hours=6)  # Cache prices for 6 hours

def fetch_single_gpu_price(gpu_name):
    """Fetch price for a single GPU (used in parallel)"""
    try:
        print(f"Fetching price for {gpu_name}...")
        price = get_gpu_price(gpu_name)
        if price:
            print(f"Found price for {gpu_name}: ${price}")
            return gpu_name, price
        else:
            print(f"✗ No price found for {gpu_name}, using fallback")
            return gpu_name, get_fallback_price(gpu_name)
    except Exception as e:
        print(f"✗ Error fetching {gpu_name}: {e}")
        return gpu_name, get_fallback_price(gpu_name)

def preload_gpu_prices():
    """Pre-fetch all GPU prices in parallel on startup"""
    print("Pre-loading GPU prices...")
    start_time = time.time()
    
    # Get list of GPUs to price
    gpu_names = list(GPU_SPECS.keys())
    
    # Use ThreadPoolExecutor for parallel requests
    with ThreadPoolExecutor(max_workers=8) as executor:
        # Submit all price fetch tasks
        future_to_gpu = {executor.submit(fetch_single_gpu_price, gpu_name): gpu_name 
                        for gpu_name in gpu_names}
        
        # Collect results as they complete
        for future in as_completed(future_to_gpu):
            gpu_name, price = future.result()
            # Store in cache with timestamp
            cache_key = gpu_name.lower().replace(" ", "_")
            price_cache[cache_key] = {
                "price": price,
                "timestamp": datetime.now()
            }
    
    end_time = time.time()
    total_time = end_time - start_time
    print(f"Loaded prices for {len(gpu_names)} GPUs in {total_time:.1f} seconds")
    print(f"Cache contains {len(price_cache)} price entries")

def start_price_preloading():
    """Start price preloading in background thread"""
    def preload_worker():
        preload_gpu_prices()
    
    # Start preloading in background
    preload_thread = threading.Thread(target=preload_worker, daemon=True)
    preload_thread.start()
    print("Price preloading started in background...")

def get_gpu_price(gpu_name):
    """Get GPU price from curated pricing data"""
    current_time = datetime.now()
    
    # Check cache first
    cache_key = gpu_name.lower().replace(" ", "_")
    if cache_key in price_cache:
        cached_data = price_cache[cache_key]
        if current_time - cached_data["timestamp"] < PRICE_CACHE_DURATION:
            return cached_data["price"]
    
    price = get_fallback_price(gpu_name)
    
    # Cache the result
    price_cache[cache_key] = {
        "price": price,
        "timestamp": current_time
    }
    
    return price


def get_fallback_price(gpu_name):
    """Curated GPU pricing data"""
    fallback_prices = {
        # Consumer RTX 30 Series
        "RTX 3060": 280,
        "RTX 3060 Ti": 320,
        "RTX 3070": 420,
        "RTX 3070 Ti": 480,
        "RTX 3080": 580,
        "RTX 3080 Ti": 720,
        "RTX 3090": 950,
        "RTX 3090 Ti": 1100,
        
        # Consumer RTX 40 Series
        "RTX 4060": 300,
        "RTX 4060 Ti": 380,
        "RTX 4060 Ti 16GB": 480,
        "RTX 4070": 580,
        "RTX 4070 Super": 680,
        "RTX 4070 Ti": 780,
        "RTX 4070 Ti Super": 880,
        "RTX 4080": 980,
        "RTX 4080 Super": 880,
        "RTX 4090": 1500,
        
        # Consumer RTX 50 Series (Expected pricing)
        "RTX 5060": 400,
        "RTX 5060 Ti": 600,
        "RTX 5070": 800,
        "RTX 5070 Ti": 1000,
        "RTX 5080": 1200,
        "RTX 5090": 2000,
        
        # Professional/Workstation GPUs
        "RTX A2000": 650,
        "RTX A4000": 1200,
        "RTX A4500": 2200,
        "RTX A5000": 2800,
        "RTX A6000": 4500,
        "RTX 6000 Ada": 6800,
        
        # Datacenter GPUs (current enterprise pricing)
        "A100 40GB": 12000,
        "A100 80GB": 15000,
        "H100 80GB": 30000,
        "H100 94GB": 35000,
        "H200 141GB": 40000,
        "B200 180GB": 50000,
        "L40": 9000,
        "L40S": 10000,
    }
    return fallback_prices.get(gpu_name, 1000)

def search_models_fast(query: str, max_results: int = 30) -> List[str]:
    if not query or len(query.strip()) < 1:
        return POPULAR_MODELS[:15]
    
    query = query.strip()
    cache_key = f"{query.lower()}_{max_results}"
    
    current_time = time.time()
    if cache_key in search_cache:
        cached_result, cache_time = search_cache[cache_key]
        if current_time - cache_time < 300:
            return cached_result
    
    try:
        print(f"Searching HF Hub for: {query}")
        
        all_matches = []
        seen_models = set()
        
        for model in POPULAR_MODELS:
            if query.lower() in model.lower() and model not in seen_models:
                all_matches.append(model)
                seen_models.add(model)
        
        models = list_models(
            search=query,
            task="text-generation", 
            library="transformers",
            sort="downloads",
            direction=-1,
            limit=max_results,
            full=False
        )
        
        for model in models:
            if model.id not in seen_models and len(all_matches) < max_results:
                all_matches.append(model.id)
                seen_models.add(model.id)
        
        result = all_matches[:max_results]
        search_cache[cache_key] = (result, current_time)
        
        if len(search_cache) > 15:
            oldest_key = min(search_cache.keys(), key=lambda k: search_cache[k][1])
            del search_cache[oldest_key]
            
        return result
        
    except Exception as e:
        print(f"Search error: {e}")
        popular_matches = [model for model in POPULAR_MODELS if query.lower() in model.lower()]
        return popular_matches if popular_matches else POPULAR_MODELS[:15]


def calculate(name: str, ctx_len: int, num_users: int, dtype: str, hf_token: str):
    if not name or not name.strip():
        raise gr.Error("Please search for and select a model first")
    
    name = name.strip()
    hf_token = hf_token.strip()
    try:
        cfg = AutoConfig.from_pretrained(
            name,
            trust_remote_code=True,
            token=hf_token or None,
        )
    except Exception as e:
        raise gr.Error(e)

    use_mla = cfg.architectures[0].startswith(("DeepseekV2", "DeepseekV3"))

    if hasattr(cfg, "text_config"):
        cfg = cfg.text_config

    num_layers = cfg.num_hidden_layers
    num_attention_heads = cfg.num_attention_heads
    num_kv_heads = getattr(cfg, "num_key_value_heads", num_attention_heads)
    
    if use_mla:
        attention_type = "MLA"
    elif num_kv_heads == num_attention_heads:
        attention_type = "MHA"
    else:
        attention_type = "GQA"
    
    model_config = [
        ["num_layers", num_layers],
        ["max_ctx_len", cfg.max_position_embeddings],
        ["attention_type", attention_type],
        ["num_attention_heads", num_attention_heads],
        ["num_kv_heads", num_kv_heads],
    ]
    if ctx_len > cfg.max_position_embeddings:
        gr.Warning(
            "Requested context length is larger than the max value supported by the model"
        )

    if use_mla:
        kv_lora_rank = cfg.kv_lora_rank
        qk_rope_head_dim = cfg.qk_rope_head_dim
        nelems_per_token = num_layers * (kv_lora_rank + qk_rope_head_dim)

        model_config.append(["kv_lora_rank", kv_lora_rank])
        model_config.append(["qk_rope_head_dim", qk_rope_head_dim])
        model_config.append(["calc_formula", f"{num_layers} * ({kv_lora_rank} + {qk_rope_head_dim})"])

    else:
        head_dim = getattr(cfg, "head_dim", cfg.hidden_size // num_attention_heads)
        nelems_per_token = num_layers * num_kv_heads * head_dim * 2

        model_config.append(["head_dim", head_dim])
        if attention_type == "GQA":
            kv_ratio = num_attention_heads // num_kv_heads
            model_config.append(["gqa_ratio", f"{kv_ratio}:1"])
        model_config.append(["calc_formula", f"{num_layers} * {num_kv_heads} * {head_dim} * 2"])

    if dtype == "fp16/bf16":
        nbytes_per_elem = 2
    elif dtype == "fp8":
        nbytes_per_elem = 1 + 2 / cfg.hidden_size  # assume per-token scaling
    elif dtype == "fp4":
        nbytes_per_elem = 0.5 + 2 / 32  # 4-bit weights + scaling factor every 32 elements (MXFP4)

    kv_cache_size = nelems_per_token * ctx_len * num_users * nbytes_per_elem / 1e9
    
    # Get GPU recommendations with complete memory analysis using actual config
    gpu_recommendations = recommend_gpus(
        kv_cache_size_gb=kv_cache_size,
        config=cfg,
        dtype=dtype,
        ctx_len=ctx_len,
        num_users=num_users
    )
    
    return kv_cache_size, model_config, gpu_recommendations


DESCRIPTION = (
    "Calculate KV cache memory requirements for transformer models. "
    "Supports MHA, GQA, and MLA attention mechanisms with fp16/bf16, fp8, and fp4 data types."
)

def search_models_on_submit(search_query):
    if not search_query or len(search_query.strip()) < 2:
        return [
            gr.Textbox(interactive=True),
            gr.Dropdown(choices=[], value="", visible=False),
            gr.Button(interactive=True)
        ]
    
    query_stripped = search_query.strip()
    
    search_results = search_models_fast(query_stripped, max_results=30)
    
    if query_stripped not in search_results:
        search_results.insert(0, query_stripped)
    
    return [
        gr.Textbox(interactive=True, value=query_stripped), 
        gr.Dropdown(
            choices=search_results, 
            value=query_stripped,
            visible=True,
            info=f"Found {len(search_results)} models - select one"
        ),
        gr.Button(interactive=True)
    ]

def update_selection_from_dropdown(dropdown_value):
    return gr.Textbox(value=dropdown_value)

def estimate_model_memory(config, dtype):
    """Estimate model weight memory requirements in GB using actual config object"""
    try:
        if not config:
            return 5.0  # Default fallback
        
        # Extract parameters for calculation
        num_layers = getattr(config, 'num_hidden_layers', getattr(config, 'num_layers', 32))
        hidden_size = getattr(config, 'hidden_size', getattr(config, 'd_model', 4096))
        vocab_size = getattr(config, 'vocab_size', 50000)
        intermediate_size = getattr(config, 'intermediate_size', hidden_size * 4)
        
        # DeepSeek V3 specific parameter calculation following the exact formula
        # Check if this is DeepSeek V3 architecture
        is_deepseek_v3 = (getattr(config, 'model_type', '') == 'deepseek_v3' or
                          any('deepseek' in arch.lower() for arch in getattr(config, 'architectures', [])))
        
        if is_deepseek_v3 and hasattr(config, 'q_lora_rank'):
            # DeepSeek V3 specific calculation
            # Config constants
            L = num_layers  # 61
            H = hidden_size  # 7168
            I = intermediate_size  # 18432
            I_moe = getattr(config, 'moe_intermediate_size', 2048)  # 2048
            n_h = getattr(config, 'num_attention_heads', 128)  # 128
            r_q = getattr(config, 'q_lora_rank', 1536)  # 1536
            r_kv = getattr(config, 'kv_lora_rank', 512)  # 512
            V = vocab_size  # 129,280
            
            # Additional config values
            qk_nope_head_dim = getattr(config, 'qk_nope_head_dim', 128)
            qk_rope_head_dim = getattr(config, 'qk_rope_head_dim', 64)
            v_head_dim = getattr(config, 'v_head_dim', 128)
            
            # Attention per layer calculation
            # W_q,a: H × r_q
            w_q_a = H * r_q
            # W_q,b: r_q × n_h × (qk_nope + qk_rope)
            w_q_b = r_q * n_h * (qk_nope_head_dim + qk_rope_head_dim)
            # W_kv,a: H × (r_kv + qk_rope)
            w_kv_a = H * (r_kv + qk_rope_head_dim)
            # W_kv,b: r_kv × n_h × (qk_nope + v)
            w_kv_b = r_kv * n_h * (qk_nope_head_dim + v_head_dim)
            # W_o: (n_h × v) × H
            w_o = (n_h * v_head_dim) * H
            
            attention_per_layer = w_q_a + w_q_b + w_kv_a + w_kv_b + w_o
            total_attention = L * attention_per_layer
            
            # Dense FFN layers (first 3 layers)
            dense_ffn_per_layer = 3 * H * I  # 3 projections: gate, up, down
            total_dense_ffn = 3 * dense_ffn_per_layer  # 3 dense layers
            
            # MoE FFN layers (remaining 58 layers)
            moe_ffn_per_expert = 3 * H * I_moe
            n_routed_experts = getattr(config, 'n_routed_experts', 256)  # 256
            n_shared_experts = getattr(config, 'n_shared_experts', 1)  # 1
            experts_per_moe_layer = n_routed_experts + n_shared_experts  # 257
            moe_ffn_per_layer = experts_per_moe_layer * moe_ffn_per_expert
            moe_layers = L - 3  # 58 MoE layers
            total_moe_ffn = moe_layers * moe_ffn_per_layer
            
            # Embeddings + LM head (untied)
            embeddings_and_head = 2 * V * H
            
            # Total parameters
            total_params = total_attention + total_dense_ffn + total_moe_ffn + embeddings_and_head
            
            print(f"DEBUG: DeepSeek V3 parameter breakdown:")
            print(f"  Attention ({L} layers): {total_attention/1e9:.2f}B")
            print(f"  Dense FFN (3 layers): {total_dense_ffn/1e9:.2f}B")
            print(f"  MoE FFN ({moe_layers} layers): {total_moe_ffn/1e9:.2f}B")
            print(f"  Embeddings + Head: {embeddings_and_head/1e9:.2f}B")
            print(f"  Total calculated: {total_params/1e9:.1f}B parameters")
            
        else:
            # Fallback to standard transformer calculation for other models
            num_attention_heads = getattr(config, 'num_attention_heads', hidden_size // 64)
            num_kv_heads = getattr(config, 'num_key_value_heads', num_attention_heads)
            head_dim = getattr(config, 'head_dim', hidden_size // num_attention_heads)
            
            # Standard attention calculation
            q_params = hidden_size * (num_attention_heads * head_dim)
            kv_params = hidden_size * (num_kv_heads * head_dim) * 2
            o_params = (num_attention_heads * head_dim) * hidden_size
            attention_params_per_layer = q_params + kv_params + o_params
            attention_params = num_layers * attention_params_per_layer
            
            # Standard FFN calculation
            ffn_params = num_layers * (2 * hidden_size * intermediate_size + intermediate_size * hidden_size)
            
            # Embeddings
            embedding_params = vocab_size * hidden_size
            
            # Other parameters
            other_params = num_layers * 2 * hidden_size + hidden_size
            
            total_params = embedding_params + attention_params + ffn_params + other_params
            
            print(f"DEBUG: Standard transformer parameter breakdown:")
            print(f"  Embeddings: {embedding_params/1e9:.1f}B")
            print(f"  Attention: {attention_params/1e9:.1f}B") 
            print(f"  FFN: {ffn_params/1e9:.1f}B")
            print(f"  Other: {other_params/1e9:.1f}B")
            print(f"  Total calculated: {total_params/1e9:.1f}B parameters")
        
        # Convert to memory based on user-selected dtype
        if dtype == "fp16/bf16":
            bytes_per_param = 2
        elif dtype == "fp8":
            bytes_per_param = 1
        elif dtype == "fp4":
            bytes_per_param = 0.5
        else:
            bytes_per_param = 4  # fp32 fallback
        
        model_memory_gb = (total_params * bytes_per_param) / (1024**3)
        
        # Add minimal overhead (5% for loading)
        model_memory_gb *= 1.05
        
        return model_memory_gb
        
    except Exception as e:
        print(f"Error estimating model memory from config: {e}")
        return 70.0  # Conservative fallback for large models


def estimate_activation_memory(ctx_len, num_users, config):
    """Estimate activation memory requirements in GB using actual config object"""
    try:
        if not config:
            return 1.0  # Default fallback
            
        # Extract parameters directly from config object
        hidden_size = getattr(config, 'hidden_size', getattr(config, 'd_model', 4096))
        
        batch_size = num_users
        
        # For inference, activations are much smaller than training
        # Only need to store activations for current forward pass, not gradients
        
        # 1. Input/output activations: batch_size * ctx_len * hidden_size
        io_activations = batch_size * ctx_len * hidden_size
        
        # 2. Intermediate activations (only a few layers worth, not all)
        # Most activations are computed and immediately used, not stored
        intermediate_size = getattr(config, 'intermediate_size', hidden_size * 4)
        stored_activations = batch_size * ctx_len * intermediate_size * 2  # Only ~2 layers worth
        
        # 3. Attention scores for current layer (not all layers stored)
        num_attention_heads = getattr(config, 'num_attention_heads', hidden_size // 64)
        attention_scores = batch_size * num_attention_heads * ctx_len * ctx_len
        
        # Total activation elements (much smaller for inference)
        total_activation_elements = io_activations + stored_activations + attention_scores
        
        # Convert to memory (fp16 = 2 bytes per element)
        activation_memory_gb = (total_activation_elements * 2) / (1024**3)
        
        # Cap at reasonable values for inference (activations shouldn't dominate)
        max_reasonable_gb = max(5.0, ctx_len * batch_size / 10000)  # Reasonable scaling
        activation_memory_gb = min(activation_memory_gb, max_reasonable_gb)
        
        return max(0.5, activation_memory_gb)  # At least 500MB
        
    except Exception as e:
        print(f"Error estimating activation memory from config: {e}")
        # Simple fallback based on context length
        try:
            # Much simpler formula for inference
            fallback_gb = (num_users * ctx_len * 4096 * 4 * 2) / (1024**3)  # Conservative
            return min(10.0, max(0.5, fallback_gb))  # Cap at 10GB
        except:
            return 2.0  # Default 2GB

def calculate_multi_gpu_configs(total_memory_needed, suitable_gpus):
    """Calculate multi-GPU configurations for large models (power-of-2 for tensor parallelism)"""
    multi_gpu_configs = []
    
    # Power-of-2 configurations for tensor parallelism (TP) - max 8 for practical use
    gpu_counts = [1, 2, 4, 8]  # Only powers of 2, max 8 GPUs
    
    # For large models, check all high-memory GPUs, not just top 3 cost-effective ones
    gpus_to_check = suitable_gpus if total_memory_needed > 500 else suitable_gpus[:3]
    
    for gpu in gpus_to_check:
        for count in gpu_counts:
            total_gpu_memory = gpu["memory_gb"] * count
            
            if total_gpu_memory >= total_memory_needed:
                # Calculate per-GPU memory utilization
                memory_per_gpu = total_memory_needed / count
                utilization = (memory_per_gpu / gpu["memory_gb"]) * 100
                
                # Skip very inefficient configurations (< 30% utilization for multi-GPU)
                if count > 1 and utilization < 30:
                    continue
                    
                # Calculate total cost
                total_cost = gpu["price_usd"] * count
                cost_per_tflop_total = total_cost / (gpu["tflops_fp32"] * count)
                
                # Format configuration name with TP indication
                if count == 1:
                    config_name = gpu['name']
                else:
                    config_name = f"{count}x {gpu['name']} (TP={count})"
                
                
                multi_gpu_configs.append({
                    "config": config_name,
                    "gpu_count": count,
                    "total_memory_gb": total_gpu_memory,
                    "memory_per_gpu": memory_per_gpu,
                    "utilization": utilization,
                    "total_cost": total_cost,
                    "cost_per_tflop": cost_per_tflop_total,
                    "base_gpu": gpu
                })
                
                # For single GPU, only add once
                if count == 1:
                    break
    
    # Sort by cost-effectiveness (total cost per TFLOP)
    multi_gpu_configs.sort(key=lambda x: x["cost_per_tflop"])
    
    return multi_gpu_configs[:8]  # Return top 8 configurations

def recommend_gpus(kv_cache_size_gb, config=None, dtype="fp16/bf16", ctx_len=128000, num_users=1):
    """Recommend cost-effective GPU configurations (single and multi-GPU with tensor parallelism) for complete memory footprint"""
    if not kv_cache_size_gb or kv_cache_size_gb <= 0:
        print("DEBUG: KV cache size is 0 or invalid")
        return []
    
    # Calculate complete memory footprint using actual config object
    model_memory_gb = estimate_model_memory(config, dtype)
    activation_memory_gb = estimate_activation_memory(ctx_len, num_users, config)
    
    # Total memory = Model weights + KV cache + Activations + Safety buffer
    total_memory_needed = model_memory_gb + kv_cache_size_gb + activation_memory_gb + 1.0  # 1GB safety buffer
    
    print(f"DEBUG: Memory breakdown - Model: {model_memory_gb:.1f}GB, KV: {kv_cache_size_gb:.1f}GB, Activations: {activation_memory_gb:.1f}GB, Total: {total_memory_needed:.1f}GB")
    
    # Get all GPUs with real pricing (from cache or live fetch)
    all_gpus = []
    
    for gpu_name, specs in GPU_SPECS.items():
        # Get real-time price (will use cache if available)
        current_price = get_gpu_price(gpu_name)
        if current_price:
            cost_per_tflop = current_price / specs["tflops_fp32"]
            all_gpus.append({
                "name": gpu_name,
                "memory_gb": specs["memory_gb"],
                "compute_capability": specs["compute_capability"],
                "tflops_fp32": specs["tflops_fp32"],
                "price_usd": current_price,
                "cost_per_tflop": cost_per_tflop,
                "category": specs.get("category", "Consumer")
            })
    
    print(f"DEBUG: Found {len(all_gpus)} GPUs with pricing")
    
    if not all_gpus:
        print("DEBUG: No GPUs found with pricing")
        return []
    
    # Sort by cost-effectiveness for single GPU evaluation
    all_gpus.sort(key=lambda x: x["cost_per_tflop"])
    
    # Calculate multi-GPU configurations
    multi_gpu_configs = calculate_multi_gpu_configs(total_memory_needed, all_gpus)
    
    print(f"DEBUG: Generated {len(multi_gpu_configs)} GPU configurations")
    
    if not multi_gpu_configs:
        print("DEBUG: No valid GPU configurations found")
        return []
    
    # Format recommendations
    recommendations = []
    for i, config in enumerate(multi_gpu_configs):
        rank = f"#{i+1}"
        
        price_source = "Live" if config["base_gpu"]["name"].lower().replace(" ", "_") in price_cache else "Est"
        
        # Format configuration display
        config_display = f"{rank} {config['config']}"
        
        # Calculate FLOP/dollar (TFLOPS per dollar)
        total_tflops = config["base_gpu"]["tflops_fp32"] * config["gpu_count"]
        flops_per_dollar = total_tflops / config['total_cost']
        
        recommendations.append([
            config_display,
            f"{flops_per_dollar:.3f}",
            f"{total_memory_needed:.1f}GB",
            f"${config['total_cost']:.0f}"
        ])
    
    return recommendations

with gr.Blocks(title="KV Cache Calculator", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# KV Cache Calculator")
    gr.Markdown(DESCRIPTION)
    
    with gr.Row():
        with gr.Column():
            model_search = gr.Textbox(
                label="🔍 Search Model", 
                placeholder="Type your model ID here.",
            )
            
            model_dropdown = gr.Dropdown(
                label="📋 Select from Results", 
                choices=[],
                value="",
                visible=False,
                info="Choose from search results"
            )
            
            with gr.Row():
                gr.Markdown("**💡 Tip:** Type model names like 'llama', 'qwen', 'mistral', then press Enter to search")
            
            ctx_len = gr.Number(label="Context Length", value=128_000, minimum=1)
            num_users = gr.Number(label="Number of Users", value=1, minimum=1)
            dtype = gr.Dropdown(
                label="KV Cache Data Type", 
                choices=["fp16/bf16", "fp8", "fp4"], 
                value="fp16/bf16"
            )
            hf_token = gr.Textbox(
                label="HuggingFace Token (optional)", 
                type="password", 
                placeholder="For gated models"
            )
            
            calculate_btn = gr.Button("Calculate KV Cache Size", variant="primary")
        
        with gr.Column():
            cache_size = gr.Number(label="KV Cache Size (GB)", precision=2)
            model_config = gr.Dataframe(
                label="Model Configuration", 
                headers=["Parameter", "Value"], 
                datatype=["str", "str"],
                wrap=True
            )
    
            gpu_recommendations = gr.Dataframe(
                label="GPU Recommendations",
                headers=["Configuration", "TFLOPS/$", "Memory", "Price"],
                datatype=["str", "str", "str", "str"],
                wrap=False,
                visible=False
            )
    
    model_search.submit(
        fn=search_models_on_submit,
        inputs=[model_search],
        outputs=[model_search, model_dropdown, calculate_btn],
        show_progress="minimal"
    )
    
    model_dropdown.change(
        fn=update_selection_from_dropdown,
        inputs=[model_dropdown],
        outputs=[model_search],
        show_progress=False
    )
    
    def calculate_and_show_gpus(model_name, ctx_len, num_users, dtype, hf_token):
        cache_size, model_config, gpu_recs = calculate(model_name, ctx_len, num_users, dtype, hf_token)
        
        print(f"DEBUG: GPU recommendations count: {len(gpu_recs) if gpu_recs else 0}")
        if gpu_recs:
            print(f"DEBUG: First recommendation: {gpu_recs[0] if gpu_recs else 'None'}")
        
        if gpu_recs:
            return (
                cache_size, 
                model_config, 
                gr.Dataframe(value=gpu_recs, visible=True)
            )
        else:
            print("DEBUG: No GPU recommendations found, showing empty table")
            return (
                cache_size, 
                model_config, 
                gr.Dataframe(value=[], visible=False)
            )
    
    calculate_btn.click(
        fn=calculate_and_show_gpus,
        inputs=[model_search, ctx_len, num_users, dtype, hf_token],
        outputs=[cache_size, model_config, gpu_recommendations]
    )

demo.css = """
.gradio-container {
    max-width: 1400px !important;
    margin: 0 auto !important;
}

/* Make dataframes wider and prevent text wrapping */
.gradio-dataframe {
    width: 100% !important;
    min-width: 800px !important;
}

.gradio-dataframe table {
    width: 100% !important;
    table-layout: auto !important;
}

.gradio-dataframe td, .gradio-dataframe th {
    white-space: nowrap !important;
    padding: 8px 12px !important;
    text-overflow: ellipsis !important;
    min-width: 120px !important;
}

/* Style disabled textboxes to be clearly disabled */
.gradio-textbox:disabled,
.gradio-textbox[aria-disabled="true"] {
    opacity: 0.6 !important;
    background-color: #f5f5f5 !important;
    color: #666 !important;
    cursor: not-allowed !important;
    border-color: #ccc !important;
}

/* Style placeholder text */
.gradio-textbox input::placeholder {
    color: #999 !important;
    font-style: italic;
}

/* Make disabled dropdowns more visually obvious */
.gradio-dropdown[data-testid="dropdown"]:disabled,
.gradio-dropdown[data-testid="dropdown"][aria-disabled="true"] {
    opacity: 0.6 !important;
    background-color: #f5f5f5 !important;
    cursor: not-allowed !important;
}

/* Make disabled buttons more obvious too */
button:disabled {
    opacity: 0.5 !important;
    background-color: #e0e0e0 !important;
    cursor: not-allowed !important;
}
"""

if __name__ == "__main__":
    # Start price preloading in background before launching the app
    start_price_preloading()
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        allowed_paths=[],
        app_kwargs={"docs_url": None, "redoc_url": None}
    )