File size: 3,887 Bytes
2415c4c | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | # SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0
"""Host-only checks for greedy token selection with logprob reporting."""
from types import SimpleNamespace
import pytest
import torch
import ttnn
from models.common.sampling.tt_sampling import TTSampling
@pytest.mark.parametrize("enabled", [False, True])
def test_argmax_logprobs_use_original_shards_and_preserve_tokens(monkeypatch, enabled):
sampler = object.__new__(TTSampling)
sampler._force_argmax_sampling = True
sampler._force_argmax_sub_core_grids = None
sampler.mesh_device = SimpleNamespace(get_num_devices=lambda: 8)
sampler.tt_ccl = SimpleNamespace(
get_and_cycle_barrier_semaphore_handle=lambda axis: None,
get_and_cycle_ag_semaphore_handles=lambda axis: None,
)
sampler.argmax_chunks_per_sync = 10
sampler.argmax_num_workers_per_link = 1
sampler._get_sampling_cluster_axis = lambda: 1
sampler._get_force_argmax_all_gather_config = lambda axis: (1, ttnn.Topology.Linear)
sampler._can_slice_valid_vocab_for_argmax = lambda: False
sampler._mask_invalid_vocab_logits = lambda value: value
sampler._untilize_for_argmax = lambda value: value
shards = SimpleNamespace(memory_config=lambda: ttnn.DRAM_MEMORY_CONFIG)
gathered, probabilities = object(), object()
tokens = SimpleNamespace(shape=(1, 1, 32))
canonical_tokens = object()
calls = []
def calculate(logits, indices):
assert logits is shards, "Logprob normalization must not count replicated full vocabularies"
assert indices is canonical_tokens
calls.append(True)
return probabilities
def argmax(value, **kwargs):
assert value is gathered
return tokens
def reshape(value, shape):
assert value is tokens and shape == (1, 1, 1, 32)
return canonical_tokens
sampler.log_probs_calculator = SimpleNamespace(enable_log_probs=enabled, calculate_log_probs=calculate)
monkeypatch.setattr(ttnn.experimental, "all_gather_async", lambda *a, **kw: gathered)
monkeypatch.setattr(ttnn, "argmax", argmax)
monkeypatch.setattr(ttnn, "reshape", reshape)
result_tokens, result_probs = sampler.forward(shards)
assert result_tokens is tokens
assert result_probs is (probabilities if enabled else None)
assert len(calls) == int(enabled)
@pytest.mark.parametrize("enabled,use_topk,expected", [(False, True, True), (True, False, True), (True, True, False)])
def test_topk_reporting_requires_candidate_pipeline(enabled, use_topk, expected):
sampler = object.__new__(TTSampling)
sampler._allow_force_argmax_sampling = True
sampler.log_probs_calculator = SimpleNamespace(enable_log_probs=enabled, _use_topk_logprobs=use_topk)
assert sampler._is_force_argmax_sampling([1], [1.0], [1.0]) is expected
def test_enabling_topk_reporting_refreshes_candidate_parameters(monkeypatch):
sampler = object.__new__(TTSampling)
sampler._allow_force_argmax_sampling = True
sampler.max_top_k = 32
sampler._sampling_dp = 1
sampler._greedy_col = object()
sampler.k_tensor, sampler.p_tensor, sampler.temp_tensor = object(), object(), object()
calculator = SimpleNamespace(enable_log_probs=False, _use_topk_logprobs=True)
calculator.set_log_probs_mode = lambda enabled, **kw: setattr(calculator, "enable_log_probs", enabled)
sampler.log_probs_calculator = calculator
copied = []
monkeypatch.setattr(ttnn, "from_torch", lambda value, **kw: value)
monkeypatch.setattr(ttnn, "copy_host_to_device_tensor", lambda source, dest: copied.append((source, dest)))
sampler.reset_params(k=[1], p=[1.0], temp=[1.0], enable_log_probs=True)
assert sampler.force_argmax_sampling is False
assert len(copied) == 4
assert copied[0][1] is sampler.k_tensor
assert torch.equal(copied[0][0], torch.tensor([1]))
|