(self)
| 157 | return replay |
| 158 | |
| 159 | def compile_generate(self): |
| 160 | |
| 161 | if self._cache is None: |
| 162 | self._cache = fast.make_cache( |
| 163 | args=self.model_args, |
| 164 | length=self.gen_args.gen_bsz * self.max_seq_length, |
| 165 | ) |
| 166 | |
| 167 | seq_lens = [1 for _ in range(self.gen_args.gen_bsz)] |
| 168 | kv_seq_lens = [self.gen_args.prompt_length for _ in range(self.gen_args.gen_bsz)] |
| 169 | |
| 170 | bias = AttnBias.from_seqlens( |
| 171 | q_seqlen=seq_lens, |
| 172 | kv_seqlen=kv_seq_lens, |
| 173 | kv_padding=self.max_seq_length, |
| 174 | ) |
| 175 | bias.q_seqinfo.to("cuda") |
| 176 | bias.k_seqinfo.to("cuda") |
| 177 | |
| 178 | tokens = torch.IntTensor([1] * self.gen_args.gen_bsz).cuda() |
| 179 | self._generate_inputs = (tokens, bias) |
| 180 | |
| 181 | s = torch.cuda.Stream() |
| 182 | s.wait_stream(torch.cuda.current_stream()) |
| 183 | |
| 184 | with torch.cuda.stream(s): |
| 185 | _ = self.decode_model.forward_with_attn_bias( |
| 186 | token_values=self._generate_inputs[0], |
| 187 | attn_bias=self._generate_inputs[1], |
| 188 | cache=self._cache, |
| 189 | ) |
| 190 | torch.cuda.current_stream().wait_stream(s) |
| 191 | |
| 192 | self._generate_cuda_graph = torch.cuda.CUDAGraph() |
| 193 | recording_kwargs = {} |
| 194 | if "capture_error_mode" in torch.cuda.graph.__init__.__annotations__: |
| 195 | # In PyTorch 2.1+ and nightlies from late Aug 2023, |
| 196 | # we can do this to maybe avoid watchdog-related crashes |
| 197 | recording_kwargs["capture_error_mode"] = "thread_local" |
| 198 | with torch.cuda.graph(self._generate_cuda_graph, **recording_kwargs): |
| 199 | self._generate_logits = self.decode_model.forward_with_attn_bias( |
| 200 | token_values=self._generate_inputs[0], |
| 201 | attn_bias=self._generate_inputs[1], |
| 202 | cache=self._cache, |
| 203 | ) |
| 204 | |
| 205 | def replay(tokens, seq_lens): |
| 206 | self._generate_inputs[0].copy_(tokens) |
| 207 | self._generate_inputs[1].k_seqinfo.seqlen.copy_(seq_lens) |
| 208 | |
| 209 | self._generate_cuda_graph.replay() |
| 210 | |
| 211 | return self._generate_logits |
| 212 | |
| 213 | return replay |
| 214 | |
| 215 | |
| 216 | @torch.inference_mode() |
no test coverage detected