feat(kda): integrate KDA attention with tokamax backend and CP support - #1

feat(kda): integrate KDA attention with tokamax backend and CP support#1

chiaotung97 wants to merge 4 commits into

mainantgroup/maxtext:mainfrom

feature_kda_integrationantgroup/maxtext:feature_kda_integrationCopy head branch name to clipboard