Grouped Query Attention › Choosing the Right Number of KV Heads

def analyze_gqa_tradeoffs(
    num_query_heads: int,
    d_model: int,
    num_layers: int,
    target_memory_gb: float,
    seq_length: int,
    batch_size: int,
):
    """Find GQA configurations that fit within a memory budget."""
    bytes_per_element = 2  # fp16
    head_dim = d_model // num_query_heads

    viable_configs = []

    for num_kv_heads in [1, 2, 4, 8, 16, num_query_heads]:
        if num_query_heads % num_kv_heads != 0:
            continue

        # KV cache memory
        cache_bytes = (
            2
            * num_layers
            * num_kv_heads
            * head_dim
            * seq_length
            * batch_size
            * bytes_per_element
        )
        cache_gb = cache_bytes / 1e9

        # KV projection parameters
        kv_params = 2 * d_model * num_kv_heads * head_dim

        # Estimated quality (simplified model based on research)
        quality_factor = min(
            1.0, 0.92 + 0.08 * (num_kv_heads / num_query_heads) ** 0.5
        )

        config = {
            "kv_heads": num_kv_heads,
            "cache_gb": cache_gb,
            "kv_params_m": kv_params / 1e6,
            "quality": quality_factor * 100,
            "fits_budget": cache_gb <= target_memory_gb,
            "memory_reduction": num_query_heads / num_kv_heads,
        }
        viable_configs.append(config)

    return viable_configs


# Example: 70B model on 80GB GPU with 8K context
configs = analyze_gqa_tradeoffs(
    num_query_heads=64,
    d_model=8192,
    num_layers=80,
    target_memory_gb=20.0,  # Reserve 60GB for model weights
    seq_length=8192,
    batch_size=8,
)