diff --git a/csrc/kernels/attn_decode.cu b/csrc/kernels/attn_decode.cu
index d85844e..2880784 100644
--- a/csrc/kernels/attn_decode.cu
+++ b/csrc/kernels/attn_decode.cu
@@ -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);
     }
diff --git a/csrc/kernels/attn_paged_decode.cu b/csrc/kernels/attn_paged_decode.cu
index d6baed4..682ebf3 100644
--- a/csrc/kernels/attn_paged_decode.cu
+++ b/csrc/kernels/attn_paged_decode.cu
@@ -12,8 +12,8 @@ torch::Tensor attn_paged_decode(
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
     double scale,
-    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();
@@ -26,14 +26,15 @@ torch::Tensor attn_paged_decode(
     auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
     p.o = (bf16*)O.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);
     }