Skip to content

Commit 5275074

Browse files
whisper : fix DTW memory access (ggml-org#2012)
* Fix DTW memory access * Memory fix - Apply changes from denersc
1 parent c15b4cd commit 5275074

File tree

1 file changed

+15
-10
lines changed

1 file changed

+15
-10
lines changed

whisper.cpp

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1148,26 +1148,31 @@ static bool aheads_masks_init(
11481148
}
11491149

11501150
// 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
11541160
std::vector<float> mask_data;
11551161
for (int64_t il = 0; il < n_text_layer; ++il) {
11561162
if (aheads_masks.m[il] != nullptr) {
11571163
auto aheads = get_alignment_heads_by_layer(cparams, il, n_text_layer, n_head);
11581164

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);
11601167
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);
11631168

1169+
std::fill(mask_data.begin(), mask_data.end(), 0);
11641170
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;
11681173
}
11691174

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);
11711176
}
11721177
}
11731178

0 commit comments

Comments
 (0)