(
self,
args: GenArgs,
model_args: fast.ModelArgs,
prefill_model: fast.Transformer,
decode_model: fast.Transformer,
tokenizer: Tokenizer,
)
| 77 | return FastGen(gen_args, model_args_prefill, prefill_model, decode_model, tokenizer) |
| 78 | |
| 79 | def __init__( |
| 80 | self, |
| 81 | args: GenArgs, |
| 82 | model_args: fast.ModelArgs, |
| 83 | prefill_model: fast.Transformer, |
| 84 | decode_model: fast.Transformer, |
| 85 | tokenizer: Tokenizer, |
| 86 | ): |
| 87 | self.gen_args = args |
| 88 | self.max_seq_length = args.prompt_length + args.gen_length |
| 89 | self.model_args = model_args |
| 90 | # self.model = model |
| 91 | self.prefill_model = prefill_model |
| 92 | self.decode_model = decode_model |
| 93 | self.tokenizer = tokenizer |
| 94 | self._prefill_cuda_graph, self._prefill_compile_model, self._prefill_inputs, self._prefill_logits = None, None, None, None |
| 95 | self._generate_cuda_graph, self._generate_compile_model, self._generate_inputs, self._generate_logits = None, None, None, None |
| 96 | self._cache = None |
| 97 | start_time = time.time() |
| 98 | self._prefill_compile_model = self.compile_prefill() |
| 99 | self._generate_compile_model = self.compile_generate() |
| 100 | print(f"compiled model in {time.time() - start_time:.2f} seconds") |
| 101 | |
| 102 | def compile_prefill(self): |
| 103 |
nothing calls this directly
no test coverage detected