olokobayusuf/$README.md Secret

## $README.md

      
    Raw
  

              $README.md
            
          
    Gemma 4 26B Compiled Model Benchmarks

This gist contains scripts that can be used to reproduce our benchmarks for compiling Gemma 4 26B A4B. We measure the following:

Image vs compiled model distribution size.
Cold-start time to first token (CTTFT)
Time to first token (TTFT)
Tokens per second (TPS)

Setup Requirements

Running these benchmarks will require authentication on Muna and Modal.

Install uv as the benchmark scripts rely on it to run.
In an empty folder, copy .env.example below and rename it to .env.
Create an access key on Muna.
Install the Modal CLI.

Running Benchmarks

The benchmark scripts are all self-contained. Run them with uv run --env-file .env <file>

  
## .env.example
# Muna
MUNA_ACCESS_KEY=

# Huggingface
# This is only required by `cttft.py` which needs  to download
# weights from HF in order to benchmark upstream SGLang.
HF_TOKEN=

## cttft.py
#
#   Muna
#   Copyright © 2026 NatML Inc. All Rights Reserved.
#

# /// script
# requires-python = ">=3.11"
# dependencies = ["muna", "modal", "truss", "typer"]
# ///

"""
Cold-start Time To First Token (CTTFT) benchmark for Gemma 4 26B A4B.

Measures the wall-clock time from a request hitting a cold system to the first
generated token across two axes: the inference `--runtime` (muna | sglang) and
the `--provider` (modal | baseten | ssh). Weights are pre-staged out of band (a
Modal Volume for Modal, a preload command for SSH, the Baseten Delivery Network
for Baseten) so the measurement captures container/process boot + model load +
prefill + first token, not a one-off weight download from Hugging Face.

For Modal we run with `single_use_containers=True` so every boot lands on a
fresh container. For SSH each boot is a fresh process and the page cache is
dropped between boots (by default) for true cold-disk reads.

Reproduce with:

    uv run --env-file .env cttft.py --provider modal

NOTE: This test requires an `HF_TOKEN` in your environment, to download the
weights.
"""

from __future__ import annotations
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from enum import StrEnum
from json import dumps
from math import ceil, floor
from modal import enable_output, App, Image, Secret, Volume
from muna import Muna
from muna.beta.openai import Message
from os import environ
from pathlib import Path
from requests import post
from rich.console import Console
from rich.table import Table
from statistics import mean
from string import Template
from subprocess import run
from sys import version_info
from tempfile import TemporaryDirectory
from time import monotonic
from typer import run as typer_run, Exit, Option
from typing import Annotated

console = Console()

DEFAULT_TAG = "@google/gemma-4-26b-a4b-it"

class Runtime(StrEnum):
    """
    Inference runtime under test.
    """
    muna = "muna"
    sglang = "sglang"

class Provider(StrEnum):
    """
    Compute provider.
    """
    modal = "modal"
    baseten = "baseten"
    ssh = "ssh"

def main(
    runtime: Annotated[list[Runtime] | None, Option(
        "--runtime", "-r",
        help="Inference runtime to benchmark; repeatable. Defaults to ALL runtimes."
    )] = None,
    provider: Annotated[list[Provider] | None, Option(
        "--provider", "-p",
        help="Compute provider to benchmark; repeatable. Defaults to ALL providers."
    )] = None,
    boots: Annotated[int, Option(
        "--boots", "-n",
        min=1,
        help="Number of cold boots per (runtime, provider)."
    )] = 5,
    prompt: Annotated[str, Option(
        "--prompt",
        help="Fixed minimal prompt sent on every cold boot."
    )] = "What is the capital of France?",
    max_tokens: Annotated[int, Option(
        "--max-tokens",
        min=1,
        help="Tokens to request; CTTFT only needs the first token."
    )] = 1,
    tag: Annotated[str, Option(
        "--tag",
        help="Muna predictor tag (for muna runtimes)."
    )] = DEFAULT_TAG,
    ssh_host: Annotated[str | None, Option(
        "--ssh-host",
        help="SSH target for the ssh provider (alias or user@host)."
    )] = None,
    ssh_command: Annotated[str | None, Option(
        "--ssh-command",
        help="Remote command timed per boot. Supports $prompt/$max_tokens/$tag placeholders."
    )] = None,
    ssh_preload_command: Annotated[str | None, Option(
        "--ssh-preload-command",
        help="Optional remote command run once (untimed) before boots to stage weights/cache."
    )] = None,
    ssh_port: Annotated[int | None, Option(
        "--ssh-port",
        help="SSH port (defaults to the system ssh default)."
    )] = None,
    ssh_identity: Annotated[Path | None, Option(
        "--ssh-identity",
        help="SSH identity (private key) file."
    )] = None,
    drop_caches: Annotated[bool, Option(
        "--drop-caches/--no-drop-caches",
        help="Drop the page cache before each ssh boot for true cold-disk reads (needs passwordless sudo)."
    )] = True,
    output_dir: Annotated[Path, Option(
        "--output-dir",
        help="Directory for results/<metric>__<target>.json output."
    )] = Path("results"),
    json_output: Annotated[bool, Option(
        "--json/--no-json",
        help="Also echo raw result JSON to stdout."
    )] = False,
    quiet: Annotated[bool, Option(
        "--quiet", "-q",
        help="Suppress non-essential output."
    )] = False,
) -> None:
    """
    Measure cold-start time to first token across `--boots` cold boots.

    Benchmarks the Cartesian product of selected runtimes and providers, writing
    one results file per (runtime, provider) so distribution statistics
    (p50/p90) can be computed later.
    """
    runtimes = runtime or list(Runtime)
    providers = provider or list(Provider)
    if Provider.ssh in providers:
        if not ssh_host or not ssh_command:
            console.print("[red]error:[/] --provider ssh requires --ssh-host and --ssh-command.")
            raise Exit(code=1)
        if len(runtimes) != 1:
            console.print(
                "[red]error:[/] --provider ssh uses a single model-agnostic --ssh-command; "
                "select exactly one --runtime (it is only a results label)."
            )
            raise Exit(code=1)
    output_dir.mkdir(parents=True, exist_ok=True)
    records: dict[str, dict] = { }
    for prov in providers:
        for rt in runtimes:
            session = _make_session(
                rt, prov,
                tag=tag,
                prompt=prompt,
                max_tokens=max_tokens,
                ssh_host=ssh_host,
                ssh_command=ssh_command,
                ssh_preload_command=ssh_preload_command,
                ssh_port=ssh_port,
                ssh_identity=ssh_identity,
                drop_caches=drop_caches,
            )
            label = f"{rt.value}-on-{prov.value}"
            if session is None:
                console.print(f"[yellow]skip[/] {label}: not implemented")
                continue
            samples: list[float] = []
            with session as boot_once:
                for index in range(boots):
                    elapsed = boot_once()
                    samples.append(elapsed)
                    print(f"{label} boot {index + 1}/{boots}: {elapsed:.2f}s")
            record = {
                "metric": "cttft",
                "target": label,
                "runtime": rt.value,
                "provider": prov.value,
                "tag": tag,
                "config": {
                    "boots": boots,
                    "prompt": prompt,
                    "max_tokens": max_tokens,
                    "batch_size": 1,
                },
                "summary": _summarize(samples),
                "samples": samples,
            }
            output = output_dir / f"cttft__{label}.json"
            output.write_text(dumps(record, indent=2))
            records[label] = record
    if not quiet and records:
        _render(records, output_dir)
    if json_output:
        console.print_json(dumps(records))
    raise Exit()

def _make_session(
    runtime: Runtime,
    provider: Provider,
    *,
    tag: str,
    prompt: str,
    max_tokens: int,
    ssh_host: str | None,
    ssh_command: str | None,
    ssh_preload_command: str | None,
    ssh_port: int | None,
    ssh_identity: Path | None,
    drop_caches: bool,
):
    """
    Return a session context manager for `(runtime, provider)`, or `None` if the
    combination is not implemented (skipped with a warning by the caller).
    """
    match (runtime, provider):
        case (Runtime.muna, Provider.modal):
            return _muna_modal_session(tag=tag, prompt=prompt, max_tokens=max_tokens)
        case (Runtime.sglang, Provider.modal):
            return _sglang_modal_session(prompt=prompt, max_tokens=max_tokens)
        case (_, Provider.ssh):
            return _ssh_session(
                host=ssh_host,
                command=ssh_command,
                preload_command=ssh_preload_command,
                port=ssh_port,
                identity=ssh_identity,
                drop_caches=drop_caches,
                prompt=prompt,
                max_tokens=max_tokens,
                tag=tag,
            )
        case _:
            return None # baseten remains a documented stub (see _profile_*_on_baseten)

@contextmanager
def _muna_modal_session(
    *,
    tag: str,
    prompt: str,
    max_tokens: int,
) -> Iterator[Callable[[], float]]:
    """
    Session for a compiled Muna model on Modal.

    Preloads the model weights onto a Modal volume (untimed) so cold boots read
    weights from the volume instead of downloading them from Hugging Face, then
    yields a `boot_once` that times one fresh-container cold boot.
    `single_use_containers=True` retires the container after every input, so each
    boot is genuinely cold.
    """
    # Pre-load the model weights on a Modal volume
    preload_app = App("muna-benchmark-preload")
    volume = Volume.from_name("muna-benchmark", create_if_missing=True, version=2)
    image = (Image
        .from_registry(f"python:{version_info.major}.{version_info.minor}-slim-bookworm") # use smallest Python image possible
        .uv_pip_install("muna")
    )
    secret = Secret.from_dict({ "MUNA_ACCESS_KEY": environ["MUNA_ACCESS_KEY"] })
    def _preload():
        muna = Muna()
        muna.predictions.create(tag, inputs={ }) # Muna preloads but doesn't execute compiled model
    preload_fn = preload_app.function(
        image=image,
        serialized=True,
        volumes={ "/muna": volume },
        env={ "MUNA_HOME": "/muna" },
        secrets=[secret],
        timeout=900,
    )(_preload)
    with enable_output(), preload_app.run():
        preload_fn.remote()
        print("Preloaded model weights on Modal volume")
    # Benchmark on a B200, one fresh container per cold boot
    benchmark_app = App("muna-benchmark")
    def _benchmark(tag: str, prompt: str, max_tokens: int):
        muna = Muna()
        message = Message(role="user", content=prompt)
        stream = muna.predictions.stream(
            tag,
            inputs={ "messages": [message], "max_output_tokens": max_tokens },
            acceleration="local_gpu"
        )
        next(stream) # block until the first token is produced
    benchmark_fn = benchmark_app.function(
        image=image,
        serialized=True,
        cpu=32,
        gpu="b200",
        volumes={ "/muna": volume },
        env={ "MUNA_HOME": "/muna" },
        secrets=[secret],
        timeout=300,
        single_use_containers=True,
    )(_benchmark)
    with enable_output(), benchmark_app.run():
        def boot_once() -> float:
            start = monotonic()
            benchmark_fn.remote(tag, prompt, max_tokens)
            return monotonic() - start
        yield boot_once

@contextmanager
def _sglang_modal_session(
    *,
    prompt: str,
    max_tokens: int,
) -> Iterator[Callable[[], float]]:
    """
    Session for SGLang on Modal.

    Mirrors `_muna_modal_session`: preloads the gated weights onto a Modal volume
    (untimed), then yields a `boot_once` that times one cold boot of a fresh
    single-use container.
    """
    repo = "google/gemma-4-26B-A4B-it"
    model_dir = "/sglang/gemma-4-26b-a4b-it"
    # Preload the model weights on a Modal volume
    preload_app = App("sglang-benchmark-preload")
    volume = Volume.from_name("sglang-benchmark", create_if_missing=True, version=2)
    preload_image = (Image
        .from_registry(f"python:{version_info.major}.{version_info.minor}-slim-bookworm")
        .uv_pip_install("huggingface_hub[hf_transfer]")
        .env({ "HF_HUB_ENABLE_HF_TRANSFER": "1" })
    )
    secret = Secret.from_dict({ "HF_TOKEN": environ["HF_TOKEN"] })
    def _preload():
        from huggingface_hub import snapshot_download
        snapshot_download(repo, local_dir=model_dir) # download the gated weights once
    preload_fn = preload_app.function(
        image=preload_image,
        serialized=True,
        volumes={ "/sglang": volume },
        secrets=[secret],
        timeout=900,
    )(_preload)
    with enable_output(), preload_app.run():
        preload_fn.remote()
        print("Preloaded model weights on Modal volume")
    # Benchmark on a B200, one fresh container per cold boot
    benchmark_app = App("sglang-benchmark")
    image = Image.from_registry(
        "lmsysorg/sglang:v0.5.13",
        add_python=f"{version_info.major}.{version_info.minor}"
    )
    def _benchmark(model_dir: str, prompt: str, max_tokens: int):
        import sglang as sgl
        engine = sgl.Engine(model_path=model_dir)
        engine.generate(prompt, sampling_params={ "max_new_tokens": max_tokens })
    benchmark_fn = benchmark_app.function(
        image=image,
        serialized=True,
        cpu=32,
        gpu="b200",
        volumes={ "/sglang": volume },
        timeout=300,
        single_use_containers=True,
    )(_benchmark)
    with enable_output(), benchmark_app.run():
        def boot_once() -> float:
            start = monotonic()
            benchmark_fn.remote(model_dir, prompt, max_tokens)
            return monotonic() - start
        yield boot_once

@contextmanager
def _ssh_session(
    *,
    host: str,
    command: str,
    preload_command: str | None,
    port: int | None,
    identity: Path | None,
    drop_caches: bool,
    prompt: str,
    max_tokens: int,
    tag: str,
) -> Iterator[Callable[[], float]]:
    """
    Session for a runtime on a baremetal box reached over SSH.

    Shells out to the system `ssh` binary (no extra deps). On enter it runs the
    optional preload once (untimed, mirroring Modal staging weights to a volume).
    Each boot is a fresh remote process; the page cache is dropped (untimed,
    unless disabled) beforehand so the timed command reads weights from disk.
    Commands are templated with `$prompt`/`$max_tokens`/`$tag`.
    """
    base = ["ssh", "-o", "BatchMode=yes"]
    if port is not None:
        base += ["-p", str(port)]
    if identity is not None:
        base += ["-i", str(identity)]
    def fill(template: str) -> str:
        return Template(template).safe_substitute(prompt=prompt, max_tokens=max_tokens, tag=tag)
    if preload_command:
        run([*base, host, fill(preload_command)], check=True) # untimed: populate weight/predictor cache once
    def boot_once() -> float:
        if drop_caches:
            # untimed: evict weights from the page cache -> true cold-disk read
            run([*base, host, "sync && echo 1 | sudo tee /proc/sys/vm/drop_caches > /dev/null"], check=True)
        start = monotonic()
        run([*base, host, fill(command)], check=True) # process cold-start: load + first token + exit
        return monotonic() - start
    yield boot_once

def _profile_muna_on_baseten():
    pass

def _profile_sglang_on_baseten():
    """
    Profile the Cold-start Time to First Token (CTTFT) of SGLang on Baseten (B10).

    Deploys a custom SGLang Docker server whose weights are delivered by the
    Baseten Delivery Network (BDN), so cold boots read weights from Baseten's
    in-cluster cache instead of downloading from Hugging Face. The Truss config
    is generated procedurally and pushed via the programmatic API (no CLI).

    Stub: not yet wired into the session/loop model (Baseten access pending).
    """
    from truss import login, push
    from truss.base.truss_config import (
        BaseImage, DockerServer, Resources,
        TrussConfig, Weights, WeightsSource
    )
    login(environ["BASETEN_API_KEY"])
    config = TrussConfig(
        model_name="gemma4-sglang-cttft",
        resources=Resources(accelerator="B200"),
        base_image=BaseImage(image="lmsysorg/sglang:v0.5.13"),
        docker_server=DockerServer(
            start_command=(
                "python3 -m sglang.launch_server --model-path /models/gemma "
                "--served-model-name gemma-4 --host 0.0.0.0 --port 8000"
            ),
            server_port=8000,
            predict_endpoint="/v1/chat/completions",
            readiness_endpoint="/health",
            liveness_endpoint="/health",
        ),
        weights=Weights([
            WeightsSource(
                source="hf://google/gemma-4-26B-A4B-it@main", # TODO: pin to a commit SHA
                mount_location="/models/gemma",
                #auth_secret_name="hf_access_token", # Baseten secret holding the HF token
                allow_patterns=["*.safetensors", "*.json", "tokenizer*"],
            )
        ]),
    )
    # Push the procedurally-generated Truss (no model.py needed for a docker_server)
    with TemporaryDirectory() as directory:
        config.write_to_yaml_file(Path(directory) / "config.yaml")
        deployment = push(
            directory,
            remote="baseten",
            publish=True,
            model_name=config.model_name
        )
    deployment.wait_for_active(timeout_seconds=30 * 60) # blocks through BDN mirroring + boot
    predict_url = deployment._baseten_service.predict_url
    print(f"Deployed SGLang on Baseten: {predict_url}")
    # Fire one request to confirm the endpoint serves
    response = post(
        predict_url,
        headers={ "Authorization": f"Api-Key {environ['BASETEN_API_KEY']}" },
        json={ "messages": [{ "role": "user", "content": "What is the capital of France?" }], "max_tokens": 1 },
    )
    print(response.json())

def _percentile(ordered: list[float], pct: float) -> float:
    if len(ordered) == 1:
        return ordered[0]
    rank = (len(ordered) - 1) * pct / 100
    low, high = floor(rank), ceil(rank)
    if low == high:
        return ordered[int(rank)]
    return ordered[low] * (high - rank) + ordered[high] * (rank - low)

def _summarize(samples: list[float]) -> dict[str, float | int]:
    if not samples:
        return { "n": 0 }
    ordered = sorted(samples)
    return {
        "n": len(ordered),
        "min": ordered[0],
        "p50": _percentile(ordered, 50),
        "p90": _percentile(ordered, 90),
        "max": ordered[-1],
        "mean": mean(ordered),
    }

def _render(records: dict[str, dict], output_dir: Path) -> None:
    table = Table(show_header=True, header_style="bold cyan", box=None)
    table.add_column("Target", style="bold green")
    table.add_column("n", style="white", justify="right")
    table.add_column("p50 (s)", style="white", justify="right")
    table.add_column("p90 (s)", style="white", justify="right")
    table.add_column("mean (s)", style="white", justify="right")
    for name, record in records.items():
        summary = record["summary"]
        if summary.get("n"):
            table.add_row(
                name,
                str(summary["n"]),
                f"{summary['p50']:.2f}",
                f"{summary['p90']:.2f}",
                f"{summary['mean']:.2f}",
            )
        else:
            table.add_row(name, "0", "[red]no samples[/]", "", "")
    console.print(table)
    console.print(f"[cyan]CTTFT[/] -> wrote results to [bold]{output_dir}[/].")

if __name__ == "__main__":
    typer_run(main)

## size.py
#
#   Muna
#   Copyright © 2026 NatML Inc. All Rights Reserved.
#

"""
Artifact size registries for Gemma 4 26B A4B.

Measures the effective size of model distributions across a container
image and a compiled binary generated by Muna.

Requirements:
1. Create a Muna access key at http://muna.ai/settings/developer.
2. Add your Muna access key to the `.env` file.
3. Login to the Modal CLI (for getting the size of a container image).

Reproduce with:

    uv run --env-file .env size.py --compiled-model @google/gemma-4-26b-a4b-it
    uv run --env-file .env size.py --image lmsysorg/sglang:latest
"""

# /// script
# requires-python = ">=3.11"
# dependencies = ["modal", "muna", "pyyaml", "typer"]
# ///

from __future__ import annotations
from enum import StrEnum
from json import dumps
from modal import enable_output, App, Image
from pathlib import Path
from requests import head
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from typer import run, Exit, Option
from typing import Annotated
from urllib.parse import urlparse
from yaml import safe_dump

console = Console()

class OutputFormat(StrEnum):
    yaml = "yaml"
    json = "json"

def main(
    image: Annotated[str | None, Option(
        "--image",
        help="SGLang Docker image reference to inspect."
    )] = None,
    compiled_model: Annotated[str | None, Option(
        help="Compiled Muna predictor tag to inspect."
    )] = None,
    output: Annotated[Path | None, Option(
        "--output", "-o",
        help="Output registry path. Defaults per mode."
    )] = None,
    output_format: Annotated[OutputFormat, Option(
        "--format",
        help="Registry output format."
    )] = OutputFormat.yaml,
    min_size: Annotated[int, Option(
        "--min-size", min=0,
        help="Omit files smaller than this many bytes (image mode only)."
    )] = 0,
) -> None:
    """
    Build a footprint size registry for either a Muna predictor or a Docker image.

    Pass exactly one of --compiled-model or --image.
    """
    if (image is None) == (compiled_model is None):
        console.print(
            "[red]Pass exactly one of[/] [bold]--compiled-model[/] [red]or[/] [bold]--image[/]."
        )
        raise Exit(code=1)
    if compiled_model is not None:
        _run_predictor(
            compiled_model,
            output or Path("registries/muna.yaml"),
            output_format
        )
    else:
        _run_image(
            image,
            output or Path("registries/container.yaml"),
            output_format,
            min_size
        )

def _run_predictor(
    tag: str,
    output: Path,
    output_format: OutputFormat
) -> None:
    """
    Emit the Muna predictor size registry (file name and size only).
    """
    from muna import Muna
    muna = Muna()
    prediction = muna.predictions.create(tag, client_id="x86_64-unknown-linux-gnu")
    resources = prediction.resources or []
    dso_resources = [res for res in resources if res.type == "dso"]
    if not dso_resources:
        console.print(f"[red]No DSO resources found for[/] [bold]{tag}[/].")
        raise Exit(code=1)
    files: list[dict[str, object]] = []
    for resource in dso_resources:
        name, size = _head_file(resource)
        files.append({ "name": name, "size": size })
    files.sort(key=lambda entry: int(entry["size"]), reverse=True)
    total = sum(int(entry["size"]) for entry in files)
    _write_registry(output, output_format, { "tag": tag, "total_size": total, "files": files })
    _render_registry(f"Muna size registry: {tag}", files, total, output)

def _run_image(
    image: str,
    output: Path,
    output_format: OutputFormat,
    min_size: int
) -> None:
    """
    Emit the SGLang image file size registry (name, size; no descriptions).

    Spins up an ephemeral Modal app whose container is built from the given
    image, walks its filesystem, and records every regular file with its size.
    """
    modal_app = App("gemma4-sglang-image-size")
    base = Image.from_registry(image)
    # `_walk_filesystem` lives at module scope; we apply the Modal decorator at
    # runtime so the container image can be driven by --image.
    walk_filesystem = modal_app.function(
        image=base,
        timeout=900,
    )(_walk_filesystem)
    with enable_output(), modal_app.run():
        result = walk_filesystem.remote(min_size)
    files = result["files"]
    total = int(result["total_size"])
    file_count = int(result["file_count"])
    _write_registry(output, output_format, {
        "image": image,
        "total_size": total,
        "file_count": file_count,
        "files": files,
    })
    _render_image_summary(image, files, total, file_count, output)

def _walk_filesystem(min_bytes: int) -> dict[str, object]:
    """
    Walk the container root filesystem and tally regular files by size.
    Runs inside the Modal container, so it relies only on the standard library.
    """
    # Skip pseudo-filesystems that do not contribute to image footprint.
    skip = { Path("/proc"), Path("/sys"), Path("/dev"), Path("/run") }
    files: list[dict[str, object]] = []
    total = 0
    count = 0
    stack = [Path("/")]
    while stack:
        try:
            entries = list(stack.pop().iterdir())
        except OSError:
            continue
        for entry in entries:
            # Skip symlinks so symlinked dirs are not descended (no loops) and
            # symlinks never count as regular files.
            if entry in skip or entry.is_symlink():
                continue
            try:
                if entry.is_dir():
                    stack.append(entry)
                elif entry.is_file():
                    size = entry.stat().st_size
                    total += size
                    count += 1
                    if size >= min_bytes:
                        files.append({ "name": str(entry), "size": size })
            except OSError:
                continue
    files.sort(key=lambda entry: int(entry["size"]), reverse=True)
    return { "total_size": total, "file_count": count, "files": files }

def _head_file(resource) -> tuple[str, int]:
    """
    Get the name and size of a Muna prediction resource with a HEAD request.
    """
    response = head(
        resource.url,
        headers={ "User-Agent": "muna-size" },
        allow_redirects=True
    )
    response.raise_for_status()
    length = response.headers.get("Content-Length")
    if length is None:
        raise RuntimeError(f"HEAD response for {resource.url} had no Content-Length")
    name = Path(urlparse(resource.url).path).name
    return name, int(length)

def _write_registry(
    output: Path,
    output_format: OutputFormat,
    data: dict[str, object]
) -> None:
    output.parent.mkdir(parents=True, exist_ok=True)
    text = (
        dumps(data, indent=2)
        if output_format is OutputFormat.json
        else safe_dump(data, sort_keys=False)
    )
    output.write_text(text, encoding="utf-8")

def _render_registry(
    title: str,
    files: list[dict[str, object]],
    total: int,
    output: Path
) -> None:
    table = Table(show_header=True, header_style="bold cyan", box=None)
    table.add_column("File", style="bold green")
    table.add_column("Size", style="white", justify="right")
    for entry in files:
        table.add_row(str(entry["name"]), _human_size(int(entry["size"])))
    table.add_row("", "")
    table.add_row("[bold magenta]Total[/]", f"[bold magenta]{_human_size(total)}[/]")
    console.print(Panel(table, title=f"[bold green]{title}[/]", border_style="green"))
    console.print(f"Wrote registry to [bold]{output}[/].")

def _render_image_summary(
    image: str,
    files: list[dict[str, object]],
    total: int,
    file_count: int,
    output: Path,
    top: int = 20,
) -> None:
    table = Table(show_header=True, header_style="bold cyan", box=None)
    table.add_column("File", style="bold green", overflow="fold")
    table.add_column("Size", style="white", justify="right")
    for entry in files[:top]:
        table.add_row(str(entry["name"]), _human_size(int(entry["size"])))
    if len(files) > top:
        table.add_row(f"[dim]... and {len(files) - top} more files[/]", "")
    table.add_row("", "")
    table.add_row("[bold magenta]Total[/]", f"[bold magenta]{_human_size(total)}[/]")
    console.print(Panel(table, title=f"[bold green]SGLang image: {image}[/]", border_style="green"))
    console.print(f"Recorded {file_count} files. Wrote registry to [bold]{output}[/].")

def _human_size(num_bytes: int) -> str:
    """
    Human readable file size.
    """
    value = float(num_bytes)
    for unit in ("B", "KiB", "MiB", "GiB"):
        if value < 1024.0 or unit == "GiB":
            return f"{value:.1f} {unit}" if unit != "B" else f"{int(value)} B"
        value /= 1024.0
    raise AssertionError("unreachable")

if __name__ == "__main__":
    run(main)

## tps.py
#
#   Muna
#   Copyright © 2026 NatML Inc. All Rights Reserved.
#

"""
Decode throughput (TPS) by prompt category benchmark for Gemma 4 26B A4B.

Measures sustained decode tokens-per-second at batch size 1. For each category
we stream a generation, time from the first content token to the last, and read
the exact completion-token count from the final usage chunk, so:

    decode_tps = (completion_tokens - 1) / (t_last_token - t_first_token)

which excludes prefill (that is the TTFT benchmark's concern).

Reproduce with:

    uv run --env-file .env tps.py
"""

# /// script
# requires-python = ">=3.11"
# dependencies = ["muna", "modal", "typer"]
# ///

from __future__ import annotations
from enum import StrEnum
from json import dumps
from modal import enable_output, App, Image, Secret, Volume
from muna import Muna
from muna.beta.openai import Message
from os import environ
from pathlib import Path
from statistics import mean, median
from sys import version_info
from typing import Annotated
from rich.console import Console
from rich.table import Table
from typer import run, Exit, Option

console = Console()

DEFAULT_TAG = "@google/gemma-4-26b-a4b-it"

class Category(StrEnum):
    """
    Prompt workload categories, each representative of a standard LLM benchmark
    so the prompt/output token shapes are reproducible.

    - reasoning     : BIG-Bench Hard / ARC
    - math          : GSM8K / MATH
    - code          : HumanEval / MBPP
    - knowledge     : MMLU / TriviaQA
    - summarization : CNN/DailyMail
    """
    reasoning = "reasoning"
    math = "math"
    code = "code"
    knowledge = "knowledge"
    summarization = "summarization"

# Representative prompts per category. Hard-coded (not sampled from the datasets)
# so a third party reproduces the exact same prompt shapes deterministically.
# Each is phrased to elicit a long answer so decode runs to `--max-tokens`.
PROMPTS: dict[Category, str] = {
    Category.reasoning: (
        "A farmer has 17 sheep. All but 9 run away. He then buys twice as many "
        "sheep as he has left, and gives a third of his total to his neighbor. "
        "Reason step by step, showing every intermediate count, and state how "
        "many sheep he ends up with."
    ),
    Category.math: (
        "Natalia sold clips to 48 of her friends in April, then sold half as "
        "many clips in May. In June she sold three times as many as May. Show "
        "your full working line by line and compute the total clips sold."
    ),
    Category.code: (
        "Write a well-documented Python function `is_palindrome(s: str) -> bool` "
        "that ignores punctuation, whitespace, and capitalization. Include a "
        "docstring, handle edge cases, and provide example usage with asserts."
    ),
    Category.knowledge: (
        "Explain the principal causes, key events, and long-term consequences "
        "of the French Revolution in detail, organized into clear paragraphs."
    ),
    Category.summarization: (
        "Summarize the following article in several detailed paragraphs.\n\n" +
        (
            "The history of computing is a long and winding road that stretches "
            "from mechanical calculators through vacuum tubes, transistors, and "
            "integrated circuits to the massively parallel accelerators of today. "
            "Each era introduced new abstractions, new bottlenecks, and new ways "
            "of thinking about how to turn electricity into useful work. "
        ) * 20
    ),
}

def main(
    tag: Annotated[str, Option(
        "--tag",
        help="Muna predictor tag (for muna targets)."
    )] = DEFAULT_TAG,
    category: Annotated[list[Category] | None, Option(
        "--category", "-c",
        help="Prompt categories to measure; repeatable. Defaults to ALL categories."
    )] = None,
    max_tokens: Annotated[int, Option(
        "--max-tokens",
        min=1,
        help="Sustained decode length used for the TPS measurement."
    )] = 512,
    rounds: Annotated[int, Option(
        "--rounds", "-r",
        min=1,
        help="Measured requests per (target, category)."
    )] = 20,
    warmup: Annotated[int, Option(
        "--warmup",
        min=0,
        help="Discarded warmup requests per (target, category)."
    )] = 3,
    output_dir: Annotated[Path, Option(
        "--output-dir",
        help="Directory for results/<metric>__<target>.json output."
    )] = Path("results"),
    json_output: Annotated[bool, Option(
        "--json/--no-json",
        help="Also echo raw result JSON to stdout."
    )] = False,
    quiet: Annotated[bool, Option(
        "--quiet", "-q",
        help="Suppress non-essential output."
    )] = False,
) -> None:
    """
    Measure decode tokens-per-second at batch size 1 across prompt categories.
    """
    categories = category or list(Category)
    prompts = { c.value: PROMPTS[c] for c in categories }
    output_dir.mkdir(parents=True, exist_ok=True)
    samples = _profile_muna_on_modal(
        tag=tag,
        prompts=prompts,
        max_tokens=max_tokens,
        rounds=rounds,
        warmup=warmup
    )
    record = {
        "metric": "tps",
        "tag": tag,
        "config": {
            "max_tokens": max_tokens,
            "rounds": rounds,
            "warmup": warmup,
            "batch_size": 1,
        },
        "categories": {
            name: { "summary": _summarize(values), "samples": values }
            for name, values in samples.items()
        },
    }
    output = output_dir / "tps.json"
    output.write_text(dumps(record, indent=2))
    if not quiet:
        _render(record["categories"], output)
    if json_output:
        console.print_json(dumps(record))
    raise Exit()

def _profile_muna_on_modal(
    *,
    tag: str,
    prompts: dict[str, str],
    max_tokens: int,
    rounds: int,
    warmup: int,
) -> dict[str, list[float]]:
    """
    Profile decode throughput (TPS) of a compiled Muna model on Modal.

    Runs the benchmark on a B200, looping every category in a single container so
    the model loads only once. The `--warmup` rounds subsume any weight download
    and JIT, so no separate preload step is needed; the volume persists the
    cache across runs.
    """
    volume = Volume.from_name("muna-benchmark", create_if_missing=True, version=2)
    image = (Image
        .debian_slim(python_version=f"{version_info.major}.{version_info.minor}")
        .uv_pip_install("muna")
    )
    secret = Secret.from_dict({ "MUNA_ACCESS_KEY": environ["MUNA_ACCESS_KEY"] })
    benchmark_app = App("muna-benchmark")
    def _benchmark() -> dict[str, list[float]]:
        from time import monotonic
        muna = Muna()
        results: dict[str, list[float]] = { }
        for name, text in prompts.items():
            tps_samples: list[float] = []
            for index in range(warmup + rounds):
                stream = muna.beta.openai.chat.completions.create(
                    model=tag,
                    messages=[
                        Message(role="user", content=text)
                    ],
                    max_completion_tokens=max_tokens,
                    stream=True,
                    acceleration="local_gpu"
                )
                t_first = t_last = None
                completion_tokens = 0
                for chunk in stream:
                    now = monotonic()
                    choice = chunk.choices[0] if chunk.choices else None
                    if choice is not None and choice.delta is not None and choice.delta.content:
                        if t_first is None:
                            t_first = now
                        t_last = now
                    if chunk.usage is not None:
                        completion_tokens = chunk.usage.completion_tokens
                if index < warmup:
                    continue
                if t_first is not None and t_last is not None and t_last > t_first and completion_tokens > 1:
                    tps_samples.append((completion_tokens - 1) / (t_last - t_first))
            results[name] = tps_samples
        return results
    benchmark_fn = benchmark_app.function(
        image=image,
        serialized=True,
        cpu=32,
        gpu="b200",
        volumes={ "/muna": volume },
        env={ "MUNA_HOME": "/muna" },
        secrets=[secret],
        timeout=3600,
    )(_benchmark)
    with enable_output(), benchmark_app.run():
        return benchmark_fn.remote()

def _summarize(samples: list[float]) -> dict[str, float | int]:
    if not samples:
        return { "n": 0 }
    ordered = sorted(samples)
    return {
        "n": len(ordered),
        "min": ordered[0],
        "p50": median(ordered),
        "max": ordered[-1],
        "mean": mean(ordered),
    }

def _render(categories: dict[str, dict], output: Path) -> None:
    table = Table(show_header=True, header_style="bold cyan", box=None)
    table.add_column("Category", style="bold green")
    table.add_column("n", style="white", justify="right")
    table.add_column("p50 TPS", style="white", justify="right")
    table.add_column("mean TPS", style="white", justify="right")
    for name, payload in categories.items():
        summary = payload["summary"]
        if summary.get("n"):
            table.add_row(name, str(summary["n"]), f"{summary['p50']:.1f}", f"{summary['mean']:.1f}")
        else:
            table.add_row(name, "0", "[red]no samples[/]", "")
    console.print(table)
    console.print(f"[cyan]Decode TPS[/] -> wrote [bold]{output}[/].")

if __name__ == "__main__":
    run(main)

## ttft.py
#
#   Muna
#   Copyright © 2026 NatML Inc. All Rights Reserved.
#

"""
Time To First Token (TTFT) vs prompt length benchmark for Gemma 4 26B A4B.

Measures prefill latency at batch size 1: for each target prompt length we
build a deterministic filler prompt, dispatch a streaming generation, and time
from the request to the first content token, so:

    ttft = t_first_token - t_request

We also read the actual `prompt_tokens` from the usage chunk, so the x-axis is
the real (tokenizer-measured) prompt length rather than the requested length.

Reproduce with:

    uv run --env-file .env ttft.py
"""

# /// script
# requires-python = ">=3.11"
# dependencies = ["muna", "modal", "typer"]
# ///

from __future__ import annotations
from json import dumps
from modal import enable_output, App, Image, Secret, Volume
from muna import Muna
from muna.beta.openai import Message
from os import environ
from pathlib import Path
from rich.console import Console
from rich.table import Table
from statistics import mean, median
from time import monotonic
from typing import Annotated
from typer import run, Exit, Option

console = Console()

DEFAULT_TAG = "@google/gemma-4-26b-a4b-it"
DEFAULT_PROMPT_LENS = [512, 2048, 4096, 8192, 16384]

# Deterministic filler so a third party reproduces the same prompt shapes.
# We size by a rough English chars-per-token heuristic, then report the actual
# `prompt_tokens` measured by the model's tokenizer.
_FILLER = (
    "The history of computing stretches from mechanical calculators through "
    "vacuum tubes and transistors to the massively parallel accelerators of "
    "today, where each era introduced new abstractions and new bottlenecks. "
)
_CHARS_PER_TOKEN = 4

def _build_prompt(target_tokens: int) -> str:
    target_chars = max(target_tokens * _CHARS_PER_TOKEN, len(_FILLER))
    repeats = target_chars // len(_FILLER) + 1
    body = (_FILLER * repeats)[:target_chars]
    return f"{body}\n\nIn one word, what is the topic of the passage above?"

def main(
    tag: Annotated[str, Option(
        "--tag",
        help="Muna predictor tag."
    )] = DEFAULT_TAG,
    prompt_lens: Annotated[list[int], Option(
        "--prompt-len",
        help="Input-token length to sweep; repeatable. Defaults to 512/2k/8k/16k."
    )] = DEFAULT_PROMPT_LENS,
    max_tokens: Annotated[int, Option(
        "--max-tokens",
        min=1,
        help="Tokens to request; TTFT only needs the first token."
    )] = 1,
    rounds: Annotated[int, Option(
        "--rounds", "-r",
        min=1,
        help="Measured requests per prompt length."
    )] = 50,
    warmup: Annotated[int, Option(
        "--warmup",
        min=0,
        help="Discarded warmup requests per prompt length."
    )] = 3,
    output_dir: Annotated[Path, Option(
        "--output-dir",
        help="Directory for results/<metric>.json output."
    )] = Path("results"),
    json_output: Annotated[bool, Option(
        "--json/--no-json",
        help="Also echo raw result JSON to stdout."
    )] = False,
    quiet: Annotated[bool, Option(
        "--quiet", "-q",
        help="Suppress non-essential output."
    )] = False,
) -> None:
    """
    Measure TTFT at batch size 1 across a sweep of input prompt lengths.
    """
    prompts = { str(n): _build_prompt(n) for n in prompt_lens }
    output_dir.mkdir(parents=True, exist_ok=True)
    volume = Volume.from_name("muna-benchmark", create_if_missing=True, version=2)
    image = (Image
        .from_registry("python:3.12-slim-bookworm")
        .uv_pip_install("muna")
    )
    secret = Secret.from_dict({ "MUNA_ACCESS_KEY": environ["MUNA_ACCESS_KEY"] })
    benchmark_app = App("muna-benchmark")
    benchmark_fn = benchmark_app.function(
        image=image,
        cpu=32,
        gpu="b200",
        volumes={ "/muna": volume },
        env={ "MUNA_HOME": "/muna" },
        secrets=[secret],
        timeout=3600,
    )(_benchmark)
    with enable_output(), benchmark_app.run():
        samples = benchmark_fn.remote(prompts, tag, max_tokens, rounds, warmup)
    record = {
        "metric": "ttft",
        "tag": tag,
        "config": {
            "max_tokens": max_tokens,
            "rounds": rounds,
            "warmup": warmup,
            "batch_size": 1,
        },
        "lengths": {
            name: {
                "requested_tokens": int(name),
                "prompt_tokens": payload["prompt_tokens"],
                "summary": _summarize(payload["ttft"]),
                "samples": payload["ttft"],
            }
            for name, payload in samples.items()
        },
    }
    output = output_dir / "ttft.json"
    output.write_text(dumps(record, indent=2))
    if not quiet:
        _render(record["lengths"], output)
    if json_output:
        console.print_json(dumps(record))
    raise Exit()

def _benchmark(
    prompts: dict[str, str],
    tag: str,
    max_tokens: int,
    rounds: int,
    warmup: int,
) -> dict[str, dict[str, object]]:
    muna = Muna()
    results: dict[str, dict[str, object]] = { }
    for name, text in prompts.items():
        ttft_samples: list[float] = []
        prompt_tokens = 0
        for index in range(warmup + rounds):
            t_request = monotonic()
            stream = muna.beta.openai.chat.completions.create(
                model=tag,
                messages=[
                    Message(role="user", content=text)
                ],
                max_completion_tokens=max_tokens,
                stream=True,
                acceleration="local_gpu"
            )
            t_first = None
            for chunk in stream:
                now = monotonic()
                choice = chunk.choices[0] if chunk.choices else None
                if t_first is None and choice is not None and choice.delta is not None and choice.delta.content:
                    t_first = now
                if chunk.usage is not None:
                    prompt_tokens = chunk.usage.prompt_tokens
            if index < warmup:
                continue
            if t_first is not None:
                ttft_samples.append(t_first - t_request)
        results[name] = { "ttft": ttft_samples, "prompt_tokens": prompt_tokens }
    return results

def _summarize(samples: list[float]) -> dict[str, float | int]:
    if not samples:
        return { "n": 0 }
    ordered = sorted(samples)
    return {
        "n": len(ordered),
        "min": ordered[0],
        "p50": median(ordered),
        "max": ordered[-1],
        "mean": mean(ordered),
    }

def _render(lengths: dict[str, dict], output: Path) -> None:
    table = Table(show_header=True, header_style="bold cyan", box=None)
    table.add_column("Requested", style="bold green", justify="right")
    table.add_column("Prompt tokens", style="white", justify="right")
    table.add_column("n", style="white", justify="right")
    table.add_column("p50 TTFT (ms)", style="white", justify="right")
    table.add_column("mean TTFT (ms)", style="white", justify="right")
    for name, payload in lengths.items():
        summary = payload["summary"]
        if summary.get("n"):
            table.add_row(
                name,
                str(payload["prompt_tokens"]),
                str(summary["n"]),
                f"{summary['p50'] * 1e3:.1f}",
                f"{summary['mean'] * 1e3:.1f}",
            )
        else:
            table.add_row(name, str(payload["prompt_tokens"]), "0", "[red]no samples[/]", "")
    console.print(table)
    console.print(f"[cyan]TTFT[/] -> wrote [bold]{output}[/].")

if __name__ == "__main__":
    run(main)
	# Muna
	MUNA_ACCESS_KEY=

	# Huggingface
	# This is only required by `cttft.py` which needs to download
	# weights from HF in order to benchmark upstream SGLang.
	HF_TOKEN=
	#
	# Muna
	# Copyright © 2026 NatML Inc. All Rights Reserved.
	#

	# /// script
	# requires-python = ">=3.11"
	# dependencies = ["muna", "modal", "truss", "typer"]
	# ///

	"""
	Cold-start Time To First Token (CTTFT) benchmark for Gemma 4 26B A4B.

	Measures the wall-clock time from a request hitting a cold system to the first
	generated token across two axes: the inference `--runtime` (muna \| sglang) and
	the `--provider` (modal \| baseten \| ssh). Weights are pre-staged out of band (a
	Modal Volume for Modal, a preload command for SSH, the Baseten Delivery Network
	for Baseten) so the measurement captures container/process boot + model load +
	prefill + first token, not a one-off weight download from Hugging Face.

	For Modal we run with `single_use_containers=True` so every boot lands on a
	fresh container. For SSH each boot is a fresh process and the page cache is
	dropped between boots (by default) for true cold-disk reads.

	Reproduce with:

	uv run --env-file .env cttft.py --provider modal

	NOTE: This test requires an `HF_TOKEN` in your environment, to download the
	weights.
	"""

	from __future__ import annotations
	from collections.abc import Callable, Iterator
	from contextlib import contextmanager
	from enum import StrEnum
	from json import dumps
	from math import ceil, floor
	from modal import enable_output, App, Image, Secret, Volume
	from muna import Muna
	from muna.beta.openai import Message
	from os import environ
	from pathlib import Path
	from requests import post
	from rich.console import Console
	from rich.table import Table
	from statistics import mean
	from string import Template
	from subprocess import run
	from sys import version_info
	from tempfile import TemporaryDirectory
	from time import monotonic
	from typer import run as typer_run, Exit, Option
	from typing import Annotated

	console = Console()

	DEFAULT_TAG = "@google/gemma-4-26b-a4b-it"

	class Runtime(StrEnum):
	"""
	Inference runtime under test.
	"""
	muna = "muna"
	sglang = "sglang"

	class Provider(StrEnum):
	"""
	Compute provider.
	"""
	modal = "modal"
	baseten = "baseten"
	ssh = "ssh"

	def main(
	runtime: Annotated[list[Runtime] \| None, Option(
	"--runtime", "-r",
	help="Inference runtime to benchmark; repeatable. Defaults to ALL runtimes."
	)] = None,
	provider: Annotated[list[Provider] \| None, Option(
	"--provider", "-p",
	help="Compute provider to benchmark; repeatable. Defaults to ALL providers."
	)] = None,
	boots: Annotated[int, Option(
	"--boots", "-n",
	min=1,
	help="Number of cold boots per (runtime, provider)."
	)] = 5,
	prompt: Annotated[str, Option(
	"--prompt",
	help="Fixed minimal prompt sent on every cold boot."
	)] = "What is the capital of France?",
	max_tokens: Annotated[int, Option(
	"--max-tokens",
	min=1,
	help="Tokens to request; CTTFT only needs the first token."
	)] = 1,
	tag: Annotated[str, Option(
	"--tag",
	help="Muna predictor tag (for muna runtimes)."
	)] = DEFAULT_TAG,
	ssh_host: Annotated[str \| None, Option(
	"--ssh-host",
	help="SSH target for the ssh provider (alias or user@host)."
	)] = None,
	ssh_command: Annotated[str \| None, Option(
	"--ssh-command",
	help="Remote command timed per boot. Supports $prompt/$max_tokens/$tag placeholders."
	)] = None,
	ssh_preload_command: Annotated[str \| None, Option(
	"--ssh-preload-command",
	help="Optional remote command run once (untimed) before boots to stage weights/cache."
	)] = None,
	ssh_port: Annotated[int \| None, Option(
	"--ssh-port",
	help="SSH port (defaults to the system ssh default)."
	)] = None,
	ssh_identity: Annotated[Path \| None, Option(
	"--ssh-identity",
	help="SSH identity (private key) file."
	)] = None,
	drop_caches: Annotated[bool, Option(
	"--drop-caches/--no-drop-caches",
	help="Drop the page cache before each ssh boot for true cold-disk reads (needs passwordless sudo)."
	)] = True,
	output_dir: Annotated[Path, Option(
	"--output-dir",
	help="Directory for results/<metric>__<target>.json output."
	)] = Path("results"),
	json_output: Annotated[bool, Option(
	"--json/--no-json",
	help="Also echo raw result JSON to stdout."
	)] = False,
	quiet: Annotated[bool, Option(
	"--quiet", "-q",
	help="Suppress non-essential output."
	)] = False,
	) -> None:
	"""
	Measure cold-start time to first token across `--boots` cold boots.

	Benchmarks the Cartesian product of selected runtimes and providers, writing
	one results file per (runtime, provider) so distribution statistics
	(p50/p90) can be computed later.
	"""
	runtimes = runtime or list(Runtime)
	providers = provider or list(Provider)
	if Provider.ssh in providers:
	if not ssh_host or not ssh_command:
	console.print("[red]error:[/] --provider ssh requires --ssh-host and --ssh-command.")
	raise Exit(code=1)
	if len(runtimes) != 1:
	console.print(
	"[red]error:[/] --provider ssh uses a single model-agnostic --ssh-command; "
	"select exactly one --runtime (it is only a results label)."
	)
	raise Exit(code=1)
	output_dir.mkdir(parents=True, exist_ok=True)
	records: dict[str, dict] = { }
	for prov in providers:
	for rt in runtimes:
	session = _make_session(
	rt, prov,
	tag=tag,
	prompt=prompt,
	max_tokens=max_tokens,
	ssh_host=ssh_host,
	ssh_command=ssh_command,
	ssh_preload_command=ssh_preload_command,
	ssh_port=ssh_port,
	ssh_identity=ssh_identity,
	drop_caches=drop_caches,
	)
	label = f"{rt.value}-on-{prov.value}"
	if session is None:
	console.print(f"[yellow]skip[/] {label}: not implemented")
	continue
	samples: list[float] = []
	with session as boot_once:
	for index in range(boots):
	elapsed = boot_once()
	samples.append(elapsed)
	print(f"{label} boot {index + 1}/{boots}: {elapsed:.2f}s")
	record = {
	"metric": "cttft",
	"target": label,
	"runtime": rt.value,
	"provider": prov.value,
	"tag": tag,
	"config": {
	"boots": boots,
	"prompt": prompt,
	"max_tokens": max_tokens,
	"batch_size": 1,
	},
	"summary": _summarize(samples),
	"samples": samples,
	}
	output = output_dir / f"cttft__{label}.json"
	output.write_text(dumps(record, indent=2))
	records[label] = record
	if not quiet and records:
	_render(records, output_dir)
	if json_output:
	console.print_json(dumps(records))
	raise Exit()

	def _make_session(
	runtime: Runtime,
	provider: Provider,
	*,
	tag: str,
	prompt: str,
	max_tokens: int,
	ssh_host: str \| None,
	ssh_command: str \| None,
	ssh_preload_command: str \| None,
	ssh_port: int \| None,
	ssh_identity: Path \| None,
	drop_caches: bool,
	):
	"""
	Return a session context manager for `(runtime, provider)`, or `None` if the
	combination is not implemented (skipped with a warning by the caller).
	"""
	match (runtime, provider):
	case (Runtime.muna, Provider.modal):
	return _muna_modal_session(tag=tag, prompt=prompt, max_tokens=max_tokens)
	case (Runtime.sglang, Provider.modal):
	return _sglang_modal_session(prompt=prompt, max_tokens=max_tokens)
	case (_, Provider.ssh):
	return _ssh_session(
	host=ssh_host,
	command=ssh_command,
	preload_command=ssh_preload_command,
	port=ssh_port,
	identity=ssh_identity,
	drop_caches=drop_caches,
	prompt=prompt,
	max_tokens=max_tokens,
	tag=tag,
	)
	case _:
	return None # baseten remains a documented stub (see _profile_*_on_baseten)

	@contextmanager
	def _muna_modal_session(
	*,
	tag: str,
	prompt: str,
	max_tokens: int,
	) -> Iterator[Callable[[], float]]:
	"""
	Session for a compiled Muna model on Modal.

	Preloads the model weights onto a Modal volume (untimed) so cold boots read
	weights from the volume instead of downloading them from Hugging Face, then
	yields a `boot_once` that times one fresh-container cold boot.
	`single_use_containers=True` retires the container after every input, so each
	boot is genuinely cold.
	"""
	# Pre-load the model weights on a Modal volume
	preload_app = App("muna-benchmark-preload")
	volume = Volume.from_name("muna-benchmark", create_if_missing=True, version=2)
	image = (Image
	.from_registry(f"python:{version_info.major}.{version_info.minor}-slim-bookworm") # use smallest Python image possible
	.uv_pip_install("muna")
	)
	secret = Secret.from_dict({ "MUNA_ACCESS_KEY": environ["MUNA_ACCESS_KEY"] })
	def _preload():
	muna = Muna()
	muna.predictions.create(tag, inputs={ }) # Muna preloads but doesn't execute compiled model
	preload_fn = preload_app.function(
	image=image,
	serialized=True,
	volumes={ "/muna": volume },
	env={ "MUNA_HOME": "/muna" },
	secrets=[secret],
	timeout=900,
	)(_preload)
	with enable_output(), preload_app.run():
	preload_fn.remote()
	print("Preloaded model weights on Modal volume")
	# Benchmark on a B200, one fresh container per cold boot
	benchmark_app = App("muna-benchmark")
	def _benchmark(tag: str, prompt: str, max_tokens: int):
	muna = Muna()
	message = Message(role="user", content=prompt)
	stream = muna.predictions.stream(
	tag,
	inputs={ "messages": [message], "max_output_tokens": max_tokens },
	acceleration="local_gpu"
	)
	next(stream) # block until the first token is produced
	benchmark_fn = benchmark_app.function(
	image=image,
	serialized=True,
	cpu=32,
	gpu="b200",
	volumes={ "/muna": volume },
	env={ "MUNA_HOME": "/muna" },
	secrets=[secret],
	timeout=300,
	single_use_containers=True,
	)(_benchmark)
	with enable_output(), benchmark_app.run():
	def boot_once() -> float:
	start = monotonic()
	benchmark_fn.remote(tag, prompt, max_tokens)
	return monotonic() - start
	yield boot_once

	@contextmanager
	def _sglang_modal_session(
	*,
	prompt: str,
	max_tokens: int,
	) -> Iterator[Callable[[], float]]:
	"""
	Session for SGLang on Modal.

	Mirrors `_muna_modal_session`: preloads the gated weights onto a Modal volume
	(untimed), then yields a `boot_once` that times one cold boot of a fresh
	single-use container.
	"""
	repo = "google/gemma-4-26B-A4B-it"
	model_dir = "/sglang/gemma-4-26b-a4b-it"
	# Preload the model weights on a Modal volume
	preload_app = App("sglang-benchmark-preload")
	volume = Volume.from_name("sglang-benchmark", create_if_missing=True, version=2)
	preload_image = (Image
	.from_registry(f"python:{version_info.major}.{version_info.minor}-slim-bookworm")
	.uv_pip_install("huggingface_hub[hf_transfer]")
	.env({ "HF_HUB_ENABLE_HF_TRANSFER": "1" })
	)
	secret = Secret.from_dict({ "HF_TOKEN": environ["HF_TOKEN"] })
	def _preload():
	from huggingface_hub import snapshot_download
	snapshot_download(repo, local_dir=model_dir) # download the gated weights once
	preload_fn = preload_app.function(
	image=preload_image,
	serialized=True,
	volumes={ "/sglang": volume },
	secrets=[secret],
	timeout=900,
	)(_preload)
	with enable_output(), preload_app.run():
	preload_fn.remote()
	print("Preloaded model weights on Modal volume")
	# Benchmark on a B200, one fresh container per cold boot
	benchmark_app = App("sglang-benchmark")
	image = Image.from_registry(
	"lmsysorg/sglang:v0.5.13",
	add_python=f"{version_info.major}.{version_info.minor}"
	)
	def _benchmark(model_dir: str, prompt: str, max_tokens: int):
	import sglang as sgl
	engine = sgl.Engine(model_path=model_dir)
	engine.generate(prompt, sampling_params={ "max_new_tokens": max_tokens })
	benchmark_fn = benchmark_app.function(
	image=image,
	serialized=True,
	cpu=32,
	gpu="b200",
	volumes={ "/sglang": volume },
	timeout=300,
	single_use_containers=True,
	)(_benchmark)
	with enable_output(), benchmark_app.run():
	def boot_once() -> float:
	start = monotonic()
	benchmark_fn.remote(model_dir, prompt, max_tokens)
	return monotonic() - start
	yield boot_once

	@contextmanager
	def _ssh_session(
	*,
	host: str,
	command: str,
	preload_command: str \| None,
	port: int \| None,
	identity: Path \| None,
	drop_caches: bool,
	prompt: str,
	max_tokens: int,
	tag: str,
	) -> Iterator[Callable[[], float]]:
	"""
	Session for a runtime on a baremetal box reached over SSH.

	Shells out to the system `ssh` binary (no extra deps). On enter it runs the
	optional preload once (untimed, mirroring Modal staging weights to a volume).
	Each boot is a fresh remote process; the page cache is dropped (untimed,
	unless disabled) beforehand so the timed command reads weights from disk.
	Commands are templated with `$prompt`/`$max_tokens`/`$tag`.
	"""
	base = ["ssh", "-o", "BatchMode=yes"]
	if port is not None:
	base += ["-p", str(port)]
	if identity is not None:
	base += ["-i", str(identity)]
	def fill(template: str) -> str:
	return Template(template).safe_substitute(prompt=prompt, max_tokens=max_tokens, tag=tag)
	if preload_command:
	run([*base, host, fill(preload_command)], check=True) # untimed: populate weight/predictor cache once
	def boot_once() -> float:
	if drop_caches:
	# untimed: evict weights from the page cache -> true cold-disk read
	run([*base, host, "sync && echo 1 \| sudo tee /proc/sys/vm/drop_caches > /dev/null"], check=True)
	start = monotonic()
	run([*base, host, fill(command)], check=True) # process cold-start: load + first token + exit
	return monotonic() - start
	yield boot_once

	def _profile_muna_on_baseten():
	pass

	def _profile_sglang_on_baseten():
	"""
	Profile the Cold-start Time to First Token (CTTFT) of SGLang on Baseten (B10).

	Deploys a custom SGLang Docker server whose weights are delivered by the
	Baseten Delivery Network (BDN), so cold boots read weights from Baseten's
	in-cluster cache instead of downloading from Hugging Face. The Truss config
	is generated procedurally and pushed via the programmatic API (no CLI).

	Stub: not yet wired into the session/loop model (Baseten access pending).
	"""
	from truss import login, push
	from truss.base.truss_config import (
	BaseImage, DockerServer, Resources,
	TrussConfig, Weights, WeightsSource
	)
	login(environ["BASETEN_API_KEY"])
	config = TrussConfig(
	model_name="gemma4-sglang-cttft",
	resources=Resources(accelerator="B200"),
	base_image=BaseImage(image="lmsysorg/sglang:v0.5.13"),
	docker_server=DockerServer(
	start_command=(
	"python3 -m sglang.launch_server --model-path /models/gemma "
	"--served-model-name gemma-4 --host 0.0.0.0 --port 8000"
	),
	server_port=8000,
	predict_endpoint="/v1/chat/completions",
	readiness_endpoint="/health",
	liveness_endpoint="/health",
	),
	weights=Weights([
	WeightsSource(
	source="hf://google/gemma-4-26B-A4B-it@main", # TODO: pin to a commit SHA
	mount_location="/models/gemma",
	#auth_secret_name="hf_access_token", # Baseten secret holding the HF token
	allow_patterns=[".safetensors", ".json", "tokenizer*"],
	)
	]),
	)
	# Push the procedurally-generated Truss (no model.py needed for a docker_server)
	with TemporaryDirectory() as directory:
	config.write_to_yaml_file(Path(directory) / "config.yaml")
	deployment = push(
	directory,
	remote="baseten",
	publish=True,
	model_name=config.model_name
	)
	deployment.wait_for_active(timeout_seconds=30 * 60) # blocks through BDN mirroring + boot
	predict_url = deployment._baseten_service.predict_url
	print(f"Deployed SGLang on Baseten: {predict_url}")
	# Fire one request to confirm the endpoint serves
	response = post(
	predict_url,
	headers={ "Authorization": f"Api-Key {environ['BASETEN_API_KEY']}" },
	json={ "messages": [{ "role": "user", "content": "What is the capital of France?" }], "max_tokens": 1 },
	)
	print(response.json())

	def _percentile(ordered: list[float], pct: float) -> float:
	if len(ordered) == 1:
	return ordered[0]
	rank = (len(ordered) - 1) * pct / 100
	low, high = floor(rank), ceil(rank)
	if low == high:
	return ordered[int(rank)]
	return ordered[low] * (high - rank) + ordered[high] * (rank - low)

	def _summarize(samples: list[float]) -> dict[str, float \| int]:
	if not samples:
	return { "n": 0 }
	ordered = sorted(samples)
	return {
	"n": len(ordered),
	"min": ordered[0],
	"p50": _percentile(ordered, 50),
	"p90": _percentile(ordered, 90),
	"max": ordered[-1],
	"mean": mean(ordered),
	}

	def _render(records: dict[str, dict], output_dir: Path) -> None:
	table = Table(show_header=True, header_style="bold cyan", box=None)
	table.add_column("Target", style="bold green")
	table.add_column("n", style="white", justify="right")
	table.add_column("p50 (s)", style="white", justify="right")
	table.add_column("p90 (s)", style="white", justify="right")
	table.add_column("mean (s)", style="white", justify="right")
	for name, record in records.items():
	summary = record["summary"]
	if summary.get("n"):
	table.add_row(
	name,
	str(summary["n"]),
	f"{summary['p50']:.2f}",
	f"{summary['p90']:.2f}",
	f"{summary['mean']:.2f}",
	)
	else:
	table.add_row(name, "0", "[red]no samples[/]", "", "")
	console.print(table)
	console.print(f"[cyan]CTTFT[/] -> wrote results to [bold]{output_dir}[/].")

	if __name__ == "__main__":
	typer_run(main)
	#
	# Muna
	# Copyright © 2026 NatML Inc. All Rights Reserved.
	#

	"""
	Artifact size registries for Gemma 4 26B A4B.

	Measures the effective size of model distributions across a container
	image and a compiled binary generated by Muna.

	Requirements:
	1. Create a Muna access key at http://muna.ai/settings/developer.
	2. Add your Muna access key to the `.env` file.
	3. Login to the Modal CLI (for getting the size of a container image).

	Reproduce with:

	uv run --env-file .env size.py --compiled-model @google/gemma-4-26b-a4b-it
	uv run --env-file .env size.py --image lmsysorg/sglang:latest
	"""

	# /// script
	# requires-python = ">=3.11"
	# dependencies = ["modal", "muna", "pyyaml", "typer"]
	# ///

	from __future__ import annotations
	from enum import StrEnum
	from json import dumps
	from modal import enable_output, App, Image
	from pathlib import Path
	from requests import head
	from rich.console import Console
	from rich.panel import Panel
	from rich.table import Table
	from typer import run, Exit, Option
	from typing import Annotated
	from urllib.parse import urlparse
	from yaml import safe_dump

	console = Console()

	class OutputFormat(StrEnum):
	yaml = "yaml"
	json = "json"

	def main(
	image: Annotated[str \| None, Option(
	"--image",
	help="SGLang Docker image reference to inspect."
	)] = None,
	compiled_model: Annotated[str \| None, Option(
	help="Compiled Muna predictor tag to inspect."
	)] = None,
	output: Annotated[Path \| None, Option(
	"--output", "-o",
	help="Output registry path. Defaults per mode."
	)] = None,
	output_format: Annotated[OutputFormat, Option(
	"--format",
	help="Registry output format."
	)] = OutputFormat.yaml,
	min_size: Annotated[int, Option(
	"--min-size", min=0,
	help="Omit files smaller than this many bytes (image mode only)."
	)] = 0,
	) -> None:
	"""
	Build a footprint size registry for either a Muna predictor or a Docker image.

	Pass exactly one of --compiled-model or --image.
	"""
	if (image is None) == (compiled_model is None):
	console.print(
	"[red]Pass exactly one of[/] [bold]--compiled-model[/] [red]or[/] [bold]--image[/]."
	)
	raise Exit(code=1)
	if compiled_model is not None:
	_run_predictor(
	compiled_model,
	output or Path("registries/muna.yaml"),
	output_format
	)
	else:
	_run_image(
	image,
	output or Path("registries/container.yaml"),
	output_format,
	min_size
	)

	def _run_predictor(
	tag: str,
	output: Path,
	output_format: OutputFormat
	) -> None:
	"""
	Emit the Muna predictor size registry (file name and size only).
	"""
	from muna import Muna
	muna = Muna()
	prediction = muna.predictions.create(tag, client_id="x86_64-unknown-linux-gnu")
	resources = prediction.resources or []
	dso_resources = [res for res in resources if res.type == "dso"]
	if not dso_resources:
	console.print(f"[red]No DSO resources found for[/] [bold]{tag}[/].")
	raise Exit(code=1)
	files: list[dict[str, object]] = []
	for resource in dso_resources:
	name, size = _head_file(resource)
	files.append({ "name": name, "size": size })
	files.sort(key=lambda entry: int(entry["size"]), reverse=True)
	total = sum(int(entry["size"]) for entry in files)
	_write_registry(output, output_format, { "tag": tag, "total_size": total, "files": files })
	_render_registry(f"Muna size registry: {tag}", files, total, output)

	def _run_image(
	image: str,
	output: Path,
	output_format: OutputFormat,
	min_size: int
	) -> None:
	"""
	Emit the SGLang image file size registry (name, size; no descriptions).

	Spins up an ephemeral Modal app whose container is built from the given
	image, walks its filesystem, and records every regular file with its size.
	"""
	modal_app = App("gemma4-sglang-image-size")
	base = Image.from_registry(image)
	# `_walk_filesystem` lives at module scope; we apply the Modal decorator at
	# runtime so the container image can be driven by --image.
	walk_filesystem = modal_app.function(
	image=base,
	timeout=900,
	)(_walk_filesystem)
	with enable_output(), modal_app.run():
	result = walk_filesystem.remote(min_size)
	files = result["files"]
	total = int(result["total_size"])
	file_count = int(result["file_count"])
	_write_registry(output, output_format, {
	"image": image,
	"total_size": total,
	"file_count": file_count,
	"files": files,
	})
	_render_image_summary(image, files, total, file_count, output)

	def _walk_filesystem(min_bytes: int) -> dict[str, object]:
	"""
	Walk the container root filesystem and tally regular files by size.
	Runs inside the Modal container, so it relies only on the standard library.
	"""
	# Skip pseudo-filesystems that do not contribute to image footprint.
	skip = { Path("/proc"), Path("/sys"), Path("/dev"), Path("/run") }
	files: list[dict[str, object]] = []
	total = 0
	count = 0
	stack = [Path("/")]
	while stack:
	try:
	entries = list(stack.pop().iterdir())
	except OSError:
	continue
	for entry in entries:
	# Skip symlinks so symlinked dirs are not descended (no loops) and
	# symlinks never count as regular files.
	if entry in skip or entry.is_symlink():
	continue
	try:
	if entry.is_dir():
	stack.append(entry)
	elif entry.is_file():
	size = entry.stat().st_size
	total += size
	count += 1
	if size >= min_bytes:
	files.append({ "name": str(entry), "size": size })
	except OSError:
	continue
	files.sort(key=lambda entry: int(entry["size"]), reverse=True)
	return { "total_size": total, "file_count": count, "files": files }

	def _head_file(resource) -> tuple[str, int]:
	"""
	Get the name and size of a Muna prediction resource with a HEAD request.
	"""
	response = head(
	resource.url,
	headers={ "User-Agent": "muna-size" },
	allow_redirects=True
	)
	response.raise_for_status()
	length = response.headers.get("Content-Length")
	if length is None:
	raise RuntimeError(f"HEAD response for {resource.url} had no Content-Length")
	name = Path(urlparse(resource.url).path).name
	return name, int(length)

	def _write_registry(
	output: Path,
	output_format: OutputFormat,
	data: dict[str, object]
	) -> None:
	output.parent.mkdir(parents=True, exist_ok=True)
	text = (
	dumps(data, indent=2)
	if output_format is OutputFormat.json
	else safe_dump(data, sort_keys=False)
	)
	output.write_text(text, encoding="utf-8")

	def _render_registry(
	title: str,
	files: list[dict[str, object]],
	total: int,
	output: Path
	) -> None:
	table = Table(show_header=True, header_style="bold cyan", box=None)
	table.add_column("File", style="bold green")
	table.add_column("Size", style="white", justify="right")
	for entry in files:
	table.add_row(str(entry["name"]), _human_size(int(entry["size"])))
	table.add_row("", "")
	table.add_row("[bold magenta]Total[/]", f"[bold magenta]{_human_size(total)}[/]")
	console.print(Panel(table, title=f"[bold green]{title}[/]", border_style="green"))
	console.print(f"Wrote registry to [bold]{output}[/].")

	def _render_image_summary(
	image: str,
	files: list[dict[str, object]],
	total: int,
	file_count: int,
	output: Path,
	top: int = 20,
	) -> None:
	table = Table(show_header=True, header_style="bold cyan", box=None)
	table.add_column("File", style="bold green", overflow="fold")
	table.add_column("Size", style="white", justify="right")
	for entry in files[:top]:
	table.add_row(str(entry["name"]), _human_size(int(entry["size"])))
	if len(files) > top:
	table.add_row(f"[dim]... and {len(files) - top} more files[/]", "")
	table.add_row("", "")
	table.add_row("[bold magenta]Total[/]", f"[bold magenta]{_human_size(total)}[/]")
	console.print(Panel(table, title=f"[bold green]SGLang image: {image}[/]", border_style="green"))
	console.print(f"Recorded {file_count} files. Wrote registry to [bold]{output}[/].")

	def _human_size(num_bytes: int) -> str:
	"""
	Human readable file size.
	"""
	value = float(num_bytes)
	for unit in ("B", "KiB", "MiB", "GiB"):
	if value < 1024.0 or unit == "GiB":
	return f"{value:.1f} {unit}" if unit != "B" else f"{int(value)} B"
	value /= 1024.0
	raise AssertionError("unreachable")

	if __name__ == "__main__":
	run(main)
	#
	# Muna
	# Copyright © 2026 NatML Inc. All Rights Reserved.
	#

	"""
	Decode throughput (TPS) by prompt category benchmark for Gemma 4 26B A4B.

	Measures sustained decode tokens-per-second at batch size 1. For each category
	we stream a generation, time from the first content token to the last, and read
	the exact completion-token count from the final usage chunk, so:

	decode_tps = (completion_tokens - 1) / (t_last_token - t_first_token)

	which excludes prefill (that is the TTFT benchmark's concern).

	Reproduce with:

	uv run --env-file .env tps.py
	"""

	# /// script
	# requires-python = ">=3.11"
	# dependencies = ["muna", "modal", "typer"]
	# ///

	from __future__ import annotations
	from enum import StrEnum
	from json import dumps
	from modal import enable_output, App, Image, Secret, Volume
	from muna import Muna
	from muna.beta.openai import Message
	from os import environ
	from pathlib import Path
	from statistics import mean, median
	from sys import version_info
	from typing import Annotated
	from rich.console import Console
	from rich.table import Table
	from typer import run, Exit, Option

	console = Console()

	DEFAULT_TAG = "@google/gemma-4-26b-a4b-it"

	class Category(StrEnum):
	"""
	Prompt workload categories, each representative of a standard LLM benchmark
	so the prompt/output token shapes are reproducible.

	- reasoning : BIG-Bench Hard / ARC
	- math : GSM8K / MATH
	- code : HumanEval / MBPP
	- knowledge : MMLU / TriviaQA
	- summarization : CNN/DailyMail
	"""
	reasoning = "reasoning"
	math = "math"
	code = "code"
	knowledge = "knowledge"
	summarization = "summarization"

	# Representative prompts per category. Hard-coded (not sampled from the datasets)
	# so a third party reproduces the exact same prompt shapes deterministically.
	# Each is phrased to elicit a long answer so decode runs to `--max-tokens`.
	PROMPTS: dict[Category, str] = {
	Category.reasoning: (
	"A farmer has 17 sheep. All but 9 run away. He then buys twice as many "
	"sheep as he has left, and gives a third of his total to his neighbor. "
	"Reason step by step, showing every intermediate count, and state how "
	"many sheep he ends up with."
	),
	Category.math: (
	"Natalia sold clips to 48 of her friends in April, then sold half as "
	"many clips in May. In June she sold three times as many as May. Show "
	"your full working line by line and compute the total clips sold."
	),
	Category.code: (
	"Write a well-documented Python function `is_palindrome(s: str) -> bool` "
	"that ignores punctuation, whitespace, and capitalization. Include a "
	"docstring, handle edge cases, and provide example usage with asserts."
	),
	Category.knowledge: (
	"Explain the principal causes, key events, and long-term consequences "
	"of the French Revolution in detail, organized into clear paragraphs."
	),
	Category.summarization: (
	"Summarize the following article in several detailed paragraphs.\n\n" +
	(
	"The history of computing is a long and winding road that stretches "
	"from mechanical calculators through vacuum tubes, transistors, and "
	"integrated circuits to the massively parallel accelerators of today. "
	"Each era introduced new abstractions, new bottlenecks, and new ways "
	"of thinking about how to turn electricity into useful work. "
	) * 20
	),
	}

	def main(
	tag: Annotated[str, Option(
	"--tag",
	help="Muna predictor tag (for muna targets)."
	)] = DEFAULT_TAG,
	category: Annotated[list[Category] \| None, Option(
	"--category", "-c",
	help="Prompt categories to measure; repeatable. Defaults to ALL categories."
	)] = None,
	max_tokens: Annotated[int, Option(
	"--max-tokens",
	min=1,
	help="Sustained decode length used for the TPS measurement."
	)] = 512,
	rounds: Annotated[int, Option(
	"--rounds", "-r",
	min=1,
	help="Measured requests per (target, category)."
	)] = 20,
	warmup: Annotated[int, Option(
	"--warmup",
	min=0,
	help="Discarded warmup requests per (target, category)."
	)] = 3,
	output_dir: Annotated[Path, Option(
	"--output-dir",
	help="Directory for results/<metric>__<target>.json output."
	)] = Path("results"),
	json_output: Annotated[bool, Option(
	"--json/--no-json",
	help="Also echo raw result JSON to stdout."
	)] = False,
	quiet: Annotated[bool, Option(
	"--quiet", "-q",
	help="Suppress non-essential output."
	)] = False,
	) -> None:
	"""
	Measure decode tokens-per-second at batch size 1 across prompt categories.
	"""
	categories = category or list(Category)
	prompts = { c.value: PROMPTS[c] for c in categories }
	output_dir.mkdir(parents=True, exist_ok=True)
	samples = _profile_muna_on_modal(
	tag=tag,
	prompts=prompts,
	max_tokens=max_tokens,
	rounds=rounds,
	warmup=warmup
	)
	record = {
	"metric": "tps",
	"tag": tag,
	"config": {
	"max_tokens": max_tokens,
	"rounds": rounds,
	"warmup": warmup,
	"batch_size": 1,
	},
	"categories": {
	name: { "summary": _summarize(values), "samples": values }
	for name, values in samples.items()
	},
	}
	output = output_dir / "tps.json"
	output.write_text(dumps(record, indent=2))
	if not quiet:
	_render(record["categories"], output)
	if json_output:
	console.print_json(dumps(record))
	raise Exit()

	def _profile_muna_on_modal(
	*,
	tag: str,
	prompts: dict[str, str],
	max_tokens: int,
	rounds: int,
	warmup: int,
	) -> dict[str, list[float]]:
	"""
	Profile decode throughput (TPS) of a compiled Muna model on Modal.

	Runs the benchmark on a B200, looping every category in a single container so
	the model loads only once. The `--warmup` rounds subsume any weight download
	and JIT, so no separate preload step is needed; the volume persists the
	cache across runs.
	"""
	volume = Volume.from_name("muna-benchmark", create_if_missing=True, version=2)
	image = (Image
	.debian_slim(python_version=f"{version_info.major}.{version_info.minor}")
	.uv_pip_install("muna")
	)
	secret = Secret.from_dict({ "MUNA_ACCESS_KEY": environ["MUNA_ACCESS_KEY"] })
	benchmark_app = App("muna-benchmark")
	def _benchmark() -> dict[str, list[float]]:
	from time import monotonic
	muna = Muna()
	results: dict[str, list[float]] = { }
	for name, text in prompts.items():
	tps_samples: list[float] = []
	for index in range(warmup + rounds):
	stream = muna.beta.openai.chat.completions.create(
	model=tag,
	messages=[
	Message(role="user", content=text)
	],
	max_completion_tokens=max_tokens,
	stream=True,
	acceleration="local_gpu"
	)
	t_first = t_last = None
	completion_tokens = 0
	for chunk in stream:
	now = monotonic()
	choice = chunk.choices[0] if chunk.choices else None
	if choice is not None and choice.delta is not None and choice.delta.content:
	if t_first is None:
	t_first = now
	t_last = now
	if chunk.usage is not None:
	completion_tokens = chunk.usage.completion_tokens
	if index < warmup:
	continue
	if t_first is not None and t_last is not None and t_last > t_first and completion_tokens > 1:
	tps_samples.append((completion_tokens - 1) / (t_last - t_first))
	results[name] = tps_samples
	return results
	benchmark_fn = benchmark_app.function(
	image=image,
	serialized=True,
	cpu=32,
	gpu="b200",
	volumes={ "/muna": volume },
	env={ "MUNA_HOME": "/muna" },
	secrets=[secret],
	timeout=3600,
	)(_benchmark)
	with enable_output(), benchmark_app.run():
	return benchmark_fn.remote()

	def _summarize(samples: list[float]) -> dict[str, float \| int]:
	if not samples:
	return { "n": 0 }
	ordered = sorted(samples)
	return {
	"n": len(ordered),
	"min": ordered[0],
	"p50": median(ordered),
	"max": ordered[-1],
	"mean": mean(ordered),
	}

	def _render(categories: dict[str, dict], output: Path) -> None:
	table = Table(show_header=True, header_style="bold cyan", box=None)
	table.add_column("Category", style="bold green")
	table.add_column("n", style="white", justify="right")
	table.add_column("p50 TPS", style="white", justify="right")
	table.add_column("mean TPS", style="white", justify="right")
	for name, payload in categories.items():
	summary = payload["summary"]
	if summary.get("n"):
	table.add_row(name, str(summary["n"]), f"{summary['p50']:.1f}", f"{summary['mean']:.1f}")
	else:
	table.add_row(name, "0", "[red]no samples[/]", "")
	console.print(table)
	console.print(f"[cyan]Decode TPS[/] -> wrote [bold]{output}[/].")

	if __name__ == "__main__":
	run(main)
	#
	# Muna
	# Copyright © 2026 NatML Inc. All Rights Reserved.
	#

	"""
	Time To First Token (TTFT) vs prompt length benchmark for Gemma 4 26B A4B.

	Measures prefill latency at batch size 1: for each target prompt length we
	build a deterministic filler prompt, dispatch a streaming generation, and time
	from the request to the first content token, so:

	ttft = t_first_token - t_request

	We also read the actual `prompt_tokens` from the usage chunk, so the x-axis is
	the real (tokenizer-measured) prompt length rather than the requested length.

	Reproduce with:

	uv run --env-file .env ttft.py
	"""

	# /// script
	# requires-python = ">=3.11"
	# dependencies = ["muna", "modal", "typer"]
	# ///

	from __future__ import annotations
	from json import dumps
	from modal import enable_output, App, Image, Secret, Volume
	from muna import Muna
	from muna.beta.openai import Message
	from os import environ
	from pathlib import Path
	from rich.console import Console
	from rich.table import Table
	from statistics import mean, median
	from time import monotonic
	from typing import Annotated
	from typer import run, Exit, Option

	console = Console()

	DEFAULT_TAG = "@google/gemma-4-26b-a4b-it"
	DEFAULT_PROMPT_LENS = [512, 2048, 4096, 8192, 16384]

	# Deterministic filler so a third party reproduces the same prompt shapes.
	# We size by a rough English chars-per-token heuristic, then report the actual
	# `prompt_tokens` measured by the model's tokenizer.
	_FILLER = (
	"The history of computing stretches from mechanical calculators through "
	"vacuum tubes and transistors to the massively parallel accelerators of "
	"today, where each era introduced new abstractions and new bottlenecks. "
	)
	_CHARS_PER_TOKEN = 4

	def _build_prompt(target_tokens: int) -> str:
	target_chars = max(target_tokens * _CHARS_PER_TOKEN, len(_FILLER))
	repeats = target_chars // len(_FILLER) + 1
	body = (_FILLER * repeats)[:target_chars]
	return f"{body}\n\nIn one word, what is the topic of the passage above?"

	def main(
	tag: Annotated[str, Option(
	"--tag",
	help="Muna predictor tag."
	)] = DEFAULT_TAG,
	prompt_lens: Annotated[list[int], Option(
	"--prompt-len",
	help="Input-token length to sweep; repeatable. Defaults to 512/2k/8k/16k."
	)] = DEFAULT_PROMPT_LENS,
	max_tokens: Annotated[int, Option(
	"--max-tokens",
	min=1,
	help="Tokens to request; TTFT only needs the first token."
	)] = 1,
	rounds: Annotated[int, Option(
	"--rounds", "-r",
	min=1,
	help="Measured requests per prompt length."
	)] = 50,
	warmup: Annotated[int, Option(
	"--warmup",
	min=0,
	help="Discarded warmup requests per prompt length."
	)] = 3,
	output_dir: Annotated[Path, Option(
	"--output-dir",
	help="Directory for results/<metric>.json output."
	)] = Path("results"),
	json_output: Annotated[bool, Option(
	"--json/--no-json",
	help="Also echo raw result JSON to stdout."
	)] = False,
	quiet: Annotated[bool, Option(
	"--quiet", "-q",
	help="Suppress non-essential output."
	)] = False,
	) -> None:
	"""
	Measure TTFT at batch size 1 across a sweep of input prompt lengths.
	"""
	prompts = { str(n): _build_prompt(n) for n in prompt_lens }
	output_dir.mkdir(parents=True, exist_ok=True)
	volume = Volume.from_name("muna-benchmark", create_if_missing=True, version=2)
	image = (Image
	.from_registry("python:3.12-slim-bookworm")
	.uv_pip_install("muna")
	)
	secret = Secret.from_dict({ "MUNA_ACCESS_KEY": environ["MUNA_ACCESS_KEY"] })
	benchmark_app = App("muna-benchmark")
	benchmark_fn = benchmark_app.function(
	image=image,
	cpu=32,
	gpu="b200",
	volumes={ "/muna": volume },
	env={ "MUNA_HOME": "/muna" },
	secrets=[secret],
	timeout=3600,
	)(_benchmark)
	with enable_output(), benchmark_app.run():
	samples = benchmark_fn.remote(prompts, tag, max_tokens, rounds, warmup)
	record = {
	"metric": "ttft",
	"tag": tag,
	"config": {
	"max_tokens": max_tokens,
	"rounds": rounds,
	"warmup": warmup,
	"batch_size": 1,
	},
	"lengths": {
	name: {
	"requested_tokens": int(name),
	"prompt_tokens": payload["prompt_tokens"],
	"summary": _summarize(payload["ttft"]),
	"samples": payload["ttft"],
	}
	for name, payload in samples.items()
	},
	}
	output = output_dir / "ttft.json"
	output.write_text(dumps(record, indent=2))
	if not quiet:
	_render(record["lengths"], output)
	if json_output:
	console.print_json(dumps(record))
	raise Exit()

	def _benchmark(
	prompts: dict[str, str],
	tag: str,
	max_tokens: int,
	rounds: int,
	warmup: int,
	) -> dict[str, dict[str, object]]:
	muna = Muna()
	results: dict[str, dict[str, object]] = { }
	for name, text in prompts.items():
	ttft_samples: list[float] = []
	prompt_tokens = 0
	for index in range(warmup + rounds):
	t_request = monotonic()
	stream = muna.beta.openai.chat.completions.create(
	model=tag,
	messages=[
	Message(role="user", content=text)
	],
	max_completion_tokens=max_tokens,
	stream=True,
	acceleration="local_gpu"
	)
	t_first = None
	for chunk in stream:
	now = monotonic()
	choice = chunk.choices[0] if chunk.choices else None
	if t_first is None and choice is not None and choice.delta is not None and choice.delta.content:
	t_first = now
	if chunk.usage is not None:
	prompt_tokens = chunk.usage.prompt_tokens
	if index < warmup:
	continue
	if t_first is not None:
	ttft_samples.append(t_first - t_request)
	results[name] = { "ttft": ttft_samples, "prompt_tokens": prompt_tokens }
	return results

	def _summarize(samples: list[float]) -> dict[str, float \| int]:
	if not samples:
	return { "n": 0 }
	ordered = sorted(samples)
	return {
	"n": len(ordered),
	"min": ordered[0],
	"p50": median(ordered),
	"max": ordered[-1],
	"mean": mean(ordered),
	}

	def _render(lengths: dict[str, dict], output: Path) -> None:
	table = Table(show_header=True, header_style="bold cyan", box=None)
	table.add_column("Requested", style="bold green", justify="right")
	table.add_column("Prompt tokens", style="white", justify="right")
	table.add_column("n", style="white", justify="right")
	table.add_column("p50 TTFT (ms)", style="white", justify="right")
	table.add_column("mean TTFT (ms)", style="white", justify="right")
	for name, payload in lengths.items():
	summary = payload["summary"]
	if summary.get("n"):
	table.add_row(
	name,
	str(payload["prompt_tokens"]),
	str(summary["n"]),
	f"{summary['p50'] * 1e3:.1f}",
	f"{summary['mean'] * 1e3:.1f}",
	)
	else:
	table.add_row(name, str(payload["prompt_tokens"]), "0", "[red]no samples[/]", "")
	console.print(table)
	console.print(f"[cyan]TTFT[/] -> wrote [bold]{output}[/].")

	if __name__ == "__main__":
	run(main)