pa_decode_reduce has two paths: a fast one where lane l owns partition l
(max_parts <= WARP_SIZE), and a slow one that stages per-partition statistics
through LDS. Only the fast path had ever run -- _SPLIT_KS in pa_decode_dense
topped out at WARP_SIZE, so _FAST was always true -- and the slow path did not
work:

  - It computed `p = step * WARP_SIZE + lane`, then used `const_expr(p <
    _MAX_PARTS)` and `arith.constant(p, ...)`. `p` depends on `lane`, so it is a
    runtime value and tracing fails with "cannot evaluate dynamic 'Boolean' as
    Python bool". The path could not compile at all.

  - The output accumulator weighted partial_out by `w * part_sum`. partial_out is
    the un-normalized numerator sum(p*v), so part_sum belongs only in the
    denominator; the fast path is correct here, accumulating `vals * norm_w`.
    Uncorrected this is wrong by a factor of the per-partition softmax sum, which
    measured two orders of magnitude outside bf16 epsilon.

The staging loop now avoids the dynamic branch rather than converting it: the LDS
arrays are padded to a whole number of waves so every lane has a slot to write,
and the load index is clamped instead, duplicating a valid entry into the
padding. Padded slots are never read back -- both accumulation loops are
range_constexpr over _MAX_PARTS.

Prerequisite for raising the decode split-K cap above one warp's worth.

The paged path in flydsl_flash_attn_func runs flash_attn_dualwave_swp_gfx950,
which maps the MFMA M dimension to query *rows*. On a single-token decode only
one M-row carries work: it is a prefill kernel running a decode shape.

Measured on MI350X (gfx950, bf16), the kernel issues orders of magnitude more
MFMA work than the maths needs, and assembling a tile that is almost entirely
padding then dominates -- thread-trace attributes 50.4% of stall cycles to
s_barrier and only 1.5% to VMEM. Achieved KV bandwidth is flat across a wide
range of problem sizes, so the kernel is instruction-bound, not memory-bound,
and occupancy work cannot reach it.

decode/pa_decode_gfx950.py already had the right shape: it packs the query heads
sharing a KV head onto the MFMA M-axis, so the matrix core stays full at any
query length. It was unreachable from flydsl_flash_attn_func, which never
imported the decode/ subpackage, and it read only a dense cache.

  - pa_decode_gfx950.py takes a `paged` compile-time flag plus block-table and
    page_size arguments. K/V addressing moves from a linear token offset to a
    per-tile page indirection. The page size is a multiple of the tile step, so
    a tile never straddles a page and the block-table read is uniform across the
    warp -- one scalar load per tile, not per lane. The split-K chunk is rounded
    up to the tile step to preserve that alignment; the trailing empty splits
    this can produce are already neutral in the combine pass (-inf max, 0 sum).

  - pa_decode_dense.py gains pa_decode_paged_launch and raises the split-K cap,
    with a minimum chunk per split. Past that cap the combine pass dominates and
    was slower on every shape tried. In a controlled A/B the higher cap is a
    clear win on long caches, but across the sweep it is within run-to-run
    noise, because most shapes never reached the old cap.

  - flash_attn_interface.py routes eligible dense paged single-token shapes on
    gfx950 to the new launcher. MSLK_DISABLE_PAGED_DECODE_HP=1 forces the old
    path. Varlen, gappy paged, return_lse and the vectorized cache layout are
    excluded and keep the dualwave path. `causal` is deliberately not a
    condition: with one query token under bottom-right alignment, causal and
    full attention cover the same range.

Over the paged decode sweep: MFMA instructions per dispatch drop by nearly two
orders of magnitude, more cases run, and page-table capacity failures fall
sharply -- the decode kernel reads the block table from memory rather than
through an LDS window, so that limit does not apply to it. Accuracy also
improves slightly on the routed shapes.

Only single-token decode is routed. The M axis holds one tile of (token, head)
pairs, so longer query blocks keep the dualwave path and show only the
pre-existing auto-split-K gain. Covering them needs multiple M-tiles inside the
KV loop, which costs loop-carried accumulator registers per tile and may spill.
Left as follow-up; note that the M axis is not full even at one query token, so
the first extra tokens per tile come for free.

Tests cover selection and decline, the kill-switch, equivalence with the
dualwave path, and a context beyond the old page-table cap.

The head-packed decode kernel put query *heads* on the MFMA M-axis and
pre-loaded Q once outside the KV loop, so flash_attn_interface could only route
single-token decode to it. Everything longer fell back to
flash_attn_dualwave_swp_gfx950, whose M-axis is query rows and which therefore
runs at a small fraction of M-utilisation on decode shapes -- the same MFMA
instruction count per dispatch however short the query block is, because a few
query rows still occupy only a few of its M slots.

M now holds (query token, head) pairs laid out as m = qtok*ratio + head, with
ceil(Sq/T_PACK) M-tiles (T_PACK = MFMA_M//ratio) looped *inside* the KV tile
loop. That nesting is the point: the K fragment (indexed by st, g, tok_lane) and
the V fragment (dpass, grp, tok_lane) do not depend on M, so each is loaded once
and fed to every tile. Extra query tokens cost MFMA issue, not memory traffic.

Measured on MI350X (gfx950): quadrupling the query block costs 2x MFMA and 1.33x
LDS at an unchanged wave count and unchanged KV traffic, since only the P staging
planes scale. Against the dualwave path that is a ~32x reduction in MFMA
instructions at the same query block.

  - MAX_M_TILES caps the tiling at 2. Each tile carries loop-carried f32
    accumulator that scales with the head dimension, and the deep tilings a long
    query block would need are left out rather than gated on a guess.

  - gqa_ratio and seqlen_q join the compile key: T_PACK and M_TILES are loop
    bounds, so they cannot stay runtime values. The ratio must divide MFMA_M for
    the tiling to be even; ratios that do not are declined to the existing
    fallback.

  - Causal masking moves inside the kernel and becomes per query token:
    min(t_end, t_full - Sq + qtok + 1), bottom-right aligned. With one query
    token this collapses to t_end, which is the degeneracy check on the term.

  - Split-K partials fold the query axis into the head axis. pa_decode_reduce is
    generic over that axis and needed no change. The folding only matches the
    output layout when there is a single KV group, and the per-token causal
    bound is the paged caller's contract but not the dense decode op's, so
    multi-token is restricted to single-group paged launches. Dense stays at one
    query token.

Over the paged decode sweep, the newly routed multi-token shapes carry the gain
and more cases run; single-token is unchanged, and the longest query blocks still
take the dualwave path. Regression guard: the single-token MFMA count is
unchanged, exactly.

Accuracy is well inside bf16 epsilon against an fp32 bottom-right-causal
reference, including cache lengths that are not a multiple of the tile step,
since a ragged tail is where an off-by-one in the per-token causal bound would
surface.

Two things that came out of those ragged cases, neither addressed here:

  - The dualwave paged path is above bf16 epsilon on a non-tile-aligned cache
    length where this kernel is two orders of magnitude tighter. That looks like
    a real masking defect. The performance sweep never exercises it, because
    every context there is tile-aligned.

  - Per-shape run-to-run variance on this hardware reaches ~25%. A few long
    query-block shapes read as regressions against main, but they are on the
    dualwave path, which this change does not touch, and they moved by a similar
    margin between two runs of identical code. Individual shape numbers there
    should not be over-read; only the aggregate is reliable.

Long query blocks remain the largest gap and need deeper tiling. Note for that
work: the register budget is what decides whether the larger head dimension is
reachable at all.

MAX_M_TILES was 2, set from hand-derived VGPR arithmetic rather than
measurement, which capped the head-packed decode kernel at a short query block
and left every longer non-SWA case on query-row-tiled kernels. Separately,
_num_kv_splits_heuristic judged "is the device full?" by counting workgroups,
which silently assumes every kernel has the same CTA width.

Measured the register cliff instead of estimating it (rocprofv3 --kernel-trace,
VGPR_Count + Scratch_Size, one compile per tile count). The no-spill point
differs by head dimension: the larger one spills several tiles earlier, and
scratch appears abruptly rather than gradually.

Benchmarking the newly-compilable configs then showed that fitting is not
sufficient. Head-packed *lost* on every deep-tiled configuration at the larger
head dimension -- the ones that spill -- and on one shape that had too few CTAs
per CU. Occupancy falls steeply as tiles are added, and on a bandwidth-bound
decode kernel that can outweigh removing MFMA waste.

So the gate needs two independent conditions, both derived from tensor shapes
and queryable hardware -- no tuned table, consistent with the heuristics-by-
default policy in mslk/flydsl/autotune.py:

  - Register budget: max_m_tiles(head_dim) from the measured no-spill point.
  - Parallelism floor: CTAs >= 2 * CU, applied only from three tiles up.
    Shallow tiling still holds high occupancy and was measured winning at the
    lowest concurrency, so the floor must not reject it. An earlier version
    applied it one tile sooner and silently took a working configuration off the
    fast path; the numeric validation caught it, by the error signature
    reverting to the dualwave path's.

_num_kv_splits_heuristic gains an optional waves_per_cta. The paged caller
passes the light kernel's value; left at 0 the original workgroup-based test
runs, so the dense path is unchanged. Where the old test counted enough
workgroups to call the device full and declined to split, those were narrow
CTAs, and the measured wave count showed the device was in fact the most starved
it gets.

Over the paged decode sweep, more cases run and every concurrency improved,
most at the low counts where the parallelism floor and the wave-aware split test
both bite. The worst remaining shapes are the deep-tiled ones at the larger head
dimension, which the spill limit excludes by design.

Accuracy stays well inside bf16 epsilon against an fp32 bottom-right-causal
reference across the full range of query block lengths at both head dimensions,
including cache lengths that are not a multiple of the tile step.

A handful of shapes are >25% slower than the previous sweep. Most are on paths
this change does not touch, which points at run-to-run variance (measured at up
to ~26% per shape on this hardware); the rest are on the dualwave path, where
the split heuristic genuinely did change. Worth separating before reading
anything into those rows.

The M-tile budget declines the larger head dimension past a few tiles because
the kernel spills there, so the longest query blocks fell through to the
query-row-tiled path. Covering such a block in two passes halves the tiles each
pass needs, which brings it under the budget: one extra KV read buys a
configuration that fits.

Grouping needs no kernel change. With Sq query tokens against a cache of length
L, token i attends to [0, L - Sq + i + 1). Split into G groups of S = Sq/G;
group g holds tokens g*S .. g*S+S-1 and is launched with

    seqlen_k_g = L - (G - 1 - g) * S

whose in-call bound L_g - S + j + 1 equals L - Sq + (g*S + j) + 1 exactly, so
the kernel's existing per-token causal masking is already correct. Derived in
query_group_seqlen and verified three ways: exhaustively over every (L, Sq, G,
g, j) used, numerically against an fp32 reference, and on ragged caches.

Selection stays derived. Grouping is only used to rescue a shape the register
budget would otherwise reject; where a single pass already fits, it measured
inside run-to-run noise and is not worth the extra KV traffic. G is 2 -- G=4 was
slower on nearly every shape tried, the extra launches costing more than the
occupancy they buy. MSLK_DISABLE_PAGED_QGROUPS=1 restores the previous
behaviour.

Fixes a latent bug this was the first caller to hit: the single-split epilogue
addresses the output with Q's strides, so a non-contiguous Q -- a slice of the
query block, here -- whose stride(0) differs from the freshly allocated output
silently corrupts every batch element after the first. It showed as an absolute
error four orders of magnitude out. The grouped path passes contiguous slices
and the precondition is now documented on pa_decode_gfx950_launch. The split-K
path was unaffected, because it computes partial-buffer offsets itself.

The rescued shapes improve; the bands the selection deliberately protects --
shallow tilings, and deep tilings that already fit in one pass -- measured
unchanged, which is the point of gating it. Only a small number of shapes are
affected, so the aggregate barely moves.

Against the pre-change path the mechanism is instruction count, not occupancy as
first supposed: issued MFMA and LDS instructions both drop by roughly an order
of magnitude at an unchanged wave count. Occupancy was the mechanism only
against a deep-tiled head-packed baseline, which the shipped selection never
reaches for these shapes.

Odd query block lengths stay on the old path: equal-span grouping cannot divide
them.

Brings in the FlyDSL paged tuning from avbokovoy#10 (e23fdca) and makes it
the primary path. Measured head-to-head over the paged decode sweep, same
harness, back-to-back, their approach is better where the two overlap. The two
branches were equal on single-token decode and within noise on short query
blocks; the entire difference was the long query-block family, where this branch
had no answer.

Two of their ideas do the work:

  - GQA head packing folds query heads into the M dimension inside the generic
    paged kernel (a Q reshape plus a Q_PACK_QLEN trait), removing the fan-out
    where each query head re-streams its KV head's whole range -- measured
    issued load traffic several times the distinct DRAM traffic. That is the
    same fix as this branch's head-packed decode kernel, but it covers a far
    wider range of query block lengths than the M-tile budget allows.

  - Split-K sized by KV-chain length rather than occupancy. Each workgroup walks
    its range as a serial chain of dependent loads at about one memory latency
    per tile, so a grid that already covers every CU can still be latency
    stalled. Occupancy cannot see that; the wave-counting heuristic this branch
    added was a proxy for it and stops helping exactly where the device looks
    full.

Routing: the head-packed decode path is now gated on `not _gqa_packed`, so it
serves only what packing declines (FLYDSL_PAGED_GQA_PACK=0, over-long query
blocks, non-causal multi-token, MHA). Its tests select it explicitly rather than
depending on which mechanism wins, and a new test pins the precedence.

Conflict resolution note. The first resolution left two split selections in the
paged path: theirs, then this branch's `_auto_paged_kv_splits` overwriting it.
That discarded the chain-length sizing, pushed shapes off the light kernel onto
dualwave, and produced an absolute error above bf16 epsilon on non-tile-aligned
cache lengths. Their branch alone was clean on the same shapes, which located
the fault here rather than upstream. The override is removed and the merged tree
now matches them exactly. The performance sweep cannot catch this, because every
context in it is tile-aligned.

38 tests pass. Case coverage over the sweep is unchanged from theirs.

The merge let GQA head packing own every shape it could take. Both mechanisms
remove the same fan-out, but only the head-packed decode route honours a
per-request seqlen_k, so preferring packing re-exposed an upstream defect.

The light/dualwave paged launch hands the kernel a scalar `seq_len_kv` -- the
batch maximum -- and never the per-request tensor, so every request is masked to
max_seqlen_kv. A request shorter than that maximum attends to stale cache, with
error growing in the gap. Measured at B=8/ctx=32000/D=64/Sq=13 against an fp32
reference, varying only the declared batch maximum:

  seqlen_k   max_seqlen_kv    before      after
     32000           32000  2.49e-04   1.47e-04
     31999           32000  2.03e-03   1.45e-04
     30001           32000  1.29e-02   1.41e-04
     29952           32000  1.30e-02   1.51e-04

bf16 epsilon is 7.81e-03, so the last two were wrong answers. Note 29952 is
tile-aligned: the trigger is the gap, not alignment, which corrects an earlier
reading of this as a ragged-seqlen problem.

The defect is not new and is not this branch's. Upstream main (d204cd8) shows
1.28e-02 and 1.30e-02 on the same cases, as does avbokovoy#10. What changed
is which route takes these shapes: the head-packed path forwards seqlen_k to the
decode kernel as seq_positions and was correct throughout, so routing away from
it made a latent upstream bug reachable.

Eligibility now resolves before the packing block rather than after, because the
packing rebinds H and Sq and the gate must see the caller's real GQA ratio and
query length. Packing takes what the head-packed budget declines.

This trades measured throughput for correctness, and the trade is not small.
Over the 268-case paged grid, gap versus the B200 baseline:

  geomean            1.23x -> 1.83x
    q_len1  (n=24)   1.39x -> 1.43x
    q_len4  (n=30)   1.24x -> 1.45x
    q_len16 (n=38)   1.14x -> 2.58x

q_len16 carries almost all of it: those shapes now take the head-packed kernel,
which is slower there than packing. The benchmark grid cannot see the defect
being avoided -- it sets seqlen_k == max_seqlen_kv on every case -- so the grid
shows only the cost and none of the benefit.

A caller that knows its batch has uniform KV lengths could safely have the
faster route. That wants an explicit opt-in rather than a guess, since deciding
it here would need a device-to-host sync and would be illegal under CUDA-graph
capture. Left as follow-up.

40 tests pass, including a regression guard that varies only max_seqlen_kv and
requires the answer not to change.

…red"

This reverts commit 0ee3fd3, and records the
defect that revert re-exposes as an executable guard.

That commit routed paged decode shapes to the head-packed kernel because it is
the only paged route that honours a per-request seqlen_k. The correctness
reasoning still holds; the price turned out to be too high to pay by default.

Measured over the paged grid, 92 cases with a B200 reference, geomean gap:

                    pack-first   head-packed-first   reverted
  q_len1  (n=24)         1.390x              1.428x     1.389x
  q_len4  (n=30)         1.242x              1.450x     1.241x
  q_len16 (n=38)         1.139x              2.585x     1.134x
  all     (n=92)         1.234x              1.834x     1.231x

q_len16 carries nearly all of it: those shapes moved onto a kernel that is
2.3x slower there than GQA head packing. Reverting returns every band to the
pack-first baseline within ~0.5%.

What this accepts. The light/dualwave paged launch hands the kernel a scalar
seq_len_kv -- the batch maximum -- and never the per-request tensor, so a
request shorter than that maximum attends to stale cache. Max abs error vs an
fp32 reference, varying only the declared batch maximum:

  seqlen_k  max_seqlen_kv     error
     32000          32000  2.49e-04
     31999          32000  2.03e-03
     30001          32000  1.29e-02
     29952          32000  1.30e-02

bf16 epsilon is 7.81e-03, so the last two are wrong answers. 29952 is
tile-aligned, so the trigger is the size of the gap, not alignment.

This defect is upstream and predates this branch: main (d204cd8) and
avbokovoy#10 both produce the same values. Reverting restores parity with
them rather than introducing anything new.

Callers must therefore pass max_seqlen_kv equal to the per-request seqlen_k.
Rather than let that constraint live only in the log, the regression test from
the reverted commit is kept as test_per_request_seqlen_k_is_honoured, marked
xfail(strict=True): it still runs, still reproduces the defect, and fails the
suite if the behaviour is ever fixed, so the guard becomes a notification
instead of dead weight. The constraint is also stated at the GQA-packing gate
in flash_attn_interface.py. 38 pass, 1 xfail.

The real fix is for the launch to forward the seqlen_k tensor it already has.
Short of that, a seqlen_k_uniform opt-in would let ragged callers take the
slower correct route without a device-to-host sync, which is illegal under
CUDA-graph capture. Both are follow-ups.

Two module-level assignments followed a function definition with a single
blank line. flake8 E305 wants two:

  decode/pa_decode_gfx950.py:86   _FX_DTYPE after max_m_tiles()
  flash_attn_interface.py:657     _PAGED_BLOCK_N_OUT after the split heuristic

Whitespace only, no behaviour change. Verified with the CI invocation
(flake8 --ignore=E501,W503,E203) over every tracked Python file outside the
external/ submodule: no remaining hits.

The window mask lives in the generic kernel's N32 path, which the paged
light route already builds -- window_left was simply never threaded
through to it. Plumb it, and route windowed shapes to the light kernel.

The other paged routes (GQA packing, MTP two-pass, head-packed decode,
dualwave) mask by their own rules and have no window term, so they drop
a window silently rather than failing. Each now declines windowed shapes.

Also scrubs benchmark, config-shape and vendor references from comments
added on this branch.

46 tests pass, 1 xfail.

* Cut paged decode fixed cost and honour per-request seqlen_k

* Stride-pack multi-token paged decode and widen the split-K combine

…_correction

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…-optimization

Clean up and harden the FlyDSL paged attention path on MI350.