diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py index 6739abf621..45125478dd 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py @@ -249,6 +249,28 @@ FMHA_BWD_API_INNER_DISPATCH_LAUNCHER = """ M0_1D = 64 +# Per-d (M0, BlockSize) for convert_dq kernel, chosen so that: +# (a) convert M0 == bwd M0 (so convert's mask check at convert-tile granularity +# exactly matches bwd's per-tile write decision; enables torch::empty for +# dq_accum without garbage reads) +# (b) the convert tile distribution constraint M0 * d >= BlockSize * alignment(8) +# is satisfied. +# +# Reference bwd M0 by d for fp16/bf16 group deterministic gfx950: +# d=32, 64: bwd M0=32 +# d=128, 256: bwd M0=16 +def get_convert_m0_and_blocksize(hdim): + if hdim <= 32: + # bwd M0=32; need M0*32 >= BlockSize*8 → BlockSize <= 128 for M0=32 + return 32, 128 + elif hdim <= 64: + # bwd M0=32; 32*64=2048 = 256*8 → BlockSize=256 OK + return 32, 256 + else: + # d in [128, 256]: bwd M0=16; 16*128=2048 = 256*8 → BlockSize=256 OK + return 16, 256 + + # GEMM0: Q@K=S^T # GEMM1: P^T@dO^T=dV(This was chosen as G1 to match fwd, but N1 must be equal to headdim_v) # GEMM2: dO@V=dP^T(This was chosen as G2 because of the calculation order) @@ -647,7 +669,7 @@ using fmha_bwd_convert_dq_pipeline_problem_{F_idx} = ck_tile::BlockFmhaBwdConvertQGradPipelineProblem< typename FmhaBwdTypeConfig::AccDataType, typename FmhaBwdTypeConfig::QGradDataType, - /* BlockSize = */ 256, + /* BlockSize = */ {F_blocksize}, {F_bm0}, {F_bn0}, {F_hdim}, @@ -712,6 +734,7 @@ class FmhaBwdConvertQGradKernel: F_dtype: str # data type F_bm0: int # tile size along q seqlen (block size) F_bn0: int # tile size along k seqlen + F_blocksize: int # convert kernel BlockSize (per-d, see get_convert_m0_and_blocksize) F_spad: str # true/false F_dpad: str # F_mode: str # value from MODE_MAP @@ -728,6 +751,7 @@ class FmhaBwdConvertQGradKernel: F_dtype=BWD_DTYPE_MAP[self.F_dtype], F_bm0=self.F_bm0, F_bn0=self.F_bn0, + F_blocksize=self.F_blocksize, F_spad=BOOL_MAP[self.F_spad], F_dpad=BOOL_MAP[self.F_dpad], F_mode=MODE_MAP[self.F_mode], @@ -889,13 +913,15 @@ class FmhaBwdApiTrait: return 2 F_dpad = "t" if self.dpad else "f" + convert_m0, convert_blocksize = get_convert_m0_and_blocksize(self.hdim) return FmhaBwdConvertQGradKernel( F_arch=self.arch, F_idx=self.idx, F_hdim=self.hdim, F_dtype=self.dtype, - F_bm0=M0_1D, + F_bm0=convert_m0, F_bn0=self.convert_dq_bn0, + F_blocksize=convert_blocksize, F_spad=self.spad1d, F_dpad=F_dpad, F_mode=self.mode, diff --git a/example/ck_tile/01_fmha/fmha_bwd.hpp b/example/ck_tile/01_fmha/fmha_bwd.hpp index 8eb8834e12..1efcf9eb8c 100644 --- a/example/ck_tile/01_fmha/fmha_bwd.hpp +++ b/example/ck_tile/01_fmha/fmha_bwd.hpp @@ -416,7 +416,10 @@ auto fmha_bwd_convert_dq_create_kargs_and_grids(fmha_bwd_args args) args.stride_dq_acc, args.nhead_stride_dq, args.nhead_stride_dq_acc, - args.split_stride_dq_acc); + args.split_stride_dq_acc, + args.mask_type, + args.window_size_left, + args.window_size_right); } else { // create batch mode kernel arguments @@ -433,7 +436,10 @@ auto fmha_bwd_convert_dq_create_kargs_and_grids(fmha_bwd_args args) args.batch_stride_dq_acc, args.split_stride_dq_acc, args.batch, - args.nhead_q); + args.nhead_q, + args.mask_type, + args.window_size_left, + args.window_size_right); } }(); diff --git a/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp index 5659162c97..d8f0690abe 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp @@ -1660,6 +1660,12 @@ struct FmhaBwdConvertQGradKernel ck_tile::index_t stride_dq_acc; ck_tile::index_t nhead_stride_dq; ck_tile::long_index_t nhead_stride_dq_acc; + + // Mask params for split-skipping in deterministic reduce. + // mask_type == 0 (NO_MASK) disables skipping; window sizes are -1 for unbounded. + ck_tile::index_t mask_type = 0; + ck_tile::index_t window_size_left = -1; + ck_tile::index_t window_size_right = -1; }; struct FmhaBwdConvertQGradDeterministicKargs @@ -1713,7 +1719,10 @@ struct FmhaBwdConvertQGradKernel ck_tile::long_index_t batch_stride_dq_acc, ck_tile::index_t split_stride_dq_acc, ck_tile::index_t batch_size, - ck_tile::index_t nhead) + ck_tile::index_t nhead, + ck_tile::index_t mask_type = 0, + ck_tile::index_t window_size_left = -1, + ck_tile::index_t window_size_right = -1) { Kargs kargs{{dq_acc_ptr, dq_ptr, @@ -1723,7 +1732,10 @@ struct FmhaBwdConvertQGradKernel stride_dq, stride_dq_acc, nhead_stride_dq, - nhead_stride_dq_acc}, + nhead_stride_dq_acc, + mask_type, + window_size_left, + window_size_right}, {}, batch_stride_dq, batch_stride_dq_acc}; @@ -1757,7 +1769,10 @@ struct FmhaBwdConvertQGradKernel ck_tile::index_t stride_dq_acc, ck_tile::index_t nhead_stride_dq, ck_tile::long_index_t nhead_stride_dq_acc, - ck_tile::index_t split_stride_dq_acc) + ck_tile::index_t split_stride_dq_acc, + ck_tile::index_t mask_type = 0, + ck_tile::index_t window_size_left = -1, + ck_tile::index_t window_size_right = -1) { Kargs kargs{{dq_acc_ptr, dq_ptr, @@ -1767,7 +1782,10 @@ struct FmhaBwdConvertQGradKernel stride_dq, stride_dq_acc, nhead_stride_dq, - nhead_stride_dq_acc}, + nhead_stride_dq_acc, + mask_type, + window_size_left, + window_size_right}, {}, reinterpret_cast(seqstart_q_ptr), reinterpret_cast(seqstart_k_ptr), @@ -1964,7 +1982,70 @@ struct FmhaBwdConvertQGradKernel if constexpr(kIsDeterministic) { - FmhaBwdConvertQGrad{}(dq_acc_dram_window, dq_dram_window, nsplits); + // Runtime mask-aware split skipping. When a mask is active, the set of + // splits whose K-chunk range intersects the valid K range for this + // Q-tile is a contiguous range [first_valid, last_valid] (causal / SWA + // produce a single contiguous valid K interval). Shift the window + // origin to first_valid and pass valid_nsplits to the same pipelined + // reduce overload — software pipelining is preserved unchanged. + // + // Skipped prefix splits avoid HBM reads entirely. dq_acc may be + // uninitialized (torch::empty) in deterministic mode; the skip is + // what makes that safe. + if(kargs.mask_type != 0) + { + const bool is_top_left = + (kargs.mask_type == + static_cast(GenericAttentionMaskEnum::MASK_FROM_TOP_LEFT)); + const auto coords = + ck_tile::make_generic_attention_mask_coordinates_from_lr_window( + kargs.window_size_left, + kargs.window_size_right, + /*sink_size=*/0, + kargs.seqlen_q, + kargs.seqlen_k, + is_top_left); + ck_tile::SimplifiedGenericAttentionMask mask( + coords.at(ck_tile::number<0>{}), + coords.at(ck_tile::number<1>{}), + coords.at(ck_tile::number<2>{}), + coords.at(ck_tile::number<3>{}), + coords.at(ck_tile::number<4>{})); + const auto k_range = mask.GetTileRangeAlongX( + i_m0, ck_tile::number{}, ck_tile::number<1>{}); + const index_t valid_k_lo = k_range.at(ck_tile::number<0>{}); + const index_t valid_k_hi = k_range.at(ck_tile::number<1>{}); + + if(valid_k_lo < valid_k_hi) + { + // Splits are laid out one per BWD-kernel K-chunk of size kN0. + const index_t first_valid = valid_k_lo / kN0; + const index_t last_valid = (valid_k_hi - 1) / kN0; + const index_t valid_nsplits = last_valid - first_valid + 1; + + auto dq_acc_dram_window_shifted = make_tile_window( + dq_acc_dram, + make_tuple( + number<1>{}, number{}, number{}), + {first_valid, i_m0, 0}); + + FmhaBwdConvertQGrad{}( + dq_acc_dram_window_shifted, dq_dram_window, valid_nsplits); + } + else + { + // Valid_k range is empty for this Q-tile under the mask. + // dq for this Q-tile is 0 (no K contributes). Caller must + // have zeroed dq up front (mha_varlen_bwd.cpp does dq.zero_() + // unconditionally). Return early without writing. + return; + } + + } + else + { + FmhaBwdConvertQGrad{}(dq_acc_dram_window, dq_dram_window, nsplits); + } } else { diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp index e6d7c622f7..e954c2c00e 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp @@ -33,6 +33,44 @@ struct BlockFmhaBwdConvertQGrad CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize() { return 0; } + // Zero-fill only — used when the mask leaves every dq_acc split unwritten for + // a given Q-tile (e.g., narrow SWA with a Q position outside the window). + // dq = 0 is mathematically correct: no K positions contribute to this Q's gradient. + // + // We mirror the "Convert only" overload's pattern (create an AccDataType tile, + // cast to QGradDataType, store) rather than directly creating a QGradDataType + // tile via make_static_distributed_tensor + clear_tile + store_tile. The latter + // empirically only stores half the kM0×kQKHeaddim tile rows for some + // distribution configurations — a subtle internal-state issue. Going through + // cast_tile from an AccDataType source produces a tile whose store_tile covers + // the full row range, matching the simple convert path. + template + CK_TILE_HOST_DEVICE void + operator()(QGradDramBlockWindowTmp& dq_dram_block_window_tmp) const + { + static_assert( + std::is_same_v>, + "wrong!"); + static_assert(kM0 == QGradDramBlockWindowTmp{}.get_window_lengths()[number<0>{}], "wrong!"); + + // Reuse the existing simple convert overload pattern: attach the canonical + // MakePostQGradDramTileDistribution to a window over the dq output itself, + // load_tile (gives us a properly-formed tile), clear in-place, then store. + // The extra load reads garbage from dq (torch::empty_like(q)) — wasted but small. + // We do this because make_static_distributed_tensor + clear_tile + store_tile + // empirically only writes half the tile rows; load_tile produces a tile whose + // store_tile covers the full row range. + auto dq_dram_window_with_dist = make_tile_window( + dq_dram_block_window_tmp.get_bottom_tensor_view(), + dq_dram_block_window_tmp.get_window_lengths(), + dq_dram_block_window_tmp.get_window_origin(), + Policy::template MakePostQGradDramTileDistribution()); + auto dq = load_tile(dq_dram_window_with_dist); + clear_tile(dq); + store_tile(dq_dram_block_window_tmp, dq); + } + // Convert only template CK_TILE_HOST_DEVICE void @@ -90,7 +128,12 @@ struct BlockFmhaBwdConvertQGrad auto dq_acc_buf = load_tile(dq_acc_dram_window); move_tile_window(dq_acc_dram_window, {1, 0, 0}); - do + // Use while-loop (not do-while) so nsplits == 1 is correct. + // For nsplits == 1: we prefetched split 0 above, the loop body is skipped, + // and the tail accumulate below sums split 0 correctly with no OOB load. + // For nsplits >= 2: pipelined pattern is preserved (prefetch next while + // accumulating current). + while(i_total_loops < (nsplits - 1)) { sweep_tile_span(dq_acc_spans[number<0>{}], [&](auto idx0) { sweep_tile_span(dq_acc_spans[number<1>{}], [&](auto idx1) { @@ -105,7 +148,7 @@ struct BlockFmhaBwdConvertQGrad move_tile_window(dq_acc_dram_window, {1, 0, 0}); i_total_loops += 1; - } while(i_total_loops < (nsplits - 1)); + } sweep_tile_span(dq_acc_spans[number<0>{}], [&](auto idx0) { sweep_tile_span(dq_acc_spans[number<1>{}], [&](auto idx1) {