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

fix: use c10::optional for o_part_buf/ml_part_buf decode kernel params

Browse Source
This commit is contained in:
ViperEkura
2026-08-06 20:50:48 +08:00
parent 6054b8dbd4
commit a59ae8f32e
2 changed files with 20 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/attn_decode.cu
+10 -9
View File
@@ -9,8 +9,8 @@ torch::Tensor attn_decode(
int64_t causal_offset,
double scale,
int64_t layout,
torch::Tensor o_part_buf,
torch::Tensor ml_part_buf
c10::optional<torch::Tensor> o_part_buf,
c10::optional<torch::Tensor> ml_part_buf
) {
const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
auto stream = at::cuda::getCurrentCUDAStream();
@@ -24,14 +24,15 @@ torch::Tensor attn_decode(
auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O;
p.o = (bf16*)O_view.data_ptr();
if (o_part_buf.defined() && ml_part_buf.defined()) {
TORCH_CHECK(o_part_buf.scalar_type() == torch::kFloat32, "o_part_buf must be f32");
TORCH_CHECK(ml_part_buf.scalar_type() == torch::kFloat32, "ml_part_buf must be f32");
if (o_part_buf.has_value() && ml_part_buf.has_value()
&& o_part_buf->defined() && ml_part_buf->defined()) {
TORCH_CHECK(o_part_buf->scalar_type() == torch::kFloat32, "o_part_buf must be f32");
TORCH_CHECK(ml_part_buf->scalar_type() == torch::kFloat32, "ml_part_buf must be f32");
int64_t o_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * p.head_dim;
TORCH_CHECK(o_part_buf.numel() >= o_needed,
"o_part_buf too small: need ", o_needed, " got ", o_part_buf.numel());
p.o_part = (float*)o_part_buf.data_ptr();
p.ml_part = (float*)ml_part_buf.data_ptr();
TORCH_CHECK(o_part_buf->numel() >= o_needed,
"o_part_buf too small: need ", o_needed, " got ", o_part_buf->numel());
p.o_part = (float*)o_part_buf->data_ptr();
p.ml_part = (float*)ml_part_buf->data_ptr();
} else {
alloc_split_partials(p);
}