Move the attention mask to each layer's device in the fast decode loops (#12290)

* Move the attention mask to each layer's device in the fast decode loops

* Keep the fix to the decode loops that are dispatched (llama family, gemma)

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

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

* Gate the two-GPU test on a real CUDA device before counting devices

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

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

---------

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

tests/python/test_fast_inference_mask_device.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team.
+"""The fast decode loops call self_attn directly, so the attention mask must follow each layer's device.
+
+With layers split across GPUs (device_map="auto" on 2 x T4), batched generation died in SDPA with
+"attn_bias is on cuda:0, different from other tensors on cuda:1".
+"""
+
+from __future__ import annotations
+
+import importlib
+
+import pytest
+from real_accelerator import has_real_cuda
+import torch
+
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.skipif(not has_real_cuda(), reason = "fast inference needs a CUDA device"),
+]
+
+PROMPTS = [
+    "Hello",
+    "The capital of France is a city called",
+    "One two three four five six seven eight",
+]
+
+
+def _load(model_id, device_map):
+    from unsloth import FastLanguageModel
+
+    model, tok = FastLanguageModel.from_pretrained(
+        model_id, max_seq_length = 128, load_in_4bit = False, dtype = torch.float16, device_map = device_map
+    )
+    FastLanguageModel.for_inference(model)
+    tok.padding_side = "left"
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    return model, tok
+
+
+def _generate(model, tok):
+    batch = tok(PROMPTS, return_tensors = "pt", padding = True).to("cuda:0")
+    with torch.no_grad():
+        out = model.generate(
+            **batch, max_new_tokens = 8, do_sample = False, pad_token_id = tok.pad_token_id
+        )
+    return out[:, batch["input_ids"].shape[1] :].tolist()
+
+
+@pytest.mark.parametrize(
+    "model_id",
+    [
+        "trl-internal-testing/tiny-Qwen3ForCausalLM",
+        "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5-Coder",
+    ],
+)
+def test_decode_loop_moves_the_mask_to_every_layer(monkeypatch, model_id):
+    mod = importlib.import_module("unsloth.models.llama")
+    real = mod.move_to_device
+    mask_moves = []
+
+    def spy(device, *tensors):
+        mask_moves.extend(t for t in tensors if t.dtype == torch.bool and t.dim() == 4)
+        return real(device, *tensors)
+
+    monkeypatch.setattr(mod, "move_to_device", spy)
+    model, tok = _load(model_id, {"": 0})
+    _generate(model, tok)
+    # 7 decode steps after prefill, each moving the mask once per layer.
+    assert len(mask_moves) >= 7 * model.config.num_hidden_layers
+
+
+@pytest.mark.skipif(
+    not has_real_cuda() or torch.cuda.device_count() < 2,
+    reason = "needs two GPUs to split the layers",
+)
+@pytest.mark.parametrize(
+    "model_id",
+    [
+        "trl-internal-testing/tiny-Qwen3ForCausalLM",
+        "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5-Coder",
+    ],
+)
+def test_split_layers_generate_like_one_gpu(model_id):
+    model, tok = _load(model_id, {"": 0})
+    single = _generate(model, tok)
+    del model
+    n = (
+        importlib.import_module("transformers")
+        .AutoConfig.from_pretrained(model_id)
+        .num_hidden_layers
+    )
+    device_map = {"model.embed_tokens": 0, "model.rotary_emb": 0, "model.norm": 1, "lm_head": 1}
+    device_map.update({f"model.layers.{i}": 0 if i < max(1, n // 2) else 1 for i in range(n)})
+    model, tok = _load(model_id, device_map)
+    assert {str(next(layer.parameters()).device) for layer in model.model.layers} == {
+        "cuda:0",
+        "cuda:1",
+    }
+    assert _generate(model, tok) == single

unsloth/models/gemma.py

@@ -193,6 +193,8 @@ def GemmaModel_fast_forward_inference(
     for idx, decoder_layer in enumerate(self.model.layers):
         layer_device, device_index = per_layer_device(decoder_layer)
         hidden_states, position_ids = move_to_device(layer_device, hidden_states, position_ids)
+        if attention_mask is not None:
+            attention_mask = move_to_device(layer_device, attention_mask)
 
         residual = hidden_states
         hidden_states = fast_rms_layernorm_inference_gemma(

unsloth/models/llama.py

@@ -1374,6 +1374,9 @@ def _LlamaModel_fast_forward_inference(
         for idx, decoder_layer in enumerate(self.model.layers):
             layer_device, device_index = per_layer_device(decoder_layer)
             X, residual, position_ids = move_to_device(layer_device, X, residual, position_ids)
+            # self_attn is called directly, so no accelerate hook moves the mask to a split layer's device.
+            if attention_mask is not None:
+                attention_mask = move_to_device(layer_device, attention_mask)
             residual.copy_(X)
             X = fast_rms_layernorm_inference(
                 decoder_layer.input_layernorm,