| 83 | self.dense = torch.nn.Linear(self.hidden_size, self.hidden_size) |
| 84 | |
| 85 | def forward( |
| 86 | self, |
| 87 | hidden_states, |
| 88 | attention_mask, |
| 89 | layer_past=None, |
| 90 | get_key_value=False, |
| 91 | prompt_length=None, |
| 92 | context_length=None, |
| 93 | ): |
| 94 | # hidden_states: [sq, b, h] |
| 95 | |
| 96 | # ===================== |
| 97 | # Query, Key, and Value |
| 98 | # ===================== |
| 99 | |
| 100 | query_layer = self.query(hidden_states) |
| 101 | key_layer = self.key(hidden_states) |
| 102 | value_layer = self.value(hidden_states) |
| 103 | |
| 104 | new_query_layer_shape = query_layer.size()[:-1] + \ |
| 105 | (self.num_attention_heads, |
| 106 | self.hidden_size_per_attention_head) |
| 107 | query_layer = query_layer.view(*new_query_layer_shape) |
| 108 | |
| 109 | new_query_layer_shape = key_layer.size()[:-1] + \ |
| 110 | (self.num_attention_heads, |
| 111 | self.hidden_size_per_attention_head) |
| 112 | key_layer = key_layer.view(*new_query_layer_shape) |
| 113 | |
| 114 | new_query_layer_shape = value_layer.size()[:-1] + \ |
| 115 | (self.num_attention_heads, |
| 116 | self.hidden_size_per_attention_head) |
| 117 | value_layer = value_layer.view(*new_query_layer_shape) |
| 118 | |
| 119 | # ================================== |
| 120 | # Adjust key and value for inference |
| 121 | # ================================== |
| 122 | |
| 123 | if layer_past is not None: |
| 124 | past_key, past_value = layer_past |
| 125 | key_layer = torch.cat((past_key.type_as(key_layer), |
| 126 | key_layer), dim=0) |
| 127 | value_layer = torch.cat((past_value.type_as(value_layer), |
| 128 | value_layer), dim=0) |
| 129 | if get_key_value: |
| 130 | present = (key_layer, value_layer) |
| 131 | |
| 132 | # =================================== |
| 133 | # Raw attention scores. [b, np, sq, sk] |
| 134 | # =================================== |
| 135 | |
| 136 | # [b, np, sq, sk] |
| 137 | output_size = (query_layer.size(1), |
| 138 | query_layer.size(2), |
| 139 | query_layer.size(0), |
| 140 | key_layer.size(0)) |
| 141 | |
| 142 | # [sq, b, np, hn] -> [sq, b * np, hn] |