Fwd optimization: "flashchunk" and "flashblock" optimization

csrc/flashmask_v2/flash_fwd_kernel_sm90.h

@@ -305,6 +305,15 @@ class FlashAttnFwdSm90 {
               // for padding 32 and padding 4: the num_chunk (pad_32) >= num_chunk (pad_4) is always true
               const int nblock_seqlen = ((seqlen_info.seqlen_k + kBlockN - 1) / kBlockN + 3) & 0xfffffffc; // umiswing: padding for int4 load
               const int num_chunk = (nblock_seqlen + CollectiveMainloop::Flashmask_n_block_buffer_valid_length - 1) / CollectiveMainloop::Flashmask_n_block_buffer_valid_length;
+              const int reverse_chunk_start = [&] {
+                if constexpr (Is_causal) {
+                    // if causal, the 'valid' sequence length is smaller. We don't need that many chunks, so we will start from chunks closer to the start
+                    const int seqlen_offset = seqlen_info.seqlen_k - seqlen_info.seqlen_q;
+                    return std::max(num_chunk - 1 - (kBlockM * (m_block + 1) + seqlen_offset) / (kBlockN * CollectiveMainloop::Flashmask_n_block_buffer_valid_length), 0);
+                } else {
+                    return 0;
+                }
+              } ();
               // reverse_chunk_idx, start from right to left: [5, 4, 3, 2, 1, 0], and fwd kernel scans from right to left
               bool valid_chunk = true;
               const int cppl_stage = scheduler.template stage<true>();      // coarse pipeline stage (offset, 0 or 2)
@@ -314,11 +323,11 @@ class FlashAttnFwdSm90 {
                             reverse_chunk_idx,                                                                                                                  \
                             num_chunk,                                                                                                                          \
                             reverse_chunk_idx == num_chunk - 1 ? CollectiveMainloop::Flashmask_n_block_finish : CollectiveMainloop::Flashmask_n_block_chunk_end,\
-                            n_block_min, n_block_max, seqlen_info.seqlen_q,                                                                                     \
+                            n_block_min, n_block_max, nblock_seqlen, seqlen_info.seqlen_q, seqlen_info.seqlen_k,                                                \
                             flashmask_maxmin_smem + 8 * CollectiveMainloop::Flashmask_n_block_buffer_length * (n_block_pipe_write.index() + cppl_stage),        \
                             n_block_smem + CollectiveMainloop::Flashmask_n_block_buffer_length * (n_block_pipe_write.index() + cppl_stage),                     \
                             extra_flags + n_block_pipe_write.index() + cppl_stage)
-              for(int reverse_chunk_idx = 0; reverse_chunk_idx < num_chunk; reverse_chunk_idx++) {
+              for(int reverse_chunk_idx = reverse_chunk_start; reverse_chunk_idx < num_chunk; reverse_chunk_idx++) {
                 if (valid_chunk)
                     pipeline_n_block.producer_acquire(n_block_pipe_write);
                 mainloop.load_max_min(params.mainloop, seqlen_info, block_coord, reverse_chunk_idx, num_chunk, flashmask_maxmin_smem +

csrc/flashmask_v2/mainloop_fwd_sm90_tma_gmma_ws.hpp

@@ -674,7 +674,20 @@ struct CollectiveMainloopFwdSm90 {
         int32_t bidh = get<1>(block_coord);
         int32_t bidb = get<2>(block_coord);
         // pad for fully 128B aligned load
-        const int nblock_seqlen = ((seqlen_info.seqlen_k + kBlockN - 1) / kBlockN + 3) & 0xfffffffc;
+
+        // Note(heqianyue): compute the part that actually needs loading and computation for causal masks
+        const int nblock_seqlen = [&] {
+            if constexpr (Is_causal) {
+                const int m_block = get<0>(block_coord);
+                // Warn(heqianyue): abs_coord might be greater than unpadded nblock_seqlen (this is mathematically prove-able) 
+                // but we don't need to clip to nblock_seqlen, since loading invalid data is fine, we can choose not to use it
+                // (mblock_id + 1) * kBlockM + diff_seqlen is the block ID of the rightmost valid block in causal, so + 1 for length
+                const int valid_nblock_seqlen = ((m_block + 1) * kBlockM + seqlen_info.seqlen_k - seqlen_info.seqlen_q) / kBlockN + 1;
+                return (valid_nblock_seqlen + 3) & 0xfffffffc;
+            } else {    // original nblock_seqlen
+                return ((seqlen_info.seqlen_k + kBlockN - 1) / kBlockN + 3) & 0xfffffffc;
+            }
+        } ();
 
         // change this to num_chunk * chunk_size (should be Flashmask_n_block_buffer_length)
         const int chunks_size = total_num_chunks * Flashmask_n_block_buffer_length;
@@ -755,7 +768,9 @@ struct CollectiveMainloopFwdSm90 {
                      int32_t const end_flag,
                      int32_t const n_block_min,
                      int32_t const n_block_max,
+                     int32_t const nblock_seqlen,
                      int32_t const seqlen_q,
+                     int32_t const seqlen_k,
                      int32_t* const __restrict__ flashmask_maxmin_smem,
                      int32_t* const __restrict__ mask_encode_n_block_smem_,
                      int32_t* const __restrict__ extra_flags) {
@@ -804,6 +819,19 @@ struct CollectiveMainloopFwdSm90 {
 
       int32_t valid_n_block_num = 0;
       const int32_t base_offset = (total_num_chunks - 1 - reverse_chunk_idx) * Flashmask_n_block_buffer_valid_length;
+      constexpr int valid_buffer_length = Flashmask_n_block_buffer_valid_length;
+      const int nblock_start = [&]() {
+        if constexpr (Is_causal) {
+            // Warn(heqianyue): abs_coord might be greater than unpadded nblock_seqlen (this is mathematically prove-able) 
+            // in order not to load invalid data (or even OOR, which is highly not possible), we clip it to nblock_seqlen
+            const int abs_coord = std::min(((m_block + 1) * kBlockM + seqlen_k - seqlen_q) / kBlockN, nblock_seqlen - 1);
+            return std::min(abs_coord - base_offset + 1, valid_buffer_length);
+        } else {
+            // Note(heqianyue): even for non-causal masks, we can skip some of the computation, for example:
+            // since buffer is 16k, if seqlen is 8k, we can skip 50% of the computation here, legit save!
+            return std::min(valid_buffer_length, nblock_seqlen);
+        }
+      } ();
 
       // explanation for the loop condition:
       // -2, -1,  0,  1,  2
@@ -813,8 +841,8 @@ struct CollectiveMainloopFwdSm90 {
       // Note(heqianyue): ute/lte will be seqlen_q (at most). Yet if m_block_e > seqlen_q, even if ute/lte are seqlen_q (masked to the end)
       // we will still consider the block as partially masked, adding unnecessary computation for those fully-masked blocks
       const int m_block_e = __viaddmin_s32(m_block_s, kBlockM, seqlen_q);       // min(a + b, c)
-      for(int32_t idx = Flashmask_n_block_buffer_valid_length - 1 - thread_idx;         // make sure thread_idx is in range [0, ProducerThreadNum)
-          idx >= (0 - (ProducerThreadNum - Flashmask_n_block_buffer_valid_length % ProducerThreadNum)); idx -= ProducerThreadNum
+      for(int32_t idx = nblock_start - 1 - thread_idx;         // make sure thread_idx is in range [0, ProducerThreadNum)
+          idx >= (0 - (ProducerThreadNum - nblock_start % ProducerThreadNum)); idx -= ProducerThreadNum
       ) {
         int32_t n_block = base_offset + idx;
         int prefix_sum = 0;