Description

The gfx950 flash-attention backward leaves work standing in windows where the matrix
pipe is idle: the dS pack and its LDS publish, GEMM3's dQ atomic burst, GEMM1's
A-fragment batch, and the second kv half's operand reads. None of it can be deleted --
every value descends from the same head-step -- but all of it can be issued where a
MFMA run already covers it.

This is a performance-only change to primus_turbo/flydsl/attention/flash_attn_bwd.py.
Every change is order-only: the same MFMAs on the same operands in the same
per-accumulator order, 16-bit throughout, no extra register, and no change to what the
op computes or to any public signature. dQ stays on the atomic path it already used, so
its existing non-determinism is unchanged (neither improved nor made worse).

Measured on MI355X against this PR's merge base, two independent runs, geomean over
seven shapes: 1.0568 / 1.0538. Table below.

Fixes # (n/a -- no issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Performance only: no API change, no numerical-semantics change, no new public knob.

Changes

dS publish → GEMM2's run

• g2_fill: hand the dS pack and its LDS publish to GEMM2, emitted between the first
  d-tile's dV and dK groups, so the rest of that run covers them while they still
  precede the dK MFMAs that read them.
• g2_weave: split the parked chain per kv tile -- tile nt's block feeds only tile
  nt's dK MFMA -- giving each block a matrix pair to issue under. A sched_barrier
  that lets MFMA, and only MFMA, cross pins the split the scheduler would otherwise undo.
• g2_hs: hand GEMM2's closing d-tiles across the dS fence, where they cover GEMM3's
  transpose-read burst; they read neither of its operands.

dQ stores

• g3_st_hs / g3_st_g2 / g3_st_w: an undeferred GEMM3 ends its head-step with an
  uncovered atomic burst. Hand the stores to the next head-step and issue them a pair of
  dQ rows at a time between GEMM2's d-tile groups.
• g3_qcarry: carry a deferred GEMM3 across the q-loop trip as well, so the last head's
  dQ rides the next trip instead of a bare [drain, barrier, GEMM3] tail.

Reads

• g1_pf: issue GEMM1's A fragments from inside the previous q-half's GEMM2 -- they were
  the one read family with no prefetch ring, so their batch opened each q-half with the
  whole LDS latency and no matrix work in flight.
• g2_kvfuse / g1_kvfuse: emit the two kv halves' d-tile loops as one. The halves
  contract the same Q/dO transpose fragments, so each is read once and the MFMA run
  standing between two reads doubles.

Dispatch / addressing

• _a16_qsplit: read the launch plan (B, Skv) instead of assuming a D=64 dispatch covers
  the machine, and widen the cut only when the GQA group is wide enough to use it.
• _wsq_ilv: band grouping used a conservative span; use the distance at which two bands'
  q BLOCK ranges stop touching. A windowed band, whose block_kv equals its window, was
  paying most for the slack.
• FQ_PAIR: walk each band phase-shifted by its own index, so a band's far half-tile and
  the next band's near half-tile share a tile and the second read hits in L2.

Measurements

MI355X (gfx950), bf16, S=8192, against the merge base 4a96767d. Both arms are loaded as
two modules in one process and timed ABBA at rep level, so clock drift and DPM state
divide out; an unmodified tree reads 1.000 on this ruler (A-vs-A over the full table:
1.0010). Two independent runs:

Backward FLOPs are 10 * B*Hq*S^2*D times the fraction of the S^2 block computed
(1/2 causal, (W+1)/S windowed); each cell shows run 1 / run 2. The windowed cells' low
TF/s is shape-determined, not a regression -- at W=128 only 1.6% of the S^2 block is
computed, so the kernel is dominated by launch and boundary work.

Two notes on reading the table honestly:

• The four D=128 cells gain 3.3-7.1%, which for most of them is within that cell's own
  run-to-run residual (2.2-4.9%). Eight of eight readings across two runs are positive,
  so the direction is solid, but do not read a single D=128 cell's magnitude as exact.
• The D=128 windowed cell is reported flat on purpose. On byte-identical source its
  two arms read up to 1.7% apart, so nothing under 2% on that shape is a measurement. It
  is the one configuration that is _fuse_d128 with the atomic dQ path off, feeding the
  DRAM-bound split-K fold; the two arms allocate their own large dQ partial workspace and
  do not land on the same pages.

Correctness

All three gradients against an fp32 eager attention, on the shipped default path:

against a 45 dB floor. dQ remains atomically accumulated and therefore
non-deterministic, exactly as before this PR.

Checklist:

• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation (n/a -- no doc covers this kernel)
• My changes generate no new warnings (ruff check and ruff format --check clean; no new
  compiler diagnostics across the JIT builds exercised by the benchmarks above)
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

The last two are not checked: this PR adds no test, and I have not run the repository's
test suite locally. Correctness was verified out of tree with the SNR gate above on both
head dims, causal and windowed. Happy to run tests/ before merge if you want that box
ticked.