处理lse为正无穷的情况

50e2de8d · zhanghj2 · 5d8e93f6 · 50e2de8d
Commit 50e2de8d authored Jan 28, 2026 by zhanghj2
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 3 deletions

csrc/smxx/decode/combine/combine.cu csrc/smxx/decode/combine/combine.cu +6 -3

No files found.
--- a/csrc/smxx/decode/combine/combine.cu
+++ b/csrc/smxx/decode/combine/combine.cu
@@ -96,9 +96,12 @@ flash_fwd_mla_combine_kernel(const CombineParams params) {
            }
            else
            {
-                float Attn_sink_exp2 = __builtin_amdgcn_exp2f(attn_sink * CUDART_L2E_F);
+                if (!flash::is_positive_infinity(global_lse))
-                float lse_exp2 = __builtin_amdgcn_exp2f(global_lse * CUDART_L2E_F);
+                {
-                o_scale = lse_exp2 / (lse_exp2 + Attn_sink_exp2);
+                    float Attn_sink_exp2 = __builtin_amdgcn_exp2f(attn_sink * CUDART_L2E_F);
+                    float lse_exp2 = __builtin_amdgcn_exp2f(global_lse * CUDART_L2E_F);
+                    o_scale = lse_exp2 / (lse_exp2 + Attn_sink_exp2);
+                }
            }