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
@@ -31,6 +31,13 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
__shared__ __align__(16) bf16 sK[Traits::STAGES * Traits::BC * Traits::LD];
__shared__ __align__(16) bf16 sV[Traits::STAGES * Traits::BC * Traits::LD];
#pragma unroll
for (int i = lane; i < Traits::STAGES * Traits::BC * Traits::LD; i += 32) {
sK[i] = __float2bfloat16(0.0f);
sV[i] = __float2bfloat16(0.0f);
}
__syncwarp();
const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
const int qra = gid;
const int qrb = gid + 8;
@@ -68,6 +75,9 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
int kc = kv0 + r;
bool valid = (kc < p.kv_len);
if constexpr (HasMask) {
valid = valid && p.mask[batch * p.mask_b_stride + kc];
}
int phys_page = valid ? p.page_table[batch * p.max_pages + kc] : 0;
valid = valid && (phys_page >= 0);
int page_off = kc % p.page_size;