"""
predict.py: use the roofline to predict decode speed limits BEFORE running anything.
Ceilings only: real engines land below them, and the KV cache (notebooks) adds more bytes.
"""
CHIPS = {                        # (memory roof, math roof): dense peaks, no sparsity
    "my CPU (float32)": (15.0e9,   164.3e9),
    "H100 SXM (BF16)":  (3.35e12,  989e12),
    "B200 (BF16)":      (8.0e12,  2250e12),
}

def step(params, bytes_per_number, tokens, bw, peak):
    """One trip through the model with `tokens` tokens on it (weights only)."""
    t_fetch = params * bytes_per_number / bw          # forklift
    t_math  = 2 * params * tokens / peak              # workbench
    return max(t_fetch, t_math), t_fetch, t_math

print("ridge point and tokens per trip needed to keep the math busy")
for name, (bw, peak) in CHIPS.items():
    bpn = 4 if "float32" in name else 2
    ridge = peak / bw
    print(f"  {name:<18} ridge {ridge:6.1f} FLOPs/byte -> about {ridge * bpn / 2:4.0f} tokens per trip")

print("\nan 8-billion-parameter model in BF16 (16 GB of weights) on one H100")
bw, peak = CHIPS["H100 SXM (BF16)"]
print(f"  {'tokens':>6} {'step ms':>8} {'fetch ms':>9} {'math ms':>8} {'tokens/s':>9}  limit")
for b in [1, 8, 64, 256, 512, 1024]:
    t, f, m = step(8e9, 2, b, bw, peak)
    print(f"  {b:>6} {t*1e3:8.2f} {f*1e3:9.2f} {m*1e3:8.2f} {b/t:9,.0f}  {'memory' if f >= m else 'math'}")

print("\nsame model, 1 user, different number sizes on the H100")
for label, bpn, pk in [("BF16, 2 bytes", 2, 989e12), ("FP8, 1 byte", 1, 1979e12)]:
    t, f, m = step(8e9, bpn, 1, bw, pk)
    print(f"  {label:<14} {t*1e3:5.2f} ms per token -> at most {1/t:4.0f} tokens/s;"
          f" ridge needs ~{(pk/bw)*bpn/2:.0f} tokens per trip")