fix: enable XPU support and update hardcoded CUDA selections for tests (#7401)

* fix: add XPU device support and update hardcoded CUDA selections

* fix: add XPU device support for pytest CUDA skipped tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix device handling for PR #7401

- perplexity_eval.py: use DEVICE_TYPE_TORCH, not DEVICE_TYPE. The latter can
  be "hip" or "mlx", which .to() rejects, so this regressed ROCm.
- test_batched_leftpad_generation_gpu.py: XPU diverges here today, so mark it
  non-strict xfail on XPU instead of reverting to a CUDA-only guard. Keeps the
  real XPU gap visible and turns green once it is fixed.
- Guard torch.xpu.is_available() with hasattr, matching device_type.py.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Re-enable the flash varlen attention test in CI for PR #7401

attention_dispatch.py now predefines flash_attn_func / flash_attn_varlen_func
as None, so test_run_attention_flash_varlen_receives_window_and_softcap no
longer needs flash_attn importable to be monkeypatched. Verified on a runner
shaped like the CPU-only one: the test fails against main's attention_dispatch
and passes at this head, so the deselect is now dead weight.

* Tighten comments for PR #7401

Drop the hasattr rationale: torch.xpu has existed since torch 2.3 and the
dependency floor is 2.4, so no supported build predates the namespace. The
guard stays as cheap defence, but the comment claimed something untrue.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>

tests/saving/language_models/test_merge_model_perplexity_llama-3.2.py

@@ -96,12 +96,11 @@ def load_and_compute_8bit_ppl(
 if __name__ == "__main__":
     mp.set_start_method("spawn", force = True)
 
-    if torch.cuda.is_bf16_supported():
-        compute_dtype = torch.bfloat16
-        attn_implementation = "flash_attention_2"
-    else:
-        compute_dtype = torch.float16
-        attn_implementation = "sdpa"
+    from unsloth import is_bfloat16_supported
+    from unsloth.models._utils import HAS_FLASH_ATTENTION
+
+    compute_dtype = torch.bfloat16 if is_bfloat16_supported() else torch.float16
+    attn_implementation = "flash_attention_2" if HAS_FLASH_ATTENTION else "sdpa"
 
     model, tokenizer = FastLanguageModel.from_pretrained(
         model_name = "unsloth/Llama-3.2-3B-Instruct",

tests/saving/language_models/test_merge_model_perplexity_mistral.py

@@ -121,12 +121,11 @@ def load_and_compute_8bit_ppl(
 if __name__ == "__main__":
     mp.set_start_method("spawn", force = True)
 
-    if torch.cuda.is_bf16_supported():
-        compute_dtype = torch.bfloat16
-        attn_implementation = "flash_attention_2"
-    else:
-        compute_dtype = torch.float16
-        attn_implementation = "sdpa"
+    from unsloth import is_bfloat16_supported
+    from unsloth.models._utils import HAS_FLASH_ATTENTION
+
+    compute_dtype = torch.bfloat16 if is_bfloat16_supported() else torch.float16
+    attn_implementation = "flash_attention_2" if HAS_FLASH_ATTENTION else "sdpa"
 
     model, tokenizer = FastLanguageModel.from_pretrained(
         model_name = "unsloth/mistral-7b-v0.3",

tests/saving/language_models/test_merge_model_perplexity_phi_4.py

@@ -98,12 +98,11 @@ def load_and_compute_8bit_ppl(
 if __name__ == "__main__":
     mp.set_start_method("spawn", force = True)
 
-    if torch.cuda.is_bf16_supported():
-        compute_dtype = torch.bfloat16
-        attn_implementation = "flash_attention_2"
-    else:
-        compute_dtype = torch.float16
-        attn_implementation = "sdpa"
+    from unsloth import is_bfloat16_supported
+    from unsloth.models._utils import HAS_FLASH_ATTENTION
+
+    compute_dtype = torch.bfloat16 if is_bfloat16_supported() else torch.float16
+    attn_implementation = "flash_attention_2" if HAS_FLASH_ATTENTION else "sdpa"
 
     model, tokenizer = FastLanguageModel.from_pretrained(
         model_name = "unsloth/Phi-4",

tests/saving/language_models/test_merged_model_perplexity_llama-3.1-8b.py

@@ -95,12 +95,11 @@ def load_and_compute_8bit_ppl(
 if __name__ == "__main__":
     mp.set_start_method("spawn", force = True)
 
-    if torch.cuda.is_bf16_supported():
-        compute_dtype = torch.bfloat16
-        attn_implementation = "flash_attention_2"
-    else:
-        compute_dtype = torch.float16
-        attn_implementation = "sdpa"
+    from unsloth import is_bfloat16_supported
+    from unsloth.models._utils import HAS_FLASH_ATTENTION
+
+    compute_dtype = torch.bfloat16 if is_bfloat16_supported() else torch.float16
+    attn_implementation = "flash_attention_2" if HAS_FLASH_ATTENTION else "sdpa"
 
     model, tokenizer = FastLanguageModel.from_pretrained(
         model_name = "unsloth/Llama-3.1-8B-Instruct",

tests/saving/language_models/test_merged_model_perplexity_qwen_2.5.py

@@ -164,12 +164,11 @@ def load_and_compute_8bit_ppl(
 if __name__ == "__main__":
     mp.set_start_method("spawn", force = True)
 
-    if torch.cuda.is_bf16_supported():
-        compute_dtype = torch.bfloat16
-        attn_implementation = "flash_attention_2"
-    else:
-        compute_dtype = torch.float16
-        attn_implementation = "sdpa"
+    from unsloth import is_bfloat16_supported
+    from unsloth.models._utils import HAS_FLASH_ATTENTION
+
+    compute_dtype = torch.bfloat16 if is_bfloat16_supported() else torch.float16
+    attn_implementation = "flash_attention_2" if HAS_FLASH_ATTENTION else "sdpa"
 
     model, tokenizer = FastLanguageModel.from_pretrained(
         model_name = "unsloth/Qwen2.5-7B-Instruct",
@@ -210,8 +209,6 @@ if __name__ == "__main__":
         loftq_config = None,
     )
 
-    from unsloth import is_bfloat16_supported
-
     trainer = SFTTrainer(
         model = model,
         tokenizer = tokenizer,

tests/test_fp8_tiny_e8m0.py

@@ -11,7 +11,11 @@ dequant reference.
 import pytest
 import torch
 
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason = "needs CUDA")
+cuda_available = torch.cuda.is_available()
+xpu_available = hasattr(torch, "xpu") and torch.xpu.is_available()
+dev = "cuda" if cuda_available else "xpu" if xpu_available else "cpu"
+
+pytestmark = pytest.mark.skipif(not (cuda_available or xpu_available), reason = "needs CUDA or XPU")
 
 
 def _reference(X, weight, scale, block):
@@ -27,7 +31,6 @@ def test_tiny_non_tileable_forward_backward_matches_reference():
     from unsloth.kernels.fp8 import FP8BlockQuantLinear
 
     torch.manual_seed(0)
-    dev = "cuda"
     block = [128, 128]
     m, n = 8, 8  # non-tileable, in-dim % 128 != 0
     weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16)  # (out=m, in=n)
@@ -50,7 +53,6 @@ def test_e8m0_scale_is_upcast_and_runs():
     if not hasattr(torch, "float8_e8m0fnu"):
         pytest.skip("torch build lacks float8_e8m0fnu")
 
-    dev = "cuda"
     m, n = 8, 8
     weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16)
     scale = (torch.rand(1, 1, device = dev) + 1.0).to(torch.float8_e8m0fnu)
@@ -70,7 +72,6 @@ def test_rectangular_block_dequant_matches_reference():
     from unsloth.kernels.fp8 import _blockwise_weight_dequant_any_shape
 
     torch.manual_seed(0)
-    dev = "cuda"
     block = [64, 128]
     m, n = 64, 256  # evenly tiled: 64 % 64 == 0, 256 % 128 == 0
     weight = torch.randn(m, n, device = dev, dtype = torch.bfloat16)
@@ -94,7 +95,6 @@ def test_e8m0_scale_preserves_non_default_block_size_attr():
         pytest.skip("torch build lacks float8_e8m0fnu")
 
     torch.manual_seed(0)
-    dev = "cuda"
     block = [64, 64]
     # in-dim 96 is not divisible by block[1]=64 -> forward takes the torch dequant
     # fallback (no fp8 matmul kernel). Scale shape (2, 2) validates for [64, 64] but

tests/utils/perplexity_eval.py

@@ -2,6 +2,9 @@ from tqdm import tqdm
 import torch
 import pandas as pd
 
+# DEVICE_TYPE_TORCH, not DEVICE_TYPE: the latter can be "hip"/"mlx", which .to() rejects.
+from unsloth.device_type import DEVICE_TYPE_TORCH
+
 model_comparison_results = {}
 
 
@@ -17,7 +20,7 @@ def ppl_model(model, tokenizer, dataset):
         for begin_loc in range(0, seq_len, stride):
             end_loc = min(begin_loc + max_length, seq_len)
             trg_len = end_loc - prev_end_loc
-            input_ids = encodings.input_ids[:, begin_loc:end_loc].to("cuda")
+            input_ids = encodings.input_ids[:, begin_loc:end_loc].to(DEVICE_TYPE_TORCH)
             target_ids = input_ids.clone()
             target_ids[:, :-trg_len] = -100
             pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0

tests/utils/test_batched_leftpad_generation_gpu.py

@@ -4,7 +4,7 @@ Greedy generation in a left-padded batch must match solo batch-size-1
 generation for the first PREFIX_TOKENS tokens (the bug makes padded rows
 diverge into garbage immediately; a full-length match would be flaky due to
 benign batch-numerics tie-flips deep in the sequence) and must not be
-gibberish. Skipped without CUDA. Run: `python -m pytest
+gibberish. Skipped without a GPU. Run: `python -m pytest
 tests/utils/test_batched_leftpad_generation_gpu.py -v`.
 """
 
@@ -12,8 +12,19 @@ import pytest
 import torch
 
 cuda_available = torch.cuda.is_available()
+xpu_available = hasattr(torch, "xpu") and torch.xpu.is_available()
+device = "cuda" if cuda_available else "xpu" if xpu_available else "cpu"
 
-pytestmark = pytest.mark.skipif(not cuda_available, reason = "requires a CUDA GPU")
+# Non-strict rather than CUDA-only: keeps the XPU divergence visible, and goes
+# green by itself once XPU generation is fixed.
+pytestmark = [
+    pytest.mark.skipif(not (cuda_available or xpu_available), reason = "requires a CUDA or XPU GPU"),
+    pytest.mark.xfail(
+        xpu_available and not cuda_available,
+        reason = "batched left-padded generation diverges on XPU",
+        strict = False,
+    ),
+]
 
 MODEL_NAME = "unsloth/Qwen2.5-0.5B-Instruct"
 MAX_NEW_TOKENS = 32
@@ -53,7 +64,7 @@ def _chat(tokenizer, prompt):
 
 def _generate(model, tokenizer, texts):
     inputs = tokenizer(texts, return_tensors = "pt", padding = True, add_special_tokens = False).to(
-        "cuda"
+        device
     )
     with torch.inference_mode():
         out = model.generate(

tests/utils/test_packing.py

@@ -44,6 +44,8 @@ def _build_packed_training_setup(tmp_path, device):
             dtype = torch.bfloat16
         else:
             dtype = torch.float16
+    elif device.type == "xpu":
+        dtype = torch.bfloat16
 
     try:
         model, tokenizer = FastLanguageModel.from_pretrained(
@@ -76,8 +78,8 @@ def _build_packed_training_setup(tmp_path, device):
         max_length = 64,
         logging_steps = 1,
         max_steps = 1,
-        fp16 = device.type == "cuda" and not torch.cuda.is_bf16_supported(),
-        bf16 = device.type == "cuda" and torch.cuda.is_bf16_supported(),
+        fp16 = dtype == torch.float16,
+        bf16 = dtype == torch.bfloat16,
         dataset_num_proc = 1,
         output_dir = str(tmp_path),
         packing = True,
@@ -974,7 +976,12 @@ def test_enable_sample_packing():
 
 
 def test_enable_sample_packing_trl_collator(tmp_path):
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    elif torch.xpu.is_available():
+        device = torch.device("xpu")
+    else:
+        device = torch.device("cpu")
     model, _, trainer, _ = _build_packed_training_setup(tmp_path, device)
 
     enable_sample_packing(model, trainer)
@@ -1030,7 +1037,12 @@ def test_enable_padding_free_metadata():
 
 
 def test_packing_sdpa(tmp_path):
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    elif torch.xpu.is_available():
+        device = torch.device("xpu")
+    else:
+        device = torch.device("cpu")
     model, batch, trainer, llama_mod = _build_packed_training_setup(tmp_path, device)
 
     assert "packed_seq_lengths" in batch

tests/utils/test_qat.py

@@ -130,8 +130,14 @@ def _test_fake_quantizers_are_called(
                     # Weight fake quantizers must always be called.
                     assert child.weight_fake_quantizer.count == 1
 
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    elif torch.xpu.is_available():
+        device = torch.device("xpu")
+    else:
+        pytest.skip("No GPU available")
     for k, v in example_inputs.items():
-        example_inputs[k] = v.cuda()
+        example_inputs[k] = v.to(device)
     model.apply(_swap_fake_quantizers)
     model(**example_inputs)
     model.apply(_assert_fake_quantizers_are_called)

tests/utils/test_rope_scaling_drift.py

@@ -15,18 +15,20 @@ import pytest
 import torch
 
 
-def _has_real_cuda():
-    try:
-        torch.zeros(1).to("cuda")
-        return True
-    except Exception:
-        return False
+def _has_real_gpu():
+    for backend in ("cuda", "xpu"):
+        try:
+            torch.zeros(1).to(backend)
+            return True
+        except Exception:
+            pass
+    return False
 
 
-HAS_REAL_CUDA = _has_real_cuda()
-requires_cuda = pytest.mark.skipif(
-    not HAS_REAL_CUDA,
-    reason = "LlamaRotaryEmbedding builds per-device CUDA caches in __init__",
+HAS_REAL_GPU = _has_real_gpu()
+requires_gpu = pytest.mark.skipif(
+    not HAS_REAL_GPU,
+    reason = "LlamaRotaryEmbedding builds per-device caches in __init__ (needs CUDA or XPU)",
 )
 
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -360,7 +362,7 @@ def _cos_at_position(rot, position):
 # --- Layer 3: CUDA behavioral guard (real instantiation needs a device) ---
 
 
-@requires_cuda
+@requires_gpu
 def test_constructor_applies_llama3_scaling():
     config = _make_config(LLAMA3_ROPE_SCALING)
     rot = _unsloth_rotary(config)
@@ -371,7 +373,7 @@ def test_constructor_applies_llama3_scaling():
     ), "LlamaRotaryEmbedding built from a llama3 config produced unscaled inv_freq (issue #2405)."
 
 
-@requires_cuda
+@requires_gpu
 def test_constructor_unscaled_config_uses_vanilla_inv_freq():
     rot = _unsloth_rotary(_make_config(None))
     got = rot.inv_freq.float().cpu()
@@ -381,7 +383,7 @@ def test_constructor_unscaled_config_uses_vanilla_inv_freq():
     ), "LlamaRotaryEmbedding with no rope_scaling must use the vanilla inv_freq"
 
 
-@requires_cuda
+@requires_gpu
 def test_cos_cache_differs_between_scaled_and_unscaled_at_long_position():
     scaled = _unsloth_rotary(_make_config(LLAMA3_ROPE_SCALING))
     unscaled = _unsloth_rotary(_make_config(None))
@@ -397,7 +399,7 @@ def test_cos_cache_differs_between_scaled_and_unscaled_at_long_position():
     )
 
 
-@requires_cuda
+@requires_gpu
 def test_extended_cache_keeps_scaling_after_growth():
     scaled = _unsloth_rotary(_make_config(LLAMA3_ROPE_SCALING))
     # Grow past the initial cache size (mirrors long-context decode).
@@ -456,7 +458,7 @@ def _build_longrope_rotary():
     return rot, config
 
 
-@requires_cuda
+@requires_gpu
 @pytest.mark.parametrize(
     "build", [_build_llama3_rotary, _build_longrope_rotary], ids = ["llama3", "longrope"]
 )