Qwen3 decode: size the RoPE read by the explicit position, not only the KV length (#14156)

## Summary

#13855 made the RoPE cache lazy: it only builds the rows that are asked
for. `Qwen3Attention_fast_forward_inference` sized the cache read by the
KV length only. A decode step at an explicit position past the KV length
therefore indexed past the built rows and hit a CUDA device-side assert.
For example: a prefill without `position_ids`, then a hand-made decode
at position >= 32768. Base (before #13855) returned finite logits.
Normal `generate` is not affected.

The model-level decode loop already computes `rotary_seq_len =
max(kv_seq_len, position_ids.max() + 1)` once per step and passes it to
every layer. Qwen3 attention now sizes the RoPE extension and read from
it, so this adds no host sync. Without `rotary_seq_len` the previous
sizing is used.

## Tests

- New `tests/test_qwen3_decode_offset_rope.py` (tiny Qwen3, GPU): decode
at an explicit position past the lazily built rows. It fails on main
with the device-side assert and passes here.

Not fixed here: Gemma2, Cohere, Granite and FalconH1 read the cache the
same way. Their model-level loops do not compute `rotary_seq_len`, so a
fix there needs one host sync per decode step. Those are left for a
follow-up.

---------

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

tests/test_qwen3_decode_offset_rope.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+"""Qwen3 decode with an explicit position past the rows the prefill built.
+
+The RoPE cache only builds the rows a forward asks for, so a prefill without position_ids on a
+model whose max_position_embeddings exceeds the 32768-row starting cache leaves later rows unbuilt.
+A decode step at such a position must still read a valid row, matching the decode after a prefill
+that built every row. Runs in a subprocess because an out-of-range index is a CUDA device assert.
+"""
+
+import os
+import subprocess
+import sys
+import textwrap
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason = "needs CUDA")
+
+MODEL = "trl-internal-testing/tiny-Qwen3ForCausalLM"  # max_position_embeddings 40960
+
+SCRIPT = textwrap.dedent(
+    """
+    import sys
+    import unsloth
+    from unsloth import FastLanguageModel
+    import torch
+
+    pos = int(sys.argv[2])
+    model, _ = FastLanguageModel.from_pretrained(
+        sys.argv[1], max_seq_length = 2048, load_in_4bit = False, dtype = torch.float16
+    )
+    FastLanguageModel.for_inference(model)
+    ids = torch.tensor([[1, 2, 3, 4]], device = "cuda")
+
+    def logits(out):
+        return out.logits if hasattr(out, "logits") else out[0]
+
+    def decode(prefill_position_ids):
+        with torch.no_grad():
+            out = model(input_ids = ids, position_ids = prefill_position_ids, use_cache = True)
+            out = model(
+                input_ids = ids[:, :1],
+                past_key_values = out.past_key_values,
+                position_ids = torch.tensor([[pos]], device = "cuda"),
+                use_cache = True,
+            )
+        torch.cuda.synchronize()
+        return logits(out).float()
+
+    lazy = decode(None)
+    full = decode(torch.arange(4, device = "cuda")[None])
+    assert torch.isfinite(lazy).all()
+    torch.testing.assert_close(lazy, full, rtol = 0, atol = 0)
+    print("OK")
+    """
+)
+
+
+def test_qwen3_decode_reads_rows_past_the_lazy_prefill_cache(tmp_path):
+    script = tmp_path / "decode_offset.py"
+    script.write_text(SCRIPT)
+    env = dict(os.environ)
+    repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    env["PYTHONPATH"] = repo_root + os.pathsep + env.get("PYTHONPATH", "")
+    env.setdefault("UNSLOTH_COMPILE_LOCATION", str(tmp_path / "unsloth_compiled_cache"))
+    result = subprocess.run(
+        [sys.executable, str(script), MODEL, "35000"],
+        capture_output = True,
+        text = True,
+        env = env,
+        timeout = 900,
+    )
+    assert result.returncode == 0 and "OK" in result.stdout, (
+        result.stdout[-2000:] + result.stderr[-4000:]
+    )

unsloth/models/qwen3.py

@@ -261,8 +261,12 @@ def Qwen3Attention_fast_forward_inference(
     Kn = Kn.transpose(1, 2)
 
     # Must be done 2 steps before hitting full on a short KV cache, or it errors.
-    self.rotary_emb.extend_rope_embedding(Vn, seq_len + 2)
-    cos, sin = self.rotary_emb.get_cached(kv_seq_len, Qn.device.index)
+    # rotary_seq_len (from the model-level decode loop, no extra sync) covers explicit positions
+    # past the KV length, since the cache only builds the rows asked for.
+    rotary_seq_len = kwargs.get("rotary_seq_len")
+    rope_len = kv_seq_len if rotary_seq_len is None else max(kv_seq_len, rotary_seq_len)
+    self.rotary_emb.extend_rope_embedding(Vn, max(seq_len + 2, rope_len + 1))
+    cos, sin = self.rotary_emb.get_cached(rope_len, Qn.device.index)
     # Transformers 5.x: position_ids may be [batch, full_seq_len]; slice to last.
     if position_ids.dim() >= 2 and position_ids.shape[-1] > 1:
         position_ids = position_ids[:, -1:]