from inferscale.latency import AnalyticalLatencyModel
from inferscale.profiles import get_accelerator, get_model


def test_prefill_grows_with_tokens():
    lm = AnalyticalLatencyModel(get_model("Llama-3.1-8B"), get_accelerator("L4"))
    assert lm.prefill_seconds([1024]) > lm.prefill_seconds([128])


def test_quantization_reduces_weight_memory():
    fp16 = AnalyticalLatencyModel(get_model("Llama-3.1-8B"), get_accelerator("L4"), "fp16")
    int4 = AnalyticalLatencyModel(get_model("Llama-3.1-8B"), get_accelerator("L4"), "int4")
    assert int4.model_weight_gb < fp16.model_weight_gb