| """Round 15 (D-fc1 go/no-go ceiling): the served fc1 (library 1D matmul, model config: in0 width-sharded on 8 cores, |
| 8x10 grid, bf16 in0, bfp8 weights, HiFi2 + fp32 acc) with its weights in DRAM (served) vs L1 interleaved (= the upper |
| bound of any weight prefetch / streaming into L1). Also qkv's library config for reference. usage: fc1_l1w.py""" |
| import sys |
| from pathlib import Path |
| from common import * |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) |
| from hamer.vit import tuned_matmul_configs |
| dev = open_dev() |
| torch.manual_seed(0) |
| L1, DR = ttnn.L1_MEMORY_CONFIG, ttnn.DRAM_MEMORY_CONFIG |
| G = dev.compute_with_storage_grid_size() |
| ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, |
| packer_l1_acc=True) |
| try: |
| for name, N, ws, xdt, fg in (("fc1", 5120, 8, ttnn.bfloat16, "8x10"), ("fc1", 5120, 8, ttnn.bfloat16, "12x10")): |
| pc = tuned_matmul_configs(G, fp32_acc=True, in0_ws=ws, fc1_grid=fg)[name] |
| sg = ttnn.num_cores_to_corerangeset(ws, G, row_wise=True) |
| hmem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1, |
| ttnn.ShardSpec(sg, [192, 1280 // ws], ttnn.ShardOrientation.ROW_MAJOR)) |
| x = ttnn.from_torch(torch.randn(1, 1, 192, 1280), dtype=xdt, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=hmem) |
| b = ttnn.from_torch(torch.randn(1, 1, 1, N) * 0.1, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev) |
| wt = torch.randn(1280, N) / 1280 ** 0.5 |
| outs = {} |
| for wloc in ("dram", "l1"): |
| w = ttnn.from_torch(wt, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=DR if wloc == "dram" else L1) |
| fn = lambda w=w: ttnn.linear(x, w, bias=b, program_config=pc, memory_config=L1, compute_kernel_config=ck) |
| o = fn() |
| outs[wloc] = ttnn.to_torch(o).float() |
| ttnn.deallocate(o) |
| print(f"[fc1w] {name} grid {fg} in0 ws{ws} w={wloc:4s} {time_op(dev, fn):7.2f} us", flush=True) |
| ttnn.deallocate(w) |
| print(f"[fc1w] dram == l1 outputs: {bool(torch.equal(outs['dram'], outs['l1']))}", flush=True) |
| ttnn.deallocate(x) |
| finally: |
| ttnn.close_device(dev) |
|
|
