From 08329b3a717229269a814531ecf533226a8808f9 Mon Sep 17 00:00:00 2001 From: Prashant Rawat Date: Mon, 24 Aug 2026 07:55:16 -0700 Subject: [PATCH] SDPA: use exp_u20 for the softmax exponential (#22082) Summary: Switch the softmax exponential in the flash-attention custom SDPA kernel (`op_sdpa_impl.h`) from `Vectorized::exp()` to `Vectorized::exp_u20()`. The file already carried this as a TODO. This is not bit-exact. `exp_u20` is a ULP-20 approximation, so it differs from `exp` outright, and fourteen autoregressive layers on that is enough to flip a near-tie argmax. Differential Revision: D117198987 --- extension/llm/custom_ops/op_sdpa_impl.h | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/extension/llm/custom_ops/op_sdpa_impl.h b/extension/llm/custom_ops/op_sdpa_impl.h index 9c479569505..dbf91b7d223 100644 --- a/extension/llm/custom_ops/op_sdpa_impl.h +++ b/extension/llm/custom_ops/op_sdpa_impl.h @@ -443,9 +443,7 @@ _exp_reduce_sum_fusion_kernel(T1* a, const int& size, T2* out, T1& val) { for (int i = 0; i < vec_size * (size / vec_size); i += vec_size) { auto tmp0 = vec::VectorizedN::loadu(a + i); auto tmp1 = tmp0 - vec_max; - // Replace with exp_u20 later - // auto tmp2 = tmp1.exp_u20(); - auto tmp2 = tmp1.exp(); + auto tmp2 = tmp1.exp_u20(); vec_tmp_sum = vec_tmp_sum + tmp2; tmp2.store(out + i); }