tests/test_gpu_offload_plan.py

"""Asking a graphics card for more than it has is not an error on Windows.

The driver backs the excess with system memory, the allocation succeeds, the
server reports itself healthy, and then every token drags weights across the
PCIe bus. On the benchmark laptop -- a 4 GB card with 2.3 GB free and a 4.5 GB
model -- switching offloading on without a ceiling took twenty workloads that
had been completing to twenty that timed out on their first generation. Nothing
in the logs said "too big". It just went quiet.

So the ceiling is computed before anything is asked for, and these are the
properties of that computation.
"""

from __future__ import annotations

import struct

from distinct_agent.gguf import ModelShape, read_shape
from distinct_agent.server_runner import (
    VRAM_BUDGET_SHARE,
    VRAM_FIXED_OVERHEAD,
    plan_offload,
)

GB = 1024**3
MB = 1024**2

#: OLMo-2 7B: 32 layers, 4096-wide, 32 heads with no grouped-query attention.
OLMO2_7B = ModelShape(layers=32, embedding_length=4096, heads=32, kv_heads=32)


def _plan(free_bytes, model_bytes=4_500_000_000, shape=OLMO2_7B, context=4096):
    return plan_offload(
        model_bytes=model_bytes,
        shape=shape,
        free_vram_bytes=free_bytes,
        context_tokens=context,
    )


# -- the ceiling -------------------------------------------------------------


def test_the_plan_never_asks_for_more_than_the_card_has_free() -> None:
    """The whole point. Every rung must fit in the budget it was given."""

    for free in (1 * GB, 2 * GB, 4 * GB, 8 * GB, 24 * GB):
        budget = int(free * VRAM_BUDGET_SHARE) - VRAM_FIXED_OVERHEAD
        per_layer = 4_500_000_000 / (OLMO2_7B.layers + 1) + OLMO2_7B.kv_bytes_per_layer(4096)
        for layers in _plan(free):
            assert layers * per_layer <= max(0, budget) + per_layer


def test_a_laptop_card_with_a_seven_billion_model_offloads_some_of_it() -> None:
    """2.3 GB free is not nothing, and it is nowhere near the whole model."""

    plan = _plan(2378 * MB)
    assert 0 < plan[0] < OLMO2_7B.layers


def test_a_card_that_fits_the_whole_model_takes_the_whole_model() -> None:
    assert _plan(24 * GB)[0] == OLMO2_7B.layers


def test_a_card_with_nothing_free_offloads_nothing() -> None:
    assert _plan(200 * MB) == (0,)


def test_no_card_at_all_offloads_nothing() -> None:
    assert _plan(0) == (0,)


def test_the_key_value_cache_is_counted_and_not_just_the_weights() -> None:
    """At 4096 tokens it can cost more per layer than the layer's weights."""

    assert OLMO2_7B.kv_bytes_per_layer(4096) > 0
    assert _plan(3 * GB, context=4096)[0] < _plan(3 * GB, context=512)[0]


def test_a_model_whose_header_cannot_be_read_offloads_nothing() -> None:
    """Unreadable is not an excuse to guess: it is the CPU, as it always was."""

    assert _plan(8 * GB, shape=ModelShape()) == (0,)


# -- the ladder under the ceiling -------------------------------------------


def test_the_ladder_descends_and_ends_at_the_cpu() -> None:
    plan = _plan(6 * GB)
    assert list(plan) == sorted(plan, reverse=True)
    assert plan[-1] == 0


def test_the_ladder_has_no_repeated_rungs() -> None:
    for free in (400 * MB, 800 * MB, 1 * GB, 2 * GB, 3 * GB, 6 * GB, 24 * GB):
        plan = _plan(free)
        assert len(plan) == len(set(plan))


# -- reading the header ------------------------------------------------------


def _gguf(tmp_path, pairs, architecture="olmo2"):
    """A minimal but valid GGUF header carrying ``pairs`` as uint32 values."""

    out = bytearray(b"GGUF")
    out += struct.pack("<I", 3)
    out += struct.pack("<Q", 0)
    out += struct.pack("<Q", len(pairs))
    for suffix, value in pairs.items():
        key = f"{architecture}{suffix}".encode()
        out += struct.pack("<Q", len(key)) + key
        out += struct.pack("<I", 4) + struct.pack("<I", value)
    path = tmp_path / "model.gguf"
    path.write_bytes(bytes(out))
    return path


def test_the_shape_is_read_out_of_the_file(tmp_path) -> None:
    path = _gguf(
        tmp_path,
        {
            ".block_count": 32,
            ".embedding_length": 4096,
            ".attention.head_count": 32,
            ".attention.head_count_kv": 8,
        },
    )
    shape = read_shape(path)
    assert (shape.layers, shape.embedding_length, shape.heads, shape.kv_heads) == (32, 4096, 32, 8)
    assert shape.head_dimension == 128
    assert shape.known


def test_multi_head_attention_omits_the_kv_head_count_and_that_is_fine(tmp_path) -> None:
    path = _gguf(
        tmp_path,
        {".block_count": 32, ".embedding_length": 4096, ".attention.head_count": 32},
    )
    assert read_shape(path).kv_heads == 32


def test_grouped_query_attention_needs_far_less_cache(tmp_path) -> None:
    """Eight key/value heads instead of thirty-two is four times less per layer."""

    many = ModelShape(layers=32, embedding_length=4096, heads=32, kv_heads=32)
    few = ModelShape(layers=32, embedding_length=4096, heads=32, kv_heads=8)
    assert few.kv_bytes_per_layer(4096) * 4 == many.kv_bytes_per_layer(4096)
    assert _plan(3 * GB, shape=few)[0] > _plan(3 * GB, shape=many)[0]


def test_a_file_that_is_not_a_gguf_reads_as_unknown_rather_than_raising(tmp_path) -> None:
    path = tmp_path / "not-a-model.gguf"
    path.write_bytes(b"this is not a model file at all")
    assert read_shape(path).known is False


def test_a_truncated_header_reads_as_unknown_rather_than_raising(tmp_path) -> None:
    path = tmp_path / "cut.gguf"
    path.write_bytes(b"GGUF" + struct.pack("<I", 3) + b"\x00\x00")
    assert read_shape(path).known is False


def test_a_missing_file_reads_as_unknown_rather_than_raising(tmp_path) -> None:
    assert read_shape(tmp_path / "absent.gguf").known is False