"""
bus_test.py: push tokens through 36 different 768 x 3072 tables (340 MB, the
same size as all 12 floors of our model), first 1 token at a time, then
many tokens per trip. Same tables, same CPU. How fast does it really work?
"""
import time
import numpy as np

rng = np.random.default_rng(0)
tables = [rng.standard_normal((768, 3072)).astype(np.float32) * 0.02
          for _ in range(36)]                     # 36 x 9.4 MB = 340 MB
print(f"tables: {sum(t.nbytes for t in tables)/1e6:.0f} MB")

def trip(x):
    for W in tables:                              # one trip through every table
        y = x @ W
    return y

for batch in [1, 16, 64, 256]:
    x = rng.standard_normal((batch, 768)).astype(np.float32)
    trip(x)                                       # warm-up
    reps = 10 if batch < 64 else 5
    t = time.perf_counter()
    for _ in range(reps):
        trip(x)
    sec = (time.perf_counter() - t) / reps
    flops = 2 * batch * sum(W.size for W in tables)
    print(f"{batch:4d} tokens per trip: {sec*1e3:7.1f} ms per trip, "
          f"{sec/batch*1e3:6.2f} ms per token, {flops/sec/1e9:6.1f} GFLOP/s")