From 31c40693d52cd4fc0da702c02db4b40c2560faa7 Mon Sep 17 00:00:00 2001 From: Enigmatisms Date: Thu, 6 Nov 2025 06:23:42 +0000 Subject: [PATCH 1/2] [Trial] Forward optimization: simple skip for causal (v2.1 causal) --- csrc/flashmask_v2/flash_fwd_kernel_sm90.h | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/csrc/flashmask_v2/flash_fwd_kernel_sm90.h b/csrc/flashmask_v2/flash_fwd_kernel_sm90.h index 3ba512689a8..9f536539d38 100644 --- a/csrc/flashmask_v2/flash_fwd_kernel_sm90.h +++ b/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(); // coarse pipeline stage (offset, 0 or 2) @@ -318,7 +327,7 @@ class FlashAttnFwdSm90 { 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 + From 26b08f0cae9cacd6c8f7ba946c0f5ceaf980adcb Mon Sep 17 00:00:00 2001 From: Enigmatisms Date: Thu, 6 Nov 2025 09:25:19 +0000 Subject: [PATCH 2/2] [Trial] Forward optimization fine-grained skipping --- csrc/flashmask_v2/flash_fwd_kernel_sm90.h | 2 +- .../mainloop_fwd_sm90_tma_gmma_ws.hpp | 34 +++++++++++++++++-- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/csrc/flashmask_v2/flash_fwd_kernel_sm90.h b/csrc/flashmask_v2/flash_fwd_kernel_sm90.h index 9f536539d38..adcc183dbdf 100644 --- a/csrc/flashmask_v2/flash_fwd_kernel_sm90.h +++ b/csrc/flashmask_v2/flash_fwd_kernel_sm90.h @@ -323,7 +323,7 @@ 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) diff --git a/csrc/flashmask_v2/mainloop_fwd_sm90_tma_gmma_ws.hpp b/csrc/flashmask_v2/mainloop_fwd_sm90_tma_gmma_ws.hpp index bce5acf0e94..f7d69c47f13 100644 --- a/csrc/flashmask_v2/mainloop_fwd_sm90_tma_gmma_ws.hpp +++ b/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;