onnxruntime
bf72f996 - Enrich logic to fuse rotaryembedding with bias and support partial RE fusion into GQA (#20300)

Commit
2 years ago
Enrich logic to fuse rotaryembedding with bias and support partial RE fusion into GQA (#20300) ### Description This PR mainly focuses on adding two functionalities: 1. Fuse RotaryEmbedding op taking output from previous layers with bias enabled. > Matmul->RotaryEmbedding -----> Matmul->Add->RotatyEmbedding 2. Fuse GQA op for partial RotaryEmbedding applied in phi-2. > # Partial rotary embedding query_rot, query_pass = ( query_states[..., : self.rotary_emb.dim], query_states[..., self.rotary_emb.dim :], ) key_rot, key_pass = ( key_states[..., : self.rotary_emb.dim], key_states[..., self.rotary_emb.dim :], ) # [batch_size, seq_length, num_heads, head_dim // config.partial_rotary_factor] query_rot, key_rot = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, position_ids) # [batch_size, seq_length, num_heads, head_dim] query_states = torch.cat((query_rot, query_pass), dim=-1) key_states = torch.cat((key_rot, key_pass), dim=-1) # Optimized graph ![image](https://github.com/microsoft/onnxruntime/assets/17421593/76fd8576-7e60-41af-9a4f-48d205fc6b56)
Parents
Loading