biondizzle bf2c7c8bb8 D1.5: Implement in-kernel O rescale via CUTLASS correction_rescale pattern ...

- Both load and store atoms built from SAME tOtO_i (composition-tiled)
- Same Repetition(corr_tile_size=16) for both copies
- pv_done_bar synchronization between MMA and softmax warps
- acc_scale computed per kt iteration, used to rescale O in TMEM
- const_expr(n_kv_tiles > 1) guards for zero overhead at s_k=128
- New test: test_d15_in_kernel_rescale.py (s_k=128/256/384)
- Minimal roundtrip test: test_tmem_roundtrip_minimal.py

2026-05-26 20:26:06 +00:00

6.4 KiB

Raw Blame History

View Raw