perf(gdr): move SM100 prepare_h CP transition-matrix recurrence to tcgen05

[Butterfingrz]

Summary

• Move the CP transition-matrix recurrence (M, Z) off Consumer-X/Y T.gemm onto the MMA warp's tcgen05 pipe, with M resident in TMEM.
• Rebalance the S/X/Y/producer register caps from 168/160/160/24 to 152/104/104/152, funded by the M fragments X and Y no longer carry.
• With store_h=False, publish the next bar_0 arrival and the left state half one iteration early, inside the tcbar_2a → tcbar_2b window.
• Store mt straight from the final fragments, dropping its shared-memory/TMA bounce.
• Add a calc_mt reachability guard to the existing long-fixed auto-CP test.

PR implementation

This PR touches flash_qla/ops/gated_delta_rule/chunk/blackwell/prepare_h.py (+144/-123 on 7c7dfe1) and tests/test_gdr_unit.py (+26/-2). Public API, tensor shapes, CP selection, state layouts, and dispatch are unchanged, and every kernel hunk is inside with T.Kernel. The test change guards calc_mt reachability inside the existing test_fwd_auto_cp long-fixed case and adds no test ids.

Correctness validation

• python -m pytest tests/ -q: 123 passed, 7 skipped. The count predates the test guard, which has not been rerun on GPU.
• 48-specialization gate vs 7c7dfe1 (store_h x h0 x state_v_first x varlen x CP): 48 / 48. h bitwise on the 24 rows that write it, ht on all 48, mt on 4390 / 4480 slices — the 90 excluded are those the baseline itself does not reproduce run to run.
• Forward integration vs 7c7dfe1: o and ht bitwise on 15 / 15 model-derived rows and 42 / 42 matrix rows.
• PrimFunc gate: sha256(prim_func.script()), TileLang's kernel-cache key, identical on 64 / 64 specializations between the benchmarked and the submitted source, so the table below applies unchanged.

Performance Results (B200 / SM100)

Across the 24 CP-active rows of a 42-case matrix on one B200: BF16, B=1, T=32768, D=128, varlen seed 42. Three interleaved CUDA-graph A/B rounds (warmup=10, rep=100), swapping only prepare_h. The table reports standalone prepare_h latency; the 18 non-CP rows are the control in the summary below. A same-source A/A control over 12 CP-active rows returns prepare_h 1.0003; clocks were not fixed.

Model config	Sequence layout	Hqk	Hv	#seq	7c7dfe1	PR	Ratio	Speedup	Time change
hk16_hv16	32k_1seq	16	16	1	153.839 us	107.526 us	0.699	1.431x	-30.10%
hk16_hv16	32k_2seq	16	16	2	152.334 us	108.499 us	0.712	1.404x	-28.78%
hk4_hv16	32k_1seq	4	16	1	149.883 us	107.499 us	0.717	1.394x	-28.28%
hk4_hv16	32k_2seq	4	16	2	151.126 us	107.225 us	0.710	1.409x	-29.05%
hk12_hv12	32k_1seq	12	12	1	118.639 us	90.480 us	0.763	1.311x	-23.74%
hk12_hv12	32k_2seq	12	12	2	118.088 us	90.808 us	0.769	1.300x	-23.10%
hk12_hv12	32k_4seq	12	12	4	116.943 us	89.413 us	0.765	1.308x	-23.54%
hk12_hv12	32k_5seq	12	12	5	115.829 us	87.885 us	0.759	1.318x	-24.13%
hk4_hv12	32k_1seq	4	12	1	118.712 us	90.112 us	0.759	1.317x	-24.09%
hk4_hv12	32k_2seq	4	12	2	118.310 us	90.168 us	0.762	1.312x	-23.79%
hk4_hv12	32k_4seq	4	12	4	118.443 us	89.486 us	0.756	1.324x	-24.45%
hk4_hv12	32k_5seq	4	12	5	116.226 us	87.250 us	0.751	1.332x	-24.93%
hk8_hv8	32k_1seq	8	8	1	90.726 us	61.362 us	0.676	1.479x	-32.36%
hk8_hv8	32k_2seq	8	8	2	89.422 us	61.619 us	0.689	1.451x	-31.09%
hk8_hv8	32k_4seq	8	8	4	90.735 us	61.498 us	0.678	1.475x	-32.22%
hk8_hv8	32k_5seq	8	8	5	89.756 us	62.524 us	0.697	1.436x	-30.34%
hk4_hv8	32k_1seq	4	8	1	90.593 us	62.299 us	0.688	1.454x	-31.23%
hk4_hv8	32k_2seq	4	8	2	88.578 us	62.162 us	0.702	1.425x	-29.82%
hk4_hv8	32k_4seq	4	8	4	90.014 us	62.169 us	0.691	1.448x	-30.93%
hk4_hv8	32k_5seq	4	8	5	89.330 us	61.425 us	0.688	1.454x	-31.24%
hk4_hv4	32k_1seq	4	4	1	55.628 us	43.684 us	0.785	1.273x	-21.47%
hk4_hv4	32k_2seq	4	4	2	55.519 us	43.742 us	0.788	1.269x	-21.21%
hk4_hv4	32k_4seq	4	4	4	55.788 us	43.893 us	0.787	1.271x	-21.32%
hk4_hv4	32k_5seq	4	4	5	55.579 us	43.740 us	0.787	1.271x	-21.30%

Performance Summary

• prepare_h geomean ratio 0.731245 (1.368x), improved in 24 / 24 rows. CP-active forward E2E 0.915596; the 18 non-CP control rows sit at 0.999925.
• A five-row model-derived rerun on a verified-idle device agrees: prepare_h 0.716833 (1.395x), CP-active E2E 0.914899.
• Backward path (store_h=True, auto-CP off): prepare_h 0.843574 (1.185x), E2E 0.967635 across 15 / 15 rows.

Change attribution

• The tcgen05 move is the bulk of the gain. M stays in TMEM; only the BF16 B-operand restages cross back through registers, because TileLang 0.1.13 has no TMEM-to-SMEM path while tcgen05 requires B in SMEM. Consumer-X restages Z alone, in a window where it would otherwise idle; X and Y each restage their own M half.
• The register rebalance removes producer spill from the seven GEMM descriptor sites now live on the MMA warp. Both cap vectors sum to 512; the set_max_nreg directions are 1/0/0/1 and a wrong direction traps at runtime, not at compile time.
• The store_h=False early publish is legal for iteration i+1 because h_shared's only GEMM reader, U = K @ S, precedes the h_L update on the same in-order pipe. The store_h=True rendezvous is untouched, and the L-half TMEM load stays adjacent to the tcbar_2a wait so the tcgen05 fence pass keeps its after_thread_sync.
• The direct mt store removes the SMEM/TMA hop at the cost of a final tcbar_4a/tcbar_4b drain.

NCU Spot Check

NCU was collected with Full + PmSampling and clock-control=base on B200 using BF16 B=1, T=32768, Hg=16, H=32, D=128, and a single sequence, replaying the spy-captured CP-preprocess fused_gdr_h call under --profile-from-start off -k regex:tilelang_prepare_h --launch-count 1.

Metric	7c7dfe1	PR	Change
sm__cycles_active.avg	276,883	197,625	-28.6%
gpu__time_duration.sum	370,816 ns	276,480 ns	-25.4%
smsp__inst_executed.sum	36,016,311	25,743,891	-28.5%
Classic HMMA ops, bf16 to fp32	16,106,127,360	0	-100%
tcgen05 UTC MMA ops, bf16 to fp32	27,404,795,904	43,769,331,712	+59.7%
Local-memory instructions (ld+st)	432,254	0	-100%
dram__bytes_write.sum	31,774,464	24,206,080	-23.8%
Global-store sectors, LSU + TMA	1,572,864	2,097,152	+33.3%

• Total tensor ops move 43.51G to 43.77G, +0.59%: the tcgen05 path absorbs the HMMA ops exactly. Same arithmetic, different pipe.
• Register spill reaches zero. The three local-memory rules NCU reports on the baseline are absent from the PR report.
• Global-store sectors rise because mt moves off the TMA path onto the LSU path. Per-sector store efficiency is unchanged at 16.0 bytes in both configurations.
• DRAM writes still fall: the spill traffic removed exceeds the store traffic added.
• Shared-memory bank conflicts rise, load 5.08M to 6.53M and store 0.30M to 1.10M. This can be further optimized.

Scope

Backward figures were taken at c18a486, before 7c7dfe1 merged PR #40's fused_bwd.py register-cap change, so they are pre-rebase; prepare_h.py is identical across those commits, so the forward figures and every gate are unaffected. Only B200 / SM100 is measured, though SM103 shares the Blackwell dispatch. Runs used TileLang 0.1.13 while setup.py pins 0.1.9.