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, int64_t causal_offset,
double scale, double scale,
int64_t layout, int64_t layout,
torch::Tensor o_part_buf, c10::optional<torch::Tensor> o_part_buf,
torch::Tensor ml_part_buf c10::optional<torch::Tensor> ml_part_buf
) { ) {
const at::cuda::OptionalCUDAGuard device_guard(device_of(q)); const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
auto stream = at::cuda::getCurrentCUDAStream(); auto stream = at::cuda::getCurrentCUDAStream();
@@ -24,14 +24,15 @@ torch::Tensor attn_decode(
auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O; auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O;
p.o = (bf16*)O_view.data_ptr(); p.o = (bf16*)O_view.data_ptr();
if (o_part_buf.defined() && ml_part_buf.defined()) { if (o_part_buf.has_value() && ml_part_buf.has_value()
TORCH_CHECK(o_part_buf.scalar_type() == torch::kFloat32, "o_part_buf must be f32"); && o_part_buf->defined() && ml_part_buf->defined()) {
TORCH_CHECK(ml_part_buf.scalar_type() == torch::kFloat32, "ml_part_buf must be f32"); 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; int64_t o_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * p.head_dim;
TORCH_CHECK(o_part_buf.numel() >= o_needed, TORCH_CHECK(o_part_buf->numel() >= o_needed,
"o_part_buf too small: need ", o_needed, " got ", o_part_buf.numel()); "o_part_buf too small: need ", o_needed, " got ", o_part_buf->numel());
p.o_part = (float*)o_part_buf.data_ptr(); p.o_part = (float*)o_part_buf->data_ptr();
p.ml_part = (float*)ml_part_buf.data_ptr(); p.ml_part = (float*)ml_part_buf->data_ptr();
} else { } else {
alloc_split_partials(p); alloc_split_partials(p);
} }
csrc/kernels/attn_paged_decode.cu
+10 -9
View File
@@ -12,8 +12,8 @@ torch::Tensor attn_paged_decode(
c10::optional<torch::Tensor> mask, c10::optional<torch::Tensor> mask,
int64_t causal_offset, int64_t causal_offset,
double scale, double scale,
torch::Tensor o_part_buf, c10::optional<torch::Tensor> o_part_buf,
torch::Tensor ml_part_buf c10::optional<torch::Tensor> ml_part_buf
) { ) {
const at::cuda::OptionalCUDAGuard device_guard(device_of(q)); const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
auto stream = at::cuda::getCurrentCUDAStream(); auto stream = at::cuda::getCurrentCUDAStream();
@@ -26,14 +26,15 @@ torch::Tensor attn_paged_decode(
auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options()); auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
p.o = (bf16*)O.data_ptr(); p.o = (bf16*)O.data_ptr();
if (o_part_buf.defined() && ml_part_buf.defined()) { if (o_part_buf.has_value() && ml_part_buf.has_value()
TORCH_CHECK(o_part_buf.scalar_type() == torch::kFloat32, "o_part_buf must be f32"); && o_part_buf->defined() && ml_part_buf->defined()) {
TORCH_CHECK(ml_part_buf.scalar_type() == torch::kFloat32, "ml_part_buf must be f32"); 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; int64_t o_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * p.head_dim;
TORCH_CHECK(o_part_buf.numel() >= o_needed, TORCH_CHECK(o_part_buf->numel() >= o_needed,
"o_part_buf too small: need ", o_needed, " got ", o_part_buf.numel()); "o_part_buf too small: need ", o_needed, " got ", o_part_buf->numel());
p.o_part = (float*)o_part_buf.data_ptr(); p.o_part = (float*)o_part_buf->data_ptr();
p.ml_part = (float*)ml_part_buf.data_ptr(); p.ml_part = (float*)ml_part_buf->data_ptr();
} else { } else {
alloc_split_partials(p); alloc_split_partials(p);
} }