Resolve the execution device instead of hardcoding CUDA

#34
by Jiqing - opened

infer() hardcoded .cuda() and torch.autocast("cuda", ...), so the model could only run on NVIDIA GPUs even though nothing in it is CUDA-specific. The device is now read back from the model own parameters, which keeps the existing CUDA behaviour identical while letting the same code run on CPU and on other accelerator backends.

ATTENTION_CLASSES only offered eager and flash_attention_2. FlashAttention is CUDA-only, which left eager as the single portable option, so the sdpa variants are registered as well and _supports_sdpa is set. MLA has no dedicated SDPA kernel here, so mla_sdpa maps to the math implementation.

Verified on an Intel Arc Pro B60 (attn_implementation="sdpa") and on CPU (attn_implementation="eager"): both produce identical output.

Jiqing changed pull request status to closed
Jiqing changed pull request status to open
Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment