Parallelize partial output-tile stores across the store warp

[morluto]

TL;DR

Keep output-pipeline and TMA ownership on the elected STORE-warp lane while distributing the boundary-safe partial-tile copy across all 32 lanes.

Closes #14.

Motivation

Full 16-row output tiles use TMA, but the final partial tile uses a manual scalar copy to avoid writing beyond the current sequence.

The manual path currently runs only on the lane selected by elect_one_sync(). For D=128, one lane can therefore perform up to:

15 rows * 128 columns = 1,920 scalar stores

This preserves sequence boundaries but serializes the entire tail copy.

Store-warp ownership

Before

lane 0: wait -> copy elements 0..n-1 -> release
lanes 1..31: --------------------------> idle

After

lane 0: wait -> warp sync -> lanes 0..31 copy -> warp sync -> lane 0: release
                              idx = lane, lane + 32, ... < n

n = actual_len * D

Pipeline wait and release remain elected-lane operations. Only the boundary-safe manual copy is distributed across the warp. Full-tile and final-state TMA stores are unchanged.

Changes

• Synchronize the STORE warp after the elected lane completes the pipeline wait.
• Distribute partial-tile elements across all 32 lanes using a lane-strided linear index.
• Synchronize before the elected lane releases the shared-memory stage.
• Keep full-tile TMA and final-state TMA operations elected-lane-only.
• Add variable-length coverage for every partial tile length from 1 through 15.
• Initialize test output with NaNs so missed stores remain observable.
• Add a focused tail-store benchmark mode for future measurements.

Suggested review order

1. csrc/smxx/fwd_kernel2.cuh — STORE-warp ownership, synchronization, and lane-strided copy.
2. tests/test_fwd_full.py — sequence-boundary and tail-length coverage.
3. benchmarks/bench_fwd.py — isolated benchmark case.

The synchronization points deliberately surround only shared-memory consumption: after the elected-lane wait and before the elected-lane release.

Test coverage

The new variable-length case places tail lengths 1 through 15 next to one another, exercising both the manual store and sequence-boundary indexing.

Checks

• python -m py_compile benchmarks/bench_fwd.py tests/test_fwd_full.py
• git diff --check origin/master...HEAD