[FEATURE REQUEST] Create sdpa_with_kv support for float16, bfloat16

pytorch / torchchat

Run PyTorch LLMs locally on servers, desktop and mobile

BSD 3-Clause "New" or "Revised" License

3.2k stars 197 forks source link

[FEATURE REQUEST] Create sdpa_with_kv support for float16, bfloat16 #550

Open mikekgfb opened 5 months ago

mikekgfb commented 5 months ago

see run here => https://github.com/pytorch/torchchat/actions/runs/8872459136/job/24356835073

We can always upcast to make things pass, but if there's an easy way to build float16 and bfloat16 flavors (iPad Pro has M-series chip with bfloat support AFAIRK?), that'd be rad.

mikekgfb commented 5 months ago

workaround in place. #553