Returns bias gradients in backward pass

Updates all FlashDMAttn autograd function classes to properly return the computed bias gradients (dbias) in their backward methods instead of returning None.

This ensures gradient computation flows correctly through the bias parameter during backpropagation.

Author: LoserCheems

flash_dmattn/flash_dmattn_interface.py

@@ -490,7 +490,7 @@ def backward(
         )
 
         dqkv = dqkv[..., : dout.shape[-1]]  # We could have padded the head dimension
-        return dqkv, None, None, None, None, None, None, None, None, None
+        return dqkv, None, dbias, None, None, None, None, None, None, None
 
 
 class FlashDMAttnVarlenQKVPackedFunc(torch.autograd.Function):
@@ -604,7 +604,7 @@ def backward(
         )
 
         dqkv = dqkv[..., : dout.shape[-1]]  # We could have padded the head dimension
-        return dqkv, None, None, None, None, None, None, None, None, None, None, None
+        return dqkv, None, dbias, None, None, None, None, None, None, None, None, None
 
 
 class FlashDMAttnKVPackedFunc(torch.autograd.Function):
@@ -712,7 +712,7 @@ def backward(
 
         dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
         dkv = dkv[..., : dout.shape[-1]]
-        return dq, dkv, None, None, None, None, None, None, None, None, None
+        return dq, dkv, None, dbias, None, None, None, None, None, None, None
 
 
 class FlashDMAttnVarlenKVPackedFunc(torch.autograd.Function):
@@ -837,7 +837,7 @@ def backward(
 
         dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
         dkv = dkv[..., : dout.shape[-1]]
-        return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None
+        return dq, dkv, None, dbias, None, None, None, None, None, None, None, None, None, None, None
 
 
 class FlashDMAttnFunc(torch.autograd.Function):
@@ -941,7 +941,7 @@ def backward(
         dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
         dk = dk[..., : dout.shape[-1]]
         dv = dv[..., : dout.shape[-1]]
-        return dq, dk, dv, None, None, None, None, None, None, None, None, None
+        return dq, dk, dv, None, dbias, None, None, None, None, None, None, None
 
 
 class FlashDMAttnVarlenFunc(torch.autograd.Function):
@@ -1063,7 +1063,7 @@ def backward(
         dq = dq[..., : dout.shape[-1]]  # We could have padded the head dimension
         dk = dk[..., : dout.shape[-1]]
         dv = dv[..., : dout.shape[-1]]
-        return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None
+        return dq, dk, dv, None, dbias, None, None, None, None, None, None, None, None, None, None, None, None
 
 
 def flash_dmattn_qkvpacked_func(