perf(sm100): async-stage incomplete packed-varlen tiles

[SakiNana7qi]

Fixes #43.

Summary

• Keep TMA as the fast path for complete 64-row tiles.
• Stage incomplete packed-varlen tails with predicated, zero-filling cp.async, so a tail never reads into the following sequence.
• Split tail staging across the two existing data-producer warps: Q/K in one warp and V/A plus scalar beta/gamma in the other.
• Apply the same policy to prepare_h: K and V/A use separate producer warps for incomplete tails, while complete tiles retain TMA.
• JIT-specialize batches with no incomplete tile, so their generated kernels omit the async-tail body.

The prior guarded elementwise tail loops made every short sequence pay a serialized staging cost. This change preserves the existing CTA shape, chunk size, tensor layouts, and complete-tile data path; it only changes how the final incomplete tile reaches shared memory.

B200 benchmark

Packed varlen, Hq=Hk=Hv=12, head dimension 128, BF16 Q/K/V, FP32 decay, TileLang 0.1.9. Times are CUDA-event medians from four balanced AB/BA rounds; backward is incremental backward latency.

For the issue workload (N=4096, L=32):

Variant	Forward (ms)	Backward (ms)	Forward + backward (ms)
Upstream main	12.7844	17.0214	29.7932
This PR	3.3603	9.9497	13.3112

The real-L32 path improves by 3.80x in forward, 1.71x in backward, and 2.24x for the combined operator.

The boundary sweep uses

L = {1, 2, 4, 8, 16, 24, 31, 32, 33, 40, 48, 56, 63, 64, 65, 127, 128, 129, 255, 256, 257, 511, 512, 513, 1023, 1024, 1025, 2047, 2048, 2049}

with N = min(4096, floor(131072 / L)). Each checkout uses an independent TileLang cache; checkout order and length order are reversed symmetrically across the four rounds.

Exact full-tile points remain at main-level training latency:

L	Main forward (ms)	PR forward (ms)	Forward change	Main fwd+bwd (ms)	PR fwd+bwd (ms)	Fwd+bwd change
64	1.4664	1.4636	-0.19%	6.7133	6.6728	-0.60%
128	1.1025	1.1115	+0.82%	4.9141	4.8966	-0.35%
256	0.9253	0.9328	+0.81%	4.0523	4.0713	+0.47%
512	0.8419	0.8474	+0.65%	3.4484	3.4380	-0.30%
1024	0.8142	0.8184	+0.52%	3.2369	3.2272	-0.30%
2048	0.8306	0.8355	+0.59%	3.2460	3.2336	-0.38%

Across these six full-tile points, combined latency changes by -0.60% to +0.47% (median -0.33%). Across all 30 points, the range is -67.72% to +0.99%; the largest positive observation is the nearly full-tile L=2049 case. The remaining sub-1% deltas are shown rather than hidden behind an acceptance threshold.

For tails immediately below successive 64-row multiples, combined latency improves by 67.72% (L=63), 59.53% (L=127), 48.14% (L=255), 35.73% (L=511), 23.18% (L=1023), and 13.39% (L=2047).

Validation

• All 8 sweep workers completed: 4 rounds on clean upstream main and 4 rounds on a clean PR checkout, with 30 cases per worker.
• In the 30-point sweep, sampled output and q/k/v/decay/beta gradients have zero max-absolute error, zero relative-L2 error, zero cosine error, and no nonfinite values between main and this PR.
• Separate full-tensor validation covers real L32 versus neutral-padded L64 at N=256; a deterministic 1,048,576-element sample covers N=4096.
• PR-versus-upstream validation also covers L1, L31–33, L63, six GVA head groupings, initial/final state, and heterogeneous L1–193 batches. Outputs, final states, and all gradients show zero max-absolute and relative-L2 error.
• Verified on TileLang 0.1.9, the version currently used by FlashQLA.

All metrics are reported as observations; no precision or performance threshold gates are used.

Implementation notes

• The TMA/cp.async split is intentional: complete tiles are safe and faster with TMA; only an incomplete packed-varlen tile needs per-row predication and zero fill. Comments state this decision next to both paths.
• In prepare_h, the K-producer warp loads K/V/A with TMA for complete tiles, while the V/A-producer warp owns only incomplete V/A tails; the branches and comments make that ownership explicit.
• Producer boundaries and logical lanes are derived from WARP_SIZE and CONSUMER_WARP_COUNT; the warp layout is documented at the start of each JIT factory.
• has_incomplete_tile is a JIT specialization argument derived from the existing chunk metadata. Pure full-tile batches do not execute or generate the larger async-tail body, and this adds no device synchronization.