(
self,
idx: torch.Tensor,
input_pos: Optional[torch.Tensor] = None,
max_seq_length: Optional[int] = None,
)
| 125 | self.mask_cache = None |
| 126 | |
| 127 | def forward( |
| 128 | self, |
| 129 | idx: torch.Tensor, |
| 130 | input_pos: Optional[torch.Tensor] = None, |
| 131 | max_seq_length: Optional[int] = None, |
| 132 | ) -> torch.Tensor: |
| 133 | B, T = idx.size() |
| 134 | use_kv_cache = input_pos is not None |
| 135 | |
| 136 | block_size = self.config.block_size |
| 137 | if max_seq_length is None: |
| 138 | max_seq_length = block_size |
| 139 | if use_kv_cache: # not relevant otherwise |
| 140 | assert ( |
| 141 | max_seq_length >= T |
| 142 | ), f'Cannot forward sequence of length {T}, max seq length is only {max_seq_length}' |
| 143 | assert ( |
| 144 | max_seq_length <= block_size |
| 145 | ), f'Cannot attend to {max_seq_length}, block size is only {block_size}' |
| 146 | assert ( |
| 147 | block_size >= T |
| 148 | ), f'Cannot forward sequence of length {T}, block size is only {block_size}' |
| 149 | |
| 150 | if self.rope_cache is None: |
| 151 | self.rope_cache = self.build_rope_cache(idx) |
| 152 | # passing `attn_mask` to SDPA downgrades it to use the inefficient implementation. since we only need the mask |
| 153 | # for the kv-cache support (only during inference), we only create it in that situation |
| 154 | # this will be resolved by https://github.com/pytorch/pytorch/issues/96099 |
| 155 | if use_kv_cache and self.mask_cache is None: |
| 156 | self.mask_cache = self.build_mask_cache(idx) |
| 157 | |
| 158 | cos, sin = self.rope_cache |
| 159 | if use_kv_cache: |
| 160 | cos = cos.index_select(0, input_pos) |
| 161 | sin = sin.index_select(0, input_pos) |
| 162 | mask = self.mask_cache.index_select(2, input_pos) |
| 163 | mask = mask[:, :, :, :max_seq_length] |
| 164 | else: |
| 165 | cos = cos[:T] |
| 166 | sin = sin[:T] |
| 167 | mask = None |
| 168 | |
| 169 | # forward the model itself |
| 170 | # token embeddings of shape (b, t, n_embd) |
| 171 | x = self.transformer.wte(idx) |
| 172 | |
| 173 | if not use_kv_cache: |
| 174 | for block in self.transformer.h: |
| 175 | x, *_ = block(x, (cos, sin), max_seq_length) |
| 176 | else: |
| 177 | self.kv_caches = self.kv_caches or self.build_kv_caches( |
| 178 | x, max_seq_length, cos.size(-1) * 2 |
| 179 | ) |
| 180 | for i, block in enumerate(self.transformer.h): |
| 181 | x, self.kv_caches[i] = block( |
| 182 | x, (cos, sin), max_seq_length, mask, input_pos, self.kv_caches[i] |
| 183 | ) |
| 184 |
nothing calls this directly
no test coverage detected