[SM100] hd256 2CTA fwd: add missing tcgen05.wait::ld before TMEM slot…

[simon-veitner-redhat]

[SM100] hd256 2CTA fwd: missing tcgen05.wait::ld before TMEM slot releases

Followup to #180.
The issue described above was already present in the kernel that we performance tuned in above PR, however only during further perf tuning i discovered this minor inconsistency with PTX docs.

---

tcgen05.ld is asynchronous. Consuming its destination registers is ordered by the register
scoreboard; the load's TMEM read is not ordered against a later TMEM writer. PTX ISA 9.7.17.6.4.5:
"a register dependency does not imply that a dependee instruction's memory accesses will be
performed before a dependent instruction's memory accesses … tcgen05.wait::ld must be used."
Required producer pattern (9.7.17.6.4.4): tcgen05.ld → tcgen05.wait::ld → mbarrier.arrive.

Three releases in sm100_hd256_2cta_fmha_forward.py skipped the wait:

site	load	release
correction_epilog O tile	Ld32x32bOp(32) ×4	o_handle.release() → MMA warp overwrites O
correction_rescale stats	Ld32x32bOp(2)	stats_handle.release() → softmax warp tcgen05.st
correction_rescale O	ld/st ring	o_handle.release() → MMA warp overwrites O

The softmax warp's S-tile release already had the wait.

Fix

cute.arch.fence_view_async_tmem_load() (tcgen05.wait::ld.sync.aligned) before each release.

The ld→scale→st ring inside correction_rescale is left alone: the st to tile i-1 consumes
every register the ld of tile i-1 produced, so that read has completed before the store issues.
CUTLASS 77_blackwell_fmha uses the same ring without a wait.

Evidence

B300, sm_103a. tcgen05.wait::ld in kernel PTX: 1 → 7; ld/st/mma counts unchanged.

correction_epilog before:

tcgen05.ld.sync.aligned.32x32b.x32.b32 {%r3622, ...}, [...];
mul.rn.f32x2 ...
st.shared.v4.b32 ...
bar.sync 2, 128;
mbarrier.arrive.shared::cluster.b64 _, [%r3710], %r3711;   // o_handle.release()

after: tcgen05.wait::ld.sync.aligned; follows each of the four loads.

correction_rescale stats before — ptxas issues the arrive ahead of the scoreboard wait:

tcgen05.ld.sync.aligned.32x32b.x2.b32 {%r2785, %r2786}, [%r2784];
sub.f32 / mul.f32 / ex2.approx.ftz.f32
mbarrier.arrive.shared.b64 %rd1219, [%r2780+240], %r2792;   // stats_handle.release()

LDTM.x2 R4, tmem[UR4+0x40] ;                        wbar=SB0
SYNCS.ARRIVE.TRANS64.A1T0 RZ, [R3+URZ+0xf0], RZ ;   waitmask=000000
FADD R4, R4, -R5 ;                                  waitmask=000001

after:

ex2.approx.ftz.f32 %r2791, %r2790;
tcgen05.wait::ld.sync.aligned;
mbarrier.arrive.shared.b64 %rd1219, [%r2780+240], %r2792;

LDTM.x2 R4, tmem[UR4+0x40] ;                        wbar=SB0
SYNCS.ARRIVE.TRANS64.A1T0 RZ, [R3+URZ+0xf0], RZ ;   waitmask=010001

correction_rescale O release after:

tcgen05.wait::st.sync.aligned;
tcgen05.wait::ld.sync.aligned;
mbarrier.arrive.shared::cluster.b64 _, [%r170], %r2792;

No SASS change here — the ring's register chain had already retired the loads.

Benchmark

Local Benchmark was run and performance stays invariant under above changes.

[simon-veitner-redhat]

additional: also fix some deprecated APIs in this pr

[MatthewBonanni]

LGTM, thanks!