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 <dzdzzd@126.com>
This commit is contained in:
Yuxiao Cheng 2025-11-14 19:06:33 +08:00 committed by GitHub
commit ed829b672a
3 changed files with 3 additions and 3 deletions

View file

@ -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

View file

@ -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

View file

@ -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