ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: improve attention kernel numerical stability and test precision checks

Browse Source 
- use fmaf() for V-accumulation in scalar decode paths to reduce rounding
- delay scale multiplication to after dot-product in scalar prefill
- unify __expf/expf across MMA and scalar paths for consistent numerics
- harmonize divide-by-zero guards to 1e-20f
- add both absolute and relative error checks in standalone tests (atol=0.01, rtol=0.01)
This commit is contained in:
ViperEkura
2026-07-21 23:05:17 +08:00
parent f7a16efc9d
commit 60d7ee614a
6 changed files with 57 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/attn_paged_decode_split_kv.cuh
+7 -7
View File
@@ -98,12 +98,12 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
+ (int64_t)kv_head * p.head_dim;
#pragma unroll
for (int i = 0; i < hd_per_thread; i++)
acc_reg[i] = acc_reg[i] * alpha
+ __bfloat162float(p.v_cache[v_base + lane * hd_per_thread + i]) * beta;
acc_reg[i] = fmaf(acc_reg[i], alpha,
__bfloat162float(p.v_cache[v_base + lane * hd_per_thread + i]) * beta);
} else {
#pragma unroll
for (int i = 0; i < hd_per_thread; i++)
acc_reg[i] = acc_reg[i] * alpha + 0.0f * beta;
acc_reg[i] = fmaf(acc_reg[i], alpha, 0.0f);
}
m = new_m;
}
@@ -140,10 +140,10 @@ __global__ void paged_attn_decode_combine_kernel(PagedAttentionParams<bf16> p) {
if (mi <= -FLT_MAX) continue;
float li = mlp[s * 2 + 1];
float nm = fmaxf(m, mi);
float corr = __expf(m - nm);
float e = __expf(mi - nm);
acc = acc * corr + op[s * p.head_dim + d] * e;
l = l * corr + li * e;
float corr = expf(m - nm);
float e = expf(mi - nm);
acc = fmaf(acc, corr, op[s * p.head_dim + d] * e);
l = fmaf(l, corr, li * e);
m = nm;
}