| 1141 | } |
| 1142 | |
| 1143 | void mha_cpu( |
| 1144 | float* xout, // (n_heads, head_dim) |
| 1145 | float* att, // (n_heads, max_seq_len) |
| 1146 | f16_t* kb, // (max_seq_len, n_heads, head_dim) |
| 1147 | f16_t* vb, // (max_seq_len, n_heads, head_dim) |
| 1148 | float* q, // (n_heads, head_dim) |
| 1149 | int head_dim, int v_head_dim, int kv_len, int max_seq_len, int n_heads |
| 1150 | ) { |
| 1151 | // Multihead attention. Iterate over all heads. |
| 1152 | int h; |
| 1153 | #pragma omp parallel for private(h) |
| 1154 | for (h = 0; h < n_heads; h++) { |
| 1155 | int k_head_offset = h * head_dim; |
| 1156 | int v_head_offset = h * v_head_dim; |
| 1157 | f16_t* kh = kb + k_head_offset; |
| 1158 | f16_t* vh = vb + v_head_offset; |
| 1159 | attn( |
| 1160 | xout + head_dim * h, att + max_seq_len * h, q + head_dim * h, |
| 1161 | kh, vh, head_dim, v_head_dim, n_heads, kv_len |
| 1162 | ); |
| 1163 | } |
| 1164 | } |
| 1165 | |
| 1166 | void ffn_cpu( |
| 1167 | float* xout, float* x, |