perf(gdr): improve SM100 KKT solve occupancy

[Butterfingrz]

Improve SM100 KKT Solve Occupancy

Summary

• Add T.annotate_min_blocks_per_sm(8) to the fixed-length and varlen KKT solve kernels on the shared SM100/SM103 path.
• Target eight resident 128-thread CTAs per SM. Keep the KKT math, layouts, launch grid, and public API unchanged.
• Improve all 15 measured B200 workloads. Reduce isolated KKT geometric-mean latency by 5.26%.

Correctness validation

Tests ran through the final production path on one B200 with Torch 2.8.0+cu128 and TileLang 0.1.13.

• python -m pytest -q tests/test_gdr_unit.py tests/test_function_signature.py -m 'not slow' — 115 passed, 6 deselected.
• python -m pytest -q tests/test_gdr_unit.py -m slow — 6 passed, 114 deselected.
• Baseline and candidate were bitwise identical in six direct KKT probes and all 30 auto_cp=True/False matrix rows (o, A, and ht).

Performance Results (B200 / SM100)

Benchmark B=1, total T=32768, D=128, and BF16 Q/K/V across five head configurations and three sequence layouts. Use three alternating A/B pairs, 10 warmups, 100 repeats, CUDA Graphs, and paired medians. Use one 148-SM B200, driver 580.126.09, Torch 2.8.0+cu128, and TileLang 0.1.13. Clocks are not fixed.

Baseline: c18a486 (unchanged in 7c7dfe1). Candidate: 44814ca, source-identical to the measured candidate.

The table reports the isolated KKT lane from the auto_cp=False campaign. auto_cp is not an input to the isolated kernel call. The auto_cp=True campaign reproduces the result with a 0.947493 geometric-mean ratio. 32k_4seq_rand uses seed 42 and lengths [7212, 7072, 13725, 4759].

Model config	Sequence layout	Hqk	Hv	#seq	c18a486	candidate	Ratio	Speedup	Time change
hk2_hv8	32k_1seq	2	8	1	56.026 µs	52.915 µs	0.944	1.059x	-5.55%
hk2_hv8	32k_4seq	2	8	4	55.980 µs	52.803 µs	0.943	1.060x	-5.68%
hk2_hv8	32k_4seq_rand	2	8	4	56.074 µs	52.929 µs	0.944	1.059x	-5.61%
hk8_hv32	32k_1seq	8	32	1	215.976 µs	207.372 µs	0.960	1.041x	-3.98%
hk8_hv32	32k_4seq	8	32	4	216.124 µs	207.546 µs	0.960	1.041x	-3.97%
hk8_hv32	32k_4seq_rand	8	32	4	217.595 µs	208.216 µs	0.957	1.045x	-4.31%
hk16_hv32	32k_1seq	16	32	1	217.624 µs	207.988 µs	0.956	1.046x	-4.43%
hk16_hv32	32k_4seq	16	32	4	217.579 µs	207.646 µs	0.954	1.048x	-4.57%
hk16_hv32	32k_4seq_rand	16	32	4	219.278 µs	208.825 µs	0.952	1.050x	-4.77%
hk16_hv64	32k_1seq	16	64	1	428.176 µs	403.333 µs	0.942	1.062x	-5.80%
hk16_hv64	32k_4seq	16	64	4	428.180 µs	403.304 µs	0.942	1.062x	-5.81%
hk16_hv64	32k_4seq_rand	16	64	4	430.417 µs	404.827 µs	0.941	1.063x	-5.95%
hk32_hv32	32k_1seq	32	32	1	217.629 µs	204.551 µs	0.940	1.064x	-6.01%
hk32_hv32	32k_4seq	32	32	4	217.590 µs	204.682 µs	0.941	1.063x	-5.93%
hk32_hv32	32k_4seq_rand	32	32	4	219.726 µs	205.420 µs	0.935	1.070x	-6.51%

Performance Summary

• 15-case geometric-mean ratio: 0.947387. Speedup: 1.0555x. Time reduction: 5.26%.
• Summed time reduction: 5.32%. Per-case ratio range: 0.934891–0.960313. Improved cases: 15 / 15.
• End-to-end latency falls by 1.43% with auto_cp=False and 1.52% with auto_cp=True. Both arms use the same fused-forward source.
• n=6 loses 15 / 15 cases to n=8 with a 1.0177 geometric-mean ratio.

NCU Occupancy Increase

NCU 2026.2.1 profiles hk16_hv32/32k_1seq on one B200. Both runs use a 16,384-CTA grid and 128-thread CTAs.

Metric	Baseline	Candidate	Change
gpu__time_duration.sum	227.872 µs	219.008 µs	-3.89%
launch__grid_size	16,384	16,384	unchanged
launch__block_size	128	128	unchanged
launch__registers_per_thread	110	64	-41.82%
launch__occupancy_limit_registers	4 blocks	8 blocks	2.00x
launch__waves_per_multiprocessor	27.676	13.838	-50.00%
sm__warps_active.avg.pct_of_peak_sustained_active	23.66%	47.74%	+24.08 pp
sm__throughput.avg.pct_of_peak_sustained_elapsed	39.35%	41.11%	+1.77 pp
gpu__compute_memory_throughput.avg.pct_of_peak_sustained_elapsed	89.65%	92.30%	+2.65 pp

The baseline is register-limited to four CTAs per SM. The candidate is register-limited to eight CTAs per SM. The grid and block size stay unchanged. Waves per SM halve. The achieved active-warp percentage doubles. The NCU duration improves by 3.89%. The paired CUDA Graph result for the same row improves by 4.43%.

Scope

• Change only two annotations in flash_qla/ops/gated_delta_rule/chunk/blackwell/kkt_solve.py. Preserve fixed-length and varlen behavior, arithmetic, layouts, masks, launch geometry, and output types.
• Measure performance on B200/SM100 only. SM103 shares the source path but is not measured. NCU covers one specialization.
• The annotation targets, but does not guarantee, eight resident CTAs per SM. The available NCU reports do not support spill-count claims

[Butterfingrz]

Hi! Could you take a look at this PR when you get a chance. Thanks!