Enhance device compatibility by adding MPS support and ensuring tensor operations respect device context

Handle LlamaFlashAttention2 import for compatibility with newer transformers versions

Initial copy