# Google Colab only: fetch the repo at this post's tag so runs/ is there; does nothing on a laptop or in VS Code.
# Colab already has pandas, numpy, matplotlib, huggingface_hub and tokenizers.
import os, sys, subprocess
if "google.colab" in sys.modules:
    REPO, TAG = "/content/Sillybench", "post-01"
    if not os.path.isdir(REPO):
        subprocess.run(["git", "-c", "advice.detachedHead=false", "clone", "-q", "--depth", "1", "--branch", TAG,
                        "https://github.com/tgarg01/Sillybench.git", REPO], check=True)
    os.chdir(f"{REPO}/notebooks")

import os, json, re, sqlite3
from pathlib import Path
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")   # quiet config/tokenizer downloads
os.environ.setdefault("HF_HUB_VERBOSITY", "error")
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# repo root = the folder that contains runs/
ROOT = Path.cwd().resolve()
while not (ROOT / "runs").is_dir() and ROOT != ROOT.parent:
    ROOT = ROOT.parent
RUN_ID = "2026-10-09_qwen3.8-27b_bf16_initial"   # the run behind the post; or your own run folder from setup.ipynb
RUN = ROOT / "runs" / RUN_ID
CFG = json.loads((RUN / "config.json").read_text())   # the exact vLLM settings of that run

MODEL = CFG["vllm"]["model"]
INPUT_LEN = CFG["workload"]["input_len"]          # 512-token prompt
OUT_TOKENS = CFG["nsys"]["out_tokens"]            # forward passes inside the traced decode range

pd.set_option("display.width", 160)
GPU = CFG.get("observed", {}).get("env", {}).get("gpu", "see the run's env/ folder")
print("run:", RUN.name, "| model:", MODEL, "| GPU:", GPU)

cfg_path = RUN / "analysis" / "model_config.json"      # stored copy of the model's config.json
try:
    from huggingface_hub import hf_hub_download
    cfg_path = Path(hf_hub_download(MODEL, "config.json"))
    print("config.json downloaded from Hugging Face:", MODEL)
except Exception as e:
    print(f"using the stored config.json ({type(e).__name__})")

model_cfg = json.loads(cfg_path.read_text())
tc = model_cfg.get("text_config", model_cfg)
LAYERS, HIDDEN, VOCAB = tc["num_hidden_layers"], tc["hidden_size"], tc["vocab_size"]
print(f"layers: {LAYERS}  |  hidden size: {HIDDEN:,} numbers per token  |  vocabulary: {VOCAB:,} tokens")

kinds = pd.Series(tc["layer_types"])
print(kinds.value_counts().rename("layers").to_string())
print("\nlayer map (F = full attention, . = linear attention):")
row = "".join("F" if t == "full_attention" else "." for t in kinds)
for i in range(0, LAYERS, 16):
    print(f"  layers {i + 1:>2}-{i + 16:<2} {row[i:i + 16]}")

ARTICLE = """Summarize this article in three bullet points:

A graphics processing unit, or GPU, began life as a chip for drawing pictures. Every frame of a video game is made of millions of pixels, and each pixel needs the same small calculation: where is the light coming from, what colour is the surface, how far away is it. A CPU handles work like this one piece at a time, with a few very fast cores. A GPU takes the opposite approach. It has thousands of simpler cores grouped into streaming multiprocessors, and it gives every pixel its own thread, so millions of them run side by side.

For many years this was only useful for graphics. That changed in 2007, when NVIDIA released CUDA, a way to write ordinary programs for the GPU. Scientists started using GPUs for weather models, protein folding and physics simulations, because all of these problems can be split into many identical pieces of work. The real turning point came in 2012, when a neural network called AlexNet, trained on two gaming GPUs, won the ImageNet competition by a wide margin. Neural networks are mostly matrix multiplications, and a matrix multiplication is exactly the kind of work a GPU is built for: the same multiply and add, repeated billions of times on different numbers.

Since then GPUs have been redesigned around this workload. Modern data-centre GPUs such as the H100 contain Tensor Cores, small units that multiply whole blocks of numbers in a single step, and they sit next to very fast memory, because a large model has to read billions of weights for every piece of text it produces. The H100 has 132 streaming multiprocessors, 80 GB of high-bandwidth memory and can move more than three terabytes of data per second. Even so, memory speed is often the real limit: when a language model writes an answer one word at a time, the GPU spends most of its time waiting for weights to arrive rather than doing arithmetic.

Large language models such as Qwen, Llama and GPT are now trained on thousands of GPUs working together, connected by high-speed links so they can share results. Serving these models to millions of users is just as demanding. Engineers batch requests, keep a KV cache and write custom kernels to keep every part of the GPU busy. The design of the chip, the memory system and the software that schedules the work all decide how fast an answer appears on your screen and how much it costs to produce."""

from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer
tok_file = hf_hub_download(MODEL, "tokenizer.json")
ids = Tokenizer.from_file(tok_file).encode(ARTICLE, add_special_tokens=False).ids
vocab = {v: k for k, v in json.loads(Path(tok_file).read_text())["model"]["vocab"].items()}
show = lambda i: f"{i} = {vocab.get(i, '?')!r}"
print(f"prompt: {len(ids)} tokens (raw text, no chat template)")
print("first 3:", ", ".join(show(i) for i in ids[:3]), "| last:", show(ids[-1]))
print(f"\nembedding lookup: each ID picks a row of {HIDDEN:,} numbers -> a {len(ids)} x {HIDDEN:,} matrix")

BENCH = RUN / "bench"
t1 = json.loads((BENCH / "lat_b1_in512_out1.json").read_text())["avg_latency"] * 1e3
t128 = json.loads((BENCH / "lat_b1_in512_out128.json").read_text())["avg_latency"] * 1e3
PREFILL_MS = t1
DECODE_MS = (t128 - t1) / 127
print(f"prompt phase: {PREFILL_MS:.1f} ms for {INPUT_LEN} tokens")
print(f"decode: ({t128:,.0f} - {t1:.1f}) / 127 = {DECODE_MS:.2f} ms per token = {1e3 / DECODE_MS:.1f} tokens/s")
print(f"a 500-token answer = 1 prompt pass + 499 decode passes = {PREFILL_MS + 499 * DECODE_MS:,.0f} ms")

TRACE_DIR = RUN / "nsys"   # written by setup.ipynb (stage "nsys")
print("traces:", sorted(p.name for p in TRACE_DIR.glob("b1_*.sqlite")))

# Kernel families, matched on the kernel name (first match wins). Same rules as notebooks/02_prediction_nsys_analysis.
CATEGORIES = [
    ("GDN conv",            r"causal_conv1d|post_conv"),
    ("GDN recurrent",       r"gated_delta|delta_rule|flashinfergdn|fused_recurrent|chunk_gated|chunk_fwd"),
    ("Full attention",      r"flash::|FlashAttn|flash_fwd|fmha|attention"),
    ("GEMM split-K reduce", r"splitKreduce|splitK_reduce"),
    ("GEMM",                r"nvjet|gemm|xmma|cublas|sm90_.*tst|cutlass.*Gemm"),
    ("RoPE / QK-norm",      r"rope|rotary|qk_rmsnorm"),
    ("Norm",                r"rms_norm|rmsnorm|layer_norm|MeanOps"),
    ("Activation",          r"silu|act_and_mul|gelu"),
    ("Sampling / logits",   r"ArgMax|topk|top_k|sampling|softmax|scatter_gather|sampled|accepted|post_update"),
    ("KV cache write",      r"reshape_and_cache|kv_cache"),
    ("Copy / elementwise",  r"copy|elementwise|Fill|index|CUDAFunctor|memcpy|batch_memcpy"),
]
GENERIC = {"device_kernel", "kernel", "Kernel", "kernel_cutlass_kernel"}   # CUTLASS kernels: use the full signature

def categorize(name):
    for cat, rx in CATEGORIES:
        if re.search(rx, name, re.I):
            return cat
    return "Other"

def load_kernels(sqlite_file):
    """Every GPU kernel in the trace: name, start, end (ns), category; plus the NVTX ranges."""
    con = sqlite3.connect(sqlite_file)
    S = pd.read_sql_query("SELECT id, value FROM StringIds", con).set_index("id")["value"]
    k = pd.read_sql_query("SELECT start, end, shortName, demangledName FROM CUPTI_ACTIVITY_KIND_KERNEL", con)
    nv = pd.read_sql_query("SELECT start, end, text, textId FROM NVTX_EVENTS WHERE end IS NOT NULL", con)
    con.close()
    short, full = k.shortName.map(S), k.demangledName.map(S)
    k["name"] = np.where(short.isin(GENERIC), full.str.replace(r"^void ", "", regex=True).str.slice(0, 110), short)
    k["cat"] = k["name"].map(categorize)
    nv["text"] = nv["text"].fillna(nv["textId"].map(S))
    return k.sort_values("start").reset_index(drop=True), nv

def decode_steps(k, nv):
    """Kernels of the decode steps in the 'prefill_decode' NVTX range, split at each LM-head kernel."""
    r = nv[nv.text.str.startswith("prefill_decode")].iloc[0]
    kk = k[(k.end > r.start) & (k.start < r.end)]
    # the LM head = the GEMM that runs once per forward pass (OUT_TOKENS times) and takes the longest
    g = (kk.end - kk.start).groupby(kk.name).agg(["size", "mean"])
    is_gemm = g.index.str.contains("nvjet|gemm|xmma|cutlass", case=False)
    lm_head = g[is_gemm & g["size"].between(OUT_TOKENS - 1, OUT_TOKENS + 1)]["mean"].idxmax()
    ends = kk[kk.name == lm_head].end.to_numpy()          # end of forward pass 1 (prompt), 2, ..., 64
    dec = kk[(kk.end > ends[0]) & (kk.start < ends[-1])]   # kernels of decode steps 1..63
    return dec, np.diff(ends), lm_head

traces = {}
for mode in ("graphs", "eager"):
    k, nv = load_kernels(TRACE_DIR / f"b1_{mode}.sqlite")
    dec, steps, lm_head = decode_steps(k, nv)
    traces[mode] = dict(dec=dec, steps=steps, lm_head=lm_head, n_steps=len(steps))
    print(f"{mode:>6}: {len(k):,} kernels in the trace, {len(steps)} decode steps, median step {np.median(steps) / 1e6:.2f} ms")

t = traces["graphs"]
dec, n = t["dec"], t["n_steps"]
KERNELS_PER_STEP = len(dec) / n
STEP_MS = np.median(t["steps"]) / 1e6

by_cat = (dec.assign(us=(dec.end - dec.start) / 1e3)
             .groupby("cat").agg(kernels=("us", "size"), us=("us", "sum")))
by_cat["kernels per step"] = (by_cat.kernels / n).round(1)
by_cat["ms per step"] = (by_cat.us / n / 1e3).round(3)
by_cat["% of GPU time"] = (100 * by_cat.us / by_cat.us.sum()).round(1)
by_cat = by_cat.drop(columns=["kernels", "us"]).sort_values("% of GPU time", ascending=False)

gemm = by_cat.loc["GEMM"]
print(f"one decode step: {KERNELS_PER_STEP:.0f} kernels, {STEP_MS:.2f} ms (Nsight median), "
      f"{KERNELS_PER_STEP / LAYERS:.1f} kernels per layer on average")
print(f"GEMMs: {gemm['kernels per step']:.0f} kernels = {gemm['% of GPU time']}% of the GPU time; "
      f"the other {KERNELS_PER_STEP - gemm['kernels per step']:.0f} kernels share the rest")
by_cat

fig, ax = plt.subplots(figsize=(7.5, 3.8))
d = by_cat["% of GPU time"].iloc[::-1]
ax.barh(d.index, d.values, color=["#2E7D4F" if c == "GEMM" else "#9C968C" for c in d.index])
for y, (v, kps) in enumerate(zip(d.values, by_cat["kernels per step"].iloc[::-1])):
    ax.text(v + 1, y, f"{v:.1f}%  ({kps:.0f} kernels)", va="center", fontsize=9)
ax.set(xlim=(0, 115), xlabel="% of GPU time in one decode step",
       title=f"One decode step on 1x NVIDIA H100 GPU: {KERNELS_PER_STEP:.0f} kernels, {STEP_MS:.1f} ms")
ax.spines[["top", "right"]].set_visible(False)
plt.tight_layout(); plt.show()

per_kernel = (dec.assign(us=(dec.end - dec.start) / 1e3)
                 .groupby(["cat", "name"]).agg(calls=("us", "size"), avg_us=("us", "mean"), total_us=("us", "sum"))
                 .reset_index())
per_kernel["per step"] = per_kernel.calls / n
per_kernel["ms per step"] = per_kernel.total_us / n / 1e3
per_kernel["% of step"] = 100 * per_kernel.total_us / per_kernel.total_us.sum()
top = per_kernel.sort_values("ms per step", ascending=False).head(10)
lm = per_kernel[per_kernel.name == t["lm_head"]].iloc[0]
print(f"LM head ({t['lm_head']}): {lm.avg_us / 1e3:.2f} ms per step, the largest single kernel")
top[["cat", "name", "per step", "avg_us", "ms per step", "% of step"]].round(2).reset_index(drop=True)

e = traces["eager"]
EAGER_KERNELS = len(e["dec"]) / e["n_steps"]
EAGER_MS = np.median(e["steps"]) / 1e6
print(f"eager: {EAGER_KERNELS:,.0f} kernels per decode step, {EAGER_MS:.1f} ms per step")
print(f"graphs: {KERNELS_PER_STEP:,.0f} kernels per decode step, {STEP_MS:.1f} ms per step")
print(f"-> {EAGER_KERNELS / KERNELS_PER_STEP:.1f}x the kernels, {EAGER_MS / STEP_MS:.1f}x the time per step")

summary = pd.DataFrame([
    ("model layers", f"{LAYERS}", "config.json"),
    ("hidden size / vocabulary", f"{HIDDEN:,} / {VOCAB:,}", "config.json"),
    ("prompt", f"{len(ids)} tokens -> {len(ids)} x {HIDDEN:,} matrix", "tokenizer.json"),
    ("prompt phase", f"{PREFILL_MS:.1f} ms", "vllm bench latency"),
    ("decode", f"{DECODE_MS:.2f} ms per token = {1e3 / DECODE_MS:.1f} tokens/s", "vllm bench latency"),
    ("kernels per decode step", f"{KERNELS_PER_STEP:.0f}", "Nsight Systems, graphs"),
    ("GEMMs", f"{by_cat.loc['GEMM', 'kernels per step']:.0f} kernels, {by_cat.loc['GEMM', '% of GPU time']}% of GPU time", "Nsight Systems, graphs"),
    ("eager (no CUDA graphs)", f"{EAGER_KERNELS:,.0f} kernels per step", "Nsight Systems, eager"),
    ("largest kernel", f"LM head, {lm.avg_us / 1e3:.2f} ms", "Nsight Systems, graphs"),
], columns=["number", "value", "from"])
print("run:", RUN.name, "| GPU:", GPU)
summary