diff --git a/src/liger_kernel/ops/backends/_ascend/ops/fused_linear_jsd.py b/src/liger_kernel/ops/backends/_ascend/ops/fused_linear_jsd.py
index 7054c58d0..20876900c 100644
--- a/src/liger_kernel/ops/backends/_ascend/ops/fused_linear_jsd.py
+++ b/src/liger_kernel/ops/backends/_ascend/ops/fused_linear_jsd.py
@@ -44,6 +44,7 @@ def fused_linear_jsd_forward(
     ignore_index,
     has_label,
     temperature,
+    accum_dtype=None,
 ):
     device = student_input.device
     dtype = student_input.dtype
@@ -202,6 +203,7 @@ def forward(
         jsd_beta: float = 0.5,
         ignore_index: int = -100,
         temperature: float = 1.0,
+        accum_dtype: Optional[torch.dtype] = None,
     ):
         """
         Args:
@@ -236,6 +238,7 @@ def forward(
             ignore_index,
             has_label,
             temperature,
+            accum_dtype,
         )
         # downcast to dtype and store for backward
         ctx.save_for_backward(
@@ -249,4 +252,4 @@ def forward(
     def backward(ctx, grad_output):
         (grad_input, grad_weight) = ctx.saved_tensors
         grad_input, grad_weight = fused_linear_jsd_backward(grad_output, grad_input, grad_weight)
-        return (grad_input, grad_weight, None, None, None, None, None, None)
+        return (grad_input, grad_weight, None, None, None, None, None, None, None)