Split-K combine: correctness and wave parallelism - #534

Split out of #531.

What changes

Split-K combine, across both kernels that implement it. Everything later in the stack raises
partition counts, and both halves of this PR are defects or limits that only appear once those
counts are high.

Reduce kernel correctness

Three independent defects in the paged decode reduce kernel:

• Out-of-bounds partition loads, both paths. The loads are issued by all WARP_SIZE lanes and
  only masked afterwards. The pm/ps buffers are built max_size=True, so nothing bounds-checks
  them and the inactive lanes walk raw memory past the end — faulting only when the overrun happens
  to clear the allocator's slab. The load index is now clamped to the last valid partition; the
  existing active mask still discards the value, so this costs nothing.
• Double-counted partition sum in the slow path. partial_out holds the un-normalised numerator
  sum(p*v), so the numerator weight is w alone and part_sum belongs only in the denominator.
  The slow path multiplied by w * ls_v. The fast path was already correct.
• NaN on an empty request. Every partition sits at -inf, so the merge computes
  exp(-inf - -inf). The global max is now floored at a large finite negative sentinel, which
  yields weight 0 and can never win the max against real data.

The slow path also staged stats in LDS behind a p < _MAX_PARTS branch and sized the LDS arrays to
the exact partition count. They are now padded to a whole number of waves so every lane has a slot
and the staging loop needs no bound branch; the padded slots are never read back, since the
accumulation loops are range_constexpr over _MAX_PARTS. AOT_CONFIGS is widened to max_parts
128 and 256 to match pa_decode_dense._SPLIT_KS.

Combine parallelism

The split-K combine in the shared DualwaveSplitKCombine helper — used by both the generic and the
gfx950 dualwave kernels — became the bottleneck as split counts grew. Three changes:

• Split groups. A lane owns combine_chunk(HEAD_DIM) head-dim values, so a row needs only
  HEAD_DIM/chunk lanes. Giving the row a whole wave leaves 64/lanes_per_row groups to divide the
  split dimension between — the only axis with work left, since every lane used to walk all
  NUM_KV_SPLITS partials itself and that walk is where the time goes. The groups are contiguous
  lane ranges, so the cross-group merge is an in-wave XOR butterfly with no LDS and no barrier. Head
  dims that cannot form power-of-two lane groups (96, for instance) keep the ungrouped layout —
  correct, just without the extra parallelism.
• Streaming accumulation. Materialising every split's m and l up front and keeping them live
  costs 2 * NUM_KV_SPLITS registers and spills hard past roughly a hundred splits, with combine
  cost growing superlinearly while the attention kernel feeding it kept getting faster. Streaming
  the max and re-reading m/l during accumulation keeps the live set O(chunk), at one extra
  4-byte read per split.
• One descriptor per workspace region instead of one per split, with the O-partial loads issued
  outside any scf.if. The per-split descriptor setup put a 4-SGPR sequence between consecutive
  loads, so they could not be in flight together and the combine paid a full memory latency per
  split — linear in split count, and enough to exceed the attention kernel itself on a
  low-concurrency decode.

FLYDSL_COMBINE_BLOCK and FLYDSL_COMBINE_PREFETCH expose the block size and the loads-in-flight
depth. The block is clamped to whole waves when split groups are in use, since a partial wave would
silently drop split groups.

A note on the sentinel: an unowned split uses a large finite negative value rather than -inf,
because a real -inf makes an all-empty group's merge compute (-inf) - (-inf) = NaN and poison
every group it then merges with.

Global→LDS copy wrappers

The direct rocdl.s_waitcnt / rocdl.raw_ptr_buffer_load_lds calls in the generic kernel are
replaced by two thin wrappers, wait_lds_copies and _buffer_load_to_lds. Both keep the existing
behaviour on the default path and add an opt-in async variant behind
FLYDSL_FLASH_ATTN_ASYNC_LDS, which selects buffer_load ... lds on the asynccnt counter instead
of the vmcnt-tracked form. Same operands, different completion counter, so the swap is mechanical —
opt-in until it is proven on this kernel.

This sits here because the wrappers are in the same two files as the combine rework and everything
later in the stack calls wait_lds_copies unconditionally.

Why it is separable

The reduce fixes and the combine rework touch disjoint files — decode/pa_decode_reduce.py against
flash_attn_{utils,generic,gfx950}.py — and no other PR in the stack touches either set. The
combine rework changes only how partials are merged, not what is produced, and the copy wrappers
preserve the existing default path exactly. All three are prerequisites for everything downstream
that raises split counts.

Merge this before PR 5, which drives the reduce slow path fixed here.

Verification

The project's lint checks pass.

No dedicated test for either half; both are covered incidentally by the split-K tests. On a
supported device, run test/attention/flydsl/flash_attn_paged_splitk_test.py in full at the top of
the stack. A direct unit test for the reduce kernel would be a worthwhile addition.