(
self,
*,
model: "Model",
draft_model: Any,
context_params: Any,
num_pred_tokens: int,
top_k: int,
p_min: float,
)
| 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) |
| 1303 | raise RuntimeError("failed to access MTP draft memory") |
| 1304 | self.mem = mem |
| 1305 | self.batch = llama_cpp.llama_batch_init(self.n_batch, self.n_embd, 1) |
| 1306 | self._batch_tokens = (llama_cpp.llama_token * self.n_batch)() |
| 1307 | self.batch.token = self._batch_tokens |
| 1308 | self.batch_embeddings = np.ctypeslib.as_array( |
| 1309 | self.batch.embd, |
| 1310 | shape=(self.n_batch * self.n_embd,), |
| 1311 | ) |
| 1312 | self._samplers: List[Any] = [] |
| 1313 | self.pending_h = np.zeros((self.n_seq_max, self.n_embd), dtype=np.float32) |
| 1314 | self.verify_h: List[np.ndarray] = [ |
| 1315 | np.empty((0, self.n_embd), dtype=np.float32) |
| 1316 | for _ in range(self.n_seq_max) |
| 1317 | ] |
| 1318 | self.verify_h_pos: List[List[int]] = [[] for _ in range(self.n_seq_max)] |
nothing calls this directly
no test coverage detected