Resolve the execution device instead of hardcoding CUDA

`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.

modeling_deepseekocr2.py

@@ -492,7 +492,7 @@ class DeepseekOCR2Model(DeepseekV2Model):
                     images_in_this_batch = torch.cat(images_in_this_batch, dim=0)
                     # exit()
 
-                    inputs_embeds[idx].masked_scatter_(images_seq_mask[idx].unsqueeze(-1).cuda(), images_in_this_batch)
+                    inputs_embeds[idx].masked_scatter_(images_seq_mask[idx].unsqueeze(-1).to(inputs_embeds.device), images_in_this_batch)
 
                 idx += 1
             
@@ -693,6 +693,9 @@ class DeepseekOCR2ForCausalLM(DeepseekV2ForCausalLM):
     def infer(self, tokenizer, prompt='', image_file='', output_path = '', base_size=1024, image_size=640, crop_mode=True, test_compress=False, save_results=False, eval_mode=False):
         self.disable_torch_init()
 
+        device = next(self.parameters()).device
+        device_type = device.type
+
         os.makedirs(output_path, exist_ok=True)
         os.makedirs(f'{output_path}/images', exist_ok=True)
 
@@ -903,12 +906,12 @@ class DeepseekOCR2ForCausalLM(DeepseekV2ForCausalLM):
 
         if not eval_mode:
             streamer = NoEOSTextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
-            with torch.autocast("cuda", dtype=torch.bfloat16):
+            with torch.autocast(device_type, dtype=torch.bfloat16):
                 with torch.no_grad():
                     output_ids = self.generate(
-                        input_ids.unsqueeze(0).cuda(),
-                        images=[(images_crop.cuda(), images_ori.cuda())],
-                        images_seq_mask = images_seq_mask.unsqueeze(0).cuda(),
+                        input_ids.unsqueeze(0).to(device),
+                        images=[(images_crop.to(device), images_ori.to(device))],
+                        images_seq_mask = images_seq_mask.unsqueeze(0).to(device),
                         images_spatial_crop = images_spatial_crop,
                         # do_sample=False,
                         # num_beams = 1,
@@ -921,12 +924,12 @@ class DeepseekOCR2ForCausalLM(DeepseekV2ForCausalLM):
                         )
 
         else:
-            with torch.autocast("cuda", dtype=torch.bfloat16):
+            with torch.autocast(device_type, dtype=torch.bfloat16):
                 with torch.no_grad():
                     output_ids = self.generate(
-                        input_ids.unsqueeze(0).cuda(),
-                        images=[(images_crop.cuda(), images_ori.cuda())],
-                        images_seq_mask = images_seq_mask.unsqueeze(0).cuda(),
+                        input_ids.unsqueeze(0).to(device),
+                        images=[(images_crop.to(device), images_ori.to(device))],
+                        images_seq_mask = images_seq_mask.unsqueeze(0).to(device),
                         images_spatial_crop = images_spatial_crop,
                         # do_sample=False,
                         # num_beams = 1,
@@ -939,7 +942,7 @@ class DeepseekOCR2ForCausalLM(DeepseekV2ForCausalLM):
                 
 
         if '<image>' in conversation[0]['content'] and eval_mode:
-                outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
+                outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).to(device).shape[1]:])
                 stop_str = '<｜end▁of▁sentence｜>'
                 if outputs.endswith(stop_str):
                     outputs = outputs[:-len(stop_str)]
@@ -949,7 +952,7 @@ class DeepseekOCR2ForCausalLM(DeepseekV2ForCausalLM):
                 return outputs
         
         if '<image>' in conversation[0]['content'] and test_compress:
-            outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
+            outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).to(device).shape[1]:])
             pure_texts_outputs_token_length = len(text_encode(tokenizer, outputs, bos=False, eos=False))
             print('='*50)
             print('image size: ', (w, h))
@@ -960,7 +963,7 @@ class DeepseekOCR2ForCausalLM(DeepseekV2ForCausalLM):
 
 
         if '<image>' in conversation[0]['content'] and save_results:
-            outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
+            outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).to(device).shape[1]:])
             stop_str = '<｜end▁of▁sentence｜>'
 
             print('='*15 + 'save results:' + '='*15)

modeling_deepseekv2.py

@@ -36,7 +36,8 @@ from transformers.cache_utils import Cache, DynamicCache
 from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
 from transformers.models.llama.modeling_llama import (
     LlamaAttention,
-    LlamaFlashAttention2
+    LlamaFlashAttention2,
+    LlamaSdpaAttention,
 )
 from transformers.modeling_outputs import (
     BaseModelOutputWithPast,
@@ -1230,12 +1231,16 @@ class DeepseekV2FlashAttention2(DeepseekV2Attention):
 ATTENTION_CLASSES = {
     "eager": DeepseekV2Attention,
     "flash_attention_2": DeepseekV2FlashAttention2,
+    "sdpa": DeepseekV2Attention,
 
     "mla_eager": DeepseekV2Attention,
     "mla_flash_attention_2": DeepseekV2FlashAttention2,
+    # MLA has no dedicated SDPA path, fall back to the math implementation
+    "mla_sdpa": DeepseekV2Attention,
 
     "mha_eager": LlamaAttention,
-    "mha_flash_attention_2": LlamaFlashAttention2
+    "mha_flash_attention_2": LlamaFlashAttention2,
+    "mha_sdpa": LlamaSdpaAttention,
 }
 
 
@@ -1361,6 +1366,7 @@ class DeepseekV2PreTrainedModel(PreTrainedModel):
     _no_split_modules = ["DeepseekV2DecoderLayer"]
     _skip_keys_device_placement = "past_key_values"
     _supports_flash_attn_2 = True
+    _supports_sdpa = True
     _supports_cache_class = True
 
     def _init_weights(self, module):