| |
| |
|
|
| """Host/fake contract tests for SamplingState1D. |
| |
| These tests deliberately avoid TT device construction. They verify controller |
| ordering, ownership, topology/config agreement, and exactly-once bookkeeping. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass, replace |
| from types import SimpleNamespace |
|
|
| import pytest |
| import torch |
|
|
| from models.common.modules.lazy_buffer import LazyBuffer |
| from models.common.modules.sampling import penalties_1d as penalties_module |
| from models.common.modules.sampling.params import PreparedSamplingParams, place_prepared_sampling_params |
| from models.common.modules.sampling.penalties_1d import Penalties1D |
| from models.common.modules.sampling.sampling_state_1d import SamplingState1D |
| from models.common.modules.sampling.seed_manager_1d import SeedManager1D |
|
|
|
|
| class FakeMesh: |
| def __init__(self, shape): |
| self.shape = tuple(shape) |
|
|
| def get_num_devices(self): |
| return self.shape[0] * self.shape[1] |
|
|
|
|
| class FakeBuffer: |
| def __init__(self, source): |
| self.source = source.clone() |
| self.handle = object() |
| self.updates = [] |
| self.released = False |
|
|
| def get_device_buffer(self): |
| return self.handle |
|
|
| def update(self, source): |
| self.source = source.clone() |
| self.updates.append(source.clone()) |
|
|
| def release(self): |
| self.released = True |
|
|
|
|
| @dataclass |
| class FakeSeedState: |
| capacity: int |
| active: list[bool] |
| seeds: list[int | None] |
|
|
| @property |
| def request_seeds(self): |
| return self.seeds |
|
|
|
|
| class FakeSeedManager: |
| def __init__(self, config, events): |
| self.max_batch_size = config.max_batch_size |
| self.seed_buffer = config.seeds |
| self.events = events |
|
|
| def create_state(self): |
| return FakeSeedState( |
| capacity=self.max_batch_size, |
| active=[False] * self.max_batch_size, |
| seeds=[None] * self.max_batch_size, |
| ) |
|
|
| def reset(self, state): |
| self.events.append(("seed.reset",)) |
| state.active[:] = [False] * state.capacity |
| state.seeds[:] = [None] * state.capacity |
|
|
| def admit(self, state, seeds, slots): |
| slots = tuple(slots) |
| seeds = tuple(seeds) |
| self.events.append(("seed.admit", seeds, slots)) |
| for slot, seed in zip(slots, seeds): |
| state.active[slot] = True |
| state.seeds[slot] = seed |
|
|
| def synchronize(self, state, seeds, active_slots, *, reset_batch): |
| active_slots = tuple(active_slots) |
| self.events.append(("seed.synchronize", active_slots, reset_batch)) |
| changed = [slot for slot in active_slots if not state.active[slot] or state.seeds[slot] != seeds[slot]] |
| if changed and not reset_batch: |
| raise RuntimeError( |
| "new or changed active seed slots require reset_batch=True or an explicit admit() call: " f"{changed}" |
| ) |
| active = set(active_slots) |
| for slot in range(state.capacity): |
| state.active[slot] = slot in active |
| state.seeds[slot] = seeds[slot] if slot in active else None |
|
|
| def refresh(self, state, active_slots, *, positions=None): |
| self.events.append(("seed.refresh", tuple(active_slots), positions)) |
|
|
| def refresh_prefill_replicated(self, state, slot, *, position=None): |
| self.events.append(("seed.prefill_replicated", int(slot), position)) |
| return 123 |
|
|
| def restore_defaults(self, state): |
| self.events.append(("seed.restore_defaults",)) |
|
|
| def apply_slot_remap(self, state, remap): |
| remap = tuple(remap) |
| self.events.append(("seed.remap", remap)) |
| old_active = tuple(state.active) |
| old_seeds = tuple(state.seeds) |
| state.active[:] = [old_active[source] for source in remap] |
| state.seeds[:] = [old_seeds[source] for source in remap] |
|
|
| def cleanup(self, state, live_slots): |
| live_slots = tuple(live_slots) |
| self.events.append(("seed.cleanup", live_slots)) |
| live = set(live_slots) |
| for slot in range(state.capacity): |
| if slot not in live: |
| state.active[slot] = False |
| state.seeds[slot] = None |
|
|
| def get_seed_device_buffer(self): |
| self.events.append(("seed.handle",)) |
| return self.seed_buffer.handle |
|
|
|
|
| class FakePenalties: |
| _BUFFER_NAMES = ( |
| "prompt_mask", |
| "output_mask", |
| "output_counts", |
| "output_counts_gathered", |
| "zeros", |
| "decode_src", |
| "presence_penalties", |
| "frequency_penalties", |
| "repetition_penalties", |
| "inverse_repetition_penalties", |
| ) |
|
|
| def __init__(self, config, events, *, config_updates=None): |
| buffers = {} |
| for name in self._BUFFER_NAMES: |
| shape = ( |
| (config.max_batch_size, config.vocab_size) |
| if name |
| in { |
| "prompt_mask", |
| "output_mask", |
| "output_counts", |
| "output_counts_gathered", |
| "zeros", |
| } |
| else (config.max_batch_size, 1) |
| ) |
| buffers[name] = FakeBuffer(torch.zeros(shape)) |
| self.config = replace(config, **buffers, **(config_updates or {})) |
| self.events = events |
| self.loaded = False |
| self.released = False |
|
|
| def load_device_buffers(self): |
| self.events.append(("penalty.load",)) |
| self.loaded = True |
|
|
| def init_prompt_penalties(self, params, accum, prompt_tokens): |
| self.events.append(("penalty.prompt", prompt_tokens.clone())) |
|
|
| def reset_output_tokens(self, accum, tokens=None): |
| self.events.append(("penalty.output_reset", None if tokens is None else tokens.clone())) |
|
|
| def decode_forward(self, logits, params, accum): |
| self.events.append(("penalty.decode", logits)) |
| return f"penalized:{logits}" |
|
|
| def update_output_tokens(self, accum, tokens): |
| self.events.append(("penalty.update", tokens)) |
|
|
| def release(self): |
| self.events.append(("penalty.release",)) |
| self.released = True |
| for name in self._BUFFER_NAMES: |
| getattr(self.config, name).release() |
|
|
|
|
| class FakeSampling: |
| def __init__(self, config, events): |
| self.config = config |
| self.events = events |
| self.released = False |
| self.raise_on_decode = False |
|
|
| def decode_forward(self, logits, **kwargs): |
| self.events.append(("sampling.decode", logits, kwargs)) |
| if self.raise_on_decode: |
| raise RuntimeError("sampling failed") |
| return "sampled-tokens", "sampled-logprobs" |
|
|
| def release(self): |
| self.released = True |
|
|
|
|
| def _make_controller(shape=(1, 4), *, penalty_config_updates=None): |
| events = [] |
| mesh = FakeMesh(shape) |
| seed_buffer = FakeBuffer(torch.arange(4, dtype=torch.int32)) |
| sub_core_grids = object() |
| config = SimpleNamespace( |
| vocab_size=128, |
| valid_vocab_size=127, |
| mesh_device=mesh, |
| max_batch_size=4, |
| max_top_k=32, |
| sub_core_grids=sub_core_grids, |
| seeds=seed_buffer, |
| ) |
| sampling = FakeSampling(config, events) |
| penalties = None |
|
|
| def penalties_factory(penalties_config): |
| nonlocal penalties |
| penalties = FakePenalties( |
| penalties_config, |
| events, |
| config_updates=penalty_config_updates, |
| ) |
| return penalties |
|
|
| controller = SamplingState1D( |
| sampling, |
| penalties_factory=penalties_factory, |
| seed_manager_factory=lambda sampling_config: FakeSeedManager(sampling_config, events), |
| ) |
| return controller, sampling, penalties, events |
|
|
|
|
| def _prepared( |
| *, |
| penalties=True, |
| sampling_path="topk", |
| log_probs=False, |
| slot_remap=None, |
| presence=(0.5, 0.0, 0.0, 0.0), |
| seeds=(11, None, None, None), |
| ): |
| repetition = (1.5, 1.0, 1.0, 1.0) if penalties else (1.0, 1.0, 1.0, 1.0) |
| frequency = (0.25, 0.0, 0.0, 0.0) if penalties else (0.0, 0.0, 0.0, 0.0) |
| presence = presence if penalties else (0.0, 0.0, 0.0, 0.0) |
| row_path = "argmax" if sampling_path == "argmax" else "topk" |
| logprob_modes = ("sampled_token", "none", "none", "none") if log_probs else ("none",) * 4 |
| return PreparedSamplingParams( |
| top_k=(1, 5, 1, 1), |
| top_p=(0.0, 0.9, 0.0, 0.0), |
| temperature=(1.0, 1.25, 1.0, 1.0), |
| presence_penalty=presence, |
| frequency_penalty=frequency, |
| repetition_penalty=repetition, |
| seeds=seeds, |
| enable_log_probs=(log_probs, False, False, False), |
| num_logprobs=(0, 0, 0, 0), |
| logprob_modes=logprob_modes, |
| greedy_mask=(sampling_path == "argmax", sampling_path == "argmax", False, False), |
| row_paths=(row_path, row_path, "inactive", "inactive"), |
| active_mask=(True, True, False, False), |
| sampling_path=sampling_path, |
| active_rows=2, |
| batch_size=4, |
| max_device_top_k=32, |
| prompt_tokens=torch.tensor([[1, 2], [3, -1]]), |
| output_tokens=torch.tensor([[7, 8], [9, -1]]), |
| slot_remap=slot_remap, |
| ) |
|
|
|
|
| def _event_names(events): |
| return [event[0] for event in events] |
|
|
|
|
| def test_constructor_derives_exact_penalty_contract_from_borrowed_sampler(): |
| controller, sampling, penalties, _ = _make_controller() |
|
|
| assert controller.sampling is sampling |
| assert penalties.config.mesh_device is sampling.config.mesh_device |
| assert penalties.config.vocab_size == sampling.config.vocab_size |
| assert penalties.config.max_batch_size == sampling.config.max_batch_size |
| assert penalties.config.sub_core_grids is sampling.config.sub_core_grids |
| assert controller.seed_manager.seed_buffer is sampling.config.seeds |
|
|
|
|
| def test_constructor_rejects_non_1d_topology_before_constructing_state(expect_error): |
| with expect_error(ValueError, "only supports 1D"): |
| _make_controller(shape=(2, 2)) |
|
|
|
|
| @pytest.mark.parametrize( |
| "updates, message", |
| [ |
| ({"vocab_size": 256}, "vocab_size"), |
| ({"max_batch_size": 8}, "max_batch_size"), |
| ({"mesh_device": FakeMesh((1, 4))}, "mesh_device"), |
| ({"sub_core_grids": object()}, "sub_core_grids"), |
| ], |
| ) |
| def test_constructor_rejects_penalty_contract_drift(updates, message, expect_error): |
| with expect_error(ValueError, message): |
| _make_controller(penalty_config_updates=updates) |
|
|
|
|
| def test_constructor_threads_unsalted_seed_policy_to_native_manager(): |
| events = [] |
| config = SimpleNamespace( |
| vocab_size=128, |
| mesh_device=FakeMesh((1, 4)), |
| max_batch_size=4, |
| max_top_k=32, |
| sub_core_grids=object(), |
| seeds=FakeBuffer(torch.arange(4, dtype=torch.int32)), |
| ) |
| controller = SamplingState1D( |
| FakeSampling(config, events), |
| penalties_factory=lambda penalties_config: FakePenalties(penalties_config, events), |
| salt_duplicate_seeds=False, |
| ) |
|
|
| assert controller.salt_duplicate_seeds is False |
| assert isinstance(controller.seed_manager, SeedManager1D) |
| assert controller.seed_manager.salt_duplicate_seeds is False |
|
|
| state = controller.seed_manager.create_state() |
| controller.seed_manager.admit(state, [1234, 1234, 1234, 1234], [0, 1, 2, 3]) |
| assert state.snapshot().salts == (0, 0, 0, 0) |
|
|
|
|
| def test_create_state_materializes_noop_params_and_clears_history(expect_error): |
| controller, _, penalties, events = _make_controller() |
| state = controller.create_state() |
|
|
| assert state.seed_state.capacity == 4 |
| assert state.active_mask == (False,) * 4 |
| assert penalties.loaded |
| assert torch.equal(penalties.config.presence_penalties.source, torch.zeros(4, 1)) |
| assert torch.equal(penalties.config.frequency_penalties.source, torch.zeros(4, 1)) |
| assert torch.equal(penalties.config.repetition_penalties.source, torch.ones(4, 1)) |
| assert torch.equal(penalties.config.inverse_repetition_penalties.source, torch.ones(4, 1)) |
| assert _event_names(events)[-3:] == ["seed.reset", "penalty.prompt", "penalty.output_reset"] |
|
|
| with expect_error(RuntimeError, "already live"): |
| controller.create_state() |
|
|
|
|
| def test_admit_updates_params_seeds_histories_and_static_identity(): |
| controller, _, penalties, events = _make_controller() |
| state = controller.create_state() |
| prepared = _prepared(log_probs=True) |
| events.clear() |
|
|
| controller.admit(state, prepared) |
|
|
| assert state.active_slots == (0, 1) |
| assert state.static_identity.sampling_path == "topk" |
| assert state.static_identity.penalties_enabled |
| assert state.static_identity.log_probs_enabled |
| assert state.static_identity.logprob_modes == ("sampled_token", "none") |
| assert torch.allclose( |
| penalties.config.presence_penalties.source.reshape(-1), |
| torch.tensor([0.5, 0.0, 0.0, 0.0]), |
| ) |
| assert torch.allclose( |
| penalties.config.inverse_repetition_penalties.source.reshape(-1), |
| torch.tensor([1.0 / 1.5, 1.0, 1.0, 1.0]), |
| ) |
| assert _event_names(events) == [ |
| "seed.admit", |
| "seed.synchronize", |
| "penalty.prompt", |
| "penalty.output_reset", |
| ] |
| assert torch.equal(events[-2][1], prepared.prompt_tokens) |
| assert torch.equal(events[-1][1], prepared.output_tokens) |
|
|
|
|
| def test_admit_validates_repetition_history_before_mutating_seed_state(expect_error): |
| controller, _, _, events = _make_controller() |
| state = controller.create_state() |
| prepared = replace(_prepared(), prompt_tokens=None) |
| events.clear() |
|
|
| with expect_error(ValueError, "prompt_tokens are required"): |
| controller.admit(state, prepared) |
|
|
| assert events == [] |
| assert not any(state.seed_state.active) |
|
|
|
|
| def test_prepared_capability_must_match_the_borrowed_sampler(expect_error): |
| controller, _, _, _ = _make_controller() |
| state = controller.create_state() |
| prepared = replace(_prepared(), max_device_top_k=64) |
|
|
| with expect_error(ValueError, "max_device_top_k"): |
| controller.admit(state, prepared) |
|
|
|
|
| def test_decode_reset_validates_history_before_mutating_seed_state(expect_error): |
| controller, _, _, events = _make_controller() |
| state = controller.create_state() |
| prepared = replace(_prepared(), prompt_tokens=None) |
| events.clear() |
|
|
| with expect_error(ValueError, "prompt_tokens are required"): |
| controller.synchronize_decode(state, prepared, reset_batch=True) |
|
|
| assert events == [] |
| assert not any(state.seed_state.active) |
|
|
|
|
| def test_refresh_updates_dynamic_params_and_seeds_without_rebuilding_history(): |
| controller, _, penalties, events = _make_controller() |
| state = controller.create_state() |
| controller.admit(state, _prepared()) |
| changed = _prepared(presence=(0.75, 0.0, 0.0, 0.0)) |
| positions = (3, 4, -1, -1) |
| events.clear() |
|
|
| controller.refresh_dynamic_inputs(state, changed, positions=positions) |
|
|
| assert _event_names(events) == ["seed.synchronize", "seed.refresh"] |
| assert events[-1][2] is positions |
| assert penalties.config.presence_penalties.source[0, 0].item() == pytest.approx(0.75) |
|
|
|
|
| def test_prefill_admission_keeps_seed_on_decode_slot_and_samples_request_order(): |
| controller, _, _, events = _make_controller() |
| state = controller.create_state() |
| prepared = replace( |
| _prepared(), |
| active_mask=(True, False, False, False), |
| active_rows=1, |
| row_paths=("topk", "inactive", "inactive", "inactive"), |
| seeds=(77, None, None, None), |
| ) |
| events.clear() |
|
|
| controller.admit_prefill(state, prepared, slots=(3,), positions=(12,)) |
| output = controller.prefill_forward( |
| "logits", |
| state, |
| prepared, |
| k=object(), |
| p=object(), |
| temp=object(), |
| ) |
|
|
| assert state.seed_state.active == [False, False, False, True] |
| assert state.seed_state.seeds[3] == 77 |
| assert ("seed.prefill_replicated", 3, 12) in events |
| assert controller.penalties.config.presence_penalties.source.reshape(-1).tolist() == [0.5] * 4 |
| assert output == ("sampled-tokens", "sampled-logprobs") |
| assert state.penalty_history_valid is False |
|
|
|
|
| def test_no_penalty_prefill_keeps_history_valid_for_gap_decode_without_reset(): |
| controller, _, _, _ = _make_controller() |
| state = controller.create_state() |
| prepared = replace( |
| _prepared(penalties=False), |
| active_mask=(True, False, False, False), |
| active_rows=1, |
| row_paths=("topk", "inactive", "inactive", "inactive"), |
| seeds=(77, None, None, None), |
| ) |
|
|
| controller.admit_prefill(state, prepared, slots=(3,), positions=(12,)) |
| controller.prefill_forward( |
| "logits", |
| state, |
| prepared, |
| k=object(), |
| p=object(), |
| temp=object(), |
| ) |
| decode_prepared = place_prepared_sampling_params(prepared, (3,)) |
|
|
| assert state.penalty_history_valid is True |
| assert controller.synchronize_decode(state, decode_prepared, reset_batch=False) is decode_prepared |
| assert state.active_slots == (3,) |
|
|
|
|
| def test_decode_identity_remap_admits_new_slot_without_reset_and_preserves_survivor(): |
| controller, _, _, events = _make_controller() |
| state = controller.create_state() |
| initial = replace( |
| _prepared(penalties=False), |
| active_mask=(True, False, False, False), |
| active_rows=1, |
| row_paths=("topk", "inactive", "inactive", "inactive"), |
| seeds=(11, None, None, None), |
| ) |
| controller.admit(state, initial) |
| events.clear() |
| expanded = _prepared( |
| penalties=False, |
| seeds=(11, 22, None, None), |
| slot_remap=(0, 1, 2, 3), |
| ) |
|
|
| consumed = controller.synchronize_decode(state, expanded, reset_batch=False) |
|
|
| assert consumed.slot_remap is None |
| assert state.seed_state.active == [True, True, False, False] |
| assert state.seed_state.seeds == [11, 22, None, None] |
| assert _event_names(events)[:3] == ["seed.remap", "seed.admit", "seed.synchronize"] |
| assert events[1] == ("seed.admit", (22,), (1,)) |
| assert events[2] == ("seed.synchronize", (0, 1), False) |
|
|
|
|
| def test_decode_new_slot_without_remap_remains_strict_and_does_not_leak(expect_error): |
| controller, _, _, events = _make_controller() |
| state = controller.create_state() |
| initial = replace( |
| _prepared(penalties=False), |
| active_mask=(True, False, False, False), |
| active_rows=1, |
| row_paths=("topk", "inactive", "inactive", "inactive"), |
| seeds=(11, None, None, None), |
| ) |
| controller.admit(state, initial) |
| before = (tuple(state.seed_state.active), tuple(state.seed_state.seeds)) |
| events.clear() |
| expanded = _prepared(penalties=False, seeds=(11, 22, None, None)) |
|
|
| with expect_error(RuntimeError, "reset_batch=True"): |
| controller.synchronize_decode(state, expanded, reset_batch=False) |
|
|
| assert (tuple(state.seed_state.active), tuple(state.seed_state.seeds)) == before |
| assert events == [("seed.synchronize", (0, 1), False)] |
|
|
|
|
| def test_decode_nonidentity_remap_changed_survivor_rejects_before_mutation(expect_error): |
| controller, _, _, events = _make_controller() |
| state = controller.create_state() |
| initial = replace( |
| _prepared(penalties=False), |
| active_mask=(True, False, False, False), |
| active_rows=1, |
| row_paths=("topk", "inactive", "inactive", "inactive"), |
| seeds=(11, None, None, None), |
| ) |
| controller.admit(state, initial) |
| before = (tuple(state.seed_state.active), tuple(state.seed_state.seeds)) |
| events.clear() |
| changed = _prepared( |
| penalties=False, |
| seeds=(22, 99, None, None), |
| slot_remap=(1, 0, 2, 3), |
| ) |
|
|
| with expect_error(RuntimeError, "changed active seed slots"): |
| controller.synchronize_decode(state, changed, reset_batch=False) |
|
|
| assert (tuple(state.seed_state.active), tuple(state.seed_state.seeds)) == before |
| assert events == [] |
|
|
|
|
| def test_decode_remap_validates_new_penalty_history_before_mutation(expect_error): |
| controller, _, _, events = _make_controller() |
| state = controller.create_state() |
| initial = replace( |
| _prepared(penalties=False), |
| active_mask=(True, False, False, False), |
| active_rows=1, |
| row_paths=("topk", "inactive", "inactive", "inactive"), |
| seeds=(11, None, None, None), |
| ) |
| controller.admit(state, initial) |
| before = (tuple(state.seed_state.active), tuple(state.seed_state.seeds)) |
| events.clear() |
| incomplete = replace( |
| _prepared(seeds=(22, 11, None, None), slot_remap=(1, 0, 2, 3)), |
| prompt_tokens=None, |
| ) |
|
|
| with expect_error(ValueError, "prompt_tokens are required"): |
| controller.synchronize_decode(state, incomplete, reset_batch=False) |
|
|
| assert (tuple(state.seed_state.active), tuple(state.seed_state.seeds)) == before |
| assert events == [] |
|
|
|
|
| def test_decode_slot_remap_moves_survivor_then_admits_new_slot(): |
| controller, _, _, events = _make_controller() |
| state = controller.create_state() |
| initial = replace( |
| _prepared(penalties=False), |
| active_mask=(True, False, False, False), |
| active_rows=1, |
| row_paths=("topk", "inactive", "inactive", "inactive"), |
| seeds=(11, None, None, None), |
| ) |
| controller.admit(state, initial) |
| events.clear() |
| remapped = _prepared( |
| penalties=False, |
| seeds=(22, 11, None, None), |
| slot_remap=(1, 0, 2, 3), |
| ) |
|
|
| consumed = controller.synchronize_decode(state, remapped, reset_batch=False) |
|
|
| assert consumed.slot_remap is None |
| assert state.seed_state.active == [True, True, False, False] |
| assert state.seed_state.seeds == [22, 11, None, None] |
| assert _event_names(events)[:3] == ["seed.remap", "seed.admit", "seed.synchronize"] |
| assert events[1] == ("seed.admit", (22,), (0,)) |
|
|
|
|
| def test_penalized_prefill_gap_decode_rebuilds_history_on_batch_reset(): |
| controller, _, _, events = _make_controller() |
| state = controller.create_state() |
| prepared = replace( |
| _prepared(), |
| active_mask=(True, False, False, False), |
| active_rows=1, |
| row_paths=("topk", "inactive", "inactive", "inactive"), |
| seeds=(77, None, None, None), |
| ) |
| controller.admit_prefill(state, prepared, slots=(3,), positions=(12,)) |
| controller.prefill_forward( |
| "logits", |
| state, |
| prepared, |
| k=object(), |
| p=object(), |
| temp=object(), |
| ) |
| decode_prepared = place_prepared_sampling_params( |
| replace(prepared, output_tokens=torch.tensor([[7, 8, 99], [9, -1, -1]])), |
| (3,), |
| ) |
| events.clear() |
|
|
| controller.synchronize_decode(state, decode_prepared, reset_batch=True) |
| controller.refresh_dynamic_inputs( |
| state, |
| decode_prepared, |
| positions=(-1, -1, -1, 13), |
| ) |
|
|
| assert state.penalty_history_valid is True |
| assert state.active_slots == (3,) |
| assert _event_names(events) == [ |
| "seed.synchronize", |
| "penalty.prompt", |
| "penalty.output_reset", |
| "seed.synchronize", |
| "seed.refresh", |
| ] |
|
|
|
|
| def test_refresh_rejects_static_identity_change_until_trace_is_reselected(expect_error): |
| controller, _, _, _ = _make_controller() |
| state = controller.create_state() |
| controller.admit(state, _prepared()) |
|
|
| with expect_error(RuntimeError, "static identity changed"): |
| controller.refresh_dynamic_inputs(state, _prepared(sampling_path="argmax")) |
|
|
|
|
| def test_slot_remap_moves_seed_state_and_rebuilds_penalty_history(): |
| controller, _, _, events = _make_controller() |
| state = controller.create_state() |
| controller.admit(state, _prepared()) |
| remapped = _prepared( |
| slot_remap=(1, 0, 2, 3), |
| seeds=(None, 11, None, None), |
| ) |
| events.clear() |
|
|
| consumed = controller.apply_slot_remap(state, remapped) |
|
|
| assert _event_names(events) == [ |
| "seed.remap", |
| "seed.synchronize", |
| "penalty.prompt", |
| "penalty.output_reset", |
| ] |
| assert consumed.slot_remap is None |
| events.clear() |
| controller.refresh_dynamic_inputs(state, consumed, positions=(1, 1, -1, -1)) |
| assert _event_names(events) == ["seed.synchronize", "seed.refresh"] |
|
|
|
|
| def test_before_after_sampling_enforces_penalty_order_and_exactly_once_update(expect_error): |
| controller, _, _, events = _make_controller() |
| state = controller.create_state() |
| prepared = _prepared() |
| controller.admit(state, prepared) |
| events.clear() |
|
|
| logits, sample_id = controller.before_sampling("logits", state, prepared, positions=(0, 0, -1, -1)) |
| assert logits == "penalized:logits" |
| controller.after_sampling(state, "token", sample_id=sample_id) |
|
|
| assert _event_names(events) == [ |
| "seed.synchronize", |
| "seed.refresh", |
| "penalty.decode", |
| "penalty.update", |
| ] |
| with expect_error(RuntimeError, "not the pending"): |
| controller.after_sampling(state, "token", sample_id=sample_id) |
|
|
|
|
| def test_atomic_decode_composes_penalties_sampling_seeds_and_history(): |
| controller, sampling, _, events = _make_controller() |
| state = controller.create_state() |
| prepared = _prepared(log_probs=True) |
| controller.admit(state, prepared) |
| events.clear() |
| k, p, temp = object(), object(), object() |
|
|
| output = controller.decode_forward( |
| "logits", |
| state, |
| prepared, |
| k=k, |
| p=p, |
| temp=temp, |
| positions=(2, 2, -1, -1), |
| ) |
|
|
| assert output == ("sampled-tokens", "sampled-logprobs") |
| assert _event_names(events) == [ |
| "seed.synchronize", |
| "seed.refresh", |
| "penalty.decode", |
| "seed.handle", |
| "sampling.decode", |
| "penalty.update", |
| ] |
| kwargs = events[-2][2] |
| assert kwargs["k"] is k and kwargs["p"] is p and kwargs["temp"] is temp |
| assert kwargs["seeds"] is sampling.config.seeds.handle |
| assert kwargs["enable_log_probs"] == [True, False, False, False] |
| assert state.pending_sample_id is None |
|
|
|
|
| def test_sampling_failure_cancels_pending_step_without_updating_history(expect_error): |
| controller, sampling, _, events = _make_controller() |
| state = controller.create_state() |
| prepared = _prepared() |
| controller.admit(state, prepared) |
| sampling.raise_on_decode = True |
| events.clear() |
|
|
| with expect_error(RuntimeError, "sampling failed"): |
| controller.decode_forward( |
| "logits", |
| state, |
| prepared, |
| k=object(), |
| p=object(), |
| temp=object(), |
| positions=(2, 2, -1, -1), |
| ) |
|
|
| assert state.pending_sample_id is None |
| assert "penalty.update" not in _event_names(events) |
|
|
|
|
| def test_compile_only_decode_does_not_count_a_phantom_output_token(): |
| controller, _, penalties, events = _make_controller() |
| state = controller.create_state() |
| prepared = _prepared() |
| controller.admit(state, prepared) |
| penalty_param_updates = len(penalties.config.presence_penalties.updates) |
| events.clear() |
|
|
| controller.decode_forward( |
| "logits", |
| state, |
| prepared, |
| k=object(), |
| p=object(), |
| temp=object(), |
| count_tokens=False, |
| advance_seeds=False, |
| ) |
|
|
| assert "sampling.decode" in _event_names(events) |
| assert "seed.restore_defaults" not in _event_names(events) |
| assert "seed.refresh" not in _event_names(events) |
| assert "seed.synchronize" not in _event_names(events) |
| assert "penalty.update" not in _event_names(events) |
| assert len(penalties.config.presence_penalties.updates) == penalty_param_updates |
| assert state.pending_sample_id is None |
|
|
|
|
| def test_trace_capture_body_records_penalty_update_without_capturing_dynamic_writes(): |
| controller, _, penalties, events = _make_controller() |
| state = controller.create_state() |
| prepared = _prepared() |
| controller.admit(state, prepared) |
| penalty_param_updates = len(penalties.config.presence_penalties.updates) |
| events.clear() |
|
|
| controller.decode_forward( |
| "logits", |
| state, |
| prepared, |
| k=object(), |
| p=object(), |
| temp=object(), |
| count_tokens=True, |
| advance_seeds=False, |
| ) |
|
|
| assert _event_names(events) == [ |
| "penalty.decode", |
| "seed.handle", |
| "sampling.decode", |
| "penalty.update", |
| ] |
| assert len(penalties.config.presence_penalties.updates) == penalty_param_updates |
|
|
|
|
| def test_argmax_restores_default_seeds_and_does_not_require_kpt(): |
| controller, _, _, events = _make_controller() |
| state = controller.create_state() |
| prepared = _prepared(penalties=False, sampling_path="argmax") |
| controller.admit(state, prepared) |
| events.clear() |
|
|
| controller.decode_forward("logits", state, prepared) |
|
|
| assert "seed.restore_defaults" in _event_names(events) |
| assert "seed.handle" not in _event_names(events) |
| assert "penalty.decode" not in _event_names(events) |
| assert "penalty.update" not in _event_names(events) |
|
|
|
|
| def test_cleanup_validates_before_mutating_and_rebuilds_live_history(expect_error): |
| controller, _, _, events = _make_controller() |
| state = controller.create_state() |
| prepared = _prepared() |
| controller.admit(state, prepared) |
| events.clear() |
|
|
| with expect_error(ValueError, "prepared sampling state is required"): |
| controller.cleanup(state, (0, 1)) |
| assert events == [] |
|
|
| controller.cleanup(state, (0, 1), prepared=prepared) |
| assert _event_names(events) == ["seed.cleanup", "penalty.prompt", "penalty.output_reset"] |
|
|
|
|
| def test_release_resets_caller_state_and_releases_only_owned_penalties(expect_error): |
| controller, sampling, penalties, events = _make_controller() |
| state = controller.create_state() |
| controller.admit(state, _prepared()) |
| events.clear() |
|
|
| controller.release(state) |
|
|
| assert state.released |
| assert penalties.released |
| assert not sampling.released |
| assert _event_names(events) == ["seed.reset", "penalty.release"] |
| with expect_error(RuntimeError, "released"): |
| controller.reset(state) |
| controller.release(state) |
|
|
|
|
| def test_penalties_release_deallocates_owned_lazy_buffers_and_slice_tensors(monkeypatch): |
| released = [] |
| monkeypatch.setattr(penalties_module.ttnn, "deallocate", released.append) |
| names = FakePenalties._BUFFER_NAMES |
| buffers = {} |
| values = [] |
| for name in names: |
| buffer = LazyBuffer(source=torch.zeros(1)) |
| value = object() |
| buffer._value = value |
| buffers[name] = buffer |
| values.append(value) |
|
|
| penalties = object.__new__(Penalties1D) |
| penalties.config = SimpleNamespace(**buffers) |
| penalties._slice_start = object() |
| penalties._slice_end = object() |
| values.extend((penalties._slice_start, penalties._slice_end)) |
| penalties._decode_src = buffers["decode_src"]._value |
| penalties._zeros = buffers["zeros"]._value |
| penalties._device_buffers_loaded = True |
|
|
| penalties.release() |
| penalties.release() |
|
|
| assert released == values |
| assert all(buffer._value is None for buffer in buffers.values()) |
| assert penalties._slice_start is None and penalties._slice_end is None |
| assert penalties._decode_src is None and penalties._zeros is None |
| assert not penalties._device_buffers_loaded |
|
|
