Add support for head dimension 128

[perklet]

Most recent models use hdim=128, it would be great to see that ThunderKittens also support that.

ThunderKittens/examples/attn/h100/h100_train_atn.py

Lines 25 to 26 in a562ed2

	# assert training does not support head dim 128 yet
	assert q.shape[2] == 64, "TK train currently supports head dim 64 only"

[benjaminfspector]

That is on the eventual to-do list!

[perklet]

Great! Thanks for the brilliant project! Are there any pointers on what the block issues are? I would like to try to fix this on my side. 来自手机回复 Benjamin Spector ***@***.***>于2024年5月16日 周四10:33写道：

…

That is on the eventual to-do list! — Reply to this email directly, view it on GitHub <#26 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAH4ZXZORXRM4YGACGHZGCLZCQLGNAVCNFSM6AAAAABHZHMMZ6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCMJTHEYDEMJXGQ> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[benjaminfspector]

looping in @Aaryan0404 for this

[dz1iang]

Could you please explain why there is a limitation on the size of the head dimension? I'm not very clear about it.

That is on the eventual to-do list!

Could you please explain why there is a limitation on the size of the head dimension? I'm not very clear about it.