kaixih 307ea87a8d support head size of 256 ...

Test large head size only on hopper+ gpus

Test large head size only on cudnn 9.5+

2024-12-19 18:38:06 +00:00

..

__init__.py

Cleanup: fix unused imports & mark exported names

2024-10-16 17:42:41 -07:00

fused_attention_stablehlo.py

support head size of 256

2024-12-19 18:38:06 +00:00

fusion.py

jax.core: deprecate a number of APIs

2024-12-10 11:11:32 -08:00