upgrade to fexp_u20 since torch has been updated to 2.9

mingfeima · mingfeima · commit fca9d7becee6 · 2025-12-18T09:19:37.000Z
diff --git a/sgl-kernel/csrc/cpu/extend.cpp b/sgl-kernel/csrc/cpu/extend.cpp
@@ -238,7 +238,7 @@ void extend_attention_kernel_impl(
 
           // s_delta <- exp(s_i - m_i)
           at::vec::map<float>(
-              [m_i](Vec x) { return (x - Vec(m_i)).exp_u20(); }, s_delta + row * BLOCK_N, s_i + row * BLOCK_N, n_size);
+              [m_i](Vec x) { return (x - Vec(m_i)).fexp_u20(); }, s_delta + row * BLOCK_N, s_i + row * BLOCK_N, n_size);
 
           // s' <- s' * m_delta + sum(s_delta)
           s_prime[row] *= m_delta;
@@ -349,7 +349,7 @@ void extend_attention_kernel_impl(
 
           // s_delta <- exp(s_i - m_i)
           at::vec::map<float>(
-              [m_i](Vec x) { return (x - Vec(m_i)).exp_u20(); }, s_delta + row * BLOCK_N, s_i + row * BLOCK_N, n_size);
+              [m_i](Vec x) { return (x - Vec(m_i)).fexp_u20(); }, s_delta + row * BLOCK_N, s_i + row * BLOCK_N, n_size);
 
           // s' <- s' * m_delta + sum(s_delta)
           s_prime[row] *= m_delta;
