"""Lab 1: the speed of light (roofline) for every scored shape, and the score formula. Run: python3 lab1_roofline.py Needs: Python 3 only (no GPU, no numpy). What it does 1. For each of the 30 scored shapes it computes the bytes that must move and the FLOPs that must be done, and the "floor" = max(bytes / 4.8 TB/s, FLOPs / 989 TFLOP/s), exactly as reference.py's speed_of_light_ns() does. 2. It prints the official times of SGLang and of our submission (check run 7e998918) next to the floor. 3. It recomputes the official score (1.343x) from those per-shape times. Things to try (edit and re-run) - Change OURS[("decode 32", "qkv_gate")] to 10.0 and see how much the score moves. - Set WEIGHT_BITS = 12 to see the floor if every weight were stored in 12 bits. - Find the arithmetic intensity at which the math term overtakes the bytes term (the "ridge"). """ import math HBM = 4.8e12 # bytes per second: H200 HBM3e bandwidth, the judge's number TC = 989e12 # FLOP per second: H200 dense BF16 tensor-core peak, the judge's number WEIGHT_BITS = None # None = as stored (16 bits, router 32 bits); try 12 # name: (N rows, K columns, bytes per weight) LAYERS = { "qkv_gate": (2304, 8192, 2), "o": (8192, 1024, 2), "pooler": (4096, 8192, 2), "unpooler": (8192, 4096, 2), "router": (256, 8192, 4), # FP32 weights and FP32 output } LOADS = [("decode 1", 1), ("decode 16", 16), ("decode 32", 32), ("verify 1", 8), ("verify 16", 128), ("verify 32", 256)] # Official judge times in microseconds (check run 7e998918, the submission as shipped). SGLANG = { ("decode 1", "qkv_gate"): 14.18, ("decode 1", "o"): 6.82, ("decode 1", "pooler"): 21.31, ("decode 1", "unpooler"): 7.07, ("decode 1", "router"): 10.11, ("decode 16", "qkv_gate"): 14.16, ("decode 16", "o"): 6.88, ("decode 16", "pooler"): 21.81, ("decode 16", "unpooler"): 18.26, ("decode 16", "router"): 14.83, ("decode 32", "qkv_gate"): 13.82, ("decode 32", "o"): 6.66, ("decode 32", "pooler"): 21.89, ("decode 32", "unpooler"): 18.94, ("decode 32", "router"): 15.15, ("verify 1", "qkv_gate"): 14.21, ("verify 1", "o"): 7.01, ("verify 1", "pooler"): 21.46, ("verify 1", "unpooler"): 7.14, ("verify 1", "router"): 26.08, ("verify 16", "qkv_gate"): 17.09, ("verify 16", "o"): 7.44, ("verify 16", "pooler"): 25.44, ("verify 16", "unpooler"): 19.95, ("verify 16", "router"): 16.42, ("verify 32", "qkv_gate"): 19.38, ("verify 32", "o"): 9.41, ("verify 32", "pooler"): 30.32, ("verify 32", "unpooler"): 24.13, ("verify 32", "router"): 17.60, } OURS = { ("decode 1", "qkv_gate"): 9.63, ("decode 1", "o"): 5.22, ("decode 1", "pooler"): 13.77, ("decode 1", "unpooler"): 4.20, ("decode 1", "router"): 5.43, ("decode 16", "qkv_gate"): 12.19, ("decode 16", "o"): 5.96, ("decode 16", "pooler"): 15.69, ("decode 16", "unpooler"): 14.29, ("decode 16", "router"): 5.90, ("decode 32", "qkv_gate"): 13.85, ("decode 32", "o"): 6.68, ("decode 32", "pooler"): 19.33, ("decode 32", "unpooler"): 17.77, ("decode 32", "router"): 6.14, ("verify 1", "qkv_gate"): 9.94, ("verify 1", "o"): 5.82, ("verify 1", "pooler"): 14.12, ("verify 1", "unpooler"): 4.26, ("verify 1", "router"): 5.84, ("verify 16", "qkv_gate"): 17.14, ("verify 16", "o"): 7.57, ("verify 16", "pooler"): 23.33, ("verify 16", "unpooler"): 20.01, ("verify 16", "router"): 9.86, ("verify 32", "qkv_gate"): 19.49, ("verify 32", "o"): 9.51, ("verify 32", "pooler"): 31.08, ("verify 32", "unpooler"): 24.24, ("verify 32", "router"): 10.48, } def out_features(layer, m): """N actually computed: at <= 8 tokens the unpooler computes only this GPU's 1024-row slice.""" return 1024 if layer == "unpooler" and m <= 8 else LAYERS[layer][0] def floor(layer, m, weight_bits=None): """(floor in us, bytes moved, flops) -- the same formula as reference.speed_of_light_ns.""" _, k, wb = LAYERS[layer] n = out_features(layer, m) w_bytes = n * k * (wb if weight_bits is None else weight_bits / 8) bytes_moved = w_bytes + m * k * 2 + m * n * wb # weights + x (BF16) + output flops = 2.0 * m * n * k # one multiply + one add per weight per token return 1e6 * max(bytes_moved / HBM, flops / TC), bytes_moved, flops def score(ours): """Geometric mean over the 6 loads of (SGLang load time / our load time). A load's time is the sum of its 5 shapes (each shape's weight is 62 in every load, so it cancels).""" logs = [] for load, _ in LOADS: sg = sum(SGLANG[(load, layer)] for layer in LAYERS) us = sum(ours[(load, layer)] for layer in LAYERS) logs.append(math.log(sg / us)) return math.exp(sum(logs) / len(logs)) def main(): print(f"ridge point: {TC / HBM:.0f} FLOP per byte (below it a kernel is limited by memory, above it by math)\n") hdr = f"{'load':<10}{'layer':<10}{'m':>4}{'MB':>8}{'GFLOP':>8}{'FLOP/B':>8}{'floor':>8} {'bound':<7}{'SGLang':>8}{'ours':>8}{'ours/floor':>11}{'speedup':>9}" print(hdr) print("-" * len(hdr)) for load, m in LOADS: for layer in LAYERS: fl, b, f = floor(layer, m, WEIGHT_BITS) bound = "memory" if b / HBM >= f / TC else "math" sg, us = SGLANG[(load, layer)], OURS[(load, layer)] print(f"{load:<10}{layer:<10}{m:>4}{b / 1e6:>8.1f}{f / 1e9:>8.2f}{f / b:>8.1f}{fl:>7.2f}u {bound:<7}" f"{sg:>7.2f}u{us:>7.2f}u{us / fl:>10.2f}x{sg / us:>8.3f}x") print() print(f"{'load':<10}{'SGLang us':>10}{'ours us':>10}{'floor us':>10}{'speedup':>9}") for load, m in LOADS: sg = sum(SGLANG[(load, layer)] for layer in LAYERS) us = sum(OURS[(load, layer)] for layer in LAYERS) fl = sum(floor(layer, m, WEIGHT_BITS)[0] for layer in LAYERS) print(f"{load:<10}{sg:>10.2f}{us:>10.2f}{fl:>10.2f}{sg / us:>8.3f}x") print(f"\nscore = geometric mean of the 6 load speedups = {score(OURS):.4f}x (official check: 1.3428x)") # A what-if: make decode 32's qkv_gate and o 1.3x faster than SGLang (they are at ~1.0x today). what_if = dict(OURS) for layer in ("qkv_gate", "o"): what_if[("decode 32", layer)] = SGLANG[("decode 32", layer)] / 1.3 print(f"what-if (decode 32 qkv_gate and o at 1.3x): score = {score(what_if):.4f}x") if __name__ == "__main__": main()