@@ -1148,26 +1148,31 @@ static bool aheads_masks_init(
1148
1148
}
1149
1149
1150
1150
// Set data on mask tensors
1151
- // Since this must be backend agnostic, we get tensor data with
1152
- // ggml_backend_tensor_get, copy our desired values and send it back
1153
- // to backend with ggml_backend_tensor_set
1151
+ // Since this must be backend agnostic, we write our desired values on mask_data,
1152
+ // and send it to backend with ggml_backend_tensor_set.
1153
+ // Each mask in N_HEADS*N_ALIGNMENT_HEADS, one per text layer containing alignment
1154
+ // heads. Each row of the mask "marks" one alignment head. E.g. if some text layer
1155
+ // has a total of 10 heads and of those, heads 0,5,6 are alignment heads, the mask
1156
+ // should read:
1157
+ // 1 0 0 0 0 0 0 0 0 0
1158
+ // 0 0 0 0 0 1 0 0 0 0
1159
+ // 0 0 0 0 0 0 1 0 0 0
1154
1160
std::vector<float > mask_data;
1155
1161
for (int64_t il = 0 ; il < n_text_layer; ++il) {
1156
1162
if (aheads_masks.m [il] != nullptr ) {
1157
1163
auto aheads = get_alignment_heads_by_layer (cparams, il, n_text_layer, n_head);
1158
1164
1159
- size_t data_size = aheads_masks.m [il]->ne [0 ] * aheads_masks.m [il]->ne [1 ] * sizeof (float );
1165
+ size_t data_size = aheads_masks.m [il]->ne [0 ] * aheads_masks.m [il]->ne [1 ];
1166
+ size_t data_size_bytes = data_size * sizeof (float );
1160
1167
mask_data.resize (data_size);
1161
- ggml_backend_tensor_get (aheads_masks.m [il], mask_data.data (), 0 , data_size);
1162
- memset (mask_data.data (), 0 , data_size);
1163
1168
1169
+ std::fill (mask_data.begin (), mask_data.end (), 0 );
1164
1170
for (size_t ih = 0 ; ih < aheads.size (); ++ih) {
1165
- size_t pos = (aheads[ih] + (ih * aheads_masks.m [il]->ne [0 ] * aheads[ih]));
1166
- float v = 1 .0f ;
1167
- memcpy (mask_data.data () + pos, &v, sizeof (float ));
1171
+ size_t pos = (aheads[ih] + (ih * aheads_masks.m [il]->ne [0 ]));
1172
+ mask_data[pos] = 1 .0f ;
1168
1173
}
1169
1174
1170
- ggml_backend_tensor_set (aheads_masks.m [il], mask_data.data (), 0 , data_size );
1175
+ ggml_backend_tensor_set (aheads_masks.m [il], mask_data.data (), 0 , data_size_bytes );
1171
1176
}
1172
1177
}
1173
1178
0 commit comments