"""
flop_counter.py: count the work for one token in ANY decoder model,
using nothing but its shape. FLOPs = 2 x (numbers in every table it uses)
+ 4 x layers x notebook_length x d   (reading the notebook).
"""

def flops_per_token(d, layers, d_ff, vocab, n_ctx, kv_dim=None, mlp_tables=2):
    kv_dim = kv_dim or d                      # smaller if K/V are shared (GQA)
    attn   = d * d + 2 * d * kv_dim + d * d   # make Q, K, V + output table
    mlp    = mlp_tables * d * d_ff            # the thinking-room tables
    parts = {
        "attention tables":   2 * layers * attn,
        "read the notebook":  4 * layers * n_ctx * d,
        "thinking tables":    2 * layers * mlp,
        "score the dictionary": 2 * d * vocab,
    }
    return parts

def show(name, ctxs, **shape):
    print(f"\n{name}")
    for n in ctxs:
        p = flops_per_token(n_ctx=n, **shape)
        total = sum(p.values())
        share = "  ".join(f"{k} {100*v/total:4.1f}%" for k, v in p.items())
        print(f"  notebook {n:>6,}: {total/1e9:6.2f} GFLOPs | {share}")

# Our tiny GPT-2-small-shaped model (should match forward_pass.py)
show("GPT-2 small shape (124M params)", [128, 1024, 8192],
     d=768, layers=12, d_ff=3072, vocab=50257)

# An 8B open model shape: Llama 3.1 8B (GQA: K and V are 1024 wide,
# and its thinking room uses 3 tables instead of 2)
show("Llama 3.1 8B shape", [128, 8192, 32768],
     d=4096, layers=32, d_ff=14336, vocab=128256, kv_dim=1024, mlp_tables=3)