Adds mask and bias handling to backward kernel block loop

Extends the memory advancement logic to include mask and bias tensors alongside the existing query tensor handling. This ensures all relevant tensors are properly synchronized when processing multiple blocks in the backward pass, maintaining consistency across attention computations.

The change mirrors the query tensor advancement pattern by updating pointers and copying data for both mask and bias tensors using the same block-based iteration approach.

Author: LoserCheems

csrc/src/flash_bwd_kernel.h

@@ -879,13 +879,27 @@ inline __device__ void compute_dq_dk_dv_1colblock(const Params &params, const in
         }
         if (!Double_buffer && m_block > m_block_min) {
             __syncthreads();
-            // Advance gQ
+            // Advance gQ, gMask, gBias
             tQgQ.data() = tQgQ.data() + (-int(kBlockM * params.q_row_stride));
+            tMaskgMask.data() = tMaskgMask.data() + (-int(kBlockM * params.mask_row_stride));
+            tBiasgBias.data() = tBiasgBias.data() + (-int(kBlockM * params.bias_row_stride));
             FLASH_NAMESPACE::copy</*Is_even_MN=*/true, Is_even_K>(
                 gmem_tiled_copy_QKV,
                 tQgQ, tQsQ,
                 tQcQ, tQpQ
             );
+            FLASH_NAMESPACE::copy_MN<Is_even_MN, /*Clear_OOB_MN=*/true>(
+                gmem_tiled_copy_Mask,
+                tMaskgMask, tMasksMask,
+                tMaskcMask,
+                binfo.actual_seqlen_q - (m_block - 1) * kBlockM, binfo.actual_seqlen_k - n_block * kBlockN
+            );
+            FLASH_NAMESPACE::copy_MN<Is_even_MN, /*Clear_OOB_MN=*/true>(
+                gmem_tiled_copy_Bias,
+                tBiasgBias, tBiassBias,
+                tBiascBias,
+                binfo.actual_seqlen_q - (m_block - 1) * kBlockM, binfo.actual_seqlen_k - n_block * kBlockN
+            );
             FLASH_NAMESPACE::cp_async_fence();
         }