()
| 27 | return args |
| 28 | |
| 29 | def main(): |
| 30 | args = parse_args() |
| 31 | config = Config() |
| 32 | |
| 33 | if not os.path.exists(args.config): |
| 34 | print("Host: Use default test values.") |
| 35 | else: |
| 36 | with open(args.config) as f: |
| 37 | config.__dict__.update(json.load(f)) |
| 38 | |
| 39 | cfg_name = os.path.splitext(os.path.basename(args.config))[0] |
| 40 | |
| 41 | P = config.P |
| 42 | bsz = config.bsz |
| 43 | group_num = config.group_num |
| 44 | dim = config.dim |
| 45 | n_heads = config.n_heads |
| 46 | n_kv_heads = config.n_kv_heads |
| 47 | head_dim = config.head_dim |
| 48 | seq_len = config.seq_len |
| 49 | ffn_dim = config.ffn_dim |
| 50 | |
| 51 | dim_p_pe = dim // P |
| 52 | pes_p_head = P // n_heads |
| 53 | pes_p_kv_head = P // n_kv_heads |
| 54 | head_dim_p_pe = head_dim // P |
| 55 | seq_len_p_pe = seq_len // P |
| 56 | ffn_dim_p_pe = ffn_dim // P |
| 57 | |
| 58 | print(f"Host: P: {P}, Batch size: {bsz}, dim_p_pe: {dim_p_pe}, pes_p_head: {pes_p_head}, pes_p_kv_head: {pes_p_kv_head}, head_dim_p_pe: {head_dim_p_pe}, seq_len_p_pe: {seq_len_p_pe}, ffn_dim_p_pe: {ffn_dim_p_pe}") |
| 59 | |
| 60 | io_dtype = MemcpyDataType.MEMCPY_16BIT |
| 61 | memcpy_order = MemcpyOrder.ROW_MAJOR |
| 62 | |
| 63 | X = np.random.rand(1, bsz*dim).astype(np.float16) |
| 64 | tensor_X = np.tile(X.reshape(P, bsz*dim_p_pe), reps=(1, P)) |
| 65 | |
| 66 | W = np.random.rand(1, dim).astype(np.float16) |
| 67 | tensor_W = np.tile(W.reshape(P, dim_p_pe), reps=(1, P)) |
| 68 | |
| 69 | tensor_q_weight = np.random.rand(dim, dim).astype(np.float16) |
| 70 | tensor_k_weight = np.random.rand(dim, dim).astype(np.float16) |
| 71 | tensor_v_weight = np.random.rand(dim, dim).astype(np.float16) |
| 72 | |
| 73 | _dim_p_pe = dim_p_pe |
| 74 | if (dim_p_pe % 2) == 1: |
| 75 | _dim_p_pe = dim_p_pe - 1 |
| 76 | |
| 77 | freqs_sin = np.random.rand(1, P*_dim_p_pe//2).astype(np.float16) |
| 78 | tensor_freqs_sin = np.tile(freqs_sin.reshape(P, _dim_p_pe//2), reps=(1, P)) |
| 79 | freqs_cos = np.random.rand(1, P*_dim_p_pe//2).astype(np.float16) |
| 80 | tensor_freqs_cos = np.tile(freqs_cos.reshape(P, _dim_p_pe//2), reps=(1, P)) |
| 81 | |
| 82 | tensor_XKCache = np.random.rand(dim, seq_len).astype(np.float16) |
| 83 | tensor_XVCache = np.random.rand(seq_len, dim).astype(np.float16) |
| 84 | |
| 85 | tensor_o_weight = np.random.rand(dim, dim).astype(np.float16) |
| 86 | tensor_up_weight = np.random.rand(dim, ffn_dim).astype(np.float16) |
no test coverage detected