Fix torch device

[deepfates]

Bug: device not set correctly, some tensors on cpu while the rest are successfully on mps.

RuntimeError: Placeholder storage has not been allocated on MPS device!

The device logic was duplicated across several files, so it was easy to miss the one line where this happened. I've moved that logic into a function in torch_device.py.

[HaileyStorm]

You left your prompts.py change in, might want to remove that for the PR (a suggestion, to be clear, I'm not important :P).

[deepfates]

That's what I get for trying to go back and isolate changes in commits instead of just doing it all at once lol. It's a good prompt for this kind of thing but i'm taking it back for clarity of PR

[xjdr-alt]

@Arrabonae

[Laz4rz]

Just wasted 2h figuring out why KVCache is misplaced on mps, and then saw this PR already exists. One thing here: the buffers are already on the correct device, hence we don't have to move the KVCache to device while KVCache.new.

[ParthSareen]

Just wasted 2h figuring out why KVCache is misplaced on mps, and then saw this PR already exists. One thing here: the buffers are already on the correct device, hence we don't have to move the KVCache to device while KVCache.new.

Dude I literally made the fix and was about to PR 💀