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

fix: stabilize paged decode attention kernels

Browse Source 
- 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:
ViperEkura
2026-07-31 21:01:12 +08:00
parent 7aa5ed09d9
commit 21ddead238
3 changed files with 23 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/attn_entry_utils.cuh
+3 -2
View File
@@ -1,4 +1,5 @@
#pragma once
#include <float.h>
#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include "attn_common.h"
@@ -23,8 +24,8 @@ using bf16 = __nv_bfloat16;
template<typename P>
inline void alloc_split_partials(P& p) {
auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
auto o_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
auto ml_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, fopt);
auto o_part = torch::zeros(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
auto ml_part = torch::full(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, -FLT_MAX, fopt);
p.o_part = (float*)o_part.data_ptr();
p.ml_part = (float*)ml_part.data_ptr();
}