[Bugfix] fix an illegal memory access was encountered of marlin kernel + act_order (#18245)

This commit is contained in:
Jinzhen Lin
2025-05-17 07:02:44 +08:00
committed by GitHub
parent 7fdfa01530
commit e73b7dfd69
3 changed files with 29 additions and 22 deletions

View File

@@ -1588,16 +1588,20 @@ __global__ void Marlin(
if constexpr (has_act_order) {
slice_k_start += tb_k * stages;
slice_k_start_shared_fetch += tb_k * stages;
int first_group_id = g_idx[slice_k_start];
int last_g_idx = slice_k_start + stages * tb_k * 2;
if (last_g_idx >= prob_k) {
last_g_idx = prob_k - 1;
}
int last_group_id = g_idx[last_g_idx];
if (last_group_id >= sh_first_group_id + sh_num_groups) {
fetch_act_order_scales_to_shared(false, first_group_id, last_group_id);
__syncthreads();
if (slice_k_start < prob_k) {
slice_k_start_shared_fetch += tb_k * stages;
int first_group_id = g_idx[slice_k_start];
int last_g_idx = slice_k_start + stages * tb_k * 2;
if (last_g_idx >= prob_k) {
last_g_idx = prob_k - 1;
}
int last_group_id = g_idx[last_g_idx];
if (last_group_id >= sh_first_group_id + sh_num_groups) {
fetch_act_order_scales_to_shared(false, first_group_id,
last_group_id);
__syncthreads();
}
}
}