refactor: simplify attention mask handling

2026-07-20 20:36:16 +08:00
parent a6e920fdb0
commit d7ac66fb73
6 changed files with 34 additions and 38 deletions
@@ -59,10 +59,10 @@ class EmbeddingEncoder(AutoModel):
        x = self.embed_tokens(input_ids)

        rotary_emb = self.rotary_embedding(x, position_ids)
-        attn_mask = process_attention_mask(x, position_ids, input_mask, is_causal=False)
+        attn_mask = process_attention_mask(input_mask)

        for layer in self.layers:
-            x = layer(x, rotary_emb, attn_mask, paged_cache=None)
+            x = layer(x, rotary_emb, attn_mask)

        hidden_states = self.norm(x)