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