Fast API on CPU + MPS GPU

[fakerybakery]

Hi @manmay-nakhashi, are you planning to support the fast API work on CPU + MPS GPU? Thanks!

[fakerybakery]

Just realized this can't work on MPS since the fast API requires autocast and autocast isn't yet supported on MPS. I guess we'll have to wait for that to get merged before the fast API will work

[SSardorf]

pytorch/pytorch#99272 (comment) Seems autocast is finally supported on MPS!

[SSardorf]

I made some changes to the codebase, which lets it utilize mps when running autocast, and based on some initial testing, it's significantly faster. Especially on the clvp step!

@fakerybakery @manmay-nakhashi

[manmay-nakhashi]

@SSardorf can you make pull request for this ?

[SSardorf]

@SSardorf can you make pull request for this ?

Sure, I'll clean up the changes and open a PR 👍

[SSardorf]

Opened a draft PR. There may be other places where mps could help, I just haven't had a chance to look
#837