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..f25c6bbad 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
@@ -63,7 +64,12 @@ def fused_linear_jsd_forward(
     chunk_size = triton.next_power_of_2(triton.cdiv(BT, inc_factor))  # (BT + inc_factor - 1) // inc_factor
     num_chunks = triton.cdiv(BT, chunk_size)  # (BT + chunk_size - 1) // chunk_size
 
-    grad_weight = torch.zeros_like(student_weight, device=device) if student_weight.requires_grad else None
+    if accum_dtype is None:
+        grad_weight = torch.zeros_like(student_weight, device=device) if student_weight.requires_grad else None
+    else:
+        grad_weight = (
+            torch.zeros_like(student_weight, dtype=accum_dtype, device=device) if student_weight.requires_grad else None
+        )
     grad_input = torch.zeros_like(student_input)
     # we use fp32 for loss accumulator
     loss_1d = torch.zeros((BT, V), dtype=torch.float32, device=device)
@@ -139,9 +145,15 @@ def fused_linear_jsd_forward(
         grad_input[start_idx:end_idx] = student_logits_chunk @ student_weight
 
         if grad_weight is not None:
-            grad_weight.add_(student_logits_chunk.t() @ student_input_chunk)
+            if accum_dtype is None:
+                grad_weight.add_(student_logits_chunk.t() @ student_input_chunk)
+            else:
+                grad_weight.add_(torch.mm(student_logits_chunk.t(), student_input_chunk).float())
 
     loss = torch.sum(loss_1d)
+    grad_weight = (
+        grad_weight.to(student_weight.dtype) if grad_weight is not None and accum_dtype is not None else grad_weight
+    )
     return loss, grad_input, grad_weight
 
 
@@ -202,6 +214,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 +249,7 @@ def forward(
             ignore_index,
             has_label,
             temperature,
+            accum_dtype,
         )
         # downcast to dtype and store for backward
         ctx.save_for_backward(
@@ -249,4 +263,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)