Problem
The recurrence kernel uses TMA for full 16-row output tiles, but falls back to a manual scalar copy for the final partial tile so it does not overwrite the next sequence.
That fallback currently executes only on the STORE warp lane selected by elect_one_sync(). With D=128, a tail of 15 rows therefore makes one lane perform:
15 rows * 128 columns = 1,920 scalar stores
The indexing is boundary-safe, but the entire tail copy is serialized.
Affected path
csrc/smxx/fwd_kernel2.cuh, in the actual_len < CHUNK output-store branch.
Desired outcome
- Retain elected-lane ownership of output-pipeline wait and release operations.
- Retain elected-lane ownership of full-tile and final-state TMA stores.
- Distribute only the partial-tile scalar copy across the STORE warp.
- Synchronize the warp before the shared-memory stage is released.
- Cover tail lengths 1 through 15 at adjacent variable-length sequence boundaries.
A focused benchmark can be used to quantify the effect separately; the ownership and serialization issue is visible directly in the current kernel.
Problem
The recurrence kernel uses TMA for full 16-row output tiles, but falls back to a manual scalar copy for the final partial tile so it does not overwrite the next sequence.
That fallback currently executes only on the STORE warp lane selected by
elect_one_sync(). WithD=128, a tail of 15 rows therefore makes one lane perform:The indexing is boundary-safe, but the entire tail copy is serialized.
Affected path
csrc/smxx/fwd_kernel2.cuh, in theactual_len < CHUNKoutput-store branch.Desired outcome
A focused benchmark can be used to quantify the effect separately; the ownership and serialization issue is visible directly in the current kernel.