Files
composable_kernel/example
Ye Wang 6f9c7df6b7 Exp1: mask-aware convert_dq skip + per-d M0_CONVERT alignment for deterministic bwd
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>
2026-06-02 22:37:17 -05:00
..
2026-01-14 07:31:45 -08:00