One-way: TMEM → registers (normalize) → SMEM → GMEM
Based on CUTLASS FMHA reference's correction_epilog pattern.
Eliminates TMEM round-trip error for O normalization.
O rescale (kt>0) still uses old atoms (separate fix).
The correction epilog (TMEM→reg→SMEM→GMEM one-way trip) is the right approach
but the TMA store from SMEM requires proper partitioning that needs more work.
Reverting to the known-working state (with 3% TMEM round-trip error) to focus
on the SMEM-P write first.
- Remove noop + normalize TMEM round-trips (3% error per trip)
- Use epilogue_tmem_copy_and_partition for TMEM→reg (paired atoms)
- Use epilogue_smem_copy_and_partition for reg→SMEM (paired atoms)
- Apply 1/row_sum normalization in register space (exact)
- TMA store from SMEM→GMEM (no TMEM write-back)
- Add iter_acc_early_release_in_epilogue attribute
- Update SMEM-P comments to reflect coordinate-indexed fallback
Each softmax thread writes its P values to sP using the (m,k) coordinates
from tTMEM_LOADcS. The k coordinate is decomposed into (k0,k1,k2) to
match sP's ((128,16),1,(4,2)) layout. CuTeDSL tensor indexing handles
the swizzle automatically. No make_tiled_copy needed.
Per CUTLASS guidance:
- make_tiled_copy_C/D encode wrong invariants for this transfer
- Build custom R→S copy where TV map comes from tTMEM_LOADcS (softmax thread
ownership) and destination addresses come from sP layout (PV A-operand swizzled SMEM)
- Use composition(sP_2d_layout, p_coord_layout) for atom_layout_tv
- Start with scalar BF16 (16-bit) stores — vectorize later
- Zero-fill source for compile test, will fill with actual P values next
Uses the CUTLASS blackwell_helpers pattern:
- get_smem_store_op creates a SMEM store atom paired with the TMEM load
- make_tiled_copy_D uses the same thread partition as the TMEM load
- Softmax warps write P to sP using the same thread mapping they use for reading S
- MMA warp reads P from sP via pv_mma.make_fragment_A(sP)
- Replaces the zero-fill stub with a proper register→SMEM copy
