File size: 1,107 Bytes
c92e699 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | # 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")
|