diff --git a/dsv4/kernels/attention/production.py b/dsv4/kernels/attention/production.py
index c2a62831..a96bc8d3 100644
--- a/dsv4/kernels/attention/production.py
+++ b/dsv4/kernels/attention/production.py
@@ -63,7 +63,7 @@ def _get_or_compile_kernel(head_dim: int, s_k: int, use_smem_p: bool = False,
     mLSE = ct.from_dlpack(lse).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse))
     mRS = ct.from_dlpack(row_sums).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums))
 
-    compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS)
+    compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, lse=mLSE)
     _kernel_cache[key] = (compiled, kernel)
     return (compiled, kernel)
 
@@ -204,7 +204,7 @@ def _attention_single_head(
             mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))
             mRS = ct.from_dlpack(row_sums_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums_tensor))
 
-            compiled(mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS)
+            compiled(mQ, mK, mV, mC, stream, lse=mLSE)
             torch.cuda.synchronize()
 
             seg_o[:, v_start:v_end] = c_tile[:, :, 0].float()