orionweller
/

test-flex-gpt

Model card Files Files and versions Community

oweller2 commited on Nov 22

Commit

9e4ff15

•

1 Parent(s): f66abc1

try again

Files changed (1) hide show

attention.py +1 -1

attention.py CHANGED Viewed

@@ -863,7 +863,7 @@ class FlexBertUnpadRopeAttention(FlexBertAttentionBase):
         qkv = self.Wqkv(hidden_states)
         # only needed for inference when we have KV cache
-        seqlen_offset = max_seqlen * (len(cu_seqlens) - 2) if len(cu_seqlens) > 1 else 0
         # (total_seqlen, 3, nheads, headdim)
         qkv = qkv.view(-1, 3, self.num_attention_heads, self.attn_head_size)

         qkv = self.Wqkv(hidden_states)
         # only needed for inference when we have KV cache
+        seqlen_offset = max_seqlen * (cu_seqlens[0].item() // max_seqlen)
         # (total_seqlen, 3, nheads, headdim)
         qkv = qkv.view(-1, 3, self.num_attention_heads, self.attn_head_size)