runtime: bf16 gemm runs a measured cuBLASLt pick per shape, not the heuristic's first answer

[xiaguan]

Summary

extern:cublaslt_bf16_tn runs cuBLASLt's first heuristic answer. That answer is not the fastest one, and which algorithm is fastest depends on the GPU generation and the library version. #25 fixes the choice to one tile (256x128, no split-K), measured on sm_89. On GB300 that same policy makes a step 4.3% slower. A measured per-shape pick is 2.6% faster than the heuristic.

GEMM is a basic capability of the runtime, so a measured choice is worth having. It should be a tuned table outside the manifest, not a tile name inside it.

Measurement (GB300, sm_103, cuBLASLt 13.7.0, 1400 W cap)

Shapes: HiDream-O1 / Qwen3-VL-8B text tower at m = 4097: qkv 6144x4096, o 4096x4096, gate_up 24576x4096, down 4096x12288. Random bf16 inputs, same column-major mapping as gemm_bf16_tn. #25's selection logic is copied exactly.

• The heuristic's candidates. It returns 8 at every shape, and all 8 are non-split-K cluster kernels (algo id 66). Top-1 is already a big tile: 256x240 with a 2x2 cluster for qkv/o/down. runtime: a cublaslt bf16 gemm launch may pin its algorithm #25's 256x128 is offered at qkv/o/gate_up and borrowed from k = n for down. On this arch it is the narrower tile.
• Sustained load, 4 s per candidate. Every candidate sits at the power cap (~1365 W, throttle reason = SW power cap). Power is flat, so time tracks J/GEMM, which confirms runtime: a cublaslt bf16 gemm launch may pin its algorithm #25's mechanism. But 256x128 is 4–6% worse than top-1 on qkv/o, in both time and energy.
• Cold timing picks the same winners as sustained load at all four shapes (qkv 128x192, o 256x224, gate_up 256x208, down 256x224). The sm_89 observation that isolated timing misranks does not reproduce here.
• Whole-step run. Each step is 36 layers x 4 GEMMs. The three policies alternate, 4 s x 6 rounds, and the round-to-round spread is < 0.5%.

policy	ms/step	vs top-1
heuristic top-1	35.12	—
#25 wide (256x128)	36.62	+4.3%
best per shape (from top-8, sustained)	34.22	−2.6%

Constraints

• The algorithm cannot live in the manifest. A cuBLASLt algorithm moves with the library version (the reason extern:cublaslt_bf16_tn_tile was removed in cd39a57), and the best one moves with the hardware (above). Static shapes in the manifest do not change this. n and k are already literals, and m is fixed per captured (program, vars). What is missing is a place to record the measured choice, not the shape.
• The choice must be deterministic. Candidates are often within noise of each other (qkv: 0.129 / 0.130 / 0.130 ms), and different tiles are not guaranteed to produce the same bits. Timing at every startup would break "same input, same bytes out". The choice should be tuned once and pinned.
• Timing cannot run inside graph capture, and the first call to a graph program happens during capture. Tuning belongs offline (or in an explicit warmup), not in the first call.

Proposal

Follow the registry's [[pick]] model: a measured choice keyed by shape, with a report behind it.

• kern bench (or kern tune) writes the table. It takes the workload's m values (the vars the service actually captures) x every GEMM call site's (n, k, ldc, A row stride) in the manifest, times cuBLASLt's top-N at each shape, and writes the winners.
• The table's key and value. Key: GPU name + cuBLASLt version + shape. Value: the opaque cublasLtMatmulAlgo_t, which is stable within one library version. The table ships with the deployment, not the manifest.
• The runtime only looks up. A hit runs the tabled algorithm; a miss runs heuristic top-1. Both are deterministic, and the manifest keeps saying extern:cublaslt_bf16_tn.
• runtime: a cublaslt bf16 gemm launch may pin its algorithm #25's _wide entry is then not needed. sm_89 and GB300 each get their own measured winner.

Open question: small m

Everything above is m = 4097, where the heuristic is already within 2.6% of the best. Decode-shaped GEMMs (m = 1..256) are where cuBLASLt heuristics are usually weakest, and they have not been measured yet. A sweep there decides how much this is worth before it goes into the roadmap.

Search space

cublasLtMatmulAlgoGetHeuristic returns only 8 candidates on GB300 even when asked for 64. cublasLtMatmulAlgoGetIds lists 20 algo ids for bf16 TN, and id 66 alone reports 645 tile ids via CUBLASLT_ALGO_CAP_TILE_IDS. Configurations beyond top-8 can be built with cublasLtMatmulAlgoConfigSetAttribute and validated with cublasLtMatmulAlgoCheck, if top-8 turns out not to be enough.

[xiaguan]

Measured: K3 16k prefill GEMMs (GB300, one TP4 rank)

To see what tuning is worth on a real workload, I took every cuBLAS GEMM in K3's 16k prefill manifest (93 layers, gen_k3.py --tp 4 --chunk 16384, tokens = ctx = 16384). That is 14 distinct shapes, 455 TFLOP per rank per prefill.

Setup. GB300, cuBLAS/cuBLASLt 13.7.0, one idle GPU.

• Baseline is what kern runs today: cublas_bf16_tn_f32 goes to cublasGemmEx with DEFAULT_TENSOR_OP, and cublaslt_bf16_tn runs the Lt heuristic's top-1. Both use a 32 MiB workspace.
• top-8: every heuristic answer, which is 8 on this GPU.
• ext: each heuristic algorithm used as a template, set to each of its CAP_TILE_IDS x 8 cluster shapes, keeping the ones AlgoCheck accepts. That is about 5000 configurations per shape. The output dtype always matches the entry, so f32 entries are tuned as Lt bf16 x bf16 -> f32.
• Timing. Operands rotate over sets larger than 3x L2 so weights never stay resident. The GPU is heated first, and each candidate's time is the median of 5 x 40 ms runs.

out	m	n	k	calls	base µs	top-8 µs	ext µs	best vs base	ms saved / prefill	winner (algo 66, no split-K)
f32	16384	12288	7168	69	1400.5	1404.6	1360.0	-2.9%	2.79	128x288, cluster 4x1
bf16	16384	7168	3072	93	343.7	342.7	332.4	-3.3%	1.05	224x128, cluster 4x2
f32	4096	12288	7168	92	358.9	354.8	355.7	-1.1%	0.38	176x128, cluster 2x2
f32	4096	7168	6144	92	177.4	177.6	173.2	-2.4%	0.39	128x256, cluster 2x2
f32	16384	5184	7168	24	625.0	614.6	600.8	-3.9%	0.58	192x192, cluster 2x2
f32	4096	3584	7168	92	104.6	105.1	104.3	-0.3%	0.03	128x256, cluster 4x2
f32	4096	7168	3584	92	106.2	107.6	102.7	-3.3%	0.32	128x256, cluster 2x2
bf16	16384	4608	1536	24	118.7	119.0	118.1	-0.5%	0.01	128x256, cluster 4x2
f32	16384	256	7168	69	47.1	47.2	45.8	-2.8%	0.09	128x224, cluster 2x1
f32	4096	67584	7168	1	1940.6	1942.7	1872.2	-3.5%	0.07	128x256, cluster 2x1
bf16	16384	7680	576	24	81.5	81.7	78.5	-3.7%	0.07	128x256, cluster 4x1
f32	4096	7168	33792	1	893.6	924.8	908.3	—	0	cublasGemmEx (every Lt candidate slower)
f32	4096	224	7168	92	17.9	15.8	14.9	-16.8%	0.28	64x112, cluster 2x1
bf16	16384	3072	128	69	20.7	20.7	18.9	-8.7%	0.12	256x192, cluster 2x2

Per rank per 16k prefill:

	GEMM time	vs baseline
baseline	226.3 ms	—
top-8 tuned	225.4 ms	-0.4%
ext tuned	220.1 ms	-2.7%

The ext-tuned table saves 6.2 ms, 0.83% of the 749.5 ms TTFT.

What this says:

• Tuning within the heuristic's top-8 is worth almost nothing on GB300. TunableOp-style "time the top-N" gets -0.4%. The gain comes from the wider tile x cluster space, so the tuner has to build configurations rather than only rank what the heuristic returns.
• The ceiling on this workload is ~1% of TTFT. It is a small, real win, not a large one. Narrow shapes gain the most relatively (-17%, -9%) but are cheap in absolute terms.
• cuBLAS can beat every Lt candidate. At 4096x7168x33792, cublasGemmEx wins. The table has to be able to say "keep the default".

Caveats:

• Isolated GEMMs on one GPU. Nothing else ran between them, and the other three ranks were idle. A replay of one prefill's GEMM sequence, baseline versus tuned, alternated, was planned but not run. The number to trust in the end is the 16k TTFT with the table applied in kern.
• lm_head (m = 1, 163840x7168) was left out. Its extended search would not finish in reasonable time, and it is 5e-6 of the FLOPs.
• The numerics check here is vacuous. The inputs were multiples of 1/8, so every summation order is exact and rel = 0 for any algorithm (the same weakness as the test in runtime: a cublaslt bf16 gemm launch may pin its algorithm #25). Whether a tile swap changes bits on real data is still open. The winners are all the same kernel family without split-K, but that is not proof.