Adds tensor reshaping for mask and bias in backward kernel

Extends the existing tensor reshaping logic to include mask and bias tensors alongside the scores tensor.

All three tensors now use the same layout conversion from MMA format to row-column format, ensuring consistent tensor structure for subsequent computations.

Author: LoserCheems

csrc/src/flash_bwd_kernel.h

@@ -651,8 +651,10 @@ inline __device__ void compute_dq_dk_dv_1colblock(const Params &params, const in
             FLASH_NAMESPACE::apply_softcap(acc_s, params.softcap);
         }
 
-        // Reshape acc_s from (MMA=4, MMA_N, MMA_N) to (row=(2, MMA_N), col=(2, MMA_N))
+        // Reshape acc_s, mask, bias from (MMA=4, MMA_N, MMA_N) to (row=(2, MMA_N), col=(2, MMA_N))
         Tensor scores = make_tensor(acc_s.data(), FLASH_NAMESPACE::convert_layout_acc_rowcol(acc_s.layout()));
+        Tensor mask = make_tensor(tSrMask.data(), FLASH_NAMESPACE::convert_layout_acc_rowcol(tSrMask.layout()));
+        Tensor bias = make_tensor(tSrBias.data(), FLASH_NAMESPACE::convert_layout_acc_rowcol(tSrBias.layout()));
         // if (cute::thread(32, 0)) { print(scores); }
 
         // Softcapping - calculating dTanh and scaling dS later with it