#!/usr/bin/env python3
"""Show where llama.cpp v0.4.0 puts a GGUF model's bytes for a given --n-gpu-layers value.

Purpose: read a GGUF file's header (never the weights), apply llama.cpp's layer placement
    rule for one or more -ngl values, and print for each how many weight and KV cache MiB
    sit on the accelerator and on the host, what share of the bytes read per generated
    token is on the host, and, given how many times slower the host path is, the decode
    slowdown Part 3's split-memory arithmetic predicts. It turns an "offloaded N/M layers"
    log line into numbers you can check against the load log's buffer lines.
Platform: all (Python standard library only; no accelerator and no model loading)
Minimum memory: 8 GB (the script itself needs a few megabytes)
Assumes: Python 3.9 or later; a GGUF file, complete or just its first megabytes (the header
    and tensor list come first); for --labbook, a notebook holding a part-05/bandwidth-test
    line written by Part 5's bandwidth-test.py with an accelerator and a cpu figure.

Usage: python3 offload-split.py <model.gguf> [--ngl 37,27,18,0] [--ctx-size 4096]
           [--cache-type f16] [--slow-ratio R | --labbook labbook.md]

Placement, from load_tensors in src/llama-model.cpp at v0.4.0: the input layer
(token_embd) always stays on the CPU; with L repeating blocks and N = --n-gpu-layers,
block i goes to the accelerator when i >= L + 1 - N, and the output layer (output.weight
and output_norm) when N >= 1; a model with tied embeddings has no output.weight, and its
output layer loads a duplicate of token_embd instead. Each block's KV cache sits on the same
device as the block (unless --no-kv-offload). Bytes read per generated token: every block
and the output layer; of token_embd only one row, counted here as nothing. The host column
counts the mapped file once. Slowdown against a full offload:
(1 - f) + f x r, with f the host's share of those bytes and r how many times slower the
host reads them. Tensor bytes use the block sizes asserted in ggml/src/ggml-common.h.
"""

import argparse
import json
import re
import struct
import sys
from collections import defaultdict

MIB = 1024 ** 2
# ggml type id: (name, bytes per block, values per block), ggml-common.h at v0.4.0
BLOCK = {0: ("F32", 4, 1), 1: ("F16", 2, 1), 2: ("Q4_0", 18, 32), 3: ("Q4_1", 20, 32),
         6: ("Q5_0", 22, 32), 7: ("Q5_1", 24, 32), 8: ("Q8_0", 34, 32), 10: ("Q2_K", 84, 256),
         11: ("Q3_K", 110, 256), 12: ("Q4_K", 144, 256), 13: ("Q5_K", 176, 256),
         14: ("Q6_K", 210, 256), 20: ("IQ4_NL", 18, 32), 23: ("IQ4_XS", 136, 256),
         30: ("BF16", 2, 1), 39: ("MXFP4", 17, 32)}
CACHE_BYTES = {"f32": 4.0, "f16": 2.0, "bf16": 2.0, "q8_0": 34 / 32, "q4_0": 18 / 32}
SCALAR = {0: "<B", 1: "<b", 2: "<H", 3: "<h", 4: "<I", 5: "<i", 6: "<f", 7: "<?",
          10: "<Q", 11: "<q", 12: "<d"}


def read_gguf(path):
    """Return (metadata, {tensor name: bytes}) from the header of a GGUF file."""
    fh = open(path, "rb")

    def get(fmt):
        data = fh.read(struct.calcsize(fmt))
        if len(data) < struct.calcsize(fmt):
            sys.exit(f"{path}: the file ends inside the header")
        return struct.unpack(fmt, data)[0]

    def text():
        return fh.read(get("<Q")).decode("utf-8", "replace")

    def value(kind):  # 8 is a string, 9 an array, the rest scalars
        if kind == 9:
            inner, count = get("<I"), get("<Q")
            return [value(inner) for _ in range(count)]
        return text() if kind == 8 else get(SCALAR[kind])

    if fh.read(4) != b"GGUF":
        sys.exit(f"{path} is not a GGUF file")
    get("<I")  # format version
    n_tensors, n_kv = get("<Q"), get("<Q")
    meta = dict((text(), value(get("<I"))) for _ in range(n_kv))
    sizes = {}
    for _ in range(n_tensors):
        name, n_dims = text(), get("<I")
        values = 1
        for _ in range(n_dims):
            values *= get("<Q")
        kind, _offset = get("<I"), get("<Q")
        if kind not in BLOCK:
            sys.exit(f"{name}: ggml type {kind} is not in BLOCK; add its block size first")
        _, block_bytes, block_values = BLOCK[kind]
        sizes[name] = values // block_values * block_bytes
    return meta, sizes


def ratio_from_labbook(path):
    """Accelerator read rate over CPU read rate from the latest part-05/bandwidth-test line."""
    latest = None
    for line in open(path, encoding="utf-8"):
        if '"part-05/bandwidth-test"' in line:
            try:
                latest = json.loads(line)
            except json.JSONDecodeError:
                pass
    rates = (latest or {}).get("gbps", {})
    accel = next((k for k in rates if k != "cpu"), None)
    if accel is None or "cpu" not in rates:
        sys.exit(f"{path}: no part-05/bandwidth-test line with an accelerator and a cpu figure; "
                 "pass --slow-ratio instead")
    return rates[accel]["read"] / rates["cpu"]["read"], accel


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("model")
    ap.add_argument("--ngl", default="",
                    help="comma-separated -ngl values (default: all, 3/4, 1/2, 0)")
    ap.add_argument("--ctx-size", type=int, default=4096, help="context length for the KV cache")
    ap.add_argument("--cache-type", default="f16", choices=sorted(CACHE_BYTES))
    ap.add_argument("--slow-ratio", type=float, help="how many times slower the host path reads")
    ap.add_argument("--labbook", help="take the ratio from Part 5's bandwidth-test line")
    args = ap.parse_args()

    meta, sizes = read_gguf(args.model)
    arch = meta["general.architecture"]
    layers = meta[f"{arch}.block_count"]
    kv_heads = meta[f"{arch}.attention.head_count_kv"]
    kv_heads = max(kv_heads) if isinstance(kv_heads, list) else kv_heads
    head_k = meta.get(f"{arch}.attention.key_length",
                      meta[f"{arch}.embedding_length"] // meta[f"{arch}.attention.head_count"])
    head_v = meta.get(f"{arch}.attention.value_length", head_k)
    kv_layer = kv_heads * (head_k + head_v) * CACHE_BYTES[args.cache_type] * args.ctx_size

    block = defaultdict(int)
    output = embd = 0
    for name, n in sizes.items():
        m = re.match(r"blk\.(\d+)\.", name)
        if m:
            block[int(m.group(1))] += n
        elif name.startswith("token_embd"):
            embd += n
        else:
            output += n  # output.weight and output_norm.weight
    tied = "output.weight" not in sizes
    if tied:  # the output layer loads a duplicate of token_embd on its own device
        output += embd

    ratio, source = args.slow_ratio, "--slow-ratio"
    if ratio is None and args.labbook:
        ratio, accel = ratio_from_labbook(args.labbook)
        source = f"{accel} read / cpu read in {args.labbook}"
    ngls = [int(x) for x in args.ngl.split(",")] if args.ngl else \
        [layers + 1, (layers + 1) * 3 // 4, (layers + 1) // 2, 0]

    total_read = sum(block.values()) + output
    file_total = sum(sizes.values())
    print(f"{args.model}: {arch}, {layers} blocks + output layer = {layers + 1} offloadable")
    print(f"token_embd {embd / MIB:.2f} MiB (always host), output layer {output / MIB:.2f} MiB"
          f"{' (tied: a copy of token_embd)' if tied else ''}, "
          f"blocks {min(block.values()) / MIB:.2f} to {max(block.values()) / MIB:.2f} MiB each")
    print(f"KV cache {kv_layer * layers / MIB:.2f} MiB at {args.ctx_size} tokens, "
          f"{args.cache_type}")
    if ratio:
        print(f"host path {ratio:.2f} times slower ({source})")
    print(f"\n{'ngl':>4} {'log line':>9} {'GPU wts':>9} {'host wts':>9} {'GPU KV':>8} "
          f"{'host KV':>8} {'host read':>9} {'slowdown':>8}")
    for ngl in ngls:
        start = max(layers + 1 - ngl, 0)
        gpu_blocks = [i for i in block if i >= start]
        gpu_w = sum(block[i] for i in gpu_blocks) + (output if ngl >= 1 else 0)
        host_w = file_total - gpu_w + (embd if tied and ngl >= 1 else 0)
        f = 1 - gpu_w / total_read
        slow = f"{(1 - f) + f * ratio:8.2f}" if ratio else f"{'-':>8}"
        print(f"{ngl:>4} {f'{min(ngl, layers + 1)}/{layers + 1}':>9} {gpu_w / MIB:9.2f} "
              f"{host_w / MIB:9.2f} {kv_layer * len(gpu_blocks) / MIB:8.2f} "
              f"{kv_layer * (layers - len(gpu_blocks)) / MIB:8.2f} {100 * f:8.1f}% {slow}")


if __name__ == "__main__":
    main()