ThunderKittens kernels for hopper

[kesavanramakrishnan]

Wrote TK kernels to replace the slower pallas path targeting the hopper arch.

Kernels are pretty straightforward, they are warp spec, 2 consumer 1 producer, using wgmma and tma and deeper pipelining than the pallas kernels.

Stage 1 kernel:
Identical to pallas kernel, except we only do 2 passes over the k-dim per query vs the 3.

we are able to decompose B into the sum of (dPa - dP*dD - ddS * D + dP * ddS) which we can expand and see these terms are independent of each other allowing us to do the reductions in the first sweep and then do the output in the second sweep (since these are just row-wise scalars it doesn't introduce much memory overhead we can keep in registers)

Stage 2 kernel:
algorithmically identical to pallas kernel

shapes supported are:
causal only
head dim: 64 (can be expanded to 128 just reduce pipeline depth)
seq len: seq % 128 == 0

performance:

end-to-end HVP times (fwd + bwd + double-backward)

batch	seq	Pallas (ms)	TK (ms)	speedup	correctness
1	512	0.18	0.22	0.82x	1.0
1	1024	0.33	0.28	1.17x	1.0
1	2048	0.88	0.66	1.34x	1.0
1	4096	2.72	1.76	1.55x	1.0
1	8192	9.46	5.56	1.70x	1.0
1	16384	34.94	19.21	1.82x	1.0
1	32768	137.54	73.63	1.87x	1.0
1	65536	555.30	289.03	1.92x	1.0
1	131072	2282.74	1143.72	2.00x	1.0

usage: see Readme

work done by me and @nick-rui

[marcelroed]

Cool! Will review this soon. Is it possible to have this build at runtime using optional CUDA packages from PyPI instead of requiring it to be prebuilt? Do you also know what the gains are from optimizing out the additional pass in Pallas?

[marcelroed]

Also if you can provide an NCU profile (with the Pallas as a baseline) and comment on where the savings are that would be great

[kesavanramakrishnan]

Hey! Just updated the commit to allow for runtime build w pypi and will address rest in dms