ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-14 19:18:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Exp1: mask-aware convert_dq skip + per-d M0_CONVERT alignment for deterministic bwd

Browse Source 
Avoids unnecessary HBM reads of unwritten dq_acc splits in the deterministic
varlen bwd path, enabling torch::empty for the dq_acc workspace (caller change).

- block_fmha_bwd_convert_dq.hpp: add zero-write operator() overload for fully
  mask-skipped Q-tiles; switch reduce loop from do-while → while so nsplits==1
  is correct (was an OOB load + accumulate of garbage).
- fmha_bwd_kernel.hpp: add mask kargs (mask_type, window_size_left/right) to
  FmhaBwdConvertQGradCommonKargs; in convert dispatch, compute valid K range
  via SimplifiedGenericAttentionMask and shift the dq_acc window origin to
  first_valid split so the pipelined reduce only reads bwd-written slots.
- fmha_bwd.hpp: plumb mask params through fmha_bwd_convert_dq_create_kargs_and_grids
  for both batch and group MakeKargs overloads.
- codegen/ops/fmha_bwd.py: per-d (M0, BlockSize) selection for convert_dq so
  convert M0 == bwd M0 (d=32→M0=32 BS=128, d=64→M0=32 BS=256, d>=128→M0=16 BS=256).
  Alignment is required because convert M0 > bwd M0 causes the convert tile to
  span multiple bwd sub-tiles, only some of which the bwd visits → garbage reads
  under torch::empty.

Verified on MI355 with full test_flash_attn_varlen_deterministic sweep (1920 cases).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Ye Wang
2026-06-02 22:37:17 -05:00
parent 791afc6465
commit 6f9c7df6b7
4 changed files with 167 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py
View File

@@ -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<fmha_dtype_{F_idx}>::AccDataType,
typename FmhaBwdTypeConfig<fmha_dtype_{F_idx}>::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,

10
example/ck_tile/01_fmha/fmha_bwd.hpp
View File

@@ -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);
}
}();

91
include/ck_tile/ops/fmha/kernel/fmha_bwd_kernel.hpp
View File

@@ -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<const int32_t*>(seqstart_q_ptr),
reinterpret_cast<const int32_t*>(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<index_t>(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</*IsMasking=*/true> 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<kM0>{}, 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<kM0>{}, number<kQKHeaddim>{}),
{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
{

47
include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_convert_dq.hpp
View File

@@ -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 <typename QGradDramBlockWindowTmp>
CK_TILE_HOST_DEVICE void
operator()(QGradDramBlockWindowTmp& dq_dram_block_window_tmp) const
{
static_assert(
std::is_same_v<QGradDataType,
remove_cvref_t<typename QGradDramBlockWindowTmp::DataType>>,
"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<Problem>());
auto dq = load_tile(dq_dram_window_with_dist);
clear_tile(dq);
store_tile(dq_dram_block_window_tmp, dq);
}
// Convert only
template <typename QGradAccDramBlockWindowTmp, typename QGradDramBlockWindowTmp>
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) {