File size: 2,329 Bytes
b025706 | 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 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | # SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0
"""Host-only state contracts for TT sampled-token log-probability mode."""
from models.common.sampling.tt_log_probs import LogProbsCalculator
def _calculator(batch_size=4):
calculator = object.__new__(LogProbsCalculator)
calculator.batch_size = batch_size
calculator.logprobs_enabled = [False] * batch_size
calculator.num_logprobs = [0] * batch_size
calculator.enable_log_probs = False
calculator.topk_logprobs_needed = False
return calculator
def test_all_false_tuple_keeps_log_probs_disabled():
calculator = _calculator()
calculator.set_log_probs_mode((False, False, False, False), num_logprobs=(0, 0, 0, 0))
assert calculator.logprobs_enabled == [False, False, False, False]
assert calculator.num_logprobs == [0, 0, 0, 0]
assert calculator.enable_log_probs is False
assert calculator.topk_logprobs_needed is False
def test_tuple_log_probs_and_counts_map_per_user_exactly():
calculator = _calculator()
calculator.set_log_probs_mode((True, False, True, False), num_logprobs=(0, 4, 7, 2))
assert calculator.logprobs_enabled == [True, False, True, False]
assert calculator.num_logprobs == [0, 4, 7, 2]
assert calculator.enable_log_probs is True
assert calculator.topk_logprobs_needed is True
def test_tuple_partial_update_respects_empty_slots():
calculator = _calculator(batch_size=6)
calculator.set_log_probs_mode((True, False), num_logprobs=(10, 15), empty_slots=[2, 5])
assert calculator.logprobs_enabled == [False, False, True, False, False, False]
assert calculator.num_logprobs == [0, 0, 10, 0, 0, 15]
assert calculator.enable_log_probs is True
def test_existing_scalar_and_list_modes_remain_unchanged():
calculator = _calculator()
calculator.set_log_probs_mode(True, num_logprobs=5)
assert calculator.logprobs_enabled == [True, True, True, True]
assert calculator.num_logprobs == [5, 5, 5, 5]
calculator.set_log_probs_mode([True, False, False, True], num_logprobs=[1, 2, 3, 4])
assert calculator.logprobs_enabled == [True, False, False, True]
assert calculator.num_logprobs == [1, 2, 3, 4]
assert calculator.enable_log_probs is True
assert calculator.topk_logprobs_needed is True
|