CCCCyx commited on
Commit
1e6a45b
·
verified ·
1 Parent(s): 25e81cb

Fix query RoPE when reusing cached vision keys

Browse files

Apply query RoPE on every cross-attention step, including cached-vision steps. Rotate new vision keys only once. Preserve the current NPU, idle-wait and text-cleanup code. No weights, configs, processors or tokenizer files changed. CPU fresh/cached regression passed.

Files changed (1) hide show
  1. modeling_moss_vl.py +6 -4
modeling_moss_vl.py CHANGED
@@ -644,6 +644,11 @@ class MossVLTextCrossAttention(nn.Module):
644
  query_states = query_states.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
645
  query_states = self.q_norm(query_states)
646
 
 
 
 
 
 
647
  if cross_attention_states is not None:
648
  # Key and Value from vision cross_attention_states
649
  key_states = self.k_proj(cross_attention_states)
@@ -653,10 +658,7 @@ class MossVLTextCrossAttention(nn.Module):
653
  key_states = self.k_norm(key_states)
654
  value_states = value_states.view(batch_size, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2)
655
 
656
- # Apply different RoPE for query (text position) and key (vision position)
657
- if query_position_embeddings is not None:
658
- cos, sin = query_position_embeddings
659
- query_states = apply_rotary_pos_emb_cross_attention(query_states, cos, sin)
660
 
661
  if vision_position_embeddings is not None:
662
  vision_cos, vision_sin = vision_position_embeddings
 
644
  query_states = query_states.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
645
  query_states = self.q_norm(query_states)
646
 
647
+ # Every new query needs RoPE, including reads from the vision KV cache.
648
+ if query_position_embeddings is not None:
649
+ cos, sin = query_position_embeddings
650
+ query_states = apply_rotary_pos_emb_cross_attention(query_states, cos, sin)
651
+
652
  if cross_attention_states is not None:
653
  # Key and Value from vision cross_attention_states
654
  key_states = self.k_proj(cross_attention_states)
 
658
  key_states = self.k_norm(key_states)
659
  value_states = value_states.view(batch_size, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2)
660
 
661
+ # Rotate new keys once; cached keys already include vision RoPE.
 
 
 
662
 
663
  if vision_position_embeddings is not None:
664
  vision_cos, vision_sin = vision_position_embeddings