Skip to content

GitLab

Commit 50e2de8d authored by zhanghj2's avatar zhanghj2
Browse files

处理lse为正无穷的情况

parent 5d8e93f6
Show whitespace changes
Inline Side-by-side
Showing with 6 additions and 3 deletions
csrc/smxx/decode/combine/combine.cu
View file @ 50e2de8d
...@@ -95,11 +95,14 @@ flash_fwd_mla_combine_kernel(const CombineParams params) { ...@@ -95,11 +95,14 @@ flash_fwd_mla_combine_kernel(const CombineParams params) {
o_scale = 0.0f; o_scale = 0.0f;
} }
else else
{
if (!flash::is_positive_infinity(global_lse))
{ {
float Attn_sink_exp2 = __builtin_amdgcn_exp2f(attn_sink * CUDART_L2E_F); float Attn_sink_exp2 = __builtin_amdgcn_exp2f(attn_sink * CUDART_L2E_F);
float lse_exp2 = __builtin_amdgcn_exp2f(global_lse * CUDART_L2E_F); float lse_exp2 = __builtin_amdgcn_exp2f(global_lse * CUDART_L2E_F);
o_scale = lse_exp2 / (lse_exp2 + Attn_sink_exp2); o_scale = lse_exp2 / (lse_exp2 + Attn_sink_exp2);
} }
}
// if (global_lse != INFINITY) { // if (global_lse != INFINITY) {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment