[AMD] Use fast exponentials in C4 and C128 ROCm kernels (#34296)

This commit is contained in:
AMD-yanfeiwang
2026-08-26 19:49:01 -07:00
committed by GitHub
parent 1c8f2b38cb
commit cbfe54fba8
2 changed files with 3 additions and 3 deletions
@@ -190,7 +190,7 @@ SGL_DEVICE void c128_forward(
float sum_product = 0.0f;
#pragma unroll
for (int32_t j = 0; j < kElementsPerWarp; ++j) {
const auto exp_score = expf(score[j] - max_value);
const auto exp_score = __expf(score[j] - max_value);
sum_product += cast<float>(kv[j][i]) * exp_score;
sum_exp_value += exp_score;
}
@@ -238,7 +238,7 @@ SGL_DEVICE void c128_forward(
float global_product = 0.0f;
#pragma unroll
for (uint32_t i = 0; i < kNumWarps; ++i) {
const auto exp_val = expf(local_val_max[i] - global_max);
const auto exp_val = __expf(local_val_max[i] - global_max);
global_exp_sum += local_exp_sum[i] * exp_val;
global_product += local_product[i] * exp_val;
}
@@ -206,7 +206,7 @@ SGL_DEVICE void c4_forward(
#pragma unroll
for (int32_t j = 0; j < 8; ++j) {
const auto fp32_score = score[j];
const auto exp_score = expf(fp32_score - max_value);
const auto exp_score = __expf(fp32_score - max_value);
sum_product += cast<float>(kv[j][i]) * exp_score;
sum_exp_value += exp_score;
}