| # SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| def decode_step_state(start_pos, iteration, prompt_length, max_seq_len): | |
| """Return the next position/input and number of cache entries written so far.""" | |
| current_position = start_pos + iteration | |
| if current_position >= max_seq_len: | |
| raise ValueError( | |
| f"decode position {current_position} exceeds the configured maximum sequence length {max_seq_len}" | |
| ) | |
| next_position = start_pos + iteration + 1 | |
| next_token_index = iteration + 1 | |
| if next_token_index >= prompt_length: | |
| next_token_index = None | |
| num_written = iteration + 1 | |
| return next_position, next_token_index, num_written | |
| def teacher_forced_decode_token(*, reference_token=None, device_token=None): | |
| """Choose one sampled token for every model participating in a comparison.""" | |
| if reference_token is not None: | |
| return reference_token | |
| if device_token is not None: | |
| return device_token | |
| raise ValueError("a reference or device token is required to continue decoding") | |