File size: 2,304 Bytes
58ba6df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | """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)
|