(
self,
batch: Any,
h_tgt_rows: np.ndarray,
start: int,
end: int,
previous_row_by_seq: Dict[int, int],
first_pos_by_seq: Dict[int, int],
target_rows_by_seq: Dict[int, List[int]],
aligned_by_seq: Dict[int, bool],
)
| 1601 | self.ready_pos[seq_id] = target_positions[-1] + 1 |
| 1602 | |
| 1603 | def _process_rows( |
| 1604 | self, |
| 1605 | batch: Any, |
| 1606 | h_tgt_rows: np.ndarray, |
| 1607 | start: int, |
| 1608 | end: int, |
| 1609 | previous_row_by_seq: Dict[int, int], |
| 1610 | first_pos_by_seq: Dict[int, int], |
| 1611 | target_rows_by_seq: Dict[int, List[int]], |
| 1612 | aligned_by_seq: Dict[int, bool], |
| 1613 | ) -> None: |
| 1614 | added_start_pos_by_seq: Dict[int, int] = {} |
| 1615 | added_end_pos_by_seq: Dict[int, int] = {} |
| 1616 | self._clear_batch() |
| 1617 | for index in range(start, end): |
| 1618 | if int(batch.n_seq_id[index]) != 1: |
| 1619 | raise RuntimeError("MTP requires one sequence id per batch token") |
| 1620 | seq_id = int(batch.seq_id[index][0]) |
| 1621 | if seq_id < 0 or seq_id >= self.n_seq_max: |
| 1622 | raise RuntimeError(f"MTP sequence id out of range: {seq_id}") |
| 1623 | pos = int(batch.pos[index]) |
| 1624 | first_pos = first_pos_by_seq.setdefault(seq_id, pos) |
| 1625 | aligned = first_pos <= 0 or ( |
| 1626 | self.ready[seq_id] and self.ready_pos[seq_id] == first_pos |
| 1627 | ) |
| 1628 | aligned_by_seq.setdefault(seq_id, aligned) |
| 1629 | previous_row = previous_row_by_seq.get(seq_id) |
| 1630 | if ( |
| 1631 | aligned |
| 1632 | and not self.is_mem_shared |
| 1633 | and pos >= 0 |
| 1634 | and pos >= self.context_pos[seq_id] |
| 1635 | ): |
| 1636 | slot = int(self.batch.n_tokens) |
| 1637 | self._add_batch_token( |
| 1638 | token=int(batch.token[index]), |
| 1639 | pos=pos, |
| 1640 | seq_id=seq_id, |
| 1641 | logits=False, |
| 1642 | ) |
| 1643 | if previous_row is None: |
| 1644 | self._set_batch_embedding_row(slot, self.pending_h[seq_id]) |
| 1645 | else: |
| 1646 | self._set_batch_embedding_row(slot, h_tgt_rows[previous_row]) |
| 1647 | added_start_pos_by_seq.setdefault(seq_id, pos) |
| 1648 | added_end_pos_by_seq[seq_id] = pos |
| 1649 | previous_row_by_seq[seq_id] = index |
| 1650 | target_rows_by_seq.setdefault(seq_id, []).append(index) |
| 1651 | |
| 1652 | if int(self.batch.n_tokens) > 0: |
| 1653 | self._decode_batch_for_mtp_heads(added_start_pos_by_seq) |
| 1654 | for seq_id, pos in added_end_pos_by_seq.items(): |
| 1655 | self.context_pos[seq_id] = max(self.context_pos[seq_id], pos + 1) |
| 1656 | |
| 1657 | def _draft_chain_heads( |
| 1658 | self, |
no test coverage detected