| 110 | }; |
| 111 | |
| 112 | void createTransformer(Context &ctx, size_t modelDim, size_t qkvDim, |
| 113 | size_t nHeads, size_t batchSize, size_t seqLen, |
| 114 | size_t hiddenWidth, Transformer &transformer, |
| 115 | Activations &activations, KVCache &kvCache) { |
| 116 | std::mt19937 gen(314159); |
| 117 | transformer = { |
| 118 | .qkv = createTensor(ctx, Shape{3 * nHeads * qkvDim, modelDim}, |
| 119 | kf32), // column-major |
| 120 | .rmsNormPre = createTensor(ctx, Shape{modelDim}, kf32), |
| 121 | .rmsNormPost = createTensor(ctx, Shape{modelDim}, kf32), |
| 122 | .out = createTensor(ctx, Shape{3 * qkvDim, modelDim}, kf32), |
| 123 | .mlp1 = createTensor(ctx, Shape{modelDim, 2 * hiddenWidth}, kf32), |
| 124 | .mlp2 = createTensor(ctx, Shape{modelDim, 2 * hiddenWidth}, kf32), |
| 125 | }; |
| 126 | |
| 127 | // Initialize values |
| 128 | std::unique_ptr<float[]> qkvInit(new float[modelDim * 3 * nHeads * qkvDim]); |
| 129 | // randint(qkvInit.get(), size(transformer.qkv.shape), gen, -2, 2); |
| 130 | range(qkvInit.get(), size(transformer.qkv.shape), 0.0); |
| 131 | LOG(kDefLog, kInfo, "%s", |
| 132 | show<float>(qkvInit.get(), transformer.qkv.shape[0], |
| 133 | transformer.qkv.shape[1], "QKV Weights") |
| 134 | .c_str()); |
| 135 | toGPU(ctx, qkvInit.get(), transformer.qkv); |
| 136 | |
| 137 | activations = { |
| 138 | .qkv = createTensor(ctx, Shape{batchSize * 3 * nHeads * qkvDim}, kf32), |
| 139 | .qk = createTensor(ctx, Shape{batchSize * nHeads}, kf32), |
| 140 | .att = createTensor(ctx, Shape{batchSize * nHeads}, kf32)}; |
| 141 | |
| 142 | kvCache = { |
| 143 | .keyCache = createTensor(ctx, Shape{seqLen, qkvDim}, kf32), |
| 144 | .valueCache = createTensor(ctx, Shape{seqLen, qkvDim}, kf32), |
| 145 | }; |
| 146 | std::unique_ptr<float[]> keyCacheInit(new float[seqLen * qkvDim]); |
| 147 | std::unique_ptr<float[]> valueCacheInit(new float[seqLen * qkvDim]); |
| 148 | range(keyCacheInit.get(), size(kvCache.keyCache.shape), 0.0); |
| 149 | range(valueCacheInit.get(), size(kvCache.valueCache.shape), 0.0); |
| 150 | toGPU(ctx, keyCacheInit.get(), kvCache.keyCache); |
| 151 | toGPU(ctx, valueCacheInit.get(), kvCache.valueCache); |
| 152 | } |
| 153 | |
| 154 | inline KernelCode createMatmul(const char *shaderTemplate, const size_t M, |
| 155 | const size_t K, const size_t N, |