(self, batch: Any, /)
| 1550 | return self.ready_pos[seq_id] == input_length - 1 |
| 1551 | |
| 1552 | def process(self, batch: Any, /) -> None: |
| 1553 | n_tokens = int(batch.n_tokens) |
| 1554 | if ( |
| 1555 | n_tokens <= 0 |
| 1556 | or not self.target_processing_enabled |
| 1557 | or not bool(batch.token) |
| 1558 | or bool(batch.embd) |
| 1559 | ): |
| 1560 | return |
| 1561 | |
| 1562 | h_tgt = llama_cpp_ext.llama_get_embeddings_nextn(self.target_ctx) |
| 1563 | if not h_tgt: |
| 1564 | raise RuntimeError("missing target nextn embeddings for MTP") |
| 1565 | h_tgt_rows = np.ctypeslib.as_array( |
| 1566 | h_tgt, |
| 1567 | shape=(n_tokens, self.n_embd), |
| 1568 | ) |
| 1569 | |
| 1570 | previous_row_by_seq: Dict[int, int] = {} |
| 1571 | first_pos_by_seq: Dict[int, int] = {} |
| 1572 | target_rows_by_seq: Dict[int, List[int]] = {} |
| 1573 | aligned_by_seq: Dict[int, bool] = {} |
| 1574 | |
| 1575 | for start in range(0, n_tokens, self.n_batch): |
| 1576 | self._process_rows( |
| 1577 | batch, |
| 1578 | h_tgt_rows, |
| 1579 | start, |
| 1580 | min(start + self.n_batch, n_tokens), |
| 1581 | previous_row_by_seq, |
| 1582 | first_pos_by_seq, |
| 1583 | target_rows_by_seq, |
| 1584 | aligned_by_seq, |
| 1585 | ) |
| 1586 | |
| 1587 | for seq_id, rows in target_rows_by_seq.items(): |
| 1588 | if not aligned_by_seq.get(seq_id, False): |
| 1589 | self.ready[seq_id] = False |
| 1590 | continue |
| 1591 | if rows[-1] - rows[0] + 1 == len(rows): |
| 1592 | target_rows = h_tgt_rows[rows[0] : rows[-1] + 1] |
| 1593 | else: |
| 1594 | target_rows = h_tgt_rows[rows].copy() |
| 1595 | target_positions = [int(batch.pos[source_index]) for source_index in rows] |
| 1596 | self.verify_h[seq_id] = target_rows |
| 1597 | self.verify_h_pos[seq_id] = target_positions |
| 1598 | self.verify_h_rows[seq_id] = len(rows) |
| 1599 | self.pending_h[seq_id] = target_rows[-1] |
| 1600 | self.ready[seq_id] = True |
| 1601 | self.ready_pos[seq_id] = target_positions[-1] + 1 |
| 1602 | |
| 1603 | def _process_rows( |
| 1604 | self, |
nothing calls this directly
no test coverage detected