Upload YangJianVLForConditionalGeneration

modeling_yangjian.py

@@ -231,9 +231,9 @@ class OptimizedCrossAttention(nn.Module):
             attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
         
         # 构造 cu_seqlens 参数（FlashAttention 必需）
-        cu_seqlens_q = torch.arange(0, (batch_size + 1) * seq_len_q, step=seq_len_q, dtype=torch.int32, device=q.device)
+        cu_seqlens_q = torch.arange(0, (batch_size*self.num_heads + 1) * seq_len_q, step=seq_len_q, dtype=torch.int32, device=q.device)
         if self.is_cross_attention and key_value_states is not None:
-            cu_seqlens_k = torch.arange(0, (batch_size + 1) * seq_len_kv, step=seq_len_kv, dtype=torch.int32, device=k.device)
+            cu_seqlens_k = torch.arange(0, (batch_size*self.num_heads + 1) * seq_len_kv, step=seq_len_kv, dtype=torch.int32, device=k.device)
         else:
             cu_seqlens_k = cu_seqlens_q