fix: stabilize paged decode attention kernels

- zero-fill split partials so combine skips unwritten splits deterministically
- skip loading masked KV in paged decode kernels to avoid 0*NaN output poisoning
- zero-fill shared memory tile buffers to prevent stale NaN leaking into softmax
This commit is contained in:
2026-07-31 21:01:12 +08:00
parent 7aa5ed09d9
commit 21ddead238
3 changed files with 23 additions and 5 deletions
+10 -3
View File
@@ -67,14 +67,17 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
partial = warp_reduce_sum(partial) * p.scale;
int kv_idx = chunk_start + s;
bool masked = false;
if constexpr (HasMask) {
if (!p.mask[mask_base + kv_idx])
partial = -FLT_MAX;
masked = true;
}
if constexpr (IsCausal) {
if (kv_idx > p.causal_offset)
partial = -FLT_MAX;
masked = true;
}
if (masked)
partial = -FLT_MAX;
float new_m = fmaxf(m, partial);
float alpha = __expf(m - new_m);
@@ -85,7 +88,11 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
int logical_page = pos / p.page_size;
int page_offset = pos % p.page_size;
int phys_page = p.page_table[batch * p.max_pages + logical_page];
if (phys_page >= 0) {
if (masked) {
#pragma unroll
for (int i = 0; i < hd_per_thread; i++)
acc_reg[i] = fmaf(acc_reg[i], alpha, 0.0f);
} else if (phys_page >= 0) {
int64_t v_base = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim
+ (int64_t)page_offset * p.kv_head * p.head_dim
+ (int64_t)kv_head * p.head_dim;