| 1243 | |
| 1244 | |
| 1245 | class MTPDraftProvider(DraftProvider): |
| 1246 | batched_draft = True |
| 1247 | sampled_batch_draft = True |
| 1248 | |
| 1249 | @dataclass |
| 1250 | class DraftManyState: |
| 1251 | result_index: int |
| 1252 | seq_id: int |
| 1253 | first_pos: int |
| 1254 | keep_len: int |
| 1255 | n_predict: int |
| 1256 | token: int |
| 1257 | drafted: List[int] |
| 1258 | embedding: np.ndarray |
| 1259 | cache_key: Tuple[Tuple[int, ...], int] |
| 1260 | |
| 1261 | def __init__( |
| 1262 | self, |
| 1263 | *, |
| 1264 | model: "Model", |
| 1265 | draft_model: Any, |
| 1266 | context_params: Any, |
| 1267 | num_pred_tokens: int, |
| 1268 | top_k: int, |
| 1269 | p_min: float, |
| 1270 | ) -> None: |
| 1271 | self.target_ctx = model.ctx |
| 1272 | self.model = draft_model |
| 1273 | self.n_seq_max = model.n_seq_max |
| 1274 | self.n_vocab = model.n_vocab |
| 1275 | self.n_embd = int(llama_cpp.llama_model_n_embd_out(self.model)) |
| 1276 | if self.n_embd <= 0: |
| 1277 | self.n_embd = int(llama_cpp.llama_model_n_embd(self.model)) |
| 1278 | if self.n_embd != model.n_embd: |
| 1279 | raise RuntimeError( |
| 1280 | "MTP draft model output embedding size must match target model " |
| 1281 | f"embedding size ({self.n_embd} != {model.n_embd})" |
| 1282 | ) |
| 1283 | self.num_pred_tokens = max(0, int(num_pred_tokens)) |
| 1284 | self.top_k = max(1, int(top_k)) |
| 1285 | self.p_min = max(0.0, min(1.0, float(p_min))) |
| 1286 | self.ctx = llama_cpp.llama_init_from_model(self.model, context_params) |
| 1287 | if self.ctx is None: |
| 1288 | raise RuntimeError("failed to create MTP draft context") |
| 1289 | ctx_other = llama_cpp_ext.llama_get_ctx_other(self.ctx) |
| 1290 | self.is_mem_shared = bool(ctx_other and ctx_other == self.target_ctx) |
| 1291 | self.n_mtp_layers = max( |
| 1292 | 1, |
| 1293 | int(llama_cpp.llama_model_n_layer_nextn(self.model)), |
| 1294 | ) |
| 1295 | self.chain_heads = self.n_mtp_layers > 1 and not self.is_mem_shared |
| 1296 | if self.chain_heads: |
| 1297 | self.num_pred_tokens = min(self.num_pred_tokens, self.n_mtp_layers) |
| 1298 | self.sampled_batch_draft = not self.is_mem_shared and not self.chain_heads |
| 1299 | self.n_batch = int(llama_cpp.llama_n_batch(self.ctx)) |
| 1300 | mem = llama_cpp.llama_get_memory(self.ctx) |
| 1301 | if mem is None: |
| 1302 | llama_cpp.llama_free(self.ctx) |
no outgoing calls
no test coverage detected
searching dependent graphs…