# Llama 2 70B Configuration
config_llama_70b = {
    "n_heads": 64,
    "n_kv_heads": 8,  # GQA with 8 groups
    "group_size": 8,  # 64/8 = 8 heads per group
    "hidden_size": 8192,
    "head_dim": 128
}

# Falcon 40B Configuration  
config_falcon_40b = {
    "n_heads": 64,
    "n_kv_heads": 1,  # MQA
    "group_size": 64,  # All heads share
    "hidden_size": 8192,
    "head_dim": 128
}

# Mistral 7B Configuration
config_mistral_7b = {
    "n_heads": 32,
    "n_kv_heads": 8,  # GQA
    "group_size": 4,
    "hidden_size": 4096,
    "head_dim": 128,
    "sliding_window": 4096  # Additional optimization
}

# Use MHA for best results
config = {
    "attention_type": "MHA",
    "n_heads": 32,
    "n_kv_heads": 32,  # Same as n_heads
    "rationale": "Quality matters most, resources available"
}

# Use GQA for balance
config = {
    "attention_type": "GQA",
    "n_heads": 32,
    "n_kv_heads": 8,  # 4:1 ratio
    "rationale": "Balance quality with serving costs"
}

# Use MQA for efficiency
config = {
    "attention_type": "MQA",
    "n_heads": 32,
    "n_kv_heads": 1,
    "rationale": "Memory/power constraints critical"
}

# Use aggressive GQA or MQA
config = {
    "attention_type": "GQA",
    "n_heads": 32,
    "n_kv_heads": 4,  # 8:1 ratio
    "rationale": "KV cache dominates at long context"
}

def convert_attention_type(model, from_type, to_type):
    """Convert model between attention types"""
    
    if from_type == "MHA" and to_type == "GQA":
        # Average K,V weights within groups
        return average_kv_weights_for_groups(model)
    
    elif from_type == "MHA" and to_type == "MQA":
        # Average all K,V weights
        return average_all_kv_weights(model)
    
    elif from_type == "GQA" and to_type == "MQA":
        # Average group K,V weights
        return average_group_weights(model)
    
    elif from_type in ["GQA", "MQA"] and to_type == "MHA":
        # Replicate K,V weights (requires fine-tuning)
        return replicate_kv_weights(model)

# Recommended fine-tuning settings
fine_tune_config = {
    "MHA_to_GQA": {
        "lr": 1e-5,
        "steps": 10000,
        "warmup": 1000
    },
    "MHA_to_MQA": {
        "lr": 5e-6,
        "steps": 20000,
        "warmup": 2000
    }
}

class OptimizedAttention(nn.Module):
    """Combine multiple optimizations"""
    
    def __init__(self, config):
        super().__init__()
        
        # Choose base attention
        if config.n_kv_heads == 1:
            self.attention = MQA(config)
        elif config.n_kv_heads < config.n_heads:
            self.attention = GQA(config)
        else:
            self.attention = MHA(config)
        
        # Add optimizations
        if config.use_flash:
            self.attention = FlashWrapper(self.attention)
        
        if config.use_sliding_window:
            self.attention = SlidingWindowWrapper(self.attention)
        
        if config.use_rope:
            self.pos_encoding = RoPE(config)

def estimate_max_batch_size(gpu_memory_gb, model_params_gb, attention_type, context_len):
    """Estimate maximum batch size for given GPU"""
    
    # Available memory for KV cache
    available = gpu_memory_gb - model_params_gb - 2  # 2GB overhead
    
    # KV cache per sequence (GB)
    kv_cache_size = {
        "MHA": context_len * n_layers * n_heads * head_dim * 2 * 4 / 1e9,
        "GQA-8": context_len * n_layers * 8 * head_dim * 2 * 4 / 1e9,
        "MQA": context_len * n_layers * 1 * head_dim * 2 * 4 / 1e9
    }
    
    return int(available / kv_cache_size[attention_type])

# Example: 80GB A100, 30GB model, 4K context
print(estimate_max_batch_size(80, 30, "MHA", 4096))    # ~15
print(estimate_max_batch_size(80, 30, "GQA-8", 4096))  # ~60  
print(estimate_max_batch_size(80, 30, "MQA", 4096))    # ~480