bwd case when offset is 0 (#4963)

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/1982

Pull Request resolved: #4963

This case is for when we are not using bottom right mask.

It should be slightly better perf in that case.

# notes

We note that backward is in general not stable. Sometimes you can get IMA. And numerics are not as good as we want it to be.

Reviewed By: q10

Differential Revision: D83076701

fbshipit-source-id: a1016b15a86d10f21d962166eae036d959befe18

Author: henrylhtsang

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp

@@ -1799,8 +1799,14 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
     } else if constexpr (std::is_base_of_v<cutlass::fmha::collective::CausalMask<false>, Mask>) {
       int offset = get<1>(problem_shape) - get<0>(problem_shape);
       iter_start = max(0, (int(get<1>(blk_coord) * TileShapeK{}) - offset) / (int)TileShapeQ{});
-    } else if constexpr (std::is_base_of_v<cutlass::fmha::collective::LocalMask<false>, Mask>) {
-      int offset = get<1>(problem_shape) - get<0>(problem_shape);
+    }
+    else if constexpr (
+        std::is_base_of_v<cutlass::fmha::collective::LocalMask<false>, Mask> ||
+        std::is_base_of_v<cutlass::fmha::collective::LocalMask<true>, Mask>
+    ) {
+      int offset = std::is_base_of_v<cutlass::fmha::collective::LocalMask<false>, Mask>
+        ? get<1>(problem_shape) - get<0>(problem_shape)
+        : 0;
 
       int k_max = (get<1>(blk_coord) + 1) * TileShapeK{};
       int q_max = min(get<0>(problem_shape), k_max - offset + params.mainloop_params.window_size_left);