From ed829b672a6e465819f2dbd1aeab862fa53022d8 Mon Sep 17 00:00:00 2001 From: Yuxiao Cheng <46640740+jarrycyx@users.noreply.github.com> Date: Fri, 14 Nov 2025 19:06:33 +0800 Subject: [PATCH] Fix: prevent rope_embedding AssertionError by checking kv_seq_len before reuse (#3578) * fix: add kv_seq_len boundary check before reusing RoPE embeddings Prevented AssertionError in rope_embedding.forward when kv_seq_len exceeds the cached rope size. Added condition to verify kv_seq_len <= position_embeddings[0].shape[0] before reuse, ensuring dynamic extension triggers correctly. Fixes #3036 #3216 * fix falcon h1 --------- Co-authored-by: jarrycyx --- unsloth/models/falcon_h1.py | 2 +- unsloth/models/llama.py | 2 +- unsloth/models/qwen3.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py index 70223a3c20..3010d37163 100644 --- a/unsloth/models/falcon_h1.py +++ b/unsloth/models/falcon_h1.py @@ -116,7 +116,7 @@ def FalconH1Attention_fast_forward( if past_key_value is not None: kv_seq_len += past_key_value[0].shape[-2] - if position_embeddings: + if position_embeddings and kv_seq_len <= position_embeddings[0].shape[0]: cos, sin = position_embeddings else: # Extend RoPE dynamically to fit in VRA diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index caf259c0bd..502cbec8bd 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -566,7 +566,7 @@ def LlamaAttention_fast_forward( if past_key_value is not None: kv_seq_len += past_key_value[0].shape[-2] - if position_embeddings: + if position_embeddings and kv_seq_len <= position_embeddings[0].shape[0]: cos, sin = position_embeddings else: # Extend RoPE dynamically to fit in VRA diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py index 3d905fe06a..1a0d641552 100644 --- a/unsloth/models/qwen3.py +++ b/unsloth/models/qwen3.py @@ -110,7 +110,7 @@ def Qwen3Attention_fast_forward( if past_key_value is not None: kv_seq_len += past_key_value[0].shape[-2] - if position_embeddings: + if position_embeddings and kv_seq_len <= position_embeddings[0].shape[0]: cos, sin = position_embeddings else: # Extend RoPE dynamically to fit in VRA