Skip to content

Commit

Permalink
Cleaned code.
Browse files Browse the repository at this point in the history
Signed-off-by: gyou2021 <ganmei.you@intel.com>
  • Loading branch information
gyou2021 committed Mar 4, 2025
1 parent fbb2a95 commit 3ca79ea
Showing 1 changed file with 2 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -1247,7 +1247,7 @@ def prefill_forward(
f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
f" {attn_weights.size()}"
)
assert attention_mask is not None

if attention_mask is not None:
if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
raise ValueError(
Expand Down Expand Up @@ -1311,7 +1311,7 @@ def decode_forward(

cos, sin = self.rotary_emb(q_pe, seq_len=kv_seq_len)
q_pe = apply_rotary_pos_emb(q_pe, cos, sin, q_position_ids)
q_nope = torch.matmul(q_nope.transpose(0, 1), self.q_absorb).transpose(0, 1) # opti
q_nope = torch.matmul(q_nope.transpose(0, 1), self.q_absorb).transpose(0, 1)
compressed_kv, k_pe = self.compress_kv(hidden_states_kv, kv_position_ids)

# update & get all compressed_kv, k_pe
Expand Down

0 comments on commit 3ca79ea

Please sign in to comment.