Skip to content

Commit

Permalink
clean
Browse files Browse the repository at this point in the history
  • Loading branch information
zaixing-wang committed Jan 6, 2025
1 parent 8ec29cc commit 5d59e33
Showing 1 changed file with 0 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -1377,7 +1377,6 @@ KERNEL(sdpa_opt_finalization_stage)(
local_max_logit = sub_group_reduce_max(local_max_logit);

// Update exp_sum with respect to the global maximum
// SOFTMAX_ACCUMULATOR_TYPE exp_sum[REG_VERSION_MAX_VALUES_PER_WI_LOWER] = {SOFTMAX_ACCUMULATOR_VAL_ZERO};
SOFTMAX_ACCUMULATOR_TYPE local_exp_sum = SOFTMAX_ACCUMULATOR_VAL_ZERO;
for (uint i = local_id; i < num_of_partitions; i+= HEAD_SIZE) {
SOFTMAX_ACCUMULATOR_TYPE exp_sum_new = cur_exp_sums[i] * native_exp(max_logits_u_exp_sum_slm[i] - local_max_logit);
Expand Down Expand Up @@ -1428,7 +1427,6 @@ KERNEL(sdpa_opt_finalization_stage)(
local_max_logit = sub_group_reduce_max(local_max_logit);

// Update exp_sum with respect to the global maximum
// SOFTMAX_ACCUMULATOR_TYPE exp_sum[REG_VERSION_MAX_VALUES_PER_WI_LOWER] = {SOFTMAX_ACCUMULATOR_VAL_ZERO};
SOFTMAX_ACCUMULATOR_TYPE local_exp_sum = SOFTMAX_ACCUMULATOR_VAL_ZERO;
for (uint i = local_id; i < num_of_partitions; i+= HEAD_SIZE) {
local_exp_sum += cur_exp_sums[i] * native_exp(cur_max_logits[i] - local_max_logit);
Expand All @@ -1450,7 +1448,6 @@ KERNEL(sdpa_opt_finalization_stage)(
b1_idx * (TARGET_SEQ_LEN * num_of_partitions * HEAD_SIZE) +
target_seq_idx * (num_of_partitions * HEAD_SIZE) +
partition_idx * (HEAD_SIZE) + local_id;
// (head_size_idx * SUBGROUP_SIZE + sglid);
OUTPUT_TYPE out_val = tmp_out[tmp_out_offset];
acc += TO_SOFTMAX_ACCUMULATOR_TYPE(tmp_out[tmp_out_offset]) * TO_SOFTMAX_ACCUMULATOR_TYPE(cur_exp_sums[partition_idx] * native_exp(cur_max_logits[partition_idx] - local_max_logit));
}
Expand Down

0 comments on commit 5d59e33

Please sign in to comment.