Choose hardware from the workload bottleneck

"""GPU roofline model and hardware selector for prefill and decode workloads."""

from __future__ import annotations

from dataclasses import asdict, dataclass
import json


@dataclass(frozen=True)
class GpuSpec:
    name: str
    memory_gb: float
    bandwidth_gbps: float
    bf16_tflops: float
    hourly_cost_usd: float

    @property
    def ridge_point_flops_per_byte(self) -> float:
        """Minimum operational intensity (FLOPs/byte) required to reach peak TFLOPS."""
        return (self.bf16_tflops * 1e12) / (self.bandwidth_gbps * 1e9)


@dataclass(frozen=True)
class ModelWorkload:
    model_name: str
    active_params_b: float
    bytes_per_param: int
    kv_bytes_per_token: int
    prompt_tokens: int
    output_tokens: int
    batch_size: int
    activation_overhead_gb: float = 2.0

    def validate(self) -> None:
        if self.active_params_b <= 0 or self.bytes_per_param <= 0:
            raise ValueError(
                "Model parameter count and bytes per parameter must be positive"
            )
        if self.prompt_tokens <= 0 or self.output_tokens <= 0 or self.batch_size <= 0:
            raise ValueError(
                "Prompt tokens, output tokens, and batch size must be positive integers"
            )

    @property
    def weight_bytes(self) -> float:
        return self.active_params_b * 1e9 * self.bytes_per_param

    @property
    def total_kv_bytes(self) -> float:
        seq_len = self.prompt_tokens + self.output_tokens
        return float(self.batch_size * seq_len * self.kv_bytes_per_token)

    @property
    def required_memory_gb(self) -> float:
        return (
            self.weight_bytes + self.total_kv_bytes
        ) / 1e9 + self.activation_overhead_gb


@dataclass(frozen=True)
class WorkloadPhaseMetrics:
    phase: str
    flops: float
    hbm_bytes: float
    operational_intensity: float
    ridge_point: float
    regime: str
    roofline_tflops_bound: float
    latency_ms: float


GPU_SPECS: dict[str, GpuSpec] = {
    "example-small": GpuSpec(
        name="example-small",
        memory_gb=48.0,
        bandwidth_gbps=864.0,
        bf16_tflops=362.0,
        hourly_cost_usd=1.55,
    ),
    "example-fast": GpuSpec(
        name="example-fast",
        memory_gb=80.0,
        bandwidth_gbps=3350.0,
        bf16_tflops=989.0,
        hourly_cost_usd=3.20,
    ),
    "example-large-memory": GpuSpec(
        name="example-large-memory",
        memory_gb=141.0,
        bandwidth_gbps=4800.0,
        bf16_tflops=989.0,
        hourly_cost_usd=3.80,
    ),
    "example-wide-bandwidth": GpuSpec(
        name="example-wide-bandwidth",
        memory_gb=192.0,
        bandwidth_gbps=8000.0,
        bf16_tflops=2250.0,
        hourly_cost_usd=6.20,
    ),
}


def compute_prefill_metrics(
    gpu: GpuSpec, workload: ModelWorkload
) -> WorkloadPhaseMetrics:
    workload.validate()
    total_prefill_tokens = workload.batch_size * workload.prompt_tokens
    flops = 2.0 * (workload.active_params_b * 1e9) * total_prefill_tokens
    kv_write_bytes = float(total_prefill_tokens * workload.kv_bytes_per_token)
    hbm_bytes = workload.weight_bytes + kv_write_bytes
    intensity = flops / hbm_bytes
    ridge = gpu.ridge_point_flops_per_byte
    regime = "compute-bound" if intensity >= ridge else "bandwidth-bound"
    compute_seconds = flops / (gpu.bf16_tflops * 1e12)
    memory_seconds = hbm_bytes / (gpu.bandwidth_gbps * 1e9)
    latency_s = max(compute_seconds, memory_seconds)
    roofline_tflops_bound = (flops / latency_s) / 1e12
    return WorkloadPhaseMetrics(
        phase="prefill",
        flops=round(flops, 2),
        hbm_bytes=round(hbm_bytes, 2),
        operational_intensity=round(intensity, 2),
        ridge_point=round(ridge, 2),
        regime=regime,
        roofline_tflops_bound=round(roofline_tflops_bound, 2),
        latency_ms=round(latency_s * 1000.0, 3),
    )


def compute_decode_step_metrics(
    gpu: GpuSpec, workload: ModelWorkload
) -> WorkloadPhaseMetrics:
    workload.validate()
    flops = 2.0 * (workload.active_params_b * 1e9) * workload.batch_size
    avg_context_len = workload.prompt_tokens + (workload.output_tokens / 2.0)
    kv_read_bytes = float(
        workload.batch_size * avg_context_len * workload.kv_bytes_per_token
    )
    hbm_bytes = workload.weight_bytes + kv_read_bytes
    intensity = flops / hbm_bytes
    ridge = gpu.ridge_point_flops_per_byte
    regime = "compute-bound" if intensity >= ridge else "bandwidth-bound"
    compute_seconds = flops / (gpu.bf16_tflops * 1e12)
    memory_seconds = hbm_bytes / (gpu.bandwidth_gbps * 1e9)
    latency_s = max(compute_seconds, memory_seconds)
    roofline_tflops_bound = (flops / latency_s) / 1e12
    return WorkloadPhaseMetrics(
        phase="decode_step",
        flops=round(flops, 2),
        hbm_bytes=round(hbm_bytes, 2),
        operational_intensity=round(intensity, 2),
        ridge_point=round(ridge, 2),
        regime=regime,
        roofline_tflops_bound=round(roofline_tflops_bound, 2),
        latency_ms=round(latency_s * 1000.0, 3),
    )


def evaluate_gpu_for_workload(
    gpu: GpuSpec,
    workload: ModelWorkload,
    require_single_card_fit: bool = True,
) -> dict[str, object]:
    workload.validate()
    required_gb = workload.required_memory_gb
    fits = required_gb <= gpu.memory_gb
    if require_single_card_fit and not fits:
        raise ValueError(
            f"{gpu.name} has {gpu.memory_gb:.1f} GB VRAM, below required {required_gb:.2f} GB"
        )
    prefill = compute_prefill_metrics(gpu, workload)
    decode = compute_decode_step_metrics(gpu, workload)
    ttft_ms = prefill.latency_ms
    tpot_ms = decode.latency_ms
    total_latency_s = (ttft_ms + tpot_ms * (workload.output_tokens - 1)) / 1000.0
    decode_tokens_per_sec = workload.batch_size / (tpot_ms / 1000.0)
    end_to_end_tokens_per_sec = (
        workload.batch_size * workload.output_tokens
    ) / total_latency_s
    tokens_per_dollar = end_to_end_tokens_per_sec * 3600 / gpu.hourly_cost_usd
    return {
        "gpu": gpu.name,
        "memory_gb": gpu.memory_gb,
        "required_memory_gb": round(required_gb, 2),
        "fits_single_gpu": fits,
        "ridge_point_flops_per_byte": round(gpu.ridge_point_flops_per_byte, 2),
        "prefill": asdict(prefill),
        "decode": asdict(decode),
        "ttft_ms": round(ttft_ms, 2),
        "tpot_ms": round(tpot_ms, 3),
        "decode_tokens_per_sec": round(decode_tokens_per_sec, 2),
        "end_to_end_tokens_per_sec": round(end_to_end_tokens_per_sec, 2),
        "tokens_per_dollar": round(tokens_per_dollar, 2),
    }


def build_report() -> dict[str, object]:
    # 8B BF16 model (16 GB weights, 128 KB/token KV cache for 32 layers GQA)
    small_model_comparisons = []
    for name, gpu in GPU_SPECS.items():
        workload = ModelWorkload(
            model_name="Llama-3-8B-BF16",
            active_params_b=8.0,
            bytes_per_param=2,
            kv_bytes_per_token=131072,
            prompt_tokens=2048,
            output_tokens=256,
            batch_size=8,
        )
        small_model_comparisons.append(evaluate_gpu_for_workload(gpu, workload))

    # 70B BF16 model (140 GB weights? Or 32B BF16 = 64 GB weights + KV cache)
    # Let's evaluate a 32B BF16 model (64 GB weights) with long context that fits H200/example-wide-bandwidth
    # and overflows example-small/H100 when KV cache grows.
    large_workload = ModelWorkload(
        model_name="Qwen-32B-BF16-LongContext",
        active_params_b=32.0,
        bytes_per_param=2,
        kv_bytes_per_token=262144,
        prompt_tokens=4096,
        output_tokens=512,
        batch_size=16,
    )
    large_model_comparisons = [
        evaluate_gpu_for_workload(gpu, large_workload, require_single_card_fit=False)
        for gpu in GPU_SPECS.values()
    ]

    batch_sweep = []
    for batch_size in (1, 8, 32, 128, 256):
        w = ModelWorkload(
            model_name="Llama-3-8B-BF16",
            active_params_b=8.0,
            bytes_per_param=2,
            kv_bytes_per_token=131072,
            prompt_tokens=1024,
            output_tokens=128,
            batch_size=batch_size,
        )
        metrics = evaluate_gpu_for_workload(GPU_SPECS["example-fast"], w)
        batch_sweep.append(
            {
                "batch_size": batch_size,
                "prefill_intensity": metrics["prefill"]["operational_intensity"],
                "prefill_regime": metrics["prefill"]["regime"],
                "decode_intensity": metrics["decode"]["operational_intensity"],
                "decode_regime": metrics["decode"]["regime"],
                "decode_roofline_tflops_bound": metrics["decode"][
                    "roofline_tflops_bound"
                ],
                "tpot_ms": metrics["tpot_ms"],
                "decode_tokens_per_sec": metrics["decode_tokens_per_sec"],
            }
        )

    return {
        "gpu_catalog": [
            {
                **asdict(gpu),
                "ridge_point_flops_per_byte": round(gpu.ridge_point_flops_per_byte, 2),
            }
            for gpu in GPU_SPECS.values()
        ],
        "llama_8b_batch8": small_model_comparisons,
        "qwen_32b_long_context": large_model_comparisons,
        "h100_batch_sweep": batch_sweep,
    }


def main() -> None:
    print(
        json.dumps(
            {
                "execution_mode": "CPU algorithm with controlled inputs; any latency/capacity coefficients are assumptions",
                **build_report(),
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()

"""Unit and failure-boundary tests for the GPU roofline and selector lab."""

from __future__ import annotations

import unittest

from lab import (
    GPU_SPECS,
    GpuSpec,
    ModelWorkload,
    build_report,
    compute_decode_step_metrics,
    compute_prefill_metrics,
    evaluate_gpu_for_workload,
)


class GpuRooflineTest(unittest.TestCase):
    def test_ridge_points_reflect_compute_to_bandwidth_ratio(self) -> None:
        self.assertAlmostEqual(
            GPU_SPECS["example-small"].ridge_point_flops_per_byte, 418.98, places=2
        )
        self.assertAlmostEqual(
            GPU_SPECS["example-fast"].ridge_point_flops_per_byte, 295.22, places=2
        )
        self.assertAlmostEqual(
            GPU_SPECS["example-large-memory"].ridge_point_flops_per_byte,
            206.04,
            places=2,
        )
        self.assertAlmostEqual(
            GPU_SPECS["example-wide-bandwidth"].ridge_point_flops_per_byte,
            281.25,
            places=2,
        )

    def test_prefill_is_compute_bound_while_small_batch_decode_is_bandwidth_bound(
        self,
    ) -> None:
        workload = ModelWorkload(
            model_name="Llama-3-8B-BF16",
            active_params_b=8.0,
            bytes_per_param=2,
            kv_bytes_per_token=131072,
            prompt_tokens=2048,
            output_tokens=256,
            batch_size=8,
        )
        h100 = GPU_SPECS["example-fast"]
        prefill = compute_prefill_metrics(h100, workload)
        decode = compute_decode_step_metrics(h100, workload)
        self.assertEqual(prefill.regime, "compute-bound")
        self.assertGreater(
            prefill.operational_intensity, h100.ridge_point_flops_per_byte
        )
        self.assertEqual(decode.regime, "bandwidth-bound")
        self.assertLess(decode.operational_intensity, 20.0)

    def test_buying_peak_tflops_for_small_batch_decode_wastes_budget(self) -> None:
        """Doubling peak TFLOPS without raising HBM bandwidth leaves small-batch decode TPOT unchanged."""
        workload = ModelWorkload(
            model_name="Llama-3-8B-BF16",
            active_params_b=8.0,
            bytes_per_param=2,
            kv_bytes_per_token=131072,
            prompt_tokens=2048,
            output_tokens=256,
            batch_size=4,
        )
        h100 = GPU_SPECS["example-fast"]
        # Hypothetical chip with 2x BF16 TFLOPS but identical 3,350 GB/s HBM bandwidth
        double_compute_same_bw = GpuSpec(
            name="Hypothetical-2x-TFLOPS",
            memory_gb=80.0,
            bandwidth_gbps=h100.bandwidth_gbps,
            bf16_tflops=h100.bf16_tflops * 2.0,
            hourly_cost_usd=h100.hourly_cost_usd * 1.5,
        )
        h100_eval = evaluate_gpu_for_workload(h100, workload)
        hyp_eval = evaluate_gpu_for_workload(double_compute_same_bw, workload)
        self.assertEqual(h100_eval["tpot_ms"], hyp_eval["tpot_ms"])
        self.assertGreater(
            GPU_SPECS["example-large-memory"].bandwidth_gbps,
            h100.bandwidth_gbps,
        )
        h200_eval = evaluate_gpu_for_workload(
            GPU_SPECS["example-large-memory"], workload
        )
        self.assertLess(h200_eval["tpot_ms"], h100_eval["tpot_ms"])

    def test_insufficient_vram_rejects_single_card_fit(self) -> None:
        """A workload whose weights + KV cache + activations exceed GPU VRAM raises ValueError."""
        large_workload = ModelWorkload(
            model_name="Qwen-32B-BF16-LongContext",
            active_params_b=32.0,
            bytes_per_param=2,
            kv_bytes_per_token=262144,
            prompt_tokens=4096,
            output_tokens=512,
            batch_size=16,
        )
        with self.assertRaisesRegex(ValueError, "below required"):
            evaluate_gpu_for_workload(
                GPU_SPECS["example-small"], large_workload, require_single_card_fit=True
            )
        with self.assertRaisesRegex(ValueError, "below required"):
            evaluate_gpu_for_workload(
                GPU_SPECS["example-fast"], large_workload, require_single_card_fit=True
            )
        h200_fit = evaluate_gpu_for_workload(
            GPU_SPECS["example-large-memory"],
            large_workload,
            require_single_card_fit=True,
        )
        self.assertTrue(h200_fit["fits_single_gpu"])

    def test_report_structure_is_deterministic(self) -> None:
        report = build_report()
        self.assertEqual(len(report["gpu_catalog"]), 4)
        self.assertEqual(len(report["llama_8b_batch8"]), 4)
        self.assertEqual(len(report["h100_batch_sweep"]), 5)


class RooflineUnitsTests(unittest.TestCase):
    def test_first_output_token_has_no_decode_interval(self):
        from lab import GpuSpec, ModelWorkload, evaluate_gpu_for_workload

        gpu = GpuSpec("test", 100, 100, 10, 2)
        work = ModelWorkload("test", 0.01, 2, 16, 8, 1, 1)
        result = evaluate_gpu_for_workload(gpu, work)
        expected = 1000 / result["prefill"]["latency_ms"]
        self.assertAlmostEqual(result["end_to_end_tokens_per_sec"], expected, delta=0.2)
        self.assertAlmostEqual(
            result["tokens_per_dollar"],
            result["end_to_end_tokens_per_sec"] * 1800,
            delta=10,
        )


if __name__ == "__main__":
    unittest.main()


class GuardBoundaryTests(unittest.TestCase):
    def test_rejects_invalid_boundary_inputs_with_specific_reason(self):
        import lab
        cases = [('ModelWorkload("test",0,2,2,1,1,1).validate()', 'Model parameter'), ('ModelWorkload("test",1,2,2,0,1,1).validate()', 'Prompt tokens')]
        for code, reason in cases:
            with self.subTest(code=code):
                with self.assertRaisesRegex((ValueError, MemoryError, KeyError, RuntimeError), reason):
                    exec(code, vars(lab).copy())