… the prologue

The in-loop mask was gated on CROSS_SEQLEN alone, on the assumption that plain
self-attention only meets a diagonal tile in the prologue and the epilogue. That
holds for full causal, but not for a sliding window: the window's LOWER edge
sweeps through the middle of the KV loop.

A q-tile's BLOCK_M rows have window floors spanning BLOCK_M-1 columns, and
split_t0 is floored to BLOCK_N and then to even parity, so 3 tiles (BLOCK_M=128)
to 5 tiles (BLOCK_M=256) starting at split_t0 still carry out-of-window columns.
Only the prologue was masked, so the rest leaked into the softmax and pulled the
result toward full causal -- worse the smaller W/S got, which is exactly how it
was reported.

vs an fp32 sliding-window reference (cos / rel_l2), JIT cache cleared:

  S=1024  W=256   bm=128  0.989033 / 0.1477  ->  0.999997 / 0.0025
  S=2048  W=512   bm=128  0.994575 / 0.1040  ->  0.999997 / 0.0026
  S=4096  W=2048  bm=128  0.999481 / 0.0322  ->  0.999997 / 0.0026
  S=1024  W=4096  bm=128  0.999997 / 0.0025  ->  0.999997 / 0.0025   (W>=S, was already fine)
  S=2048  W=512   bm=256  0.984625 / 0.1747  ->  0.999997 / 0.0026   (8-wave, stagger on)

0.999997 is the bf16 floor -- what the already-correct paths (cross_seqlen,
non-causal) score. cos against a FULL-causal reference stays at 0.88-0.89, so the
window is landing on the right boundary rather than simply being widened.

causal_mask_prologue_if_needed's guard is a runtime compare, so tiles past the
edge remain no-ops; non-SWA builds fold the branch away at const_expr time.

flash_attn_fwd.py imports 27 symbols from primus_turbo.flydsl.utils.attn_helper.
That module has been in the repo since 89b31f2 and was extended to its current
form by 03357f9 (the Meta hd64 integration) -- but b914822 deleted it, so
cloning this branch and importing the forward kernel raised ImportError.

I deleted it. b914822 was where I rebuilt the branch to strip the 32 scratch
harnesses the bwd campaign had left in the tree; I reconstructed it by keeping
"just the two product files" (flash_attn_fwd.py, flash_attn_bwd.py) and dropped
this module along with the harnesses, without noting it in that commit message.
It then lived on only as an untracked file in my working tree, which is why every
subsequent check looked clean -- `git diff --name-only` does not report untracked
files.

Restored byte-for-byte from 03357f9. My working copy had drifted 184 lines ahead
of it with knobs the fwd campaign added and never validated (LPT_QORDER,
PACKED_SOFTMAX, CLUSTER_NOP, COMPUTE_BARRIER, P_ANCHOR and a packed-softmax helper
set); that campaign stopped at round 5 when the node went away, none of it was
ever kept, and neither kernel references any of it. Restoring the original rather
than my copy keeps unvalidated code out of the branch.

Also restores benchmark/ops/training/bench_flydsl_flash_attn_fwd.py, the
square-causal forward acceptance table added in 390c0f8 and lost in the same sweep.

Verified from a clean `git archive HEAD` tree: every primus_turbo import in both
flash_attn_fwd.py and flash_attn_bwd.py resolves, and all 27 symbols the forward
kernel needs are present.

…d windowed

MI355X (gfx950), bf16, D=64, B=1, medians of same-card back-to-back runs.

S=16384 Sq=Skv full-causal, Hq=128/Hkv=16: 983.7 -> 1213.4 conv TF/s (+23.4%);
S=8192 950 -> 1177 (+23.9%). Across Meta's 10 base shapes x {full-causal,
SWA(W)}: mean full-causal 819.2 -> 1041.1 (+27.1%), mean SWA 678.0 -> 845.4
(+24.7%), and the 3 SWA : 1 full-causal blend those shapes ship at
759.8 -> 952.4 (+25.3%). O SNR 51.61 dB vs an fp32 torch reference, bitwise
deterministic, and 22 ragged-varlen shapes (full-causal plus SWA at W = 128,
512, 2048 and 8192; sub-tile / 10x-spread / non-tile-multiple lengths, 4- and
8-wave) pass at cos >= 0.999996.

Four changes, each measured on its own:

* Fixed-max softmax (DUALWAVE_SWP_FIXED_MAX, the single largest win at +13.8%).
  Softmax is shift-invariant, so causal tiles use a zero reference max and the
  final l_row normalization absorbs it: per-tile reduce_max, s - m and the O
  rescale all disappear. Per iteration that drops 28 v_max3_f32, 60 v_sub_f32
  and 23 v_pk_mul_f32, taking VALU from ~1512 to ~1050 cycles against 992 cycles
  of MFMA. A window changes which tiles are masked, not that invariance, so it
  takes the same path; split-K keeps the online path because its partial m/l are
  combined across chunks.

* MFMA row-sum (DUALWAVE_SWP_MFMA_ROWSUM). With nothing rebasing l_row mid-loop,
  each packed bf16 P slice is summed by a 16x16x32 MFMA against a ones operand
  that also folds the half-wave partner, so the permlane pair-reduce and the
  scalar add chain both go away.

* GQA sharer merge (Q_HEADS_PER_WG / WAVE_ROW_GROUPS). One CTA runs the same
  q-tile for two q-heads of one kv-head, so both sharers read a single K/V LDS
  tile and the per-CU K/V request rate per unit of work halves; merged CTAs also
  let four of the eight loop barriers fold away.

* Longest-processing-time-first dispatch. Dispatch order is the list-schedule
  order and causal work grows with q_block (a window flattens but never inverts
  that), so q_block walks descending. Derived from seq_len rather than
  grid_dim.y to stay wave-uniform.

Scored ops replace inline asm on the LDS-read and pack paths: rocdl
ds_read_b64_tr_b16 and a vector fptrunc, so SIInsertWaitcnts scores them and
places the waits itself instead of needing a hand lgkmcnt plus sched_barrier at
every consumer. That freed the registers that let the kernel sit at 3 waves/SIMD.

The kernel sits exactly on the gfx950 allocation-granule boundary at 168 VGPRs
(floor(512/168) = 3 waves/SIMD, 176 would be 2), so two independently useful
changes had to be rejected for costing 2 registers: a packed row-sum and
removing a provably-unreachable main-loop causal mask each measured -6%.

…and env knobs

No functional change: an interleaved same-card A/B of the forward at S=16384 reads
+0.12% / 0.00% / -0.07% across three rounds, SNR and determinism are unchanged, the
22 ragged-varlen shapes still pass at W = 128, 512, 2048 and 8192, and the backward
square-causal table is unmoved.

* build_flash_attn_bwd_module (962 lines) had no callers: its three modes -- "delta",
  "dq" and "fused_dq_delta" -- are the experimental path that the odo/dq/dkdv split
  replaced, and _get_bwd only ever builds the latter three.
* Removed the DQ_QSPLIT and DQ_BLOCK_KV sweep overrides; the production values are
  the fallbacks they shadowed. The forward benchmark's five knobs are likewise
  hardcoded to the configuration it is meant to reproduce.
* Dropped scope-local helpers and constants with no readers (_fadd/_fsub/_fmul in
  three kernels, _uniform_len, two vector-pair converters, the leftover Schraudolph
  scale/floor/type triple) and a duplicate math import with its _LOG2E redefinition.
* Comments: 716 -> 347 lines. Derivations, measured deltas and experiment history
  belong in the commit log, not the source; what is left states the non-obvious
  constraint (MFMA operand layout, LDS swizzle algebra, register budget, waitcnt
  hazards) in one or two sentences, and no block exceeds five lines.

Mirrors bench_flydsl_flash_attn_fwd.py (same shapes, hot-steady protocol and
MI350-equivalent verdict) for the deterministic backward. It times
flydsl_varlen_backward rather than the three kernels directly, so block_kv,
q_split and dq_block_kv always come from the deployed entry point and the number
cannot drift from what ships; that also folds in the split-K dK/dV reduce, which
costs ~1.6% against timing the kernels alone.

MI355X, B=1, Hq=128/Hkv=16, D=64, bf16, conv TF/s over the 5 backward GEMMs:
539 / 771 / 844 / 876 at S = 2048 / 4096 / 8192 / 16384, i.e. 1.99x / 1.70x /
1.51x / 1.44x H100 FA-v3 in MI350-equivalent terms -- 4/4 over the 1.4x bar.

…arks

The square-causal table only covers Sq==Skv at B=1; Meta's matrix is 10 mostly
rectangular shapes x {full-causal, SWA(W)} at B=4, and a change can move those
two regimes in opposite directions. Both scripts now print the square table first
and then the 20 configs, compared in wall time against the measured H100 numbers
for the same shapes (forward: FA-3; backward: whichever of FA-3 / FA-2 is faster
in deterministic mode, which is FA-2 on SWA). Iteration counts are derived from a
single timed call so the short and the 50 ms shapes both land near a fixed block.

MI355X, bf16, D=64, MI350-equivalent = ours x 1.2:
  forward  18/20 over 1.4x, 1.82-2.08x except the degenerate 1024x1024 pair
  backward 19/20 over 1.4x, full-causal 1.44-1.59x and SWA 1.48-3.17x

…re sizing bench iterations

_blockkv_for chose 64 for Sq <= 2048 on the reasoning that a short sequence needs
small tiles to fill the CU array. But dkdv is KV-outer: its grid is
B * Hkv * ceil(Skv / BLOCK_KV) * q_split, so Skv is what fills it, and on a
rectangular shape the two disagree. At Sq=2048, Skv=16384, B=4 the Sq-keyed choice
cost 19%: 15.03 -> 12.49 ms, i.e. 1.37x -> 1.65x the faster of H100's deterministic
backends, which was the one Meta config the backward missed. Square shapes are
unaffected (Sq == Skv), and Sq=1024, Skv=16384 also gains (4.09 -> 3.39 ms). SNR on
the rectangular shape is dq 50.3 / dk 49.2 / dv 49.2 dB against an fp32 reference and
the result is bitwise reproducible.

The benchmarks derived their iteration count from the very first call, which carries
the JIT compile, so it collapsed to the floor of 3 and the per-block event and sync
overhead landed inside the reported time. Harmless on the 50 ms shapes, but it
inflated the 1024x1024 forward by 18% (0.098 vs 0.083 ms). One warm call first.

Meta's 20 configs are now 19/20 forward and 20/20 backward over the 1.4x bar.

Short sequences were host-bound, not kernel-bound. At Hq=64, Sq=Skv=1024, B=4 the
wall clock sat at 84 us while rocprof put the kernel at 50 us and the Python enqueue
at 85 us: FlyDSL re-resolves the JIT signature on every call (inspect.Signature.bind,
the globals-drift check, per-argument cache-key construction), and the GPU waited.
The tell was that 128, 256, 512 and 1024 all took the same ~80 us -- 64x the work for
the same time, which no tiling or occupancy knob can explain.

flyc.compile returns an artifact whose call path only refreshes data pointers, so
each launcher now keeps one per scalar signature (batch, seq_len, strides ...) and
tensor pointers are updated per call. Host time per forward launch drops 93 -> 7 us.
The backward pays the dispatch three times (odo, dq, dkdv) and routes all three
through the same helper. Cache is capped at 64 entries and cleared wholesale, which
is ample for a training loop's handful of shapes.

Forward 1024x1024 goes 0.084 -> 0.045 ms and the square S=2048 811 -> 929 TF/s;
backward 1024x1024 0.232 -> 0.202 ms and square S=16384 875 -> 891 TF/s. Long
sequences are unchanged, being GPU-bound all along. Meta's 20 configs are now 20/20
over the 1.4x bar in both directions. SNR, determinism and the 22 ragged-varlen
shapes are untouched.

Issue a buffer's K and V overwrite DMAs from the same P*V cluster, so each
barrier region reads one whole K/V buffer and fills the other. That region's
single rendezvous then covers both edges of the buffer it fills, leaving two
barriers per iteration instead of four.

S=16384 causal 1211.8 -> 1217.3 TF/s on MI355X, bit-identical output.
Square-causal 4/4 and Meta's 20 configs 20/20 still clear 1.4x H100.

Rebase the SBHD [S,B,H,D] zero-permute flash-attention delta onto the new
traits-based dualwave-swp forward. gpt-oss Hq64/Hkv8/D64 pure-causal bf16:

* fwd: SBHD trait threads through DualwaveSwpTraits/cache_tag and fixes the
  per-batch base to H*D in DualwaveKernelContext (seq-step B*H*D lives in the
  runtime stride). Picks up the new dualwave fwd -> 1135 conv-TF/s @s=8192
  (SNR o/lse 51.6/84.8, det).
* bwd: self-contained 16x16x32 deterministic SBHD backward -> 859 conv-TF/s
  @s=8192 (dq/dk/dv SNR 50.3/50.4/50.8 bit-identical, det).
* pytorch entries: flash_attn_sbhd_flydsl_forward_impl / _backward_impl.

…h + block_kv-from-Skv)

SBHD [S,B,H,D] addressed natively via compile-time sbhd trait + RD_STRIDE=B*H*D
seq step; per-batch base folds only H*D. THD path (sbhd=False) takes original
else-legs. Gate: SBHD 829.2TF SNR 50.26/50.36/50.77 det; THD self-check
SNR 50.26/49.20/49.53 det.

dq is Q-outer and already streams dO; it now reduces DELTA=-rowsum(O.dO) for its
owned q rows via a 2-step ds_bpermute butterfly (O enters the freed K16 slot, the
O pack dies in the B-operand loop leaving only the f32 partial), stores it for
dkdv under the _q_store_end single-owner mask, and _defer_delta drops the standalone
odo launch. Drift-immune same-process A/B: 842.6 vs 828.9 TF (+1.64%); SNR
50.26/50.36/50.77 bit-identical to odo, det=True.

…ller)

Splits the dkdv MT=4 GEMM1 into PV_K_STEPS halves and stages the whole GQA
group's (lse, -delta) into LDS once per q-block (re-read at use points) instead
of carrying them in a ~32-VGPR register prefetch. This halves the live
S/dP/P/dS transient and drops the register carry, eliminating the spill on the
new-meta refined body (VGPR 256 spill 2 / sgpr-spill 20 -> VGPR 241 spill 0 /
sgpr-spill 8). Layout-agnostic: the compute body is shared by THD and SBHD.

Cheap levers that only transfer once the body is VGPR-lean: exp_intrin=True
(exp2-intrinsic FOLD anchor, +0.71% here; was -1.2% on the spilling body).
Levers left off after drift-immune A/B: agpr=64 (+0.18% noise), g2d=2 (-0.44%).

drift-immune A/B same-process, S=8192:
  gpt-oss SBHD bwd best-of-3 848.8 -> 860.7 TF (+1.4%), SNR 50.3/50.4/50.8 det
  THD (sbhd=False) meta bwd SNR 50.26/49.20/49.53 det -- general path correct
  THD square 4/4 + 20-config 19/20 perf, no regression

Remove the lds_ld=False full-MT register-prefetch backward arm (_head_step /
_load_ld / _qbase and the lds_ld knob): production always traces the LDS-staged
path (lds_ld=True), so the alternate arm was probe-only dead code. The removal is
byte-neutral -- const_expr(lds_ld=True) traced identically to the now-unconditional
body. Also drop the unused fast_exp2 impl-layer param (never forwarded to the
backward). Comment cleanup: keep WHY-only rationale, no bench numbers in source.

Compress the FlyDSL forward footprint (flash_attn_fwd.py + attn_helper.py)
without changing a single traced instruction. Every step was gated by a
byte-exact diff of the traced MLIR (00_origin.mlir) AND final ISA
(21_final_isa.s) across five causal configs (sbhd / thd / swa / lse / d128,
S=8192) vs frozen baselines -- all SAME. Code is ruff-check + ruff-format
clean (line-length 110), so the counts below are the push-compliant ones.

What changed (all provably trace-neutral):
- Drop dead code unreachable on the causal-only deployment path: the entire
  online/lazy-rescale softmax machinery and the non-MFMA-rowsum VALU-fold
  branches (deployment hardcodes causal=True => fixed_max => mfma_rowsum),
  plus dead split-K/paged/vectorized-KV traits and the debug-counts feature.
- Remove the now-unused dualwave_swp_lazy_rescale knob and the BLOCK_N_OUT
  trait (both were write-only after the rescale path was deleted), including
  the two shipped bench call-sites that still passed lazy_rescale=True.
- Flatten the mixin context-class hierarchy: fold single-use helper methods
  and free functions into their sole callers; merge helper subclasses and the
  init_* setup methods into one DualwaveKernelContext._setup.
- Dedup: shared K/V DMA loaders, _batch_views for q/o + k/v view setup,
  _store_o_128 for the O epilogue, _fill_defaults for launch/compile, and
  fold _sched_barrier_exp_pairs into _sched_barrier_pairs (mask param).

Combined fwd 3848 -> 2092 lines (attn_helper 1279 + flash_attn_fwd 813),
below the backward reference (flash_attn_bwd.py = 2093). No numeric/
scheduling/VGPR change: traced MLIR + ISA byte-identical.

Add true variable-length (packed document masking) to the hd64 flash-attn
backward: ragged segments, per-segment bottom-right causal, cross-segment
masking, and causal-SWA. Kernel + impl layer only (no public dispatch).

- VARLEN build-trait const_expr branch on dq/dkdv (in the JIT cache key);
  VARLEN=False keeps the rect16 compute loop byte-identical.
- Shadow seq_len_q/k to the per-segment length so base/SRD/loop-bounds
  follow the segment; packed tok_base addressing; grid tiles by max_seqlen
  with out-of-segment early-exit.
- Split-K workspace packed [q_split,total_kv,Hkv,D]; packed lse/delta
  [total_q,Hq]; driver dispatches uniform (placeholder cu) vs ragged.
- impl layer: dispatch uniform->rect16 / ragged->varlen; unified packed
  lse contract (uniform transposes internally).
- Fix pre-existing causal-SWA off-by-one (W+1 keys) across the bwd.

Add a gfx950-only pytest exercising the flydsl hd64 THD ragged / block-causal
(packed document-masking) fwd+bwd over ragged and uniform segment layouts x
{full-causal, causal-SWA}: SNR vs a fp32 block-diagonal reference plus bitwise
determinism. Add a separate ragged backward perf table (not on the uniform
20-config 1.4x-H100 metric, since ragged tiles by max_seqlen).

…len < kv_len)

The dq kernel splits its causal kv-loop at q_start + causal_offset: unmasked
below, masked across the diagonal. Both loops walk one BLOCK_KV tile grid
anchored at 0 (the LDS hand-over prefetches kv_start + BLOCK_KV), so that split
point has to sit on the grid. With Sq == Skv it always did -- q_start and the
zero offset are both BLOCK_KV multiples -- but under context parallelism the
bottom-right offset Skv - Sq is arbitrary, and an unaligned split made the last
unmasked tile cross the diagonal unmasked, made the masked loop redo that tile's
tail (double-counted kv), and desynced the LDS hand-over (the masked loop's
first tile was never DMA'd, so it consumed the previous tile's K/V). dQ came out
wrong; out/dK/dV were unaffected. Floor-align the split; the diagonal tile goes
to the masked loop, whose mask is elementwise.

Also route uniform-q + ragged-kv through the ragged path: the rect16 fast path
takes a single (Sq, Skv) pair, so it asserted on uneven CP documents.

Verified on MI355X (gfx950), Hq=8 Hkv=1 D=64 bf16: the reported Sq=96 sweep is
now at the bf16 noise floor for every Skv-Sq (was wrong for every offset that is
not a multiple of 64); 108 dense Sq x offset x window configs and 6 ragged
q < kv document layouts match an fp32 reference. Bit-exact on Sq == Skv (15
uniform/SWA shapes, out/dq/dk/dv hashes unchanged) and the bwd bench is
unchanged within run-to-run noise.

Two cache keys made an e2e run rebuild and recompile the attention kernels on
every micro-batch whose document count or max seqlen differed.

_get_bwd keyed the module cache on batch_size, and the THD path passes num_seq
there. batch_size is baked into the kernel only for SBHD, where it sets the
seq-step stride RD_STRIDE_* (batch is interleaved inside the seq axis); THD
takes it as a runtime fx.Int32 argument, so the module does not depend on it.
Normalise it to None for THD before building the key -- SBHD is unchanged and
THD emits exactly the module it did before, only the key is coarser.

Both launch caches (bwd _cached_launch, fwd _launch) then keyed the compiled
artifact on the scalar argument values. Also unnecessary: the launch entry
declares those scalars as fx.Int32 runtime arguments and computes its grid from
them inside the compiled launch, and FlyDSL's own artifact cache folds only the
*type* of a scalar into its key (Numeric.__cache_signature__ returns (type,)).
It was not free either -- flyc.compile ends in a real launch
(JitFunction.__call__ runs the kernel), so each miss cost a duplicate
fwd/dq/dkdv execution on top of the artifact rebuild. Key on the argument types
instead. Shape-driven module differences (q_split / block_kv / dq_block_kv /
varlen / sbhd) are unaffected: those are separate modules, each with its own
cache dict.

Measured on MI355X (gfx950):
- 12 varlen steps with drifting num_seq / max_seqlen: 30 flyc.compile calls -> 5
  (three at the first step, two more when a step's segments happen to be equal
  and take the uniform rather than the ragged module), steady state zero.
- 20 cases over B in {1,2,3,5} x (Sq,Skv) covering 512/1000/2048 square and
  rectangular: 60 -> 4, and all 80 output hashes are byte-identical to the
  value-keyed build -- one artifact, grids from tens to hundreds of work-groups.
- Compile wall-clock saved depends on FlyDSL's disk cache: warm, 6.5s -> 0.1s
  over the 12 steps; cold, both arms pay the one real compile and land near
  5.5s. The steadier win is the duplicate launch each miss used to cost.

tests/pytorch/ops/test_attention_varlen.py -k "flydsl or no_grad": 40 passed,
plus a 114-case q<kv sweep against an fp32 reference -- every shape in one
process, so the single artifact is reused across all of them.

Block-diagonal backward launches early-exit workgroups for every empty q-segment, cost growing linearly with segment count. Empty q-segments contribute zero dk/dv, so compute only non-empty segments as rect16 single-seg sub-problems when sparse (<=1/8). CP 64-seg/1-nonempty: 0.740 -> 0.153 ms, 2.8x faster than FAv2; dense ragged layouts fall through unchanged.