Part B: Self-hosting

"""Memory arithmetic for self-hosting: weights plus KV cache, from a model's config.

TinyLM is computed exactly and checked against the real tensors. Real open
models are computed from their published config.json files (values copied
below, with the source named) so you can size a GPU before renting one.
"""
from __future__ import annotations

from dataclasses import dataclass

import torch

from supportdesk.tinylm import load

GIB = 1024 ** 3
BYTES = {"fp32": 4.0, "bf16": 2.0, "int8": 1.0, "int4": 0.5, "mxfp4": 4.25 / 8}  # MXFP4: 4 bits + 1 shared 8-bit scale per 32 values


@dataclass(frozen=True)
class Arch:
    """The config.json fields that decide memory. Values come from each model's published config."""

    name: str
    hidden: int
    layers: int
    heads: int
    kv_heads: int
    head_dim: int
    intermediate: int
    vocab: int
    tied: bool
    experts: int = 1            # mixture-of-experts: experts per layer (1 = dense)
    active_experts: int = 1     # experts each token actually uses
    full_attention_layers: int | None = None   # others use a sliding window
    window: int | None = None


# Sources: config.json on Hugging Face for meta-llama/Llama-3.1-8B, Qwen/Qwen3-8B, and
# openai/gpt-oss-120b (read 21 Sep 2026). gpt-oss-120b alternates sliding-window (128 tokens)
# and full-attention layers, and routes each token to 4 of its 128 experts.
LLAMA_31_8B = Arch("Llama 3.1 8B", 4096, 32, 32, 8, 128, 14336, 128256, False)
QWEN3_8B = Arch("Qwen3-8B", 4096, 36, 32, 8, 128, 12288, 151936, False)
GPT_OSS_120B = Arch("gpt-oss-120b", 2880, 36, 64, 8, 64, 2880, 201088, False, experts=128,
                    active_experts=4, full_attention_layers=18, window=128)


def params(a: Arch) -> dict[str, int]:
    """Parameter count by part, for a Llama-style block (gated MLP, grouped-query attention).

    Small terms (norm weights, attention biases, router, Qwen3's q/k norms) are
    left out; they change the total by far less than 0.1 percent.
    """
    embed = a.vocab * a.hidden * (1 if a.tied else 2)
    attn = a.hidden * a.heads * a.head_dim * 2 + a.hidden * a.kv_heads * a.head_dim * 2
    mlp = 3 * a.hidden * a.intermediate * a.experts
    active = embed // (1 if a.tied else 2) + (attn + 3 * a.hidden * a.intermediate * a.active_experts) * a.layers
    return {"embeddings": embed, "attention": attn * a.layers, "mlp": mlp * a.layers,
            "total": embed + (attn + mlp) * a.layers, "active": active}


def kv_bytes_per_token(a: Arch, dtype: str = "bf16") -> float:
    """K and V, for every layer, for every KV head: 2 * layers * kv_heads * head_dim * bytes."""
    layers = a.full_attention_layers if a.full_attention_layers is not None else a.layers
    return 2 * layers * a.kv_heads * a.head_dim * BYTES[dtype]


def kv_bytes_for_sequence(a: Arch, tokens: int, dtype: str = "bf16") -> float:
    full = kv_bytes_per_token(a, dtype) * tokens
    if a.window is not None:  # sliding-window layers keep at most `window` tokens
        sliding_layers = a.layers - (a.full_attention_layers or 0)
        full += 2 * sliding_layers * a.kv_heads * a.head_dim * BYTES[dtype] * min(tokens, a.window)
    return full


def tinylm_check() -> None:
    model, tokenizer = load()
    cfg = model.cfg
    d, v, c, n = cfg.d_model, cfg.vocab_size, cfg.context, cfg.n_layers
    per_layer = (2 * d) + (d * 3 * d + 3 * d) + (d * d + d) + (2 * d) + (d * 4 * d + 4 * d) + (4 * d * d + d)
    by_hand = v * d + c * d + n * per_layer + 2 * d        # head is tied to tok_emb, so it adds nothing
    weights_fp32 = sum(p.numel() * p.element_size() for p in model.parameters())
    print(f"TinyLM parameters: by hand {by_hand:,}  model.num_parameters() {model.num_parameters():,}")
    print(f"TinyLM weights in fp32: {weights_fp32:,} bytes = {weights_fp32 / 1024**2:.2f} MiB")
    used = tokenizer.get_vocab_size()
    print(f"embedding rows: {v:,} in the config, {used:,} used by the tokenizer, "
          f"{(v - used) * d * 4:,} bytes of rows no token can reach")

    kv_formula = 2 * n * cfg.n_heads * (d // cfg.n_heads) * 4   # TinyLM has no grouped-query attention
    ids = tokenizer.encode("Customer (Ben): I was charged twice for the Team plan.\nAgent (Dara):").ids
    caches = [dict() for _ in model.blocks]
    with torch.no_grad():
        model(torch.tensor([ids]), caches)
    measured = sum(cache["k"].nbytes + cache["v"].nbytes for cache in caches)
    print(f"TinyLM KV cache: formula {kv_formula:,} bytes/token; measured {measured:,} bytes "
          f"for {len(ids)} tokens = {measured // len(ids):,} bytes/token")
    print(f"TinyLM full context ({c} tokens): {kv_formula * c / 1024:.0f} KiB per sequence")


def real_models() -> None:
    print(f"\n{'model':<14}{'params (B)':>11}{'bf16 GiB':>10}{'int4 GiB':>10}{'KV KiB/tok':>12}{'KV GiB @8k':>12}{'KV GiB @128k':>14}")
    for a in (LLAMA_31_8B, QWEN3_8B):
        p = params(a)["total"]
        print(f"{a.name:<14}{p / 1e9:>11.2f}{p * 2 / GIB:>10.1f}{p * 0.5 / GIB:>10.1f}"
              f"{kv_bytes_per_token(a) / 1024:>12.0f}{kv_bytes_for_sequence(a, 8192) / GIB:>12.2f}"
              f"{kv_bytes_for_sequence(a, 131072) / GIB:>14.2f}")
    a = GPT_OSS_120B
    p = params(a)
    experts, rest = p["mlp"], p["total"] - p["mlp"]
    size = experts * BYTES["mxfp4"] + rest * BYTES["bf16"]
    print(f"{a.name:<14}{p['total'] / 1e9:>11.2f}{p['total'] * 2 / GIB:>10.1f}{'':>10}"
          f"{kv_bytes_per_token(a) / 1024:>12.0f}{kv_bytes_for_sequence(a, 8192) / GIB:>12.2f}"
          f"{kv_bytes_for_sequence(a, 131072) / GIB:>14.2f}")
    print(f"gpt-oss-120b as shipped (experts MXFP4, rest bf16): {size / GIB:.1f} GiB of weights")
    print(f"gpt-oss-120b parameters used per token (4 of 128 experts): {p['active'] / 1e9:.2f} B "
          f"(counting the {a.vocab * a.hidden / 1e9:.2f} B output head, not the input lookup table)")


def fits(a: Arch, gpu_gib: float, weight_dtype: str, context: int, overhead_gib: float = 2.0) -> int:
    """How many sequences of `context` tokens fit next to the weights (overhead: activations, runtime)."""
    weights = params(a)["total"] * BYTES[weight_dtype] / GIB
    free = gpu_gib - weights - overhead_gib
    return max(int(free // (kv_bytes_for_sequence(a, context) / GIB)), 0)


if __name__ == "__main__":
    tinylm_check()
    real_models()
    print("\nConcurrent sequences that fit (2 GiB runtime overhead assumed):")
    for gpu, gib in (("L4 24 GB", 22.5), ("A100 80 GB", 80.0)):
        for dtype in ("bf16", "int4"):
            print(f"  Qwen3-8B on {gpu:<10} weights {dtype}: "
                  f"{fits(QWEN3_8B, gib, dtype, 4096):>4} at 4k tokens, {fits(QWEN3_8B, gib, dtype, 32768):>3} at 32k")

"""Quantize TinyLM's linear layers to int8 and measure what it saves and what it costs.

Measures, on this machine: size on disk, forward-pass latency, validation
perplexity, agreement of the top next-token prediction, and greedy generations.
"""
from __future__ import annotations

import copy
import io
import math
import os
import statistics
import time
import warnings

import torch

from supportdesk.data import DATA_DIR
from supportdesk.tinylm import SamplingParams, generate, load, loss_on

warnings.filterwarnings("ignore")          # torch.ao.quantization prints deprecation notices
torch.manual_seed(0)
torch.set_num_threads(1)          # one thread: steadier timings on a shared 2-core machine


def validation_ids(tokenizer) -> torch.Tensor:
    """The same held-out tail (last 5 percent of tokens) that pretraining used for validation."""
    ids = torch.tensor(tokenizer.encode((DATA_DIR / "corpus.txt").read_text(encoding="utf-8")).ids)
    return ids[int(len(ids) * 0.95):]


@torch.no_grad()
def evaluate(model, ids: torch.Tensor, context: int = 128) -> tuple[float, list[int]]:
    """Perplexity over non-overlapping windows, plus the argmax prediction at every position."""
    total, count, predictions = 0.0, 0, []
    for start in range(0, len(ids) - 1, context):
        chunk = ids[start: start + context + 1]
        if len(chunk) < 2:
            break
        x, y = chunk[:-1].unsqueeze(0), chunk[1:].unsqueeze(0)
        total += loss_on(model, x, y).item() * y.numel()
        count += y.numel()
        predictions += model(x)[0].argmax(-1).tolist()
    return math.exp(total / count), predictions


def size_on_disk(model) -> int:
    buffer = io.BytesIO()
    torch.save(model.state_dict(), buffer)
    return buffer.getbuffer().nbytes


@torch.no_grad()
def forward_ms(models: dict, batch: torch.Tensor, repeats: int = 30) -> dict[str, float]:
    """Median forward time per model. Runs alternate between models so background load hits both."""
    times: dict[str, list[float]] = {name: [] for name in models}
    for i in range(repeats + 3):
        for name, model in models.items():
            started = time.perf_counter()
            model(batch)
            if i >= 3:                        # the first rounds are warm-up
                times[name].append((time.perf_counter() - started) * 1000)
    return {name: statistics.median(t) for name, t in times.items()}


def main() -> None:
    print(f"machine load (1-minute average) {os.getloadavg()[0]:.2f} on {os.cpu_count()} cores; "
          f"torch threads {torch.get_num_threads()}")
    fp32, tokenizer = load()
    int8 = torch.ao.quantization.quantize_dynamic(copy.deepcopy(fp32), {torch.nn.Linear}, dtype=torch.qint8)
    ids = validation_ids(tokenizer)
    batch = ids[:128].unsqueeze(0)

    models = {"fp32": fp32, "int8": int8}
    speed = forward_ms(models, batch)
    results = {}
    for name, model in models.items():
        ppl, preds = evaluate(model, ids)
        results[name] = preds
        print(f"{name}: size {size_on_disk(model) / 1024:5.0f} KiB   forward (1 x 128 tokens) "
              f"{speed[name]:6.2f} ms   validation perplexity {ppl:.4f}")
    same = sum(a == b for a, b in zip(results["fp32"], results["int8"]))
    print(f"top-1 next-token agreement: {same} of {len(results['fp32'])} positions "
          f"({same / len(results['fp32']):.2%})")

    prompts = ["Customer (Ben): How much does the Team plan cost?\nAgent (Dara):",
               "Customer (Mo): I forgot my password and the reset link expired.\nAgent (Lena):",
               "Customer (Ana): Can I export my tasks to CSV?\nAgent (Kofi):"]
    params = SamplingParams(max_new_tokens=30, temperature=0.0, stop=["\n"])
    for prompt in prompts:
        a = generate(fp32, tokenizer, prompt, params).text
        b = generate(int8, tokenizer, prompt, params).text
        print(f"greedy identical: {a == b!s:<5} | fp32: {a.strip()[:70]}")
        if a != b:
            print(f"                     | int8: {b.strip()[:70]}")


if __name__ == "__main__":
    main()

"""Throughput vs batch size on TinyLM: decode many requests together, measure the tradeoff.

Each request is a 64-token prompt from the held-out corpus tail followed by
32 greedy decode steps with the KV cache. A batch of B requests runs those
steps together, the way a serving engine batches concurrent users.
"""
from __future__ import annotations

import os
import statistics
import time

import torch

from supportdesk.data import DATA_DIR
from supportdesk.tinylm import load

torch.manual_seed(0)
torch.set_num_threads(1)          # one thread: steadier timings on a shared 2-core machine
PROMPT, DECODE = 64, 32


@torch.no_grad()
def run_batch(model, prompts: torch.Tensor) -> tuple[float, float]:
    """Return (prefill seconds, decode seconds) for one batch of prompts."""
    caches = [dict() for _ in model.blocks]
    started = time.perf_counter()
    logits = model(prompts, caches)[:, -1]
    prefill = time.perf_counter() - started
    started = time.perf_counter()
    for step in range(DECODE):
        next_ids = logits.argmax(-1, keepdim=True)
        logits = model(next_ids, caches, start_pos=PROMPT + step)[:, -1]
    return prefill, time.perf_counter() - started


def main() -> None:
    print(f"machine load (1-minute average) {os.getloadavg()[0]:.2f} on {os.cpu_count()} cores; "
          f"torch threads {torch.get_num_threads()}")
    model, tokenizer = load()
    ids = torch.tensor(tokenizer.encode((DATA_DIR / "corpus.txt").read_text(encoding="utf-8")).ids)
    tail = ids[int(len(ids) * 0.95):]
    pool = torch.stack([tail[i * PROMPT:(i + 1) * PROMPT] for i in range(32)])
    run_batch(model, pool[:4])  # warm-up

    print(f"{'batch':>5} {'decode tok/s':>13} {'ms per step':>12} {'request latency ms':>19} {'prefill ms':>11}")
    rows = []
    for batch in (1, 2, 4, 8, 16, 32):
        trials = [run_batch(model, pool[:batch]) for _ in range(5)]
        prefill = statistics.median(t[0] for t in trials)
        decode = statistics.median(t[1] for t in trials)
        tps = batch * DECODE / decode
        latency = (prefill + decode) * 1000
        rows.append((batch, tps, latency))
        print(f"{batch:>5} {tps:>13.0f} {decode / DECODE * 1000:>12.2f} {latency:>19.0f} {prefill * 1000:>11.1f}")
    base_tps, base_latency = rows[0][1], rows[0][2]
    print("\nrelative to batch 1:")
    for batch, tps, latency in rows:
        print(f"  batch {batch:>2}: throughput x{tps / base_tps:5.1f}, each request waits x{latency / base_latency:4.1f}")


if __name__ == "__main__":
    main()

"""Diagnose slow decoding on a shared machine: time TinyLM decode steps with 1 and with 2 threads."""
import os
import statistics
import time

import torch

from supportdesk.tinylm import load

model, _ = load()
x = torch.randint(0, 1503, (1, 64), generator=torch.Generator().manual_seed(0))
print(f"machine load (1-minute average) {os.getloadavg()[0]:.2f} on {os.cpu_count()} cores")
results: dict[int, list[float]] = {1: [], 2: []}
with torch.no_grad():
    for _ in range(5):                                   # alternate settings so background load hits both
        for threads in (2, 1):
            torch.set_num_threads(threads)
            model(x)                                     # warm-up after switching
            started = time.perf_counter()
            for _ in range(100):
                model(x[:, :1])
            results[threads].append((time.perf_counter() - started) / 100 * 1000)
for threads, runs in results.items():
    print(f"{threads} thread(s): median {statistics.median(runs):.2f} ms per decode step, "
          f"worst {max(runs):.2f} ms (5 rounds of 100 steps)")

"""Total cost of ownership: hosted API vs a rented GPU, for Brightlane's ticket volume.

Token counts per ticket are measured on the real dataset. API prices come from
supportdesk.pricing. GPU prices and every operating assumption are labeled
below; replace them with your own quotes and, above all, with throughput you
measured on your own serving engine.
"""
from __future__ import annotations

import statistics
from dataclasses import dataclass

from supportdesk.data import get_article, load_tickets
from supportdesk.kb_search import KBSearch
from supportdesk.llm import Usage
from supportdesk.pricing import PRICES, Price, cost_usd
from supportdesk.tokens import count_messages

HOURS_PER_MONTH = 730

# GPU on-demand prices, USD per GPU-hour, read on 21 Sep 2026 (verify before use):
#   RunPod secure cloud (runpod.io/pricing): L4 0.49, A100 SXM 80 GB 1.59, H100 SXM 3.49
#   Lambda (lambda.ai/pricing): H100 SXM 3.99, A100 SXM 80 GB 2.79
GPU_PRICES = {"L4 (RunPod)": 0.49, "A100 80GB (RunPod)": 1.59, "H100 (RunPod)": 3.49, "H100 (Lambda)": 3.99}

# Assumptions (ours, for teaching; change them):
TICKETS_PER_MONTH = 30_000
OUTPUT_TOKENS_PER_TICKET = 350        # triage JSON (about 60) + draft reply (about 290)
OPS_HOURS_PER_MONTH = 20              # patching, upgrades, on-call for the serving stack
OPS_RATE_USD = 100                    # loaded cost of an engineer hour
REPLICAS = 2                          # one GPU is a single point of failure

# supportdesk.pricing bills Groq's cached input at the full input price. Groq's prompt-caching
# page (console.groq.com/docs/prompt-caching, read 21 Sep 2026) lists a 50 percent discount on
# cached input tokens for gpt-oss models, which does not stack with the 50 percent batch discount.
# We correct it here instead of editing the shared price table.
GROQ_CACHED = {m: Price(PRICES[m].input, PRICES[m].output, PRICES[m].input * 0.5)
               for m in ("openai/gpt-oss-20b", "openai/gpt-oss-120b")}


def groq_cost(usage: Usage, model: str, batch: bool = False) -> float:
    """Cost with Groq's cached-input discount; in batch mode the cache discount does not stack."""
    if batch:
        return cost_usd(Usage(usage.input_tokens, usage.output_tokens), model, batch=True)
    p = GROQ_CACHED[model]
    uncached = usage.input_tokens - usage.cached_tokens
    return (uncached * p.input + usage.cached_tokens * p.cached_input + usage.output_tokens * p.output) / 1e6


def measured_input_tokens() -> list[int]:
    """Input tokens per ticket for triage + draft, built from the real tickets and KB articles."""
    search = KBSearch()
    system_triage = "Classify the Brightlane support ticket. Reply with JSON: category, priority, language, summary, needs_human."
    system_draft = "Draft a reply for a human agent to review. Use only the articles below and cite their ids."
    totals = []
    for t in load_tickets():
        triage = [{"role": "system", "content": system_triage}, {"role": "user", "content": t.text}]
        articles = "\n\n".join(f"[{h.article_id}]\n{get_article(h.article_id).body}" for h in search.search(t.text, k=3))
        draft = [{"role": "system", "content": system_draft + "\n\n" + articles}, {"role": "user", "content": t.text}]
        totals.append(count_messages(triage) + count_messages(draft))
    return totals


@dataclass
class SelfHost:
    gpu: str
    tokens_per_second: float   # sustained output tokens/s for the whole GPU at your latency target
    utilization: float         # share of the month the GPU does useful work

    def monthly_usd(self) -> float:
        return GPU_PRICES[self.gpu] * HOURS_PER_MONTH * REPLICAS + OPS_HOURS_PER_MONTH * OPS_RATE_USD

    def capacity_tickets(self) -> float:
        return self.tokens_per_second * 3600 * HOURS_PER_MONTH * self.utilization / OUTPUT_TOKENS_PER_TICKET * REPLICAS


def main() -> None:
    inputs = measured_input_tokens()
    mean_in = statistics.mean(inputs)
    print(f"measured input tokens per ticket (o200k estimate, n={len(inputs)}): "
          f"mean {mean_in:.0f}, median {statistics.median(inputs):.0f}, max {max(inputs)}")
    usage = Usage(input_tokens=round(mean_in), output_tokens=OUTPUT_TOKENS_PER_TICKET)

    print(f"\nHosted API at {TICKETS_PER_MONTH:,} tickets per month:")
    for model in ("openai/gpt-oss-20b", "openai/gpt-oss-120b", "gemini-3.5-flash-lite", "gemini-3.5-flash"):
        per_ticket = cost_usd(usage, model)
        print(f"  {model:<24} {per_ticket * 1000:6.3f} USD per 1k tickets   {per_ticket * TICKETS_PER_MONTH:9.2f} USD per month")

    cached = Usage(round(mean_in), OUTPUT_TOKENS_PER_TICKET, cached_tokens=round(mean_in * 0.9))
    print("\nIf 90% of input tokens hit the prompt cache (USD per 1k tickets):")
    for model in ("openai/gpt-oss-20b", "openai/gpt-oss-120b"):
        print(f"  {model:<24} pricing.py {cost_usd(cached, model) * 1000:.3f}   with Groq's 50% cached rate "
              f"{groq_cost(cached, model) * 1000:.3f}   batch (no stacking) {groq_cost(cached, model, batch=True) * 1000:.3f}")

    print(f"\nSelf-hosted, {REPLICAS} replicas, plus {OPS_HOURS_PER_MONTH} ops hours at {OPS_RATE_USD} USD:")
    for gpu in GPU_PRICES:
        fixed = SelfHost(gpu, 1, 1).monthly_usd()
        print(f"  {gpu:<20} {fixed:8.2f} USD per month, whether you send 0 tickets or a million")

    print("\nBreak-even: the volume where self-hosting the same open model costs what the API costs")
    for gpu, model in (("L4 (RunPod)", "openai/gpt-oss-20b"), ("H100 (RunPod)", "openai/gpt-oss-120b")):
        api = cost_usd(usage, model)
        volume = SelfHost(gpu, 1, 1).monthly_usd() / api
        tps = volume * OUTPUT_TOKENS_PER_TICKET / (HOURS_PER_MONTH * 3600) / REPLICAS
        print(f"  {model:<20} on {gpu:<14} {volume:>12,.0f} tickets per month, "
              f"which needs {tps:,.0f} output tok/s per GPU around the clock")

    print("\nWhat one self-hosted ticket costs, by sustained throughput and utilization (L4 pair):")
    print(f"  {'tok/s per GPU':>13} " + " ".join(f"{'util ' + format(u, '.0%'):>12}" for u in (0.1, 0.3, 0.6)))
    for tps in (100, 300, 1000):
        cells = []
        for u in (0.1, 0.3, 0.6):
            host = SelfHost("L4 (RunPod)", tps, u)
            cells.append(f"{host.monthly_usd() / host.capacity_tickets() * 1000:>9.3f}/1k")
        print(f"  {tps:>13} " + " ".join(f"{c:>12}" for c in cells))
    need = TICKETS_PER_MONTH * OUTPUT_TOKENS_PER_TICKET / (HOURS_PER_MONTH * 3600)
    print(f"\nBrightlane's average load is only {need:.1f} output tokens/s: the GPUs would sit mostly idle.")


if __name__ == "__main__":
    main()