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:
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user